From f1e4ed6a329ce8d7e9cc0d29d091dc320da054ad Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Tue, 30 Jun 2026 21:17:07 +0200 Subject: [PATCH] Wire node memory and shard slot capabilities --- packages/node/meshnet_node/cli.py | 8 ++++++++ packages/node/meshnet_node/startup.py | 11 +++++++++++ packages/tracker/meshnet_tracker/server.py | 10 +++++++--- tests/test_node_startup.py | 9 ++++++++- tests/test_tracker_routing.py | 19 +++++++++++++++++++ 5 files changed, 53 insertions(+), 4 deletions(-) diff --git a/packages/node/meshnet_node/cli.py b/packages/node/meshnet_node/cli.py index ea77e4f..4250677 100644 --- a/packages/node/meshnet_node/cli.py +++ b/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() diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index b791fd1..753b251 100644 --- a/packages/node/meshnet_node/startup.py +++ b/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, }, ) diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 3344638..8f5eef5 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/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, diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index d560e84..e77a554 100644 --- a/tests/test_node_startup.py +++ b/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 diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index bd1d3fe..4bb8b2c 100644 --- a/tests/test_tracker_routing.py +++ b/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={