[API server][Conductor][Critical fix] Hold a finished request until every emitted chunk has arrived

mstar/api_server/data_worker.py

@@ -137,7 +137,6 @@ def _kv_page_sizes(model: Model, model_config: dict) -> dict[str, int]:
     }
 
 
-NameToLoopIndices = dict[str, NestedLoopIndices]
 
 
 class PreprocessWorker:
@@ -166,7 +165,8 @@ def __init__(
         self.wakeup = EventWakeup()
 
         self.per_request_reading_tensors = {}
-        self.output_loop_idxs: dict[str, NameToLoopIndices] = {}
+        # request id -> output name -> chunks received, for received_final_chunks
+        self.received_chunk_counts: dict[str, dict[str, int]] = {}
 
         # Build the communicator + tensor manager here (main thread) and hand
         # them to the worker thread, rather than constructing them inside it.
@@ -219,7 +219,7 @@ def _signal(self):
             wakeup.signal()
 
     def new_request(self, input: PreprocessInput):
-        self.output_loop_idxs[input.request_id] = {}
+        self.received_chunk_counts[input.request_id] = {}
         self.per_request_reading_tensors[input.request_id] = 0
         self.request_input_queue.put(input)
         self._signal()
@@ -232,7 +232,7 @@ def abort_request(self, request_id: str):
         # out from under it. Drop only the main-thread bookkeeping here.
         self.abort_request_queue.put(request_id)
         self._signal()
-        self.output_loop_idxs.pop(request_id, None)
+        self.received_chunk_counts.pop(request_id, None)
         self.per_request_reading_tensors.pop(request_id, None)
 
     def finished_reading(self, request_id: str, drained: bool = True):
@@ -244,21 +244,20 @@ def finished_reading(self, request_id: str, drained: bool = True):
         for any in-flight read."""
         self.reads_done_queue.put((request_id, drained))
         self._signal()
-        self.output_loop_idxs.pop(request_id, None)
+        self.received_chunk_counts.pop(request_id, None)
         self.per_request_reading_tensors.pop(request_id, None)
 
     def new_result_tensors(self, input: ResultTensors):
         name = input.graph_edge.name
-        if input.request_id not in self.output_loop_idxs:
+        if input.request_id not in self.received_chunk_counts:
             # Request was removed while this output was still in flight; ack the
             # tensors so the producing worker can reclaim them rather than leak.
             logger.debug("Late result_tensors for cleaned-up request %s, acking and dropping", input.request_id)
             self.discard_result_tensors(input)
             return
 
-        self.output_loop_idxs[input.request_id][name] = input.loop_indices.max(
-            self.output_loop_idxs[input.request_id].get(name, None)
-        )
+        counts = self.received_chunk_counts[input.request_id]
+        counts[name] = counts.get(name, 0) + 1
 
         self.per_request_reading_tensors[input.request_id] += len(input.graph_edge.tensor_info)
         logger.debug(
@@ -277,13 +276,11 @@ def new_result_token(
         transport read and nothing to ack. ``wake=False`` queues without
         ending the thread's idle wait: a frame's caller hands over every
         request's value first and calls ``wake`` once."""
-        if request_id not in self.output_loop_idxs:
+        if request_id not in self.received_chunk_counts:
             logger.debug("Late token for cleaned-up request %s dropped", request_id)
             return
-        if loop_indices is not None:
-            self.output_loop_idxs[request_id][signal] = loop_indices.max(
-                self.output_loop_idxs[request_id].get(signal, None)
-            )
+        counts = self.received_chunk_counts[request_id]
+        counts[signal] = counts.get(signal, 0) + 1
         self.per_request_reading_tensors[request_id] += 1
         self.result_tensor_input_queue.put(
             InlineToken(request_id, value, loop_indices, modality)
@@ -309,18 +306,18 @@ def has_pending_tensors(self, request_id: str):
 
     def received_final_chunks(
         self, request_id: str,
-        final_outputs: dict[str, NestedLoopIndices],
+        output_chunk_counts: dict[str, int],
     ):
         # Every serving walk reports at least one client-facing output, so an
         # empty dict means the walk emitted nothing (or completion raced ahead
         # of every result). Report not-done and let the TTL backstop close the
         # request rather than completing it instantly with zero chunks.
-        if not final_outputs:
+        if not output_chunk_counts:
             return False
+        received = self.received_chunk_counts.get(request_id, {})
         return all(
-            not loop_iters.label_context_gt( # recv'd loop iters is not less than the final_fwd
-                self.output_loop_idxs[request_id].get(name, None)
-            ) for name, loop_iters in final_outputs.items()
+            received.get(name, 0) >= count
+            for name, count in output_chunk_counts.items()
         )
 
     def get_result_chunks(self)-> list[ResultChunk]:
@@ -373,7 +370,7 @@ def get_rx_info(self, request_id: str) -> list[RxInfo]:
     def cleanup_request(self, request_id: str):
         self.cleanup_request_queue.put(request_id)
         self._signal()
-        self.output_loop_idxs.pop(request_id, None)
+        self.received_chunk_counts.pop(request_id, None)
         self.per_request_reading_tensors.pop(request_id, None)
 
     def shutdown(self):

mstar/api_server/entrypoint.py

@@ -224,7 +224,7 @@ class PendingRequest:
     profile: RequestProfile
     event: threading.Event = field(default_factory=threading.Event)
     chunks: list[ResultChunk] = field(default_factory=list)
-    final_outputs: dict = field(default_factory=dict)
+    output_chunk_counts: dict[str, int] = field(default_factory=dict)
     consumed_chunks: int = 0
     error: Any | None = None
     error_status: int = 500
@@ -629,7 +629,7 @@ def _prune_recently_completed(self) -> None:
                 req is not None
                 and not self.preprocess_worker.has_pending_tensors(rid)
                 and self.preprocess_worker.received_final_chunks(
-                    rid, req.final_outputs
+                    rid, req.output_chunk_counts
                 )
             )
             if drained or req is None:
@@ -762,14 +762,14 @@ def _process_messages(self) -> None:
                                 logger.info("API server received %s done", rid)
                                 self.recently_completed[rid] = time.time()
 
-                                if not message.body.final_outputs:
+                                if not message.body.output_chunk_counts:
                                     logger.warning(
                                         "Request %s completed with no reported "
                                         "outputs; holding for late results "
                                         "until the cleanup TTL", rid,
                                     )
                                 req = self.pending_requests[rid]
-                                req.final_outputs = message.body.final_outputs
+                                req.output_chunk_counts = message.body.output_chunk_counts
                                 req.profile.timing.conductor_ingest_time = \
                                     message.body.conductor_ingest_time
                                 req.profile.timing.conductor_finish_time = \

mstar/api_server/request_types.py

@@ -43,10 +43,10 @@ class ResultTokens:
 class RequestComplete:
     """Signals that a request has finished processing."""
     request_id: str
-    # Maps output signal name to its final forward pass number.
-    # The API server waits until all entries are received before
-    # completing the request.
-    final_outputs: dict[str, NestedLoopIndices]
+    # Output signal name -> chunks the workers emitted to the client. The API
+    # server waits until it has received that many of each before completing
+    # the request: completion can overtake the last chunks on the wire.
+    output_chunk_counts: dict[str, int]
     conductor_ingest_time: float
     conductor_finish_time: float
     graph_timings: GraphTimings = field(default_factory=dict)

mstar/conductor/conductor.py

@@ -28,7 +28,6 @@
 from mstar.distributed.communication import GlobalParallelConfig, WorkerParallelGroups
 from mstar.engine.resources import KVReqConfig, ResourceReqConfig
 from mstar.graph.base import GraphEdge, NodeAndGraphWalk, TensorPointerInfo
-from mstar.graph.loop_indices import NestedLoopIndices
 from mstar.model.base import ForwardPassArgs, Model, WorkerGraph
 from mstar.platforms import apply_process_env_defaults
 from mstar.profile.format import RxInfo, TxInfo
@@ -263,8 +262,9 @@ class RequestData:
     # Per-streaming-connection state, keyed by (from_partition, to_partition, edge_name)
     streaming_connections: dict[tuple[str, str, str], StreamingConnectionState] = field(default_factory=dict)
 
-    # for api server recv bookeeping
-    final_outputs: dict[str, NestedLoopIndices] = field(default_factory=dict)
+    # output name -> chunks the workers emitted to the client; the API server
+    # holds the request open until it has received this many of each
+    output_chunk_counts: dict[str, int] = field(default_factory=dict)
 
     # Conductor-side profiling timestamps, in ``time.perf_counter()`` seconds.
     # Comparable with the api-server stamps because perf_counter is
@@ -1438,7 +1438,7 @@ def _process_request_done(
                 message_type="request_complete",
                 body=RequestComplete(
                     request_id=request_id,
-                    final_outputs=request_data.final_outputs,
+                    output_chunk_counts=request_data.output_chunk_counts,
                     conductor_ingest_time=request_data.conductor_ingest_time,
                     conductor_finish_time=request_data.conductor_finish_time,
                     graph_timings=request_data.graph_timings,
@@ -1499,8 +1499,13 @@ def _process_worker_graphs_done(
             for tx in body.tx_info:
                 request_data.tx_info[(tx.source_entity, tx.edge_name)] = tx
 
+        # From every rank: a rank reports only the chunks it sent itself, and
+        # each of those reaches the API server as its own message.
+        counts = request_data.output_chunk_counts
+        for name in body.output_signal_names:
+            counts[name] = counts.get(name, 0) + 1
+
         pstate = request_data.partition_states.get(partition_name)
-        request_data.final_outputs.update(body.output_loop_indices)
         if pstate is None:
             logger.warning(
                 "WorkerGraphsDone for unknown partition %s (request %s)",
@@ -1538,10 +1543,6 @@ def _process_worker_graphs_done(
                     consumed = body.stream_tokens_consumed.get(conn.edge_name, 0)
                     conn.consumed_count = max(conn.consumed_count, consumed)
 
-            request_data.final_outputs.update(body.output_loop_indices)
-
-            pstate.curr_forward_outputs += body.output_signal_names
-
         # Each wg is only marked complete when all its TP ranks have reported.
         for wg_id in body.worker_graph_ids:
             count = pstate.wg_rank_completions.get(wg_id, 0) + 1

mstar/conductor/request_info.py

@@ -122,7 +122,6 @@ class PartitionState:
     # wg_id -> count of distinct TP ranks that have reported completion
     wg_rank_completions: dict[int, int] = field(default_factory=dict)
     num_output_tokens: int = 0
-    curr_forward_outputs: list[str] = field(default_factory=list)
     # resource label -> PublishedInfo, accumulated across all TP ranks and
     # handed back out on the next forward
     resource_publish_info: dict[str, PublishedInfo] = field(default_factory=dict)

mstar/graph/loop_indices.py

@@ -38,10 +38,3 @@ def label_context_gt(self, other: "NestedLoopIndices | None", target_loop_name:
             if our_idx < their_idx:
                 return False
         return False
-
-    def max(self, other: "NestedLoopIndices | None") -> "NestedLoopIndices":
-        if other is None:
-            return self
-        if other.label_context_gt(self):
-            return other
-        return self

mstar/graph/runtime/python.py

@@ -108,10 +108,8 @@ class GraphRuntimeRequestInfo:
     # so a persist signal cannot race the message that announces it.
     pending_persist_signals: list[GraphEdge] = field(default_factory=list)
     pending_new_token_counts: dict[str, dict[str, int]] = field(default_factory=dict)  # by partition
+    # one name per chunk emitted to the client; the conductor counts them
     current_output_chunks: list[str] = field(default_factory=list)
-    output_loop_indices: dict[str, NestedLoopIndices] = field(
-        default_factory=dict
-    )
 
 
 class PythonGraphRuntime(GraphRuntime):
@@ -1509,7 +1507,6 @@ def send_outputs(
                     edge.name for edge in routing.emit_to_client
                 )
                 for edge in routing.emit_to_client:
-                    info.output_loop_indices[edge.name] = nested_idxs.get(rid)
                     self._communicator.send("api_server", APIServerMessage(
                         message_type="result_tensors",
                         body=ResultTensors(
@@ -1618,7 +1615,6 @@ def _send_worker_graphs_done(
                 partition_name=partition,
                 partition_done=partition_done,
                 stream_tokens_consumed=stream_tokens_consumed,
-                output_loop_indices=info.output_loop_indices,
                 graph_timings=graph_timings,
                 rx_info=rx_info,
                 tx_info=tx_info,

mstar/utils/ipc_format.py

@@ -235,12 +235,12 @@ class WorkerGraphsDone(MessageBody):
     is_first_tp_rank: bool
     persist_signals: dict[str, list[TensorPointerInfo]] = field(default_factory=dict)
     new_token_counts: dict[str, int] = field(default_factory=dict) # name to token counts
+    # one name per chunk this worker emitted to the client since its last report
     output_signal_names: list[str] = field(default_factory=list)
     resource_publish_info: dict[str, PublishedInfo] = field(default_factory=dict)
     partition_name: str = field(default="default")
     partition_done: bool = field(default=False)
     stream_tokens_consumed: dict[str, int] = field(default_factory=dict)  # edge_name -> tokens consumed from stream
-    output_loop_indices: dict[str, NestedLoopIndices] = field(default_factory=dict)
     graph_timings: GraphTimings = field(default_factory=dict)
     rx_info: list[RxInfo] = field(default_factory=list)
     tx_info: list[TxInfo] = field(default_factory=list)

mstar/worker/worker.py

@@ -3245,7 +3245,7 @@ def _pp_stage(name: str) -> None:
             range_push("worker.send_outputs", synchronize=False)
 
         # The consuming pass (not the earlier ingest) reports the partition
-        # done, so it rides this pass's WGD with the final output loop index.
+        # done, so it rides the WGD that also counts this pass's last chunks.
         for rid in batch_N.node_batch.stream_partition_done_rids:
             self._graph_runtime.mark_stream_partition_done(rid, batch_N.partition)
 

rust/src/graph/frames.rs

@@ -239,7 +239,6 @@ pub struct WorkerGraphsDone<'a> {
     pub partition_name: &'a str,
     pub partition_done: bool,
     pub stream_tokens_consumed: &'a [(String, i64)],
-    pub output_loop_indices: Vec<(Sym, (Vec<Sym>, Vec<(Sym, u32)>, u32))>,
     /// `dict[str, PublishedInfo]` -- abstract, so it stays Python's. Read
     /// encoded by Python and forwarded without interpreting resource types.
     pub resource_publish_info: Option<Value>,
@@ -275,14 +274,6 @@ impl WorkerGraphsDone<'_> {
             str_counts(
                 m.key("stream_tokens_consumed"), self.stream_tokens_consumed,
             );
-            w_map(
-                m.key("output_loop_indices"),
-                self.output_loop_indices.len() as u32,
-            );
-            for (sig, (order, idxs, fwd)) in &self.output_loop_indices {
-                w_str(m.buf(), it.name(*sig));
-                nested_loop_indices(m.buf(), it, order, idxs, *fwd);
-            }
             // Omitted when absent, so the decoder's default stands.
             if let Some(v) = &self.resource_publish_info {
                 w_value(m.key("resource_publish_info"), v);

rust/src/graph/request.rs

@@ -106,15 +106,14 @@ pub struct PendingOutputs {
     pub persist: Vec<(Sym, Vec<u64>)>,
     /// partition -> signal -> count; each report carries its own partition's.
     pub new_tokens: Vec<(PartitionId, Vec<(String, i64)>)>,
-    /// Emitted signal names, in emission order.
+    /// One name per chunk emitted to the client, in emission order. The
+    /// conductor counts them to tell the API server how many chunks to wait
+    /// for.
     pub output_signals: Vec<Sym>,
-    /// signal -> the loop context it was emitted at.
-    pub output_loop_indices: Vec<(Sym, (Vec<Sym>, Vec<(Sym, u32)>, u32))>,
 }
 
 impl PendingOutputs {
-    /// Everything but `output_loop_indices`, which the conductor keeps for the
-    /// life of the request -- Python re-sends it on every WGD.
+    /// Everything pending, with only `partition`'s token counts.
     pub fn take_for_send(
         &mut self, partition: Option<PartitionId>,
     ) -> (Vec<(Sym, Vec<u64>)>, Vec<(String, i64)>, Vec<Sym>) {
@@ -147,8 +146,7 @@ impl PendingOutputs {
         }
     }
 
-    /// Accumulate, matching Python: a repeated signal's count adds up, and a
-    /// repeated loop index overwrites.
+    /// Accumulate, matching Python: a repeated signal's count adds up.
     pub fn add_new_tokens(&mut self, partition: PartitionId, counts: &[(String, i64)]) {
         let i = match self.new_tokens.iter().position(|(p, _)| *p == partition) {
             Some(i) => i,
@@ -165,15 +163,6 @@ impl PendingOutputs {
             }
         }
     }
-
-    pub fn set_loop_indices(
-        &mut self, signal: Sym, idx: (Vec<Sym>, Vec<(Sym, u32)>, u32),
-    ) {
-        match self.output_loop_indices.iter_mut().find(|(s, _)| *s == signal) {
-            Some(slot) => slot.1 = idx,
-            None => self.output_loop_indices.push((signal, idx)),
-        }
-    }
 }
 
 impl RequestInfo {