feat: add node startup flow

packages/tracker/meshnet_tracker/server.py

@@ -4,7 +4,8 @@ HTTP API contract:
 - POST /v1/nodes/register
   Request JSON: {
     "endpoint": "http://host:port", "shard_start": int, "shard_end": int,
-    "hardware_profile": object, "score": number optional
+    "hardware_profile": object, "wallet_address": str optional,
+    "score": number optional
   }
   Response 200: {"node_id": str}
   Response 400: {"error": str}
@@ -32,7 +33,7 @@ DEFAULT_MODEL_PRESETS: dict[str, dict] = {
 class _NodeEntry:
     __slots__ = (
         "node_id", "endpoint", "shard_start", "shard_end",
-        "hardware_profile", "score", "last_heartbeat",
+        "hardware_profile", "wallet_address", "score", "last_heartbeat",
     )
 
     def __init__(
@@ -42,6 +43,7 @@ class _NodeEntry:
         shard_start: int,
         shard_end: int,
         hardware_profile: dict,
+        wallet_address: str | None,
         score: float,
     ) -> None:
         self.node_id = node_id
@@ -49,6 +51,7 @@ class _NodeEntry:
         self.shard_start = shard_start
         self.shard_end = shard_end
         self.hardware_profile = hardware_profile
+        self.wallet_address = wallet_address
         self.score = score
         self.last_heartbeat: float = time.monotonic()
 
@@ -146,6 +149,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
         parsed = urllib.parse.urlparse(self.path)
         if parsed.path == "/v1/route":
             self._handle_route(parsed)
+        elif parsed.path == "/v1/nodes/assign":
+            self._handle_assign(parsed)
         else:
             self.send_response(404)
             self.end_headers()
@@ -180,6 +185,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
         if not isinstance(hardware_profile, dict):
             self._send_json(400, {"error": "hardware_profile must be an object"})
             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"})
+            return
 
         node_id = str(uuid.uuid4())
         entry = _NodeEntry(
@@ -188,6 +197,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             shard_start=shard_start,
             shard_end=shard_end,
             hardware_profile=hardware_profile,
+            wallet_address=wallet_address,
             score=score,
         )
         with server.lock:
@@ -207,6 +217,81 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
             entry.last_heartbeat = time.monotonic()
         self._send_json(200, {})
 
+    def _handle_assign(self, parsed: urllib.parse.ParseResult):
+        """Return an optimal shard assignment for a node given its hardware profile.
+
+        Query params:
+          model   — model preset name (default: first preset)
+          device  — "cuda" | "cpu"
+          vram_mb — integer VRAM in MB (0 for CPU)
+
+        The greedy strategy: find the first gap in current layer coverage
+        and assign it.  If no gap exists, assign the full model range so the
+        node provides redundant coverage.
+        """
+        server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
+        params = urllib.parse.parse_qs(parsed.query)
+
+        model_list = params.get("model")
+        if not model_list:
+            model = next(iter(server.model_presets), None)
+            if model is None:
+                self._send_json(503, {"error": "no model presets configured"})
+                return
+        else:
+            model = model_list[0]
+
+        preset = server.model_presets.get(model)
+        if preset is None:
+            self._send_json(404, {"error": f"unknown model preset: {model!r}"})
+            return
+
+        required_start: int = preset["layers_start"]
+        required_end: int = preset["layers_end"]
+
+        with server.lock:
+            self._purge_expired_nodes()
+            alive = list(server.registry.values())
+
+        device = params.get("device", ["cpu"])[0]
+        try:
+            vram_mb = int(params.get("vram_mb", ["0"])[0])
+        except ValueError:
+            vram_mb = 0
+        max_layers = required_end - required_start + 1
+        if device != "cuda" or vram_mb < 8192:
+            max_layers = min(max_layers, 16)
+
+        # Collect covered intervals sorted by start layer.
+        covered = sorted(
+            [(n.shard_start, n.shard_end) for n in alive],
+            key=lambda t: t[0],
+        )
+
+        # Walk from required_start to find the first uncovered layer.
+        gap_start = required_start
+        for s, e in covered:
+            if s <= gap_start:
+                gap_start = max(gap_start, e + 1)
+            else:
+                break  # gap found before this node
+
+        if gap_start > required_end:
+            # Full coverage exists — assign the full range for redundancy.
+            shard_start = required_start
+            shard_end = min(required_end, shard_start + max_layers - 1)
+        else:
+            shard_start = gap_start
+            shard_end = min(required_end, shard_start + max_layers - 1)
+
+        self._send_json(200, {
+            "shard_start": shard_start,
+            "shard_end": shard_end,
+            "model": model,
+            "model_layers_end": required_end,
+            **({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}),
+        })
+
     def _handle_route(self, parsed: urllib.parse.ParseResult):
         server: _TrackerHTTPServer = self.server  # type: ignore[assignment]
         params = urllib.parse.parse_qs(parsed.query)