From e9a094b6200b775518b52f3787b18ea123a9b327 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Tue, 7 Jul 2026 18:35:36 +0300 Subject: [PATCH] ram pool map --- billing.sqlite | Bin 24576 -> 24576 bytes meshnet_registry.sqlite3 | Bin 12288 -> 20480 bytes packages/node/meshnet_node/startup.py | 35 +++++- packages/node/meshnet_node/torch_server.py | 120 +++++++++++++++------ packages/tracker/meshnet_tracker/server.py | 41 ++++--- tests/test_tracker_routing.py | 111 ++++++++++++++++++- 6 files changed, 253 insertions(+), 54 deletions(-) diff --git a/billing.sqlite b/billing.sqlite index eec52b48119325239da5cdb198f48228ddf7ec9e..00972758407fae36802e8957678889f43999e487 100644 GIT binary patch delta 233 zcmZoTz}Rqrae_3X#6%fqMv09HOXQg)xwlScS5Rl--nv;&p_N-ln>m^>IU}(sJykc& zIME`-+{84^*eKD&#L_S^$<)jsEy*l3F)7W!G|5nM^FwWA6)tA(6%2e6cxAX(Y%H9` z9oJMS$r|0t&Q7Eu)k>KuN>&O=L@J;B&?IZ}Ja^T}oN>h>#>S?4mL{g=2FB*Kj<+Ve O{r}UJxmhOm2@?PbzDG|0 delta 58 zcmV-A0LA}+zyW~30gxL37m*x80T;1gq#pwq3*nOoASbghAcYH&fG)G)D%* zVV-Z3*%f5Dc)szk=jZ0z$=|lI@C46fd0kEg3sVEL#6%N|RD(nl^HdWH15*owL}LSk zR0C7fG*c7d&GNbrnHC8MaPfjIP-Ea%0~tQCAvvNkT$nYqDUv-aKS?){fH~DlnJG$E j3QEMNo%~NAPsG^RRL|1H)Xdz}q}K7)WVio+T0a2*>$*dn delta 55 zcmZozz}S#5L7G*Ffq{W}qJll6(8h!X{EG|(*tvirLcH>m*%f3t_}BAu^X=qs+gSLS KXS2NSLnZ)Z0}ast diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index a184e9b..61e950a 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -372,6 +372,27 @@ def _start_heartbeat( except Exception: 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: if not directives: return @@ -384,6 +405,9 @@ def _start_heartbeat( print(f" [node] WARNING: failed to apply tracker directives: {exc}", flush=True) return 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")) register_payload["model"] = str(model_id).split("/")[-1] register_payload["hf_repo"] = model_id @@ -423,11 +447,12 @@ def _start_heartbeat( new_asgn = resp.get("new_assignment") if new_asgn: print( - f" [node] tracker reassignment received: " - f"model={new_asgn.get('model')!r} " + f" [node] tracker assignment received: " + 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')}", flush=True, ) + _apply_directives([new_asgn]) except urllib.error.HTTPError as exc: if exc.code == 404: # Node was purged (e.g. long gap before restart noticed) — re-register now. @@ -668,6 +693,7 @@ def run_startup( route_timeout=route_timeout, cache_dir=cache_dir, debug=debug, + max_loaded_shards=max_loaded_shards, ) _node_start_time = time.monotonic() actual_port = node.start() @@ -810,6 +836,7 @@ def run_startup( route_timeout=route_timeout, cache_dir=cache_dir, debug=debug, + max_loaded_shards=max_loaded_shards, ) _node_start_time = time.monotonic() actual_port = node.start() @@ -883,13 +910,13 @@ def run_startup( ) return node - if model != "stub-model": + 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 ." ) - # 3b. Stub preset path (tests / explicit stub-model only) + # 3b. Stub preset path (tests / explicit stub-model) or named preset models. print("Querying tracker for shard assignment...", flush=True) assign_qs = urllib.parse.urlencode({ "model": model or "stub-model", diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 50a4985..e8cf923 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -75,20 +75,39 @@ class _TorchHTTPServer(http.server.HTTPServer): tracker_url: str | None = None, route_timeout: float = 30.0, debug: bool = False, + max_loaded_shards: int = 1, ): super().__init__(addr, handler) self.backend = backend + self.backends: dict[str, TorchModelShard] = {backend.model_id: backend} self.received_activations = False self.forward_chunk_count = 0 self.tracker_mode = tracker_mode self.tracker_url = tracker_url self.route_timeout = route_timeout self.debug = debug + self.max_loaded_shards = max(1, max_loaded_shards) self.total_requests: int = 0 self.failed_requests: int = 0 self.queue_depth: int = 0 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): def log_message(self, fmt, *args): # noqa: suppress request logs in tests @@ -100,7 +119,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self._handle_forward() elif self.path == "/v1/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() else: self.send_response(404) @@ -284,22 +303,26 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): messages = [] stream = bool(body.get("stream", False)) 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) temperature = float(body.get("temperature") 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. # 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: if stream: 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, ) else: - text = server.backend.generate_text(messages, max_tokens, temperature, top_p) - self._send_openai_response(text, model_name, False, messages) + text = backend.generate_text(messages, max_tokens, temperature, top_p) + self._send_openai_response(text, model_name, False, messages, backend=backend) except Exception as exc: self._record_failed_request() 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 # but correct. Each step: head encodes current sequence → forwards through route # → 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( f" [node] chat route model={model_name!r} max_tokens={max_tokens} " f"downstream={remaining_route}", @@ -318,11 +341,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): if not remaining_route: self._send_openai_response( "error: no downstream route — check tracker connectivity", - model_name, False, messages, + model_name, False, messages, backend=backend, ) return - backend = server.backend # Format with chat template so the model knows it's in assistant mode. try: if hasattr(backend.tokenizer, "apply_chat_template"): @@ -352,7 +374,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): except Exception as exc: print(f" [node] distributed encode error: {exc}", flush=True) 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: break # Stop on error responses or EOS. @@ -369,9 +391,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): if stream_emit is not None: stream_emit(None) 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. 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). server: _TorchHTTPServer = self.server # type: ignore[assignment] + active_backend = backend or server.backend if server.tracker_url is None: return [] - route_model = getattr(server.backend, "model_id", None) or model + route_model = getattr(active_backend, "model_id", None) or model try: url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}" 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) 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] + active_backend = backend or server.backend if not route: # Partial shard at tail: decode the activation from the previous node. # Full single-node (head+tail) is handled before entering this method. - if server.backend.is_tail: + if active_backend.is_tail: try: - tensor = server.backend.torch.frombuffer( + tensor = active_backend.torch.frombuffer( bytearray(payload.body), # type: ignore[union-attr] - dtype=server.backend.torch.bfloat16, - ).reshape(payload.shape).to(server.backend.device) # type: ignore[union-attr] - return server.backend.decode_tail(tensor) + dtype=active_backend.torch.bfloat16, + ).reshape(payload.shape).to(active_backend.device) # type: ignore[union-attr] + return active_backend.decode_tail(tensor) except Exception as exc: return f"decode error: {exc}" return "no downstream route available for non-tail shard" @@ -591,11 +615,13 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): model: str, stream: bool, messages: list[dict] | None = None, + backend: TorchModelShard | None = None, ) -> None: chunk_id = "chatcmpl-node" created = int(time.time()) + active_backend = backend or self.server.backend # type: ignore[attr-defined] 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, { "id": chunk_id, "object": "chat.completion", @@ -706,9 +732,11 @@ class TorchNodeServer: route_timeout: float = 30.0, cache_dir: Path | None = None, debug: bool = False, + max_loaded_shards: int = 1, ) -> None: self._host = host self._requested_port = port + self._max_loaded_shards = max(1, max_loaded_shards) self._backend = backend or _load_backend( model_id, shard_start, @@ -716,6 +744,7 @@ class TorchNodeServer: quantization, 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 self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0) self._tracker_url = tracker_url @@ -754,41 +783,64 @@ class TorchNodeServer: def queue_depth(self) -> int: 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: - """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( (directive for directive in reversed(directives) if directive.get("action") == "LOAD_SHARD"), None, ) - if load_directive is None: + directive = add_directive or load_directive + if directive is None: return None - shard_start = int(load_directive["shard_start"]) - shard_end = int(load_directive["shard_end"]) - quantization = str(load_directive.get("quantization") or self._backend.quantization) - model_id = str(load_directive.get("model") or self._backend.model_id) + shard_start = int(directive["shard_start"]) + shard_end = int(directive["shard_end"]) + quantization = str(directive.get("quantization") or self._backend.quantization) + 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( - 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, ) try: new_backend = _load_backend(model_id, shard_start, shard_end, quantization, self._cache_dir) except TypeError: new_backend = _load_backend(model_id, shard_start, shard_end, quantization) - self._backend = new_backend - self._tracker_mode = shard_start == 0 - if self._server is not None: - self._server.backend = new_backend - self._server.tracker_mode = self._tracker_mode + self._backends[model_id] = new_backend + if replacing or shard_start == 0: + self._backend = new_backend + self._tracker_mode = shard_start == 0 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, ) + 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 { + "action": directive.get("action"), "model": model_id, "shard_start": shard_start, "shard_end": shard_end, "quantization": quantization, - "tracker_mode": self._tracker_mode, + "tracker_mode": shard_start == 0, } def start(self) -> int: @@ -802,7 +854,9 @@ class TorchNodeServer: self._tracker_url, self._route_timeout, self._debug, + self._max_loaded_shards, ) + self._server.backends = dict(self._backends) self.port = self._server.server_address[1] self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) self._thread.start() diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 69b77cc..db330cd 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -944,6 +944,20 @@ 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) @@ -1392,20 +1406,6 @@ def _assign_redundant_managed_nodes( _emit_shard_change_directives(node, model, previous_range, preset) -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 _relay_http_request_frames( relay_addr: str, path: str, @@ -2587,6 +2587,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): with server.lock: self._purge_expired_nodes() nodes = list(server.registry.values()) + memory_pool = _memory_pool_map(server) def capacity_for(node: _NodeEntry) -> dict: preset = None @@ -2611,6 +2612,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._send_json(200, { "relay_url": server.relay_url, "pool": _pool_summary(nodes), + "memory_pool": memory_pool, "recommended_models": [ { "id": name, @@ -3538,10 +3540,16 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): ) with server.lock: 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 = [ 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]] = [] for eid in stale_ids: @@ -5436,6 +5444,7 @@ class TrackerServer: with self._lock: _purge_expired_nodes_locked(server) _rebalance_all_locked(server) + _scale_demanded_models_locked(server) def _push_to_peers(self, path: str, body: bytes) -> bool: """POST a hive-signed payload to every cluster peer; True if all succeeded.""" diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index ab3fa0a..ef934a6 100644 --- a/tests/test_tracker_routing.py +++ b/tests/test_tracker_routing.py @@ -13,7 +13,13 @@ from meshnet_gateway.server import GatewayServer, _banned_route_wallet from meshnet_node.server import StubNodeServer from meshnet_contracts import LocalSolanaContracts 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" @@ -2340,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 applied == { + "action": "LOAD_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct", "shard_start": 0, "shard_end": 23, @@ -2456,3 +2463,105 @@ def test_shard_heal_cycle_surviving_node_covers_dead_peers_gap(monkeypatch): ) finally: 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()