ram pool map

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 <HF_REPO>."
         )
 
-    # 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",