[dghosef]

What this does

This is part 3 of 6 of native full duplex (#186): client streams. A model can now declare a live input stream that the client feeds into a running request:

Connection(from_partition=CLIENT, to_partition="Encoder", edge_name="audio_in",
           chunk_policy_factory=lambda: FixedChunkPolicy(chunk_size=1))

There is one request per conversation:

• The client opens a call over /generate/ws, appends chunks while it runs, and then ends its input.
• The consuming node reads the stream through the normal StreamBuffer and chunk policy, the same as a stream from another partition.
• There are no sessions, and append behaviour stays inside the node.

This PR builds on #766 (the preprocess thread's non-blocking sends). It makes no engine change. It makes two worker changes:

• A producer_done now ends only the edges it counts. Every multi-partition model sees this (Orpheus, Qwen3-Omni, Qwen3-TTS, Chatterbox, Cosmos3 windowed). See "Changes every model sees".
• The worker acks a streaming edge that has no buffer, without reading it.

How it works

• Topology (mstar/streaming/topology.py):
  • CLIENT is a reserved producer. It is not a partition.
  • New: PartitionTopology.client_connections().
• Model hook: Model.process_stream_chunk(chunk: ClientChunk) turns one append into stream items, keyed by client edge name. A ClientChunk carries:
  • the append's text;
  • its loaded files (tensors and input_metadata by modality, and each file's target streams);
  • the append's own model_kwargs, and the opening's (request_kwargs).
  • The hook is stateless. The consuming node keeps its own state.
  • If process_prompt returns tensors under a client edge's name, they become the stream's first chunk. So a one-shot upload is a stream of one chunk.
• API server (entrypoint.py): /generate/ws serves concurrent live calls.
  • Messages:
    • {"request_id", "append": {"text"?, "files"?: [{"name", "data", "stream"?}], "model_kwargs"?}} feeds a call. It takes the same fields an opening does. "append": [files] is shorthand for {"files": [...]}.
    • end_input is either true (every client stream of the call) or a list of client stream names (only those). It can be sent alone, on an append, or on the opening message for a one-shot.
    • A file that names a stream the model doesn't have gets a 400, on an append or on the opening. A text-only append that produces no stream item gets a 400.
    • If the model routes items to a stream the client already ended, the client gets the late-append reply. Those items are neither charged nor sent.
    • A null field counts as absent.
  • Ownership:
    • A live call needs a request_id. Only the socket that opened the call can append to it.
    • Opening a request under a live call's id gets a 409, and the call goes on.
    • A closed socket aborts its calls.
  • Replies:
    • A rejected append ends the call with its error (status 400) and no finish.
    • An append after the input or the call has ended gets {"request_id", "append": true, "error"}, and the call goes on.
    • APIServer.stream_input returns False for a call that is not taking input.
  • Input budget: each call has a lifetime cap on the file bytes it may send (max_call_input_mb, docs/serving.rst).
  • Timeout: input resets the request timeout.
  • Python frontend only: the Rust frontend serves no /generate/ws.
• Preprocess thread (data_worker.py):
  • Appends: it loads each append, stores its tensors pinned, and sends CLIENT_STREAM_INPUT to the conductor.
  • Queue: one FIFO, as on main. A pass takes up to 32 appends (MAX_APPENDS_PER_PASS) in arrival order. It stops at the first new request, so output delivery gets its turn between requests.
  • Pending check: before it processes a queued opening or append, the thread asks the API server whether that submission is still pending (APIServer.is_pending: held, and not finished, failed or aborted).
    • If it is not, the thread drops the item without sending it, and deletes an append's files.
    • Keys are unique per submission since [API server][Bug fix] A reused request_id no longer gets the previous request's teardown #519, so the check needs no marker in the queue.
    • An abort can land after the check. It then goes out behind what the thread just sent, from the same thread and outbox. The conductor then tears the request down as usual.
  • Failures: if an append fails to load or process, its call ends with that failure's status and the request is aborted (main's _fail_request abort, [Data worker][API server][Bug fix] Abort a request when its output fails, and stop at the first error #471).
  • Media limits: appends go through the same loader as an opening's files. So main's per-model file counts and audio-length charges apply to each append.
• Conductor:
  • It routes client chunks to the consuming partition's worker. It counts them per worker, so that the end of input reports how many items each worker was sent ([Conductor][Bug fix] Fix three partition bookkeeping slips: token counts, persist order, stream edges #616's producer_done_items).
  • At max_concurrent_requests: an opening that leaves any client stream open (a live call) is refused with a 503, instead of waiting for a slot. A call that waited would start, and stay, as far behind its client as it waited. A one-shot still queues. The opening's own client input (its first chunk and the streams it ends) rides on NEW_REQUEST (NewRequestConductor.client_input) and goes out after admission, so nothing is held for a queued request.
  • It keeps failures separate per request.
  • It drains workers at completion when the model has client streams.
  • It discovers client streams after get_worker_graphs, because bagel shapes its graphs from the deployment there.
  • At startup it checks that each client stream has exactly one consumer, and that the model implements process_stream_chunk. A partition may take several client streams. Each one ends on its own (per-edge producer_done).
• Per-stream limits (ClientStreamLimits on a client Connection, mstar/streaming/topology.py):
  • max_items, in the model's units (item_size / unit). It is counted in the data worker before anything is stored. A call that goes past it ends with a 400 that names the limit.
  • A deployment may lower a declared limit, or set one the model leaves open, under the serving config's client_streams:. It can't raise or remove one, because buffers are sized from the model's own bound.
• Worker:
  • A producer_done ends only the edges it counts. Before, it ended every stream from that producer. That also closed a co-located consumer's sibling stream, with a count of 0.
  • A streaming edge that has no buffer on this worker is acked without being read.
• Docs:
  • docs/clients.rst: the protocol.
  • docs/adding_models.rst: declaring a client stream.
  • docs/serving.rst: max_call_input_mb and client_streams: limits.
  • docs/environment_variables.rst: the Rust frontend's limits.

Changes every model sees

• Pending check: a request that is aborted while it waits for preprocessing is now dropped without being sent. Main preprocessed it and sent it to the conductor after its abort, and relied on the conductor's abort tombstone to drop it.
• Queue: the thread can take several queued items per pass. Without client streams, the queue only holds new requests, so it still takes one per pass, as on main.
• Streaming handler: it lets go of the chunks it has delivered (PendingRequest.consumed_chunks is gone). With --log-stats, the per-modality output counts come from what it delivered.
• producer_done: it ends only the edges it counts. The conductor is its only sender. The conductor's producer_done_items names every edge between the two partitions that it ends, with 0 for an edge that carried nothing. So the only streams that no longer end are another consumer partition's streams from the same producer on the same worker. The old code closed those with a count of 0.
• Reads done: finished_reading(drained=True) for a request the thread is already draining now waits for in-flight reads before READS_DONE.
• /generate/ws routing: messages are routed by shape (append vs opening). Null fields count as absent. Replies to refusals are sent from tasks, so they never hold up reading.

How it was tested

Rebased onto main d519553db (one commit). The rebase adapted to main's per-edge connection keys and item counts (#616), output-failure aborts (#471), delivery accounting (#551) and media caps (#644). Stub models in main's newer data-worker tests gained get_partition_topology, which the thread now reads when it is built.

On this PR's head e741cfe90 (with #766's fix merged in): CPU only, on a 24-vCPU Xeon Platinum 8488C, Python 3.12.13, torch 2.12.1+cu132, pytest 9.1.1, with mstar_rust built from this tree.

• ruff check .: clean.
• CPU Core (OMP_NUM_THREADS=2 python -m pytest <CPU_TEST_PATHS> -p xdist -n 4): 6050 passed, 3 failed, 612 skipped. Main abdfd5802 in the same environment: 6061 / 3 / 612, with the same 3 failures.
  • The 3 failures come from the box, not the code. test_serving_metrics.py::test_refusals_count_from_when_metrics_are_turned_on needs prometheus_client, which [all] doesn't install (CI installs [metrics]). The two test_eval_compare_multigpu.py::test_dump_stacks_retries_with_sudo_only_when_asked cases find no processes under nohup + xdist. All three pass on main with the environment fixed.
  • test/modular/test_piecewise_config_signature.py ran outside xdist. It parametrizes over a set, so xdist workers collect it in different orders (on main too). It passes.
  • Main's count is higher because main is 26 commits ahead of this PR's base.
• The client-stream conductor tests: 52 passed, 0 skipped. The Rust cases ran.
• test/rust/test_rust_wire_frames.py and test/rust/test_runtime_parity.py: 193 passed, 1 skipped. The skip is test_a_finished_worker_graph_asks_for_request_info[python]: only the Rust runtime computes that flag, by design.
• [Data worker][Bug fix] Keep a torn-down request out of the drain sets #481's tests in test_data_worker_late_release.py run as on main (a thread with no is_pending), next to variants with the API server's pending set wired in.
• CI runs only on PRs into main, so it runs on this PR once the PRs below it merge.

On d566acd52 (a live call refused at the cap; max_pending_items removed): CPU only, on macOS (arm64), Python 3.12.10, torch 2.12.1, without mstar_rust.

• ruff check .: clean.
• The client-stream tests (test_client_stream_api.py, test_client_stream_conductor.py): 114 passed.
• The 152 test/modular files that touch the changed modules, each in its own process, on fb9056948 and on d566acd52: every file has the same outcome on both. The failures they share come from the Mac (no CUDA or triton, IPC socket paths longer than macOS allows). Run as one process, the suite segfaults partway through, on both.
• The refusal fails without the fix. With d566acd52's conductor test run on fb9056948 (minus the new client_input field), a live call at the cap is queued, not refused: assert 'b' not in ... and [NewRequestConductor(request_id='b', ...)] == [].

GPU, not rerun on this head. The latest GPU run is of the whole stack through #770, at its previously pushed head. That is before d566acd52 and before #766's ff5002e7e. The load run below used 36 slots for 32 live calls, so no live call met the cap. One H100 80GB HBM3, torch 2.12.1+cu132 (CUDA 13.2), against main d519553db, with the correctness gate harness correctness.py, kept outside the repo.

• Correctness gate (one-shot vs main, live vs one-shot, batching up to 32 calls, graphs off, edge cases): every row that ran passes, except one. One-shot text vs main on interruptions differs at frame 329 (, vs .). This is a bf16 near-tie, and the stack's output is unchanged from the previous revision's run (details in nemotron_duplex: live input over a client stream, cache-aware streaming encoder (full duplex 4/6) #594). Not run this time: encoder parity, served vs standalone, the cuFFT check, listening.
• Edge rows (all pass):
  • The append record and its shorthand stream end to end.
  • A text append to nemotron_duplex ends the call with a 400.
  • end_input: ["audio_in"] ends the call like true.
  • An unknown stream on an append or on an opening gets a 400.
  • A live call past 160 s gets a 400 naming client_streams.audio_in.max_items.
• Load: 32 paced live calls, plus a 20 s upload every 2 s, with 36 slots so uploads share steps with live calls. 32/32 calls and 21/21 uploads finish with no error, and every reply has words.

Earlier revisions

• Previous revision, after the rebase, on one H200 (torch 2.11.0+cu128, CUDA 12.8): the whole stack through nemotron_duplex: per-frame codec in the talker's step (full duplex 6/6) #770 (321763dd8), against main d519553db, with the same gate harness.
  • Every required row passes except two, and neither tests the stack:
    • One-shot text vs main on 2 of 3 clips. Main's own output moved at near-tie frames, while the stack's is unchanged.
    • The cuFFT plan-eviction check. It exercises raw torch.stft / irfft, and the stack's captured regions run no FFT.
  • Passing rows include: batching up to 32 calls; graphs on vs off identical; the 11 protocol edge cases (hang up before any audio, append after end_input, append from another socket, the input budget, the session cap, abort mid-call and recovery, 409 on a live id, null fields, upload channel / rate / length refusals).
  • Live calls: 0 failed over 5 clips × 6 calls and a 1–32 session sweep.
• Previous revision on the H100: the CPU Core run on the box failed exactly main's 4 environment tests, with 6135 passed against main's 5887.
• Before the rebase:
  • uvx ruff@0.16.0 check .: clean.
  • CPU core in the CI image (python -m pytest @test/cpu-core.txt, before main dropped the list): 3004 passed, 255 skipped, 21 failed. The 21 are exactly the set main fa37d227c fails in the same container (test_e2e_validators, test_flex_attention_resource, test_waypoint_*).
  • The API, worker and conductor modules outside CI (39 files plus test_whisper_kv_sizing.py): 502 passed, 2 failed. The 2 are the ones main fails (test_graph_runtime_factory.py::test_rust_refuses_the_pyzmq_communicator and test_kv_prefix_root.py::test_only_a_node_whose_output_is_cached_under_a_digest_is_compiled_dynamic).

Performance

The path every model takes. One H200, torch 2.11.0+cu128, MSTAR_SHM_ARENA=1. The whole stack through #770 at 321763dd8, against main d519553db. Fresh server processes, and main run twice to show its own spread. Harness perf_final.py, phase regress_llm, kept outside the repo. Qwen3.5-0.8B, stack vs main:

• /generate c=32: 3960 / 4001 tok/s vs main 3952–4070.
• /generate c=64: 5441 / 4984 vs 4961–5169.
• /generate/ws with 32 sockets: 3753 / 3729 vs 3645–3780.
• No failed request.

Not rerun after d566acd52. For a model without client streams it adds one unset field to NEW_REQUEST.

Next steps

• Known problem: every client chunk goes through the conductor. A follow-up PR will fix this shortly.
  • Streams between partitions go straight from the producing worker to the consuming one. A client chunk instead goes from the API server to the preprocess thread. The thread stores it and sends a CLIENT_STREAM_INPUT control message to the conductor, which routes it to the consuming worker.
  • So the conductor is on the data path. That is one control message per call per append (about 400 a second at 32 calls with 80 ms appends), on top of the per-append ClientAppend and the file round trip below.
  • It keeps up at the scales measured here (32 live calls on one H200). But it puts per-chunk work on the one process every request shares (AGENTS.md invariant 6).
  • The follow-up sends client chunks straight to the consuming partition's worker, as streams between partitions already go. The conductor then handles only the call's start and end.
• Live input shares the preprocess thread with request decoding. The API server has one preprocess thread. It decodes new requests' media as well as live appends. So a burst of large uploads (images, video) queued ahead of a live call's append delays that append by the uploads' preprocess time. My estimate is 0.1–0.4 s. That is not measured on GPU. Fixes, in order of size:
  1. Decode a new request's media in a separate worker pool before it enters the queue, so the shared queue holds only cheap items.
  2. Several preprocess workers, each call pinned to one, so a call's appends stay in order on one worker.
  3. Frontend scale-out. The Rust frontend has no /generate/ws yet, so live calls need the Python frontend today.
• Timeouts for live calls: today input resets the request timeout, and output does not. A call that only listens could also want an idle-input timeout (Q3).
• WebSocket default mode: should a /generate/ws opening to a client-stream model, without end_input, be live by default? (Q5. Today it is, and a one-shot must say end_input: true.)
• Models that end a call mid-stream: the data worker only learns of the end through the API server. So an append can be processed in that window, which is one conductor round trip (Q6).
• Worker failure statuses: a per-request failure raised in a worker carries no client status, so it reaches the client as a 500. Session caps no longer hit this (they are stream limits now, with a 400), but other worker-side refusals still do. This needs a system change to FailRequests.
• Chunk policy:
  • We need a policy that releases everything buffered, with a cap and a concatenating collate. Today one append is one item. So appends shorter than the consumer's step (80 ms for nemotron_duplex) make a call fall behind.
  • Until then, the docs ask for appends of 80 ms or more.
• Input handling costs:
  • Appends are awaited in order inside the socket's receive loop. So one slow append decode holds up that socket's other calls.
  • Each append makes a round trip through a file.
  • Neither has been measured as a cost.

Stack

1. Keep a fresh request's final stream chunk in a speculative step (full duplex 1/6) #769: worker: keep a fresh request's final stream chunk in a speculative step.
2. Non-blocking sends from the preprocess thread (full duplex 2/6) #766: the preprocess thread's non-blocking sends.
3. This PR: client streams.
4. nemotron_duplex: live input over a client stream, cache-aware streaming encoder (full duplex 4/6) #594: nemotron_duplex live input and a cache-aware streaming encoder.
5. nemotron_duplex: live calls at one-shot cost (full duplex 5/6) #595: duplex performance.
6. nemotron_duplex: per-frame codec in the talker's step (full duplex 6/6) #770: the per-frame codec in the talker's step.

🤖 Generated with Claude Code

[mstar-review]

Automated review

🟡 No findings on a1bf86a, light review

The previous run stopped before posting any findings, so this is effectively the first review.

Description

• The opening says "The only worker change is that the worker acks, unread, a streaming edge that has no buffer", but the Worker bullet further down (and worker.py:1085-1089) also changes producer_done: it now ends only the edges it counts. That second change affects every multi-partition model (Orpheus, Qwen3-Omni, Qwen3-TTS, Chatterbox, Cosmos3 windowed), so the summary sentence should name it too.
• The CPU Core and correctness results for the current revision come from the whole stack through nemotron_duplex: per-frame codec in the talker's step (full duplex 6/6) #770, not from this PR's head on its own. The description also says this PR's CPU suite has not run since the rebase. No test result is reported for this commit by itself.
• The Tests section spends most of its length on the previous revision's H200 runs and per-frame near-tie analysis. Those are hard to map to this commit and bury the current revision's results.

Questions

• refuse under max_pending_items compares token_count against consumed_count, and consumed_count only moves when a WorkerGraphsDone reports stream_tokens_consumed. Does a consumer that reads its client stream inside a worker-driven dynamic loop report consumption on every step, or only when the walk ends? If only at the walk's end, a well-behaved live call would be refused with a 400 once it has sent max_pending_items items.

What the review checked

• Client stream buffers on the consumer worker: _my_consumer_connections includes CLIENT connections and buffers are built at NEW_REQUEST (worker.py:819-832), so the first client chunk is not acked unread by the new unbuffered path.
• The producer_done narrowing is safe for existing models: the conductor is the only sender (conductor.py:1928), and its producer_done_items names every edge it ends, with 0 for edges that carried nothing.
• Conductor streaming_connections gets a (CLIENT, partition, edge) entry from topology.connections, and no code indexes partition_states[conn.from_partition], so the non-partition producer name never reaches partition lookups.
• _register_draining and _remove_request both accept failure_status; the 400 from ClientInputRefusedError gets through.
• Held chunks for waiting requests: every removal from waiting_queue either goes through _drop_waiting or _admit_request, and both pop _held_stream_inputs.
• NEW_REQUEST and the opening's CLIENT_STREAM_INPUT leave the preprocess thread in order through one outbox; an append is queued only after submit_request has enqueued its opening.
• /generate/ws: an append's await live[request_id] resolves on every serve_one path (success, exception before the send, finally), and wait_for_request_id is bounded by DUPLICATE_ID_GRACE_S, so the receive loop cannot hang on it.
• Aborting on socket close: abort_submission is a no-op for keys already finished, so the later generator finally abort does nothing twice.
• The streaming handler now clears req.chunks: completion and drain accounting use output_chunk_counts and data-worker counts rather than req.chunks, and the non-streaming path still returns the full list.
• _finish_reading(drained=True) gating only applies to rids already in _draining_rids; normal completion sends the preprocess worker no DrainRequest, so the happy path still ACKs immediately.
• get_partition_topology on the API server's model instance only reads config fields for every model (nemotron_duplex, cosmos3, chatterbox, orpheus, qwen3_tts, qwen3_omni).
• ClientStreamInput needs no Rust mirror: the Rust tree has no body registry for conductor messages (invariant 5).
• New test modules under test/modular guard torch with importorskip and need no GPU.

Not reviewed: docs/adding_models.rst, docs/clients.rst, test/modular/test_client_stream_api.py (body), test/modular/test_client_stream_conductor.py (body), test/modular/test_data_worker_late_release.py, test/modular/test_generate_ws.py, the remaining small test fixture updates under test/modular/.

_{Advisory, not a merge gate. Rules: AGENTS.md. Reply to discuss.}

[dghosef]

Re the automated review:

• Opening sentence: it now names the producer_done change and the models that see it, and the new consumption report. "Changes every model sees" has a producer_done bullet.
• No result for this commit by itself: "Tests run" now opens with this head's own runs, on e741cfe90: ruff clean; CPU Core 6050 passed, 3 failed, 612 skipped (the same 3 environment-only failures as main abdfd5802); test_client_stream_conductor.py 52 passed, Rust cases included; test_rust_wire_frames.py + test_runtime_parity.py 193 passed, 1 skipped by design.
• Tests section length: the previous revision's H200 runs moved into a collapsed "Earlier revisions", and the near-tie details are left to nemotron_duplex: live input over a client stream, cache-aware streaming encoder (full duplex 4/6) #594.
• Question, does consumed_count move only when the walk ends: Good catch — it was only at walk end, so a loop consumer would be refused. Fixed in 0ed4266: after a step whose consumed count moved, the worker sends one columnar CLIENT_STREAM_CONSUMED (lockstep leader only; only for client streams with a max_pending_items bound; skipping rids whose WORKER_GRAPHS_DONE carried the count that step), and the conductor folds it in with max through the same path as WORKER_GRAPHS_DONE's counts, sending a held keep_newest item once there is room. New tests drive a real graph runtime (Python and Rust) through a looped consumer and the real Worker._stream_and_send into the real Conductor: a call that sends three times its limit while being consumed is not refused, and one whose consumer stops reading still is. With conductor.py, worker.py and stream_buffer.py reverted, both loop-consumer cases fail.

Posted by mistake

[Gaurav-Shah05]

Overall good, and should be ready after the bug below is fixed.

Two things for follow-ups:

1. One append is one stream item of any length, so appends under 80 ms fall behind, and the scheduler can't tell a 20 s upload from an 80 ms append. The precedent from other engines is cutting the client audio into model frames as it arrives. Stream items in the model's own units would fix both.

2. Other engines also serve duplex on the OpenAI Realtime API, where we only have /generate/ws. Main's nemotron_duplex/realtime_api.py still runs DuplexStream beside the engine, and a thin /v1/realtime adapter over client streams would let that file go.

[Gaurav-Shah05]

at max_concurrent_requests, a live call holds up to 256 appends (20 s at 80 ms) until it gets a slot, so it starts and stays that far behind. The runs in #770 leave no headroom at 32 calls, so I think a call with open input should be refused at that cap.

[dghosef]

Done

[Gaurav-Shah05]

No model or config in the stack sets max_pending_items, so refuse, keep_newest, the held waiting item and the per-step CLIENT_STREAM_CONSUMED reports only run in tests. Can you leave them out until the first model needs a bound on pending input?

[dghosef]

Done

[Gaurav-Shah05]

Need a rebase; approving