[verified] feat: complete Ralph task workstreams

packages/node/meshnet_node/server.py

@@ -8,6 +8,7 @@ import urllib.parse
 from pathlib import Path
 
 from .downloader import compute_shard_checksum, write_shard_archive
+from .activation_compression import CompressionPolicies, compress_activation, decompress_activation
 
 # Binary activation wire format (contract for all shard nodes):
 # POST /forward with raw tensor bytes in the body and tensor/session/chunk
@@ -21,6 +22,7 @@ _DTYPE_SIZES = {
     "bfloat16": 2,
     "float32": 4,
 }
+_COMPRESSION_POLICIES = CompressionPolicies()
 
 
 def _make_stub_binary_activation(shape: list[int], dtype: str) -> bytes:
@@ -42,16 +44,7 @@ def _parse_shape(value: str | None) -> list[int]:
 
 
 def _decompress_body(body: bytes, encoding: str | None) -> bytes:
-    if not encoding:
-        return body
-    if encoding != "zstd":
-        raise ValueError("unsupported X-Meshnet-Encoding")
-    import zstandard as zstd
-
-    try:
-        return zstd.ZstdDecompressor().decompress(body)
-    except zstd.ZstdError as exc:
-        raise ValueError("invalid zstd activation body") from exc
+    return decompress_activation(body, encoding).body
 
 
 def _compress_body(body: bytes, encoding: str | None) -> bytes:
@@ -307,7 +300,14 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
             server.received_activations = True
 
         raw_payload = _make_stub_binary_activation(shape, dtype)
-        payload = _compress_body(raw_payload, encoding)
+        route_condition = self.headers.get("X-Meshnet-Compression-Route", "lan")
+        phase_condition = self.headers.get("X-Meshnet-Cache", "prefill")
+        if phase_condition not in {"prefill", "decode"}:
+            phase_condition = "prefill"
+        compression = compress_activation(
+            raw_payload, _COMPRESSION_POLICIES.for_condition(route_condition, phase_condition),
+        )
+        payload = compression.body
         self.send_response(200)
         self.send_header("Content-Type", "application/octet-stream")
         self.send_header("Content-Length", str(len(payload)))
@@ -317,8 +317,8 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
         self.send_header("X-Meshnet-Session", session)
         self.send_header("X-Meshnet-Chunk-Index", chunk_index)
         self.send_header("X-Meshnet-Chunk-Total", chunk_total)
-        if encoding:
-            self.send_header("X-Meshnet-Encoding", encoding)
+        if compression.encoding:
+            self.send_header("X-Meshnet-Encoding", compression.encoding)
         if server.is_last_shard:
             self.send_header("X-Meshnet-Stub-Response-Prefix", server.response_prefix)
         self.end_headers()