diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index c8b8992..7dac320 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -31,28 +31,50 @@ from typing import Any DEFAULT_MODEL_PRESETS: dict[str, dict] = { - "stub-model": {"layers_start": 0, "layers_end": 31}, + "stub-model": { + "layers_start": 0, + "layers_end": 31, + "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}, + }, + "openai-community/gpt2": { + "layers_start": 0, + "layers_end": 11, + "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024}, + }, } +DEFAULT_VRAM_BYTES = 8 * 1024 * 1024 * 1024 +DEFAULT_RAM_BYTES = 16 * 1024 * 1024 * 1024 +DEFAULT_QUANTIZATIONS = ["bfloat16"] +DEFAULT_BENCHMARK_TOKENS_PER_SEC = 1.0 + class _NodeEntry: __slots__ = ( "node_id", "endpoint", "shard_start", "shard_end", "model", "shard_checksum", "hardware_profile", "wallet_address", - "score", "last_heartbeat", + "score", "vram_bytes", "ram_bytes", "quantizations", + "benchmark_tokens_per_sec", "quantization", "managed_assignment", + "pending_directives", "last_heartbeat", ) def __init__( self, node_id: str, endpoint: str, - shard_start: int, - shard_end: int, + shard_start: int | None, + shard_end: int | None, model: str | None, shard_checksum: str | None, hardware_profile: dict, wallet_address: str | None, score: float, + vram_bytes: int = DEFAULT_VRAM_BYTES, + ram_bytes: int = DEFAULT_RAM_BYTES, + quantizations: list[str] | None = None, + benchmark_tokens_per_sec: float = DEFAULT_BENCHMARK_TOKENS_PER_SEC, + quantization: str | None = None, + managed_assignment: bool = False, ) -> None: self.node_id = node_id self.endpoint = endpoint @@ -63,6 +85,13 @@ class _NodeEntry: self.hardware_profile = hardware_profile self.wallet_address = wallet_address self.score = score + self.vram_bytes = vram_bytes + self.ram_bytes = ram_bytes + self.quantizations = quantizations or list(DEFAULT_QUANTIZATIONS) + self.benchmark_tokens_per_sec = benchmark_tokens_per_sec + self.quantization = quantization + self.managed_assignment = managed_assignment + self.pending_directives: list[dict] = [] self.last_heartbeat: float = time.monotonic() @@ -72,7 +101,10 @@ def _select_route( required_end: int, ) -> tuple[list[_NodeEntry], str]: """Greedy interval-cover. Returns (ordered route, error_message).""" - candidates = sorted(nodes, key=lambda n: (n.shard_start, -n.shard_end)) + candidates = sorted( + [node for node in nodes if node.shard_start is not None and node.shard_end is not None], + key=lambda n: (n.shard_start, -n.shard_end), # type: ignore[operator] + ) route: list[_NodeEntry] = [] covered_up_to = required_start - 1 @@ -105,6 +137,8 @@ def _coverage_percentage( ( (max(required_start, node.shard_start), min(required_end, node.shard_end)) for node in nodes + if node.shard_start is not None + and node.shard_end is not None if node.shard_end >= required_start and node.shard_start <= required_end ), key=lambda interval: interval[0], @@ -122,6 +156,197 @@ def _coverage_percentage( return round((covered / required_layers) * 100, 2) +def _preset_layer_bounds(preset: dict) -> tuple[int, int]: + start = int(preset.get("layers_start", 0)) + if "layers_end" in preset: + return start, int(preset["layers_end"]) + return start, start + int(preset["total_layers"]) - 1 + + +def _preset_bytes_per_layer(preset: dict) -> dict[str, int]: + raw = preset.get("bytes_per_layer", preset.get("bytes_per_layer_at_quant", {})) + if isinstance(raw, dict) and raw: + return {str(quant): int(value) for quant, value in raw.items()} + return {"bfloat16": 30 * 1024 * 1024} + + +def _node_quantization(node: _NodeEntry, preset: dict) -> str: + bytes_per_layer = _preset_bytes_per_layer(preset) + if node.quantization in bytes_per_layer: + return node.quantization + for quantization in node.quantizations: + if quantization in bytes_per_layer: + return quantization + return next(iter(bytes_per_layer)) + + +def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int: + bytes_per_layer = _preset_bytes_per_layer(preset) + quantization = _node_quantization(node, preset) + layer_bytes = bytes_per_layer[quantization] + if layer_bytes <= 0: + return 0 + return int((node.vram_bytes * 0.8) // layer_bytes) + + +def _coverage_map( + nodes: list[_NodeEntry], + required_start: int, + required_end: int, +) -> list[dict]: + layer_counts = [] + for layer in range(required_start, required_end + 1): + count = 0 + for node in nodes: + if node.shard_start is None or node.shard_end is None: + continue + if node.shard_start <= layer <= node.shard_end: + count += 1 + layer_counts.append((layer, count)) + + coverage: list[dict] = [] + for layer, count in layer_counts: + if coverage and coverage[-1]["node_count"] == count and coverage[-1]["end_layer"] == layer - 1: + coverage[-1]["end_layer"] = layer + else: + coverage.append({"start_layer": layer, "end_layer": layer, "node_count": count}) + return coverage + + +def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]: + return [ + (segment["start_layer"], segment["end_layer"]) + for segment in coverage + if segment["node_count"] == 0 + ] + + +def _load_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict: + return { + "action": "LOAD_SHARD", + "model": model, + "start_layer": start, + "end_layer": end, + "shard_start": start, + "shard_end": end, + "quantization": quantization, + } + + +def _drop_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict: + return { + "action": "DROP_SHARD", + "model": model, + "start_layer": start, + "end_layer": end, + "shard_start": start, + "shard_end": end, + "quantization": quantization, + } + + +def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]: + now = time.monotonic() + expired_ids = [ + node_id for node_id, entry in server.registry.items() + if (now - entry.last_heartbeat) > server.heartbeat_timeout + ] + for node_id in expired_ids: + del server.registry[node_id] + if expired_ids: + _rebalance_all_locked(server) + return expired_ids + + +def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None: + preset = server.model_presets.get(model) + if preset is None: + return + required_start, required_end = _preset_layer_bounds(preset) + total_layers = required_end - required_start + 1 + model_nodes = [node for node in server.registry.values() if node.model == model] + managed_nodes = [node for node in model_nodes if node.managed_assignment] + if not managed_nodes: + return + + previous_ranges = { + node.node_id: (node.shard_start, node.shard_end, node.quantization) + for node in managed_nodes + } + for node in managed_nodes: + node.shard_start = None + node.shard_end = None + + managed_nodes.sort( + key=lambda node: ( + -node.benchmark_tokens_per_sec, + -_node_layer_capacity(node, preset), + node.node_id, + ) + ) + base_nodes = [node for node in model_nodes if not node.managed_assignment] + coverage = _coverage_map(base_nodes, required_start, required_end) + gaps = _coverage_gaps(coverage) + if not gaps: + gaps = [(required_start, required_end)] + + eligible_nodes = [ + node for node in managed_nodes + if _node_layer_capacity(node, preset) > 0 + ] + node_index = 0 + for gap_start, gap_end in gaps: + cursor = gap_start + while cursor <= gap_end and node_index < len(eligible_nodes): + node = eligible_nodes[node_index] + remaining_layers = gap_end - cursor + 1 + remaining_nodes_after = len(eligible_nodes) - node_index - 1 + capacity = min( + _node_layer_capacity(node, preset), + total_layers, + max(1, remaining_layers - remaining_nodes_after), + ) + if capacity <= 0: + node_index += 1 + continue + quantization = _node_quantization(node, preset) + node.quantization = quantization + node.shard_start = cursor + node.shard_end = min(gap_end, cursor + capacity - 1) + cursor = node.shard_end + 1 + node_index += 1 + + for node in managed_nodes: + previous_start, previous_end, previous_quantization = previous_ranges[node.node_id] + 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_ranges[node.node_id]: + continue + 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 _rebalance_all_locked(server: "_TrackerHTTPServer") -> None: + for model in list(server.model_presets): + _rebalance_model_locked(server, model) + + def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None: if contracts is None or not wallet_address: return None @@ -177,13 +402,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): def _purge_expired_nodes(self) -> None: server: _TrackerHTTPServer = self.server # type: ignore[assignment] - now = time.monotonic() - expired_ids = [ - node_id for node_id, entry in server.registry.items() - if (now - entry.last_heartbeat) > server.heartbeat_timeout - ] - for node_id in expired_ids: - del server.registry[node_id] + _purge_expired_nodes_locked(server) def do_POST(self): if self.path == "/v1/nodes/register": @@ -207,6 +426,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._handle_assign(parsed) elif parsed.path == "/v1/models": self._handle_models() + elif parsed.path.startswith("/v1/coverage/"): + model = urllib.parse.unquote(parsed.path.removeprefix("/v1/coverage/")) + self._handle_coverage(model) else: self.send_response(404) self.end_headers() @@ -225,10 +447,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): data = [] for name, preset in server.model_presets.items(): model_nodes = [node for node in alive if node.model == name] + if not model_nodes: + continue + required_start, required_end = _preset_layer_bounds(preset) coverage = _coverage_percentage( model_nodes, - preset["layers_start"], - preset["layers_end"], + required_start, + required_end, ) data.append({ "id": name, @@ -239,6 +464,24 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): }) self._send_json(200, {"object": "list", "data": data}) + def _handle_coverage(self, model: str): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + preset = server.model_presets.get(model) + if preset is None: + self._send_json(404, {"error": f"unknown model preset: {model!r}"}) + return + required_start, required_end = _preset_layer_bounds(preset) + with server.lock: + self._purge_expired_nodes() + alive = [node for node in server.registry.values() if node.model == model] + if server.contracts is not None: + alive = [ + node for node in alive + if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned + ] + coverage = _coverage_map(alive, required_start, required_end) + self._send_json(200, {"model": model, "coverage": coverage}) + def _handle_register(self): server: _TrackerHTTPServer = self.server # type: ignore[assignment] body = self._read_json_body() @@ -254,15 +497,26 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._send_json(400, {"error": "endpoint must be an http(s) URL"}) return + shard_start: int | None + shard_end: int | None + explicit_shard = "shard_start" in body or "shard_end" in body + if explicit_shard: + try: + shard_start = int(body["shard_start"]) + shard_end = int(body["shard_end"]) + except (KeyError, TypeError, ValueError): + self._send_json(400, {"error": "shard_start and shard_end must be numeric"}) + return + if shard_start < 0 or shard_end < 0 or shard_start > shard_end: + self._send_json(400, {"error": "shard range must be non-negative and ordered"}) + return + else: + shard_start = None + shard_end = None try: - shard_start = int(body["shard_start"]) - shard_end = int(body["shard_end"]) score = float(body.get("score", 1.0)) - except (KeyError, TypeError, ValueError): - self._send_json(400, {"error": "shard_start, shard_end, and score must be numeric"}) - return - if shard_start < 0 or shard_end < 0 or shard_start > shard_end: - self._send_json(400, {"error": "shard range must be non-negative and ordered"}) + except (TypeError, ValueError): + self._send_json(400, {"error": "score must be numeric"}) return hardware_profile = body.get("hardware_profile", {}) @@ -279,6 +533,31 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): if shard_checksum is not None and not isinstance(shard_checksum, str): self._send_json(400, {"error": "shard_checksum must be a string"}) return + try: + vram_bytes = int(body.get("vram_bytes", DEFAULT_VRAM_BYTES)) + ram_bytes = int(body.get("ram_bytes", DEFAULT_RAM_BYTES)) + benchmark_tokens_per_sec = float( + body.get("benchmark_tokens_per_sec", DEFAULT_BENCHMARK_TOKENS_PER_SEC) + ) + except (TypeError, ValueError): + self._send_json(400, {"error": "vram_bytes, ram_bytes, and benchmark_tokens_per_sec must be numeric"}) + return + if vram_bytes < 0 or ram_bytes < 0 or benchmark_tokens_per_sec <= 0: + self._send_json(400, {"error": "capability values must be positive"}) + return + quantizations_body = body.get("quantizations", DEFAULT_QUANTIZATIONS) + if not ( + isinstance(quantizations_body, list) + and quantizations_body + and all(isinstance(item, str) and item for item in quantizations_body) + ): + self._send_json(400, {"error": "quantizations must be a non-empty string array"}) + return + quantizations = list(quantizations_body) + quantization = body.get("quantization") + if quantization is not None and not isinstance(quantization, str): + self._send_json(400, {"error": "quantization must be a string"}) + return wallet_address = body.get("wallet_address") if wallet_address is not None and not isinstance(wallet_address, str): self._send_json(400, {"error": "wallet_address must be a string"}) @@ -299,12 +578,26 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): hardware_profile=hardware_profile, wallet_address=wallet_address, score=score, + vram_bytes=vram_bytes, + ram_bytes=ram_bytes, + quantizations=quantizations, + benchmark_tokens_per_sec=benchmark_tokens_per_sec, + quantization=quantization, + managed_assignment=not explicit_shard, ) with server.lock: self._purge_expired_nodes() server.registry[node_id] = entry + if entry.managed_assignment: + _rebalance_model_locked(server, model) + assignment_directive = entry.pending_directives[-1] if entry.pending_directives else None + if assignment_directive is not None: + entry.pending_directives.clear() - self._send_json(200, {"node_id": node_id}) + payload = {"node_id": node_id} + if assignment_directive is not None: + payload["directive"] = assignment_directive + self._send_json(200, payload) def _handle_heartbeat(self, node_id: str): server: _TrackerHTTPServer = self.server # type: ignore[assignment] @@ -315,7 +608,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._send_json(404, {"error": "node not found"}) return entry.last_heartbeat = time.monotonic() - self._send_json(200, {}) + _rebalance_model_locked(server, entry.model or "stub-model") + directives = list(entry.pending_directives) + entry.pending_directives.clear() + if directives: + self._send_json(200, {"directives": directives}) + else: + self._send_json(200, {}) def _handle_assign(self, parsed: urllib.parse.ParseResult): """Return an optimal shard assignment for a node given its hardware profile. @@ -346,8 +645,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._send_json(404, {"error": f"unknown model preset: {model!r}"}) return - required_start: int = preset["layers_start"] - required_end: int = preset["layers_end"] + required_start, required_end = _preset_layer_bounds(preset) with server.lock: self._purge_expired_nodes() @@ -369,7 +667,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): # Collect covered intervals sorted by start layer. covered = sorted( - [(n.shard_start, n.shard_end) for n in alive], + [ + (n.shard_start, n.shard_end) + for n in alive + if n.shard_start is not None and n.shard_end is not None + ], key=lambda t: t[0], ) @@ -421,8 +723,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._send_json(404, {"error": f"unknown model preset: {model!r}"}) return - required_start: int = preset["layers_start"] - required_end: int = preset["layers_end"] + required_start, required_end = _preset_layer_bounds(preset) with server.lock: self._purge_expired_nodes() @@ -477,8 +778,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._send_json(404, {"error": f"unknown model preset: {model!r}"}) return - required_start: int = preset["layers_start"] - required_end: int = preset["layers_end"] + required_start, required_end = _preset_layer_bounds(preset) with server.lock: self._purge_expired_nodes() @@ -530,6 +830,7 @@ class TrackerServer: host: str = "127.0.0.1", port: int = 0, heartbeat_timeout: float = 30.0, + rebalance_interval: float = 30.0, model_presets: dict | None = None, contracts: Any | None = None, minimum_stake: int = 0, @@ -537,6 +838,7 @@ class TrackerServer: self._host = host self._requested_port = port self._heartbeat_timeout = heartbeat_timeout + self._rebalance_interval = rebalance_interval self._model_presets: dict = ( model_presets if model_presets is not None else dict(DEFAULT_MODEL_PRESETS) ) @@ -546,6 +848,8 @@ class TrackerServer: self._lock = threading.Lock() self._server: _TrackerHTTPServer | None = None self._thread: threading.Thread | None = None + self._rebalance_stop = threading.Event() + self._rebalance_thread: threading.Thread | None = None self.port: int | None = None def start(self) -> int: @@ -562,17 +866,33 @@ class TrackerServer: self._minimum_stake, ) self.port = self._server.server_address[1] + self._rebalance_stop.clear() self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) self._thread.start() + self._rebalance_thread = threading.Thread(target=self._rebalance_loop, daemon=True) + self._rebalance_thread.start() return self.port + def _rebalance_loop(self) -> None: + while not self._rebalance_stop.wait(self._rebalance_interval): + server = self._server + if server is None: + return + with self._lock: + _purge_expired_nodes_locked(server) + _rebalance_all_locked(server) + def stop(self) -> None: if self._server is None: return + self._rebalance_stop.set() self._server.shutdown() self._server.server_close() if self._thread is not None: self._thread.join(timeout=1) + if self._rebalance_thread is not None: + self._rebalance_thread.join(timeout=1) self._server = None self._thread = None + self._rebalance_thread = None self.port = None diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index 39b6b8c..ab63a31 100644 --- a/tests/test_tracker_routing.py +++ b/tests/test_tracker_routing.py @@ -10,7 +10,7 @@ import urllib.request from meshnet_gateway.server import GatewayServer, _banned_route_wallet from meshnet_node.server import StubNodeServer from meshnet_contracts import LocalSolanaContracts -from meshnet_tracker.server import TrackerServer, _registration_ban_error +from meshnet_tracker.server import TrackerServer, _NodeEntry, _registration_ban_error def _post_json(url: str, payload: dict) -> dict: @@ -100,6 +100,294 @@ def test_tracker_route_error_no_coverage(): tracker.stop() +def test_tracker_coverage_endpoint_reports_uncovered_ranges(): + """Coverage endpoint returns compressed layer ranges with node counts.""" + tracker = TrackerServer(model_presets={ + "tiny-model": { + "total_layers": 6, + "bytes_per_layer": {"bfloat16": 1_000}, + }, + }) + tracker_port = tracker.start() + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9001", "model": "tiny-model", + "shard_start": 0, "shard_end": 2, "hardware_profile": {}, "score": 1.0}, + ) + + coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model") + assert coverage_resp["coverage"] == [ + {"start_layer": 0, "end_layer": 2, "node_count": 1}, + {"start_layer": 3, "end_layer": 5, "node_count": 0}, + ] + + try: + _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model") + raise AssertionError("Expected 503 for zero-coverage range") + except urllib.error.HTTPError as exc: + assert exc.code == 503 + finally: + tracker.stop() + + +def test_tracker_auto_assigns_new_node_to_uncovered_range_first(): + """Capability-driven registration fills the first uncovered layer gap.""" + tracker = TrackerServer(model_presets={ + "tiny-model": { + "total_layers": 8, + "bytes_per_layer": {"bfloat16": 1_000}, + }, + }) + tracker_port = tracker.start() + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9001", "model": "tiny-model", + "shard_start": 0, "shard_end": 3, "hardware_profile": {}, "score": 1.0}, + ) + reg = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9002", "model": "tiny-model", + "vram_bytes": 10_000, "ram_bytes": 20_000, "quantizations": ["bfloat16"], + "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, + ) + + assert reg["directive"]["action"] == "LOAD_SHARD" + assert reg["directive"]["start_layer"] == 4 + assert reg["directive"]["end_layer"] == 7 + finally: + tracker.stop() + + +def test_tracker_speed_weighted_vram_assignment_covers_model(): + """Three auto-assigned nodes reach 100% coverage with widest range on largest VRAM.""" + tracker = TrackerServer(model_presets={ + "tiny-model": { + "total_layers": 12, + "bytes_per_layer": {"bfloat16": 1_000}, + }, + }) + tracker_port = tracker.start() + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9001", "model": "tiny-model", + "vram_bytes": 4_000, "ram_bytes": 10_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:9002", "model": "tiny-model", + "vram_bytes": 5_000, "ram_bytes": 10_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:9003", "model": "tiny-model", + "vram_bytes": 7_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], + "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, + ) + + route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model") + widths = { + node["endpoint"]: node["shard_end"] - node["shard_start"] + 1 + for node in route_resp["nodes"] + } + assert widths["http://127.0.0.1:9003"] == 5 + assert widths["http://127.0.0.1:9002"] == 4 + assert widths["http://127.0.0.1:9001"] == 3 + + coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model") + assert all(segment["node_count"] >= 1 for segment in coverage_resp["coverage"]) + finally: + tracker.stop() + + +def test_tracker_speed_is_primary_when_both_nodes_can_cover_gap(): + tracker = TrackerServer(model_presets={ + "tiny-model": { + "total_layers": 12, + "bytes_per_layer": {"bfloat16": 1_000}, + }, + }) + tracker_port = tracker.start() + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9011", "model": "tiny-model", + "vram_bytes": 20_000, "ram_bytes": 10_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:9012", "model": "tiny-model", + "vram_bytes": 15_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], + "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") + widths = { + node["endpoint"]: node["shard_end"] - node["shard_start"] + 1 + for node in route_resp["nodes"] + } + assert widths["http://127.0.0.1:9012"] > widths["http://127.0.0.1:9011"] + finally: + tracker.stop() + + +def test_tracker_registration_directive_is_not_replayed_on_heartbeat(): + tracker = TrackerServer(model_presets={ + "tiny-model": { + "total_layers": 2, + "bytes_per_layer": {"bfloat16": 1_000}, + }, + }) + tracker_port = tracker.start() + try: + reg = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9013", "model": "tiny-model", + "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], + "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, + ) + assert reg["directive"]["action"] == "LOAD_SHARD" + + hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{reg['node_id']}/heartbeat", {}) + assert hb == {} + finally: + tracker.stop() + + +def test_tracker_reassignment_emits_drop_before_load(): + tracker = TrackerServer(model_presets={ + "tiny-model": { + "total_layers": 4, + "bytes_per_layer": {"bfloat16": 1_000}, + }, + }) + tracker_port = tracker.start() + try: + slow = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9015", "model": "tiny-model", + "vram_bytes": 10_000, "ram_bytes": 10_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:9016", "model": "tiny-model", + "vram_bytes": 10_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], + "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", {}) + assert [directive["action"] for directive in hb["directives"]] == ["DROP_SHARD", "LOAD_SHARD"] + finally: + tracker.stop() + + +def test_tracker_periodic_rebalance_purges_expired_nodes_without_requests(): + tracker = TrackerServer( + heartbeat_timeout=0.05, + rebalance_interval=0.02, + model_presets={"tiny-model": {"total_layers": 1, "bytes_per_layer": {"bfloat16": 1_000}}}, + ) + tracker_port = tracker.start() + try: + reg = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9014", "model": "tiny-model", + "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], + "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, + ) + assert reg["node_id"] in tracker._registry + time.sleep(0.12) + assert reg["node_id"] not in tracker._registry + finally: + tracker.stop() + + +def test_tracker_faster_node_receives_wider_range_when_capacity_ties(): + tracker = TrackerServer(model_presets={ + "tiny-model": { + "total_layers": 12, + "bytes_per_layer": {"bfloat16": 1_000}, + }, + }) + tracker_port = tracker.start() + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9001", "model": "tiny-model", + "vram_bytes": 20_000, "ram_bytes": 10_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:9002", "model": "tiny-model", + "vram_bytes": 20_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], + "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") + widths = { + node["endpoint"]: node["shard_end"] - node["shard_start"] + 1 + for node in route_resp["nodes"] + } + assert widths["http://127.0.0.1:9002"] > widths["http://127.0.0.1:9001"] + finally: + tracker.stop() + + +def test_tracker_rebalances_after_middle_range_node_timeout(): + """Killing a middle shard queues LOAD_SHARD and restores coverage.""" + tracker = TrackerServer( + heartbeat_timeout=0.15, + model_presets={ + "tiny-model": { + "total_layers": 12, + "bytes_per_layer": {"bfloat16": 1_000}, + }, + }, + ) + tracker_port = tracker.start() + ids: dict[str, str] = {} + try: + for port in (9001, 9002, 9003, 9004): + reg = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": f"http://127.0.0.1:{port}", "model": "tiny-model", + "vram_bytes": 5_000, "ram_bytes": 10_000, "quantizations": ["bfloat16"], + "benchmark_tokens_per_sec": 1.0, "hardware_profile": {}, "score": 1.0}, + ) + ids[str(port)] = reg["node_id"] + + for node_id in ids.values(): + _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{node_id}/heartbeat", {}) + + time.sleep(0.10) + for port in ("9001", "9003", "9004"): + _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{ids[port]}/heartbeat", {}) + time.sleep(0.10) + + coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model") + assert all(segment["node_count"] >= 1 for segment in coverage_resp["coverage"]) + + heartbeat_responses = [ + _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{ids[port]}/heartbeat", {}) + for port in ("9001", "9003", "9004") + ] + directives = [ + directive + for response in heartbeat_responses + for directive in response.get("directives", []) + ] + assert any(directive["action"] == "LOAD_SHARD" for directive in directives) + finally: + tracker.stop() + + def test_tracker_route_error_no_nodes(): """Tracker returns 503 with clear error when the registry is empty.""" tracker = TrackerServer()