Bug fix: dynamic loop iter counts

mstar/engine/engine.py

@@ -177,6 +177,26 @@ def set_piecewise_slot(self, slot: int) -> None:
         for runner in self.piecewise_runners.values():
             runner.set_slot(slot)
 
+@dataclass
+class ForwardPassInfoWrapper:
+    """
+    Stop-gap solution: the dynamic loop iters in CurrentForwardPassInfo may
+    be used in two places at once: in _postprocess_batch (for batch N), and
+    in the forward pass for batch N+1, which require different iterations.
+    The InputMetadata is built per executing batch, so its value is guaranteed
+    to be up-to-date.
+
+    TODO: the long-term solution is to maybe move the dynamic loop iters off
+    of CurrentForwardPassInfo entirely.
+    """
+    info: CurrentForwardPassInfo
+    metadata: InputMetadata
+
+    def __getattr__(self, name):
+        if name == "dynamic_loop_iter_counts":
+            return self.metadata.dynamic_loop_iter_counts
+        return getattr(self.info, name)
+
 
 @dataclass
 class ExecutingBatch:
@@ -264,6 +284,17 @@ class ExecutingBatch:
     # Per-step wall-clock, for the worker's profiler
     exec_timings: ExecTimings = field(default_factory=ExecTimings)
 
+    def __post_init__(self):
+        self.per_request_info_wrapped = {
+            rid: ForwardPassInfoWrapper(
+                # Sparse: only rids with stream chunks or loop counters get an
+                # entry, so default like the .get() at prepare_inputs does.
+                info=info, metadata=self.per_request_input_metadata.get(
+                    rid, EMPTY_INPUT_METADATA,
+                ),
+            ) for rid, info in self.per_request_info.items()
+        }
+
     @property
     def request_ids(self):
         return self.step_context.request_ids
@@ -306,6 +337,10 @@ def drop_rids(self, rids: set[str]):
             rid: info for rid, info in self.per_request_info.items()
             if rid not in rids
         }
+        self.per_request_info_wrapped = {
+            rid: info for rid, info in self.per_request_info_wrapped.items()
+            if rid not in rids
+        }
 
     def register_admit_error(
         self, reason: AdmitFailedReason, failed_resource: str | None = None,
@@ -666,7 +701,7 @@ def _prepare_inputs(self, batch: ExecutingBatch) -> None:
             try:
                 req_inputs = submodule.prepare_inputs(
                     graph_walk=batch.step_context.graph_walk,
-                    fwd_info=batch.per_request_info[rid],
+                    fwd_info=batch.per_request_info_wrapped[rid],
                     inputs=batch.per_request_input_tensors.get(rid, {}),
                     resources=self._submodules[batch.node_name].resources,
                     # the step that ends the last of the node's streams: it must
@@ -731,7 +766,7 @@ def _skip_cached_prefix(
         if matched <= 0:
             return inputs
         return self._submodules[batch.node_name].submodule.split_inputs(
-            walk, batch.per_request_info[rid], inputs,
+            walk, batch.per_request_info_wrapped[rid], inputs,
             matched, inputs.input_seq_len,
         )
 
@@ -806,11 +841,11 @@ def _exec_single(
         real_bs = len(batch.request_ids)
 
         inputs = batch.inputs
-        req_info = batch.per_request_info
+        req_info = batch.per_request_info_wrapped
         if lease is not None:
             inputs = cg_runner.pad_inputs(lease, inputs)
             req_info = cg_runner.step_metadata(
-                lease, batch.request_ids, batch.per_request_info
+                lease, batch.request_ids, req_info
             )
             batch.step_context.set_padded_rids(
                 cg_runner.step_ids(lease, batch.request_ids)
@@ -919,7 +954,7 @@ def _exec_per_request(
         slot_events: dict[int, torch.cuda.Event] = {}
 
         for rid, inp in zip(batch.request_ids, batch.inputs, strict=True):
-            req_info = {rid: batch.per_request_info[rid]}
+            req_info = {rid: batch.per_request_info_wrapped[rid]}
             slot = ctxs[rid].slot
             submodule_mgmt.set_piecewise_slot(slot)
             in_flight = slot_events.get(slot)
@@ -1208,7 +1243,7 @@ def _postprocess_batch(
             try:
                 submodule.postprocess(
                     request_id=rid,
-                    request_info=batch.per_request_info[rid],
+                    request_info=batch.per_request_info_wrapped[rid],
                     outputs=outputs.get(rid, {}),
                     inputs=node_inputs,
                 )
@@ -1255,7 +1290,7 @@ def check_stop_for_batch(
                 continue
             try:
                 rid_stops = submodule.check_stop(
-                    rid, batch.per_request_info[rid], rid_outputs
+                    rid, batch.per_request_info_wrapped[rid], rid_outputs
                 )
             except Exception as error:
                 logger.exception(
@@ -1487,7 +1522,7 @@ def reserve_replay_slot(self, batch: ExecutingBatch) -> SlotLease | None:
         # step is declared, so the submodule is asked directly. It answers from
         # the same per-request facts its `declare_step` stamps on the step.
         batch.cg_key_info = submodule_mgmt.submodule.cg_key_info(
-            batch.step_context.graph_walk, batch.per_request_info,
+            batch.step_context.graph_walk, batch.per_request_info_wrapped,
             per_request_input_metadata=batch.per_request_input_metadata,
         )
         lease = cg_runner.lease_slot(

mstar/graph/base.py

@@ -242,7 +242,9 @@ class GraphNode(GraphSection):
     # Whether this node is currently speculatively being executed: if so, we
     # don't want to queue it in a "ready" queue because it's already being
     # executed, but we do want to update everything else as normal.
-    _speculatively_scheduled: bool = False
+    # Set while a batch carrying this node is in flight (speculative or not):
+    # gates ready-queue adds so the node is not double-queued underneath it.
+    _in_flight: bool = False
 
     _streaming_inputs: set[str] = field(default_factory=set)
 
@@ -796,25 +798,25 @@ def __init__(self, graph_section: GraphSection):
 
     def register_ingested_input(self, graph_edge: GraphEdge):
         node = self.nodes[graph_edge.next_node]
-        # If node._speculatively_scheduled, the node is already executing as
-        # a spec batch. We don't want to double-queue it (either for the
+        # If node._in_flight, the node is already executing as part of a
+        # batch. We don't want to double-queue it (either for the
         # current iter via ``ready_names`` or for the next iter via
         # ``ready_next_iter`` — both eventually feed the scheduler) so we
         # gate every queue add on the flag.
         if node.ready_signals.is_ready:
-            if not node._speculatively_scheduled:
+            if not node._in_flight:
                 self.ready_names.add(node.name)
             self.ready_for_streaming.discard(node.name)
         elif node.ready_signals.is_ready_for_streaming:
-            if not node._speculatively_scheduled:
+            if not node._in_flight:
                 self.ready_for_streaming.add(node.name)
 
         if node.ready_next_iter.is_ready:
-            if not node._speculatively_scheduled:
+            if not node._in_flight:
                 self.ready_next_iter.add(node.name)
             self.ready_streaming_next_iter.discard(node.name)
         elif node.ready_next_iter.is_ready_for_streaming:
-            if not node._speculatively_scheduled:
+            if not node._in_flight:
                 self.ready_streaming_next_iter.add(node.name)
 
     def mark_entity_complete(self, entity_name: str) -> NodeCompletionOutput:

mstar/graph/runtime/base.py

@@ -561,17 +561,17 @@ def set_walk(self, rid: int, partition: str, walk: str):
         pass
 
     @abstractmethod
-    def set_speculatively_scheduled(
+    def set_in_flight(
         self, node: str, wg_id: int, rids: list[int],
-        speculatively_scheduled: bool
+        in_flight: bool,
     ):
         pass
 
     @abstractmethod
-    def is_speculatively_scheduled(
+    def is_in_flight(
         self, node: str, wg_id: int, rid: int,
     ) -> bool:
-        """Whether this rid's node is marked speculatively scheduled.
+        """Whether this rid's node is marked in flight.
 
         The flag must SURVIVE node completion: its rids are still in flight
         for the speculative N+1 step, so the node must stay out of the ready

mstar/graph/runtime/python.py

@@ -325,21 +325,21 @@ def set_walk(self, rid: int, partition: str, walk: str):
             if walk in self._all_wg_ids_to_graph_walks[wg_id]
         ]
 
-    def set_speculatively_scheduled(
+    def set_in_flight(
         self, node: str, wg_id: int, rids: list[int],
-        speculatively_scheduled: bool
+        in_flight: bool,
     ):
         queues = self._queues[wg_id].per_request_queues
         for rid in rids:
             if rid not in queues:
                 continue
-            queues[rid].get_node(node)._speculatively_scheduled = speculatively_scheduled
+            queues[rid].get_node(node)._in_flight = in_flight
 
-    def is_speculatively_scheduled(
+    def is_in_flight(
         self, node: str, wg_id: int, rid: int,
     ) -> bool:
         wgio = self._queues[wg_id].per_request_queues.get(rid)
-        return wgio is not None and wgio.get_node(node)._speculatively_scheduled
+        return wgio is not None and wgio.get_node(node)._in_flight
 
     def get_dynamic_loop_iters(
         self, request_ids: list[int],
@@ -860,7 +860,7 @@ def commit_speculation(
         scheduled_rids: list[int] = (),
     ):
         if success and node is not None and wg_id is not None:
-            self.set_speculatively_scheduled(
+            self.set_in_flight(
                 node, wg_id, list(scheduled_rids), True,
             )
         staged = self._staged_specs.pop(spec_id, None)
@@ -897,7 +897,7 @@ def _undo_spec_ingest(
 
         The two slots are removed from separately, which is why the ingest
         tracks which one each chunk landed in. Registry state was never
-        touched -- _speculatively_scheduled was held True across the ingest --
+        touched -- _in_flight was held True across the ingest --
         so the caller only has to return the chunks to their StreamBuffers.
         """
         for _idx, name in into_signals:
@@ -916,8 +916,11 @@ def _prep_one_spec_rid(
         """
         node = wgio.nodes[spec_node_name]
         # Held True across the ingest so a streaming input cannot re-add the
-        # node to the ready queue underneath us.
-        node._speculatively_scheduled = True
+        # node to the ready queue underneath us. Saved rather than assumed
+        # False: the node may already be in flight from a committed batch, and
+        # forcing it down here would drop that batch's protection.
+        was_in_flight = node._in_flight
+        node._in_flight = True
 
         # ingest_input reports success without saying WHICH slot it used, so
         # the slot is inferred by peeking before the call. The rollback below
@@ -940,7 +943,7 @@ def _prep_one_spec_rid(
             check_next_iter=same_node, allow_streaming=False,
         )
         wgio.clear_speculative_inputs()
-        node._speculatively_scheduled = False  # reset in case the rid is dropped
+        node._in_flight = was_in_flight  # restore; the node may still be in flight
 
         if not fully_ready:
             self._undo_spec_ingest(node, into_signals, into_next_iter)
@@ -1482,7 +1485,7 @@ def send_outputs(
                 ].per_request_queues.get(rid)
                 speculative = node is not None and node.get_node(
                     completion.node_name
-                )._speculatively_scheduled
+                )._in_flight
                 self._send_worker_graphs_done(
                     rid, routing, info, fwd_info, partition,
                     resource_publish_info=publications.get(rid, {}),

mstar/graph/runtime/rust.py

@@ -326,18 +326,16 @@ def get_rid_handle(self, rid: str) -> int | None:
     def set_walk(self, rid: int, partition: str, walk: str):
         self._rust.set_walk(rid, partition, walk)
 
-    def set_speculatively_scheduled(
+    def set_in_flight(
         self, node: str, wg_id: int, rids: list[int],
-        speculatively_scheduled: bool,
+        in_flight: bool,
     ):
-        self._rust.set_speculatively_scheduled(
-            node, wg_id, rids, speculatively_scheduled
-        )
+        self._rust.set_in_flight(node, wg_id, rids, in_flight)
 
-    def is_speculatively_scheduled(
+    def is_in_flight(
         self, node: str, wg_id: int, rid: int,
     ) -> bool:
-        return self._rust.is_speculatively_scheduled(node, wg_id, rid)
+        return self._rust.is_in_flight(node, wg_id, rid)
 
     def mark_stream_partition_done(self, rid: int, partition: str):
         self._rust.mark_stream_partition_done(rid, partition)

mstar/model/submodule_base.py

@@ -29,6 +29,8 @@ class InputMetadata(NamedTuple):
     """
     # input_name -> where that streamed input's chunk sits in its stream
     stream_chunks: Mapping[str, StreamChunkInfo] = MappingProxyType({})
+    dynamic_loop_iter_counts: Mapping[str, int] = MappingProxyType({})
+
     # TODO: fold in prepare_inputs' is_final_stream_chunk kwarg once the Chatterbox work settles.
     # TODO: move step_metadata and dynamic_loop_iter_counts here off CurrentForwardPassInfo (move, don't copy).