feat(us-016): smart shard gap detection for auto-join and --model-id

Tracker /v1/network/assign now accepts an optional `hf_repo` query param
to restrict assignment to a specific model, and returns `gap_found: bool`
so callers know whether they received a real gap vs a redundancy slot.

Node startup with --model-id (no explicit shard args) now queries the
tracker first for an uncovered gap for that model before defaulting to
full coverage (0..n-1). This means a second node with --model-id will
serve only the missing layers, not the whole model again.

Auto-join fallback (no --model-id) now prints why it fell through
instead of silently switching to stub-model.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
2026-06-30 00:48:34 +03:00
parent 97eefd3d5e
commit 3286e42783
3 changed files with 104 additions and 7 deletions

View File

@@ -112,6 +112,23 @@ def run_startup(
f"Could not read num_hidden_layers from {model_id} config. "
"Pass --shard-start and --shard-end explicitly."
)
# When no explicit shard range given, ask the tracker if there's a gap for this model.
if shard_start is None and shard_end is None:
try:
qs = urllib.parse.urlencode({
"device": device, "vram_mb": vram_mb, "hf_repo": model_id,
})
net_asgn = _get_json(f"{tracker_url}/v1/network/assign?{qs}", timeout=5.0)
if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
shard_start = net_asgn["shard_start"]
shard_end = net_asgn["shard_end"]
print(
f" Tracker found uncovered shard: "
f"layers {shard_start}{shard_end} (of {detected})",
flush=True,
)
except Exception:
pass # No other nodes registered yet — default to full model below
shard_start = shard_start if shard_start is not None else 0
shard_end = shard_end if shard_end is not None else detected - 1
print(f" Auto-detected {detected} layers → shard {shard_start}{shard_end}", flush=True)
@@ -179,13 +196,15 @@ def run_startup(
# 3a. Auto-join: query tracker for network-wide HF model assignment.
print("Querying tracker for network assignment...", flush=True)
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": vram_mb})
net_assignment: dict = {}
try:
net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}")
assigned_hf_repo: str | None = net_assignment.get("hf_repo")
except Exception:
assigned_hf_repo = None
except Exception as exc:
print(f" (auto-join unavailable: {exc})", flush=True)
assigned_hf_repo: str | None = net_assignment.get("hf_repo")
_gap_found: bool = bool(net_assignment.get("gap_found", False))
if assigned_hf_repo:
if assigned_hf_repo and _gap_found:
assigned_shard_start: int = net_assignment["shard_start"]
assigned_shard_end: int = net_assignment["shard_end"]
assigned_num_layers: int = net_assignment["num_layers"]