6 Commits

Author SHA1 Message Date
Dobromir Popov
c38e36f685 Retry tracker registration when initial connect fails.
Start background re-registration when the tracker is unreachable at startup so nodes do not stay permanently unregistered.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-07 17:59:27 +02:00
Dobromir Popov
50e8904f1c ignore 2026-07-07 17:57:33 +02:00
Dobromir Popov
7e289fef2e Fix meshnet-node model and shard flag parsing.
Unify --model and --model-id so catalog names use the tracker path, and allow --shard-start/--shard-end with --model instead of requiring --model-id.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-07 17:54:30 +02:00
Dobromir Popov
e9a094b620 ram pool map 2026-07-07 18:35:36 +03:00
Dobromir Popov
1299a6bb1c balancing improvements 2026-07-07 18:30:30 +03:00
Dobromir Popov
f220fd2210 tracker rebalancing tweaks 2026-07-07 18:24:09 +03:00
10 changed files with 1073 additions and 171 deletions

1
.gitignore vendored
View File

@@ -19,3 +19,4 @@ dist/
!.env.testnet !.env.testnet
.rocm-local/* .rocm-local/*
billing.sqlite billing.sqlite
.pytest-tmp/*

Binary file not shown.

Binary file not shown.

View File

@@ -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")

View File

@@ -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,
} }

View File

@@ -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()

View File

@@ -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."""

View File

@@ -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

View File

@@ -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,

View File

@@ -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()