popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[verified] feat: complete Ralph task workstreams

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-12 11:17:03 +03:00
parent 9a1b15c020
commit 377346c301
37 changed files with 5862 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
packages/gateway/meshnet_gateway/server.py
View File

@@ -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: