Compare commits
6 Commits
fdeb881c83
...
cursor/fix
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c38e36f685 | ||
|
|
50e8904f1c | ||
|
|
7e289fef2e | ||
|
|
e9a094b620 | ||
|
|
1299a6bb1c | ||
|
|
f220fd2210 |
1
.gitignore
vendored
1
.gitignore
vendored
@@ -19,3 +19,4 @@ dist/
|
|||||||
!.env.testnet
|
!.env.testnet
|
||||||
.rocm-local/*
|
.rocm-local/*
|
||||||
billing.sqlite
|
billing.sqlite
|
||||||
|
.pytest-tmp/*
|
||||||
BIN
billing.sqlite
BIN
billing.sqlite
Binary file not shown.
Binary file not shown.
@@ -52,7 +52,7 @@ def _run_node(cfg: dict) -> None:
|
|||||||
node = run_startup(
|
node = run_startup(
|
||||||
tracker_url=cfg["tracker_url"],
|
tracker_url=cfg["tracker_url"],
|
||||||
port=cfg.get("port", 7000),
|
port=cfg.get("port", 7000),
|
||||||
model=cfg.get("model_name") or "stub-model",
|
model=cfg.get("model_name") or None,
|
||||||
model_id=cfg.get("model_hf_repo") or None,
|
model_id=cfg.get("model_hf_repo") or None,
|
||||||
shard_start=cfg.get("shard_start"),
|
shard_start=cfg.get("shard_start"),
|
||||||
shard_end=cfg.get("shard_end"),
|
shard_end=cfg.get("shard_end"),
|
||||||
@@ -90,6 +90,19 @@ def _run_node(cfg: dict) -> None:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _resolve_model_flags(
|
||||||
|
model: str | None,
|
||||||
|
model_id: str | None,
|
||||||
|
) -> tuple[str | None, str | None]:
|
||||||
|
"""Return (model_name, hf_repo_or_none) from --model / --model-id flags."""
|
||||||
|
explicit = model_id or model
|
||||||
|
if not explicit:
|
||||||
|
return None, None
|
||||||
|
if "/" in explicit:
|
||||||
|
return explicit.split("/")[-1], explicit
|
||||||
|
return explicit, None
|
||||||
|
|
||||||
|
|
||||||
def _first_available_port(host: str, start: int = 7000, attempts: int = 100) -> int:
|
def _first_available_port(host: str, start: int = 7000, attempts: int = 100) -> int:
|
||||||
"""Return the first TCP port bindable on host, starting at start."""
|
"""Return the first TCP port bindable on host, starting at start."""
|
||||||
bind_host = "" if host == "0.0.0.0" else host
|
bind_host = "" if host == "0.0.0.0" else host
|
||||||
@@ -122,9 +135,10 @@ def _cmd_default(args) -> int:
|
|||||||
|
|
||||||
# Apply CLI overrides on top of saved config
|
# Apply CLI overrides on top of saved config
|
||||||
overrides: dict = {}
|
overrides: dict = {}
|
||||||
if args.model:
|
model_name, hf_repo = _resolve_model_flags(args.model, getattr(args, "model_id", None))
|
||||||
overrides["model_hf_repo"] = args.model
|
if model_name is not None:
|
||||||
overrides["model_name"] = args.model.split("/")[-1]
|
overrides["model_name"] = model_name
|
||||||
|
overrides["model_hf_repo"] = hf_repo or ""
|
||||||
if args.quantization:
|
if args.quantization:
|
||||||
overrides["quantization"] = args.quantization
|
overrides["quantization"] = args.quantization
|
||||||
if args.download_dir:
|
if args.download_dir:
|
||||||
@@ -215,16 +229,15 @@ def _cmd_start(args) -> int:
|
|||||||
if args.tracker:
|
if args.tracker:
|
||||||
cfg["tracker_url"] = args.tracker
|
cfg["tracker_url"] = args.tracker
|
||||||
cfg["port"] = args.port if args.port is not None else _first_available_port(args.host)
|
cfg["port"] = args.port if args.port is not None else _first_available_port(args.host)
|
||||||
model = args.model or cfg.get("model_hf_repo") or cfg.get("model_name") or "stub-model"
|
model_name, hf_repo = _resolve_model_flags(
|
||||||
if args.model_id is None and "/" in model:
|
args.model or cfg.get("model_hf_repo") or cfg.get("model_name") or None,
|
||||||
cfg["model_hf_repo"] = model
|
args.model_id,
|
||||||
cfg["model_name"] = model.split("/")[-1]
|
)
|
||||||
else:
|
if model_name is not None:
|
||||||
cfg["model_name"] = model
|
cfg["model_name"] = model_name
|
||||||
|
cfg["model_hf_repo"] = hf_repo or ""
|
||||||
cfg["quantization"] = args.quantization
|
cfg["quantization"] = args.quantization
|
||||||
cfg["host"] = args.host
|
cfg["host"] = args.host
|
||||||
if args.model_id:
|
|
||||||
cfg["model_hf_repo"] = args.model_id
|
|
||||||
if args.shard_start is not None:
|
if args.shard_start is not None:
|
||||||
cfg["shard_start"] = args.shard_start
|
cfg["shard_start"] = args.shard_start
|
||||||
if args.shard_end is not None:
|
if args.shard_end is not None:
|
||||||
@@ -242,7 +255,7 @@ def _cmd_start(args) -> int:
|
|||||||
tracker_url=cfg["tracker_url"],
|
tracker_url=cfg["tracker_url"],
|
||||||
port=cfg["port"],
|
port=cfg["port"],
|
||||||
model=cfg["model_name"],
|
model=cfg["model_name"],
|
||||||
model_id=cfg.get("model_hf_repo"),
|
model_id=cfg.get("model_hf_repo") or None,
|
||||||
shard_start=cfg.get("shard_start"),
|
shard_start=cfg.get("shard_start"),
|
||||||
shard_end=cfg.get("shard_end"),
|
shard_end=cfg.get("shard_end"),
|
||||||
quantization=cfg["quantization"].replace("bf16", "bfloat16"),
|
quantization=cfg["quantization"].replace("bf16", "bfloat16"),
|
||||||
@@ -288,7 +301,8 @@ def main() -> None:
|
|||||||
)
|
)
|
||||||
|
|
||||||
# Flags that apply to the no-subcommand (default) path
|
# Flags that apply to the no-subcommand (default) path
|
||||||
parser.add_argument("--model", metavar="HF_REPO", help="HuggingFace repo ID to serve")
|
parser.add_argument("--model", metavar="MODEL", help="Model name or HuggingFace repo ID to serve")
|
||||||
|
parser.add_argument("--model-id", metavar="MODEL", help="Alias for --model (catalog name or HuggingFace repo)")
|
||||||
parser.add_argument("--quantization", "-q", choices=["bf16", "int8", "nf4", "bfloat16"],
|
parser.add_argument("--quantization", "-q", choices=["bf16", "int8", "nf4", "bfloat16"],
|
||||||
help="Quantization level")
|
help="Quantization level")
|
||||||
parser.add_argument("--download-dir", metavar="PATH", help="Model download directory")
|
parser.add_argument("--download-dir", metavar="PATH", help="Model download directory")
|
||||||
@@ -329,8 +343,8 @@ def main() -> None:
|
|||||||
start_cmd = subparsers.add_parser("start", help="Start node (legacy flags)")
|
start_cmd = subparsers.add_parser("start", help="Start node (legacy flags)")
|
||||||
start_cmd.add_argument("--tracker")
|
start_cmd.add_argument("--tracker")
|
||||||
start_cmd.add_argument("--port", type=int)
|
start_cmd.add_argument("--port", type=int)
|
||||||
start_cmd.add_argument("--model")
|
start_cmd.add_argument("--model", help="Model name or HuggingFace repo ID")
|
||||||
start_cmd.add_argument("--model-id", help="HuggingFace repo ID")
|
start_cmd.add_argument("--model-id", help="Alias for --model (catalog name or HuggingFace repo)")
|
||||||
start_cmd.add_argument("--shard-start", type=int)
|
start_cmd.add_argument("--shard-start", type=int)
|
||||||
start_cmd.add_argument("--shard-end", type=int)
|
start_cmd.add_argument("--shard-end", type=int)
|
||||||
start_cmd.add_argument("--quantization", choices=["auto", "bfloat16", "int8", "nf4", "bf16"], default="auto")
|
start_cmd.add_argument("--quantization", choices=["auto", "bfloat16", "int8", "nf4", "bf16"], default="auto")
|
||||||
|
|||||||
@@ -331,6 +331,9 @@ def _attach_relay_bridge(node: StubNodeServer | TorchNodeServer, bridge: RelayHt
|
|||||||
node.stop = _stop_with_bridge # type: ignore[method-assign]
|
node.stop = _stop_with_bridge # type: ignore[method-assign]
|
||||||
|
|
||||||
|
|
||||||
|
_PENDING_NODE_ID = "pending"
|
||||||
|
|
||||||
|
|
||||||
def _start_heartbeat(
|
def _start_heartbeat(
|
||||||
tracker_url: str,
|
tracker_url: str,
|
||||||
node_id: str,
|
node_id: str,
|
||||||
@@ -368,10 +371,33 @@ def _start_heartbeat(
|
|||||||
try:
|
try:
|
||||||
resp = _post_json(f"{tracker_url}/v1/nodes/register", register_payload)
|
resp = _post_json(f"{tracker_url}/v1/nodes/register", register_payload)
|
||||||
node_id = resp.get("node_id", node_id)
|
node_id = resp.get("node_id", node_id)
|
||||||
|
if node_ref is not None:
|
||||||
|
setattr(node_ref, "tracker_node_id", node_id)
|
||||||
return True
|
return True
|
||||||
except Exception:
|
except Exception:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
def _register_additional_assignment(applied: dict) -> None:
|
||||||
|
model_id = str(applied.get("model") or register_payload.get("hf_repo") or register_payload.get("model"))
|
||||||
|
extra_payload = {
|
||||||
|
**register_payload,
|
||||||
|
"model": model_id.split("/")[-1],
|
||||||
|
"hf_repo": model_id if "/" in model_id else register_payload.get("hf_repo"),
|
||||||
|
"shard_start": applied["shard_start"],
|
||||||
|
"shard_end": applied["shard_end"],
|
||||||
|
"quantization": applied.get("quantization", register_payload.get("quantization")),
|
||||||
|
"tracker_mode": bool(applied.get("tracker_mode", False)),
|
||||||
|
"managed_assignment": True,
|
||||||
|
}
|
||||||
|
try:
|
||||||
|
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", extra_payload)
|
||||||
|
print(
|
||||||
|
f" [node] registered additional model — node ID: {reg_resp.get('node_id')}",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
|
except Exception as exc:
|
||||||
|
print(f" [node] WARNING: additional model registration failed: {exc}", flush=True)
|
||||||
|
|
||||||
def _apply_directives(directives: list[dict]) -> None:
|
def _apply_directives(directives: list[dict]) -> None:
|
||||||
if not directives:
|
if not directives:
|
||||||
return
|
return
|
||||||
@@ -384,6 +410,9 @@ def _start_heartbeat(
|
|||||||
print(f" [node] WARNING: failed to apply tracker directives: {exc}", flush=True)
|
print(f" [node] WARNING: failed to apply tracker directives: {exc}", flush=True)
|
||||||
return
|
return
|
||||||
if applied:
|
if applied:
|
||||||
|
if applied.get("action") == "ADD_SHARD":
|
||||||
|
_register_additional_assignment(applied)
|
||||||
|
return
|
||||||
model_id = applied.get("model", register_payload.get("hf_repo") or register_payload.get("model"))
|
model_id = applied.get("model", register_payload.get("hf_repo") or register_payload.get("model"))
|
||||||
register_payload["model"] = str(model_id).split("/")[-1]
|
register_payload["model"] = str(model_id).split("/")[-1]
|
||||||
register_payload["hf_repo"] = model_id
|
register_payload["hf_repo"] = model_id
|
||||||
@@ -395,7 +424,7 @@ def _start_heartbeat(
|
|||||||
def _loop() -> None:
|
def _loop() -> None:
|
||||||
nonlocal node_id
|
nonlocal node_id
|
||||||
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
||||||
outage_streak = 0 # consecutive intervals where tracker was unreachable
|
outage_streak = 1 if node_id == _PENDING_NODE_ID else 0
|
||||||
|
|
||||||
while True:
|
while True:
|
||||||
time.sleep(interval)
|
time.sleep(interval)
|
||||||
@@ -423,11 +452,12 @@ def _start_heartbeat(
|
|||||||
new_asgn = resp.get("new_assignment")
|
new_asgn = resp.get("new_assignment")
|
||||||
if new_asgn:
|
if new_asgn:
|
||||||
print(
|
print(
|
||||||
f" [node] tracker reassignment received: "
|
f" [node] tracker assignment received: "
|
||||||
f"model={new_asgn.get('model')!r} "
|
f"action={new_asgn.get('action')!r} model={new_asgn.get('model')!r} "
|
||||||
f"shards={new_asgn.get('shard_start')}-{new_asgn.get('shard_end')}",
|
f"shards={new_asgn.get('shard_start')}-{new_asgn.get('shard_end')}",
|
||||||
flush=True,
|
flush=True,
|
||||||
)
|
)
|
||||||
|
_apply_directives([new_asgn])
|
||||||
except urllib.error.HTTPError as exc:
|
except urllib.error.HTTPError as exc:
|
||||||
if exc.code == 404:
|
if exc.code == 404:
|
||||||
# Node was purged (e.g. long gap before restart noticed) — re-register now.
|
# Node was purged (e.g. long gap before restart noticed) — re-register now.
|
||||||
@@ -449,6 +479,34 @@ def _start_heartbeat(
|
|||||||
return t
|
return t
|
||||||
|
|
||||||
|
|
||||||
|
def _register_with_tracker(
|
||||||
|
tracker_url: str,
|
||||||
|
reg_payload: dict,
|
||||||
|
node: Any,
|
||||||
|
start_time: float,
|
||||||
|
) -> str | None:
|
||||||
|
"""Register with the tracker, or start background retries when it is unreachable."""
|
||||||
|
try:
|
||||||
|
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload)
|
||||||
|
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
||||||
|
setattr(node, "tracker_node_id", tracker_node_id)
|
||||||
|
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
||||||
|
_start_heartbeat(tracker_url, tracker_node_id, reg_payload, node_ref=node, start_time=start_time)
|
||||||
|
return tracker_node_id
|
||||||
|
except Exception as exc:
|
||||||
|
setattr(node, "tracker_node_id", None)
|
||||||
|
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||||||
|
print(" [node] will retry registration in the background", flush=True)
|
||||||
|
_start_heartbeat(
|
||||||
|
tracker_url,
|
||||||
|
_PENDING_NODE_ID,
|
||||||
|
reg_payload,
|
||||||
|
node_ref=node,
|
||||||
|
start_time=start_time,
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
def _warn_virtual_network_ip(ip: str | None) -> None:
|
def _warn_virtual_network_ip(ip: str | None) -> None:
|
||||||
"""Print a warning when the auto-detected advertise IP is in a known virtual-network range.
|
"""Print a warning when the auto-detected advertise IP is in a known virtual-network range.
|
||||||
|
|
||||||
@@ -482,7 +540,7 @@ def _warn_virtual_network_ip(ip: str | None) -> None:
|
|||||||
def run_startup(
|
def run_startup(
|
||||||
tracker_url: str,
|
tracker_url: str,
|
||||||
port: int = 0,
|
port: int = 0,
|
||||||
model: str = "stub-model",
|
model: str | None = None,
|
||||||
model_id: str | None = None,
|
model_id: str | None = None,
|
||||||
shard_start: int | None = None,
|
shard_start: int | None = None,
|
||||||
shard_end: int | None = None,
|
shard_end: int | None = None,
|
||||||
@@ -608,8 +666,11 @@ def run_startup(
|
|||||||
if probationary_line is not None:
|
if probationary_line is not None:
|
||||||
print(f" {probationary_line}", flush=True)
|
print(f" {probationary_line}", flush=True)
|
||||||
|
|
||||||
|
pinned_shard_start = shard_start
|
||||||
|
pinned_shard_end = shard_end
|
||||||
|
user_pinned_shard = pinned_shard_start is not None or pinned_shard_end is not None
|
||||||
|
|
||||||
if model_id: # treat "" the same as None — no explicit model given
|
if model_id: # treat "" the same as None — no explicit model given
|
||||||
user_pinned_shard = shard_start is not None or shard_end is not None
|
|
||||||
full_sources: list[dict] = []
|
full_sources: list[dict] = []
|
||||||
# Auto-detect shard range from model config if not explicitly provided
|
# Auto-detect shard range from model config if not explicitly provided
|
||||||
if shard_start is None or shard_end is None:
|
if shard_start is None or shard_end is None:
|
||||||
@@ -668,6 +729,7 @@ def run_startup(
|
|||||||
route_timeout=route_timeout,
|
route_timeout=route_timeout,
|
||||||
cache_dir=cache_dir,
|
cache_dir=cache_dir,
|
||||||
debug=debug,
|
debug=debug,
|
||||||
|
max_loaded_shards=max_loaded_shards,
|
||||||
)
|
)
|
||||||
_node_start_time = time.monotonic()
|
_node_start_time = time.monotonic()
|
||||||
actual_port = node.start()
|
actual_port = node.start()
|
||||||
@@ -720,16 +782,9 @@ def run_startup(
|
|||||||
**registration_capabilities,
|
**registration_capabilities,
|
||||||
**relay_fields,
|
**relay_fields,
|
||||||
}
|
}
|
||||||
tracker_node_id: str | None = None
|
tracker_node_id = _register_with_tracker(
|
||||||
try:
|
tracker_url, reg_payload, node, _node_start_time,
|
||||||
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload)
|
)
|
||||||
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
|
||||||
setattr(node, "tracker_node_id", tracker_node_id)
|
|
||||||
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
|
||||||
_start_heartbeat(tracker_url, tracker_node_id, reg_payload, node_ref=node, start_time=_node_start_time)
|
|
||||||
except Exception as exc:
|
|
||||||
setattr(node, "tracker_node_id", None)
|
|
||||||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
|
||||||
|
|
||||||
print(
|
print(
|
||||||
f"\n{'=' * 32}\n"
|
f"\n{'=' * 32}\n"
|
||||||
@@ -747,16 +802,17 @@ def run_startup(
|
|||||||
flush=True,
|
flush=True,
|
||||||
)
|
)
|
||||||
return node
|
return node
|
||||||
if shard_start is not None or shard_end is not None:
|
if user_pinned_shard and not model:
|
||||||
raise ValueError("--shard-start / --shard-end require --model-id")
|
raise ValueError("--shard-start / --shard-end require --model")
|
||||||
|
|
||||||
# 3a. Auto-join: query tracker for network-wide HF model assignment.
|
# 3a. Auto-join: query tracker for network-wide HF model assignment.
|
||||||
# Skipped when the user explicitly requested a model — the shard-assignment
|
# Skipped when the user explicitly requested a model — the shard-assignment
|
||||||
# query below (/v1/nodes/assign?model=…) is authoritative there, and a fresh
|
# query below (/v1/nodes/assign?model=…) is authoritative there, and a fresh
|
||||||
# tracker would otherwise print a scary 503 for the model-less auto-join.
|
# tracker would otherwise print a scary 503 for the model-less auto-join.
|
||||||
net_assignment: dict = {}
|
net_assignment: dict = {}
|
||||||
if model and model != "stub-model":
|
if model_id or (model and model != "stub-model"):
|
||||||
print(f"Model {model!r} requested explicitly — skipping network auto-join.", flush=True)
|
if model:
|
||||||
|
print(f"Model {model!r} requested explicitly — skipping network auto-join.", flush=True)
|
||||||
else:
|
else:
|
||||||
print("Querying tracker for network assignment...", flush=True)
|
print("Querying tracker for network assignment...", flush=True)
|
||||||
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb})
|
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": assignment_vram_mb, "ram_mb": ram_mb})
|
||||||
@@ -767,17 +823,25 @@ def run_startup(
|
|||||||
assigned_hf_repo: str | None = net_assignment.get("hf_repo")
|
assigned_hf_repo: str | None = net_assignment.get("hf_repo")
|
||||||
_gap_found: bool = bool(net_assignment.get("gap_found", False))
|
_gap_found: bool = bool(net_assignment.get("gap_found", False))
|
||||||
|
|
||||||
if assigned_hf_repo and _gap_found:
|
if assigned_hf_repo:
|
||||||
assigned_shard_start: int = net_assignment["shard_start"]
|
assigned_shard_start: int = net_assignment["shard_start"]
|
||||||
assigned_shard_end: int = net_assignment["shard_end"]
|
assigned_shard_end: int = net_assignment["shard_end"]
|
||||||
assigned_num_layers: int = net_assignment["num_layers"]
|
assigned_num_layers: int = net_assignment["num_layers"]
|
||||||
assigned_model_sources: list[dict] = net_assignment.get("model_sources", [])
|
assigned_model_sources: list[dict] = net_assignment.get("model_sources", [])
|
||||||
print(
|
if _gap_found:
|
||||||
f" Assigned: {assigned_hf_repo} "
|
print(
|
||||||
f"layers {assigned_shard_start}–{assigned_shard_end} "
|
f" Assigned gap: {assigned_hf_repo} "
|
||||||
f"(of {assigned_num_layers})",
|
f"layers {assigned_shard_start}–{assigned_shard_end} "
|
||||||
flush=True,
|
f"(of {assigned_num_layers})",
|
||||||
)
|
flush=True,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
print(
|
||||||
|
f" Assigned redundant copy: {assigned_hf_repo} "
|
||||||
|
f"layers {assigned_shard_start}–{assigned_shard_end} "
|
||||||
|
f"(of {assigned_num_layers})",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
full_sources = [] if tracker_source_disabled else _full_model_sources(assigned_model_sources)
|
full_sources = [] if tracker_source_disabled else _full_model_sources(assigned_model_sources)
|
||||||
if full_sources:
|
if full_sources:
|
||||||
print("Downloading assigned model snapshot...", flush=True)
|
print("Downloading assigned model snapshot...", flush=True)
|
||||||
@@ -801,6 +865,7 @@ def run_startup(
|
|||||||
route_timeout=route_timeout,
|
route_timeout=route_timeout,
|
||||||
cache_dir=cache_dir,
|
cache_dir=cache_dir,
|
||||||
debug=debug,
|
debug=debug,
|
||||||
|
max_loaded_shards=max_loaded_shards,
|
||||||
)
|
)
|
||||||
_node_start_time = time.monotonic()
|
_node_start_time = time.monotonic()
|
||||||
actual_port = node.start()
|
actual_port = node.start()
|
||||||
@@ -845,16 +910,9 @@ def run_startup(
|
|||||||
**registration_capabilities,
|
**registration_capabilities,
|
||||||
**relay_fields,
|
**relay_fields,
|
||||||
}
|
}
|
||||||
tracker_node_id = None
|
tracker_node_id = _register_with_tracker(
|
||||||
try:
|
tracker_url, auto_reg_payload, node, _node_start_time,
|
||||||
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", auto_reg_payload)
|
)
|
||||||
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
|
||||||
setattr(node, "tracker_node_id", tracker_node_id)
|
|
||||||
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
|
||||||
_start_heartbeat(tracker_url, tracker_node_id, auto_reg_payload, node_ref=node, start_time=_node_start_time)
|
|
||||||
except Exception as exc:
|
|
||||||
setattr(node, "tracker_node_id", None)
|
|
||||||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
|
||||||
shard_count = assigned_shard_end - assigned_shard_start + 1
|
shard_count = assigned_shard_end - assigned_shard_start + 1
|
||||||
print(
|
print(
|
||||||
f"\n{'=' * 32}\n"
|
f"\n{'=' * 32}\n"
|
||||||
@@ -874,10 +932,16 @@ def run_startup(
|
|||||||
)
|
)
|
||||||
return node
|
return node
|
||||||
|
|
||||||
# 3b. Shard assignment from tracker (stub-model / preset-based path)
|
if not assigned_hf_repo and model is None:
|
||||||
|
raise RuntimeError(
|
||||||
|
"Tracker did not assign a model. Join a network that already serves one, "
|
||||||
|
"or start with --model <HF_REPO>."
|
||||||
|
)
|
||||||
|
|
||||||
|
# 3b. Stub preset path (tests / explicit stub-model) or named preset models.
|
||||||
print("Querying tracker for shard assignment...", flush=True)
|
print("Querying tracker for shard assignment...", flush=True)
|
||||||
assign_qs = urllib.parse.urlencode({
|
assign_qs = urllib.parse.urlencode({
|
||||||
"model": model,
|
"model": model or "stub-model",
|
||||||
"device": device,
|
"device": device,
|
||||||
# CPU-mode nodes must be sized by RAM: a detected-but-unusable GPU's
|
# CPU-mode nodes must be sized by RAM: a detected-but-unusable GPU's
|
||||||
# VRAM would otherwise cap the shard (e.g. 8 GB VRAM → 3 layers on a
|
# VRAM would otherwise cap the shard (e.g. 8 GB VRAM → 3 layers on a
|
||||||
@@ -891,14 +955,25 @@ def run_startup(
|
|||||||
print(f" ERROR: Cannot reach tracker at {tracker_url}: {exc}", file=sys.stderr, flush=True)
|
print(f" ERROR: Cannot reach tracker at {tracker_url}: {exc}", file=sys.stderr, flush=True)
|
||||||
raise
|
raise
|
||||||
|
|
||||||
shard_start: int = assignment["shard_start"]
|
shard_start = assignment["shard_start"]
|
||||||
shard_end: int = assignment["shard_end"]
|
shard_end = assignment["shard_end"]
|
||||||
|
if user_pinned_shard:
|
||||||
|
if pinned_shard_start is not None:
|
||||||
|
shard_start = pinned_shard_start
|
||||||
|
if pinned_shard_end is not None:
|
||||||
|
shard_end = pinned_shard_end
|
||||||
assigned_model: str = assignment.get("model", model)
|
assigned_model: str = assignment.get("model", model)
|
||||||
hf_repo: str | None = assignment.get("hf_repo")
|
hf_repo: str | None = assignment.get("hf_repo")
|
||||||
peers: list[dict] = assignment.get("peers", [])
|
peers: list[dict] = assignment.get("peers", [])
|
||||||
model_sources: list[dict] = [] if tracker_source_disabled else assignment.get("model_sources", [])
|
model_sources: list[dict] = [] if tracker_source_disabled else assignment.get("model_sources", [])
|
||||||
assignment_bytes_per_layer = _assignment_bytes_per_layer(assignment, quantization)
|
assignment_bytes_per_layer = _assignment_bytes_per_layer(assignment, quantization)
|
||||||
print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
|
if user_pinned_shard:
|
||||||
|
print(
|
||||||
|
f" Shard: layers {shard_start}-{shard_end} of {assigned_model} (pinned)",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
print(f" Shard: layers {shard_start}-{shard_end} of {assigned_model}", flush=True)
|
||||||
|
|
||||||
# 4. Download shard
|
# 4. Download shard
|
||||||
print("Downloading shard...", flush=True)
|
print("Downloading shard...", flush=True)
|
||||||
@@ -960,6 +1035,7 @@ def run_startup(
|
|||||||
"hardware_profile": hw,
|
"hardware_profile": hw,
|
||||||
"wallet_address": address,
|
"wallet_address": address,
|
||||||
"score": 1.0,
|
"score": 1.0,
|
||||||
|
"managed_assignment": not user_pinned_shard,
|
||||||
**registration_capabilities,
|
**registration_capabilities,
|
||||||
**relay_fields,
|
**relay_fields,
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -75,20 +75,39 @@ class _TorchHTTPServer(http.server.HTTPServer):
|
|||||||
tracker_url: str | None = None,
|
tracker_url: str | None = None,
|
||||||
route_timeout: float = 30.0,
|
route_timeout: float = 30.0,
|
||||||
debug: bool = False,
|
debug: bool = False,
|
||||||
|
max_loaded_shards: int = 1,
|
||||||
):
|
):
|
||||||
super().__init__(addr, handler)
|
super().__init__(addr, handler)
|
||||||
self.backend = backend
|
self.backend = backend
|
||||||
|
self.backends: dict[str, TorchModelShard] = {backend.model_id: backend}
|
||||||
self.received_activations = False
|
self.received_activations = False
|
||||||
self.forward_chunk_count = 0
|
self.forward_chunk_count = 0
|
||||||
self.tracker_mode = tracker_mode
|
self.tracker_mode = tracker_mode
|
||||||
self.tracker_url = tracker_url
|
self.tracker_url = tracker_url
|
||||||
self.route_timeout = route_timeout
|
self.route_timeout = route_timeout
|
||||||
self.debug = debug
|
self.debug = debug
|
||||||
|
self.max_loaded_shards = max(1, max_loaded_shards)
|
||||||
self.total_requests: int = 0
|
self.total_requests: int = 0
|
||||||
self.failed_requests: int = 0
|
self.failed_requests: int = 0
|
||||||
self.queue_depth: int = 0
|
self.queue_depth: int = 0
|
||||||
self._stats_lock = threading.Lock()
|
self._stats_lock = threading.Lock()
|
||||||
|
|
||||||
|
def resolve_backend(self, model_name: str | None) -> TorchModelShard | None:
|
||||||
|
if not model_name:
|
||||||
|
return self.backend
|
||||||
|
wanted = model_name.strip().lower()
|
||||||
|
for key, shard_backend in self.backends.items():
|
||||||
|
key_l = key.lower()
|
||||||
|
if key_l == wanted or key_l.rsplit("/", 1)[-1] == wanted:
|
||||||
|
return shard_backend
|
||||||
|
return self.backend
|
||||||
|
|
||||||
|
def chat_enabled(self) -> bool:
|
||||||
|
return any(
|
||||||
|
shard_backend.is_head
|
||||||
|
for shard_backend in self.backends.values()
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||||
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
|
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
|
||||||
@@ -100,7 +119,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self._handle_forward()
|
self._handle_forward()
|
||||||
elif self.path == "/v1/infer":
|
elif self.path == "/v1/infer":
|
||||||
self._handle_infer()
|
self._handle_infer()
|
||||||
elif self.path == "/v1/chat/completions" and server.tracker_mode:
|
elif self.path == "/v1/chat/completions" and server.chat_enabled():
|
||||||
self._handle_chat_completions()
|
self._handle_chat_completions()
|
||||||
else:
|
else:
|
||||||
self.send_response(404)
|
self.send_response(404)
|
||||||
@@ -284,22 +303,26 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
messages = []
|
messages = []
|
||||||
stream = bool(body.get("stream", False))
|
stream = bool(body.get("stream", False))
|
||||||
model_name = str(body.get("model", ""))
|
model_name = str(body.get("model", ""))
|
||||||
|
backend = server.resolve_backend(model_name)
|
||||||
|
if backend is None or not backend.is_head:
|
||||||
|
self._send_json(400, {"error": "model not loaded on this node"})
|
||||||
|
return
|
||||||
max_tokens = int(body.get("max_tokens") or body.get("max_new_tokens") or 256)
|
max_tokens = int(body.get("max_tokens") or body.get("max_new_tokens") or 256)
|
||||||
temperature = float(body.get("temperature") or 1.0)
|
temperature = float(body.get("temperature") or 1.0)
|
||||||
top_p = float(body.get("top_p") or 1.0)
|
top_p = float(body.get("top_p") or 1.0)
|
||||||
|
|
||||||
# Fast path: this node owns the complete model — use HF generate() with KV cache.
|
# Fast path: this node owns the complete model — use HF generate() with KV cache.
|
||||||
# Avoids the single-token-per-forward-pass limitation of the distributed path.
|
# Avoids the single-token-per-forward-pass limitation of the distributed path.
|
||||||
if server.backend.is_head and server.backend.is_tail:
|
if backend.is_head and backend.is_tail:
|
||||||
try:
|
try:
|
||||||
if stream:
|
if stream:
|
||||||
self._stream_openai_response(
|
self._stream_openai_response(
|
||||||
server.backend.generate_text_streaming(messages, max_tokens, temperature, top_p),
|
backend.generate_text_streaming(messages, max_tokens, temperature, top_p),
|
||||||
model_name,
|
model_name,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
text = server.backend.generate_text(messages, max_tokens, temperature, top_p)
|
text = backend.generate_text(messages, max_tokens, temperature, top_p)
|
||||||
self._send_openai_response(text, model_name, False, messages)
|
self._send_openai_response(text, model_name, False, messages, backend=backend)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
self._record_failed_request()
|
self._record_failed_request()
|
||||||
self._send_json(500, {"error": f"generation failed: {exc}"})
|
self._send_json(500, {"error": f"generation failed: {exc}"})
|
||||||
@@ -309,7 +332,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
# We do N single-step forward passes (no cross-node KV cache), which is slow
|
# We do N single-step forward passes (no cross-node KV cache), which is slow
|
||||||
# but correct. Each step: head encodes current sequence → forwards through route
|
# but correct. Each step: head encodes current sequence → forwards through route
|
||||||
# → tail returns the next token string → append → repeat.
|
# → tail returns the next token string → append → repeat.
|
||||||
remaining_route = self._get_remaining_route(model_name)
|
remaining_route = self._get_remaining_route(model_name, backend=backend)
|
||||||
print(
|
print(
|
||||||
f" [node] chat route model={model_name!r} max_tokens={max_tokens} "
|
f" [node] chat route model={model_name!r} max_tokens={max_tokens} "
|
||||||
f"downstream={remaining_route}",
|
f"downstream={remaining_route}",
|
||||||
@@ -318,11 +341,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
if not remaining_route:
|
if not remaining_route:
|
||||||
self._send_openai_response(
|
self._send_openai_response(
|
||||||
"error: no downstream route — check tracker connectivity",
|
"error: no downstream route — check tracker connectivity",
|
||||||
model_name, False, messages,
|
model_name, False, messages, backend=backend,
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
|
|
||||||
backend = server.backend
|
|
||||||
# Format with chat template so the model knows it's in assistant mode.
|
# Format with chat template so the model knows it's in assistant mode.
|
||||||
try:
|
try:
|
||||||
if hasattr(backend.tokenizer, "apply_chat_template"):
|
if hasattr(backend.tokenizer, "apply_chat_template"):
|
||||||
@@ -352,7 +374,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
print(f" [node] distributed encode error: {exc}", flush=True)
|
print(f" [node] distributed encode error: {exc}", flush=True)
|
||||||
break
|
break
|
||||||
token_str = self._run_downstream_pipeline(payload, remaining_route)
|
token_str = self._run_downstream_pipeline(payload, remaining_route, backend=backend)
|
||||||
if not token_str:
|
if not token_str:
|
||||||
break
|
break
|
||||||
# Stop on error responses or EOS.
|
# Stop on error responses or EOS.
|
||||||
@@ -369,9 +391,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
if stream_emit is not None:
|
if stream_emit is not None:
|
||||||
stream_emit(None)
|
stream_emit(None)
|
||||||
return
|
return
|
||||||
self._send_openai_response(result_text, model_name, stream, messages)
|
self._send_openai_response(result_text, model_name, stream, messages, backend=backend)
|
||||||
|
|
||||||
def _get_remaining_route(self, model: str) -> list[dict]:
|
def _get_remaining_route(self, model: str, *, backend: TorchModelShard | None = None) -> list[dict]:
|
||||||
"""Return downstream hops as dicts with endpoint, start_layer, and optional relay_addr.
|
"""Return downstream hops as dicts with endpoint, start_layer, and optional relay_addr.
|
||||||
|
|
||||||
Fast path reads X-Meshnet-Route header injected by the tracker.
|
Fast path reads X-Meshnet-Route header injected by the tracker.
|
||||||
@@ -404,9 +426,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
|
|
||||||
# Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject).
|
# Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject).
|
||||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||||
|
active_backend = backend or server.backend
|
||||||
if server.tracker_url is None:
|
if server.tracker_url is None:
|
||||||
return []
|
return []
|
||||||
route_model = getattr(server.backend, "model_id", None) or model
|
route_model = getattr(active_backend, "model_id", None) or model
|
||||||
try:
|
try:
|
||||||
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
|
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
|
||||||
with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
|
with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
|
||||||
@@ -433,18 +456,19 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
|
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
|
||||||
return []
|
return []
|
||||||
|
|
||||||
def _run_downstream_pipeline(self, payload: object, route: list[dict]) -> str:
|
def _run_downstream_pipeline(self, payload: object, route: list[dict], *, backend: TorchModelShard | None = None) -> str:
|
||||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||||
|
active_backend = backend or server.backend
|
||||||
if not route:
|
if not route:
|
||||||
# Partial shard at tail: decode the activation from the previous node.
|
# Partial shard at tail: decode the activation from the previous node.
|
||||||
# Full single-node (head+tail) is handled before entering this method.
|
# Full single-node (head+tail) is handled before entering this method.
|
||||||
if server.backend.is_tail:
|
if active_backend.is_tail:
|
||||||
try:
|
try:
|
||||||
tensor = server.backend.torch.frombuffer(
|
tensor = active_backend.torch.frombuffer(
|
||||||
bytearray(payload.body), # type: ignore[union-attr]
|
bytearray(payload.body), # type: ignore[union-attr]
|
||||||
dtype=server.backend.torch.bfloat16,
|
dtype=active_backend.torch.bfloat16,
|
||||||
).reshape(payload.shape).to(server.backend.device) # type: ignore[union-attr]
|
).reshape(payload.shape).to(active_backend.device) # type: ignore[union-attr]
|
||||||
return server.backend.decode_tail(tensor)
|
return active_backend.decode_tail(tensor)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
return f"decode error: {exc}"
|
return f"decode error: {exc}"
|
||||||
return "no downstream route available for non-tail shard"
|
return "no downstream route available for non-tail shard"
|
||||||
@@ -591,11 +615,13 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
model: str,
|
model: str,
|
||||||
stream: bool,
|
stream: bool,
|
||||||
messages: list[dict] | None = None,
|
messages: list[dict] | None = None,
|
||||||
|
backend: TorchModelShard | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
chunk_id = "chatcmpl-node"
|
chunk_id = "chatcmpl-node"
|
||||||
created = int(time.time())
|
created = int(time.time())
|
||||||
|
active_backend = backend or self.server.backend # type: ignore[attr-defined]
|
||||||
if not stream:
|
if not stream:
|
||||||
usage = _usage_for_response(self.server.backend, messages or [], text) # type: ignore[attr-defined]
|
usage = _usage_for_response(active_backend, messages or [], text)
|
||||||
self._send_json(200, {
|
self._send_json(200, {
|
||||||
"id": chunk_id,
|
"id": chunk_id,
|
||||||
"object": "chat.completion",
|
"object": "chat.completion",
|
||||||
@@ -706,9 +732,11 @@ class TorchNodeServer:
|
|||||||
route_timeout: float = 30.0,
|
route_timeout: float = 30.0,
|
||||||
cache_dir: Path | None = None,
|
cache_dir: Path | None = None,
|
||||||
debug: bool = False,
|
debug: bool = False,
|
||||||
|
max_loaded_shards: int = 1,
|
||||||
) -> None:
|
) -> None:
|
||||||
self._host = host
|
self._host = host
|
||||||
self._requested_port = port
|
self._requested_port = port
|
||||||
|
self._max_loaded_shards = max(1, max_loaded_shards)
|
||||||
self._backend = backend or _load_backend(
|
self._backend = backend or _load_backend(
|
||||||
model_id,
|
model_id,
|
||||||
shard_start,
|
shard_start,
|
||||||
@@ -716,6 +744,7 @@ class TorchNodeServer:
|
|||||||
quantization,
|
quantization,
|
||||||
cache_dir,
|
cache_dir,
|
||||||
)
|
)
|
||||||
|
self._backends: dict[str, TorchModelShard] = {self._backend.model_id: self._backend}
|
||||||
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
|
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
|
||||||
self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0)
|
self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0)
|
||||||
self._tracker_url = tracker_url
|
self._tracker_url = tracker_url
|
||||||
@@ -754,41 +783,64 @@ class TorchNodeServer:
|
|||||||
def queue_depth(self) -> int:
|
def queue_depth(self) -> int:
|
||||||
return self._server.queue_depth if self._server is not None else 0
|
return self._server.queue_depth if self._server is not None else 0
|
||||||
|
|
||||||
|
@property
|
||||||
|
def loaded_model_ids(self) -> list[str]:
|
||||||
|
return list(self._backends.keys())
|
||||||
|
|
||||||
def apply_tracker_directives(self, directives: list[dict]) -> dict | None:
|
def apply_tracker_directives(self, directives: list[dict]) -> dict | None:
|
||||||
"""Apply tracker LOAD_SHARD directives by hot-swapping the loaded backend."""
|
"""Apply tracker shard directives (LOAD_SHARD replace, ADD_SHARD load-more)."""
|
||||||
|
add_directive = next(
|
||||||
|
(directive for directive in reversed(directives) if directive.get("action") == "ADD_SHARD"),
|
||||||
|
None,
|
||||||
|
)
|
||||||
load_directive = next(
|
load_directive = next(
|
||||||
(directive for directive in reversed(directives) if directive.get("action") == "LOAD_SHARD"),
|
(directive for directive in reversed(directives) if directive.get("action") == "LOAD_SHARD"),
|
||||||
None,
|
None,
|
||||||
)
|
)
|
||||||
if load_directive is None:
|
directive = add_directive or load_directive
|
||||||
|
if directive is None:
|
||||||
return None
|
return None
|
||||||
shard_start = int(load_directive["shard_start"])
|
shard_start = int(directive["shard_start"])
|
||||||
shard_end = int(load_directive["shard_end"])
|
shard_end = int(directive["shard_end"])
|
||||||
quantization = str(load_directive.get("quantization") or self._backend.quantization)
|
quantization = str(directive.get("quantization") or self._backend.quantization)
|
||||||
model_id = str(load_directive.get("model") or self._backend.model_id)
|
model_id = str(directive.get("model") or self._backend.model_id)
|
||||||
|
replacing = directive.get("action") == "LOAD_SHARD"
|
||||||
|
if not replacing and len(self._backends) >= self._max_loaded_shards:
|
||||||
|
print(
|
||||||
|
f" [node] WARNING: ignoring ADD_SHARD for {model_id!r} — "
|
||||||
|
f"loaded {len(self._backends)}/{self._max_loaded_shards} slots full",
|
||||||
|
flush=True,
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
action_label = "reassigned" if replacing else "additional"
|
||||||
print(
|
print(
|
||||||
f" [node] loading reassigned shard: {model_id} layers {shard_start}-{shard_end}",
|
f" [node] loading {action_label} shard: {model_id} layers {shard_start}-{shard_end}",
|
||||||
flush=True,
|
flush=True,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
new_backend = _load_backend(model_id, shard_start, shard_end, quantization, self._cache_dir)
|
new_backend = _load_backend(model_id, shard_start, shard_end, quantization, self._cache_dir)
|
||||||
except TypeError:
|
except TypeError:
|
||||||
new_backend = _load_backend(model_id, shard_start, shard_end, quantization)
|
new_backend = _load_backend(model_id, shard_start, shard_end, quantization)
|
||||||
self._backend = new_backend
|
self._backends[model_id] = new_backend
|
||||||
self._tracker_mode = shard_start == 0
|
if replacing or shard_start == 0:
|
||||||
if self._server is not None:
|
self._backend = new_backend
|
||||||
self._server.backend = new_backend
|
self._tracker_mode = shard_start == 0
|
||||||
self._server.tracker_mode = self._tracker_mode
|
|
||||||
print(
|
print(
|
||||||
f" [node] loaded reassigned shard: {model_id} layers {shard_start}-{shard_end}",
|
f" [node] loaded {action_label} shard: {model_id} layers {shard_start}-{shard_end}",
|
||||||
flush=True,
|
flush=True,
|
||||||
)
|
)
|
||||||
|
if self._server is not None:
|
||||||
|
self._server.backends = dict(self._backends)
|
||||||
|
if replacing or shard_start == 0:
|
||||||
|
self._server.backend = new_backend
|
||||||
|
self._server.tracker_mode = self._tracker_mode
|
||||||
return {
|
return {
|
||||||
|
"action": directive.get("action"),
|
||||||
"model": model_id,
|
"model": model_id,
|
||||||
"shard_start": shard_start,
|
"shard_start": shard_start,
|
||||||
"shard_end": shard_end,
|
"shard_end": shard_end,
|
||||||
"quantization": quantization,
|
"quantization": quantization,
|
||||||
"tracker_mode": self._tracker_mode,
|
"tracker_mode": shard_start == 0,
|
||||||
}
|
}
|
||||||
|
|
||||||
def start(self) -> int:
|
def start(self) -> int:
|
||||||
@@ -802,7 +854,9 @@ class TorchNodeServer:
|
|||||||
self._tracker_url,
|
self._tracker_url,
|
||||||
self._route_timeout,
|
self._route_timeout,
|
||||||
self._debug,
|
self._debug,
|
||||||
|
self._max_loaded_shards,
|
||||||
)
|
)
|
||||||
|
self._server.backends = dict(self._backends)
|
||||||
self.port = self._server.server_address[1]
|
self.port = self._server.server_address[1]
|
||||||
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
|
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
|
||||||
self._thread.start()
|
self._thread.start()
|
||||||
|
|||||||
@@ -944,6 +944,243 @@ def _pool_summary(nodes: list[_NodeEntry], preset: dict | None = None) -> dict:
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _repo_demand_rpm(
|
||||||
|
server: "_TrackerHTTPServer",
|
||||||
|
repo: str,
|
||||||
|
local_rpms: dict[str, dict],
|
||||||
|
) -> float:
|
||||||
|
"""Look up recent request rate by hf_repo or preset alias."""
|
||||||
|
if repo in local_rpms:
|
||||||
|
return float(local_rpms[repo].get("rpm_last_hour", 0.0) or 0.0)
|
||||||
|
resolved_name, _preset = _resolve_model_preset(server.model_presets, repo)
|
||||||
|
if resolved_name and resolved_name in local_rpms:
|
||||||
|
return float(local_rpms[resolved_name].get("rpm_last_hour", 0.0) or 0.0)
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
|
||||||
|
def _preset_for_node(server: "_TrackerHTTPServer", node: _NodeEntry) -> dict | None:
|
||||||
|
if node.hf_repo:
|
||||||
|
_resolved, preset = _resolve_model_preset(server.model_presets, node.hf_repo)
|
||||||
|
if preset is not None:
|
||||||
|
return preset
|
||||||
|
if node.model:
|
||||||
|
preset = server.model_presets.get(node.model)
|
||||||
|
if preset is not None:
|
||||||
|
return preset
|
||||||
|
if node.hf_repo and node.num_layers:
|
||||||
|
return _hf_rebalance_preset([node])
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def _assignment_memory_bytes(node: _NodeEntry, preset: dict | None) -> int:
|
||||||
|
if preset is None or node.shard_start is None or node.shard_end is None:
|
||||||
|
return 0
|
||||||
|
layers = node.shard_end - node.shard_start + 1
|
||||||
|
if layers <= 0:
|
||||||
|
return 0
|
||||||
|
quantization = _node_quantization(node, preset)
|
||||||
|
layer_bytes = _preset_bytes_per_layer(preset).get(quantization)
|
||||||
|
if not layer_bytes:
|
||||||
|
layer_bytes = next(iter(_preset_bytes_per_layer(preset).values()))
|
||||||
|
return layers * int(layer_bytes)
|
||||||
|
|
||||||
|
|
||||||
|
def _endpoint_memory_pool(
|
||||||
|
server: "_TrackerHTTPServer",
|
||||||
|
nodes: list[_NodeEntry],
|
||||||
|
) -> dict:
|
||||||
|
"""Per-host RAM/VRAM budget, usage, spare capacity, and loaded assignments."""
|
||||||
|
if not nodes:
|
||||||
|
return {}
|
||||||
|
host = nodes[0]
|
||||||
|
budget_bytes, budget_source = _node_memory_budget_bytes(host)
|
||||||
|
reserve_bytes = int(budget_bytes * 0.8)
|
||||||
|
loaded: list[dict] = []
|
||||||
|
used_bytes = 0
|
||||||
|
for node in nodes:
|
||||||
|
preset = _preset_for_node(server, node)
|
||||||
|
assignment_bytes = _assignment_memory_bytes(node, preset)
|
||||||
|
used_bytes += assignment_bytes
|
||||||
|
loaded.append({
|
||||||
|
"node_id": node.node_id,
|
||||||
|
"model": node.model,
|
||||||
|
"hf_repo": node.hf_repo,
|
||||||
|
"shard_start": node.shard_start,
|
||||||
|
"shard_end": node.shard_end,
|
||||||
|
"memory_bytes": assignment_bytes,
|
||||||
|
"status": node.status,
|
||||||
|
})
|
||||||
|
max_slots = max(node.max_loaded_shards for node in nodes)
|
||||||
|
loaded_slots = len([node for node in nodes if node.shard_start is not None])
|
||||||
|
spare_bytes = max(0, reserve_bytes - used_bytes)
|
||||||
|
return {
|
||||||
|
"endpoint": host.endpoint,
|
||||||
|
"memory_budget_bytes": budget_bytes,
|
||||||
|
"memory_budget_source": budget_source,
|
||||||
|
"memory_reserve_bytes": reserve_bytes,
|
||||||
|
"memory_used_bytes": used_bytes,
|
||||||
|
"memory_spare_bytes": spare_bytes,
|
||||||
|
"loaded_slots": loaded_slots,
|
||||||
|
"max_loaded_shards": max_slots,
|
||||||
|
"spare_slots": max(0, max_slots - loaded_slots),
|
||||||
|
"loaded": loaded,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _memory_pool_map(server: "_TrackerHTTPServer") -> dict:
|
||||||
|
"""Aggregate and per-endpoint view of assignable RAM across the hive."""
|
||||||
|
from collections import defaultdict
|
||||||
|
|
||||||
|
groups: dict[str, list[_NodeEntry]] = defaultdict(list)
|
||||||
|
for node in server.registry.values():
|
||||||
|
groups[node.endpoint.rstrip("/")].append(node)
|
||||||
|
|
||||||
|
hosts = [_endpoint_memory_pool(server, group) for group in groups.values()]
|
||||||
|
hosts.sort(key=lambda item: (-item.get("memory_spare_bytes", 0), item.get("endpoint", "")))
|
||||||
|
total_budget = sum(item["memory_budget_bytes"] for item in hosts)
|
||||||
|
total_used = sum(item["memory_used_bytes"] for item in hosts)
|
||||||
|
total_spare = sum(item["memory_spare_bytes"] for item in hosts)
|
||||||
|
total_slots = sum(item["max_loaded_shards"] for item in hosts)
|
||||||
|
loaded_slots = sum(item["loaded_slots"] for item in hosts)
|
||||||
|
return {
|
||||||
|
"total_memory_budget_bytes": total_budget,
|
||||||
|
"total_memory_used_bytes": total_used,
|
||||||
|
"total_memory_spare_bytes": total_spare,
|
||||||
|
"total_loaded_slots": loaded_slots,
|
||||||
|
"total_max_loaded_shards": total_slots,
|
||||||
|
"total_spare_slots": max(0, total_slots - loaded_slots),
|
||||||
|
"hosts": hosts,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _add_shard_directive(
|
||||||
|
node: _NodeEntry,
|
||||||
|
model: str,
|
||||||
|
start: int,
|
||||||
|
end: int,
|
||||||
|
quantization: str,
|
||||||
|
*,
|
||||||
|
model_sources: list[dict] | None = None,
|
||||||
|
) -> dict:
|
||||||
|
directive = {
|
||||||
|
"action": "ADD_SHARD",
|
||||||
|
"model": model,
|
||||||
|
"start_layer": start,
|
||||||
|
"end_layer": end,
|
||||||
|
"shard_start": start,
|
||||||
|
"shard_end": end,
|
||||||
|
"quantization": quantization,
|
||||||
|
}
|
||||||
|
if model_sources:
|
||||||
|
directive["model_sources"] = model_sources
|
||||||
|
return directive
|
||||||
|
|
||||||
|
|
||||||
|
def _model_demand_and_supply(
|
||||||
|
server: "_TrackerHTTPServer",
|
||||||
|
model_key: str,
|
||||||
|
preset: dict,
|
||||||
|
) -> dict:
|
||||||
|
resolved_name, _ = _resolve_model_preset(server.model_presets, model_key)
|
||||||
|
lookup = preset.get("hf_repo") or resolved_name or model_key
|
||||||
|
required_start, required_end = _preset_layer_bounds(preset)
|
||||||
|
model_nodes = [
|
||||||
|
node for node in server.registry.values()
|
||||||
|
if _node_matches_preset(node, resolved_name or model_key, preset) # type: ignore[arg-type]
|
||||||
|
]
|
||||||
|
health = _model_health_summary(server, model_key, preset.get("hf_repo"))
|
||||||
|
local_rpms = server.stats.get_local_rpms() if server.stats is not None else {}
|
||||||
|
demand_rpm = _repo_demand_rpm(server, str(lookup), local_rpms)
|
||||||
|
return {
|
||||||
|
"model": resolved_name or model_key,
|
||||||
|
"hf_repo": preset.get("hf_repo"),
|
||||||
|
"preset": preset,
|
||||||
|
"demand_rpm": demand_rpm,
|
||||||
|
"served_model_copies": float(health.get("served_model_copies") or 0.0),
|
||||||
|
"coverage_percentage": float(health.get("coverage_percentage") or 0.0),
|
||||||
|
"required_start": required_start,
|
||||||
|
"required_end": required_end,
|
||||||
|
"model_nodes": model_nodes,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _scale_demanded_models_locked(server: "_TrackerHTTPServer") -> None:
|
||||||
|
"""Use spare host RAM/slots to add copies of high-demand models."""
|
||||||
|
pool = _memory_pool_map(server)
|
||||||
|
candidates: list[tuple[float, str, dict]] = []
|
||||||
|
for model_key, preset in server.model_presets.items():
|
||||||
|
if not isinstance(preset, dict) or not preset.get("hf_repo"):
|
||||||
|
continue
|
||||||
|
info = _model_demand_and_supply(server, model_key, preset)
|
||||||
|
demand = info["demand_rpm"]
|
||||||
|
copies = info["served_model_copies"]
|
||||||
|
coverage = info["coverage_percentage"]
|
||||||
|
if coverage < 100.0:
|
||||||
|
continue
|
||||||
|
target_copies = max(1.0, demand / 5.0)
|
||||||
|
if copies >= target_copies:
|
||||||
|
continue
|
||||||
|
candidates.append((demand, model_key, info))
|
||||||
|
|
||||||
|
if not candidates:
|
||||||
|
return
|
||||||
|
|
||||||
|
candidates.sort(key=lambda item: (-item[0], item[2]["served_model_copies"]))
|
||||||
|
for _demand, model_key, info in candidates:
|
||||||
|
preset = info["preset"]
|
||||||
|
hf_repo = str(info["hf_repo"])
|
||||||
|
required_start = info["required_start"]
|
||||||
|
required_end = info["required_end"]
|
||||||
|
total_layers = required_end - required_start + 1
|
||||||
|
layer_bytes = next(iter(_preset_bytes_per_layer(preset).values()))
|
||||||
|
full_copy_bytes = total_layers * layer_bytes
|
||||||
|
|
||||||
|
for host in pool["hosts"]:
|
||||||
|
if host["spare_slots"] <= 0 or host["memory_spare_bytes"] < full_copy_bytes:
|
||||||
|
continue
|
||||||
|
host_nodes = [
|
||||||
|
server.registry[node["node_id"]]
|
||||||
|
for node in host["loaded"]
|
||||||
|
if node["node_id"] in server.registry
|
||||||
|
]
|
||||||
|
if not host_nodes:
|
||||||
|
continue
|
||||||
|
if any((n.hf_repo or n.model) == hf_repo for n in host_nodes):
|
||||||
|
continue
|
||||||
|
anchor = max(host_nodes, key=lambda n: n.benchmark_tokens_per_sec)
|
||||||
|
if anchor.status != "ready" or anchor.pending_new_assignment is not None:
|
||||||
|
continue
|
||||||
|
capacity = min(_node_layer_capacity(anchor, preset), total_layers)
|
||||||
|
if capacity <= 0:
|
||||||
|
continue
|
||||||
|
quantization = _node_quantization(anchor, preset)
|
||||||
|
shard_end = min(required_end, required_start + capacity - 1)
|
||||||
|
assignment = _add_shard_directive(
|
||||||
|
anchor,
|
||||||
|
hf_repo,
|
||||||
|
required_start,
|
||||||
|
shard_end,
|
||||||
|
quantization,
|
||||||
|
)
|
||||||
|
anchor.pending_new_assignment = assignment
|
||||||
|
anchor.pending_directives.append(assignment)
|
||||||
|
_tracker_log(
|
||||||
|
server,
|
||||||
|
"info",
|
||||||
|
"scale demanded model copy",
|
||||||
|
endpoint=anchor.endpoint,
|
||||||
|
node_id=anchor.node_id,
|
||||||
|
model=model_key,
|
||||||
|
hf_repo=hf_repo,
|
||||||
|
shard=f"{required_start}-{shard_end}",
|
||||||
|
demand_rpm=round(info["demand_rpm"], 4),
|
||||||
|
served_model_copies=info["served_model_copies"],
|
||||||
|
target_copies=round(max(1.0, info["demand_rpm"] / 5.0), 2),
|
||||||
|
host_spare_bytes=host["memory_spare_bytes"],
|
||||||
|
)
|
||||||
|
break
|
||||||
|
|
||||||
|
|
||||||
def _deployment_summary(nodes: list[_NodeEntry], preset: dict | None) -> dict:
|
def _deployment_summary(nodes: list[_NodeEntry], preset: dict | None) -> dict:
|
||||||
if preset is None:
|
if preset is None:
|
||||||
return {"recommended": False}
|
return {"recommended": False}
|
||||||
@@ -1099,6 +1336,76 @@ def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]:
|
|||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def _unassigned_managed_nodes(nodes: list["_NodeEntry"]) -> list["_NodeEntry"]:
|
||||||
|
return [
|
||||||
|
node for node in nodes
|
||||||
|
if node.managed_assignment
|
||||||
|
and (node.shard_start is None or node.shard_end is None)
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def _emit_shard_change_directives(
|
||||||
|
node: "_NodeEntry",
|
||||||
|
model: str,
|
||||||
|
previous_range: tuple[int | None, int | None, str | None],
|
||||||
|
preset: dict,
|
||||||
|
) -> None:
|
||||||
|
"""Queue DROP/LOAD directives when a managed node's shard assignment changes."""
|
||||||
|
previous_start, previous_end, previous_quantization = previous_range
|
||||||
|
current_range = (node.shard_start, node.shard_end, node.quantization)
|
||||||
|
if node.shard_start is None or node.shard_end is None or current_range == previous_range:
|
||||||
|
return
|
||||||
|
if previous_start is not None and previous_end is not None:
|
||||||
|
node.pending_directives.append(
|
||||||
|
_drop_directive(
|
||||||
|
node,
|
||||||
|
model,
|
||||||
|
previous_start,
|
||||||
|
previous_end,
|
||||||
|
previous_quantization or _node_quantization(node, preset),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
node.pending_directives.append(
|
||||||
|
_load_directive(
|
||||||
|
node,
|
||||||
|
model,
|
||||||
|
node.shard_start,
|
||||||
|
node.shard_end,
|
||||||
|
node.quantization or _node_quantization(node, preset),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _assign_redundant_managed_nodes(
|
||||||
|
managed_nodes: list["_NodeEntry"],
|
||||||
|
model: str,
|
||||||
|
preset: dict,
|
||||||
|
required_start: int,
|
||||||
|
required_end: int,
|
||||||
|
) -> None:
|
||||||
|
"""Give newly joined managed nodes their own copy without reshuffling incumbents."""
|
||||||
|
total_layers = required_end - required_start + 1
|
||||||
|
for node in sorted(
|
||||||
|
managed_nodes,
|
||||||
|
key=lambda entry: (
|
||||||
|
-entry.benchmark_tokens_per_sec,
|
||||||
|
-_node_layer_capacity(entry, preset),
|
||||||
|
entry.node_id,
|
||||||
|
),
|
||||||
|
):
|
||||||
|
if node.shard_start is not None and node.shard_end is not None:
|
||||||
|
continue
|
||||||
|
capacity = min(_node_layer_capacity(node, preset), total_layers)
|
||||||
|
if capacity <= 0:
|
||||||
|
continue
|
||||||
|
previous_range = (node.shard_start, node.shard_end, node.quantization)
|
||||||
|
quantization = _node_quantization(node, preset)
|
||||||
|
node.quantization = quantization
|
||||||
|
node.shard_start = required_start
|
||||||
|
node.shard_end = min(required_end, required_start + capacity - 1)
|
||||||
|
_emit_shard_change_directives(node, model, previous_range, preset)
|
||||||
|
|
||||||
|
|
||||||
def _relay_http_request_frames(
|
def _relay_http_request_frames(
|
||||||
relay_addr: str,
|
relay_addr: str,
|
||||||
path: str,
|
path: str,
|
||||||
@@ -1439,6 +1746,17 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
|
|||||||
if not managed_nodes:
|
if not managed_nodes:
|
||||||
return
|
return
|
||||||
|
|
||||||
|
coverage = _coverage_map(model_nodes, required_start, required_end)
|
||||||
|
gaps = _coverage_gaps(coverage)
|
||||||
|
unassigned = _unassigned_managed_nodes(managed_nodes)
|
||||||
|
if not gaps and not unassigned:
|
||||||
|
return
|
||||||
|
if not gaps and unassigned:
|
||||||
|
_assign_redundant_managed_nodes(
|
||||||
|
unassigned, model, preset, required_start, required_end,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
previous_ranges = {
|
previous_ranges = {
|
||||||
node.node_id: (node.shard_start, node.shard_end, node.quantization)
|
node.node_id: (node.shard_start, node.shard_end, node.quantization)
|
||||||
for node in managed_nodes
|
for node in managed_nodes
|
||||||
@@ -1455,10 +1773,8 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
base_nodes = [node for node in model_nodes if not node.managed_assignment]
|
base_nodes = [node for node in model_nodes if not node.managed_assignment]
|
||||||
coverage = _coverage_map(base_nodes, required_start, required_end)
|
base_coverage = _coverage_map(base_nodes, required_start, required_end)
|
||||||
gaps = _coverage_gaps(coverage)
|
gaps = _coverage_gaps(base_coverage)
|
||||||
if not gaps:
|
|
||||||
gaps = [(required_start, required_end)]
|
|
||||||
|
|
||||||
eligible_nodes = [
|
eligible_nodes = [
|
||||||
node for node in managed_nodes
|
node for node in managed_nodes
|
||||||
@@ -1487,28 +1803,11 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
|
|||||||
node_index += 1
|
node_index += 1
|
||||||
|
|
||||||
for node in managed_nodes:
|
for node in managed_nodes:
|
||||||
previous_start, previous_end, previous_quantization = previous_ranges[node.node_id]
|
_emit_shard_change_directives(
|
||||||
current_range = (node.shard_start, node.shard_end, node.quantization)
|
node,
|
||||||
if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]:
|
model,
|
||||||
continue
|
previous_ranges[node.node_id],
|
||||||
if previous_start is not None and previous_end is not None:
|
preset,
|
||||||
node.pending_directives.append(
|
|
||||||
_drop_directive(
|
|
||||||
node,
|
|
||||||
model,
|
|
||||||
previous_start,
|
|
||||||
previous_end,
|
|
||||||
previous_quantization or _node_quantization(node, preset),
|
|
||||||
)
|
|
||||||
)
|
|
||||||
node.pending_directives.append(
|
|
||||||
_load_directive(
|
|
||||||
node,
|
|
||||||
model,
|
|
||||||
node.shard_start,
|
|
||||||
node.shard_end,
|
|
||||||
node.quantization or _node_quantization(node, preset),
|
|
||||||
)
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -1537,6 +1836,17 @@ def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> No
|
|||||||
if total_layers <= 0:
|
if total_layers <= 0:
|
||||||
return
|
return
|
||||||
|
|
||||||
|
coverage = _coverage_map(model_nodes, required_start, required_end)
|
||||||
|
gaps = _coverage_gaps(coverage)
|
||||||
|
unassigned = _unassigned_managed_nodes(managed_nodes)
|
||||||
|
if not gaps and not unassigned:
|
||||||
|
return
|
||||||
|
if not gaps and unassigned:
|
||||||
|
_assign_redundant_managed_nodes(
|
||||||
|
unassigned, hf_repo, preset, required_start, required_end,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
previous_ranges = {
|
previous_ranges = {
|
||||||
node.node_id: (node.shard_start, node.shard_end, node.quantization)
|
node.node_id: (node.shard_start, node.shard_end, node.quantization)
|
||||||
for node in managed_nodes
|
for node in managed_nodes
|
||||||
@@ -1553,10 +1863,8 @@ def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> No
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
base_nodes = [node for node in model_nodes if not node.managed_assignment]
|
base_nodes = [node for node in model_nodes if not node.managed_assignment]
|
||||||
coverage = _coverage_map(base_nodes, required_start, required_end)
|
base_coverage = _coverage_map(base_nodes, required_start, required_end)
|
||||||
gaps = _coverage_gaps(coverage)
|
gaps = _coverage_gaps(base_coverage)
|
||||||
if not gaps:
|
|
||||||
gaps = [(required_start, required_end)]
|
|
||||||
|
|
||||||
eligible_nodes = [
|
eligible_nodes = [
|
||||||
node for node in managed_nodes
|
node for node in managed_nodes
|
||||||
@@ -1585,28 +1893,11 @@ def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> No
|
|||||||
node_index += 1
|
node_index += 1
|
||||||
|
|
||||||
for node in managed_nodes:
|
for node in managed_nodes:
|
||||||
previous_start, previous_end, previous_quantization = previous_ranges[node.node_id]
|
_emit_shard_change_directives(
|
||||||
current_range = (node.shard_start, node.shard_end, node.quantization)
|
node,
|
||||||
if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]:
|
hf_repo,
|
||||||
continue
|
previous_ranges[node.node_id],
|
||||||
if previous_start is not None and previous_end is not None:
|
preset,
|
||||||
node.pending_directives.append(
|
|
||||||
_drop_directive(
|
|
||||||
node,
|
|
||||||
hf_repo,
|
|
||||||
previous_start,
|
|
||||||
previous_end,
|
|
||||||
previous_quantization or _node_quantization(node, preset),
|
|
||||||
)
|
|
||||||
)
|
|
||||||
node.pending_directives.append(
|
|
||||||
_load_directive(
|
|
||||||
node,
|
|
||||||
hf_repo,
|
|
||||||
node.shard_start,
|
|
||||||
node.shard_end,
|
|
||||||
node.quantization or _node_quantization(node, preset),
|
|
||||||
)
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -2296,6 +2587,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
with server.lock:
|
with server.lock:
|
||||||
self._purge_expired_nodes()
|
self._purge_expired_nodes()
|
||||||
nodes = list(server.registry.values())
|
nodes = list(server.registry.values())
|
||||||
|
memory_pool = _memory_pool_map(server)
|
||||||
|
|
||||||
def capacity_for(node: _NodeEntry) -> dict:
|
def capacity_for(node: _NodeEntry) -> dict:
|
||||||
preset = None
|
preset = None
|
||||||
@@ -2320,6 +2612,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self._send_json(200, {
|
self._send_json(200, {
|
||||||
"relay_url": server.relay_url,
|
"relay_url": server.relay_url,
|
||||||
"pool": _pool_summary(nodes),
|
"pool": _pool_summary(nodes),
|
||||||
|
"memory_pool": memory_pool,
|
||||||
"recommended_models": [
|
"recommended_models": [
|
||||||
{
|
{
|
||||||
"id": name,
|
"id": name,
|
||||||
@@ -3247,17 +3540,23 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
)
|
)
|
||||||
with server.lock:
|
with server.lock:
|
||||||
self._purge_expired_nodes()
|
self._purge_expired_nodes()
|
||||||
# Dedup: if this endpoint is already registered, remove the old entry first.
|
# Dedup: replace the same node id or the same endpoint+model assignment.
|
||||||
|
endpoint_key = entry.endpoint.rstrip("/")
|
||||||
|
model_key = entry.hf_repo or entry.model
|
||||||
stale_ids = [
|
stale_ids = [
|
||||||
eid for eid, e in server.registry.items()
|
eid for eid, e in server.registry.items()
|
||||||
if e.endpoint == entry.endpoint.rstrip("/")
|
if eid == node_id
|
||||||
|
or (
|
||||||
|
e.endpoint.rstrip("/") == endpoint_key
|
||||||
|
and (e.hf_repo or e.model) == model_key
|
||||||
|
)
|
||||||
]
|
]
|
||||||
stale_entries: list[tuple[str, _NodeEntry]] = []
|
stale_entries: list[tuple[str, _NodeEntry]] = []
|
||||||
for eid in stale_ids:
|
for eid in stale_ids:
|
||||||
old = server.registry.pop(eid)
|
old = server.registry.pop(eid)
|
||||||
stale_entries.append((eid, old))
|
stale_entries.append((eid, old))
|
||||||
server.registry[node_id] = entry
|
server.registry[node_id] = entry
|
||||||
if entry.managed_assignment:
|
if entry.managed_assignment and not explicit_shard:
|
||||||
if entry.hf_repo:
|
if entry.hf_repo:
|
||||||
_rebalance_hf_model_locked(server, entry.hf_repo)
|
_rebalance_hf_model_locked(server, entry.hf_repo)
|
||||||
else:
|
else:
|
||||||
@@ -3358,10 +3657,6 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
entry.uptime_seconds = float(body["uptime_seconds"])
|
entry.uptime_seconds = float(body["uptime_seconds"])
|
||||||
if "status" in body and body["status"] in ("ready", "loading"):
|
if "status" in body and body["status"] in ("ready", "loading"):
|
||||||
entry.status = body["status"]
|
entry.status = body["status"]
|
||||||
if entry.hf_repo:
|
|
||||||
_rebalance_hf_model_locked(server, entry.hf_repo)
|
|
||||||
else:
|
|
||||||
_rebalance_model_locked(server, entry.model or "stub-model")
|
|
||||||
directives = list(entry.pending_directives)
|
directives = list(entry.pending_directives)
|
||||||
entry.pending_directives.clear()
|
entry.pending_directives.clear()
|
||||||
new_assignment = entry.pending_new_assignment
|
new_assignment = entry.pending_new_assignment
|
||||||
@@ -4487,7 +4782,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
if server.stats is not None:
|
if server.stats is not None:
|
||||||
local_rpms = server.stats.get_local_rpms()
|
local_rpms = server.stats.get_local_rpms()
|
||||||
for repo in repo_groups:
|
for repo in repo_groups:
|
||||||
demand_rpms[repo] = local_rpms.get(repo, {}).get("rpm_last_hour", 0.0)
|
demand_rpms[repo] = _repo_demand_rpm(server, repo, local_rpms)
|
||||||
|
|
||||||
best_repo = None
|
best_repo = None
|
||||||
best_score = -1.0
|
best_score = -1.0
|
||||||
@@ -4497,6 +4792,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
|
|
||||||
for repo, nodes in repo_groups.items():
|
for repo, nodes in repo_groups.items():
|
||||||
total = repo_layers[repo]
|
total = repo_layers[repo]
|
||||||
|
served_copies = _served_model_copies(nodes, 0, total - 1)
|
||||||
covered = sorted(
|
covered = sorted(
|
||||||
[(n.shard_start, n.shard_end) for n in nodes],
|
[(n.shard_start, n.shard_end) for n in nodes],
|
||||||
key=lambda t: t[0],
|
key=lambda t: t[0],
|
||||||
@@ -4511,8 +4807,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
gap_size = max(0, (total - 1) - gap_start + 1)
|
gap_size = max(0, (total - 1) - gap_start + 1)
|
||||||
coverage_deficit = gap_size / max(total, 1)
|
coverage_deficit = gap_size / max(total, 1)
|
||||||
demand = demand_rpms.get(repo, 0.0)
|
demand = demand_rpms.get(repo, 0.0)
|
||||||
# +1.0 floor on demand so models with no recorded traffic still compete by coverage.
|
# Prefer demanded models already being served; coverage gaps still dominate.
|
||||||
score = (demand + 1.0) * (coverage_deficit + 0.01)
|
score = (demand + 1.0) * (coverage_deficit + 0.01) + served_copies * 0.01
|
||||||
if score > best_score or (score == best_score and gap_size > best_gap_size):
|
if score > best_score or (score == best_score and gap_size > best_gap_size):
|
||||||
best_score = score
|
best_score = score
|
||||||
best_gap_size = gap_size
|
best_gap_size = gap_size
|
||||||
@@ -4522,8 +4818,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
|
|
||||||
gap_found = best_gap_size > 0
|
gap_found = best_gap_size > 0
|
||||||
if not gap_found:
|
if not gap_found:
|
||||||
# All shards covered — assign to highest-demand model for redundancy.
|
# All shards covered — add redundancy on the most demanded/served model.
|
||||||
best_repo = max(repo_groups, key=lambda r: demand_rpms.get(r, 0.0))
|
best_repo = max(
|
||||||
|
repo_groups,
|
||||||
|
key=lambda repo: (
|
||||||
|
demand_rpms.get(repo, 0.0),
|
||||||
|
_served_model_copies(repo_groups[repo], 0, repo_layers[repo] - 1),
|
||||||
|
len(repo_groups[repo]),
|
||||||
|
),
|
||||||
|
)
|
||||||
best_gap_start = 0
|
best_gap_start = 0
|
||||||
best_num_layers = repo_layers[best_repo]
|
best_num_layers = repo_layers[best_repo]
|
||||||
|
|
||||||
@@ -5141,6 +5444,7 @@ class TrackerServer:
|
|||||||
with self._lock:
|
with self._lock:
|
||||||
_purge_expired_nodes_locked(server)
|
_purge_expired_nodes_locked(server)
|
||||||
_rebalance_all_locked(server)
|
_rebalance_all_locked(server)
|
||||||
|
_scale_demanded_models_locked(server)
|
||||||
|
|
||||||
def _push_to_peers(self, path: str, body: bytes) -> bool:
|
def _push_to_peers(self, path: str, body: bytes) -> bool:
|
||||||
"""POST a hive-signed payload to every cluster peer; True if all succeeded."""
|
"""POST a hive-signed payload to every cluster peer; True if all succeeded."""
|
||||||
|
|||||||
@@ -388,6 +388,107 @@ def test_legacy_start_treats_repo_model_as_model_id(monkeypatch):
|
|||||||
assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||||
|
|
||||||
|
|
||||||
|
def test_legacy_start_catalog_model_with_pinned_shards(monkeypatch):
|
||||||
|
"""Catalog model names accept --shard-start/--shard-end without --model-id."""
|
||||||
|
from meshnet_node.cli import main
|
||||||
|
|
||||||
|
captured = {}
|
||||||
|
|
||||||
|
def fake_run_startup(*args, **kwargs):
|
||||||
|
captured.update(kwargs)
|
||||||
|
class _FakeNode:
|
||||||
|
chat_completion_count = 0
|
||||||
|
def stop(self): pass
|
||||||
|
return _FakeNode()
|
||||||
|
|
||||||
|
monkeypatch.setattr(sys, "argv", [
|
||||||
|
"meshnet-node", "start",
|
||||||
|
"--tracker", "http://192.168.0.179:8080",
|
||||||
|
"--model", "Qwen3.6-35B-A3B",
|
||||||
|
"--shard-start", "0",
|
||||||
|
"--shard-end", "44",
|
||||||
|
"--port", "0",
|
||||||
|
])
|
||||||
|
|
||||||
|
with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
|
||||||
|
with patch("time.sleep", side_effect=KeyboardInterrupt):
|
||||||
|
try:
|
||||||
|
main()
|
||||||
|
except SystemExit as exc:
|
||||||
|
assert exc.code == 0
|
||||||
|
|
||||||
|
assert captured["model"] == "Qwen3.6-35B-A3B"
|
||||||
|
assert captured["model_id"] is None
|
||||||
|
assert captured["shard_start"] == 0
|
||||||
|
assert captured["shard_end"] == 44
|
||||||
|
|
||||||
|
|
||||||
|
def test_legacy_start_model_id_alias_for_catalog_name(monkeypatch):
|
||||||
|
"""--model-id with a catalog name routes through the tracker preset path."""
|
||||||
|
from meshnet_node.cli import main
|
||||||
|
|
||||||
|
captured = {}
|
||||||
|
|
||||||
|
def fake_run_startup(*args, **kwargs):
|
||||||
|
captured.update(kwargs)
|
||||||
|
class _FakeNode:
|
||||||
|
chat_completion_count = 0
|
||||||
|
def stop(self): pass
|
||||||
|
return _FakeNode()
|
||||||
|
|
||||||
|
monkeypatch.setattr(sys, "argv", [
|
||||||
|
"meshnet-node", "start",
|
||||||
|
"--tracker", "http://192.168.0.179:8080",
|
||||||
|
"--model-id", "Qwen3.6-35B-A3B",
|
||||||
|
"--port", "0",
|
||||||
|
])
|
||||||
|
|
||||||
|
with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
|
||||||
|
with patch("time.sleep", side_effect=KeyboardInterrupt):
|
||||||
|
try:
|
||||||
|
main()
|
||||||
|
except SystemExit as exc:
|
||||||
|
assert exc.code == 0
|
||||||
|
|
||||||
|
assert captured["model"] == "Qwen3.6-35B-A3B"
|
||||||
|
assert captured["model_id"] is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_legacy_start_hf_repo_with_pinned_shards(monkeypatch):
|
||||||
|
"""HF repo --model with pinned shards still enters the torch startup path."""
|
||||||
|
from meshnet_node.cli import main
|
||||||
|
|
||||||
|
captured = {}
|
||||||
|
|
||||||
|
def fake_run_startup(*args, **kwargs):
|
||||||
|
captured.update(kwargs)
|
||||||
|
class _FakeNode:
|
||||||
|
chat_completion_count = 0
|
||||||
|
def stop(self): pass
|
||||||
|
return _FakeNode()
|
||||||
|
|
||||||
|
monkeypatch.setattr(sys, "argv", [
|
||||||
|
"meshnet-node", "start",
|
||||||
|
"--tracker", "http://192.168.0.179:8081",
|
||||||
|
"--model", "Qwen/Qwen2.5-0.5B-Instruct",
|
||||||
|
"--shard-start", "12",
|
||||||
|
"--shard-end", "23",
|
||||||
|
"--port", "0",
|
||||||
|
])
|
||||||
|
|
||||||
|
with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
|
||||||
|
with patch("time.sleep", side_effect=KeyboardInterrupt):
|
||||||
|
try:
|
||||||
|
main()
|
||||||
|
except SystemExit as exc:
|
||||||
|
assert exc.code == 0
|
||||||
|
|
||||||
|
assert captured["model"] == "Qwen2.5-0.5B-Instruct"
|
||||||
|
assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||||
|
assert captured["shard_start"] == 12
|
||||||
|
assert captured["shard_end"] == 23
|
||||||
|
|
||||||
|
|
||||||
def test_legacy_start_falls_back_to_env_tracker_and_model(monkeypatch):
|
def test_legacy_start_falls_back_to_env_tracker_and_model(monkeypatch):
|
||||||
"""`meshnet-node start` uses env defaults when tracker/model flags are omitted."""
|
"""`meshnet-node start` uses env defaults when tracker/model flags are omitted."""
|
||||||
import importlib
|
import importlib
|
||||||
|
|||||||
@@ -5,8 +5,10 @@ import io
|
|||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
import tarfile
|
import tarfile
|
||||||
|
import threading
|
||||||
import time
|
import time
|
||||||
import types
|
import types
|
||||||
|
import urllib.error
|
||||||
import urllib.request
|
import urllib.request
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
@@ -1431,6 +1433,84 @@ def test_later_node_auto_joins_existing_public_hf_model_with_only_tracker_url(
|
|||||||
assert route_resp["route"] == ["http://203.0.113.20:8001", "http://203.0.113.21:8002"]
|
assert route_resp["route"] == ["http://203.0.113.20:8001", "http://203.0.113.21:8002"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_later_node_auto_joins_redundant_copy_when_model_is_fully_covered(
|
||||||
|
tmp_path,
|
||||||
|
monkeypatch,
|
||||||
|
):
|
||||||
|
"""Model-less joins should load the served HF model even when gap_found=false."""
|
||||||
|
import meshnet_node.startup as startup_mod
|
||||||
|
|
||||||
|
captured = {}
|
||||||
|
|
||||||
|
class FakeBackend:
|
||||||
|
total_layers = 24
|
||||||
|
|
||||||
|
class FakeTorchNodeServer:
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
captured.update(kwargs)
|
||||||
|
self.backend = FakeBackend()
|
||||||
|
self.port = None
|
||||||
|
self.chat_completion_count = 0
|
||||||
|
self.total_requests = 0
|
||||||
|
self.failed_requests = 0
|
||||||
|
self.queue_depth = 0
|
||||||
|
|
||||||
|
def start(self):
|
||||||
|
self.port = 8003
|
||||||
|
return self.port
|
||||||
|
|
||||||
|
def stop(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
startup_mod,
|
||||||
|
"detect_hardware",
|
||||||
|
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
|
||||||
|
|
||||||
|
tracker = TrackerServer()
|
||||||
|
tracker_port = tracker.start()
|
||||||
|
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||||
|
try:
|
||||||
|
for endpoint, shard_start, shard_end in (
|
||||||
|
("http://203.0.113.30:8001", 0, 11),
|
||||||
|
("http://203.0.113.31:8001", 12, 23),
|
||||||
|
):
|
||||||
|
data = json.dumps({
|
||||||
|
"endpoint": endpoint,
|
||||||
|
"model": "Qwen2.5-0.5B-Instruct",
|
||||||
|
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||||
|
"num_layers": 24,
|
||||||
|
"shard_start": shard_start,
|
||||||
|
"shard_end": shard_end,
|
||||||
|
"tracker_mode": shard_start == 0,
|
||||||
|
"hardware_profile": {},
|
||||||
|
"score": 1.0,
|
||||||
|
}).encode()
|
||||||
|
req = urllib.request.Request(
|
||||||
|
f"{tracker_url}/v1/nodes/register",
|
||||||
|
data=data,
|
||||||
|
headers={"Content-Type": "application/json"},
|
||||||
|
method="POST",
|
||||||
|
)
|
||||||
|
with urllib.request.urlopen(req) as resp:
|
||||||
|
resp.read()
|
||||||
|
|
||||||
|
node = run_startup(
|
||||||
|
tracker_url=tracker_url,
|
||||||
|
advertise_host="203.0.113.32",
|
||||||
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||||
|
assert captured["shard_start"] == 0
|
||||||
|
finally:
|
||||||
|
node.stop()
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# Full startup integration test
|
# Full startup integration test
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
@@ -1525,6 +1605,120 @@ def test_preset_model_startup_starts_heartbeat(tmp_path, monkeypatch):
|
|||||||
tracker.stop()
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
|
def test_preset_model_startup_honors_pinned_shard_range(tmp_path, monkeypatch):
|
||||||
|
"""Explicit --shard-start/--shard-end override tracker auto-assignment."""
|
||||||
|
import meshnet_node.startup as startup_mod
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
startup_mod,
|
||||||
|
"detect_hardware",
|
||||||
|
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16 * 1024},
|
||||||
|
)
|
||||||
|
heartbeat_calls = []
|
||||||
|
monkeypatch.setattr(
|
||||||
|
startup_mod,
|
||||||
|
"_start_heartbeat",
|
||||||
|
lambda *args, **kwargs: heartbeat_calls.append((args, kwargs)),
|
||||||
|
)
|
||||||
|
|
||||||
|
tracker = TrackerServer(model_presets={"stub-model": {"layers_start": 0, "layers_end": 15}})
|
||||||
|
tracker_port = tracker.start()
|
||||||
|
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||||
|
try:
|
||||||
|
node = run_startup(
|
||||||
|
tracker_url=tracker_url,
|
||||||
|
model="stub-model",
|
||||||
|
shard_start=0,
|
||||||
|
shard_end=5,
|
||||||
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
cache_dir=tmp_path / "shards",
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
assert len(heartbeat_calls) == 1
|
||||||
|
args, kwargs = heartbeat_calls[0]
|
||||||
|
reg_payload = args[2]
|
||||||
|
assert reg_payload["shard_start"] == 0
|
||||||
|
assert reg_payload["shard_end"] == 5
|
||||||
|
assert reg_payload["managed_assignment"] is False
|
||||||
|
finally:
|
||||||
|
node.stop()
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
|
def test_torch_startup_retries_registration_when_tracker_unreachable(
|
||||||
|
tmp_path,
|
||||||
|
monkeypatch,
|
||||||
|
):
|
||||||
|
"""Failed initial registration should start background retry, not stay unregistered."""
|
||||||
|
import meshnet_node.startup as startup_mod
|
||||||
|
|
||||||
|
class FakeBackend:
|
||||||
|
total_layers = 24
|
||||||
|
|
||||||
|
class FakeTorchNodeServer:
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
self.backend = FakeBackend()
|
||||||
|
self.port = None
|
||||||
|
self.chat_completion_count = 0
|
||||||
|
self.tracker_node_id = None
|
||||||
|
|
||||||
|
def start(self):
|
||||||
|
self.port = 7000
|
||||||
|
return self.port
|
||||||
|
|
||||||
|
def stop(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
startup_mod,
|
||||||
|
"detect_hardware",
|
||||||
|
lambda: {"device": "cuda", "gpu_name": "Test GPU", "vram_mb": 8192, "ram_mb": 16 * 1024},
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
startup_mod,
|
||||||
|
"_detect_num_layers",
|
||||||
|
lambda *_args, **_kwargs: 24,
|
||||||
|
)
|
||||||
|
|
||||||
|
heartbeat_calls = []
|
||||||
|
monkeypatch.setattr(
|
||||||
|
startup_mod,
|
||||||
|
"_start_heartbeat",
|
||||||
|
lambda *args, **kwargs: heartbeat_calls.append((args, kwargs)) or threading.Thread(),
|
||||||
|
)
|
||||||
|
|
||||||
|
register_calls = {"count": 0}
|
||||||
|
|
||||||
|
def flaky_register(url, payload):
|
||||||
|
register_calls["count"] += 1
|
||||||
|
raise urllib.error.URLError("connection refused")
|
||||||
|
|
||||||
|
monkeypatch.setattr(startup_mod, "_post_json", flaky_register)
|
||||||
|
|
||||||
|
tracker = TrackerServer()
|
||||||
|
tracker_port = tracker.start()
|
||||||
|
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||||
|
try:
|
||||||
|
node = run_startup(
|
||||||
|
tracker_url=tracker_url,
|
||||||
|
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
||||||
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
assert register_calls["count"] == 1
|
||||||
|
assert node.tracker_node_id is None
|
||||||
|
assert len(heartbeat_calls) == 1
|
||||||
|
args, kwargs = heartbeat_calls[0]
|
||||||
|
assert args[1] == startup_mod._PENDING_NODE_ID
|
||||||
|
assert kwargs["node_ref"] is node
|
||||||
|
finally:
|
||||||
|
node.stop()
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
def test_real_model_startup_registers_downloaded_inventory_without_checksum(
|
def test_real_model_startup_registers_downloaded_inventory_without_checksum(
|
||||||
tmp_path,
|
tmp_path,
|
||||||
monkeypatch,
|
monkeypatch,
|
||||||
|
|||||||
@@ -13,7 +13,13 @@ from meshnet_gateway.server import GatewayServer, _banned_route_wallet
|
|||||||
from meshnet_node.server import StubNodeServer
|
from meshnet_node.server import StubNodeServer
|
||||||
from meshnet_contracts import LocalSolanaContracts
|
from meshnet_contracts import LocalSolanaContracts
|
||||||
from meshnet_tracker.auth import sign_hive_request
|
from meshnet_tracker.auth import sign_hive_request
|
||||||
from meshnet_tracker.server import TrackerServer, _NodeEntry, _registration_ban_error
|
from meshnet_tracker.server import (
|
||||||
|
TrackerServer,
|
||||||
|
_NodeEntry,
|
||||||
|
_memory_pool_map,
|
||||||
|
_registration_ban_error,
|
||||||
|
_scale_demanded_models_locked,
|
||||||
|
)
|
||||||
|
|
||||||
_TEST_HIVE_SECRET = "test-hive-secret"
|
_TEST_HIVE_SECRET = "test-hive-secret"
|
||||||
|
|
||||||
@@ -1114,12 +1120,14 @@ def test_tracker_speed_is_primary_when_both_nodes_can_cover_gap():
|
|||||||
"benchmark_tokens_per_sec": 3.0, "hardware_profile": {}, "score": 1.0},
|
"benchmark_tokens_per_sec": 3.0, "hardware_profile": {}, "score": 1.0},
|
||||||
)
|
)
|
||||||
|
|
||||||
route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
|
net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
|
||||||
widths = {
|
widths = {
|
||||||
node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
|
node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
|
||||||
for node in route_resp["nodes"]
|
for node in net["nodes"]
|
||||||
}
|
}
|
||||||
assert widths["http://127.0.0.1:9012"] > widths["http://127.0.0.1:9011"]
|
assert widths["http://127.0.0.1:9011"] == 12
|
||||||
|
assert widths["http://127.0.0.1:9012"] == 12
|
||||||
|
assert net["nodes"][0]["model_supply"]["served_model_copies"] == 2.0
|
||||||
finally:
|
finally:
|
||||||
tracker.stop()
|
tracker.stop()
|
||||||
|
|
||||||
@@ -1147,7 +1155,8 @@ def test_tracker_registration_directive_is_not_replayed_on_heartbeat():
|
|||||||
tracker.stop()
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
def test_tracker_reassignment_emits_drop_before_load():
|
def test_tracker_pool_join_adds_redundant_copy_without_splitting_incumbent():
|
||||||
|
"""A second managed node with capacity for the full model keeps the first copy intact."""
|
||||||
tracker = TrackerServer(model_presets={
|
tracker = TrackerServer(model_presets={
|
||||||
"tiny-model": {
|
"tiny-model": {
|
||||||
"total_layers": 4,
|
"total_layers": 4,
|
||||||
@@ -1156,21 +1165,69 @@ def test_tracker_reassignment_emits_drop_before_load():
|
|||||||
})
|
})
|
||||||
tracker_port = tracker.start()
|
tracker_port = tracker.start()
|
||||||
try:
|
try:
|
||||||
slow = _post_json(
|
first = _post_json(
|
||||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||||
{"endpoint": "http://127.0.0.1:9015", "model": "tiny-model",
|
{"endpoint": "http://127.0.0.1:9015", "model": "tiny-model",
|
||||||
"vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
"vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
||||||
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
||||||
)
|
)
|
||||||
_post_json(
|
second = _post_json(
|
||||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||||
{"endpoint": "http://127.0.0.1:9016", "model": "tiny-model",
|
{"endpoint": "http://127.0.0.1:9016", "model": "tiny-model",
|
||||||
"vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
"vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
||||||
"benchmark_tokens_per_sec": 3.0, "hardware_profile": {}, "score": 1.0},
|
"benchmark_tokens_per_sec": 3.0, "hardware_profile": {}, "score": 1.0},
|
||||||
)
|
)
|
||||||
|
|
||||||
hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{slow['node_id']}/heartbeat", {})
|
net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
|
||||||
assert [directive["action"] for directive in hb["directives"]] == ["DROP_SHARD", "LOAD_SHARD"]
|
widths = {
|
||||||
|
node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
|
||||||
|
for node in net["nodes"]
|
||||||
|
}
|
||||||
|
assert widths["http://127.0.0.1:9015"] == 4
|
||||||
|
assert widths["http://127.0.0.1:9016"] == 4
|
||||||
|
assert net["nodes"][0]["model_supply"]["served_model_copies"] == 2.0
|
||||||
|
|
||||||
|
hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{first['node_id']}/heartbeat", {})
|
||||||
|
assert hb.get("directives", []) == []
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
|
def test_tracker_explicit_full_copy_join_preserves_existing_serving_node():
|
||||||
|
"""--model style joins with explicit shards add redundancy instead of reshuffling."""
|
||||||
|
tracker = TrackerServer(heartbeat_timeout=10.0)
|
||||||
|
tracker_port = tracker.start()
|
||||||
|
try:
|
||||||
|
base_reg = {
|
||||||
|
"model": "Qwen2.5-0.5B-Instruct",
|
||||||
|
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||||
|
"num_layers": 24,
|
||||||
|
"shard_start": 0,
|
||||||
|
"shard_end": 23,
|
||||||
|
"managed_assignment": True,
|
||||||
|
"vram_bytes": 2_000_000_000,
|
||||||
|
"hardware_profile": {},
|
||||||
|
"score": 1.0,
|
||||||
|
}
|
||||||
|
first = _post_json(
|
||||||
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||||
|
{**base_reg, "endpoint": "http://127.0.0.1:9201"},
|
||||||
|
)
|
||||||
|
second = _post_json(
|
||||||
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||||
|
{**base_reg, "endpoint": "http://127.0.0.1:9202"},
|
||||||
|
)
|
||||||
|
|
||||||
|
coverage = _get_json(
|
||||||
|
f"http://127.0.0.1:{tracker_port}/v1/network/map"
|
||||||
|
)
|
||||||
|
assert coverage["nodes"][0]["model_supply"]["served_model_copies"] == 2.0
|
||||||
|
|
||||||
|
hb = _post_json(
|
||||||
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/{first['node_id']}/heartbeat", {}
|
||||||
|
)
|
||||||
|
assert hb.get("directives", []) == []
|
||||||
|
assert second["node_id"] in tracker._registry
|
||||||
finally:
|
finally:
|
||||||
tracker.stop()
|
tracker.stop()
|
||||||
|
|
||||||
@@ -1208,22 +1265,20 @@ def test_tracker_faster_node_receives_wider_range_when_capacity_ties():
|
|||||||
_post_json(
|
_post_json(
|
||||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||||
{"endpoint": "http://127.0.0.1:9001", "model": "tiny-model",
|
{"endpoint": "http://127.0.0.1:9001", "model": "tiny-model",
|
||||||
"vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
"vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
||||||
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
||||||
)
|
)
|
||||||
_post_json(
|
_post_json(
|
||||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||||
{"endpoint": "http://127.0.0.1:9002", "model": "tiny-model",
|
{"endpoint": "http://127.0.0.1:9002", "model": "tiny-model",
|
||||||
"vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
"vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"],
|
||||||
"benchmark_tokens_per_sec": 2.0, "hardware_profile": {}, "score": 1.0},
|
"benchmark_tokens_per_sec": 2.0, "hardware_profile": {}, "score": 1.0},
|
||||||
)
|
)
|
||||||
|
|
||||||
route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
|
net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
|
||||||
widths = {
|
heads = [node for node in net["nodes"] if node["shard_start"] == 0]
|
||||||
node["endpoint"]: node["shard_end"] - node["shard_start"] + 1
|
assert len(heads) == 1
|
||||||
for node in route_resp["nodes"]
|
assert heads[0]["endpoint"] == "http://127.0.0.1:9002"
|
||||||
}
|
|
||||||
assert widths["http://127.0.0.1:9002"] > widths["http://127.0.0.1:9001"]
|
|
||||||
finally:
|
finally:
|
||||||
tracker.stop()
|
tracker.stop()
|
||||||
|
|
||||||
@@ -2291,6 +2346,7 @@ def test_torch_node_applies_tracker_load_shard_directive(monkeypatch):
|
|||||||
|
|
||||||
assert loaded == [("Qwen/Qwen2.5-0.5B-Instruct", 0, 23, "bfloat16")]
|
assert loaded == [("Qwen/Qwen2.5-0.5B-Instruct", 0, 23, "bfloat16")]
|
||||||
assert applied == {
|
assert applied == {
|
||||||
|
"action": "LOAD_SHARD",
|
||||||
"model": "Qwen/Qwen2.5-0.5B-Instruct",
|
"model": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||||
"shard_start": 0,
|
"shard_start": 0,
|
||||||
"shard_end": 23,
|
"shard_end": 23,
|
||||||
@@ -2407,3 +2463,105 @@ def test_shard_heal_cycle_surviving_node_covers_dead_peers_gap(monkeypatch):
|
|||||||
)
|
)
|
||||||
finally:
|
finally:
|
||||||
tracker.stop()
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
|
def test_network_map_exposes_memory_pool():
|
||||||
|
tracker = TrackerServer(model_presets={
|
||||||
|
"tiny-model": {
|
||||||
|
"total_layers": 8,
|
||||||
|
"bytes_per_layer": {"bfloat16": 1_000},
|
||||||
|
"hf_repo": "org/TinyModel",
|
||||||
|
},
|
||||||
|
})
|
||||||
|
tracker_port = tracker.start()
|
||||||
|
try:
|
||||||
|
_post_json(
|
||||||
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||||
|
{"endpoint": "http://127.0.0.1:9050", "model": "tiny-model",
|
||||||
|
"hf_repo": "org/TinyModel", "num_layers": 8,
|
||||||
|
"shard_start": 0, "shard_end": 7, "max_loaded_shards": 2,
|
||||||
|
"vram_bytes": 20_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
|
||||||
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
||||||
|
)
|
||||||
|
net = _get_json(f"http://127.0.0.1:{tracker_port}/v1/network/map")
|
||||||
|
pool = net["memory_pool"]
|
||||||
|
assert pool["total_spare_slots"] == 1
|
||||||
|
assert pool["hosts"][0]["loaded_slots"] == 1
|
||||||
|
assert pool["hosts"][0]["max_loaded_shards"] == 2
|
||||||
|
assert pool["hosts"][0]["memory_spare_bytes"] > 0
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
|
def test_same_endpoint_can_register_multiple_models():
|
||||||
|
tracker = TrackerServer()
|
||||||
|
tracker_port = tracker.start()
|
||||||
|
try:
|
||||||
|
base = {
|
||||||
|
"endpoint": "http://127.0.0.1:9055",
|
||||||
|
"num_layers": 24,
|
||||||
|
"shard_start": 0,
|
||||||
|
"shard_end": 23,
|
||||||
|
"hardware_profile": {},
|
||||||
|
"score": 1.0,
|
||||||
|
"max_loaded_shards": 2,
|
||||||
|
"vram_bytes": 50_000_000,
|
||||||
|
"ram_bytes": 50_000_000,
|
||||||
|
}
|
||||||
|
first = _post_json(
|
||||||
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||||
|
{**base, "model": "Qwen2.5-0.5B-Instruct", "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct"},
|
||||||
|
)
|
||||||
|
second = _post_json(
|
||||||
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||||
|
{**base, "model": "OtherModel", "hf_repo": "org/OtherModel"},
|
||||||
|
)
|
||||||
|
assert first["node_id"] != second["node_id"]
|
||||||
|
assert len(tracker._registry) == 2
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|
||||||
|
|
||||||
|
def test_scale_demanded_models_queues_add_shard_on_spare_host():
|
||||||
|
tracker = TrackerServer(model_presets={
|
||||||
|
"model-a": {
|
||||||
|
"total_layers": 4,
|
||||||
|
"bytes_per_layer": {"bfloat16": 1_000},
|
||||||
|
"hf_repo": "org/ModelA",
|
||||||
|
},
|
||||||
|
"model-b": {
|
||||||
|
"total_layers": 4,
|
||||||
|
"bytes_per_layer": {"bfloat16": 1_000},
|
||||||
|
"hf_repo": "org/ModelB",
|
||||||
|
},
|
||||||
|
})
|
||||||
|
tracker_port = tracker.start()
|
||||||
|
try:
|
||||||
|
reg_b = _post_json(
|
||||||
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||||
|
{"endpoint": "http://127.0.0.1:9061", "model": "model-b",
|
||||||
|
"hf_repo": "org/ModelB", "num_layers": 4,
|
||||||
|
"shard_start": 0, "shard_end": 3, "max_loaded_shards": 2,
|
||||||
|
"vram_bytes": 20_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
|
||||||
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
||||||
|
)
|
||||||
|
_post_json(
|
||||||
|
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||||
|
{"endpoint": "http://127.0.0.1:9060", "model": "model-a",
|
||||||
|
"hf_repo": "org/ModelA", "num_layers": 4,
|
||||||
|
"shard_start": 0, "shard_end": 3, "max_loaded_shards": 1,
|
||||||
|
"vram_bytes": 20_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"],
|
||||||
|
"benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0},
|
||||||
|
)
|
||||||
|
assert tracker._stats is not None
|
||||||
|
for _ in range(400):
|
||||||
|
tracker._stats.record_request("org/ModelA")
|
||||||
|
with tracker._lock:
|
||||||
|
_scale_demanded_models_locked(tracker._server) # type: ignore[arg-type]
|
||||||
|
node_b = tracker._registry[reg_b["node_id"]]
|
||||||
|
assignment = node_b.pending_new_assignment
|
||||||
|
assert assignment is not None
|
||||||
|
assert assignment["action"] == "ADD_SHARD"
|
||||||
|
assert assignment["model"] == "org/ModelA"
|
||||||
|
finally:
|
||||||
|
tracker.stop()
|
||||||
|
|||||||
Reference in New Issue
Block a user