feat: add p2p shard swarm

packages/node/meshnet_node/server.py

@@ -5,6 +5,10 @@ import http.server
 import json
 import struct
 import threading
+import urllib.parse
+from pathlib import Path
+
+from .downloader import compute_shard_checksum, write_shard_archive
 
 # Activation wire format (contract for all shard nodes):
 # JSON object: {"shape": [B, S, D], "dtype": "float32",
@@ -36,12 +40,16 @@ class _StubHTTPServer(http.server.HTTPServer):
         shard_end: int,
         is_last_shard: bool,
         response_prefix: str,
+        model: str,
+        shard_path: Path | None,
     ):
         super().__init__(addr, handler)
         self.shard_start = shard_start
         self.shard_end = shard_end
         self.is_last_shard = is_last_shard
         self.response_prefix = response_prefix
+        self.model = model
+        self.shard_path = shard_path
         self.received_activations: bool = False
 
 
@@ -56,6 +64,43 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
             self.send_response(404)
             self.end_headers()
 
+    def do_GET(self):
+        parsed = urllib.parse.urlparse(self.path)
+        if parsed.path == "/v1/shards/download":
+            self._handle_shard_download(parsed)
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+    def _handle_shard_download(self, parsed: urllib.parse.ParseResult):
+        server: _StubHTTPServer = self.server  # type: ignore[assignment]
+        params = urllib.parse.parse_qs(parsed.query)
+        model = params.get("model", [""])[0]
+        try:
+            shard_start = int(params.get("shard_start", [""])[0])
+            shard_end = int(params.get("shard_end", [""])[0])
+        except ValueError:
+            self.send_response(400)
+            self.end_headers()
+            return
+
+        if (
+            server.shard_path is None
+            or not server.shard_path.exists()
+            or model != server.model
+            or shard_start != server.shard_start
+            or shard_end != server.shard_end
+        ):
+            self.send_response(404)
+            self.end_headers()
+            return
+
+        self.send_response(200)
+        self.send_header("Content-Type", "application/x-tar")
+        self.send_header("X-Meshnet-Shard-Checksum", compute_shard_checksum(server.shard_path))
+        self.end_headers()
+        write_shard_archive(server.shard_path, self.wfile)
+
     def _handle_infer(self):
         server: _StubHTTPServer = self.server  # type: ignore[assignment]
         length = int(self.headers.get("Content-Length", 0))
@@ -102,6 +147,8 @@ class StubNodeServer:
         shard_end: int = 31,
         is_last_shard: bool = True,
         response_prefix: str = "stub response to:",
+        model: str = "stub-model",
+        shard_path: Path | None = None,
     ):
         self._host = host
         self._requested_port = port
@@ -111,6 +158,8 @@ class StubNodeServer:
         self._shard_end = shard_end
         self._is_last_shard = is_last_shard
         self._response_prefix = response_prefix
+        self._model = model
+        self._shard_path = shard_path
         self._server: _StubHTTPServer | None = None
         self._thread: threading.Thread | None = None
         self.port: int | None = None
@@ -131,6 +180,8 @@ class StubNodeServer:
             self._shard_end,
             self._is_last_shard,
             self._response_prefix,
+            self._model,
+            self._shard_path,
         )
         self.port = self._server.server_address[1]
         self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)