131 lines
4.6 KiB
Python
131 lines
4.6 KiB
Python
"""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
|