NemotronLabs VoiceChat-11B PR 1: streaming consumers inside loops, stream-terminated partition end

mstar/conductor/conductor.py

@@ -1383,6 +1383,19 @@ def _process_done_forward(
         pstate.metadata = fwd_args.full_metadata
         pstate.metadata.kwargs.update(fwd_args.step_metadata)
 
+        # The worker's signal is authoritative for a stream-terminated
+        # partition: the pass that consumed the stream's final chunk has run.
+        # Treat it as the partition's own `request_done` so its downstream
+        # connections get `producer_done` and the chain closes. A model that
+        # inferred the end from the connection counters instead
+        # (`consumed_count >= token_count`) could race: `consumed_count` counts
+        # chunks popped for execution, and every worker-graphs-done report from
+        # a worker carries every buffer's count, so a colocated consumer's report
+        # could finish this partition while its last step was still running,
+        # and that step's output arrived for a request already gone.
+        if incoming_connections and partition_done_from_worker:
+            fwd_args.request_done = True
+
         # Check max output tokens for partitions that produce tokens
         if pstate.num_output_tokens >= request_data.max_output_tokens:
             logger.info(

mstar/graph/runtime/base.py

@@ -55,6 +55,9 @@ class InputTensors(NamedTuple):
     # Per rid, the edges that carried a stream's final chunk. The CONSUMING
     # pass reports the partition done, which is why this rides with the inputs.
     final_stream_edges: dict[int, set[str]]
+    # Rids with an edge whose tensor the store no longer has, under
+    # ``skip_missing``. Empty otherwise, because the lookup raises instead.
+    unresolved_rids: frozenset[int] = frozenset()
 
 
 @dataclass
@@ -172,6 +175,7 @@ def to_edges(
 
     def to_input_tensors(
         self, get_tensor, rids: list[int] | None = None,
+        skip_missing: bool = False,
     ) -> InputTensors:
         """One walk for both things a batch build needs: each rid's inputs as
         ``{signal: [tensor]}``, and per rid the edges that carried a stream's
@@ -180,11 +184,20 @@ def to_input_tensors(
         ``rids`` seeds the result so a rid with no ready edges still gets an
         empty mapping -- what slicing a zero-length run used to give. Left
         None, only the rids that actually have edges appear.
+
+        ``skip_missing`` makes a uuid the store has forgotten take only its own
+        rid out, reported in ``unresolved_rids``, instead of raising out of the
+        whole walk. That a live rid reaches here holding a freed uuid is always
+        a bug, but this resolves the batch, so without the per-rid skip one such
+        rid fails every other request in the step with it. Off by default: the
+        speculative preps roll back and retry serially, so for them raising is
+        the better answer.
         """
         by_rid: dict[int, NameToTensorList] = (
             {} if rids is None else {rid: {} for rid in rids}
         )
         final_edges: dict[int, set[str]] = {}
+        unresolved: set[int] = set()
         names = self.signal_names
         uuids = self.uuids
         at = 0
@@ -194,11 +207,22 @@ def to_input_tensors(
             if slot is None:
                 slot = by_rid[rid] = {}
             name = names[self.signal_name_idxs[i]]
-            slot[name] = [get_tensor(uuids[j]) for j in range(at, end)]
+            try:
+                tensors = [get_tensor(uuids[j]) for j in range(at, end)]
+            except KeyError:
+                if not skip_missing:
+                    raise
+                # Leave the signal absent rather than half-filled: the caller
+                # drops the rid, and a partial input dict would otherwise
+                # surface as a missing-signal error further down.
+                unresolved.add(rid)
+                at = end
+                continue
+            slot[name] = tensors
             if self.is_final_streaming_chunk[i]:
                 final_edges.setdefault(rid, set()).add(name)
             at = end
-        return InputTensors(by_rid, final_edges)
+        return InputTensors(by_rid, final_edges, frozenset(unresolved))
 
     def edge_tuples(self) -> list[EdgeTuple]:
         """One ``EdgeTuple`` per edge.
@@ -362,6 +386,10 @@ class SpeculationPrepOutput(NamedTuple):
     # edges simply has none here, so callers that want an entry for it pass
     # ready_rids to ``to_input_tensors``.
     input_edges: ColumnarEdgeSpecs
+    # Handle for ``commit_speculation``, which every prep that returns one must
+    # eventually be settled with. 0 when the prep staged no streaming ingest,
+    # and so has nothing to settle.
+    spec_id: int = 0
 
 
 class ReadyNodeSpec(NamedTuple):
@@ -762,6 +790,26 @@ def prep_follow_spec_rids(
         """
         pass
 
+    @abstractmethod
+    def commit_speculation(
+        self, spec_id: int, success: bool, dropped_rids: list[int] = (),
+        node: str | None = None, wg_id: int | None = None,
+        scheduled_rids: list[int] = (),
+    ):
+        """Settle the streaming ingests a prep staged under ``spec_id``.
+        On ``success``, ``scheduled_rids`` are also marked speculatively
+        scheduled on ``node``.
+
+        The scheduled rids are passed rather than read off the stage in order
+        to also include the fresh rids rolled in from the ready queue. For
+        dropped rids, stream chunks are returned to the stream buffers and
+        un-ingested from the graph.
+
+        Idempotent, and a no-op for an unknown or 0 id, so a caller settling a
+        prep that staged nothing (or settling twice on an error path) is safe.
+        """
+        pass
+
     # --------- Postprocess ----------
     @abstractmethod
     def stop_loops_batched(

mstar/graph/runtime/python.py

@@ -158,6 +158,11 @@
         self._completions: dict[int, CompletionState] = {}
         self._completion_counter = 0
 
+        # Streaming ingests a prep staged, parked until commit_speculation
+        # settles them: spec_id -> [(rid, node, into_signals, into_next_iter)].
+        self._staged_specs: dict[int, list[tuple]] = {}
+        self._spec_counter = 0
+
         # worker graph info
         self._queues = {
             worker_graph.worker_graph_id: WorkerGraphQueues(
@@ -272,6 +277,16 @@
             del self._completions[cid]
         for c in self._completions.values():
             c.routing.pop(rid, None)
+        # Same recycling hazard for a staged speculation nobody settled: its
+        # undo would otherwise run against whichever request gets this handle
+        # next. The node objects go away with the queues below anyway.
+        for sid in [
+            sid for sid, staged in self._staged_specs.items()
+            if all(s[0] == rid for s in staged)
+        ]:
+            del self._staged_specs[sid]
+        for sid, staged in self._staged_specs.items():
+            self._staged_specs[sid] = [s for s in staged if s[0] != rid]
 
         info = self._request_info.pop(rid)
         for wg_id in info.worker_graph_ids:
@@ -759,6 +774,7 @@
             ready_rids=[p.rid for p in prepped],
             wg_ids=[wg_id] * len(prepped),
             input_edges=edges,
+            spec_id=self._stage_spec(prepped),
         )
 
     def prep_spec_rids(
@@ -782,6 +798,7 @@
         consumed_idxs: list[int] = []
         ready_rids: list[int] = []
         wg_ids: list[int] = []
+        staged: list["_SpecRidPrep"] = []
         edges = ColumnarEdgeSpecs.empty()
 
         for rid, indexed_edges in per_rid_edges:
@@ -809,6 +826,7 @@
             consumed_idxs.extend(prepped.consumed_idxs)
             ready_rids.append(rid)
             wg_ids.append(wg_id)
+            staged.append(prepped)
             for signal, uuids, final_chunk in prepped.input_edges:
                 edges.add(rid, signal, uuids, final_chunk)
 
@@ -817,8 +835,43 @@
             ready_rids=ready_rids,
             wg_ids=wg_ids,
             input_edges=edges,
+            spec_id=self._stage_spec(staged),
         )
 
+    def _stage_spec(self, prepped: list["_SpecRidPrep"]) -> int:
+        """Park what a prep ingested, for ``commit_speculation`` to settle.
+
+        0 when nothing was ingested: there is no stage to undo, so the caller
+        has nothing to settle and does not need an id to forget.
+        """
+        staged = [
+            (p.rid, p.node, p.into_signals, p.into_next_iter) for p in prepped
+            if p.into_signals or p.into_next_iter
+        ]
+        if not staged:
+            return 0
+        self._spec_counter += 1
+        self._staged_specs[self._spec_counter] = staged
+        return self._spec_counter
+
+    def commit_speculation(
+        self, spec_id: int, success: bool, dropped_rids: list[int] = (),
+        node: str | None = None, wg_id: int | None = None,
+        scheduled_rids: list[int] = (),
+    ):
+        if success and node is not None and wg_id is not None:
+            self.set_speculatively_scheduled(
+                node, wg_id, list(scheduled_rids), True,
+            )
+        staged = self._staged_specs.pop(spec_id, None)
+        if staged is None:
+            return
+        dropped = set(dropped_rids)
+        for rid, node, into_signals, into_next_iter in staged:
+            if success and rid not in dropped:
+                continue
+            self._undo_spec_ingest(node, into_signals, into_next_iter)
+
     def _can_continue_loop(
         self, rid: int, wgio, spec_info, graph_walk: str,
     ) -> bool:

mstar/graph/runtime/rust.py

@@ -514,6 +514,17 @@ def _spec_prep_out(out) -> SpeculationPrepOutput:
             ready_rids=out.ready_rids,
             wg_ids=out.wg_ids,
             input_edges=_columns(out),
+            spec_id=out.spec_id,
+        )
+
+    def commit_speculation(
+        self, spec_id: int, success: bool, dropped_rids: list[int] = (),
+        node: str | None = None, wg_id: int | None = None,
+        scheduled_rids: list[int] = (),
+    ):
+        self._rust.commit_speculation(
+            spec_id, success, list(dropped_rids),
+            node, wg_id, list(scheduled_rids),
         )
 
     def prep_spec_rids(