[verified] feat: complete Ralph task workstreams

packages/gateway/meshnet_gateway/prefill_backpressure.py

@@ -0,0 +1,130 @@
+"""Bounded, ordered prefill transfer primitives.
+
+Prefill chunks mutate the downstream shard's session cache, so they must reach a
+route in order.  This deliberately uses a serial acknowledgement window: it is
+the safe default for both current peers and old peers which do not advertise a
+windowing capability.  The configured in-flight limit is still explicit so a
+future ordered transport can widen the window without changing callers.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from threading import Event
+from typing import Callable, Iterable, TypeVar
+
+
+DEFAULT_PREFILL_CHUNK_TOKENS = 128
+DEFAULT_PREFILL_MAX_IN_FLIGHT = 1
+DEFAULT_PREFILL_MAX_CHUNK_BYTES = 8 * 1024 * 1024
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+@dataclass(frozen=True)
+class PrefillTransferLimits:
+    """Configuration for one ordered prefill seam."""
+
+    chunk_tokens: int = DEFAULT_PREFILL_CHUNK_TOKENS
+    max_in_flight: int = DEFAULT_PREFILL_MAX_IN_FLIGHT
+    max_chunk_bytes: int = DEFAULT_PREFILL_MAX_CHUNK_BYTES
+
+    @property
+    def effective_in_flight(self) -> int:
+        """Current peers require ordered session-cache mutation, hence one ack."""
+        return 1
+
+    @property
+    def max_buffered_bytes(self) -> int:
+        """Hard accounting bound, including any future wider ack window."""
+        return self.max_chunk_bytes * self.max_in_flight
+
+    @classmethod
+    def from_env(cls) -> "PrefillTransferLimits":
+        # MESHNET_CHUNK_TOKENS was the pre-DIP-007 name.  Keep it as a fallback
+        # so existing deployments retain their chunk shape while upgrading.
+        return cls(
+            chunk_tokens=_positive_env(
+                "MESHNET_PREFILL_CHUNK_TOKENS",
+                _positive_env("MESHNET_CHUNK_TOKENS", DEFAULT_PREFILL_CHUNK_TOKENS),
+            ),
+            max_in_flight=_positive_env(
+                "MESHNET_PREFILL_MAX_IN_FLIGHT", DEFAULT_PREFILL_MAX_IN_FLIGHT,
+            ),
+            max_chunk_bytes=_positive_env(
+                "MESHNET_PREFILL_MAX_CHUNK_BYTES", DEFAULT_PREFILL_MAX_CHUNK_BYTES,
+            ),
+        )
+
+
+class BoundedPrefillSender:
+    """Send lazily-produced chunks with bounded ownership and ordered acks."""
+
+    def __init__(self, limits: PrefillTransferLimits) -> None:
+        self.limits = limits
+        self.buffered_bytes = 0
+        self.peak_buffered_bytes = 0
+        self.in_flight = 0
+        self.peak_in_flight = 0
+        self.closed = False
+
+    def send(
+        self,
+        chunks: Iterable[T],
+        *,
+        body_size: Callable[[T], int],
+        forward: Callable[[T], R],
+        cancelled: Event | None = None,
+    ) -> list[R]:
+        """Forward chunks in source order, releasing each body after its ack.
+
+        ``forward`` is synchronous by design: a slow consumer therefore blocks
+        production of the next chunk instead of accumulating an unbounded queue.
+        Every retained body is dropped on cancellation or route failure.
+        """
+        results: list[R] = []
+        try:
+            for chunk in chunks:
+                if self.closed or (cancelled is not None and cancelled.is_set()):
+                    break
+                size = body_size(chunk)
+                if size < 0:
+                    raise ValueError("prefill chunk size cannot be negative")
+                if size > self.limits.max_chunk_bytes:
+                    raise ValueError(
+                        f"prefill chunk exceeds {self.limits.max_chunk_bytes} byte limit"
+                    )
+                self.buffered_bytes += size
+                self.in_flight += 1
+                self.peak_buffered_bytes = max(self.peak_buffered_bytes, self.buffered_bytes)
+                self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
+                try:
+                    results.append(forward(chunk))
+                finally:
+                    # Do not retain a body while waiting for the next chunk.
+                    self.buffered_bytes -= size
+                    self.in_flight -= 1
+        except BaseException:
+            self.close()
+            raise
+        return results
+
+    def close(self) -> None:
+        """Release accounting after cancellation or route failure.
+
+        The sender deliberately owns no queued chunk references; callers must
+        discard their iterator on close rather than trying to drain it.
+        """
+        self.buffered_bytes = 0
+        self.in_flight = 0
+        self.closed = True
+
+
+def _positive_env(name: str, default: int) -> int:
+    try:
+        value = int(os.environ.get(name, default))
+    except (TypeError, ValueError):
+        return default
+    return value if value > 0 else default

packages/gateway/meshnet_gateway/server.py

@@ -14,6 +14,8 @@ import urllib.request
 import uuid
 from typing import Any
 
+from .prefill_backpressure import BoundedPrefillSender, PrefillTransferLimits
+
 _STUB_HIDDEN_DIM = 64
 _STUB_DTYPE = "bfloat16"
 _WIRE_VERSION = "2"
@@ -653,24 +655,27 @@ def _last_message_content(messages: object) -> str:
 
 def _run_binary_pipeline(route: list[str], prompt: str, timeout: float = 5.0) -> list[_BinaryActivation]:
     session = str(uuid.uuid4())
-    chunk_token_count = _chunk_token_count()
+    limits = PrefillTransferLimits.from_env()
+    chunk_token_count = limits.chunk_tokens
     total_tokens = max(1, _prompt_token_count(prompt))
     chunk_total = max(1, (total_tokens + chunk_token_count - 1) // chunk_token_count)
-    responses: list[_BinaryActivation] = []
 
-    for chunk_index in range(chunk_total):
-        remaining_tokens = total_tokens - (chunk_index * chunk_token_count)
-        seq_len = min(chunk_token_count, remaining_tokens)
-        activation = _BinaryActivation(
-            body=_make_stub_binary_activation([1, seq_len, _STUB_HIDDEN_DIM], _STUB_DTYPE),
-            shape=[1, seq_len, _STUB_HIDDEN_DIM],
-            dtype=_STUB_DTYPE,
-            session=session,
-            chunk_index=chunk_index,
-            chunk_total=chunk_total,
-            encoding=_preferred_binary_encoding(),
-            headers={},
-        )
+    def chunks():
+        for chunk_index in range(chunk_total):
+            remaining_tokens = total_tokens - (chunk_index * chunk_token_count)
+            seq_len = min(chunk_token_count, remaining_tokens)
+            yield _BinaryActivation(
+                body=_make_stub_binary_activation([1, seq_len, _STUB_HIDDEN_DIM], _STUB_DTYPE),
+                shape=[1, seq_len, _STUB_HIDDEN_DIM],
+                dtype=_STUB_DTYPE,
+                session=session,
+                chunk_index=chunk_index,
+                chunk_total=chunk_total,
+                encoding=_preferred_binary_encoding(),
+                headers={},
+            )
+
+    def forward(activation: _BinaryActivation) -> _BinaryActivation:
         for hop_index, node_url in enumerate(route):
             activation = _post_binary_forward(
                 f"{node_url}/forward",
@@ -678,17 +683,15 @@ def _run_binary_pipeline(route: list[str], prompt: str, timeout: float = 5.0) ->
                 hop_index=hop_index,
                 timeout=timeout,
             )
-        responses.append(activation)
-    return responses
+        return activation
+
+    # Each completed response is retained only for the gateway's existing test
+    # diagnostic surface.  At every hop the sender owns one chunk at a time.
+    return BoundedPrefillSender(limits).send(chunks(), body_size=lambda item: len(item.body), forward=forward)
 
 
 def _chunk_token_count() -> int:
-    raw_value = os.environ.get("MESHNET_CHUNK_TOKENS", "128")
-    try:
-        value = int(raw_value)
-    except ValueError:
-        return 128
-    return value if value > 0 else 128
+    return PrefillTransferLimits.from_env().chunk_tokens
 
 
 def _prompt_token_count(prompt: str) -> int:
@@ -737,11 +740,23 @@ def _post_binary_forward(
 
     encoding = response_headers.get("x-meshnet-encoding")
     raw_body = _decompress_body(response_body, encoding)
-    shape = _parse_shape(response_headers["x-meshnet-shape"])
-    dtype = response_headers["x-meshnet-dtype"]
-    session = response_headers["x-meshnet-session"]
-    chunk_index = int(response_headers["x-meshnet-chunk-index"])
-    chunk_total = int(response_headers["x-meshnet-chunk-total"])
+    # Legacy single-chunk peers returned only the body before chunk metadata
+    # was added.  Treat absent metadata as the caller's single chunk, while a
+    # peer which partially implements the new protocol still fails closed.
+    if not response_headers.get("x-meshnet-shape"):
+        if activation.chunk_total != 1:
+            raise ValueError("legacy peer cannot acknowledge multi-chunk prefill")
+        shape = activation.shape
+        dtype = activation.dtype
+        session = activation.session
+        chunk_index = activation.chunk_index
+        chunk_total = activation.chunk_total
+    else:
+        shape = _parse_shape(response_headers["x-meshnet-shape"])
+        dtype = response_headers["x-meshnet-dtype"]
+        session = response_headers["x-meshnet-session"]
+        chunk_index = int(response_headers["x-meshnet-chunk-index"])
+        chunk_total = int(response_headers["x-meshnet-chunk-total"])
     if session != activation.session:
         raise ValueError("binary activation response changed session")
     if chunk_index != activation.chunk_index or chunk_total != activation.chunk_total: