From e3446659649545899b1cf0d338c6a9fece4e9116 Mon Sep 17 00:00:00 2001 From: Naomi Sagan Date: Sun, 1 Nov 2026 18:02:01 +0000 Subject: [PATCH 1/2] Stop policies: the graph runtime checks a step's declared stops and runs its postprocess #857 as reviewed, squashed for the rebase onto main: its commits predate #794 (END_PARTITION) and #858 (TokenStop), and the ones that conflict are intermediate designs later commits replace. Reconciled with both: - END_PARTITION from check_stop: handled in Worker._check_stops, on both paths that call the model's check (before the runtime, or on its STOPS break): a row inside a loop fails, any other ends its partition through the request's partition_done (renamed from stream_partition_done, as #794 renamed the rest), or straight on the runtime once the request is built. - An inline route skips a step with withheld outputs, as main's _route_and_send_fast did. - Qwen3.5 declares no policies of its own: TokenStop(runtime_check=True) derives them (ARNodeSubmodule.get_stop_policies), as StopKind.EOS, since main's loop cap is the token budget; its prompt walks keep check_stop, which can end the partition there. Other TokenStop models are unchanged. Co-Authored-By: Claude Opus 5.5 --- docs/adding_models.rst | 19 + mstar/engine/engine.py | 17 + mstar/graph/runtime/base.py | 330 ++++++++- mstar/graph/runtime/python.py | 70 +- mstar/graph/runtime/rust.py | 171 ++++- mstar/model/orpheus/submodules.py | 14 +- mstar/model/qwen3_5/submodules.py | 15 +- mstar/model/submodule_base.py | 98 ++- mstar/worker/worker.py | 648 +++++++++--------- rust/src/bindings.rs | 2 +- rust/src/graph/request.rs | 5 + rust/src/graph/runtime.rs | 381 +++++----- rust/src/graph/runtime/postprocess.rs | 572 ++++++++++++++++ test/modular/test_dropped_row_stream_stop.py | 59 +- test/modular/test_ingest_rollback.py | 1 + test/modular/test_inline_tokens.py | 27 +- test/modular/test_output_spans.py | 8 +- test/modular/test_postprocess_request.py | 56 ++ test/modular/test_rid_handles.py | 1 + .../modular/test_send_registration_failure.py | 12 +- test/modular/test_stop_policies.py | 68 ++ .../test_tp_read_failure_acceptance.py | 4 +- test/modular/test_withheld_outputs.py | 8 +- test/modular/test_worker_drain.py | 15 +- test/modular/test_worker_host_copies.py | 9 +- test/rust/test_postprocess.py | 231 +++++++ test/rust/test_route_and_send_inline.py | 168 ----- test/rust/test_runtime_check_stops.py | 114 +++ 28 files changed, 2325 insertions(+), 798 deletions(-) create mode 100644 rust/src/graph/runtime/postprocess.rs create mode 100644 test/modular/test_postprocess_request.py create mode 100644 test/modular/test_stop_policies.py create mode 100644 test/rust/test_postprocess.py delete mode 100644 test/rust/test_route_and_send_inline.py create mode 100644 test/rust/test_runtime_check_stops.py diff --git a/docs/adding_models.rst b/docs/adding_models.rst index 00f72102b2..3195ccb542 100644 --- a/docs/adding_models.rst +++ b/docs/adding_models.rst @@ -1226,6 +1226,25 @@ methods are: starting early on a thread of the submodule's own; it must not touch the device. The DiT scaffold draws a request's seeded noise here (``PREFETCH_SEEDS``). +``get_stop_policies(self, graph_walk) -> dict[str, StopPolicy] | None`` (optional) + Declare what ``check_stop`` amounts to on a walk, per loop it can stop, when it is + only the common checks; the graph runtime then runs them over the step's token column + without calling ``check_stop``. That takes the batched output carrying the policies' + ``token_signal`` as a stop buffer (``BatchedModelOutput(check_stop_buffers={...})``, + as Qwen3.5 and Orpheus return); a step without one, or one that mixes walks, + still calls ``check_stop``. ``StopPolicy(StopKind.AR, token_signal=..., + eos_tokens=..., sampler=...)`` stops its loop on an EOS token (unless the ``sampler`` + resource's request config sets ``ignore_eos``) or when the loop's iteration count plus + one reaches ``max_iters`` (``info.max_tokens`` unless given). ``StopKind.ITER`` checks + only the count, and ``StopKind.EOS`` only the token, for a loop whose cap is its budget. + The limits are computed from the request's forward-pass info whenever + the worker receives it, not per step. Return ``{}`` for a + walk on which ``check_stop`` never stops a loop, and ``None`` (the default) to keep + calling ``check_stop``. The policies must agree with ``check_stop`` for every row; + ``test/modular/test_stop_policies.py`` checks the models that declare them. A + ``TokenStop`` with ``runtime_check=True`` declares them for you: its loop's stop ids, + as ``StopKind.EOS``, on ``walk``; its ``prefill_walks`` keep ``check_stop``. + ``cleanup_request(self, request_id)`` (optional) Free per-request state held inside the submodule when a request finishes: buffers, per-request caches, counters. See Qwen3-Omni's ``Code2WavSubmodule`` for an example. diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 7696bcd0dc..669daf3ddf 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -62,6 +62,7 @@ NodeSubmodule, StepOrder, StopCheck, + StopPolicy, ) from mstar.platforms.compile import with_inductor_overrides from mstar.profile.worker import ExecTimings @@ -1138,6 +1139,15 @@ def after_output_copy(self, batch: ExecutingBatch) -> None: if batch.committed is not None: self._runner.after_output_copy(batch.committed) + def step_finished( + self, batch: ExecutingBatch, outputs: dict[str, NameToTensorList], + host_rows: HostRows | None = None, + ) -> None: + """A step's outputs are on the host (the stop check's copy) and its + stops not yet decided: so a request ending here has its state up to + date, e.g. its pages indexed for prefix reuse.""" + self.extend_prefix_chains(batch, outputs, host_rows=host_rows) + def extend_prefix_chains( self, batch: ExecutingBatch, outputs: dict[str, NameToTensorList], host_rows: HostRows | None = None, @@ -1902,6 +1912,13 @@ def exec_and_postprocess( laps.lap("tail") return outputs + def get_stop_policies(self, node_name: str, graph_walk: str) -> dict[str, StopPolicy] | None: + """The node's declared stop policies on ``graph_walk`` + (``NodeSubmodule.get_stop_policies``); None when ``check_stop`` + decides, or the node has no submodule here.""" + mgmt = self._submodules.get(node_name) + return None if mgmt is None else mgmt.submodule.get_stop_policies(graph_walk) + def check_stop_for_batch( self, batch: ExecutingBatch, outputs: dict[str, NameToTensorList], host_rows: HostRows | None = None, diff --git a/mstar/graph/runtime/base.py b/mstar/graph/runtime/base.py index 0d7a1d2b88..ad0acc157e 100644 --- a/mstar/graph/runtime/base.py +++ b/mstar/graph/runtime/base.py @@ -1,8 +1,9 @@ +import logging from abc import ABC, abstractmethod -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from dataclasses import dataclass, field from enum import IntEnum -from typing import NamedTuple +from typing import TYPE_CHECKING, NamedTuple from mstar.communication.tensors import NameToTensorList, TensorStore from mstar.conductor.request_info import CurrentForwardPassInfo @@ -13,6 +14,13 @@ from mstar.profile.format import GraphTiming, RxInfo, TxInfo from mstar.utils.containers import ParallelList +if TYPE_CHECKING: + import torch + + from mstar.model.submodule_base import StopPolicy + +logger = logging.getLogger(__name__) + class SpeculationOutput(NamedTuple): node_name: str @@ -424,28 +432,6 @@ class SpecRows(NamedTuple): input_edges: ColumnarEdgeSpecs | None -class InlineRouteStatus(IntEnum): - """What ``route_and_send_inline`` did.""" - # nothing (a rid with no host row, an unsupported dtype): the caller takes - # the general path - NOT_ROUTED = 0 - # routed and sent - SENT = 1 - # routed, but the send needs Python-side payloads: the caller finishes - # with ``send_outputs`` - ROUTED = 2 - - -class InlineRoute(NamedTuple): - """``route_and_send_inline``'s result.""" - status: InlineRouteStatus - completion_id: int - # ``send_outputs``' return, under SENT - sent_rids: list[int] - # the route for ``send_outputs`` to finish, under ROUTED - route_output: "RouteOutput | None" - - class ReadyNodeSpec(NamedTuple): node_name: str graph_walk: str @@ -558,7 +544,270 @@ class SendInput(NamedTuple): failed_rids: frozenset[int] = frozenset() -class GraphRuntime(ABC): +class PostprocessNeed(IntEnum): + """What a step's postprocess (``begin_postprocess``) waits on from the + worker, or that it is done.""" + DONE = 0 + # the step's stops, when the runtime's own check failed or a row had no + # limits (the model's check_stop decides): resume with ``stops`` and the + # ``dropped`` rows + STOPS = 1 + # the outputs stored as tensors, to route: resume with ``stored`` + STORE = 2 + # the send's Python payloads (forward-pass info, publication, new-token + # counts, stream consumption, profiling): resume with ``send`` + SEND = 3 + + +class InlineColumn(NamedTuple): + """A device-loopback step's client signal, one host scalar per row: its + emit edge rides inline, and nothing is stored.""" + signal: str + # 1-D host view, row i for ``PostprocessRequest.row_rids[i]`` + column: "torch.Tensor" + + +class PostprocessRequest(NamedTuple): + """One step's postprocess: check its stops, route its outputs, send them.""" + partition: str + node_name: str + graph_walk: str + output_signals: list[str] + # rid -> the worker graph it ran this node in + wg_ids: ParallelList[int, int] + # several walks in one step: wg_ids.keys[i] ran walks[rid_walk_idx[i]] + walks: list[str] | None = None + rid_walk_idx: list[int] | None = None + # The step's stops known before the runtime's check: stream-terminated + # loops, and with runtime_stops False the model's check_stop too, so they + # are all of them and the step routes without stopping for STOPS. + stops: Mapping[int, set[str]] | None = None + # whether the runtime checks the walk's declared stop policies itself + runtime_stops: bool = False + # the token column they read (None when none does) and the rid of each + # row of the step's host rows; the caller keeps them alive for the step + tokens: "torch.Tensor | None" = None + row_rids: tuple[int, ...] | list[int] = () + # a device-loopback step: routed without storing anything + inline: InlineColumn | None = None + # rids whose partition this pass ends, reported on its WORKER_GRAPHS_DONE: the + # pass consumed its stream's final chunk, or check_stop returned END_PARTITION + partition_done: tuple[int, ...] | list[int] = () + # the send needs Python whatever the route needs (profiling) + python_send: bool = False + + +class RouteTensors(NamedTuple): + """``resume_postprocess(stored=...)``: what ``RouteInput`` takes of the + stored outputs.""" + tensors: list[int] + num_tensors: list[int] + inline_signal: str | None = None + inline_values: tuple[int, ...] | list[int] = () + logprobs: tuple[float | None, ...] | list[float | None] = () + + +class PostprocessStep(NamedTuple): + """Where a step's postprocess stands.""" + need: PostprocessNeed + # the rows it routes, without the dropped ones (STORE and SEND) + wg_ids: ParallelList[int, int] | None = None + rid_walk_idx: list[int] | None = None + # the rows whose loops it stopped (from STORE on) + stopped_rids: list[int] | tuple[()] = () + # the route, for SEND's payloads (SEND) + route_output: RouteOutput | None = None + # the rids a WORKER_GRAPHS_DONE went out for (DONE) + completed_rids: list[int] | tuple[()] = () + + +class _PostprocessPhase(IntEnum): + STOPS = 0 + ROUTE = 1 + SEND = 2 + + +@dataclass(slots=True) +class _Postprocess: + """A step's postprocess between ``begin`` and ``resume``.""" + request: PostprocessRequest + phase: _PostprocessPhase = _PostprocessPhase.STOPS + stops: dict[int, set[str]] | None = None + wg_ids: ParallelList[int, int] | None = None + rid_walk_idx: list[int] | None = None + stopped: list[int] = field(default_factory=list) + # the request's partition_done, less the dropped rows + partition_done: list[int] = field(default_factory=list) + stored: RouteTensors | None = None + route: RouteOutput | None = None + send: SendInput | None = None + + def step(self, need: PostprocessNeed, completed: list[int] | None = None) -> PostprocessStep: + if need is PostprocessNeed.DONE: + return PostprocessStep(need, stopped_rids=self.stopped, completed_rids=completed or []) + return PostprocessStep(need, self.wg_ids, self.rid_walk_idx, self.stopped, self.route) + + +def _merged_stops(*parts: Mapping[int, set[str]] | None) -> dict[int, set[str]]: + out: dict[int, set[str]] = {} + for part in parts: + for rid, loops in (part or {}).items(): + out.setdefault(rid, set()).update(loops) + return out + + +def _route_needs_python(route: RouteOutput) -> bool: + """Whether the send of ``route`` needs what only the worker has.""" + return bool( + route.rids_needing_request_info is None or route.rids_needing_request_info + or route.register_uuids or route.new_token_output_idxs + or route.local_streaming_by_signal or route.freed_inputs.uuids + ) + + +class PhasedPostprocess: + """A step's postprocess as ``begin_postprocess`` / ``resume_postprocess`` + over a runtime's three phases: the stops (``check_stops``, then + ``stop_loops_batched``), the route (``complete_and_route_batch``) and the + send (``send_outputs``). It stops for the worker where a phase needs + Python (``PostprocessNeed``). The Python runtime's postprocess, and the + reference the Rust runtime's own follows.""" + + _postprocess: _Postprocess | None = None + + def begin_postprocess(self, request: PostprocessRequest) -> PostprocessStep: + self._postprocess = _Postprocess(request, partition_done=list(request.partition_done)) + return self._run_postprocess() + + def resume_postprocess( + self, stops: Mapping[int, set[str]] | None = None, dropped: Iterable[int] = (), + stored: RouteTensors | None = None, send: SendInput | None = None, + ) -> PostprocessStep: + """Continue the waiting step with what it asked for: ``stops`` and the + rows ``dropped`` (STOPS), ``stored`` (STORE) or ``send`` (SEND).""" + pp = self._postprocess + if pp is None: + raise RuntimeError("resume_postprocess: no step is waiting") + if stops is not None: + pp.stops = _merged_stops(pp.request.stops, stops) + self._keep_rows(pp, set(dropped)) + if stored is not None: + pp.stored = stored + if send is not None: + pp.send = send + return self._run_postprocess() + + def _run_postprocess(self) -> PostprocessStep: + pp = self._postprocess + req = pp.request + if pp.phase is _PostprocessPhase.STOPS: + if pp.stops is None and not req.runtime_stops: + # checked by the worker: the request's stops are all of them + pp.stops = _merged_stops(req.stops) + self._keep_rows(pp, set()) + if pp.stops is None: + found = None + if req.runtime_stops: + # a check that fails, or misses a row's limits, leaves + # the decision to check_stop + try: + found = self.check_stops( + req.partition, req.node_name, req.graph_walk, + list(req.wg_ids.keys), req.tokens, list(req.row_rids), + ) + except Exception: + logger.debug("runtime stop check failed (node=%s)", req.node_name, + exc_info=True) + if found is None: + return pp.step(PostprocessNeed.STOPS) + pp.stops = _merged_stops(req.stops, found) + self._keep_rows(pp, set()) + pp.stopped = self._stop_step_loops(pp) + pp.phase = _PostprocessPhase.ROUTE + if pp.phase is _PostprocessPhase.ROUTE: + if not pp.wg_ids.keys: + self._postprocess = None # every row was dropped + return pp.step(PostprocessNeed.DONE) + inline_values = None if req.inline is None else self._inline_values(pp) + if inline_values is None and pp.stored is None: + return pp.step(PostprocessNeed.STORE) + n = len(pp.wg_ids.keys) * len(req.output_signals) + stored = pp.stored or RouteTensors([], [0] * n, req.inline.signal, inline_values) + pp.route = self.complete_and_route_batch(RouteInput( + partition=req.partition, graph_walk=req.graph_walk, node_name=req.node_name, + output_signals=req.output_signals, wg_ids=pp.wg_ids, + tensors=stored.tensors, num_tensors=stored.num_tensors, + inline_signal=stored.inline_signal, inline_values=stored.inline_values, + logprobs=stored.logprobs, + walks=req.walks if pp.rid_walk_idx is not None else None, + rid_walk_idx=pp.rid_walk_idx, + )) + pp.phase = _PostprocessPhase.SEND + if pp.send is None: + if req.python_send or pp.stopped or _route_needs_python(pp.route): + return pp.step(PostprocessNeed.SEND) + pp.send = SendInput( + pp.route.completion_id, ParallelList([], []), ParallelList([], []), + ) + # the consuming pass reports the partition done, on its WORKER_GRAPHS_DONE + for rid in pp.partition_done: + self.mark_partition_done(rid, req.partition) + completed = self.send_outputs(pp.send) + self._postprocess = None + return pp.step(PostprocessNeed.DONE, completed) + + @staticmethod + def _keep_rows(pp: _Postprocess, dropped: set[int]) -> None: + """The rows the step routes: the request's, less ``dropped``.""" + req = pp.request + keys, values = req.wg_ids.keys, req.wg_ids.values + idx = req.rid_walk_idx + keep = [i for i, rid in enumerate(keys) if rid not in dropped] + pp.wg_ids = ParallelList([keys[i] for i in keep], [values[i] for i in keep]) + pp.rid_walk_idx = None if idx is None else [idx[i] for i in keep] + for rid in dropped: + pp.stops.pop(rid, None) + # a dropped row's streams do not end here + pp.partition_done = [rid for rid in pp.partition_done if rid not in dropped] + + def _stop_step_loops(self, pp: _Postprocess) -> list[int]: + """Stop the step's loops, per walk; the rids whose loops stopped.""" + req = pp.request + if not pp.stops: + return [] + walk_of = dict.fromkeys(pp.wg_ids.keys, req.graph_walk) + if pp.rid_walk_idx is not None: + walk_of = { + rid: req.walks[w] for rid, w in zip(pp.wg_ids.keys, pp.rid_walk_idx, strict=True) + } + by_walk: dict[str, dict[int, set[str]]] = {} + for rid, loops in pp.stops.items(): + if rid in walk_of: + by_walk.setdefault(walk_of[rid], {})[rid] = loops + stopped: list[int] = [] + for walk, walk_stops in by_walk.items(): + stopped += self.stop_loops_batched( + partition=req.partition, graph_walk=walk, last_node_run=req.node_name, + loop_names=ParallelList(list(walk_stops), [list(v) for v in walk_stops.values()]), + ) + return stopped + + @staticmethod + def _inline_values(pp: _Postprocess) -> list[int] | None: + """The inline column in route order; None when a row has no value.""" + req = pp.request + values = req.inline.column.tolist() + row_of = {rid: i for i, rid in enumerate(req.row_rids)} + out = [] + for rid in pp.wg_ids.keys: + i = row_of.get(rid) + if i is None or i >= len(values): + return None + out.append(values[i]) + return out + + +class GraphRuntime(PhasedPostprocess, ABC): """Owns a worker's graph state: its worker graphs' per-request queues, the routing and sharding derived from them, loop state, and the request id <-> handle table. The worker drives it through this contract and keeps @@ -690,6 +939,37 @@ def update_dynamic_loop_iters( ): per_request_info[rid].dynamic_loop_iter_counts.update(new_iters) + # --------- Declared stop policies ---------- + @abstractmethod + def set_stop_policies( + self, node_name: str, graph_walk: str, policies: Mapping[str, "StopPolicy"], + ) -> None: + """(node, walk)'s declared stop policies, per loop + (``NodeSubmodule.get_stop_policies``). Once, at startup.""" + + @abstractmethod + def register_stop_limits(self, rid: int, info: CurrentForwardPassInfo) -> None: + """A request's limits (``StopPolicy.limits``) under every declared + walk's policies, from its forward-pass info as it arrives. Dropped + with the request.""" + + @abstractmethod + def check_stops( + self, partition: str, node_name: str, graph_walk: str, rids: list[int], + tokens: "torch.Tensor | None", row_rids: list[int], + ) -> dict[int, set[str]] | None: + """(node, walk)'s declared stop policies over a step's rows: per rid, + the loops it stops (``StopPolicy.stops``); None when a row has no + limits for the walk, and the worker's ``check_stop`` decides. + + Each loop's count is this runtime's as of the step: counts move when + the step is routed, which comes after. ``tokens`` is the step's token + column, row i for ``row_rids[i]``, a 1-D host view the caller keeps + alive, or None when no policy reads a token. With a column, a rid with + no row is skipped, as ``check_stop`` skips a request with no output; + without one (``ITER`` only), every rid is checked. + """ + @abstractmethod def is_async_schedulable(self, node_name: str, graph_walk: str) -> bool: """Whether this node opts into async scheduling. Structural, so it diff --git a/mstar/graph/runtime/python.py b/mstar/graph/runtime/python.py index f7ecae1bd4..8157c49954 100644 --- a/mstar/graph/runtime/python.py +++ b/mstar/graph/runtime/python.py @@ -1,5 +1,7 @@ import logging +from collections.abc import Mapping from dataclasses import dataclass, field +from typing import TYPE_CHECKING from mstar.api_server.request_types import APIServerMessage, ResultTensors, ResultTokens from mstar.communication.communicator import BaseCommunicator @@ -39,6 +41,7 @@ SPECIAL_DESTINATIONS, ) from mstar.model.base import WorkerGraph +from mstar.model.submodule_base import StopLimits, StopPolicy from mstar.utils.ipc_format import ( ConductorMessage, ConductorMessageType, @@ -55,6 +58,9 @@ ) from mstar.worker.rid_table import RidTable +if TYPE_CHECKING: + import torch + logger = logging.getLogger(__name__) @@ -115,6 +121,9 @@ class GraphRuntimeRequestInfo: # stream edge -> destination worker -> items sent, local ones included pending_stream_items: dict[str, dict[str, int]] = field(default_factory=dict) current_output_chunks: list[str] = field(default_factory=list) + # (node, walk) -> loop -> this request's limits under its declared stop + # policies (``register_stop_limits``) + stop_limits: dict[tuple[str, str], dict[str, StopLimits]] = field(default_factory=dict) class PythonGraphRuntime(GraphRuntime): @@ -165,6 +174,8 @@ def __init__( # Loop stops from this iteration's check_stop; cleared every iteration. self._pending_loop_stops: set[PendingLoopStop] = set() + # (node, walk) -> its declared stop policies, per loop + self._stop_policies: dict[tuple[str, str], dict[str, StopPolicy]] = {} # Routing parked between complete_and_route_batch and send_outputs. self._completions: dict[int, CompletionState] = {} @@ -372,14 +383,57 @@ def get_dynamic_loop_iters( self, request_ids: list[int], partition: str, ) -> ParallelList[int, dict[str, int]]: - values = [] - for rid in request_ids: - iter_counts: dict[str, int] = {} - part_info = self._request_info[rid].partition_info[partition] - for wg_id in part_info.graph_walk_worker_graph_ids: - iter_counts.update(self._queues[wg_id].get_dynamic_loop_iters(rid)) - values.append(iter_counts) - return ParallelList(list(request_ids), values) + return ParallelList( + list(request_ids), [self._loop_counts(rid, partition) for rid in request_ids], + ) + + def _loop_counts(self, rid: int, partition: str) -> dict[str, int]: + """``rid``'s loop counts in ``partition``; the last worker graph with + a loop has its count.""" + iter_counts: dict[str, int] = {} + part_info = self._request_info[rid].partition_info[partition] + for wg_id in part_info.graph_walk_worker_graph_ids: + iter_counts.update(self._queues[wg_id].get_dynamic_loop_iters(rid)) + return iter_counts + + def set_stop_policies( + self, node_name: str, graph_walk: str, policies: Mapping[str, StopPolicy], + ) -> None: + self._stop_policies[(node_name, graph_walk)] = dict(policies) + + def register_stop_limits(self, rid: int, info: CurrentForwardPassInfo) -> None: + limits = self._request_info[rid].stop_limits + for key, policies in self._stop_policies.items(): + limits[key] = {loop: p.limits(info) for loop, p in policies.items()} + + def check_stops( + self, partition: str, node_name: str, graph_walk: str, rids: list[int], + tokens: "torch.Tensor | None", row_rids: list[int], + ) -> dict[int, set[str]] | None: + policies = self._stop_policies[(node_name, graph_walk)] + stops: dict[int, set[str]] = {} + if not policies: + return stops + values = None if tokens is None else tokens.tolist() + row_of = {rid: i for i, rid in enumerate(row_rids)} + for rid in rids: + info = self._request_info.get(rid) + if info is None: + continue + token = None + if values is not None: + i = row_of.get(rid) + if i is None or i >= len(values): + continue + token = values[i] + limits = info.stop_limits.get((node_name, graph_walk)) + if limits is None: + return None + counts = self._loop_counts(rid, partition) + for loop, p in policies.items(): + if p.stops(counts.get(loop, 0), token, limits[loop]): + stops.setdefault(rid, set()).add(loop) + return stops def get_walk(self, rid: int, partition: str) -> str: return self._request_info[rid].partition_info[partition].graph_walk diff --git a/mstar/graph/runtime/rust.py b/mstar/graph/runtime/rust.py index 94a0366ea9..8f9671d25f 100644 --- a/mstar/graph/runtime/rust.py +++ b/mstar/graph/runtime/rust.py @@ -19,6 +19,7 @@ from __future__ import annotations import time as _time +from collections.abc import Iterable, Mapping from copy import deepcopy from typing import TYPE_CHECKING, get_args @@ -37,15 +38,17 @@ ColumnarEdgeSpecs, FreedTensors, GraphRuntime, - InlineRoute, - InlineRouteStatus, ParallelList, PopRidsOutput, + PostprocessNeed, + PostprocessRequest, + PostprocessStep, Profiling, ReadyNodeSpec, RequestWalks, RouteInput, RouteOutput, + RouteTensors, SendInput, SpecRows, SpeculationOutput, @@ -53,6 +56,7 @@ SpeculationPrepOutput, ) from mstar.model.base import WorkerGraph +from mstar.model.submodule_base import StopPolicy from mstar.utils.profiler import PHASE_PERIOD, phase_record if TYPE_CHECKING: @@ -211,6 +215,14 @@ def _columns(out) -> ColumnarEdgeSpecs: ) +def _column(t: torch.Tensor | None) -> tuple[int, int, int, str]: + """A 1-D host view as Rust reads it in place: pointer, rows, stride and + dtype; a null pointer for none.""" + if t is None: + return 0, 0, 0, "" + return t.data_ptr(), t.shape[0], t.stride(0), _dtype_name(t.dtype) + + class RustGraphRuntime(GraphRuntime): """Delegates to ``mstar_rust.GraphRuntime``. @@ -240,6 +252,11 @@ def __init__( ): mine = {wg.worker_graph_id for wg in my_worker_graphs} combined_walk_of = combined_walk_of or {} + # (node, walk) -> its declared stop policies, per loop, in the order + # the Rust rules index them + self._stop_policies: dict[tuple[str, str], dict[str, StopPolicy]] = {} + self._postprocess_request: PostprocessRequest | None = None + self._postprocess_nogil_ns = 0 self._rust = _RustGraphRuntime( worker_graphs=[worker_graph_args(wg) for wg in my_worker_graphs], remote_worker_graphs=[ @@ -652,34 +669,112 @@ def _route_output(out) -> RouteOutput: ), ) - supports_route_and_send_inline = True + def set_stop_policies( + self, node_name: str, graph_walk: str, policies: Mapping[str, StopPolicy], + ) -> None: + self._stop_policies[(node_name, graph_walk)] = dict(policies) + self._rust.set_stop_rules( + node_name, graph_walk, list(policies), + [p.reads_token for p in policies.values()], + [sorted(p.eos_tokens) for p in policies.values()], + ) - def route_and_send_inline( - self, partition: str, graph_walk: str, node_name: str, - output_signals: list[str], rids: list[int], - inline_signal: str, inline_column: torch.Tensor, - row_rids: list[int], - ) -> InlineRoute: - """``complete_and_route_batch`` + ``send_outputs`` for a device-loopback - step, one crossing with the GIL released (see the Rust side). + def register_stop_limits(self, rid: int, info: CurrentForwardPassInfo) -> None: + cap = (1 << 32) - 1 # also "no limit": a loop count never reaches it + for (node, walk), policies in self._stop_policies.items(): + if not policies: + continue + limits = [p.limits(info) for p in policies.values()] + self._rust.set_stop_limits( + node, walk, [rid], + [[cap if lim.max_iters is None else min(lim.max_iters, cap)] for lim in limits], + [[lim.ignore_eos] for lim in limits], + ) - ``inline_column`` is a 1-D host view (row i for ``row_rids[i]``) the - caller keeps alive; Rust reads it in place. - """ - out = self._rust.route_and_send_inline( - partition=partition, graph_walk=graph_walk, node_name=node_name, - output_signals=list(output_signals), rids=rids, inline_signal=inline_signal, - inline_ptr=inline_column.data_ptr(), inline_rows=inline_column.shape[0], - inline_stride=inline_column.stride(0), - inline_dtype=_dtype_name(inline_column.dtype), row_rids=row_rids, + def check_stops( + self, partition: str, node_name: str, graph_walk: str, rids: list[int], + tokens: torch.Tensor | None, row_rids: list[int], + ) -> dict[int, set[str]] | None: + policies = self._stop_policies[(node_name, graph_walk)] + if not policies: + return {} + ptr, rows, stride, dtype = _column(tokens) + stop_rids, stop_rules, missing = self._rust.check_stops_inline( + partition, node_name, graph_walk, rids, ptr, rows, stride, dtype, row_rids, + ) + if missing: + return None + loops = list(policies) + stops: dict[int, set[str]] = {} + for rid, k in zip(stop_rids, stop_rules, strict=True): + stops.setdefault(rid, set()).add(loops[k]) + return stops + + def begin_postprocess(self, request: PostprocessRequest) -> PostprocessStep: + """Natively (``postprocess.rs``).""" + self._postprocess_request = request + self._postprocess_nogil_ns = 0 + stops = request.stops or {} + inline = request.inline + ptr, rows, stride, dtype = _column(request.tokens) + iptr, irows, istride, idtype = _column(None if inline is None else inline.column) + out = self._rust.begin_postprocess( + partition=request.partition, node_name=request.node_name, + graph_walk=request.graph_walk, output_signals=list(request.output_signals), + rids=request.wg_ids.keys, walks=request.walks, rid_walk_idx=request.rid_walk_idx, + stop_rids=list(stops), stop_loops=[sorted(v) for v in stops.values()], + runtime_stops=request.runtime_stops, + token_ptr=ptr, token_rows=rows, token_stride=stride, token_dtype=dtype, + row_rids=list(request.row_rids), + inline_signal=None if inline is None else inline.signal, + inline_ptr=iptr, inline_rows=irows, inline_stride=istride, inline_dtype=idtype, + partition_done=list(request.partition_done), python_send=request.python_send, + ) + return self._postprocess_step(out) + + def resume_postprocess( + self, stops: Mapping[int, set[str]] | None = None, dropped: Iterable[int] = (), + stored: RouteTensors | None = None, send: SendInput | None = None, + ) -> PostprocessStep: + if stops is not None: + out = self._rust.resume_postprocess_stops( + list(stops), [sorted(v) for v in stops.values()], list(dropped), + ) + elif stored is not None: + out = self._rust.resume_postprocess_store( + tensors=list(stored.tensors), num_tensors=list(stored.num_tensors), + inline_signal=stored.inline_signal, inline_values=list(stored.inline_values), + logprobs=list(stored.logprobs), + ) + elif send is not None: + # The parked step's own route output carries the completion id. + args = self._send_args(send) + del args["completion_id"] + out = self._rust.resume_postprocess_send(**args) + else: + raise ValueError("resume_postprocess: nothing to resume with") + return self._postprocess_step(out) + + def _postprocess_step(self, out) -> PostprocessStep: + self._postprocess_nogil_ns += out.nogil_ns + need = PostprocessNeed(out.need) + if need is PostprocessNeed.DONE: + if PHASE_PERIOD: # per step + phase_record("rust.postprocess.nogil", self._postprocess_nogil_ns * 1e-9) + self._postprocess_nogil_ns = 0 + self._postprocess_request = None + return PostprocessStep(need, stopped_rids=out.stopped_rids, + completed_rids=out.completed_rids) + if need is PostprocessNeed.STOPS: + return PostprocessStep(need) + wg_ids = self._postprocess_request.wg_ids + wg_of = dict(zip(wg_ids.keys, wg_ids.values, strict=True)) + rids = out.rids + route = out.take_route() + return PostprocessStep( + need, ParallelList(rids, [wg_of[r] for r in rids]), out.rid_walk_idx, + out.stopped_rids, None if route is None else self._route_output(route), ) - status = InlineRouteStatus(out.status) - route = None - if status is InlineRouteStatus.ROUTED: - route = self._route_output(out.take_route()) - if PHASE_PERIOD: - phase_record("rust.route_and_send_inline.nogil", out.nogil_ns * 1e-9) - return InlineRoute(status, out.completion_id, out.sent_rids, route) def stop_loops_batched( self, partition: str, graph_walk: str, last_node_run: str, @@ -737,12 +832,24 @@ def send_outputs(self, input: SendInput) -> list[int]: """ # The wrapper's own per-rid rebuilding, timed apart from the crossing. _t0 = _time.perf_counter() if PHASE_PERIOD else 0.0 + args = self._send_args(input) + if PHASE_PERIOD: + _t1 = _time.perf_counter() + phase_record("rust.send_outputs.marshal", _t1 - _t0) + sent_rids = self._rust.send_outputs(**args) + phase_record("rust.send_outputs.call", _time.perf_counter() - _t1) + return sent_rids + return self._rust.send_outputs(**args) + + @staticmethod + def _send_args(input: SendInput) -> dict: + """``send_outputs``' arguments as Rust takes them.""" # Straight through as struct-of-arrays. `ParallelList` already holds # keys and values as two flat lists -- which is the whole reason it # exists -- so zipping them into tuples here only to have Rust unzip # them is a per-rid interpreted loop per argument, per forward pass. # The count dicts go over as dicts; PyO3 extracts them. - args = dict( + return dict( completion_id=input.completion_id, info_rids=input.per_request_info.keys, request_infos=[ @@ -782,11 +889,3 @@ def send_outputs(self, input: SendInput) -> list[int]: ), skip_rids=list(input.failed_rids), ) - if PHASE_PERIOD: - _t1 = _time.perf_counter() - phase_record("rust.send_outputs.marshal", _t1 - _t0) - sent_rids = self._rust.send_outputs(**args) - phase_record("rust.send_outputs.call", _time.perf_counter() - _t1) - return sent_rids - else: - return self._rust.send_outputs(**args) diff --git a/mstar/model/orpheus/submodules.py b/mstar/model/orpheus/submodules.py index e0211d5c48..74ebc010a9 100644 --- a/mstar/model/orpheus/submodules.py +++ b/mstar/model/orpheus/submodules.py @@ -26,6 +26,8 @@ ModelInputsFromEngine, NodeInputs, NodeSubmodule, + StopKind, + StopPolicy, ) from mstar.streaming.stream_buffer import StreamChunkInfo @@ -191,9 +193,9 @@ def forward_batched( engine_inputs=engine_inputs, text_inputs=text_inputs ) - # row i is request i; a per-rid dict here would guard the compiled frame on the ids. - # As stop buffers too, the stop check reads one host copy of the rows - # instead of one per request. + # Row i is request i; a per-rid dict here would guard the compiled frame + # on the ids. As the stop buffers too, the stop check reads one host + # copy of the rows, and the declared stop policies run in the runtime. return BatchedModelOutput( row_outputs={"new_token": new_tokens}, check_stop_buffers={"new_token": new_tokens}, @@ -211,6 +213,12 @@ def postprocess( return outputs["text_inputs"] = outputs["new_token"] + def get_stop_policies(self, graph_walk: str) -> dict[str, StopPolicy]: + # check_stop's EOS / token-budget test, on every walk + return {"decode_loop": StopPolicy( + StopKind.AR, token_signal="new_token", eos_tokens=(self.config.stop_token_id,), sampler=SAMPLER, + )} + def check_stop( self, request_id: str, request_info: CurrentForwardPassInfo, diff --git a/mstar/model/qwen3_5/submodules.py b/mstar/model/qwen3_5/submodules.py index ac8ba3b800..112d46ff32 100644 --- a/mstar/model/qwen3_5/submodules.py +++ b/mstar/model/qwen3_5/submodules.py @@ -82,6 +82,15 @@ logger = logging.getLogger(__name__) +def token_stop(stop_token_ids) -> TokenStop: + """The LLM's stop: a decode row's stop id, checked by the graph runtime.""" + return TokenStop( + loop="decode_loop", buffer="new_token", stop_ids=frozenset(stop_token_ids), + sampler=SAMPLER, walk="decode", prefill_walks=frozenset({"prefill_text", "prefill_vision"}), + last_prefill_flag="last_prefill", runtime_check=True, + ) + + class LLMSubmodule(ARNodeSubmodule): # the decode walk reads its input ids off the sampler's slot masters reads_device_loopback = True @@ -139,11 +148,7 @@ def __init__( self.model = model self.config = config self.vision_config = vision_config - self.token_stop = TokenStop( - loop="decode_loop", buffer="new_token", stop_ids=config.stop_token_ids, - sampler=SAMPLER, walk="decode", prefill_walks=frozenset({"prefill_text", "prefill_vision"}), - last_prefill_flag="last_prefill", - ) + self.token_stop = token_stop(config.stop_token_ids) self._vision_sentinels: tuple[torch.Tensor, torch.Tensor] | None = None # ------------------------------------------------------------------ diff --git a/mstar/model/submodule_base.py b/mstar/model/submodule_base.py index df195c697d..745daadf00 100644 --- a/mstar/model/submodule_base.py +++ b/mstar/model/submodule_base.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from collections import defaultdict -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, field, fields, replace from enum import Enum from types import MappingProxyType @@ -72,6 +72,10 @@ class TokenStop(NamedTuple): # for a prompt prefilled over several passes: the ``step_metadata`` flag of # the pass whose token is the reply's (earlier passes' tokens are not) last_prefill_flag: str | None = None + # the graph runtime checks ``walk``'s stop ids itself, without calling + # check_stop (``ARNodeSubmodule.get_stop_policies``); the prompt walks keep + # check_stop, which can end the partition + runtime_check: bool = False class StopCheck(NamedTuple): @@ -131,6 +135,76 @@ class InputMetadata(NamedTuple): EMPTY_INPUT_METADATA = InputMetadata() +class StopKind(Enum): + """How a loop's stop is decided after each step (``StopPolicy``).""" + # the sampled token is one of the policy's EOS ids (unless the request + # ignores EOS), or the loop reached the request's iteration limit + AR = "ar" + # only the iteration limit + ITER = "iter" + # only the EOS ids: the loop's cap (``Model.get_loop_caps``) is its budget + EOS = "eos" + + +class StopLimits(NamedTuple): + """One request's limits under a ``StopPolicy``.""" + # None: no iteration limit (``StopKind.EOS``) + max_iters: int | None + ignore_eos: bool + + +@dataclass(frozen=True) +class StopPolicy: + """A loop's stop check, declared so the graph runtime can run it without + calling ``check_stop`` (``NodeSubmodule.get_stop_policies``). + + A row stops the loop when ``stops`` says so, given the loop's iteration + count as of the step (what ``check_stop`` reads as + ``dynamic_loop_iter_counts.get(loop, 0)``) and, for ``AR`` and ``EOS``, + the row's ``token_signal`` value. + """ + kind: StopKind + # AR, EOS: the output signal carrying the sampled token, read off the host + # rows; the first value of the row + token_signal: str | None = None + # any iterable of ids; held as a frozenset + eos_tokens: frozenset[int] = frozenset() + # AR, EOS: the resource label whose request config carries ``ignore_eos``; + # None means EOS is never ignored + sampler: str | None = None + # the request's iteration limit for this loop; evaluated once per request + # and walk, not per step. None is ``info.max_tokens``. + max_iters: Callable[[CurrentForwardPassInfo], int] | None = None + + def __post_init__(self): + if self.reads_token and self.token_signal is None: + raise ValueError(f"a {self.kind.name} StopPolicy needs a token_signal") + object.__setattr__(self, "eos_tokens", frozenset(self.eos_tokens)) + + @property + def reads_token(self) -> bool: + return self.kind is not StopKind.ITER + + def limits(self, info: CurrentForwardPassInfo) -> StopLimits: + """One request's limits, from its forward-pass info.""" + max_iters = None + if self.kind is not StopKind.EOS: + max_iters = int(info.max_tokens if self.max_iters is None else self.max_iters(info)) + ignore_eos = False + if self.reads_token and self.sampler is not None: + config = info.resource_configs.get(self.sampler) + ignore_eos = bool(getattr(config, "ignore_eos", False)) + return StopLimits(max_iters, ignore_eos) + + def stops(self, count: int, token: int | None, limits: StopLimits) -> bool: + """Whether a row with the loop at ``count`` and sampled ``token`` + stops it: the count plus one reaches the limit, or (``AR``, ``EOS``) + the token is an EOS id the request does not ignore.""" + return (limits.max_iters is not None and count + 1 >= limits.max_iters) or ( + self.reads_token and not limits.ignore_eos and token in self.eos_tokens + ) + + @dataclass class BatchedModelOutput: """A step's outputs, split by how they are addressed. @@ -1141,6 +1215,17 @@ def check_stop( """ return set() + def get_stop_policies(self, graph_walk: str) -> dict[str, StopPolicy] | None: + """The stop check ``check_stop`` amounts to on ``graph_walk``, per + loop it can stop, if it is only these, so the graph runtime can run + them without calling it. + + None (the default) means ``check_stop`` decides; an empty dict means + it never stops a loop on this walk. Asked once per walk, at startup. + Every policy must agree with ``check_stop`` for every row. + """ + return None + def check_stop_batched( self, request_ids: list[str], request_infos: dict[str, CurrentForwardPassInfo], @@ -1181,6 +1266,17 @@ class ARNodeSubmodule(NodeSubmodule): # sampled token ends a dynamic loop; the stop checks below follow from it. token_stop: TokenStop | None = None + def get_stop_policies(self, graph_walk: str) -> dict[str, StopPolicy] | None: + stop = self.token_stop + if stop is None or not stop.runtime_check or graph_walk in stop.prefill_walks: + return super().get_stop_policies(graph_walk) + if stop.walk is not None and graph_walk != stop.walk: + return {} # as `_token_stops`: only a row of ``walk`` ends the loop + # the budget is the loop's cap, which the runtime enforces + return {stop.loop: StopPolicy( + StopKind.EOS, token_signal=stop.buffer, eos_tokens=stop.stop_ids, sampler=stop.sampler, + )} + def check_stop( self, request_id: str, request_info: CurrentForwardPassInfo, diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index afdfdbba67..7d48bcc172 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -46,10 +46,13 @@ ColumnarEdgeSpecs, EdgeSpec, GraphRuntime, - InlineRouteStatus, + InlineColumn, + PostprocessNeed, + PostprocessRequest, + PostprocessStep, RequestWalks, - RouteInput, RouteOutput, + RouteTensors, SendInput, SpeculationOutput, SpeculationPrepInput, @@ -89,6 +92,8 @@ PhaseLaps, nvtx_enabled, phase_buffer, + phase_end, + phase_start, range_pop, range_push, ) @@ -127,6 +132,10 @@ _REMOVED_RIDS_KEEP = 4096 +# worker_phases span per message type handled in _process_message_list +_MSG_SPANS = {t: f"worker.msg.{t.name.lower()}" for t in WorkerMessageType} + + def _thread_role() -> str: """The worker thread this runs on, by the executors' name prefixes.""" name = threading.current_thread().name @@ -279,6 +288,14 @@ class EvictionPolicy(Enum): # question. Needs eviction-policy metadata naming which resource to consult. +class ModelStops(NamedTuple): + """``Worker._check_stops``: the model's stops by rid, the rows it failed + (which leave the batch), and the rows whose partition it ended.""" + stops: dict[int, set[str]] + dropped: list[int] + ended: list[int] + + class SpeculationTarget(NamedTuple): info: SpeculationOutput walk_rids: list[int] # the step's rids on ``info.graph_walk`` @@ -289,6 +306,14 @@ class SpeculationTarget(NamedTuple): _HOST_SPANS_PURPOSE = tuple(f"output_spans{slot}" for slot in range(_HOST_ROWS_RING)) +def _host_column(host_rows: HostRows | None, name: str) -> torch.Tensor | None: + """Host buffer ``name``'s first value per row, a view; None without one.""" + buf = None if host_rows is None else host_rows.buffers.get(name) + if not torch.is_tensor(buf) or buf.dim() == 0 or buf.numel() == 0: + return None + return buf.reshape(buf.shape[0], -1)[:, 0] + + class _LazyPerRid: """``Worker._rows_to_per_rid`` on demand: the per-rid host dicts are built the first time a consumer reads them. The batched stop check, the row store @@ -528,6 +553,7 @@ def __init__( device_capability=self.device_capability, ) self._never_publish_unshared(all_worker_graph_ids_to_instances) + self._declare_stop_policies(my_walks) self.request_state = RequestStateManager( node_to_partition=node_to_partition, @@ -843,6 +869,7 @@ def _add_new_request(self, body: NewRequest) -> None: if not self._takes_input(body.request_id): return logger.debug("Worker %s received request %s", self.worker_id, body.request_id) + laps = PhaseLaps("worker.msg.new_request.") # The one place a handle is minted; a request with several partitions # on this worker gets the same one for each. Everything below keys on it. request_id = self._graph_runtime.add_request( @@ -858,17 +885,22 @@ def _add_new_request(self, body: NewRequest) -> None: # Drained before admission: take no new work, as ``_begin_drain`` # does for an admitted rid, while committed TP steps still run. self.scheduler.fail_rids({request_id}) + laps.lap("runtime_add") self.request_state.add_request(request_id, body.request_info) self.scheduler.note_arrival(request_id) + self._graph_runtime.register_stop_limits(request_id, body.request_info) self.tensor_manager.register_request( request_id, self._graph_runtime.get_sharding_config(request_id), ) + laps.lap("worker_state") try: self.engine_manager.add_request( request_id, body.request_info.resource_configs, body.request_info, ) + laps.lap("engine_add") self._skip_unshared_publish(request_id, body.worker_graph_to_workers) + laps.lap("skip_publish") except Exception as err: # the engine rolled the ingest back; the conductor's REMOVE tears down the rest text, _ = client_error_text( @@ -882,12 +914,14 @@ def _add_new_request(self, body: NewRequest) -> None: self.streams.add_request(request_id, body.request_info) + laps.lap("stream_buffers") # Start RDMA reads for tensors that have tensor_info futures = self.tensor_manager.start_read_tensors( request_id, body.initial_inputs, graph_walk=body.request_info.graph_walk ) self.wakeup_event.register_futures(futures) + laps.lap("start_reads") # Signal-only edges (tensor_info is None) can be processed immediately signal_only = [ @@ -897,6 +931,7 @@ def _add_new_request(self, body: NewRequest) -> None: self._graph_runtime.ingest_inputs_batch( self._one_rids_edges(request_id, signal_only), can_buffer=True, ) + laps.lap("ingest_signals") self._process_unprocessed(body.request_id) def _process_unprocessed(self, wire_rid: str) -> None: @@ -1011,24 +1046,30 @@ def _remove_request(self, body: RemoveRequest) -> None: # REMOVE_REQUEST once every reader has confirmed drained (READS_DONE) — # on the abort/fail path via a prior DrainRequest, on the happy path # once the api server finished reading the outputs. + laps = PhaseLaps("worker.msg.remove_request.") self._mark_removed(body.request_id) freed = self.engine_manager.remove_request(request_id) if freed: self.scheduler.lift_batch_caps(freed) + laps.lap("engine") if self._lockstep_led: self._lockstep_unread.pop(request_id, None) self.streams.remove_request(request_id) self.request_state.remove_request(request_id) self.tensor_manager.force_cleanup_request(request_id) + laps.lap("worker_state_tensors") self.profile_info.pop_request(request_id) self.streaming_buffers.pop(request_id, None) self.scheduler.clear_rid(request_id, body.request_id) self._pending_removes.pop(request_id, None) + laps.lap("misc") logger.info("Request cleanup complete: %s", body.request_id) + laps.lap("log") # Last: frees the handle for reuse, so nothing above may run after it. self._graph_runtime.remove_request(request_id) + laps.lap("runtime") def _drain_request(self, body: DrainRequest) -> None: """Phase-1 teardown (abort/fail): stop scheduling this rid, then ACK @@ -1219,6 +1260,7 @@ def _process_new_inputs(self, body: InputSignals) -> None: request_id, current_fwd_info=body.request_info, partition_name=body.partition_name ) + self._graph_runtime.register_stop_limits(request_id, body.request_info) if self.enable_nvtx: range_pop(synchronize=False) @@ -1304,6 +1346,7 @@ def _process_message_list(self, messages: list[WorkerMessage]): wire_rid, [] ).append(message) continue + t0 = phase_start() if message.message_type == WorkerMessageType.NEW_REQUEST: self._add_new_request(message.body) elif message.message_type == WorkerMessageType.DRAIN_REQUEST: @@ -1325,6 +1368,8 @@ def _process_message_list(self, messages: list[WorkerMessage]): elif message.message_type == WorkerMessageType.ENABLE_METRICS: # sent only with metrics on, so _metrics is set self._metrics.start_counting() + if t0 is not None: + phase_end(_MSG_SPANS[message.message_type], t0) def _process_messages(self) -> None: self._process_message_list(self.communicator.get_all_new_messages()) @@ -3370,55 +3415,200 @@ def _pp_stage(name: str) -> None: _pp_stage("crop_spans") if self.enable_nvtx: range_pop(synchronize=False) - range_push("worker.postprocess.check_stop", synchronize=False) + range_push("worker.postprocess.prematerialize", synchronize=False) - # Check for stops. Prematerialising pulls sampled tokens to the host, - # so this can carry a device transfer as well as the stop logic. - _t_stop = _time.perf_counter() if self._phase_period else 0.0 + # Prematerialising pulls the stop check's rows to the host, so this can + # carry a device transfer. engine = self.engine_manager.get_engine(batch_N.node_name) cpu_outputs, host_rows = self._prematerialize_for_check_stop( outputs, batch_N.node_batch.completion_event, request_ids=batch_N.node_batch.request_ids, ) + engine.step_finished(batch_N.node_batch, cpu_outputs, host_rows=host_rows) _pp_stage("prematerialize") - stops = engine.check_stop_for_batch( - batch_N.node_batch, cpu_outputs, host_rows=host_rows, + runtime_stops, tokens = self._runtime_stop_column(batch_N, host_rows) + checked = None + if not runtime_stops: + # the model's check_stop decides, before the runtime, which then + # routes without stopping for it + checked = self._check_stops(batch_N, outputs, cpu_outputs, host_rows, engine).stops + _pp_stage("check_stop") + request = self._get_postprocess_request( + batch_N, host_rows, engine, runtime_stops, tokens, checked, ) + if self.enable_nvtx: + range_pop(synchronize=False) + range_push("worker.postprocess.runtime", synchronize=False) + + # The runtime checks the step's stops, routes its outputs and sends + # them, stopping for what only the worker has (PostprocessNeed). + # Main loop only: the runtime holds `&mut self` across its GIL + # release, so a concurrent caller gets "Already mutably borrowed" + # rather than blocking. + rt = self._graph_runtime + _t_rt = _time.perf_counter() if self._phase_period else 0.0 + step = rt.begin_postprocess(request) + rt_time = _time.perf_counter() - _t_rt if self._phase_period else 0.0 + stored = None + sent_here = False + while step.need is not PostprocessNeed.DONE: + _t = _time.perf_counter() if self._phase_period else 0.0 + if step.need is PostprocessNeed.STOPS: + model_stops = self._check_stops(batch_N, outputs, cpu_outputs, host_rows, engine) + # after the request was built, which carries the step's other + # partition ends; the runtime reports them on the same send + for rid in model_stops.ended: + rt.mark_partition_done(rid, batch_N.partition) + stage = "check_stop" + _t_rt = _time.perf_counter() if self._phase_period else 0.0 + step = rt.resume_postprocess(stops=model_stops.stops, dropped=model_stops.dropped) + elif step.need is PostprocessNeed.STORE: + # What the stop check withheld (``StopCheck.withheld``), and an + # output cut to an empty span, is not stored, so not routed: a + # streaming edge then carries no item for the row, and the + # producer's item count leaves it out as well. + for withheld in (batch_N.node_batch.withheld_outputs, empty_spans): + for rid, names in withheld.items(): + rid_outputs = outputs.get(rid) + if rid_outputs: + for name in names: + rid_outputs.pop(name, None) + stored, route_tensors = self._store_outputs( + batch_N, outputs, request, step, cpu_outputs, host_rows, engine, + ) + stage = "store_tensors" + _t_rt = _time.perf_counter() if self._phase_period else 0.0 + step = rt.resume_postprocess(stored=route_tensors) + else: + send = self._send_input(batch_N, request, step, stored, engine) + sent_here = True + stage = "send_prep" + _t_rt = _time.perf_counter() if self._phase_period else 0.0 + step = rt.resume_postprocess(send=send) + if self._phase_period: + now = _time.perf_counter() + self._phase_record(f"worker.postprocess.{stage}", _t_rt - _t) + rt_time += now - _t_rt if self._phase_period: - self._phase_record( - "worker.postprocess.check_stop", _time.perf_counter() - _t_stop, + self._phase_record("worker.postprocess.runtime", rt_time) + _pp_t = _time.perf_counter() + + if self.enable_nvtx: + range_pop(synchronize=False) + if not sent_here and batch_N.batch.request_to_worker_graph: + # sent without the worker (SEND does these before its send) + self._buffer_publish(batch_N) + rt.update_dynamic_loop_iters( + batch_N.node_batch.per_request_info, batch_N.partition, ) - # the same host copy, before stops, so a request ending here still indexes its pages - engine.extend_prefix_chains(batch_N.node_batch, cpu_outputs, host_rows=host_rows) - # What the stop check withheld (``StopCheck.withheld``), and an output - # cut to an empty span, is not stored, so not routed: a streaming edge - # then carries no item for the row, and the producer's item count - # leaves it out as well. - for withheld in (batch_N.node_batch.withheld_outputs, empty_spans): - for rid, names in withheld.items(): - rid_outputs = outputs.get(rid) - if rid_outputs: - for name in names: - rid_outputs.pop(name, None) + for rid in step.completed_rids: + self.request_state.flush_publish_info(rid, batch_N.partition) + # Those sends were the last readers of this step's host copies. + self.tensor_manager.drop_host_copies() + _pp_stage("tail") + engine.after_output_copy(batch_N.node_batch) + _pp_stage("after_output_copy") + def _runtime_stop_column( + self, batch_N: PendingBatch, host_rows: "HostRows | None", + ) -> tuple[bool, "torch.Tensor | None"]: + """Whether the runtime checks the step's declared stop policies, and + the token column they read (None when none reads one). A step that + mixes walks, has no declared policies, or has no column keeps + check_stop.""" + key = (batch_N.node_name, batch_N.graph_walk) + if key not in self._stop_token_signal or len(batch_N.batch.rids_by_walk()) > 1: + return False, None + token_signal = self._stop_token_signal[key] + if token_signal is None: + return True, None + tokens = _host_column(host_rows, token_signal) + return tokens is not None, tokens + + def _get_postprocess_request( + self, batch_N: PendingBatch, host_rows: "HostRows | None", engine, + runtime_stops: bool, tokens: "torch.Tensor | None", + checked: dict[int, set[str]] | None, + ) -> PostprocessRequest: + """The step's ``PostprocessRequest``: its rows, its stops known before + the runtime's check (``checked``: the model's, when the runtime does + not check), and the host columns the runtime reads in place.""" + batch = batch_N.batch + by_walk = batch.rids_by_walk() + walks = list(by_walk) or [batch_N.graph_walk] + if len(by_walk) <= 1: + rids, rid_walk_idx = list(by_walk.get(walks[0], ())), None + # a combined batch keeps its first walk's names, which may not be this walk's + signals = ( + list(self._walk_signals(batch, walks[0])) if batch.request_walks + else batch.output_signals + ) + else: + rids, rid_walk_idx = [], [] + for i, walk_rids in enumerate(by_walk.values()): + rids.extend(walk_rids) + rid_walk_idx.extend([i] * len(walk_rids)) + # several walks route under the union of their names; a walk's + # node ignores the ones it lacks + signals = list(dict.fromkeys( + signal for walk in walks for signal in self._walk_signals(batch, walk) + )) # Stream-terminated loop: a stream-consuming node inside a loop has no - # internal stop signal (unlike a self-EOS loop), so when it consumes the - # terminal chunk of its stream (``final_stream_rids``: the StreamBuffer - # popped ``is_final``, which a ``continue_after_done`` policy never - # sets) end its innermost enclosing loop. Without this the loop would - # spin to ``max_iters`` and the partition would never report done. + # internal stop signal (unlike a self-EOS loop), so when it consumes + # the terminal chunk of its stream (``final_stream_rids``: the + # StreamBuffer popped ``is_final``, which a ``continue_after_done`` + # policy never sets) end its innermost enclosing loop. Without this + # the loop would spin to ``max_iters`` and the partition would never + # report done. + stops: dict[int, set[str]] = {} for rid in batch_N.node_batch.final_stream_rids: - loop_name = self._innermost_loop_by_node.get( - (batch_N.batch.walk_of(rid), batch_N.node_name) - ) + loop_name = self._innermost_loop_by_node.get((batch.walk_of(rid), batch_N.node_name)) if loop_name is not None: stops.setdefault(rid, set()).add(loop_name) + for rid, loops in (checked or {}).items(): + stops.setdefault(rid, set()).update(loops) + # Device loop-back: every routed signal stays on the device and the + # client's copy rides inline, so the route stores nothing. + # A step with log-probs stores instead: STORE passes them to the route. + inline = None + if ( + host_rows is not None and rid_walk_idx is None + and SAMPLED_LOGPROB not in host_rows.buffers + # an inline route carries every output, a withheld one too + and not batch_N.node_batch.withheld_outputs + and self._all_on_device(engine, batch_N.node_name, batch_N.graph_walk, signals) + ): + client = engine.inline_client_signals(batch_N.node_name, batch_N.graph_walk) + signal = next((s for s in client if s in signals), None) + column = None if signal is None else _host_column(host_rows, client[signal]) + if column is not None: + inline = InlineColumn(signal, column) + return PostprocessRequest( + partition=batch_N.partition, node_name=batch_N.node_name, graph_walk=walks[0], + output_signals=signals, + wg_ids=ParallelList(rids, [batch.request_to_worker_graph[r] for r in rids]), + walks=walks if rid_walk_idx is not None else None, rid_walk_idx=rid_walk_idx, + stops=stops, runtime_stops=runtime_stops, tokens=tokens, + row_rids=host_rows.request_ids if host_rows is not None else (), + inline=inline, + # the consuming pass (not the earlier ingest) reports the partition + # done, so it rides the WGD that also counts this pass's last chunks + partition_done=list(batch_N.node_batch.partition_done_rids), + python_send=self.enable_prof, + ) - # END_PARTITION (e.g. a prefill that sampled a stop token) ends the - # partition on this pass's report. Only a row outside any loop can end - # its pass here: inside one it would stop only the innermost loop, and - # a speculated next iteration would carry no report of it. Returning it - # from a loop is a model bug, so the client sees an internal error. + def _check_stops( + self, batch_N: PendingBatch, outputs: BatchedModelOutput, cpu_outputs, + host_rows: "HostRows | None", engine, + ) -> ModelStops: + """The model's ``check_stop``, less the ``END_PARTITION`` stops, which + end their rows' partitions instead (``partition_done_rids``).""" + stops = engine.check_stop_for_batch(batch_N.node_batch, cpu_outputs, host_rows=host_rows) + ended = [] + # Only a row outside any loop can end its pass here: inside one it + # would stop only the innermost loop, and a speculated next iteration + # would carry no report of it. Returning it from a loop is a model + # bug, so the client sees an internal error. for rid in [rid for rid, s in stops.items() if END_PARTITION in s]: walk = batch_N.batch.walk_of(rid) if (walk, batch_N.node_name) in self._innermost_loop_by_node: @@ -3430,244 +3620,67 @@ def _pp_stage(name: str) -> None: else: del stops[rid] batch_N.node_batch.partition_done_rids.add(rid) - - if batch_N.node_batch.failed_requests: - # A rid whose stop check raised (or was refused above) has no - # trustworthy stop decision: routing it would either run its loop - # forever or end it early. Fail it here and take it out of the - # batch before the routing loops below touch it. Only ever the - # ones this stop stage just added: the caller already reported - # (and cleared) the rids that failed in prepare_inputs / postprocess. - failed = dict(batch_N.node_batch.failed_requests) - self._drop_failed_rids(batch_N, outputs, failed) - self._fail_requests(failed) - # `stops` was filled above, before these rows left the batch; - # stopping one would finalize state the drop just removed - for rid in failed: - stops.pop(rid, None) - if not batch_N.node_batch.request_ids and not batch_N.node_batch.cached_rids: - if self.enable_nvtx: - range_pop(synchronize=False) - return - - # the stop check itself is the direct window above; this stage is its - # tail (prefix chains, stream-terminated loops, failed rids, partition ends) - _pp_stage("check_stop_tail") - if self.enable_nvtx: - range_pop(synchronize=False) - range_push("worker.postprocess.stop_loops", synchronize=False) - - - # Stop loops, if applicable. The runtime filters rids whose walk does - # not contain the loop, snapshots the stop times, records the pending - # stops and fans out to peers. - # Main loop only: the runtime holds `&mut self` across its GIL - # release, so a concurrent caller gets "Already mutably - # borrowed" rather than blocking. - stopped_rids = [] - stops_by_walk: dict[str, dict[int, set[str]]] = {} - for rid, loops in stops.items(): - stops_by_walk.setdefault(batch_N.batch.walk_of(rid), {})[rid] = loops - for walk, walk_stops in stops_by_walk.items(): - stopped_rids += self._graph_runtime.stop_loops_batched( - partition=batch_N.partition, - graph_walk=walk, - last_node_run=batch_N.node_name, - loop_names=ParallelList( - list(walk_stops), [list(v) for v in walk_stops.values()], - ), - ) - - # Ordinary publication precedes stop detection on the GPU thread. - # Export final-only state now, after the last iteration committed and - # before routing reports the completed loop to the conductor. The - # ordinary publication of a stopping request is finished first, so the - # completion carries both, as it did when every step published. - if stopped_rids: - engine.materialize_publish(batch_N.node_batch, stopped_rids) - engine.finalize_stopped_requests(batch_N.node_batch, stopped_rids) - - # CurrentForwardPassInfo also contains publication inherited from peer - # ranks. Buffer only what this worker produced for the conductor. - buffered_publish: set[int] = set() + ended.append(rid) + if not batch_N.node_batch.failed_requests: + return ModelStops(stops, [], ended) + # A rid whose stop check raised (or was refused above) has no + # trustworthy stop decision: routing it would either run its loop + # forever or end it early. Only ever the ones this stop stage just + # added: the caller already reported (and cleared) the rids that + # failed in prepare_inputs / postprocess. + failed = dict(batch_N.node_batch.failed_requests) + self._drop_failed_rids(batch_N, outputs, failed) + self._fail_requests(failed) + for rid in failed: + stops.pop(rid, None) + return ModelStops(stops, list(failed), ended) + + def _declare_stop_policies(self, walks: set[tuple[str, str]]) -> None: + """Give the graph runtime each local (node, walk)'s declared stop + policies (``NodeSubmodule.get_stop_policies``), and note the token + signal each one reads. A walk whose policies read two token signals + keeps ``check_stop``.""" + # (node, walk) -> the token signal its policies read, or None + self._stop_token_signal: dict[tuple[str, str], str | None] = {} + for node, walk in sorted(walks): + policies = self.engine_manager.get_engine(node).get_stop_policies(node, walk) + if policies is None: + continue + signals = {p.token_signal for p in policies.values() if p.reads_token} + if len(signals) > 1: + logger.warning( + "%s/%s: stop policies read several token signals %s; check_stop decides", + node, walk, sorted(signals), + ) + continue + self._graph_runtime.set_stop_policies(node, walk, policies) + self._stop_token_signal[(node, walk)] = next(iter(signals), None) + + def _buffer_publish(self, batch_N: PendingBatch) -> set[int]: + """Buffer this step's publication for the conductor; the rids that + published. CurrentForwardPassInfo also contains publication inherited + from peer ranks, so only what this worker produced goes.""" + buffered: set[int] = set() for rid, published in batch_N.node_batch.resource_publish_info.items(): if published: - self.request_state.buffer_publish_info( - rid, batch_N.partition, published, - ) - buffered_publish.add(rid) - - _pp_stage("stop_loops") - if self.enable_nvtx: - range_pop(synchronize=False) - range_push("worker.postprocess.route_outputs", synchronize=False) - # Store this batch's output tensors, then hand the runtime their - # uuids. The store keeps the descriptors, so routing needs no - # TensorPointerInfo objects. - # One store and route for the step, whatever walks its rows ran. - by_walk = batch_N.batch.rids_by_walk() - walks = list(by_walk) - if len(walks) == 1: - rids, rid_walk_idx = by_walk[walks[0]], None - else: - rids, rid_walk_idx = [], [] - for i, walk_rids in enumerate(by_walk.values()): - rids.extend(walk_rids) - rid_walk_idx.extend([i] * len(walk_rids)) - if ( - rid_walk_idx is None - and self._route_and_send_fast( - batch_N, host_rows, engine, rids, buffered_publish, - ) - ): - _pp_stage("route_and_send") - if self.enable_nvtx: - range_pop(synchronize=False) - engine.after_output_copy(batch_N.node_batch) - _pp_stage("after_output_copy") - return - stored, route_output, signals = self._store_and_route( - batch_N, outputs, rids, walks, rid_walk_idx, cpu_outputs, - host_rows, engine, - ) - # One runtime call rather than a Python loop over the rows (#368). - self._graph_runtime.update_dynamic_loop_iters( - batch_N.node_batch.per_request_info, batch_N.partition, - ) - - # Normally empty: cleanup_consumed_inputs ran above and took them. Not - # empty if a completion ever precedes it, and then nobody else will. - if route_output.freed_inputs.uuids: - self.tensor_manager.cleanup_collectable(*route_output.freed_inputs) - - _pp_stage("route_outputs") - if self.enable_nvtx: - range_pop(synchronize=False) - range_push("worker.postprocess.register_outputs", synchronize=False) - failed_rids = self._register_outputs(route_output) - # Those sends were the last readers of this step's host copies. - self.tensor_manager.drop_host_copies() - _pp_stage("register_outputs") - - # send outputs - if self.enable_nvtx: - range_pop(synchronize=False) - range_push("worker.send_outputs", synchronize=False) - - # The consuming pass (not the earlier ingest) reports the partition - # done, so it rides the WGD that also counts this pass's last chunks; - # an END_PARTITION stop rides this pass's WGD the same way. - for rid in batch_N.node_batch.partition_done_rids: - self._graph_runtime.mark_partition_done(rid, batch_N.partition) - - # set this before send_outputs so that we can send updated profiling info to the conductor - if self.enable_prof: - self.profile_info.register_end( - batch_N.node_batch.node_name, - batch_N.node_batch.graph_walk, - batch_N.node_batch.request_ids, - batch_N.node_batch.exec_timings, - ) - - self._stream_and_send( - batch_N, rids, stored, route_output, signals, - engine, buffered_publish, failed_rids, - ) - - _pp_stage("send_outputs") - if self.enable_nvtx: - range_pop(synchronize=False) - engine.after_output_copy(batch_N.node_batch) - _pp_stage("after_output_copy") - - def _route_and_send_fast( - self, batch_N: PendingBatch, host_rows: "HostRows | None", engine, - rids: list[int], buffered_publish: set[int], - ) -> bool: - """Route and send a device-loopback step in one runtime call, with the - GIL released, reading the inline token column off the host rows. - - Qualifies only where ``_store_and_route`` would take its device-loopback - branch and ``_stream_and_send`` would add nothing Python-side before - the send. Returns False, having done nothing, otherwise. A step whose - stop check withheld an output (``StopCheck.withheld``) does not, as - this path routes every output. - """ - rt = self._graph_runtime - batch = batch_N.batch - signals = batch.output_signals - if ( - # log-probs ride the route input, which this path doesn't build - host_rows is None or SAMPLED_LOGPROB in host_rows.buffers - or batch.request_walks or self.enable_prof - or batch_N.node_batch.partition_done_rids - or batch_N.node_batch.withheld_outputs - or not getattr(rt, "supports_route_and_send_inline", False) - or not self._all_on_device( - engine, batch_N.node_name, batch_N.graph_walk, signals, - ) - ): - return False - inline = engine.inline_client_signals(batch_N.node_name, batch_N.graph_walk) - signal = next((s for s in inline if s in signals), None) - if signal is None: - return False - buf = host_rows.buffers.get(inline[signal]) - if not torch.is_tensor(buf) or buf.dim() == 0 or buf.numel() == 0: - return False - column = buf.reshape(buf.shape[0], -1)[:, 0] - inline = rt.route_and_send_inline( - partition=batch_N.partition, graph_walk=batch.graph_walk, - node_name=batch_N.node_name, output_signals=signals, rids=rids, - inline_signal=signal, inline_column=column, - row_rids=list(host_rows.request_ids), - ) - if inline.status is InlineRouteStatus.NOT_ROUTED: - return False - rt.update_dynamic_loop_iters( - batch_N.node_batch.per_request_info, batch_N.partition, - ) - route_output = inline.route_output - if inline.status is InlineRouteStatus.ROUTED: - # Routed; the send needs what only Python has: the general tail. - if route_output.freed_inputs.uuids: - self.tensor_manager.cleanup_collectable(*route_output.freed_inputs) - self._register_outputs(route_output) - stored = StoredOutputs([], [], [], [0] * (len(rids) * len(signals))) - self._stream_and_send( - batch_N, rids, stored, route_output, signals, engine, - buffered_publish, - ) - return True - for rid in inline.sent_rids: - self.request_state.flush_publish_info(rid, batch_N.partition) - return True + self.request_state.buffer_publish_info(rid, batch_N.partition, published) + buffered.add(rid) + return buffered - def _store_and_route( + def _store_outputs( self, batch_N: PendingBatch, outputs: BatchedModelOutput, - rids: list[int], walks: list[str], rid_walk_idx: list[int] | None, cpu_outputs, + request: PostprocessRequest, step: PostprocessStep, cpu_outputs, host_rows: "HostRows | None" = None, engine=None, - ): - """Store and route ``rids``; ``rids[i]`` ran ``walks[rid_walk_idx[i]]``, - or ``walks[0]`` when ``rid_walk_idx`` is None.""" - # Recorded by the pop that scheduled this batch, not asked for again -- - # and taken from the GRAPH, so a model returning a tensor under a name - # no edge carries cannot change what gets routed. Stale outputs are - # dropped by complete_and_route_batch itself. Several walks route under - # the union of their names; a walk's node ignores the ones it lacks. - batch = batch_N.batch - graph_walk = walks[0] - if not batch.request_walks: - signals = batch.output_signals - else: - signals = list(dict.fromkeys( - signal for walk in walks for signal in self._walk_signals(batch, walk) - )) + ) -> tuple[StoredOutputs, RouteTensors]: + """STORE: store the step's routed rows' outputs, held for the route; + what the route takes of them.""" + rids = step.wg_ids.keys + signals = request.output_signals + graph_walk = request.graph_walk # One store call for the batch rather than one per request, and the # flat columns come back already built: the manager fills them as it # mints, so nothing is keyed by request and signal only to be taken # apart again here. - _t_store = _time.perf_counter() if self._phase_period else 0.0 # Any of the walks names the source's TP group: a combined walk's walks # share one (refuse_combined_walks_across_tp_groups). # Pass the stop check's host copies so a host-memory transport needn't @@ -3733,11 +3746,10 @@ def store(rids): stored = store(rids) if self._lockstep_led and batch_N.node_name not in self.parallel_nodes: self._note_local_lockstep_inputs( - stored, signals, batch_N.node_name, rids, walks, rid_walk_idx, + stored, signals, batch_N.node_name, rids, + request.walks or [graph_walk], step.rid_walk_idx, ) flat_uuids = stored.flat_uuids - flat_rids = stored.flat_rids - signal_idxs = stored.signal_idxs num_tensors = stored.num_tensors # Safety hold: ref=1 until the real fanout is known, which # complete_and_route_batch settles. One call for the batch rather than @@ -3746,56 +3758,53 @@ def store(rids): # so it crosses as a scalar rather than a list built per batch. if flat_uuids: self.tensor_manager.increment_ref_batch_uniform(flat_uuids, 1) - if self._phase_period: - self._phase_record( - "worker.postprocess.store_tensors", - _time.perf_counter() - _t_store, - ) - # The graph runtime's own share of postprocess: the routing call. - _t_route = _time.perf_counter() if self._phase_period else 0.0 - route_output = self._graph_runtime.complete_and_route_batch( - RouteInput( - partition=batch_N.partition, - graph_walk=graph_walk, - node_name=batch_N.node_name, - output_signals=signals, - wg_ids=ParallelList( - rids, - [batch_N.batch.request_to_worker_graph[r] for r in rids], - ), - tensors=flat_uuids, - num_tensors=num_tensors, - inline_signal=inline_signal, - inline_values=inline_values, - logprobs=logprobs, - walks=walks if rid_walk_idx is not None else None, - rid_walk_idx=rid_walk_idx, - ), + return stored, RouteTensors(flat_uuids, num_tensors, inline_signal, inline_values, logprobs) + + def _send_input( + self, batch_N: PendingBatch, request: PostprocessRequest, step: PostprocessStep, + stored: StoredOutputs | None, engine, + ) -> SendInput: + """SEND: what the send needs of the worker, after the stopped rows' + finalize, the step's publication and the route's tensor work.""" + rids = step.wg_ids.keys + signals = request.output_signals + route_output = step.route_output + if stored is None: # routed inline: nothing stored + stored = StoredOutputs([], [], [], [0] * (len(rids) * len(signals))) + # Ordinary publication precedes stop detection on the GPU thread. + # Export final-only state now, after the last iteration committed and + # before the send reports the completed loop to the conductor. The + # ordinary publication of a stopping request is finished first, so the + # completion carries both, as it did when every step published. + if step.stopped_rids: + engine.materialize_publish(batch_N.node_batch, step.stopped_rids) + engine.finalize_stopped_requests(batch_N.node_batch, step.stopped_rids) + buffered_publish = self._buffer_publish(batch_N) + # One runtime call rather than a Python loop over the rows (#368). + self._graph_runtime.update_dynamic_loop_iters( + batch_N.node_batch.per_request_info, batch_N.partition, ) - if self._phase_period: - self._phase_record( - "worker.postprocess.route", _time.perf_counter() - _t_route, + # Normally empty: cleanup_consumed_inputs took them before the route. + # Not empty if a completion ever precedes it, and then nobody else will. + if route_output.freed_inputs.uuids: + self.tensor_manager.cleanup_collectable(*route_output.freed_inputs) + # a request whose outputs the transfer engine refused failed alone: none of it is sent + failed_rids = self._register_outputs(route_output) + if self.enable_prof: + self.profile_info.register_end( + batch_N.node_batch.node_name, + batch_N.node_batch.graph_walk, + batch_N.node_batch.request_ids, + batch_N.node_batch.exec_timings, ) - return stored, route_output, signals - - def _stream_and_send( - self, batch_N: PendingBatch, rids: list[int], stored, route_output, signals, - engine=None, buffered_publish: set[int] | None = None, - failed_rids: frozenset[int] = frozenset(), - ) -> None: - # Shared across the per-walk calls, so a rid is not buffered twice. - buffered_publish = set() if buffered_publish is None else buffered_publish - flat_uuids = stored.flat_uuids - flat_rids = stored.flat_rids - signal_idxs = stored.signal_idxs # Before the local-streaming dereference below: numel() needs the # tensors, and a new_token whose only other destination is a local # stream holds a single reference, which that dereference drops to zero # -- collecting the tensor this count is about to read. new_token_counts = self._count_new_tokens( route_output.new_token_output_idxs, - flat_rids, flat_uuids, signals, signal_idxs, + stored.flat_rids, stored.flat_uuids, signals, stored.signal_idxs, ) # Before the sends below: remote consumers read the walks off the @@ -3821,13 +3830,7 @@ def _stream_and_send( # streamed tensor. self.tensor_manager.dereference_batch_uniform(streamed) - send_rids = ( - [rid for rid in rids if rid not in failed_rids] if failed_rids - else list(rids) - ) - # Timed apart from the send: these comprehensions are per rid, and - # building them is the worker's cost, not the runtime's. - _t_prep = _time.perf_counter() if self._phase_period else 0.0 + send_rids = [rid for rid in rids if rid not in failed_rids] if failed_rids else list(rids) needs_info = route_output.rids_needing_request_info info_rids = ( send_rids if needs_info is None @@ -3858,7 +3861,7 @@ def _stream_and_send( # which is what an empty entry did. ntc_rids = [rid for rid in send_rids if new_token_counts.get(rid)] consumed = self.streams.consumption_batch(send_rids) - send_input = SendInput( + return SendInput( completion_id=route_output.completion_id, per_request_info=ParallelList( info_rids, @@ -3886,19 +3889,6 @@ def _stream_and_send( ), failed_rids=failed_rids, ) - _t_send = _time.perf_counter() if self._phase_period else 0.0 - if self._phase_period: - self._phase_record("worker.postprocess.send_prep", _t_send - _t_prep) - # Main loop only: the runtime holds `&mut self` across its GIL - # release, so a concurrent caller gets "Already mutably - # borrowed" rather than blocking. - completed_rids = self._graph_runtime.send_outputs(send_input) - for rid in completed_rids: - self.request_state.flush_publish_info(rid, batch_N.partition) - if self._phase_period: - self._phase_record( - "worker.postprocess.send", _time.perf_counter() - _t_send, - ) def _walk_signals(self, batch: ScheduledBatch, walk: str): signals = batch.walk_output_signals.get(walk) diff --git a/rust/src/bindings.rs b/rust/src/bindings.rs index 3dedb6b491..fae241fa1e 100644 --- a/rust/src/bindings.rs +++ b/rust/src/bindings.rs @@ -321,7 +321,7 @@ fn mstar_rust(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; - m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/rust/src/graph/request.rs b/rust/src/graph/request.rs index c4b4f62a89..b8c46a6a15 100644 --- a/rust/src/graph/request.rs +++ b/rust/src/graph/request.rs @@ -97,6 +97,11 @@ pub struct RequestInfo { /// persist signals and token counts to whichever request draws that /// integer next. Here, remove_request frees them with everything else. pub pending: PendingOutputs, + + /// Per (node, walk), this request's limits for its declared stop rules + /// (`GraphRuntime.stop_rules`): (max_iters, ignore_eos) per rule. Set by + /// `set_stop_limits` as the request arrives. + pub stop_limits: Vec<((Sym, Sym), Vec<(u32, bool)>)>, } /// What a request has accumulated since its last WORKER_GRAPHS_DONE. diff --git a/rust/src/graph/runtime.rs b/rust/src/graph/runtime.rs index 12a1aadb16..426755bb11 100644 --- a/rust/src/graph/runtime.rs +++ b/rust/src/graph/runtime.rs @@ -15,6 +15,9 @@ use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyValueError}; use pyo3::prelude::*; use rustc_hash::{FxHashMap, FxHashSet}; +mod postprocess; + +pub use postprocess::PostprocessOut; /// Python's `NestedLoopIndices`. #[pyclass] @@ -163,6 +166,8 @@ pub struct GraphRuntime { node_owner: FxHashMap<(Sym, Sym), WgIndex>, /// Loop stops from this iteration's check_stop; cleared every iteration. pending_loop_stops: FxHashSet<(u32, Sym, Sym)>, + /// (node, walk) -> its declared stop rules (`set_stop_rules`). + stop_rules: FxHashMap<(Sym, Sym), Vec>, /// Routing parked between complete_and_route_batch and send_outputs. completions: FxHashMap, completion_counter: u64, @@ -186,6 +191,9 @@ pub struct GraphRuntime { /// handed back to Python one at a time. None where the worker built the /// pyzmq communicator, which has no shareable object. communicator: Option>, + + /// A postprocess parked at a break (`begin_postprocess`). + postprocess: Option, } impl GraphRuntime { @@ -862,6 +870,21 @@ pub struct GraphRuntime { } } + /// `rid`'s iteration of loop `lp` in partition `p`, as + /// `get_dynamic_loop_iters` reports it (the last worker graph with the + /// loop wins); 0 when no live worker graph has it. + fn loop_count(&self, rid: u32, p: Sym, lp: Sym) -> u32 { + let Some(part) = self.info(rid).and_then(|i| i.partitions.get(&p)) else { return 0 }; + let mut count = 0; + for &wg in &part.walk_worker_graphs { + let Some(state) = self.state(wg, rid) else { continue }; + if let Some(&lid) = self.graphs[wg as usize].loop_by_name.get(&lp) { + count = state.loop_iter(lid); + } + } + count + } + /// Whether this request's current walk actually contains the loop. A stop /// for one it does not is a model bug: logged by the caller and dropped. fn check_dyn_loop(&self, rid: u32, partition: &str, loop_name: &str) -> bool { @@ -1287,25 +1310,13 @@ pub struct RouteOut { #[pyo3(get)] pub rids_needing_request_info: Vec, } -/// `route_and_send_inline`'s result; see there. -#[pyclass] -#[derive(Default)] -pub struct FusedRouteOut { - #[pyo3(get)] pub status: u8, - #[pyo3(get)] pub completion_id: u64, - #[pyo3(get)] pub sent_rids: Vec, - /// Wall time spent with the GIL released, for measuring the overlap. - #[pyo3(get)] pub nogil_ns: u64, - pub route: Option, -} - -#[pymethods] -impl FusedRouteOut { - /// The parked route, for a status-2 caller to finish with send_outputs. - /// Taken, not cloned: it is read once. - fn take_route(&mut self) -> Option { - self.route.take() - } +/// One of a walk's declared stop checks (Python's `StopPolicy`): stop `lp` +/// when its count plus one reaches the request's limit or, for an AR rule, +/// when the row's token is in `eos` and the request does not ignore EOS. +struct StopRule { + lp: Sym, + ar: bool, + eos: Vec, } /// Routing parked between complete_and_route_batch and send_outputs. @@ -1614,11 +1625,13 @@ impl GraphRuntime { walk_to_local_wgs, node_owner, pending_loop_stops: FxHashSet::default(), + stop_rules: FxHashMap::default(), completions: FxHashMap::default(), completion_counter: 0, staged_specs: FxHashMap::default(), spec_counter: 0, combined_walk_of, + postprocess: None, bookkeeping: bookkeeping.share(), communicator: communicator.map(|c| c.share()), }) @@ -3158,132 +3171,163 @@ impl GraphRuntime { Ok(out) } - /// `complete_and_route_batch` then `send_outputs`, in one crossing and - /// with the GIL released throughout, for a step whose every output stays - /// on the device and whose client copy rides inline (a device-loopback - /// decode step). Exactly the two calls the worker would make, with the - /// arguments its device-loopback branch builds: no tensors, and the - /// inline values read here from the stop check's host rows rather than - /// handed over as a list. - /// - /// The inline column is ``inline_rows`` scalars of ``inline_dtype`` at - /// ``inline_ptr`` with ``inline_stride`` elements between rows, row i - /// belonging to ``row_rids[i]``. The caller keeps the tensor alive and - /// landed for the call (it has waited on the step's completion event). + /// Declares (node, walk)'s stop rules, as `NodeSubmodule.get_stop_policies` + /// gives them, one column entry per rule. + fn set_stop_rules( + &mut self, node_name: &str, graph_walk: &str, + loops: Vec, ar: Vec, eos: Vec>, + ) -> PyResult<()> { + if loops.len() != ar.len() || loops.len() != eos.len() { + return Err(PyValueError::new_err("set_stop_rules: columns differ in length")); + } + let key = (self.interner.intern(node_name), self.interner.intern(graph_walk)); + let mut rules = Vec::with_capacity(loops.len()); + for ((name, ar), eos) in loops.iter().zip(ar).zip(eos) { + rules.push(StopRule { lp: self.interner.intern(name), ar, eos }); + } + self.stop_rules.insert(key, rules); + Ok(()) + } + + /// Each rid's limits for (node, walk)'s rules: `max_iters[k][i]` and + /// `ignore_eos[k][i]` are rule k's for `rids[i]`. Kept on the request, so + /// `remove_request` drops them with it. + fn set_stop_limits( + &mut self, node_name: &str, graph_walk: &str, rids: Vec, + max_iters: Vec>, ignore_eos: Vec>, + ) -> PyResult<()> { + if max_iters.len() != ignore_eos.len() + || max_iters.iter().any(|c| c.len() != rids.len()) + || ignore_eos.iter().any(|c| c.len() != rids.len()) + { + return Err(PyValueError::new_err("set_stop_limits: columns differ in length")); + } + let key = (self.interner.intern(node_name), self.interner.intern(graph_walk)); + for (i, &rid) in rids.iter().enumerate() { + let limits: Vec<(u32, bool)> = + max_iters.iter().zip(&ignore_eos).map(|(m, e)| (m[i], e[i])).collect(); + let Some(Some(info)) = self.requests.get_mut(rid as usize) else { continue }; + match info.stop_limits.iter_mut().find(|(k, _)| *k == key) { + Some(entry) => entry.1 = limits, + None => info.stop_limits.push((key, limits)), + } + } + Ok(()) + } + + /// (node, walk)'s declared stop rules over a step's rows, reading the + /// token column in place and each loop's count from the walk state. /// - /// Returns ``status``: - /// 0 -- nothing done (a rid has no row, or an unsupported dtype): take - /// the general path, exactly as ``_inline_emit_values`` falling - /// back does; - /// 1 -- routed and sent; ``sent_rids`` are send_outputs' return; - /// 2 -- routed, but the send needs Python-side payloads (a frame - /// carries ``per_request_info``, a tensor to register, new-token - /// counts, local streaming, freed inputs): ``route`` holds the - /// RouteOut and the caller finishes with ``send_outputs``. + /// At a step's stop check its loop counts are the ones the step ran at: + /// they move when the step is routed, which comes after. Returns + /// `(stop_rids, stop_rules, missing)`: a stop per (rid, rule index), and + /// the rids with no limits for this walk yet, which the caller registers + /// with `set_stop_limits` and checks again. A rid with no token row is + /// skipped, as `check_stop` skips a request with no output. + #[pyo3(signature = ( + partition, node_name, graph_walk, rids, + token_ptr, token_rows, token_stride, token_dtype, row_rids, + ))] + #[allow(clippy::too_many_arguments, clippy::type_complexity)] + fn check_stops_inline( + &self, + partition: &str, + node_name: &str, + graph_walk: &str, + rids: Vec, + token_ptr: usize, + token_rows: usize, + token_stride: usize, + token_dtype: &str, + row_rids: Vec, + ) -> PyResult<(Vec, Vec, Vec)> { + let tokens = postprocess::TokenColumn::new(token_ptr, token_rows, token_stride, token_dtype); + self.stops_over(partition, node_name, graph_walk, &rids, tokens.as_ref(), &row_rids) + } + + + /// Begin a step's postprocess (`PostprocessRequest`): check its stops, + /// route its outputs, send them, stopping where a phase needs Python + /// (`postprocess.rs`). `stop_rids[i]` stops + /// `stop_loops[i]` before any check (stream-terminated loops). The token + /// and inline columns are host views the caller keeps alive until the + /// step is done: `rows` integers of `dtype`, `stride` elements apart, row + /// i belonging to `row_rids[i]`; a null pointer is no column. #[pyo3(signature = ( - partition, graph_walk, node_name, output_signals, rids, + partition, node_name, graph_walk, output_signals, rids, walks, rid_walk_idx, + stop_rids, stop_loops, runtime_stops, + token_ptr, token_rows, token_stride, token_dtype, row_rids, inline_signal, inline_ptr, inline_rows, inline_stride, inline_dtype, - row_rids, + partition_done, python_send, ))] #[allow(clippy::too_many_arguments)] - fn route_and_send_inline( - &mut self, - py: Python<'_>, - partition: String, - graph_walk: String, - node_name: String, - output_signals: Vec, - rids: Vec, - inline_signal: String, - inline_ptr: usize, - inline_rows: usize, - inline_stride: usize, - inline_dtype: &str, + fn begin_postprocess( + &mut self, py: Python<'_>, + partition: String, node_name: String, graph_walk: String, output_signals: Vec, + rids: Vec, walks: Option>, rid_walk_idx: Option>, + stop_rids: Vec, stop_loops: Vec>, runtime_stops: bool, + token_ptr: usize, token_rows: usize, token_stride: usize, token_dtype: &str, row_rids: Vec, - ) -> PyResult { - let itemsize = match inline_dtype { - "int64" | "int32" => if inline_dtype == "int64" { 8 } else { 4 }, - _ => return Ok(FusedRouteOut::default()), - }; - if inline_ptr == 0 || row_rids.len() > inline_rows { - return Ok(FusedRouteOut::default()); - } - let t_nogil = std::time::Instant::now(); - let outcome = py.detach(|| -> PyResult { - // The inline values in `rids` order: a slice when the rows are in - // the batch's order (the usual case), else looked up per rid. - let read = |row: usize| -> i64 { - // SAFETY: row < inline_rows, checked by the callers below; - // the caller holds the buffer for the duration of the call. - unsafe { - let at = (inline_ptr as *const u8).add(row * inline_stride * itemsize); - if itemsize == 8 { - std::ptr::read_unaligned(at as *const i64) - } else { - i64::from(std::ptr::read_unaligned(at as *const i32)) - } - } - }; - let n = rids.len(); - let in_order = row_rids.len() >= n && row_rids[..n] == rids[..]; - let values: Vec = if in_order { - (0..n).map(read).collect() - } else { - let row_of: FxHashMap = - row_rids.iter().enumerate().map(|(i, &r)| (r, i)).collect(); - let mut v = Vec::with_capacity(n); - for r in &rids { - match row_of.get(r) { - Some(&i) => v.push(read(i)), - None => return Ok(FusedRouteOut::default()), - } - } - v - }; - let n_sig = output_signals.len(); - let route = self.complete_and_route_batch(RouteArg { - partition, - graph_walk, - node_name, - output_signals, - rids, - tensors: Vec::new(), - num_tensors: vec![0; n * n_sig], - inline_signal: Some(inline_signal), - inline_values: values, - logprobs: Vec::new(), - walks: None, - rid_walk_idx: None, - })?; - let completion_id = route.completion_id; - let needs_python = !route.register_uuids.is_empty() - || !route.new_token_output_idxs.is_empty() - || !route.local_streaming_signals.is_empty() - || !route.rids_needing_request_info.is_empty() - || !route.freed_input_uuids.is_empty(); - if needs_python { - return Ok(FusedRouteOut { - status: 2, completion_id, sent_rids: Vec::new(), nogil_ns: 0, - route: Some(route), - }); - } - // Nothing in the frames comes from Python: no rid carries - // per_request_info, so none completes a worker graph, and the - // count, consumption, profiling and publication payloads are all - // read only for those. - let (sent_rids, frames) = self.build_send_frames( - completion_id, vec![], vec![], vec![], vec![], - vec![], vec![], vec![], vec![], vec![], vec![], vec![], - )?; - self.dispatch_frames_nogil(frames)?; - Ok(FusedRouteOut { - status: 1, completion_id, sent_rids, nogil_ns: 0, route: None, - }) + inline_signal: Option, inline_ptr: usize, inline_rows: usize, + inline_stride: usize, inline_dtype: &str, + partition_done: Vec, python_send: bool, + ) -> PyResult { + let tokens = postprocess::TokenColumn::new(token_ptr, token_rows, token_stride, token_dtype); + let inline = inline_signal.and_then(|signal| { + postprocess::TokenColumn::new(inline_ptr, inline_rows, inline_stride, inline_dtype) + .map(|col| (signal, col)) }); - let nogil_ns = t_nogil.elapsed().as_nanos() as u64; - let mut outcome = outcome?; - outcome.nogil_ns = nogil_ns; - Ok(outcome) + let pp = postprocess::Postprocess::new( + partition, node_name, graph_walk, output_signals, rids, walks, rid_walk_idx, + stop_rids.into_iter().zip(stop_loops).collect(), runtime_stops, + tokens, row_rids, inline, partition_done, python_send, + ); + self.postprocess = None; + self.pp_run(py, pp) + } + + /// Resume a step parked for its stops (STOPS): the loops each rid stops, + /// and the rows to drop (their stop check failed). + fn resume_postprocess_stops( + &mut self, py: Python<'_>, rids: Vec, loops: Vec>, dropped: Vec, + ) -> PyResult { + self.pp_resume_stops(py, rids, loops, dropped) + } + + /// Resume a step parked for its stored outputs (STORE): `RouteTensors`. + #[pyo3(signature = (tensors, num_tensors, inline_signal, inline_values, logprobs = Vec::new()))] + fn resume_postprocess_store( + &mut self, py: Python<'_>, tensors: Vec, num_tensors: Vec, + inline_signal: Option, inline_values: Vec, logprobs: Vec>, + ) -> PyResult { + let stored = postprocess::Stored { + tensors, num_tensors, inline_signal, inline_values, logprobs, + }; + self.pp_resume_store(py, stored) + } + + /// Resume a step parked for its send's Python payloads (SEND), given as + /// `send_outputs` takes them. + #[pyo3(signature = ( + info_rids, request_infos, ntc_rids, new_token_counts, + consumed_rids, stream_tokens_consumed, prof_rids, profiling, + publish_rids, publications, skip_rids = Vec::new(), + ))] + #[allow(clippy::too_many_arguments)] + fn resume_postprocess_send( + &mut self, py: Python<'_>, + info_rids: Vec, request_infos: Vec>>, + ntc_rids: Vec, new_token_counts: Vec>, + consumed_rids: Vec, stream_tokens_consumed: Vec>, + prof_rids: Vec, profiling: Vec>>, + publish_rids: Vec, publications: Vec>>, + skip_rids: Vec, + ) -> PyResult { + let args = postprocess::SendArgs { + info_rids, request_infos, ntc_rids, new_token_counts, consumed_rids, + stream_tokens_consumed, prof_rids, profiling, publish_rids, publications, skip_rids, + }; + self.pp_resume_send(py, args) } /// Stop loops, record the pending stops, notify peers, and return the rids @@ -3298,65 +3342,9 @@ impl GraphRuntime { rids: Vec, loop_names: Vec>, ) -> PyResult> { - // One release for the whole batch; see send_outputs. // bookkeeping with the GIL held; the frames go out through send_frames - let mut frames: Vec<(String, Vec)> = Vec::new(); - let out = (|| -> PyResult> { - if rids.len() != loop_names.len() { - return Err(PyValueError::new_err( - "stop_loops_batched: rids and loop_names must be the same length", - )); - } - let Some(walk) = self.interner.get(graph_walk) else { - return Ok(Vec::new()); - }; - let mut stopped_rids = Vec::new(); - for (rid, names) in rids.into_iter().zip(loop_names) { - let wanted: Vec = names - .into_iter() - .filter(|n| self.check_dyn_loop(rid, partition, n)) - .collect(); - if wanted.is_empty() { - continue; - } - self.stop_loops_for_rid(rid, partition, &wanted, Some(last_node_run)); - for name in &wanted { - if let Some(l) = self.interner.get(name) { - self.pending_loop_stops.insert((rid, walk, l)); - } - } - let mut per_worker: FxHashMap> = FxHashMap::default(); - for name in &wanted { - for w in self.dyn_loop_workers(rid, partition, name) { - per_worker.entry(w).or_default().push(name.clone()); - } - } - // Never to ourselves: this rank originated the stop and has - // already applied it. A self-send would land in apply_peer_loop_stops - // and stop the loops a second time. - let me = self.shard.me; - let request_id = self.rid_name(rid)?; - // Snapshotted BEFORE the sends: stop_loops_for_rid above already - // recorded this stop, and every peer must see the same context. - let stop_times = self.loop_stop_times_of(rid); - for (worker, names) in per_worker { - if worker == me { - continue; - } - let bytes = frames::StopLoops { - request_id: &request_id, - partition_name: partition, - loop_names: &names, - loop_stop_times: stop_times.clone(), - } - .encode(&self.interner); - frames.push((self.interner.name(worker).to_string(), bytes)); - } - stopped_rids.push(rid); - } - Ok(stopped_rids) - })(); - let out = out?; + let (out, frames) = + self.stop_loops_frames(partition, graph_walk, last_node_run, rids, loop_names)?; self.send_frames(py, frames)?; Ok(out) } @@ -4019,6 +4007,7 @@ impl GraphRuntime { } } + #[cfg(test)] mod tests { use super::*; diff --git a/rust/src/graph/runtime/postprocess.rs b/rust/src/graph/runtime/postprocess.rs new file mode 100644 index 0000000000..4c9c8acc85 --- /dev/null +++ b/rust/src/graph/runtime/postprocess.rs @@ -0,0 +1,572 @@ +//! A step's postprocess -- its stops, its route, its send -- run here, +//! stopping for the worker where a phase needs Python. +//! +//! The GIL stays held where the separate calls held it (the stop check, a +//! stored step's route and send), with peer frames sent try-first +//! (`send_frames`). Only an inline step's route and send release it, as the +//! one call they replaced did: nothing in them touches Python. +//! +//! Python's `PhasedPostprocess` (`mstar/graph/runtime/base.py`) is the +//! reference: the phases below follow it, and +//! `test/rust/test_postprocess.py` drives both over the same steps. + +use super::*; + +/// What a step waits on (Python's `PostprocessNeed`). +pub(super) const NEED_DONE: u8 = 0; +pub(super) const NEED_STOPS: u8 = 1; +pub(super) const NEED_STORE: u8 = 2; +pub(super) const NEED_SEND: u8 = 3; + +/// A step's token column, read in place: `rows` integers of `itemsize` +/// bytes, `stride` elements apart, row i belonging to the caller's +/// `row_rids[i]`. The caller holds the tensor for the step. +pub(super) struct TokenColumn { + ptr: usize, + rows: usize, + stride: usize, + itemsize: usize, +} + +impl TokenColumn { + /// None for a null pointer or a dtype other than int32 / int64. + pub(super) fn new(ptr: usize, rows: usize, stride: usize, dtype: &str) -> Option { + let itemsize = match dtype { + "int64" => 8, + "int32" => 4, + _ => return None, + }; + (ptr != 0).then_some(TokenColumn { ptr, rows, stride, itemsize }) + } + + /// Row `row`'s value; the callers check `row < self.rows`. + fn read(&self, row: usize) -> i64 { + debug_assert!(row < self.rows); + // SAFETY: row < rows, which the callers check against the column the + // caller holds for the step. + unsafe { + let at = (self.ptr as *const u8).add(row * self.stride * self.itemsize); + if self.itemsize == 8 { + std::ptr::read_unaligned(at as *const i64) + } else { + i64::from(std::ptr::read_unaligned(at as *const i32)) + } + } + } + + /// The column in `rids` order; None when a rid has no row. + fn values_for(&self, rids: &[u32], row_rids: &[u32]) -> Option> { + if row_rids.len() > self.rows { + return None; + } + if row_rids.len() >= rids.len() && row_rids[..rids.len()] == rids[..] { + return Some((0..rids.len()).map(|i| self.read(i)).collect()); + } + let row_of: FxHashMap = + row_rids.iter().enumerate().map(|(i, &r)| (r, i)).collect(); + rids.iter().map(|r| row_of.get(r).map(|&i| self.read(i))).collect() + } +} + +#[derive(Clone, Copy, PartialEq)] +enum Phase { + Stops, + Route, +} + +/// The stored outputs a STORE break resumes with (`RouteTensors`). +pub(super) struct Stored { + pub tensors: Vec, + pub num_tensors: Vec, + pub inline_signal: Option, + pub inline_values: Vec, + pub logprobs: Vec>, +} + +/// A step's postprocess (`PostprocessRequest`), and how far it got. +pub(super) struct Postprocess { + pub partition: String, + pub node_name: String, + pub graph_walk: String, + pub output_signals: Vec, + pub rids: Vec, + pub walks: Option>, + pub rid_walk_idx: Option>, + /// stops known before the runtime's check: stream-terminated loops, and + /// without `runtime_stops` the model's too (all of the step's stops) + pub known_stops: Vec<(u32, Vec)>, + pub runtime_stops: bool, + pub tokens: Option, + pub row_rids: Vec, + pub inline: Option<(String, TokenColumn)>, + pub partition_done: Vec, + pub python_send: bool, + phase: Phase, + stops: Option)>>, + stopped: Vec, + stored: Option, + completion_id: u64, +} + +impl Postprocess { + #[allow(clippy::too_many_arguments)] + pub(super) fn new( + partition: String, node_name: String, graph_walk: String, output_signals: Vec, + rids: Vec, walks: Option>, rid_walk_idx: Option>, + known_stops: Vec<(u32, Vec)>, runtime_stops: bool, + tokens: Option, row_rids: Vec, inline: Option<(String, TokenColumn)>, + partition_done: Vec, python_send: bool, + ) -> Self { + Postprocess { + partition, node_name, graph_walk, output_signals, rids, walks, rid_walk_idx, + known_stops, runtime_stops, tokens, row_rids, inline, partition_done, python_send, + phase: Phase::Stops, stops: None, stopped: Vec::new(), stored: None, completion_id: 0, + } + } + + /// `_keep_rows`: the rows the step routes, less `dropped`. + fn keep_rows(&mut self, dropped: &FxHashSet) { + if dropped.is_empty() { + return; + } + let keep: Vec = (0..self.rids.len()).filter(|&i| !dropped.contains(&self.rids[i])).collect(); + self.rids = keep.iter().map(|&i| self.rids[i]).collect(); + if let Some(idx) = &self.rid_walk_idx { + self.rid_walk_idx = Some(keep.iter().map(|&i| idx[i]).collect()); + } + if let Some(stops) = &mut self.stops { + stops.retain(|(rid, _)| !dropped.contains(rid)); + } + // a dropped row's streams do not end here + self.partition_done.retain(|rid| !dropped.contains(rid)); + } + + /// Whether the send needs what only the worker has (`_route_needs_python`). + fn send_needs_python(&self, route: &RouteOut) -> bool { + self.python_send + || !self.stopped.is_empty() + || !route.register_uuids.is_empty() + || !route.new_token_output_idxs.is_empty() + || !route.local_streaming_signals.is_empty() + || !route.rids_needing_request_info.is_empty() + || !route.freed_input_uuids.is_empty() + } + + fn route_arg(&self, s: Stored) -> RouteArg { + RouteArg { + partition: self.partition.clone(), + graph_walk: self.graph_walk.clone(), + node_name: self.node_name.clone(), + output_signals: self.output_signals.clone(), + rids: self.rids.clone(), + tensors: s.tensors, + num_tensors: s.num_tensors, + inline_signal: s.inline_signal, + inline_values: s.inline_values, + logprobs: s.logprobs, + walks: if self.rid_walk_idx.is_some() { self.walks.clone() } else { None }, + rid_walk_idx: self.rid_walk_idx.clone(), + } + } + + fn done(&self, completed: Vec) -> PostprocessOut { + PostprocessOut { completed_rids: completed, ..self.out(NEED_DONE) } + } + + fn out(&self, need: u8) -> PostprocessOut { + PostprocessOut { + need, + rids: if need == NEED_DONE { Vec::new() } else { self.rids.clone() }, + rid_walk_idx: if need == NEED_DONE { None } else { self.rid_walk_idx.clone() }, + stopped_rids: self.stopped.clone(), + ..Default::default() + } + } +} + +/// `_merged_stops`: the loops each rid stops, in first-seen order. +fn merged_stops(parts: [&[(u32, Vec)]; 2]) -> Vec<(u32, Vec)> { + let mut out: Vec<(u32, Vec)> = Vec::new(); + let mut at: FxHashMap = FxHashMap::default(); + for part in parts { + for (rid, loops) in part { + let i = *at.entry(*rid).or_insert_with(|| { + out.push((*rid, Vec::new())); + out.len() - 1 + }); + for l in loops { + if !out[i].1.contains(l) { + out[i].1.push(l.clone()); + } + } + } + } + out +} + +/// `PostprocessStep`, as columns. +#[pyclass] +#[derive(Default)] +pub struct PostprocessOut { + #[pyo3(get)] pub need: u8, + /// The rows the step routes, less the dropped ones (STORE, SEND). + #[pyo3(get)] pub rids: Vec, + #[pyo3(get)] pub rid_walk_idx: Option>, + #[pyo3(get)] pub stopped_rids: Vec, + #[pyo3(get)] pub completed_rids: Vec, + /// Wall time this crossing ran with the GIL released (an inline route + /// and send). + #[pyo3(get)] pub nogil_ns: u64, + pub route: Option, +} + +#[pymethods] +impl PostprocessOut { + /// The route a SEND break's payloads finish; taken, not cloned. + fn take_route(&mut self) -> Option { + self.route.take() + } +} + +impl GraphRuntime { + pub(super) fn sym_or_none(&self, s: &str) -> Sym { + // a name no graph has matches nothing + self.interner.get(s).unwrap_or(Sym::MAX) + } + + /// `check_stops_inline`'s check: `(stop_rids, stop_rules, missing)`. + #[allow(clippy::type_complexity)] + pub(super) fn stops_over( + &self, partition: &str, node_name: &str, graph_walk: &str, rids: &[u32], + tokens: Option<&TokenColumn>, row_rids: &[u32], + ) -> PyResult<(Vec, Vec, Vec)> { + let (Some(node), Some(walk)) = (self.interner.get(node_name), self.interner.get(graph_walk)) + else { + return Err(PyKeyError::new_err("check_stops: no rules for this (node, walk)")); + }; + let Some(rules) = self.stop_rules.get(&(node, walk)) else { + return Err(PyKeyError::new_err("check_stops: no rules for this (node, walk)")); + }; + let (mut stop_rids, mut stop_rules, mut missing) = (Vec::new(), Vec::new(), Vec::new()); + if rules.is_empty() { + return Ok((stop_rids, stop_rules, missing)); + } + let tokens = if rules.iter().any(|r| r.ar) { + match tokens { + Some(t) if row_rids.len() <= t.rows => Some(t), + _ => return Err(PyValueError::new_err( + "check_stops: no int32 / int64 token column covering the rows", + )), + } + } else { + None + }; + let in_order = row_rids.len() >= rids.len() && row_rids[..rids.len()] == rids[..]; + let row_of: FxHashMap = if in_order || tokens.is_none() { + FxHashMap::default() + } else { + row_rids.iter().enumerate().map(|(i, &r)| (r, i)).collect() + }; + let p = self.interner.get(partition); + for (i, &rid) in rids.iter().enumerate() { + let token = match tokens { + None => 0, + Some(t) if in_order => t.read(i), + Some(t) => match row_of.get(&rid) { + Some(&row) => t.read(row), + None => continue, + }, + }; + let Some(info) = self.info(rid) else { continue }; + let Some((_, limits)) = info.stop_limits.iter().find(|(k, _)| *k == (node, walk)) else { + missing.push(rid); + continue; + }; + for (k, (rule, &(max_iters, ignore_eos))) in rules.iter().zip(limits).enumerate() { + let count = p.map_or(0, |p| self.loop_count(rid, p, rule.lp)); + if count + 1 >= max_iters || (rule.ar && !ignore_eos && rule.eos.contains(&token)) { + stop_rids.push(rid); + stop_rules.push(k as u32); + } + } + } + Ok((stop_rids, stop_rules, missing)) + } + + /// `stop_loops_batched` without the send: the rids whose loops stopped, + /// and the STOP_LOOPS frames for their peers. + #[allow(clippy::type_complexity)] + pub(super) fn stop_loops_frames( + &mut self, partition: &str, graph_walk: &str, last_node_run: &str, + rids: Vec, loop_names: Vec>, + ) -> PyResult<(Vec, Vec<(String, Vec)>)> { + if rids.len() != loop_names.len() { + return Err(PyValueError::new_err( + "stop_loops_batched: rids and loop_names must be the same length", + )); + } + let mut frames: Vec<(String, Vec)> = Vec::new(); + let Some(walk) = self.interner.get(graph_walk) else { + return Ok((Vec::new(), frames)); + }; + let mut stopped_rids = Vec::new(); + for (rid, names) in rids.into_iter().zip(loop_names) { + let wanted: Vec = names + .into_iter() + .filter(|n| self.check_dyn_loop(rid, partition, n)) + .collect(); + if wanted.is_empty() { + continue; + } + self.stop_loops_for_rid(rid, partition, &wanted, Some(last_node_run)); + for name in &wanted { + if let Some(l) = self.interner.get(name) { + self.pending_loop_stops.insert((rid, walk, l)); + } + } + let mut per_worker: FxHashMap> = FxHashMap::default(); + for name in &wanted { + for w in self.dyn_loop_workers(rid, partition, name) { + per_worker.entry(w).or_default().push(name.clone()); + } + } + // Never to ourselves: this rank originated the stop and has + // already applied it. A self-send would land in apply_peer_loop_stops + // and stop the loops a second time. + let me = self.shard.me; + let request_id = self.rid_name(rid)?; + // Snapshotted BEFORE the sends: stop_loops_for_rid above already + // recorded this stop, and every peer must see the same context. + let stop_times = self.loop_stop_times_of(rid); + for (worker, names) in per_worker { + if worker == me { + continue; + } + let bytes = frames::StopLoops { + request_id: &request_id, + partition_name: partition, + loop_names: &names, + loop_stop_times: stop_times.clone(), + } + .encode(&self.interner); + frames.push((self.interner.name(worker).to_string(), bytes)); + } + stopped_rids.push(rid); + } + Ok((stopped_rids, frames)) + } + + /// The loop name of each of (node, walk)'s stop rules, by index. + fn stop_rule_loops(&self, node_name: &str, graph_walk: &str) -> Vec { + let key = (self.sym_or_none(node_name), self.sym_or_none(graph_walk)); + self.stop_rules.get(&key).map_or_else(Vec::new, |rules| { + rules.iter().map(|r| self.interner.name(r.lp).to_string()).collect() + }) + } + + /// `_stop_step_loops`: stop the step's loops, per walk, and send the + /// peers' STOP_LOOPS; the rids whose loops stopped. + fn pp_stop_loops(&mut self, py: Python<'_>, pp: &Postprocess) -> PyResult> { + let stops = pp.stops.as_deref().unwrap_or(&[]); + if stops.is_empty() { + return Ok(Vec::new()); + } + let walk_of: FxHashMap = match (&pp.rid_walk_idx, &pp.walks) { + (Some(idx), Some(walks)) => { + pp.rids.iter().zip(idx).map(|(&r, &w)| (r, walks[w as usize].as_str())).collect() + } + _ => pp.rids.iter().map(|&r| (r, pp.graph_walk.as_str())).collect(), + }; + let mut by_walk: Vec<(&str, Vec, Vec>)> = Vec::new(); + for (rid, loops) in stops { + let Some(&walk) = walk_of.get(rid) else { continue }; + match by_walk.iter_mut().find(|(w, _, _)| *w == walk) { + Some((_, rids, names)) => { + rids.push(*rid); + names.push(loops.clone()); + } + None => by_walk.push((walk, vec![*rid], vec![loops.clone()])), + } + } + let mut stopped = Vec::new(); + for (walk, rids, names) in by_walk { + let walk = walk.to_string(); + let (rids, frames) = self.stop_loops_frames(&pp.partition, &walk, &pp.node_name, rids, names)?; + self.send_frames(py, frames)?; + stopped.extend(rids); + } + Ok(stopped) + } + + /// `_run_postprocess`, from wherever `pp` stands: DONE, or a break with + /// the step parked in `self.postprocess`. + pub(super) fn pp_run(&mut self, py: Python<'_>, mut pp: Postprocess) -> PyResult { + if pp.phase == Phase::Stops { + if pp.stops.is_none() && !pp.runtime_stops { + // checked by the worker: the request's stops are all of them + pp.stops = Some(merged_stops([&pp.known_stops, &[]])); + } + if pp.stops.is_none() && pp.runtime_stops { + // a check that fails, or misses a row's limits, leaves the + // decision to check_stop + let checked = self.stops_over( + &pp.partition, &pp.node_name, &pp.graph_walk, &pp.rids, + pp.tokens.as_ref(), &pp.row_rids, + ); + if let Ok((stop_rids, stop_rules, missing)) = checked { + if missing.is_empty() { + let loops = self.stop_rule_loops(&pp.node_name, &pp.graph_walk); + let found: Vec<(u32, Vec)> = stop_rids.iter().zip(&stop_rules) + .map(|(&r, &k)| (r, vec![loops[k as usize].clone()])) + .collect(); + pp.stops = Some(merged_stops([&pp.known_stops, &found])); + } + } + } + if pp.stops.is_none() { + return Ok(self.pp_park(pp, NEED_STOPS)); + } + pp.stopped = self.pp_stop_loops(py, &pp)?; + pp.phase = Phase::Route; + } + if pp.rids.is_empty() { + return Ok(pp.out(NEED_DONE)); // every row was dropped + } + if let Some(s) = pp.stored.take() { + let route = self.complete_and_route_batch(pp.route_arg(s))?; + pp.completion_id = route.completion_id; + if pp.send_needs_python(&route) { + return Ok(self.pp_park_route(pp, route)); + } + let (completed, frames) = self.pp_send_frames(&pp, SendArgs::default())?; + self.send_frames(py, frames)?; + return Ok(pp.done(completed)); + } + let values = pp.inline.as_ref().and_then(|(_, col)| col.values_for(&pp.rids, &pp.row_rids)); + let Some(values) = values else { + return Ok(self.pp_park(pp, NEED_STORE)); + }; + // Inline: nothing stored, so nothing here touches Python. + let n = pp.rids.len() * pp.output_signals.len(); + let signal = pp.inline.as_ref().map(|(s, _)| s.clone()); + let t = std::time::Instant::now(); + let sent = py.detach(|| -> PyResult, RouteOut>> { + let route = self.complete_and_route_batch(pp.route_arg(Stored { + tensors: Vec::new(), num_tensors: vec![0; n], inline_signal: signal, + inline_values: values, logprobs: Vec::new(), + }))?; + pp.completion_id = route.completion_id; + if pp.send_needs_python(&route) { + return Ok(Err(route)); + } + // Nothing in the frames comes from Python: no rid carries + // per_request_info, so none completes a worker graph, and the + // count, consumption, profiling and publication payloads are all + // read only for those. + let (completed, frames) = self.pp_send_frames(&pp, SendArgs::default())?; + self.dispatch_frames_nogil(frames)?; + Ok(Ok(completed)) + })?; + let nogil_ns = t.elapsed().as_nanos() as u64; + let mut out = match sent { + Ok(completed) => pp.done(completed), + Err(route) => self.pp_park_route(pp, route), + }; + out.nogil_ns = nogil_ns; + Ok(out) + } + + fn pp_park(&mut self, pp: Postprocess, need: u8) -> PostprocessOut { + let out = pp.out(need); + self.postprocess = Some(pp); + out + } + + /// Park a routed step for its send's Python payloads (SEND). + fn pp_park_route(&mut self, pp: Postprocess, route: RouteOut) -> PostprocessOut { + let mut out = self.pp_park(pp, NEED_SEND); + out.route = Some(route); + out + } + + /// The send: the stream partitions this pass reports done, then its + /// frames; the rids a WORKER_GRAPHS_DONE goes out for. + #[allow(clippy::type_complexity)] + fn pp_send_frames( + &mut self, pp: &Postprocess, args: SendArgs, + ) -> PyResult<(Vec, Vec<(String, Vec)>)> { + for &rid in &pp.partition_done { + self.mark_partition_done(rid, &pp.partition); + } + self.build_send_frames( + pp.completion_id, args.info_rids, args.request_infos, args.ntc_rids, + args.new_token_counts, args.consumed_rids, args.stream_tokens_consumed, + args.prof_rids, args.profiling, args.publish_rids, args.publications, args.skip_rids, + ) + } + + /// The parked step, for a resume. + pub(super) fn pp_take(&mut self) -> PyResult { + self.postprocess.take().ok_or_else(|| { + PyRuntimeError::new_err("resume_postprocess: no step is waiting") + }) + } + + pub(super) fn pp_resume_stops( + &mut self, py: Python<'_>, rids: Vec, loops: Vec>, dropped: Vec, + ) -> PyResult { + let mut pp = self.pp_take()?; + let given: Vec<(u32, Vec)> = rids.into_iter().zip(loops).collect(); + pp.stops = Some(merged_stops([&pp.known_stops, &given])); + pp.keep_rows(&dropped.into_iter().collect()); + self.pp_run(py, pp) + } + + pub(super) fn pp_resume_store(&mut self, py: Python<'_>, stored: Stored) -> PyResult { + let mut pp = self.pp_take()?; + pp.stored = Some(stored); + self.pp_run(py, pp) + } + + pub(super) fn pp_resume_send(&mut self, py: Python<'_>, args: SendArgs) -> PyResult { + let pp = self.pp_take()?; + let (completed, frames) = self.pp_send_frames(&pp, args)?; + self.send_frames(py, frames)?; + Ok(pp.done(completed)) + } +} + +/// `send_outputs`' Python payloads, as columns. +#[derive(Default)] +pub(super) struct SendArgs { + pub info_rids: Vec, + pub request_infos: Vec>>, + pub ntc_rids: Vec, + pub new_token_counts: Vec>, + pub consumed_rids: Vec, + pub stream_tokens_consumed: Vec>, + pub prof_rids: Vec, + pub profiling: Vec>>, + pub publish_rids: Vec, + pub publications: Vec>>, + /// rids that failed after routing: none of their frames go out + pub skip_rids: Vec, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn keep_rows_drops_a_rows_stops_and_partition_done() { + let mut pp = Postprocess::new( + "p".into(), "n".into(), "w".into(), vec!["s".into()], vec![1, 2, 3], None, None, + vec![(2, vec!["l".into()])], false, None, Vec::new(), None, vec![1, 2], false, + ); + pp.stops = Some(vec![(2, vec!["l".into()]), (3, vec!["l".into()])]); + pp.keep_rows(&[2].into_iter().collect()); + assert_eq!(pp.rids, vec![1, 3]); + assert_eq!(pp.stops, Some(vec![(3, vec!["l".to_string()])])); + assert_eq!(pp.partition_done, vec![1]); + } +} diff --git a/test/modular/test_dropped_row_stream_stop.py b/test/modular/test_dropped_row_stream_stop.py index 97326005ea..f817929005 100644 --- a/test/modular/test_dropped_row_stream_stop.py +++ b/test/modular/test_dropped_row_stream_stop.py @@ -25,9 +25,11 @@ torch = pytest.importorskip("torch") +from mstar.communication.tensors import StoredOutputs # noqa: E402 from mstar.engine.engine import Engine, ExecutingBatch # noqa: E402 from mstar.engine.resources.runner import StepRunner # noqa: E402 from mstar.engine.resources.step import StepContext # noqa: E402 +from mstar.graph.runtime.base import PhasedPostprocess, RouteOutput # noqa: E402 from mstar.model.submodule_base import END_PARTITION, BatchedModelOutput # noqa: E402 from mstar.worker.micro_scheduler import ScheduledBatch # noqa: E402 from mstar.worker.worker import PendingBatch, Worker # noqa: E402 @@ -67,13 +69,16 @@ def check_stop(self, rid, info, outputs): return set(self.stops.get(rid, ())) -class _Runtime: +class _Runtime(PhasedPostprocess): """Answers ``stop_loops_batched`` like the real runtime for a running dynamic loop: it does not know a rid failed until the REMOVE arrives.""" + supports_inline_emit = False + def __init__(self, pending_stops=()): self.stopped: list[int] = [] self.loop_names: dict[int, set[str]] = {} + self.partition_done: list[int] = [] self.pending_stops = set(pending_stops) def pending_loop_stop_rids(self, graph_walk, loop_name): @@ -85,6 +90,21 @@ def cleanup_consumed_inputs(self, node, rids, wgs): def clear_pending_loop_stops(self): pass + def mark_partition_done(self, rid, partition): + self.partition_done.append(rid) + + def complete_and_route_batch(self, route_input): + return RouteOutput( + completion_id=1, register_uuids=[], register_rids=[], + new_token_output_idxs=[], local_streaming_by_signal={}, rids_needing_request_info=set(), + ) + + def send_outputs(self, send_input): + return [] + + def update_dynamic_loop_iters(self, per_request_info, partition): + pass + def stop_loops_batched(self, partition, graph_walk, last_node_run, loop_names): rids = [rid for rid, _ in loop_names] self.stopped.extend(rids) @@ -116,12 +136,17 @@ def _batch(rids, streams=True): ) -def _worker(eng, runtime, in_loop=True): - def _store(*a, **k): - raise _ReachedRoutingError +def _worker(eng, runtime, in_loop=True, through_send=False): + def _store(rids, outputs, signals, **k): + if not through_send: + raise _ReachedRoutingError + return StoredOutputs([], [], [], [0] * (len(rids) * len(signals))) w = Worker.__new__(Worker) w._phase_period = 0 + w._stop_token_signal = {} # no declared stop policies: check_stop decides + w._walk_drivers = {} # no producer-driven walk transitions here + w._walk_driven_edges = {} w.enable_nvtx = False w.enable_prof = False w.device = torch.device("cpu") @@ -130,16 +155,23 @@ def _store(*a, **k): w.tensor_manager = SimpleNamespace( cleanup_collectable=lambda *a: None, store_and_return_tensor_info_batch=_store, + drop_host_copies=lambda: None, + dereference_batch_uniform=lambda uuids: None, ) + w.streams = SimpleNamespace(consumption_batch=lambda rids: []) w.engine_manager = SimpleNamespace(get_engine=lambda n: eng) w._innermost_loop_by_node = {(WALK, NODE): LOOP} if in_loop else {} w.reported: dict[int, str] = {} w._fail_requests = w.reported.update - w.request_state = SimpleNamespace(buffer_publish_info=lambda *a: None) + w.request_state = SimpleNamespace( + buffer_publish_info=lambda *a: None, flush_publish_info=lambda *a: None, + get_fwd_info=lambda rid, partition: None, + get_pending_publish_info=lambda rid, partition: {}, + ) return w -def _postprocess(sub, rids=(1, 2), pending_stops=(), streams=True, in_loop=True): +def _postprocess(sub, rids=(1, 2), pending_stops=(), streams=True, in_loop=True, through_send=False): """What ``Worker.run`` does with a finished step: drop and report the rows a per-rid stage failed, then postprocess the rest. ``pending_stops`` makes the step a speculative new loop iteration whose rows in it already @@ -150,7 +182,7 @@ def _postprocess(sub, rids=(1, 2), pending_stops=(), streams=True, in_loop=True) outputs = BatchedModelOutput(per_rid_outputs={ r: {"audio": [torch.zeros(1)]} for r in batch.request_ids }) - w = _worker(eng, runtime, in_loop) + w = _worker(eng, runtime, in_loop, through_send) pending = PendingBatch( batch=ScheduledBatch( node_name=NODE, graph_walk=WALK, @@ -164,8 +196,11 @@ def _postprocess(sub, rids=(1, 2), pending_stops=(), streams=True, in_loop=True) failed = dict(batch.failed_requests) w._drop_failed_rids(pending, outputs, failed) w._fail_requests(failed) - with pytest.raises(_ReachedRoutingError): + if through_send: w._postprocess_batch(pending, outputs) + else: + with pytest.raises(_ReachedRoutingError): + w._postprocess_batch(pending, outputs) return batch, runtime, w @@ -218,3 +253,11 @@ def test_an_end_partition_stop_inside_a_loop_fails_the_request(): assert set(w.reported) == {1} assert batch.request_ids == [2] and batch.partition_done_rids == set() assert runtime.stopped == [] + + +def test_a_failed_check_stop_row_does_not_report_its_partition_done(): + """The send reports the step's partitions done; not a row the stop check + failed, which the STOPS break drops after the request was built.""" + _, runtime, w = _postprocess(_Sub(bad={2}, stage="check_stop"), through_send=True) + assert set(w.reported) == {2} + assert runtime.partition_done == [1] diff --git a/test/modular/test_ingest_rollback.py b/test/modular/test_ingest_rollback.py index 459ed91de8..44d8044249 100644 --- a/test/modular/test_ingest_rollback.py +++ b/test/modular/test_ingest_rollback.py @@ -50,6 +50,7 @@ def test_the_worker_fails_only_the_request_whose_ingest_raised(): processed, failed, registered = [], [], [] w._graph_runtime = SimpleNamespace( add_request=lambda **kw: 5, get_sharding_config=lambda rid: None, + register_stop_limits=lambda rid, info: None, ) w.engine_manager = SimpleNamespace( engine=SimpleNamespace(publisher_nodes={}, skip_publish=lambda rid, keys: None), diff --git a/test/modular/test_inline_tokens.py b/test/modular/test_inline_tokens.py index 8c6c21eb7a..e0f3ef383e 100644 --- a/test/modular/test_inline_tokens.py +++ b/test/modular/test_inline_tokens.py @@ -349,9 +349,24 @@ def test_row_logprobs_follow_the_batch_order(): assert Worker._row_logprobs(HostRows((7,), {SAMPLED_LOGPROB: torch.tensor([nan])}), [7]) == [] -def test_a_step_with_logprobs_skips_the_fused_route(): - # the fused route builds no route input, which is where log-probs ride - worker = _worker(SimpleNamespace(supports_route_and_send_inline=True)) - rows = HostRows((7,), {"new_token": torch.tensor([70]), SAMPLED_LOGPROB: torch.tensor([-0.5])}) - batch_N = SimpleNamespace(batch=SimpleNamespace(output_signals=["text_inputs"])) - assert worker._route_and_send_fast(batch_N, rows, None, [7], set()) is False +def test_a_step_with_logprobs_routes_through_the_store(): + # the in-runtime route builds no log-probs: such a step stores, and STORE + # passes them on with the route tensors + w = _worker() + w._stop_token_signal, w._device_loopback, w._innermost_loop_by_node = {}, {}, {} + w.enable_prof = False + engine = SimpleNamespace( + device_loopback_signals=lambda node, walk: ("text_inputs",), + inline_client_signals=lambda node, walk: {"text_inputs": "new_token"}, + ) + batch = _Batch() + batch.partition = "p" + batch.batch.output_signals = ["text_inputs"] + batch.batch.request_to_worker_graph = {7: 0} + batch.node_batch = SimpleNamespace(final_stream_rids=(), partition_done_rids=()) + batch.batch.rids_by_walk = lambda: {"decode": [7]} + tokens = torch.tensor([70]) + plain = w._get_postprocess_request(batch, HostRows((7,), {"new_token": tokens}), engine, False, None, None) + assert plain.inline is not None and plain.inline.signal == "text_inputs" + rows = HostRows((7,), {"new_token": tokens, SAMPLED_LOGPROB: torch.tensor([-0.5])}) + assert w._get_postprocess_request(batch, rows, engine, False, None, None).inline is None diff --git a/test/modular/test_output_spans.py b/test/modular/test_output_spans.py index af7f5b5c88..8ea569ec0e 100644 --- a/test/modular/test_output_spans.py +++ b/test/modular/test_output_spans.py @@ -27,6 +27,7 @@ SubmoduleStep, ) from mstar.graph.base import GraphEdge, GraphNode, Loop # noqa: E402 +from mstar.graph.runtime.base import PhasedPostprocess # noqa: E402 from mstar.graph.special_destinations import EMIT_TO_CLIENT # noqa: E402 from mstar.model.submodule_base import BatchedModelOutput, NodeInputs, NodeSubmodule # noqa: E402 from mstar.worker.engine_manager import _refuse_speculated_spans # noqa: E402 @@ -141,6 +142,10 @@ def _batch(samples: dict[int, torch.Tensor]) -> ExecutingBatch: ) +class _Runtime(SimpleNamespace, PhasedPostprocess): + """Stubbed phase methods under the reference begin / resume.""" + + def _worker(engine: Engine) -> tuple[Worker, dict]: stored: dict = {} @@ -153,7 +158,7 @@ def store(rids, outputs, signals, **kwargs): worker.enable_nvtx = False worker.enable_prof = False worker.device = torch.device("cpu") - worker._graph_runtime = SimpleNamespace( + worker._graph_runtime = _Runtime( cleanup_consumed_inputs=lambda node, rids, wgs: ([], []), clear_pending_loop_stops=lambda: None, stop_loops_batched=lambda **kwargs: [], @@ -165,6 +170,7 @@ def store(rids, outputs, signals, **kwargs): ) worker.engine_manager = SimpleNamespace(get_engine=lambda node: engine) worker._innermost_loop_by_node = {} + worker._stop_token_signal = {} # no declared stop policies: check_stop decides worker.request_state = SimpleNamespace(buffer_publish_info=lambda *args: None) worker.scheduler = SimpleNamespace(failed_rids=set()) return worker, stored diff --git a/test/modular/test_postprocess_request.py b/test/modular/test_postprocess_request.py new file mode 100644 index 0000000000..803b13cb3a --- /dev/null +++ b/test/modular/test_postprocess_request.py @@ -0,0 +1,56 @@ +"""``Worker._get_postprocess_request`` routes each walk under its own output +signals.""" +from types import SimpleNamespace + +from mstar.worker.micro_scheduler import ScheduledBatch +from mstar.worker.worker import Worker + +NODE, MIXED = "LLM", "LLM_MIXED" +SIGNALS = {"decode": ["text_inputs"], "prefill_text": ["text_inputs", "new_token"]} + + +def _worker() -> Worker: + w = Worker.__new__(Worker) + w._graph_runtime = SimpleNamespace( + get_output_signals=lambda node, walk: SIGNALS[walk], supports_inline_emit=False, + ) + w._stop_token_signal = {} + w._innermost_loop_by_node = {} + w._device_loopback = {} + w.enable_prof = False + return w + + +def _pending(request_walks: dict[int, str]): + batch = ScheduledBatch( + node_name=NODE, graph_walk=MIXED, + request_to_worker_graph=dict.fromkeys(request_walks, 0), + # a merged batch keeps its first walk's names + output_signals=SIGNALS["decode"], request_walks=dict(request_walks), + ) + return SimpleNamespace( + batch=batch, node_name=NODE, graph_walk=MIXED, partition="p", + node_batch=SimpleNamespace(final_stream_rids=set(), partition_done_rids=set()), + ) + + +def test_a_combined_batch_down_to_one_walk_routes_under_that_walks_signals(): + # the decode row left the batch (failed in prepare_inputs, say) + req = _worker()._get_postprocess_request(_pending({2: "prefill_text"}), None, None, False, None, None) + assert (req.graph_walk, req.walks) == ("prefill_text", None) + assert req.output_signals == SIGNALS["prefill_text"] + + +def test_a_combined_batch_of_several_walks_routes_under_their_union(): + req = _worker()._get_postprocess_request( + _pending({1: "decode", 2: "prefill_text"}), None, None, False, None, None, + ) + assert req.walks == ["decode", "prefill_text"] and req.rid_walk_idx == [0, 1] + assert req.output_signals == ["text_inputs", "new_token"] + + +def test_the_models_stops_ride_the_request_when_the_runtime_does_not_check(): + req = _worker()._get_postprocess_request( + _pending({2: "prefill_text"}), None, None, False, None, {2: {"decode_loop"}}, + ) + assert not req.runtime_stops and req.stops == {2: {"decode_loop"}} diff --git a/test/modular/test_rid_handles.py b/test/modular/test_rid_handles.py index 379069c3bc..136c99c3d3 100644 --- a/test/modular/test_rid_handles.py +++ b/test/modular/test_rid_handles.py @@ -137,6 +137,7 @@ def _runtime_remove(h): w._graph_runtime = SimpleNamespace( # The runtime mints the handle now; interning is still the whole of it. add_request=lambda request_id, **kw: t.intern(request_id), + register_stop_limits=lambda rid, info: None, get_rid_handle=t.handle, get_rid_string=t.name, get_sharding_config=lambda h: SimpleNamespace(groups=[]), remove_request=_runtime_remove, diff --git a/test/modular/test_send_registration_failure.py b/test/modular/test_send_registration_failure.py index 6f5612bf28..c9a094bb4a 100644 --- a/test/modular/test_send_registration_failure.py +++ b/test/modular/test_send_registration_failure.py @@ -27,7 +27,7 @@ TensorCommunicationManager, ) from mstar.graph.base import GraphEdge, GraphNode, TensorPointerInfo -from mstar.graph.runtime.base import FreedTensors, RouteOutput, SendInput +from mstar.graph.runtime.base import FreedTensors, PhasedPostprocess, RouteOutput, SendInput from mstar.graph.special_destinations import EMIT_TO_CLIENT from mstar.model.submodule_base import BatchedModelOutput from mstar.utils.containers import ParallelList @@ -51,6 +51,10 @@ def poll_for_messages(self, timeout_ms=20): return False +class _StubRuntime(SimpleNamespace, PhasedPostprocess): + """Stubbed phase methods under the reference begin / resume.""" + + class _Engine: """Registers like a transfer engine, but refuses the given pointers the way Mooncake refuses an overlapping region (-7).""" @@ -128,6 +132,7 @@ def test_the_worker_fails_only_the_refused_request(): worker = Worker.__new__(Worker) worker._walk_drivers = {} worker._walk_driven_edges = {} + worker._stop_token_signal = {} # no declared stop policies: check_stop decides worker.streams = StreamBufferManager(SimpleNamespace(per_request_info={}), None, []) # no streams worker.tensor_manager = mgr worker.scheduler = SimpleNamespace(failed_rids=set()) @@ -180,6 +185,7 @@ def _postprocess_worker(refuse, follower=False): w = Worker.__new__(Worker) w._walk_drivers = {} w._walk_driven_edges = {} + w._stop_token_signal = {} # no declared stop policies: check_stop decides w.streams = StreamBufferManager(SimpleNamespace(per_request_info={}), None, []) # no streams w._phase_period = 0 w.enable_nvtx = w.enable_prof = False @@ -216,6 +222,7 @@ def check_stop(node_batch, cpu_outputs, host_rows=None): materialize_publish=lambda *a: None, # like the engine's, reads each stopped rid's info finalize_stopped_requests=lambda nb, rids: [nb.per_request_info[r] for r in rids], + step_finished=lambda batch, outputs, **kwargs: None, ) w.engine_manager = SimpleNamespace(get_engine=lambda node: engine) @@ -255,7 +262,7 @@ def send(send_input): ) return [] - w._graph_runtime = SimpleNamespace( + w._graph_runtime = _StubRuntime( cleanup_consumed_inputs=lambda *a: FreedTensors.none(), clear_pending_loop_stops=lambda: None, stop_loops_batched=lambda **kw: list(kw["loop_names"].keys), @@ -468,6 +475,7 @@ def test_remove_request_finishes_when_an_unregister_fails(monkeypatch, rust): w = Worker.__new__(Worker) w._walk_drivers = {} w._walk_driven_edges = {} + w._stop_token_signal = {} # no declared stop policies: check_stop decides w.streams = StreamBufferManager(SimpleNamespace(per_request_info={}), None, []) # no streams w.is_tp_follower = False w._rid = lambda wire: 3 diff --git a/test/modular/test_stop_policies.py b/test/modular/test_stop_policies.py new file mode 100644 index 0000000000..53cc887bae --- /dev/null +++ b/test/modular/test_stop_policies.py @@ -0,0 +1,68 @@ +"""A submodule's declared stop policies agree with its own ``check_stop``. + +The graph runtime runs ``get_stop_policies`` in place of ``check_stop`` when +a walk declares them, so for every model that declares any, each row a +policy stops (``StopPolicy.stops``) must be one ``check_stop`` stops, on +every walk, over random tokens, iteration counts, limits and ``ignore_eos`` +settings. +""" +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from mstar.model.orpheus.submodules import OrpheusLLMSubmodule +from mstar.model.qwen3_5.submodules import LLMSubmodule, token_stop + +LOOP = "decode_loop" +EOS = (7, 9) + +MODELS = { + # (submodule, its config, its walks, its TokenStop if it declares one) + "qwen3_5": (LLMSubmodule, SimpleNamespace(stop_token_ids=frozenset(EOS)), + ("decode", "prefill_text", "prefill_vision"), token_stop(EOS)), + "orpheus": (OrpheusLLMSubmodule, SimpleNamespace(stop_token_id=EOS[0]), ("prefill", "decode"), None), +} + + +def _sampler_key(cls): + return __import__(cls.__module__, fromlist=["SAMPLER"]).SAMPLER + + +def _info(rng, walk, sampler): + max_tokens = int(rng.integers(1, 6)) + return SimpleNamespace( + graph_walk=walk, step_metadata={}, max_tokens=max_tokens, + resource_configs={sampler: SimpleNamespace(ignore_eos=bool(rng.integers(2)))}, + dynamic_loop_iter_counts={LOOP: int(rng.integers(0, max_tokens + 1))}, + ) + + +@pytest.mark.parametrize("model", sorted(MODELS)) +@pytest.mark.parametrize("seed", range(8)) +def test_declared_policies_stop_the_rows_check_stop_does(model, seed): + cls, config, walks, stop = MODELS[model] + sampler = _sampler_key(cls) + sub = cls.__new__(cls) + torch.nn.Module.__init__(sub) + sub.config = config + if stop is not None: + sub.token_stop = stop + rng = np.random.default_rng(seed) + declared_walks = 0 + for walk in walks: + policies = sub.get_stop_policies(walk) + if policies is None: + continue # check_stop decides (a TokenStop's prompt walks) + declared_walks += 1 + for rid in range(16): + token = int(rng.choice([*EOS, 1, 2, 3])) + info = _info(rng, walk, sampler) + expected = sub.check_stop(rid, info, {"new_token": [torch.tensor([token])]}) or set() + declared = { + loop for loop, p in policies.items() + if p.stops(info.dynamic_loop_iter_counts.get(loop, 0), token, p.limits(info)) + } + assert declared == set(expected), (walk, rid) + assert declared_walks diff --git a/test/modular/test_tp_read_failure_acceptance.py b/test/modular/test_tp_read_failure_acceptance.py index 1fd3246956..a836454500 100644 --- a/test/modular/test_tp_read_failure_acceptance.py +++ b/test/modular/test_tp_read_failure_acceptance.py @@ -59,6 +59,7 @@ import pytest import torch +from mstar.graph.runtime.base import PhasedPostprocess from mstar.utils.ipc_format import WorkerMessageType HERE = Path(__file__).resolve().parent @@ -375,7 +376,7 @@ def _engine(rank, group, kv, prefix_cache, graph_runtime): return engine -class _Graph: +class _Graph(PhasedPostprocess): """The graph runtime and request state of one rank, for one node: a row is ready once each of its walk's inputs landed, and what the rank routes and sends is reported to the parent.""" @@ -696,6 +697,7 @@ def _worker(rank, scenario, comm, graph, engine, tm, group): w._stream_partition = {} w.streams = StreamBufferManager(w.request_state, w.tensor_manager, []) w._innermost_loop_by_node = {} + w._stop_token_signal = {} # no declared stop policies: check_stop decides w._metrics = None for setup in _DESIGN_SETUP: if hasattr(w, setup): diff --git a/test/modular/test_withheld_outputs.py b/test/modular/test_withheld_outputs.py index 87fd7b4944..ce5cb0daf7 100644 --- a/test/modular/test_withheld_outputs.py +++ b/test/modular/test_withheld_outputs.py @@ -17,7 +17,7 @@ from mstar.communication.tensors import SharedMemoryCommunicationManager from mstar.engine.engine import Engine, ExecutingBatch from mstar.engine.resources import StepContext -from mstar.graph.runtime.base import FreedTensors, RouteOutput +from mstar.graph.runtime.base import FreedTensors, PhasedPostprocess, RouteOutput from mstar.model.submodule_base import BatchedModelOutput, BatchedStopCheck, HostRows, NodeSubmodule, StopCheck from mstar.worker.micro_scheduler import ScheduledBatch from mstar.worker.stream_buffers import StreamBufferManager @@ -153,7 +153,9 @@ def poll_for_messages(self, timeout_ms=20): return False -class _RecordingRuntime: +class _RecordingRuntime(PhasedPostprocess): + """Recording phase stubs under the reference begin / resume.""" + def __init__(self): self.routes = [] self.stops = [] @@ -190,6 +192,7 @@ def test_the_worker_routes_no_tensor_for_a_withheld_output(tmp_path): w.device = torch.device("cpu") w._innermost_loop_by_node = {} w._walk_drivers = {} # no producer-driven walk transitions here + w._stop_token_signal = {} # no declared stop policies: check_stop decides w._graph_runtime = _RecordingRuntime() w.scheduler = SimpleNamespace(failed_rids=set()) w.tensor_manager = SharedMemoryCommunicationManager( @@ -201,6 +204,7 @@ def test_the_worker_routes_no_tensor_for_a_withheld_output(tmp_path): extend_prefix_chains=lambda batch, outputs, **kwargs: None, materialize_publish=lambda batch, rids: None, after_output_copy=lambda batch: None, + step_finished=lambda batch, outputs, **kwargs: None, )) w.request_state = SimpleNamespace( per_request_info={}, diff --git a/test/modular/test_worker_drain.py b/test/modular/test_worker_drain.py index 6b14460928..c14a48108a 100644 --- a/test/modular/test_worker_drain.py +++ b/test/modular/test_worker_drain.py @@ -15,7 +15,7 @@ from mstar.engine.engine import ExecutingBatch from mstar.engine.resources import StepContext from mstar.graph.base import GraphEdge, TensorPointerInfo -from mstar.graph.runtime.base import FreedTensors, RouteOutput +from mstar.graph.runtime.base import FreedTensors, PhasedPostprocess, RouteOutput from mstar.model.submodule_base import BatchedModelOutput from mstar.streaming.chunk_policy import FixedChunkPolicy from mstar.streaming.stream_buffer import StreamBuffer @@ -35,6 +35,12 @@ # ── worker ────────────────────────────────────────────────────────────────── +class _StubRuntime(SimpleNamespace, PhasedPostprocess): + """Stubbed phase methods under the reference begin / resume.""" + + supports_inline_emit = False + + def _worker( known_rids=("X",), in_flight=(), draining=(), reads_done=(), pending_drains=(), inflight_reads=False, is_follower=False, tp_follow=(), @@ -44,6 +50,7 @@ def _worker( w._walk_driven_edges = {} w.worker_id = "w0" w._phase_period = 0 + w._stop_token_signal = {} # no declared stop policies: check_stop decides w.is_tp_follower = is_follower w.sent = [] w.cleared = [] @@ -59,8 +66,9 @@ def _worker( w._pending_removes = {} # Identity interning: these tests use the rid string as its own handle, so # the string/handle split is exercised without a real runtime. - w._graph_runtime = SimpleNamespace( + w._graph_runtime = _StubRuntime( get_rid_handle=lambda r: r if r in known_rids else None, + register_stop_limits=lambda rid, info: None, get_rid_string=lambda h: h, remove_request=lambda h: None, # The TP fan-out reads the sharding config off the runtime now; no @@ -212,7 +220,7 @@ def test_a_walk_the_cache_served_whole_is_still_routed(): w.engine_manager.get_engine = lambda node: SimpleNamespace( check_stop_for_batch=lambda batch, outputs, **kwargs: {}, after_output_copy=lambda batch: None, - extend_prefix_chains=lambda batch, outputs, **kwargs: None, + step_finished=lambda batch, outputs, **kwargs: None, materialize_publish=lambda batch, rids: None, ) w.request_state.get_fwd_info = lambda rid, partition: None @@ -489,6 +497,7 @@ def test_add_new_request_hands_off_the_handle_not_the_string(): add_request=lambda request_id, **kw: rids.intern(request_id), get_rid_handle=rids.handle, get_rid_string=rids.name, get_sharding_config=lambda h: None, + register_stop_limits=lambda rid, info: None, ) seen = {} diff --git a/test/modular/test_worker_host_copies.py b/test/modular/test_worker_host_copies.py index 186537b60c..87625cab9c 100644 --- a/test/modular/test_worker_host_copies.py +++ b/test/modular/test_worker_host_copies.py @@ -23,7 +23,7 @@ from mstar.engine.engine import ExecutingBatch from mstar.engine.resources import StepContext from mstar.graph.base import GraphEdge -from mstar.graph.runtime.base import FreedTensors, RouteOutput +from mstar.graph.runtime.base import FreedTensors, PhasedPostprocess, RouteOutput from mstar.model.submodule_base import BatchedModelOutput from mstar.worker.micro_scheduler import ScheduledBatch from mstar.worker.stream_buffers import StreamBufferManager @@ -46,9 +46,11 @@ def poll_for_messages(self, timeout_ms=20): return False -class _AccumulatingRuntime: +class _AccumulatingRuntime(PhasedPostprocess): """Holds every iteration's output and registers all of them on the last.""" + supports_inline_emit = False + def __init__(self): self.held: list[int] = [] @@ -88,6 +90,7 @@ def _worker(tmp_path): w._walk_driven_edges = {} w._phase_period = 0 w.enable_nvtx = w.enable_prof = False + w._stop_token_signal = {} # no declared stop policies: check_stop decides w.device = torch.device("cpu") w._innermost_loop_by_node = {} w._graph_runtime = _AccumulatingRuntime() @@ -96,7 +99,7 @@ def _worker(tmp_path): w.engine_manager = SimpleNamespace(get_engine=lambda node: SimpleNamespace( check_stop_for_batch=lambda batch, outputs, **kwargs: {}, after_output_copy=lambda batch: None, - extend_prefix_chains=lambda batch, outputs, **kwargs: None, + step_finished=lambda batch, outputs, **kwargs: None, materialize_publish=lambda batch, rids: None, )) w.request_state = SimpleNamespace( diff --git a/test/rust/test_postprocess.py b/test/rust/test_postprocess.py new file mode 100644 index 0000000000..ff29fa1f32 --- /dev/null +++ b/test/rust/test_postprocess.py @@ -0,0 +1,231 @@ +"""The Rust runtime's native ``begin_postprocess`` / ``resume_postprocess`` +must match ``PhasedPostprocess``, the Python reference, run over the same +runtime's phase methods (``check_stops``, ``stop_loops_batched``, +``complete_and_route_batch``, ``send_outputs``). + +Two runtimes are driven through the same decode loop, one per path, and +everything observable is compared: where each step stops for the worker and +what it hands over, the frames each sends (read off real inboxes), and the +loop counters. One row samples EOS partway through, so steps cover the native +send, the send with stopped rows, and the completion of the worker graph. +""" +import sys +from types import SimpleNamespace + +sys.path.insert(0, ".") + +import pytest + +pytest.importorskip("mstar_rust") + +import torch +from mstar_rust import ZmqCommunicator + +import mstar.communication.wire_types # noqa: F401 (registers the tags) +from mstar.communication.tensor_store import RustTensorBookkeeping +from mstar.communication.wire import decode +from mstar.distributed.base import ShardingConfig +from mstar.graph.base import GraphEdge, GraphNode, Loop +from mstar.graph.runtime import rust as rust_runtime +from mstar.graph.runtime.base import ( + ColumnarEdgeSpecs, + InlineColumn, + PhasedPostprocess, + PostprocessNeed, + PostprocessRequest, + RouteTensors, + SendInput, +) +from mstar.graph.special_destinations import EMIT_TO_CLIENT +from mstar.model.base import WorkerGraph +from mstar.model.submodule_base import StopKind, StopPolicy +from mstar.utils.containers import ParallelList + +WALK, NODE, LOOP, SIG, PART, ME = "decode", "LLM", "decode_loop", "text_inputs", "default", "w0" +MAX_ITERS = 4 +EOS = 7 +POLICY = StopPolicy(StopKind.AR, token_signal=SIG, eos_tokens=(EOS,)) + + +class _Reference(rust_runtime.RustGraphRuntime): + """The Rust runtime with the Python reference's begin / resume.""" + begin_postprocess = PhasedPostprocess.begin_postprocess + resume_postprocess = PhasedPostprocess.resume_postprocess + + +def _runtime(cls, sock_dir): + section = Loop( + name=LOOP, + section=GraphNode( + name=NODE, input_names=[SIG], + outputs=[ + GraphEdge(next_node=NODE, name=SIG), + GraphEdge(next_node=EMIT_TO_CLIENT, name=SIG, output_modality="text"), + ], + ), + max_iters=MAX_ITERS, outputs=[], + ) + wg = WorkerGraph(section=section, graph_walks={WALK}, ranks=[0], worker_graph_id=0) + rt = cls( + my_worker_id=ME, my_worker_graphs=[wg], + all_wg_ids_to_graph_walks={0: {WALK}}, all_wg_ids_to_dyn_loops={0: {LOOP}}, + all_wg_ids_to_nodes={0: {NODE}}, node_to_partition={NODE: PART}, + sharding_config=ShardingConfig(groups=[], tp_enabled_nodes=set(), shard_dim={}), + bookkeeping=RustTensorBookkeeping(), + communicator=SimpleNamespace(_inner=ZmqCommunicator(ME, sock_dir)), + ) + rt.set_stop_policies(NODE, WALK, {LOOP: POLICY}) + return rt + + +def _collect(box): + out = [] + first = box.recv_timeout(1000) + if first is None: + return out + out.append(decode(first)) + while (more := box.recv_timeout(50)) is not None: + out.append(decode(more)) + return out + + +def _start(rt, n, limits=True): + rids = [] + for i in range(n): + rid = rt.add_request( + request_id=f"r{i}", partition=PART, graph_walk=WALK, + partition_worker_graph_ids=[0], + worker_graph_to_workers=ParallelList([0], [[ME]]), loop_caps={}) + if limits: + # a limit no step reaches: EOS and the loop's cap end the rows + rt.register_stop_limits(rid, SimpleNamespace(max_tokens=1 << 20, resource_configs={})) + rids.append(rid) + block = ColumnarEdgeSpecs.empty() + for rid in rids: + block.add(rid, SIG, [], False, next_node=NODE) + rt.ingest_inputs_batch(block) + return rids + + +def _drive(rt, request, values, model_stops, dropped=()): + """One step's postprocess as the worker runs it; each break, summarized.""" + trace = [] + step = rt.begin_postprocess(request) + while True: + trace.append(( + step.need, None if step.wg_ids is None else list(step.wg_ids.keys), + sorted(step.stopped_rids), sorted(step.completed_rids), + )) + if step.need is PostprocessNeed.DONE: + return trace + if step.need is PostprocessNeed.STOPS: + step = rt.resume_postprocess(stops=model_stops, dropped=dropped) + elif step.need is PostprocessNeed.STORE: + keys = step.wg_ids.keys + step = rt.resume_postprocess(stored=RouteTensors( + [], [0] * len(keys), SIG, [values[r] for r in keys], + )) + else: + info = sorted(step.route_output.rids_needing_request_info) + step = rt.resume_postprocess(send=SendInput( + completion_id=step.route_output.completion_id, + per_request_info=ParallelList(info, [None] * len(info)), + new_token_counts=ParallelList([], []), + )) + + +@pytest.mark.parametrize( + "mode", ["runtime", "model_stops", "stored"], + ids=["runtime_stops_inline", "model_stops_inline", "runtime_stops_stored"], +) +@pytest.mark.parametrize("dtype", [torch.int64, torch.int32]) +@pytest.mark.parametrize("reverse_rows", [False, True], ids=["in_order", "permuted"]) +def test_native_matches_reference(tmp_path, mode, dtype, reverse_rows): + n = 5 + boxes, rts = {}, {} + for k, cls in (("native", rust_runtime.RustGraphRuntime), ("reference", _Reference)): + d = tmp_path / k + d.mkdir() + boxes[k] = {p: ZmqCommunicator(p, str(d)) for p in ("api_server", "conductor")} + rts[k] = _runtime(cls, str(d)) + rids = [_start(rt, n) for rt in rts.values()] + assert rids[0] == rids[1] + active = rids[0] + needs = set() + for step in range(MAX_ITERS): + for rt in rts.values(): + assert rt.pop_rids(NODE, WALK, active) is not None + # row 1 samples EOS on the second step + values = {r: EOS if (step, i) == (1, 1) else 100 * step + i for i, r in enumerate(active)} + # host rows as the stop check stages them: [rows, 1], row i for row_rids[i] + row_rids = list(reversed(active)) if reverse_rows else list(active) + rows = torch.tensor([[values[r]] for r in row_rids], dtype=dtype) + column = rows.reshape(rows.shape[0], -1)[:, 0] + # the model's stops, checked before the runtime when it does not check + model_stops = {r: {LOOP} for r, v in values.items() if v == EOS} + request = PostprocessRequest( + partition=PART, node_name=NODE, graph_walk=WALK, output_signals=[SIG], + wg_ids=ParallelList(list(active), [0] * len(active)), + stops=model_stops if mode == "model_stops" else None, + runtime_stops=mode != "model_stops", tokens=column, row_rids=row_rids, + inline=None if mode == "stored" else InlineColumn(SIG, column), + ) + traces = {k: _drive(rt, request, values, model_stops) for k, rt in rts.items()} + assert traces["native"] == traces["reference"], step + needs |= {t[0] for t in traces["native"]} + stopped = set(traces["native"][-1][2]) + assert ( + rts["native"].get_dynamic_loop_iters(active, PART).values + == rts["reference"].get_dynamic_loop_iters(active, PART).values + ) + active = [r for r in active if r not in stopped] + expected = {PostprocessNeed.DONE, PostprocessNeed.SEND} + expected |= {"stored": {PostprocessNeed.STORE}}.get(mode, set()) + assert needs == expected + for peer in ("api_server", "conductor"): + got = {k: _collect(boxes[k][peer]) for k in boxes} + assert got["native"] == got["reference"], peer + assert got["native"], peer + for rt in rts.values(): + assert rt._rust.num_parked_completions() == 0 + + +def test_dropped_rows_are_not_routed(tmp_path): + """When the runtime's own check cannot decide (here: no limits), a row + ``check_stop`` then fails (STOPS' ``dropped``) leaves the step: it is + neither stopped nor routed, in either implementation.""" + traces = {} + for k, cls in (("native", rust_runtime.RustGraphRuntime), ("reference", _Reference)): + d = tmp_path / k + d.mkdir() + rt = _runtime(cls, str(d)) + rids = _start(rt, 3, limits=False) + assert rt.pop_rids(NODE, WALK, rids) is not None + column = torch.tensor([1, 2, 3], dtype=torch.int64) + request = PostprocessRequest( + partition=PART, node_name=NODE, graph_walk=WALK, output_signals=[SIG], + wg_ids=ParallelList(list(rids), [0] * 3), runtime_stops=True, tokens=column, + row_rids=rids, inline=InlineColumn(SIG, column), + ) + traces[k] = _drive( + rt, request, dict(zip(rids, [1, 2, 3], strict=True)), + model_stops={rids[1]: {LOOP}}, dropped=[rids[1]], + ) + assert traces["native"] == traces["reference"] + need, routed, stopped, _ = traces["native"][-1] + assert need is PostprocessNeed.DONE and stopped == [] + assert traces["native"][0][0] is PostprocessNeed.STOPS + + +def test_missing_limits_leave_the_check_to_the_model(tmp_path): + """A row with no registered limits: the runtime decides nothing (STOPS).""" + rt = _runtime(rust_runtime.RustGraphRuntime, str(tmp_path)) + rid = rt.add_request( + request_id="r", partition=PART, graph_walk=WALK, partition_worker_graph_ids=[0], + worker_graph_to_workers=ParallelList([0], [[ME]]), loop_caps={}) + column = torch.tensor([EOS], dtype=torch.int64) + step = rt.begin_postprocess(PostprocessRequest( + partition=PART, node_name=NODE, graph_walk=WALK, output_signals=[SIG], + wg_ids=ParallelList([rid], [0]), runtime_stops=True, tokens=column, row_rids=[rid], + )) + assert step.need is PostprocessNeed.STOPS diff --git a/test/rust/test_route_and_send_inline.py b/test/rust/test_route_and_send_inline.py deleted file mode 100644 index 57a464bf50..0000000000 --- a/test/rust/test_route_and_send_inline.py +++ /dev/null @@ -1,168 +0,0 @@ -"""``route_and_send_inline`` must be exactly ``complete_and_route_batch`` then -``send_outputs`` with the device-loopback arguments the worker builds. - -Two runtimes are driven through the same decode steps, one per path, and -everything observable is compared: the frames each sends (read off real -inboxes), the next pop, and the loop counters. The loop is short, so the last -step completes the worker graph and exercises the fallback (``ROUTED``) where a -frame needs Python's per-request payloads. -""" -import sys -from types import SimpleNamespace - -sys.path.insert(0, ".") - -import pytest - -pytest.importorskip("mstar_rust") - -import torch -from mstar_rust import ZmqCommunicator - -import mstar.communication.wire_types # noqa: F401 (registers the tags) -from mstar.communication.tensor_store import RustTensorBookkeeping -from mstar.communication.wire import decode -from mstar.distributed.base import ShardingConfig -from mstar.graph.base import GraphEdge, GraphNode, Loop -from mstar.graph.runtime import rust as rust_runtime -from mstar.graph.runtime.base import ColumnarEdgeSpecs, InlineRouteStatus, RouteInput, SendInput -from mstar.graph.special_destinations import EMIT_TO_CLIENT -from mstar.model.base import WorkerGraph -from mstar.utils.containers import ParallelList - -WALK, NODE, LOOP, SIG, PART, ME = "decode", "LLM", "decode_loop", "text_inputs", "default", "w0" -MAX_ITERS = 4 - - -def _runtime(sock_dir): - section = Loop( - name=LOOP, - section=GraphNode( - name=NODE, input_names=[SIG], - outputs=[ - GraphEdge(next_node=NODE, name=SIG), - GraphEdge(next_node=EMIT_TO_CLIENT, name=SIG, output_modality="text"), - ], - ), - max_iters=MAX_ITERS, outputs=[], - ) - wg = WorkerGraph(section=section, graph_walks={WALK}, ranks=[0], worker_graph_id=0) - return rust_runtime.RustGraphRuntime( - my_worker_id=ME, my_worker_graphs=[wg], - all_wg_ids_to_graph_walks={0: {WALK}}, all_wg_ids_to_dyn_loops={0: {LOOP}}, - all_wg_ids_to_nodes={0: {NODE}}, node_to_partition={NODE: PART}, - sharding_config=ShardingConfig(groups=[], tp_enabled_nodes=set(), shard_dim={}), - bookkeeping=RustTensorBookkeeping(), - communicator=SimpleNamespace(_inner=ZmqCommunicator(ME, sock_dir)), - ) - - -def _collect(box): - out = [] - first = box.recv_timeout(1000) - if first is None: - return out - out.append(decode(first)) - while (more := box.recv_timeout(50)) is not None: - out.append(decode(more)) - return out - - -def _start(rt, n): - rids = [ - rt.add_request( - request_id=f"r{i}", partition=PART, graph_walk=WALK, - partition_worker_graph_ids=[0], - worker_graph_to_workers=ParallelList([0], [[ME]]), loop_caps={}) - for i in range(n) - ] - block = ColumnarEdgeSpecs.empty() - for rid in rids: - block.add(rid, SIG, [], False, next_node=NODE) - rt.ingest_inputs_batch(block) - return rids - - -def _general_step(rt, rids, values): - out = rt.complete_and_route_batch(RouteInput( - partition=PART, graph_walk=WALK, node_name=NODE, output_signals=[SIG], - wg_ids=ParallelList(rids, [0] * len(rids)), tensors=[], - num_tensors=[0] * len(rids), inline_signal=SIG, inline_values=values, - )) - info = sorted(out.rids_needing_request_info) - return rt.send_outputs(SendInput( - completion_id=out.completion_id, - per_request_info=ParallelList(info, [None] * len(info)), - new_token_counts=ParallelList([], []), - )) - - -def _fused_step(rt, rids, column, row_rids): - status, _cid, sent, route = rt.route_and_send_inline( - PART, WALK, NODE, [SIG], rids, SIG, column, row_rids, - ) - assert status in (InlineRouteStatus.SENT, InlineRouteStatus.ROUTED) - if status is InlineRouteStatus.SENT: - assert route is None - return sent - info = sorted(route.rids_needing_request_info) - assert info, "ROUTED only when a frame needs Python's payloads" - return rt.send_outputs(SendInput( - completion_id=route.completion_id, - per_request_info=ParallelList(info, [None] * len(info)), - new_token_counts=ParallelList([], []), - )) - - -@pytest.mark.parametrize("dtype", [torch.int64, torch.int32]) -@pytest.mark.parametrize("reverse_rows", [False, True], ids=["in_order", "permuted"]) -def test_fused_matches_route_then_send(tmp_path, dtype, reverse_rows): - n = 5 - dirs = {k: tmp_path / k for k in ("g", "f")} - boxes = {} - rts = {} - for k, d in dirs.items(): - d.mkdir() - boxes[k] = {p: ZmqCommunicator(p, str(d)) for p in ("api_server", "conductor")} - rts[k] = _runtime(str(d)) - rids = {k: _start(rt, n) for k, rt in rts.items()} - assert rids["g"] == rids["f"] - rids = rids["g"] - statuses = [] - for step in range(MAX_ITERS): - for rt in rts.values(): - assert rt.pop_rids(NODE, WALK, rids) is not None - values = [1000 * step + i for i in range(n)] - # host rows as the stop check stages them: [rows, 1], row i for row_rids[i] - row_rids = list(reversed(rids)) if reverse_rows else list(rids) - rows = torch.tensor([values[rids.index(r)] for r in row_rids], dtype=dtype) - column = rows.reshape(-1, 1).reshape(rows.shape[0], -1)[:, 0] - sent_g = _general_step(rts["g"], rids, values) - sent_f = _fused_step(rts["f"], rids, column, row_rids) - assert sent_g == sent_f - statuses.append(bool(sent_f)) - assert ( - rts["g"].get_dynamic_loop_iters(rids, PART).values - == rts["f"].get_dynamic_loop_iters(rids, PART).values - ) - assert statuses[-1], "the last step completes the worker graph" - for peer in ("api_server", "conductor"): - got = {k: _collect(boxes[k][peer]) for k in boxes} - assert got["g"] == got["f"], peer - assert got["g"], peer - for rt in rts.values(): - assert rt._rust.num_parked_completions() == 0 - - -def test_a_rid_without_a_row_does_nothing(tmp_path): - rt = _runtime(str(tmp_path)) - rids = _start(rt, 3) - assert rt.pop_rids(NODE, WALK, rids) is not None - column = torch.tensor([7, 8], dtype=torch.int64) - status, _cid, sent, route = rt.route_and_send_inline( - PART, WALK, NODE, [SIG], rids, SIG, column, rids[:2], - ) - assert (status, sent, route) == (InlineRouteStatus.NOT_ROUTED, [], None) - # nothing was routed: the general path still completes this step - assert rt._rust.num_parked_completions() == 0 - assert _general_step(rt, rids, [1, 2, 3]) == [] diff --git a/test/rust/test_runtime_check_stops.py b/test/rust/test_runtime_check_stops.py new file mode 100644 index 0000000000..678b2fddac --- /dev/null +++ b/test/rust/test_runtime_check_stops.py @@ -0,0 +1,114 @@ +"""Each graph runtime's stop check (``GraphRuntime.check_stops``), over the +token column the worker's ``PostprocessRequest`` hands it, stops what the +declared policy does (``StopPolicy.stops``) on the loop count the step was +built with, through the worker's async decode loop. + +A runtime reads each loop's count from its walk state at postprocess time, by +when step N+1 has been speculated; this drives a speculated decode loop in the +worker's stage order and compares at every step's stop check. +""" +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from test_spec_rows import ( + KINDS, + LOOP, + NODE, + PART, + WALK, + _admit, + _advance, + _first_step, + _rows_only, + _speculate, + _worker, +) + +from mstar.model.submodule_base import HostRows, StopKind, StopPolicy + +EOS = 7 +SAMPLER = "sampler" +POLICY = StopPolicy(StopKind.AR, token_signal="new_token", eos_tokens=(EOS,), sampler=SAMPLER) + + +# no signal stays on the device, so nothing rides inline +_ENGINE = SimpleNamespace( + device_loopback_signals=lambda node, walk: (), inline_client_signals=lambda node, walk: {}, +) + + +def _declared_worker(kind, policy=POLICY): + w = _worker(kind) + w._graph_runtime.set_stop_policies(NODE, WALK, {LOOP: policy}) + w._stop_token_signal = {(NODE, WALK): policy.token_signal} + w._device_loopback = {} + w._innermost_loop_by_node = {} + return w + + +def _expected(nb, host_rows) -> dict[int, set[str]]: + tokens = dict(zip(host_rows.request_ids, host_rows.buffers["new_token"][:, 0].tolist(), strict=True)) + out = {} + for rid in nb.request_ids: + info = nb.per_request_info_wrapped[rid] + if rid in tokens and POLICY.stops( + info.dynamic_loop_iter_counts.get(LOOP, 0), tokens[rid], POLICY.limits(info), + ): + out[rid] = {LOOP} + return out + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("seed", range(4)) +def test_the_runtime_stops_what_the_policy_does_through_a_speculated_loop(kind, seed, monkeypatch): + rng = np.random.default_rng(seed) + w = _declared_worker(kind) + rids = _admit(w, 6) + for i, rid in enumerate(rids): + info = w.request_state.get_fwd_info(rid, PART) + info.max_tokens = 2 + i + info.resource_configs[SAMPLER] = SimpleNamespace(ignore_eos=bool(rng.integers(2))) + w._graph_runtime.register_stop_limits(rid, info) + pending = _first_step(w, rids) + stopped = 0 + for _ in range(8): + spec = _speculate(w, pending, rows=True, monkeypatch=monkeypatch) + assert spec is not None + w._thread_outputs_to_speculative(spec, _rows_only(pending)) + # postprocess N: N+1 is speculated, N is not routed yet + nb = pending.node_batch + # host rows in another order than the step's, and sometimes one row + # with no output at all + ids = list(nb.request_ids) + order = rng.permutation(len(ids))[: len(ids) - int(rng.random() < 0.5)] + tokens = torch.tensor(rng.choice([EOS, 1, 2], size=len(order)))[:, None] + host_rows = HostRows(tuple(ids[i] for i in order), {"new_token": tokens}) + runtime_stops, tokens = w._runtime_stop_column(pending, host_rows) + assert runtime_stops + req = w._get_postprocess_request(pending, host_rows, _ENGINE, runtime_stops, tokens, None) + assert req.inline is None + got = w._graph_runtime.check_stops( + req.partition, req.node_name, req.graph_walk, list(req.wg_ids.keys), + req.tokens, list(req.row_rids), + ) + assert got == _expected(nb, host_rows) + stopped += len(got) + pending = _advance(w, pending, spec) + assert stopped, "the limits are small enough that rows reach them" + + +@pytest.mark.parametrize("kind", KINDS) +def test_limits_are_dropped_with_the_request(kind): + w = _declared_worker(kind, StopPolicy(StopKind.ITER, max_iters=lambda info: 3)) + rt = w._graph_runtime + (rid,) = _admit(w, 1) + rt.register_stop_limits(rid, w.request_state.get_fwd_info(rid, PART)) + assert rt.check_stops(PART, NODE, WALK, [rid], None, []) == {} + rt.remove_request(rid) + w.request_state.remove_request(rid) + (again,) = _admit(w, 1, start=1) + assert again == rid, "the handle is recycled" + # until its own limits are registered, the runtime decides nothing + assert rt.check_stops(PART, NODE, WALK, [again], None, []) is None From 1d4cc87a9ddca336f933cc01ce66356954e6cc55 Mon Sep 17 00:00:00 2001 From: Naomi Sagan Date: Sun, 1 Nov 2026 20:21:28 +0000 Subject: [PATCH 2/2] test_inline_tokens: a withheld output routes through the store main's test called _route_and_send_fast, which this PR replaces; the same check is now the request's inline column. The log-probs case's batch stub gains the withheld_outputs the inline test reads. Co-Authored-By: Claude Opus 5.5 --- test/modular/test_inline_tokens.py | 49 +++++++++++++----------------- 1 file changed, 21 insertions(+), 28 deletions(-) diff --git a/test/modular/test_inline_tokens.py b/test/modular/test_inline_tokens.py index e0f3ef383e..28c188f596 100644 --- a/test/modular/test_inline_tokens.py +++ b/test/modular/test_inline_tokens.py @@ -269,30 +269,6 @@ def test_inline_values_follow_the_batch_order_off_the_host_rows(): assert sig == "text_inputs" and vals == [90, 70, 80] -def test_a_withheld_output_skips_the_fused_route_and_send(): - """The fused device-loopback call routes every row's outputs, so a step - whose stop check withheld one takes the general path.""" - class _Fused(_Runtime): - supports_route_and_send_inline = True - - def route_and_send_inline(self, **kwargs): - raise AssertionError("a withheld row would ride inline") - - class _LoopbackEngine(_Engine): - def device_loopback_signals(self, node_name, graph_walk): - return ["new_token"] - - w = _worker(_Fused()) - w.enable_prof, w._device_loopback = False, {} - batch = _Batch(withheld={8: ("new_token",)}) - batch.partition = "p" - batch.batch.output_signals, batch.batch.graph_walk = ["new_token"], "decode" - batch.node_batch.partition_done_rids = () - rows = HostRows((7, 8), {"new_token": torch.tensor([[70], [80]])}) - eng = _LoopbackEngine({"new_token": "new_token"}) - assert w._route_and_send_fast(batch, rows, eng, [7, 8], set()) is False - - def test_inline_values_fall_back_when_a_rid_has_no_row_or_nothing_applies(): rows = HostRows((7, 8), {"new_token": torch.tensor([70, 80])}) w = _worker() @@ -349,9 +325,9 @@ def test_row_logprobs_follow_the_batch_order(): assert Worker._row_logprobs(HostRows((7,), {SAMPLED_LOGPROB: torch.tensor([nan])}), [7]) == [] -def test_a_step_with_logprobs_routes_through_the_store(): - # the in-runtime route builds no log-probs: such a step stores, and STORE - # passes them on with the route tensors +def _inline_step(withheld=None): + """A one-row device-loopback decode step whose token rides inline, as + `_get_postprocess_request` sees it.""" w = _worker() w._stop_token_signal, w._device_loopback, w._innermost_loop_by_node = {}, {}, {} w.enable_prof = False @@ -363,10 +339,27 @@ def test_a_step_with_logprobs_routes_through_the_store(): batch.partition = "p" batch.batch.output_signals = ["text_inputs"] batch.batch.request_to_worker_graph = {7: 0} - batch.node_batch = SimpleNamespace(final_stream_rids=(), partition_done_rids=()) + batch.node_batch = SimpleNamespace( + final_stream_rids=(), partition_done_rids=(), withheld_outputs=withheld or {}, + ) batch.batch.rids_by_walk = lambda: {"decode": [7]} + return w, engine, batch + + +def test_a_step_with_logprobs_routes_through_the_store(): + # the in-runtime route builds no log-probs: such a step stores, and STORE + # passes them on with the route tensors + w, engine, batch = _inline_step() tokens = torch.tensor([70]) plain = w._get_postprocess_request(batch, HostRows((7,), {"new_token": tokens}), engine, False, None, None) assert plain.inline is not None and plain.inline.signal == "text_inputs" rows = HostRows((7,), {"new_token": tokens, SAMPLED_LOGPROB: torch.tensor([-0.5])}) assert w._get_postprocess_request(batch, rows, engine, False, None, None).inline is None + + +def test_a_withheld_output_routes_through_the_store(): + """An inline route carries every row's outputs, so a step whose stop check + withheld one stores, and the store leaves the withheld output out.""" + rows = HostRows((7,), {"new_token": torch.tensor([70])}) + w, engine, batch = _inline_step(withheld={7: ("text_inputs",)}) + assert w._get_postprocess_request(batch, rows, engine, False, None, None).inline is None