Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai

.scratch/distributed-inference-network/issues/13-coverage-first-shard-assignment.md

@@ -47,3 +47,7 @@ Model preset metadata (stored in tracker config) includes `bytes_per_layer` at e
 - Rebalance directives are delivered as responses to node heartbeat POSTs (node polls tracker every 10s anyway) — no new push channel needed
 - `bytes_per_layer` is read from a `model_presets.json` config file; add GPT-2 (12 layers, ~30MB/layer bfloat16) as the test preset
 - The current tracker `_select_route` function is extended, not replaced — backward compat with stub nodes that don't send VRAM data (default to 8GB / bfloat16 if omitted)
+
+## Comments
+
+- 2026-06-30: Follow-up capacity hardening lives in `20-memory-budget-shard-slots-and-dropout-relocation.md`. US-013 remains the base coverage-first assignment and dropout rebalance story; US-020 owns operator `--memory`, `--max-shards`, shard-slot enforcement, and relocation limit hardening so the two scopes do not conflict.

.scratch/distributed-inference-network/issues/20-memory-budget-shard-slots-and-dropout-relocation.md

@@ -0,0 +1,65 @@
+Status: ready-for-agent
+
+# US-020 - Memory budget, shard slots, and dropout relocation hardening
+
+## Goal
+
+Make node capacity limits explicit and enforce them consistently when the tracker assigns, rebalances, and relocates shards after a node dropout.
+
+This is a follow-up to US-013, not a replacement. US-013 owns the coverage-first assignment and rebalance algorithm. This issue hardens the capacity contract around that algorithm: operator memory budget, maximum loaded shard slots, and relocation behavior when one node must absorb or split ranges after another node disappears.
+
+## Context
+
+Recent work added the first part of the contract:
+
+- `meshnet-node --memory MB` is registered with the tracker as `vram_bytes` when explicitly set.
+- CPU nodes without `--memory` keep the tracker default capacity, preserving old behavior.
+- `meshnet-node --max-shards N` is accepted and registered as `max_loaded_shards`.
+- Tracker registration validates `max_loaded_shards >= 1`.
+
+The current runtime still effectively has one active backend shard per node. A node may advertise `max_loaded_shards`, but the tracker does not yet use multiple shard slots in bin-packing, and the node does not yet host multiple concurrently loaded shard ranges.
+
+## Scope
+
+- Make tracker rebalance logic account for `max_loaded_shards` as a capacity multiplier or explicit shard-slot list.
+- Ensure a node is never assigned more total layers than its memory budget can support across all loaded shard slots.
+- Decide and implement the runtime behavior for multiple loaded shards:
+  - either support multiple concurrently loaded shard backends on one node, or
+  - keep one backend active and treat `max_loaded_shards` as future metadata, with tracker enforcement preventing multi-range assignment for now.
+- On heartbeat timeout, relocate the dropped node's uncovered layer range to eligible managed nodes while respecting both memory and shard-slot limits.
+- Surface the effective memory budget and shard slot count in tracker/network inspection output so operators can diagnose why a node did or did not receive a range.
+
+## Non-Goals
+
+- Do not redesign the US-013 coverage-first algorithm from scratch.
+- Do not change relay, `/ws`, or `/rpc` behavior.
+- Do not change the token/reward model.
+- Do not require public internet verification; all behavior must be locally testable.
+
+## Acceptance Criteria
+
+- Tracker stores and exposes `max_loaded_shards` for registered nodes.
+- Assignment/rebalance never exceeds:
+  - `assigned_layers_total <= floor((vram_bytes * 0.8) / bytes_per_layer_at_quant)`
+  - `assigned_range_count <= max_loaded_shards`
+- A managed node with `max_loaded_shards=1` only receives one active shard range.
+- A managed node with `max_loaded_shards=2` can absorb two non-contiguous uncovered ranges only if the node runtime supports serving both; otherwise tracker must keep assigning at most one range and document `max_loaded_shards` as reserved.
+- Dropout test: register nodes covering a model, let a middle/tail node heartbeat-expire, and assert the tracker queues `LOAD_SHARD` directives that restore full coverage without violating memory or shard-slot limits.
+- CLI test: `--memory` and `--max-shards` are reflected in the registration payload.
+- `python -m pytest tests/test_tracker_routing.py tests/test_node_startup.py` passes in the project virtualenv, aside from any pre-existing platform-specific wallet permission assertion documented in the final notes.
+
+## Implementation Notes
+
+- Existing files likely involved:
+  - `packages/node/meshnet_node/cli.py`
+  - `packages/node/meshnet_node/startup.py`
+  - `packages/node/meshnet_node/torch_server.py`
+  - `packages/tracker/meshnet_tracker/server.py`
+  - `tests/test_tracker_routing.py`
+  - `tests/test_node_startup.py`
+- Keep backward compatibility: nodes that omit `vram_bytes` default to tracker defaults; nodes that omit `max_loaded_shards` default to `1`.
+- Prefer a small internal representation for assigned ranges if multiple ranges become real, for example `assigned_shards: list[tuple[int, int]]`, while preserving `shard_start`/`shard_end` in public responses for single-range nodes.
+
+## Comments
+
+- 2026-06-30: Created after implementing the initial registration plumbing in commit `f1e4ed6` (`--memory`, `--max-shards`, tracker validation). This issue captures the remaining end-to-end behavior so it does not conflict with US-013.

packages/node/meshnet_node/cli.py

@@ -30,6 +30,7 @@ def _run_node(cfg: dict) -> None:
             advertise_host=cfg.get("advertise_host"),
             route_timeout=float(cfg.get("route_timeout", 30.0)),
             vram_mb_override=cfg.get("vram_mb_override"),
+            max_loaded_shards=int(cfg.get("max_loaded_shards", 1)),
             debug=bool(cfg.get("debug", False)),
         )
     except Exception as exc:
@@ -110,6 +111,8 @@ def _cmd_default(args) -> int:
         overrides["route_timeout"] = args.route_timeout
     if getattr(args, "memory", None) is not None:
         overrides["vram_mb_override"] = args.memory
+    if getattr(args, "max_shards", None) is not None:
+        overrides["max_loaded_shards"] = args.max_shards
     if args.debug:
         overrides["debug"] = True
 
@@ -205,6 +208,7 @@ def _cmd_start(args) -> int:
             advertise_host=getattr(args, "advertise_host", None),
             route_timeout=getattr(args, "route_timeout", 30.0),
             vram_mb_override=getattr(args, "memory", None),
+            max_loaded_shards=getattr(args, "max_shards", 1),
             debug=getattr(args, "debug", False),
         )
     except Exception as exc:
@@ -250,6 +254,8 @@ def main() -> None:
                         help="Seconds to wait for tracker route lookup (default 30)")
     parser.add_argument("--memory", type=int, metavar="MB", default=None,
                         help="Override autodetected VRAM/RAM budget in MB used for shard assignment")
+    parser.add_argument("--max-shards", type=int, metavar="N", default=None,
+                        help="Maximum shard slots this node advertises to the tracker (default 1)")
     parser.add_argument("--debug", action="store_true", help="Enable verbose node debug logging")
     parser.add_argument("--no-tui", action="store_true", help="Plain-text output (no rich dashboard)")
     parser.add_argument("--compact", action="store_true", help="Single-line status output")
@@ -283,6 +289,8 @@ def main() -> None:
                            help="Seconds to wait for tracker route lookup (default 30)")
     start_cmd.add_argument("--memory", type=int, default=None, metavar="MB",
                            help="Override autodetected VRAM/RAM budget in MB used for shard assignment")
+    start_cmd.add_argument("--max-shards", type=int, default=1, metavar="N",
+                           help="Maximum shard slots this node advertises to the tracker (default 1)")
     start_cmd.add_argument("--debug", action="store_true", help="Enable verbose node debug logging")
 
     args = parser.parse_args()

packages/node/meshnet_node/startup.py

@@ -265,6 +265,7 @@ def run_startup(
     contracts: Any | None = None,
     route_timeout: float = 30.0,
     vram_mb_override: int | None = None,
+    max_loaded_shards: int = 1,
     debug: bool = False,
 ) -> StubNodeServer | TorchNodeServer:
     """Execute the full startup sequence and return a running node server.
@@ -282,6 +283,8 @@ def run_startup(
 
     tracker_url = tracker_url.rstrip("/")
     relay_url = _discover_relay_url(tracker_url)
+    if max_loaded_shards < 1:
+        raise ValueError("--max-shards must be at least 1")
 
     # 1. Hardware detection
     if advertise_host is None and host == "0.0.0.0":
@@ -315,6 +318,11 @@ def run_startup(
         print("  WARNING: No CUDA GPU detected — running in CPU mode", flush=True)
     else:
         print(f"  GPU: {gpu_name} ({vram_mb} MB VRAM)", flush=True)
+    registration_capabilities = {
+        "max_loaded_shards": max_loaded_shards,
+    }
+    if vram_mb_override is not None or vram_mb > 0:
+        registration_capabilities["vram_bytes"] = max(0, int(vram_mb)) * 1024 * 1024
 
     # 2. Wallet
     print("Loading wallet...", flush=True)
@@ -404,6 +412,7 @@ def run_startup(
             "score": 1.0,
             "tracker_mode": (shard_start == 0),
             "managed_assignment": not user_pinned_shard,
+            **registration_capabilities,
             **relay_fields,
         }
         tracker_node_id: str | None = None
@@ -493,6 +502,7 @@ def run_startup(
             "score": 1.0,
             "tracker_mode": (assigned_shard_start == 0),
             "managed_assignment": True,
+            **registration_capabilities,
             **relay_fields,
         }
         tracker_node_id = None
@@ -593,6 +603,7 @@ def run_startup(
                 "hardware_profile": hw,
                 "wallet_address": address,
                 "score": 1.0,
+                **registration_capabilities,
                 **relay_fields,
             },
         )

packages/tracker/meshnet_tracker/server.py

@@ -259,7 +259,7 @@ class _NodeEntry:
     __slots__ = (
         "node_id", "endpoint", "shard_start", "shard_end",
         "model", "hf_repo", "num_layers", "shard_checksum", "hardware_profile", "wallet_address",
-        "score", "vram_bytes", "ram_bytes", "quantizations",
+        "score", "vram_bytes", "ram_bytes", "quantizations", "max_loaded_shards",
         "benchmark_tokens_per_sec", "quantization", "managed_assignment",
         "pending_directives", "last_heartbeat", "tracker_mode",
         "relay_addr", "cert_fingerprint", "peer_id",
@@ -285,6 +285,7 @@ class _NodeEntry:
         vram_bytes: int = DEFAULT_VRAM_BYTES,
         ram_bytes: int = DEFAULT_RAM_BYTES,
         quantizations: list[str] | None = None,
+        max_loaded_shards: int = 1,
         benchmark_tokens_per_sec: float = DEFAULT_BENCHMARK_TOKENS_PER_SEC,
         quantization: str | None = None,
         managed_assignment: bool = False,
@@ -307,6 +308,7 @@ class _NodeEntry:
         self.vram_bytes = vram_bytes
         self.ram_bytes = ram_bytes
         self.quantizations = quantizations or list(DEFAULT_QUANTIZATIONS)
+        self.max_loaded_shards = max_loaded_shards
         self.benchmark_tokens_per_sec = benchmark_tokens_per_sec
         self.quantization = quantization
         self.managed_assignment = managed_assignment
@@ -1431,13 +1433,14 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
         try:
             vram_bytes = int(body.get("vram_bytes", DEFAULT_VRAM_BYTES))
             ram_bytes = int(body.get("ram_bytes", DEFAULT_RAM_BYTES))
+            max_loaded_shards = int(body.get("max_loaded_shards", 1))
             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"})
+            self._send_json(400, {"error": "vram_bytes, ram_bytes, max_loaded_shards, and benchmark_tokens_per_sec must be numeric"})
             return
-        if vram_bytes < 0 or ram_bytes < 0 or benchmark_tokens_per_sec <= 0:
+        if vram_bytes < 0 or ram_bytes < 0 or max_loaded_shards < 1 or benchmark_tokens_per_sec <= 0:
             self._send_json(400, {"error": "capability values must be positive"})
             return
         quantizations_body = body.get("quantizations", DEFAULT_QUANTIZATIONS)
@@ -1501,6 +1504,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             vram_bytes=vram_bytes,
             ram_bytes=ram_bytes,
             quantizations=quantizations,
+            max_loaded_shards=max_loaded_shards,
             benchmark_tokens_per_sec=benchmark_tokens_per_sec,
             quantization=quantization,
             managed_assignment=managed_assignment or not explicit_shard,

tests/test_node_startup.py

@@ -442,6 +442,7 @@ def test_public_https_tracker_infers_relay_when_network_map_omits_relay_url(
 def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, capsys):
     """Real-model startup summary prints the shard range plus total model layers."""
     import meshnet_node.startup as startup_mod
+    captured_registration = {}
 
     class FakeBackend:
         total_layers = 24
@@ -465,7 +466,9 @@ def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, ca
     monkeypatch.setattr(
         startup_mod,
         "_post_json",
-        lambda _url, _payload, timeout=10.0: {"node_id": "node-test-123"},
+        lambda _url, _payload, timeout=10.0: (
+            captured_registration.update(_payload) or {"node_id": "node-test-123"}
+        ),
     )
 
     node = run_startup(
@@ -473,11 +476,15 @@ def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, ca
         model_id="Qwen/Qwen2.5-0.5B-Instruct",
         shard_start=0,
         shard_end=23,
+        vram_mb_override=6144,
+        max_loaded_shards=2,
         wallet_path=tmp_path / "wallet.json",
     )
 
     assert node.backend.total_layers == 24
     assert node.tracker_node_id == "node-test-123"
+    assert captured_registration["vram_bytes"] == 6144 * 1024 * 1024
+    assert captured_registration["max_loaded_shards"] == 2
     output = capsys.readouterr().out
     assert "Shard:        layers 0–23; 24 of 24" in output
     assert "Node ID:      node-test-123" in output

tests/test_tracker_routing.py

@@ -845,6 +845,25 @@ def test_tracker_registration_rejects_invalid_payload():
         tracker.stop()
 
 
+def test_tracker_registration_rejects_invalid_max_loaded_shards():
+    tracker = TrackerServer()
+    tracker_port = tracker.start()
+    try:
+        try:
+            _post_json(
+                f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
+                {"endpoint": "http://127.0.0.1:9001", "shard_start": 0, "shard_end": 31,
+                 "max_loaded_shards": 0, "hardware_profile": {}, "score": 1.0},
+            )
+            raise AssertionError("Expected 400 for invalid max_loaded_shards")
+        except urllib.error.HTTPError as exc:
+            assert exc.code == 400
+            body = json.loads(exc.read())
+            assert "capability values" in body["error"]
+    finally:
+        tracker.stop()
+
+
 def test_tracker_routes_only_nodes_for_requested_model():
     """A node registered for one model cannot satisfy another model's route."""
     tracker = TrackerServer(model_presets={