From 006f3c728e1861ee540e235c4fc3a83712a79ffa Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Mon, 19 Oct 2026 18:29:48 +0000 Subject: [PATCH 01/15] Fail only the request whose async (Mooncake) tensor read failed, and free its buffers once its transfers stop Fixes #418. get_ready_tensors re-raised a failed async read on every poll and never reached the line that removes it, so every request behind it on that worker or data worker stalled until the failed request was torn down. A failed or cancelled read is now dropped and reported once through take_failed_reads. The worker fails its request through FAIL_REQUESTS; the data worker answers it with one error chunk in the output slot of each tensor the read would have delivered. Tensor reads go through TrackedMooncakeReader: one batch per tensor, read with batch_transfer_async_read and polled with transfer_check_status, with the sync read's retries and deadline. A failed read hands over the batch ids still running; the manager keeps the read's buffers (out of the store, so teardown cannot free them) and frees them once every id has reported COMPLETED or FAILED. A failed submit, a status other than those four, an error that names no transfers, or a failed unregister keeps the buffers for good. KV transfers keep the batched sync read. --- mstar/api_server/data_worker.py | 30 +- mstar/communication/tensors.py | 352 ++++++++++- mstar/worker/worker.py | 7 + test/modular/test_api_completion_guard.py | 1 + test/modular/test_api_result_delivery.py | 5 + test/modular/test_data_worker_late_release.py | 5 + test/modular/test_data_worker_wakeup.py | 5 + test/modular/test_empty_tensor_transport.py | 14 +- test/modular/test_failed_tensor_read.py | 589 ++++++++++++++++++ test/modular/test_inline_tokens.py | 5 + test/modular/test_mooncake_flat_layout.py | 5 +- test/modular/test_process_liveness.py | 1 + .../test_request_failure_propagation.py | 2 + test/modular/test_request_id_reuse.py | 5 + test/modular/test_video_frame_protocol.py | 5 + 15 files changed, 1006 insertions(+), 25 deletions(-) create mode 100644 test/modular/test_failed_tensor_read.py diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index 79e48949a0..113066fcca 100644 --- a/mstar/api_server/data_worker.py +++ b/mstar/api_server/data_worker.py @@ -27,7 +27,11 @@ ) from mstar.communication.communicator import BaseCommunicator, CommProtocol, make_communicator from mstar.communication.event import EventWakeup -from mstar.communication.tensors import NameToTensorList, create_tensor_communication_manager +from mstar.communication.tensors import ( + FailedRead, + NameToTensorList, + create_tensor_communication_manager, +) from mstar.engine.resources.kv.config import KVSpec, PagedKVConfig from mstar.engine.resources.kv.keys import PageItem, chain, end_key, fingerprint, page_items from mstar.engine.resources.spec import apply_yaml_overrides @@ -902,7 +906,7 @@ def _fail_request( ``abort`` (False for a request the conductor has not seen) stops the request's walk, since its client already has its answer. """ - logger.exception("%s failed for request %s", stage, request_id) + logger.error("%s failed for request %s", stage, request_id, exc_info=exc) status = 400 if isinstance(exc, (ValueError, TypeError)) else 500 chunks = [ ResultChunk( @@ -1040,9 +1044,29 @@ def _discard_result_tensor( graph_edges=[result.graph_edge], ) + def _fail_unread_outputs(self, failed: FailedRead) -> None: + """Fail the request of a read that will not be delivered: one error + chunk in the output slot of each tensor it would have delivered. The + manager already acked or kept its buffers.""" + request_id = failed.request_id + state = self.request_output_state.get(request_id) + metadata = self.tensor_uuid_to_metadata_per_request.get(request_id, {}) + for edge in failed.graph_edges: + for info in edge.tensor_info: + metadata.pop(info.uuid, None) + slot = state.order.pop(info.uuid, None) if state is not None else None + self._fail_request( + request_id, failed.error, "output tensor read", + sequence=slot[0] if slot is not None else None, + ) + def _process_read_tensors(self): did_work = False - for request_id, graph_edges in self.tensor_manager.get_ready_tensors().items(): + ready = self.tensor_manager.get_ready_tensors() + for failed in self.tensor_manager.take_failed_reads(): + did_work = True + self._fail_unread_outputs(failed) + for request_id, graph_edges in ready.items(): did_work = True for graph_edge in graph_edges: modality = graph_edge.name.replace("_output", "") diff --git a/mstar/communication/tensors.py b/mstar/communication/tensors.py index d9f696f977..f9619c9b19 100644 --- a/mstar/communication/tensors.py +++ b/mstar/communication/tensors.py @@ -1,6 +1,7 @@ import errno import fcntl import hashlib +import json import logging import math import os @@ -13,9 +14,9 @@ import time from abc import ABC, abstractmethod from collections.abc import Callable -from concurrent.futures import Future, ThreadPoolExecutor +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor from contextlib import nullcontext as _nullcontext -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import NamedTuple from mstar.communication.tensor_store import ( # NameToTensorList, Rid re-exported @@ -76,6 +77,8 @@ class FutureAndPointers: graph_edges: list[GraphEdge] request_id: Rid = "" rx_time: float | None = None # seconds + # uuids of the local buffers this read writes into + transfer_uuids: list[int] = field(default_factory=list) # What the read moved over the transport (see count_read), set only while # transfer metrics are on and cleared once counted. Unannotated, so class # attributes rather than fields: a read with them off stores nothing more. @@ -109,6 +112,32 @@ def count_read(self, info: TensorPointerInfo) -> None: row[1] += 1 +@dataclass +class FailedRead: + """Edges of a request that will not be delivered because one of its + reads failed. ``error`` is that read's exception.""" + request_id: Rid + graph_edges: list[GraphEdge] + error: BaseException + + +@dataclass +class _KeptRead: + """The local buffers of a read that failed or was torn down mid-transfer, + kept until every transfer into them has reported that it stopped. + + ``future`` is set while the read still runs in the reader thread, which + owns its batch ids until it finishes. ``running`` holds the ids handed + over after that (local pointer -> batch id), polled by the manager. + ``registered`` is index-parallel to ``buffers``.""" + request: str + buffers: list[torch.Tensor] + registered: list[bool] + nbytes: int + future: Future | None = None + running: dict[int, int] = field(default_factory=dict) + + #: Earlier Mooncake releases let any host that reaches the TCP data port read #: or write arbitrary process memory (CVE-2026-103764). NVD lists "before #: 0.3.13" because it points at a later refactor; the fix (Mooncake #2314) @@ -254,6 +283,57 @@ class TransferReadInfo: fp: FutureAndPointers | None = None +# transfer_check_status: COMPLETED and FAILED free the batch; WAITING, and +# TIMEOUT (a slice past MC_SLICE_TIMEOUT), mean slices may still be moving +_TRANSFER_COMPLETED = 1 +_TRANSFER_FAILED = -1 +_TRANSFER_RUNNING = (0, -2) + +# between a read's status polls. Each poll takes the GIL, which the sync read's +# wait never did; polling faster slows a main thread that releases it often. +_READ_POLL_INTERVAL_S = 1e-4 + + +def _submit_failed(batch_id: int) -> bool: + """A failed ``batch_transfer_async_read``: 0, or uint64 -1 from Mooncake + releases before 0.3.13.""" + return batch_id == 0 or batch_id == 2**64 - 1 + + +def _mooncake_transfer_timeout_s() -> int: + """The timeout Mooncake's sync read applies before its 1 ns per byte: + ``MC_TRANSFER_TIMEOUT`` (read with atoi, at least 5), else 30.""" + raw = os.environ.get("MC_TRANSFER_TIMEOUT") + if raw is None: + return 30 + match = re.match(r"\s*[+-]?\d+", raw) + return max(5, int(match.group())) if match else 5 + + +def _local_hca_count(engine) -> int: + """HCAs in the engine's local topology: Mooncake's sync read tries a + failed batch ``numContexts() + 1`` times.""" + try: + topology = json.loads(engine.get_local_topology()) + return len({hca for tiers in topology.values() for tier in tiers for hca in tier}) + except Exception: # noqa: BLE001 — a topology we can't read means no retries + return 0 + + +class AsyncReadError(RuntimeError): + """A Mooncake read that failed. ``running`` maps the local pointer of each + buffer whose transfer had not reported COMPLETED or FAILED to its batch + id; the reader stops polling them when it raises, and the manager that + keeps the buffers takes them over. ``unknown`` holds pointers whose + transfer state cannot be learned: a failed submit (which may have posted + part of the transfer) or a status this code does not know.""" + + def __init__(self, message: str, running: dict[int, int], unknown: set[int]): + super().__init__(message) + self.running = running + self.unknown = unknown + + class AsyncMooncakeReader: """Background thread for non-blocking mooncake READ operations. @@ -298,7 +378,15 @@ def _do_read(self, read_info: list["TransferReadInfo"], event: torch.cuda.Event self._copy_stream.synchronize() start_time = time.perf_counter() + self._read(read_info) + end_time = time.perf_counter() + # The caller points every read of one edge at its FutureAndPointers + # only when it wants the time (--log-stats or transfer metrics). + if read_info[0].fp is not None: + for info in read_info: + info.fp.rx_time = (end_time - start_time) + def _read(self, read_info: list["TransferReadInfo"]) -> None: # group read_info by session id for batch read grouped_read = {} for info in read_info: @@ -316,12 +404,6 @@ def _do_read(self, read_info: list["TransferReadInfo"], event: torch.cuda.Event ) if status < 0: raise RuntimeError(f"Mooncake read failed. Status: {status}") - end_time = time.perf_counter() - # The caller points every read of one edge at its FutureAndPointers - # only when it wants the time (--log-stats or transfer metrics). - if read_info[0].fp is not None: - for info in read_info: - info.fp.rx_time = (end_time - start_time) def wait_all(self): """Block until all pending writes complete. Re-raises exceptions.""" @@ -335,6 +417,73 @@ def shutdown(self): self._executor.shutdown(wait=True) +class TrackedMooncakeReader(AsyncMooncakeReader): + """Tensor reads that say, when they fail, which transfers may still be + writing into their buffers. + + Each tensor is its own batch, read with ``batch_transfer_async_read`` and + polled with ``transfer_check_status``, which reports on a batch's first + task only. The sync read gives up on a timeout or on one failed tensor + without waiting for the rest of its batch, and returns no handle to them. + Retries and the deadline are the sync read's. + """ + + def __init__(self, engine, device, max_workers: int = 3): + super().__init__(engine, device, max_workers=max_workers) + self._read_attempts = _local_hca_count(engine) + 1 + self._timeout_s = _mooncake_transfer_timeout_s() + + def _read(self, read_info: list["TransferReadInfo"]) -> None: + """Raises ``AsyncReadError``, after which this thread polls none of + the batch ids it hands over.""" + engine = self._engine + deadline = time.monotonic() + self._timeout_s + sum(i.nbytes for i in read_info) * 1e-9 + running: dict[int, tuple[int, TransferReadInfo]] = {} + attempts: dict[int, int] = {} + + def fail(message: str, unknown: set[int] = frozenset()) -> AsyncReadError: + return AsyncReadError( + message, {ptr: batch_id for ptr, (batch_id, _) in running.items()}, set(unknown), + ) + + def submit(info: TransferReadInfo) -> None: + batch_id = engine.batch_transfer_async_read( + info.source_session_id, [info.local_ptr], [info.remote_ptr], [info.nbytes], + ) + if _submit_failed(batch_id): + raise fail("Mooncake read submit failed", {info.local_ptr}) + running[info.local_ptr] = (batch_id, info) + attempts[info.local_ptr] = attempts.get(info.local_ptr, 0) + 1 + + for info in read_info: + submit(info) + while running: + for ptr, (batch_id, info) in list(running.items()): + status = engine.transfer_check_status(batch_id) + if status in _TRANSFER_RUNNING: + # the read needs every tensor, so one still running is + # enough to wait on; polling the rest costs GIL hand-offs + break + del running[ptr] + if status == _TRANSFER_COMPLETED: + continue + if status != _TRANSFER_FAILED: + # may or may not have freed the batch: never poll it again + raise fail(f"Mooncake read returned unknown status {status}", {ptr}) + if attempts[ptr] >= self._read_attempts or time.monotonic() > deadline: + raise fail(f"Mooncake read failed. Status: {status}") + submit(info) + if not running: + break + if time.monotonic() > deadline: + raise fail(f"Mooncake read timed out after {self._timeout_s} s") + time.sleep(_READ_POLL_INTERVAL_S) + + def transfer_status(self, batch_id: int) -> int: + """``transfer_check_status`` of one batch this reader handed over.""" + return self._engine.transfer_check_status(batch_id) + + #: Mooncake's wheel per CUDA major. All of them install the same ``mooncake`` #: package; the default one is built for CUDA 12. _MOONCAKE_WHEELS = { @@ -537,6 +686,15 @@ def __init__( self.rid_to_str: Callable[[Rid], str] = str self._uuid_minter = TensorUuidMinter(my_entity_id) self.pending: list[FutureAndPointers] = [] + # rid -> the error of its first failed read; held until cleanup_request + self._failed_read_errors: dict[Rid, BaseException] = {} + self._failed_reads: list[FailedRead] = [] + # buffers of failed or torn-down reads, freed once their transfers stop + self._kept_reads: list[_KeptRead] = [] + # buffers whose transfers can't be known to have stopped: never freed + self._unreleasable_buffers: list[torch.Tensor] = [] + # batch id -> transfer_check_status; None: the transport reports nothing + self._transfer_status: Callable[[int], int] | None = None self.read_finished: dict[Rid, set[int]] = {} # req_id -> {(producer uuid, offset, nbytes, dims) -> local uuid} for # the received slices ``_take_cached_read`` re-pointed, and req_id -> @@ -1422,7 +1580,149 @@ def _record_transfer(self, stats: EdgeTransfers, ep: FutureAndPointers) -> None: for src, row in rows.items(): stats.record(src, dst, row[0], row[1], seconds * row[col] / total) + @staticmethod + def _poll_read(ep: FutureAndPointers) -> tuple[bool, BaseException | None]: + """``(done, error)`` for a read. A read that finished cleanly drops its + future, so later polls of it (it can wait behind an earlier read of + its request) query nothing.""" + if ep.future is None: + return True, None + if not ep.future.done(): + return False, None + # exception(), not result(): a re-raise would put this frame on the + # stored traceback, and with it the read's buffers + error = CancelledError() if ep.future.cancelled() else ep.future.exception() + if error is None: + ep.future = None + return True, error + + def _drop_failed_read(self, ep: FutureAndPointers, error: BaseException | None) -> None: + """Retire a finished read that will not be delivered. + + ``error`` set: this read failed. Its transfers may still write into its + buffers, so they are kept until they report that they stopped, and the + producer is not acked; it reclaims its copy when the request is torn + down. ``error`` None: an earlier read of the request failed and this + one completed, so it is acked and its buffers go with the request. + """ + if error is None: + self._collect_and_send_acks(ep.request_id, ep.graph_edges) + else: + if ep.request_id not in self._failed_read_errors: + logger.error( + "Tensor read for request %s failed: %s: %s", + self.rid_to_str(ep.request_id), type(error).__name__, error, + ) + self._keep_read_buffers(ep, error) + error = self._failed_read_errors.setdefault(ep.request_id, error) + self._failed_reads.append(FailedRead(ep.request_id, ep.graph_edges, error)) + + def take_failed_reads(self) -> list[FailedRead]: + """Reads dropped by :meth:`get_ready_tensors` since the last call. The + caller fails their requests; their other reads are dropped as they + finish, until :meth:`cleanup_request`.""" + failed, self._failed_reads = self._failed_reads, [] + return failed + + def _keep_read_buffers(self, ep: FutureAndPointers, error: BaseException | None) -> None: + """Take a read's buffers out of the store, so the request's teardown + cannot free them, and keep them until its transfers stop. ``error`` is + how the read failed; None for one still running at teardown.""" + buffers, registered = [], [] + for uuid in ep.transfer_uuids: + if not self.tensor_store.check_uuid_presence(uuid): + continue + buffers.append(self.tensor_store.get_tensor(uuid)) + registered.append(self.tensor_store.is_registered(uuid)) + self.uuid_to_shard_dim.pop(uuid, None) + self.uuid_to_edge_name.pop(uuid, None) + self.tensor_store.remove_tensor(uuid) + if not buffers: + return + kept = _KeptRead( + self.rid_to_str(ep.request_id), buffers, registered, + sum(b.nbytes for b in buffers), + ) + logger.warning( + "Keeping %d bytes of read buffers for request %s until its " + "transfers report that they stopped", kept.nbytes, kept.request, + ) + if error is None and ep.future is not None and not ep.future.done(): + kept.future = ep.future + elif not self._take_over(kept, error): + return + self._kept_reads.append(kept) + + def _take_over(self, kept: _KeptRead, error: BaseException | None) -> bool: + """Adopt the batch ids a finished read handed over for ``kept``. False + if whether its transfers stopped cannot be learned: then its buffers + are never freed.""" + if error is None or isinstance(error, CancelledError): + return True # completed after all, or never started + if self._transfer_status is None or not isinstance(error, AsyncReadError): + why = f"its {type(error).__name__} does not say whether the transfer stopped" + return self._never_free(kept, why) + ptrs = {buffer.data_ptr() for buffer in kept.buffers} + if error.unknown & ptrs: + return self._never_free(kept, f"{error}, so part of the transfer may be running") + kept.running = {ptr: batch_id for ptr, batch_id in error.running.items() if ptr in ptrs} + return True + + def _never_free(self, kept: _KeptRead, why: str) -> bool: + self._unreleasable_buffers.extend(kept.buffers) + logger.warning( + "Read buffers of request %s (%d bytes) stay allocated for the life of " + "the process: %s", kept.request, kept.nbytes, why, + ) + return False + + def _release_kept_reads(self) -> None: + """Free each kept read whose transfers have all reported COMPLETED or + FAILED. Never blocks; any other status keeps the buffers.""" + still_kept = [] + for kept in self._kept_reads: + if kept.future is not None: + if not kept.future.done(): + still_kept.append(kept) + continue + future, kept.future = kept.future, None + error = CancelledError() if future.cancelled() else future.exception() + if not self._take_over(kept, error): + continue + for ptr, batch_id in list(kept.running.items()): + status = self._transfer_status(batch_id) + if status in (_TRANSFER_COMPLETED, _TRANSFER_FAILED): + # both free the batch, so the id is never polled again + del kept.running[ptr] + elif status not in _TRANSFER_RUNNING: + kept.running.clear() + self._never_free(kept, f"its transfer returned unknown status {status}") + break + else: + if not kept.running: + self._free_kept_read(kept) + else: + still_kept.append(kept) + self._kept_reads = still_kept + + def _free_kept_read(self, kept: _KeptRead) -> None: + for buffer, registered in zip(kept.buffers, kept.registered, strict=True): + if not registered or not buffer.nbytes: + continue + ret_value = self.transfer_engine.unregister_memory(buffer.data_ptr()) + if ret_value != 0: + self._never_free(kept, f"unregistering its buffers failed ({ret_value})") + return + logger.info( + "Freed %d bytes of read buffers kept for request %s: its transfers " + "stopped", kept.nbytes, kept.request, + ) + # the reader may still hold this read's future and its error + kept.buffers.clear() + def get_ready_tensors(self, graph_walk: str | None=None) -> dict[Rid, list[GraphEdge]]: + if self._kept_reads: + self._release_kept_reads() ready: dict[Rid, list[GraphEdge]] = {} still_pending = [] uuid_to_time = {} @@ -1432,10 +1732,11 @@ def get_ready_tensors(self, graph_walk: str | None=None) -> dict[Rid, list[Graph # must not reach the client ahead of the one before it waiting: set[Rid] = set() for ep in self.pending: - done = ep.future is None or ep.future.done() + done, error = self._poll_read(ep) + if error is not None or (done and ep.request_id in self._failed_read_errors): + self._drop_failed_read(ep, error) + continue if done and ep.request_id not in waiting: - if ep.future is not None: - ep.future.result() if stats is not None and ep.read_src is not None: self._record_transfer(stats, ep) for edge in ep.graph_edges: @@ -1683,7 +1984,22 @@ def cleanup_request(self, request_id: Rid, force: bool = False): referenced or persisted. Does NOT force-drop persisted signals — that is the conductor-coordinated hard cleanup, see :meth:`force_cleanup_request`. + + Reads of the request still in ``pending`` that failed, or have not + finished, keep their buffers and are not acked: their transfers may + still write into the buffers freed below. """ + leftover = [ep for ep in self.pending if ep.request_id == request_id] + self.pending = [ep for ep in self.pending if ep.request_id != request_id] + ack_edges: list[GraphEdge] = [] + for ep in leftover: + done, error = self._poll_read(ep) + if not done or error is not None: + self._keep_read_buffers(ep, error) + else: + ack_edges.extend(ep.graph_edges) + self._failed_read_errors.pop(request_id, None) + self._failed_reads = [f for f in self._failed_reads if f.request_id != request_id] self.read_finished.pop(request_id, None) self.buffered_shards.pop(request_id, None) self.sharding_configs.pop(request_id, None) @@ -1699,11 +2015,7 @@ def cleanup_request(self, request_id: Rid, force: bool = False): continue self._cleanup_by_uuid(uuid) - self._collect_and_send_acks( - request_id, - sum([ep.graph_edges for ep in self.pending if ep.request_id == request_id], start=[]), - ) - self.pending = [ep for ep in self.pending if ep.request_id != request_id] + self._collect_and_send_acks(request_id, ack_edges) self.slice_uuids.pop(request_id, None) self.slice_origins.pop(request_id, None) @@ -1761,7 +2073,8 @@ def __init__( # the consumer has the engine copy ``nbytes`` from ``address + offset`` # into a buffer it views with the descriptor's strides self.needs_flat_layout = protocol in (CommProtocol.RDMA, CommProtocol.TCP) - self._async_reader = AsyncMooncakeReader(engine._engine, device=device) + self._async_reader = TrackedMooncakeReader(engine._engine, device=device) + self._transfer_status = self._async_reader.transfer_status @property def transport_label(self) -> str: @@ -1855,12 +2168,14 @@ def start_read_tensors( ) read_info = [] + transfer_uuids = [] # Created before the reads and pointed at by each of them, so the # async reader can stamp ``rx_time`` on it when the read completes # (the reader runs on another thread; the object identity is # stable, only ``.future`` is filled in after submit). fp = FutureAndPointers( - future=None, graph_edges=[graph_edge], request_id=request_id + future=None, graph_edges=[graph_edge], request_id=request_id, + transfer_uuids=transfer_uuids, ) for info in graph_edge.tensor_info: if info.source_entity == self.my_entity_id: @@ -1902,6 +2217,7 @@ def start_read_tensors( remote_ptr=info.address + info.offset, nbytes=info.nbytes, )) + transfer_uuids.append(info.uuid) logger.debug("Started transfer read for uuid %s", info.uuid) if self.enable_prof or fp.read_src is not None: for ri in read_info: diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index f6ea14a4cf..164fea98ea 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -1460,6 +1460,13 @@ def _check_ready_tensors(self) -> None: """Poll for completed RDMA transfers, feed ready graph edges to worker graph queues.""" self.wakeup_event.drain() ready = self.tensor_manager.get_ready_tensors() + failed = { + f.request_id: f"Error in worker: input tensor read failed: {type(f.error).__name__}: {f.error}" + for f in self.tensor_manager.take_failed_reads() + if f.request_id not in self.scheduler.failed_rids + } + if failed: + self._fail_requests(failed) for request_id, edges in ready.items(): # Separate streaming edges from normal edges streaming = [e for e in edges if e.is_streaming] diff --git a/test/modular/test_api_completion_guard.py b/test/modular/test_api_completion_guard.py index e16321c007..74585585bc 100644 --- a/test/modular/test_api_completion_guard.py +++ b/test/modular/test_api_completion_guard.py @@ -61,6 +61,7 @@ def process_prompt(self, *args, **kwargs): force_cleanup_request=cleaned.append, has_inflight_reads=lambda rid: False, get_ready_tensors=lambda: {}, + take_failed_reads=lambda: [], ) wt.model = _RejectingModel() wt.device = "cpu" diff --git a/test/modular/test_api_result_delivery.py b/test/modular/test_api_result_delivery.py index d330290996..b98baeee7c 100644 --- a/test/modular/test_api_result_delivery.py +++ b/test/modular/test_api_result_delivery.py @@ -32,6 +32,11 @@ def start_read_tensors(self, request_id, graph_edges, graph_walk=None): self.read_started.append(request_id) return [] + def take_failed_reads(self): + + return [] + + def get_ready_tensors(self): return {} diff --git a/test/modular/test_data_worker_late_release.py b/test/modular/test_data_worker_late_release.py index 483fc65bf8..3b2ec5d552 100644 --- a/test/modular/test_data_worker_late_release.py +++ b/test/modular/test_data_worker_late_release.py @@ -90,6 +90,11 @@ def store_and_return_tensor_info(self, request_id, tensors): def register_for_send(self, request_id, tensor_infos): pass + def take_failed_reads(self): + + return [] + + def get_ready_tensors(self): # the loop polls this after its release queues and before in_queue hook, self.on_ready = getattr(self, "on_ready", None), None diff --git a/test/modular/test_data_worker_wakeup.py b/test/modular/test_data_worker_wakeup.py index 4b4c9e68ba..d772e0ecea 100644 --- a/test/modular/test_data_worker_wakeup.py +++ b/test/modular/test_data_worker_wakeup.py @@ -57,6 +57,11 @@ def __init__(self): def start_read_tensors(self, request_id, graph_edges, graph_walk=None): return list(self.futures) + def take_failed_reads(self): + + return [] + + def get_ready_tensors(self): return {} diff --git a/test/modular/test_empty_tensor_transport.py b/test/modular/test_empty_tensor_transport.py index 36dc650b77..900cec7ef5 100644 --- a/test/modular/test_empty_tensor_transport.py +++ b/test/modular/test_empty_tensor_transport.py @@ -59,6 +59,8 @@ def __init__(self): self.regions: dict[int, int] = {} self.calls: list[tuple] = [] self.gate: threading.Event | None = None + self._batch_ids = itertools.count(1) + self.batches: dict[int, int] = {} self.instances.append(self) def initialize(self, hostname, metadata_server, protocol, device): @@ -87,16 +89,22 @@ def unregister_memory(self, ptr): self.calls.append(("unregister", ptr, ret)) return ret - def batch_transfer_sync_read(self, session, local_ptrs, remote_ptrs, lengths): + def batch_transfer_async_read(self, session, local_ptrs, remote_ptrs, lengths): self.calls.append(("read", tuple(lengths))) if self.gate is not None: self.gate.wait(10) peer = self.sessions[session] + batch_id = next(self._batch_ids) if any(n and not peer._covers(r, n) for r, n in zip(remote_ptrs, lengths, strict=True)): - return -1 + self.batches[batch_id] = -1 + return batch_id for local, remote, n in zip(local_ptrs, remote_ptrs, lengths, strict=True): ctypes.memmove(local, remote, n) - return 0 + self.batches[batch_id] = 1 + return batch_id + + def transfer_check_status(self, batch_id): + return self.batches.pop(batch_id) @pytest.fixture diff --git a/test/modular/test_failed_tensor_read.py b/test/modular/test_failed_tensor_read.py new file mode 100644 index 0000000000..bb5c856486 --- /dev/null +++ b/test/modular/test_failed_tensor_read.py @@ -0,0 +1,589 @@ +"""A failed async tensor read (a Mooncake transfer error) fails only its own +request, and its buffers are freed once Mooncake reports that every transfer +into them stopped. + +Before, ``get_ready_tensors`` re-raised the same failed future on every poll, +so no other request on that worker or data worker received its inputs until +the failed request was torn down (#418).""" + +import collections +import ctypes +import itertools +import json +import threading +from concurrent.futures import Future +from types import SimpleNamespace + +import pytest +import torch +from test_request_failure_propagation import _preprocess_thread, _worker + +from mstar.api_server.request_types import ResultChunk +from mstar.communication.communicator import BaseCommunicator, CommProtocol +from mstar.communication.tensors import ( + MOONCAKE_IMPORT_ERROR, + AsyncMooncakeReader, + AsyncReadError, + FutureAndPointers, + MooncakeCommunicationManager, + SharedMemoryCommunicationManager, + TensorCommunicationManager, + TrackedMooncakeReader, + TransferReadInfo, + create_tensor_communication_manager, +) +from mstar.graph.base import GraphEdge, TensorPointerInfo +from mstar.utils.ipc_format import ConductorMessageType, WorkerMessageType + + +class _Communicator(BaseCommunicator): + def __init__(self): + self.sent = [] + + def send(self, entity_id, msg): + self.sent.append((entity_id, msg)) + + def get_all_new_messages(self): + return [] + + +def _acked(comm): + return [ + (msg.body.request_id, dict(msg.body.successful_tensors)) + for _, msg in comm.sent + if msg.message_type == WorkerMessageType.TENSOR_RECEIVED + ] + + +# ── the read fix: one failed read fails one request ─────────────────────── + + +def _shm_manager(tmp_path): + comm = _Communicator() + mgr = SharedMemoryCommunicationManager("worker_1", "host", "cpu", comm, shm_dir=str(tmp_path)) + return mgr, comm + + +def _edge(*uuids, name="x"): + edge = GraphEdge(next_node="node", name=name) + edge.tensor_info = [SimpleNamespace(uuid=u, source_entity="worker_0") for u in uuids] + return edge + + +def _failed_future(message="Mooncake read failed. Status: -1"): + f = Future() + f.set_exception(RuntimeError(message)) + return f + + +def _read(rid, edge, future=None): + return FutureAndPointers(future=future, graph_edges=[edge], request_id=rid) + + +def test_a_failed_read_does_not_block_other_requests(tmp_path): + mgr, comm = _shm_manager(tmp_path) + bad, healthy = _edge(11), _edge() + mgr.pending = [_read(7, bad, _failed_future()), _read(8, healthy)] + + assert mgr.get_ready_tensors() == {8: [healthy]} + assert mgr.pending == [] + (failed,) = mgr.take_failed_reads() + assert failed.request_id == 7 and failed.graph_edges == [bad] + assert "Mooncake read failed" in str(failed.error) + # the transfer may not have stopped, so the producer keeps its copy + assert _acked(comm) == [] + + assert mgr.get_ready_tensors() == {} + assert mgr.take_failed_reads() == [] + + +def test_a_cancelled_read_counts_as_failed(tmp_path): + mgr, _ = _shm_manager(tmp_path) + cancelled = Future() + cancelled.cancel() + mgr.pending = [_read(7, _edge(11), cancelled)] + + assert mgr.get_ready_tensors() == {} + (failed,) = mgr.take_failed_reads() + assert failed.request_id == 7 and mgr.pending == [] + + +def test_later_reads_of_a_failed_request_are_dropped_until_cleanup(tmp_path): + mgr, comm = _shm_manager(tmp_path) + in_flight = Future() + mgr.pending = [ + _read(7, _edge(11), _failed_future()), + _read(7, _edge(12), in_flight), + _read(7, _edge(13)), + ] + + assert mgr.get_ready_tensors() == {} + first, later = mgr.take_failed_reads() + assert first.error is later.error + # a read still touching the producer keeps the drain waiting + assert [ep.future for ep in mgr.pending] == [in_flight] + assert mgr.has_inflight_reads(7) + + in_flight.set_result(None) + assert mgr.get_ready_tensors() == {} + (dropped,) = mgr.take_failed_reads() + assert dropped.error is first.error + assert mgr.pending == [] and not mgr.has_inflight_reads(7) + # the reads that completed are acked; the failed one is not + assert _acked(comm) == [("7", {13: 1}), ("7", {12: 1})] + + mgr.cleanup_request(7) + fresh = _edge() + mgr.pending = [_read(7, fresh)] + assert mgr.get_ready_tensors() == {7: [fresh]} + assert mgr.take_failed_reads() == [] + + +class _CountingFuture(Future): + """A read future that counts the queries made on it.""" + + def __init__(self): + super().__init__() + self.done_calls = 0 + self.exception_calls = 0 + + def done(self): + self.done_calls += 1 + return super().done() + + def exception(self, timeout=None): + self.exception_calls += 1 + return super().exception(timeout) + + +def test_a_finished_read_is_queried_once_while_it_waits_behind_an_earlier_one(tmp_path): + mgr, _ = _shm_manager(tmp_path) + earlier, later = _CountingFuture(), _CountingFuture() + later.set_result(None) + first, second = _edge(), _edge() + mgr.pending = [_read(7, first, earlier), _read(7, second, later)] + + for _ in range(3): + assert mgr.get_ready_tensors() == {} + assert (earlier.done_calls, earlier.exception_calls) == (3, 0) + assert (later.done_calls, later.exception_calls) == (1, 1) + + earlier.set_result(None) + assert mgr.get_ready_tensors() == {7: [first, second]} + assert (later.done_calls, later.exception_calls) == (1, 1) + + +def test_the_worker_fails_a_request_once_when_its_input_read_fails(tmp_path): + w = _worker() + w.enable_nvtx = False + w.wakeup_event = SimpleNamespace(drain=lambda: None) + w.tensor_manager, _ = _shm_manager(tmp_path) + w.tensor_manager.pending = [_read("r1", _edge(11), _failed_future())] + + w._check_ready_tensors() + w._check_ready_tensors() + + (sent,) = w.sent + assert sent[1].message_type == ConductorMessageType.FAIL_REQUESTS + assert list(sent[1].body.errors) == ["r1"] + assert "input tensor read failed: RuntimeError: Mooncake read failed" in sent[1].body.errors["r1"] + assert w.scheduler.failed_rids == {"r1"} + + +def test_the_data_worker_answers_a_failed_output_read_in_its_slot(tmp_path): + """The read of output 0 failed while output 1 waits in the reorder buffer: + the error takes slot 0, so 1 follows it instead of waiting for the TTL, + and the API server gets one chunk per tensor it counted.""" + wt = _preprocess_thread(model=None) + held = ResultChunk(request_id="r1", modality="text", data=b"later", metadata={}) + wt.request_output_state["r1"].pending[1] = held + wt.request_output_state["r1"].next_sequence = 2 + wt.tensor_manager, _ = _shm_manager(tmp_path) + wt.tensor_manager.pending = [_read("r1", _edge("u1", name="text_output"), _failed_future())] + + assert wt._process_read_tensors() is True + chunks = [wt.out_queue.get_nowait(), wt.out_queue.get_nowait()] + assert [c.modality for c in chunks] == ["error", "text"] + assert b"output tensor read failed: RuntimeError: Mooncake read failed" in chunks[0].data + assert chunks[0].metadata["status"] == 500 + assert wt.out_queue.empty() + assert wt.request_output_state["r1"].order == {} + assert wt.tensor_uuid_to_metadata_per_request["r1"] == {} + + assert wt._process_read_tensors() is False + assert wt.out_queue.empty() + + +# ── releasing a failed read's buffers ───────────────────────────────────── + + +class _Engine: + """Mooncake's async read API, one task per batch, with scripted statuses. + + ``plan[remote_ptr]`` lists, per submit of that tensor, the statuses its + batch reports to successive ``transfer_check_status`` calls; the last one + repeats. ``"submit-fails"`` makes that submit return ``submit_failure``. + COMPLETED copies the bytes; COMPLETED and FAILED free the batch, as in + Mooncake, and a freed batch must never be polled again.""" + + def __init__(self, hcas=0): + self.plan: dict[int, list] = {} + self.hcas = hcas + self.submit_failure = 0 + self.unregister_result = 0 + self.unregistered: list[int] = [] + self.polls: list[tuple[int, int, str]] = [] # (batch, remote, thread) + self.after_free: list[int] = [] + self._batches: dict[int, list] = {} + self._freed: set[int] = set() + self._submits = collections.Counter() + self._ids = itertools.count(1) + self._lock = threading.Lock() + + def get_local_topology(self): + return json.dumps({"cpu:0": [[f"mlx5_{i}" for i in range(self.hcas)], []]}) + + def register_memory(self, ptr, nbytes): + return 0 + + def unregister_memory(self, ptr): + self.unregistered.append(ptr) + return self.unregister_result + + def batch_transfer_async_read(self, session, local, remote, lengths): + (local,), (remote,), (nbytes,) = local, remote, lengths + with self._lock: + attempts = self.plan[remote] + statuses = attempts[min(self._submits[remote], len(attempts) - 1)] + self._submits[remote] += 1 + if statuses == "submit-fails": + return self.submit_failure + batch_id = next(self._ids) + self._batches[batch_id] = [local, remote, nbytes, list(statuses)] + return batch_id + + def transfer_check_status(self, batch_id): + with self._lock: + if batch_id in self._freed: + self.after_free.append(batch_id) + return 0 + local, remote, nbytes, statuses = self._batches[batch_id] + self.polls.append((batch_id, remote, threading.current_thread().name)) + status = statuses.pop(0) if len(statuses) > 1 else statuses[0] + if status == 1: + ctypes.memmove(local, remote, nbytes) + if status in (1, -1): + self._freed.add(batch_id) + return status + + def report(self, remote, status): + """From now on, the live batch reading ``remote`` reports ``status``.""" + with self._lock: + for batch_id, batch in self._batches.items(): + if batch[1] == remote and batch_id not in self._freed: + batch[3] = [status] + + def submits(self, remote): + return self._submits[remote] + + def polled_by(self, remote): + return {thread for _, r, thread in self.polls if r == remote} + + +@pytest.fixture +def engine(): + engine = _Engine() + managers = [] + + def make(): + mgr = MooncakeCommunicationManager.__new__(MooncakeCommunicationManager) + TensorCommunicationManager.__init__(mgr, "receiver", "receiver", "cpu", _Communicator(), engine) + mgr.protocol = CommProtocol.RDMA + mgr._async_reader = TrackedMooncakeReader(engine, "cpu") + mgr._transfer_status = mgr._async_reader.transfer_status + managers.append(mgr) + return mgr + + engine.manager = make + yield engine + # let any reader still polling a running batch finish + for remote in list(engine.plan): + engine.report(remote, 1) + for mgr in managers: + mgr._async_reader._executor.shutdown(wait=True) + assert engine.after_free == [], "a freed batch was polled" + + +def _sources(n, nbytes=16): + return [torch.full((nbytes,), i + 1, dtype=torch.uint8) for i in range(n)] + + +def _remote_edge(sources, first_uuid=100): + infos = [ + TensorPointerInfo( + list(src.shape), src.dtype, src.nbytes, src.data_ptr(), + list(src.stride()), first_uuid + i, "producer:1", "producer", + ) + for i, src in enumerate(sources) + ] + return GraphEdge(next_node="node", name="x", tensor_info=infos) + + +def _start(mgr, rid, sources): + edge = _remote_edge(sources) + (future,) = mgr.start_read_tensors(rid, [edge]) + buffers = [mgr.get_tensor(info.uuid) for info in edge.tensor_info] + return future, [b.data_ptr() for b in buffers] + + +def _kept(mgr): + return [b.data_ptr() for kept in mgr._kept_reads for b in kept.buffers] + + +def _never_freed(mgr): + return [b.data_ptr() for b in mgr._unreleasable_buffers] + + +@pytest.mark.parametrize("final", [1, -1], ids=["completed", "failed"]) +def test_a_failed_reads_buffers_are_freed_once_its_last_transfer_stops(engine, final): + """Tensor 0 fails while tensor 1 is still moving. Its buffers stay through + teardown while tensor 1 reports WAITING (0) or TIMEOUT (-2), and are freed + once it reports COMPLETED or FAILED.""" + mgr = engine.manager() + src = _sources(2) + engine.plan = {src[0].data_ptr(): [[0, 0, -1]], src[1].data_ptr(): [[0]]} + future, ptrs = _start(mgr, "bad", src) + + error = future.exception(timeout=5) + assert isinstance(error, AsyncReadError) + assert list(error.running) == [ptrs[1]] and error.unknown == set() + assert mgr.get_ready_tensors() == {} + assert [f.request_id for f in mgr.take_failed_reads()] == ["bad"] + mgr.force_cleanup_request("bad") + assert not mgr.tensor_store.get_all_uuids("bad") + + for status in (0, -2, 0): + engine.report(src[1].data_ptr(), status) + mgr.get_ready_tensors() + assert _kept(mgr) == ptrs and engine.unregistered == [] + + engine.report(src[1].data_ptr(), final) + mgr.get_ready_tensors() + assert mgr._kept_reads == [] and _never_freed(mgr) == [] + assert sorted(engine.unregistered) == sorted(ptrs) + assert _acked(mgr.communicator) == [] + # the reader polled only tensor 0; tensor 1 was polled here, after hand-over + assert engine.polled_by(src[0].data_ptr()) != {threading.current_thread().name} + assert engine.polled_by(src[1].data_ptr()) == {threading.current_thread().name} + + +@pytest.mark.parametrize("first", [1, -1], ids=["completes", "fails"]) +def test_a_read_torn_down_mid_transfer_is_freed_once_it_stops(engine, first): + """Teardown finds the read running. The reader owns its batches until it + finishes; the manager never polls them before that, and frees the buffers + once every transfer stopped, whether the read completed or failed.""" + mgr = engine.manager() + src = _sources(2) + engine.plan = {s.data_ptr(): [[0]] for s in src} + future, ptrs = _start(mgr, "r", src) + + mgr.force_cleanup_request("r") + assert _kept(mgr) == ptrs and _acked(mgr.communicator) == [] + mgr.get_ready_tensors() + assert _kept(mgr) == ptrs + + engine.report(src[0].data_ptr(), first) + if first == -1: + assert list(future.exception(timeout=5).running) == [ptrs[1]] + mgr.get_ready_tensors() + assert _kept(mgr) == ptrs, "tensor 1 is still moving" + engine.report(src[1].data_ptr(), 1) + if first == 1: + assert future.exception(timeout=5) is None + mgr.get_ready_tensors() + assert mgr._kept_reads == [] and sorted(engine.unregistered) == sorted(ptrs) + me = threading.current_thread().name + assert me not in engine.polled_by(src[0].data_ptr()) + if first == 1: + assert me not in engine.polled_by(src[1].data_ptr()) + + +@pytest.mark.parametrize("failure", [0, 2**64 - 1], ids=["zero", "legacy-uint64-minus-1"]) +@pytest.mark.parametrize("when", ["first-submit", "retry-submit"]) +def test_a_failed_submit_keeps_the_buffers_for_good(engine, failure, when): + """A failed submit may have posted part of its transfer and gives no batch + id to ask about, so nothing of the read is freed, and nothing is polled + for it after the read fails.""" + engine.hcas = 1 # two attempts per tensor + engine.submit_failure = failure + mgr = engine.manager() + src = _sources(2) + second = [[0, -1], "submit-fails"] if when == "retry-submit" else ["submit-fails"] + engine.plan = {src[0].data_ptr(): [[0, 0, 1]], src[1].data_ptr(): second} + future, ptrs = _start(mgr, "bad", src) + + error = future.exception(timeout=5) + assert isinstance(error, AsyncReadError) and "submit failed" in str(error) + assert error.unknown == {ptrs[1]} + mgr.get_ready_tensors() + mgr.force_cleanup_request("bad") + for _ in range(3): + mgr.get_ready_tensors() + assert mgr._kept_reads == [] and _never_freed(mgr) == ptrs + assert engine.unregistered == [] + assert threading.current_thread().name not in engine.polled_by(src[0].data_ptr()) + + +def test_an_unknown_status_keeps_the_buffers_for_good(engine): + """Only COMPLETED and FAILED say a transfer stopped. Any other status may + or may not have freed its batch: the buffers stay and it is never polled + again, in the reader or in the manager.""" + mgr = engine.manager() + src = _sources(3) + engine.plan = { + src[0].data_ptr(): [[-1]], src[1].data_ptr(): [[0]], src[2].data_ptr(): [[0]], + } + future, ptrs = _start(mgr, "a", src) + assert set(future.exception(timeout=5).running) == {ptrs[1], ptrs[2]} + mgr.get_ready_tensors() + engine.report(src[1].data_ptr(), 7) + mgr.get_ready_tensors() + mgr.get_ready_tensors() + assert mgr._kept_reads == [] and _never_freed(mgr) == ptrs + assert [r for _, r, _ in engine.polls].count(src[1].data_ptr()) == 1 + assert engine.unregistered == [] + + reader_src = _sources(1, nbytes=8) + engine.plan[reader_src[0].data_ptr()] = [[0, 9]] + future, reader_ptrs = _start(mgr, "b", reader_src) + error = future.exception(timeout=5) + assert "unknown status 9" in str(error) and error.unknown == set(reader_ptrs) + mgr.get_ready_tensors() + assert _never_freed(mgr) == ptrs + reader_ptrs + + +def test_buffers_whose_unregister_fails_are_kept(engine): + engine.unregister_result = -3 + mgr = engine.manager() + src = _sources(1) + engine.plan = {src[0].data_ptr(): [[-1]]} + future, ptrs = _start(mgr, "a", src) + future.exception(timeout=5) + mgr.get_ready_tensors() + mgr.get_ready_tensors() + assert engine.unregistered == ptrs + assert mgr._kept_reads == [] and _never_freed(mgr) == ptrs + + +@pytest.mark.parametrize("hcas, outcome", [(1, "read"), (0, "failed")]) +def test_a_failed_tensor_is_retried_as_often_as_the_sync_read_retries(engine, hcas, outcome): + engine.hcas = hcas + mgr = engine.manager() + src = _sources(1) + engine.plan = {src[0].data_ptr(): [[-1], [1]]} + future, _ = _start(mgr, "a", src) + error = future.exception(timeout=5) + assert engine.submits(src[0].data_ptr()) == hcas + 1 + if outcome == "read": + assert error is None + (edge,) = mgr.get_ready_tensors()["a"] + assert torch.equal(mgr.get_tensor(edge.tensor_info[0].uuid), src[0]) + else: + assert "Status: -1" in str(error) + + +def test_a_read_past_its_deadline_hands_over_its_running_transfers(engine): + mgr = engine.manager() + mgr._async_reader._timeout_s = 0 + src = _sources(1) + engine.plan = {src[0].data_ptr(): [[0]]} + future, ptrs = _start(mgr, "slow", src) + error = future.exception(timeout=5) + assert "timed out" in str(error) and list(error.running) == ptrs + + mgr.get_ready_tensors() + mgr.force_cleanup_request("slow") + mgr.get_ready_tensors() + assert _kept(mgr) == ptrs + engine.report(src[0].data_ptr(), 1) + mgr.get_ready_tensors() + assert mgr._kept_reads == [] and engine.unregistered == ptrs + + +def test_kv_reads_keep_the_batched_sync_read(): + calls = [] + engine = SimpleNamespace( + batch_transfer_sync_read=lambda session, local, remote, n: calls.append(len(local)) or 0, + ) + reader = AsyncMooncakeReader(engine, "cpu") + try: + infos = [TransferReadInfo("peer", 1000 + i, 2000 + i, 1) for i in range(3)] + assert reader.submit(infos).result(timeout=5) is None + assert calls == [3] + finally: + reader.shutdown() + + +# ── real Mooncake TCP ───────────────────────────────────────────────────── + + +class _Recording: + """Records what ``transfer_check_status`` returned.""" + + def __init__(self, engine): + self._engine = engine + self.statuses = [] + + def transfer_check_status(self, batch_id): + status = self._engine.transfer_check_status(batch_id) + self.statuses.append(status) + return status + + def __getattr__(self, name): + return getattr(self._engine, name) + + +@pytest.mark.skipif(MOONCAKE_IMPORT_ERROR is not None, reason=f"mooncake unavailable: {MOONCAKE_IMPORT_ERROR}") +def test_real_mooncake_tcp_frees_a_failed_reads_buffer(): + """A read of producer memory that was never registered: Mooncake's TCP + transport rejects it, ``transfer_check_status`` reports FAILED, the + request fails alone, and its buffer is unregistered and freed at the + next poll. A good read on the same pair still delivers its bytes.""" + def make(entity): + return create_tensor_communication_manager( + protocol=CommProtocol.TCP, my_entity_id=entity, hostname="localhost", + device="cpu", communicator=_Communicator(), tcp_transfer_device="lo", + ) + + producer, consumer = make("producer"), make("consumer") + recording = _Recording(consumer._async_reader._engine) + consumer._async_reader._engine = recording + unregistered = [] + unregister = consumer.transfer_engine.unregister_memory + consumer.transfer_engine.unregister_memory = lambda ptr: unregistered.append(ptr) or unregister(ptr) + try: + good, bad = torch.arange(4096, dtype=torch.int32), torch.ones(4096, dtype=torch.int32) + good_edge = GraphEdge(next_node="n", name="x") + good_edge.tensor_info = producer.store_and_return_tensor_info("good", {"x": [good]})["x"] + producer.register_for_send("good", good_edge.tensor_info) + bad_edge = GraphEdge(next_node="n", name="x") + bad_edge.tensor_info = producer.store_and_return_tensor_info("bad", {"x": [bad]})["x"] + + (bad_future,) = consumer.start_read_tensors("bad", [bad_edge]) + buffer_ptr = consumer.get_tensor(bad_edge.tensor_info[0].uuid).data_ptr() + (good_future,) = consumer.start_read_tensors("good", [good_edge]) + assert bad_future.exception(timeout=30) is not None + assert good_future.exception(timeout=30) is None + assert -1 in recording.statuses + + ready = consumer.get_ready_tensors() + assert list(ready) == ["good"] + assert torch.equal(consumer.get_tensor(ready["good"][0].tensor_info[0].uuid), good) + assert [f.request_id for f in consumer.take_failed_reads()] == ["bad"] + consumer.force_cleanup_request("bad") + consumer.get_ready_tensors() + assert buffer_ptr in unregistered + assert consumer._kept_reads == [] and consumer._unreleasable_buffers == [] + finally: + for mgr in (producer, consumer): + mgr._async_reader._executor.shutdown(wait=True) diff --git a/test/modular/test_inline_tokens.py b/test/modular/test_inline_tokens.py index 388aabf95e..e44f999cba 100644 --- a/test/modular/test_inline_tokens.py +++ b/test/modular/test_inline_tokens.py @@ -63,6 +63,11 @@ class _TM: def start_read_tensors(self, request_id, graph_edges, graph_walk=None): return [] + def take_failed_reads(self): + + return [] + + def get_ready_tensors(self): return {} diff --git a/test/modular/test_mooncake_flat_layout.py b/test/modular/test_mooncake_flat_layout.py index 8f7ba248cf..75f2b3e3d6 100644 --- a/test/modular/test_mooncake_flat_layout.py +++ b/test/modular/test_mooncake_flat_layout.py @@ -65,11 +65,14 @@ def submit(self, reads): for r in reads: ctypes.memmove(r.local_ptr, r.remote_ptr, r.nbytes) + def transfer_status(self, batch_id): + raise AssertionError("no read fails here, so no transfer is handed over") + @pytest.fixture(autouse=True) def _no_transfer_engine(monkeypatch): monkeypatch.setattr(tensors, "MooncakeTransferEngine", _NoEngine) - monkeypatch.setattr(tensors, "AsyncMooncakeReader", _FlatCopyReader) + monkeypatch.setattr(tensors, "TrackedMooncakeReader", _FlatCopyReader) def _mooncake(entity, bookkeeping="python"): diff --git a/test/modular/test_process_liveness.py b/test/modular/test_process_liveness.py index a7ab5df44e..5930346d32 100644 --- a/test/modular/test_process_liveness.py +++ b/test/modular/test_process_liveness.py @@ -300,6 +300,7 @@ def test_stopping_data_worker_drops_tracked_requests(): force_cleanup_request=cleaned.append, has_inflight_reads=lambda rid: False, get_ready_tensors=lambda: {}, + take_failed_reads=lambda: [], ) wt.in_flight_requests = {"r1", "r2"} wt.tensor_uuid_to_metadata_per_request = {"r2": {}} diff --git a/test/modular/test_request_failure_propagation.py b/test/modular/test_request_failure_propagation.py index fcf773d35e..044a662dc4 100644 --- a/test/modular/test_request_failure_propagation.py +++ b/test/modular/test_request_failure_propagation.py @@ -291,6 +291,7 @@ def postprocess(self, tensor, modality, request_kwargs=None): ) wt.tensor_manager = SimpleNamespace( get_ready_tensors=lambda: {"r1": [edge]}, + take_failed_reads=lambda: [], get_tensor=lambda uuid: object(), dereference=dereferenced.append, ) @@ -323,6 +324,7 @@ def postprocess(self, tensor, modality, request_kwargs=None): ) wt.tensor_manager = SimpleNamespace( get_ready_tensors=lambda: {"r1": [edge]}, + take_failed_reads=lambda: [], get_tensor=lambda uuid: object(), dereference=lambda uuid: None, ) diff --git a/test/modular/test_request_id_reuse.py b/test/modular/test_request_id_reuse.py index 2b11e1f016..83901a1fe1 100644 --- a/test/modular/test_request_id_reuse.py +++ b/test/modular/test_request_id_reuse.py @@ -95,6 +95,11 @@ def store_and_return_tensor_info(self, request_id, tensors): def register_for_send(self, request_id, tensor_infos): pass + def take_failed_reads(self): + + return [] + + def get_ready_tensors(self): return {} diff --git a/test/modular/test_video_frame_protocol.py b/test/modular/test_video_frame_protocol.py index 3d65fe9df2..a715e724e4 100644 --- a/test/modular/test_video_frame_protocol.py +++ b/test/modular/test_video_frame_protocol.py @@ -205,6 +205,11 @@ def __init__(self, tensors_by_request): self.dereferenced = [] self.cleaned = [] + def take_failed_reads(self): + + return [] + + def get_ready_tensors(self): ready, self.ready = self.ready, {} return { From b9622ba49f7f0ee84d74b1f3017a9ab11a8ded32 Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Mon, 19 Oct 2026 18:44:51 +0000 Subject: [PATCH 02/15] Hold READS_DONE while a failed read's transfers run, and read a torn-down read's outcome from its future --- mstar/api_server/data_worker.py | 7 +-- mstar/communication/tensors.py | 27 ++++++++-- test/modular/test_failed_tensor_read.py | 70 +++++++++++++++++++++++++ test/modular/test_shm_tensor_comm.py | 5 +- test/modular/test_worker_drain.py | 6 +-- 5 files changed, 104 insertions(+), 11 deletions(-) diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index 113066fcca..50308cffe7 100644 --- a/mstar/api_server/data_worker.py +++ b/mstar/api_server/data_worker.py @@ -1203,12 +1203,13 @@ def _send_enable_metrics_if_requested(self) -> None: def _finish_reading(self, request_id: str, drained: bool) -> None: """The API server is done with this rid's outputs. ``drained`` (every chunk delivered) means no read can still be in flight, so ACK straight - away — the happy path pays nothing. Otherwise delivery was abandoned - (TTL, client gone) with reads possibly still running, so gate the ACK on + away, unless a failed read's transfers are still running: its error + chunks count as delivered. Otherwise delivery was abandoned (TTL, + client gone) with reads possibly still running, so gate the ACK on them; ACKing there would let the conductor unlink under a read.""" if request_id not in self._live_requests: return - if drained: + if drained and not self.tensor_manager.has_inflight_reads(request_id): self._send_reads_done(request_id) else: self._begin_drain(request_id) diff --git a/mstar/communication/tensors.py b/mstar/communication/tensors.py index f9619c9b19..bb35c6dabc 100644 --- a/mstar/communication/tensors.py +++ b/mstar/communication/tensors.py @@ -129,7 +129,9 @@ class _KeptRead: ``future`` is set while the read still runs in the reader thread, which owns its batch ids until it finishes. ``running`` holds the ids handed over after that (local pointer -> batch id), polled by the manager. - ``registered`` is index-parallel to ``buffers``.""" + ``registered`` is index-parallel to ``buffers``. ``rid`` is None once the + request is torn down; ``request`` names it in logs.""" + rid: Rid | None request: str buffers: list[torch.Tensor] registered: list[bool] @@ -294,6 +296,14 @@ class TransferReadInfo: _READ_POLL_INTERVAL_S = 1e-4 +def _may_be_reading(future: Future) -> bool: + """A read still running, or one that failed and so may have left + transfers running until :meth:`get_ready_tensors` takes them over.""" + if not future.done(): + return True + return not future.cancelled() and future.exception() is not None + + def _submit_failed(batch_id: int) -> bool: """A failed ``batch_transfer_async_read``: 0, or uint64 -1 from Mooncake releases before 0.3.13.""" @@ -1640,14 +1650,15 @@ def _keep_read_buffers(self, ep: FutureAndPointers, error: BaseException | None) if not buffers: return kept = _KeptRead( - self.rid_to_str(ep.request_id), buffers, registered, + ep.request_id, self.rid_to_str(ep.request_id), buffers, registered, sum(b.nbytes for b in buffers), ) logger.warning( "Keeping %d bytes of read buffers for request %s until its " "transfers report that they stopped", kept.nbytes, kept.request, ) - if error is None and ep.future is not None and not ep.future.done(): + if error is None and ep.future is not None: + # unfinished when polled: how it ends is read from the future later kept.future = ep.future elif not self._take_over(kept, error): return @@ -1970,11 +1981,16 @@ def has_inflight_reads(self, request_id: Rid) -> bool: segment. Synchronous (SHM) reads complete inside ``start_read_tensors`` with ``future=None``, so they never count here; only outstanding async (RDMA) futures do. Gates the READS_DONE ACK during a drain. + So does a failed read whose transfers have not all stopped: they may + still be reading the producer's memory. """ return any( ep.request_id == request_id - and ep.future is not None and not ep.future.done() + and ep.future is not None and _may_be_reading(ep.future) for ep in self.pending + ) or any( + kept.rid == request_id and (kept.future is not None or kept.running) + for kept in self._kept_reads ) def cleanup_request(self, request_id: Rid, force: bool = False): @@ -1998,6 +2014,9 @@ def cleanup_request(self, request_id: Rid, force: bool = False): self._keep_read_buffers(ep, error) else: ack_edges.extend(ep.graph_edges) + for kept in self._kept_reads: + if kept.rid == request_id: + kept.rid = None self._failed_read_errors.pop(request_id, None) self._failed_reads = [f for f in self._failed_reads if f.request_id != request_id] self.read_finished.pop(request_id, None) diff --git a/test/modular/test_failed_tensor_read.py b/test/modular/test_failed_tensor_read.py index bb5c856486..b0fdd55e8e 100644 --- a/test/modular/test_failed_tensor_read.py +++ b/test/modular/test_failed_tensor_read.py @@ -408,6 +408,76 @@ def test_a_read_torn_down_mid_transfer_is_freed_once_it_stops(engine, first): assert me not in engine.polled_by(src[1].data_ptr()) +def test_a_read_that_fails_while_teardown_polls_it_keeps_its_buffers(engine, monkeypatch): + """Teardown polls the read just before the reader fails and hands over a + running transfer: the buffers wait for that transfer, not for the poll.""" + mgr = engine.manager() + src = _sources(2) + engine.plan = {src[0].data_ptr(): [[-1]], src[1].data_ptr(): [[0]]} + future, ptrs = _start(mgr, "r", src) + assert list(future.exception(timeout=5).running) == [ptrs[1]] + poll = type(mgr)._poll_read + monkeypatch.setattr(mgr, "_poll_read", lambda ep: (False, None)) + mgr.force_cleanup_request("r") + monkeypatch.setattr(mgr, "_poll_read", poll) + + mgr.get_ready_tensors() + assert _kept(mgr) == ptrs and engine.unregistered == [] + engine.report(src[1].data_ptr(), 1) + mgr.get_ready_tensors() + assert mgr._kept_reads == [] and sorted(engine.unregistered) == sorted(ptrs) + + +def test_the_drain_waits_for_a_failed_reads_running_transfers(engine): + """READS_DONE lets the producer free what the read is reading, so a + failed read counts as in flight until its last transfer stops, and not + after its request is torn down.""" + mgr = engine.manager() + src = _sources(2) + engine.plan = {src[0].data_ptr(): [[0, -1]], src[1].data_ptr(): [[0]]} + future, _ = _start(mgr, "r", src) + assert mgr.has_inflight_reads("r") + future.exception(timeout=5) + assert mgr.has_inflight_reads("r"), "the reader is done, the transfer is not" + mgr.get_ready_tensors() + assert mgr.take_failed_reads() and mgr.has_inflight_reads("r") + engine.report(src[1].data_ptr(), 1) + mgr.get_ready_tensors() + assert not mgr.has_inflight_reads("r") + + +def test_a_torn_down_requests_kept_read_no_longer_holds_a_drain(engine): + mgr = engine.manager() + src = _sources(2) + engine.plan = {src[0].data_ptr(): [[-1]], src[1].data_ptr(): [[0]]} + future, ptrs = _start(mgr, "r", src) + future.exception(timeout=5) + mgr.get_ready_tensors() + assert mgr.has_inflight_reads("r") + mgr.force_cleanup_request("r") + assert not mgr.has_inflight_reads("r") and _kept(mgr) == ptrs + + +def test_the_data_worker_holds_reads_done_for_a_failed_reads_running_transfer(tmp_path): + """The failed read's error chunks count as delivered, so the API server + reports the request drained; READS_DONE still waits for the transfer.""" + from mstar.utils.ipc_format import ConductorMessageType as T + + wt = _preprocess_thread(model=None) + wt._live_requests = {"r1": 1} + wt._draining_rids, wt._reads_done_sent = set(), set() + wt.sent = [] + wt.communicator = SimpleNamespace(send=lambda entity, msg: wt.sent.append(msg)) + running = {"r1"} + wt.tensor_manager = SimpleNamespace(has_inflight_reads=lambda rid: rid in running) + + wt._finish_reading("r1", drained=True) + assert [m.message_type for m in wt.sent] == [] and "r1" in wt._draining_rids + running.clear() + wt._complete_drain_if_ready("r1") + assert [m.message_type for m in wt.sent] == [T.READS_DONE] + + @pytest.mark.parametrize("failure", [0, 2**64 - 1], ids=["zero", "legacy-uint64-minus-1"]) @pytest.mark.parametrize("when", ["first-submit", "retry-submit"]) def test_a_failed_submit_keeps_the_buffers_for_good(engine, failure, when): diff --git a/test/modular/test_shm_tensor_comm.py b/test/modular/test_shm_tensor_comm.py index c7ad2b3f71..0d77da9733 100644 --- a/test/modular/test_shm_tensor_comm.py +++ b/test/modular/test_shm_tensor_comm.py @@ -611,7 +611,10 @@ def __init__(self, done: bool): def done(self): return self.finished - def result(self): + def cancelled(self): + return False + + def exception(self): return None def _read(future, rid: str, name: str) -> FutureAndPointers: diff --git a/test/modular/test_worker_drain.py b/test/modular/test_worker_drain.py index 998fee0b0c..83ef8a8b9b 100644 --- a/test/modular/test_worker_drain.py +++ b/test/modular/test_worker_drain.py @@ -497,9 +497,9 @@ def test_preprocess_reads_done_idempotent(): def test_preprocess_finished_reading_acks_immediately_when_drained(): - """Happy path: delivery completed, so the ACK goes out without paying for a - drain check (and without gating on reads that cannot exist).""" - wt = _preprocess(inflight_reads=True) # would block a gated ACK + """Happy path: delivery completed and no read is in flight, so the ACK + goes out at once, without a drain.""" + wt = _preprocess(inflight_reads=False) wt._finish_reading("X", drained=True) assert len(_pw_reads_done(wt)) == 1 assert "X" not in wt._draining_rids From b15fba9bf76d0349a244b4557bb39dfff7152adf Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Mon, 19 Oct 2026 21:10:23 +0000 Subject: [PATCH 03/15] Poll a read's future through done() and result() only, as the SHM arena's event future offers no more --- mstar/communication/tensors.py | 21 ++++++++---- test/modular/test_failed_tensor_read.py | 43 +++++++++++++++++++++---- test/modular/test_shm_tensor_comm.py | 5 +-- 3 files changed, 52 insertions(+), 17 deletions(-) diff --git a/mstar/communication/tensors.py b/mstar/communication/tensors.py index bb35c6dabc..f3c55432f8 100644 --- a/mstar/communication/tensors.py +++ b/mstar/communication/tensors.py @@ -296,12 +296,23 @@ class TransferReadInfo: _READ_POLL_INTERVAL_S = 1e-4 -def _may_be_reading(future: Future) -> bool: +def _read_error(future) -> BaseException | None: + """How a finished read failed, or None. Only ``done()`` and ``result()`` + are used: the SHM arena's ``_CudaEventFuture`` has nothing else.""" + try: + future.result() + except Exception as exc: # noqa: BLE001 — includes a cancelled read + return exc + return None + + +def _may_be_reading(future) -> bool: """A read still running, or one that failed and so may have left transfers running until :meth:`get_ready_tensors` takes them over.""" if not future.done(): return True - return not future.cancelled() and future.exception() is not None + error = _read_error(future) + return error is not None and not isinstance(error, CancelledError) def _submit_failed(batch_id: int) -> bool: @@ -1599,9 +1610,7 @@ def _poll_read(ep: FutureAndPointers) -> tuple[bool, BaseException | None]: return True, None if not ep.future.done(): return False, None - # exception(), not result(): a re-raise would put this frame on the - # stored traceback, and with it the read's buffers - error = CancelledError() if ep.future.cancelled() else ep.future.exception() + error = _read_error(ep.future) if error is None: ep.future = None return True, error @@ -1697,7 +1706,7 @@ def _release_kept_reads(self) -> None: still_kept.append(kept) continue future, kept.future = kept.future, None - error = CancelledError() if future.cancelled() else future.exception() + error = _read_error(future) if not self._take_over(kept, error): continue for ptr, batch_id in list(kept.running.items()): diff --git a/test/modular/test_failed_tensor_read.py b/test/modular/test_failed_tensor_read.py index b0fdd55e8e..8d0ca07fcf 100644 --- a/test/modular/test_failed_tensor_read.py +++ b/test/modular/test_failed_tensor_read.py @@ -145,15 +145,15 @@ class _CountingFuture(Future): def __init__(self): super().__init__() self.done_calls = 0 - self.exception_calls = 0 + self.result_calls = 0 def done(self): self.done_calls += 1 return super().done() - def exception(self, timeout=None): - self.exception_calls += 1 - return super().exception(timeout) + def result(self, timeout=None): + self.result_calls += 1 + return super().result(timeout) def test_a_finished_read_is_queried_once_while_it_waits_behind_an_earlier_one(tmp_path): @@ -165,12 +165,41 @@ def test_a_finished_read_is_queried_once_while_it_waits_behind_an_earlier_one(tm for _ in range(3): assert mgr.get_ready_tensors() == {} - assert (earlier.done_calls, earlier.exception_calls) == (3, 0) - assert (later.done_calls, later.exception_calls) == (1, 1) + assert (earlier.done_calls, earlier.result_calls) == (3, 0) + assert (later.done_calls, later.result_calls) == (1, 1) earlier.set_result(None) assert mgr.get_ready_tensors() == {7: [first, second]} - assert (later.done_calls, later.exception_calls) == (1, 1) + assert (later.done_calls, later.result_calls) == (1, 1) + + +class _EventFuture: + """Shaped like the SHM arena's ``_CudaEventFuture``: ``done()`` and + ``result()`` only.""" + + def __init__(self, done): + self.finished = done + + def done(self): + return self.finished + + def result(self): + return None + + +def test_a_future_with_only_done_and_result_is_polled_like_any_other(tmp_path): + """The arena attaches such a future to every host-to-device read.""" + mgr, _ = _shm_manager(tmp_path) + running, finished = _EventFuture(False), _EventFuture(True) + first, second = _edge(), _edge() + mgr.pending = [_read(7, first, running), _read(8, second, finished)] + + assert mgr.get_ready_tensors() == {8: [second]} + assert mgr.has_inflight_reads(7) and not mgr.has_inflight_reads(8) + running.finished = True + assert not mgr.has_inflight_reads(7) + assert mgr.get_ready_tensors() == {7: [first]} + assert mgr.take_failed_reads() == [] and mgr.pending == [] def test_the_worker_fails_a_request_once_when_its_input_read_fails(tmp_path): diff --git a/test/modular/test_shm_tensor_comm.py b/test/modular/test_shm_tensor_comm.py index 0d77da9733..c7ad2b3f71 100644 --- a/test/modular/test_shm_tensor_comm.py +++ b/test/modular/test_shm_tensor_comm.py @@ -611,10 +611,7 @@ def __init__(self, done: bool): def done(self): return self.finished - def cancelled(self): - return False - - def exception(self): + def result(self): return None def _read(future, rid: str, name: str) -> FutureAndPointers: From 3e014451c928c1d096c256ece740550379c83090 Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Tue, 20 Oct 2026 04:10:54 +0000 Subject: [PATCH 04/15] test: the output-failure teardown fake tensor manager reports no failed reads --- test/modular/test_output_failure_teardown.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/modular/test_output_failure_teardown.py b/test/modular/test_output_failure_teardown.py index 1604ec5ea9..5aecaaffea 100644 --- a/test/modular/test_output_failure_teardown.py +++ b/test/modular/test_output_failure_teardown.py @@ -62,6 +62,9 @@ def get_ready_tensors(self): ready, self.ready = self.ready, {} return ready + def take_failed_reads(self): + return [] + def get_tensor(self, uuid): return object() From dc83e301905bd3b970d90022fea989a17abb7adf Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Tue, 20 Oct 2026 21:17:02 +0000 Subject: [PATCH 05/15] Transfer metrics: a failed read is dropped and reported, not raised, so the read delivered before it counts once --- mstar/communication/tensors.py | 3 +-- test/modular/test_transfer_metrics.py | 12 +++++++----- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/mstar/communication/tensors.py b/mstar/communication/tensors.py index f3c55432f8..7328baf601 100644 --- a/mstar/communication/tensors.py +++ b/mstar/communication/tensors.py @@ -1586,8 +1586,7 @@ def _record_transfer(self, stats: EdgeTransfers, ep: FutureAndPointers) -> None: is read as a whole, so an edge joined from several producers splits its read time between them by bytes (by tensors if it moved none).""" src = ep.read_src - # Counted once: a later future that raises leaves this entry pending - # for the next poll. + # Cleared, so this read is never counted twice. ep.read_src = None dst = self._transfer_dst or ep.graph_edges[0].next_node seconds = ep.rx_time or 0.0 diff --git a/test/modular/test_transfer_metrics.py b/test/modular/test_transfer_metrics.py index e8aefe3169..8e223d7ef3 100644 --- a/test/modular/test_transfer_metrics.py +++ b/test/modular/test_transfer_metrics.py @@ -263,8 +263,8 @@ def batch_transfer_sync_read(self, *args): def test_a_failed_read_does_not_count_the_reads_before_it_again(): - """A read whose future raises leaves the poll's earlier entries pending, - so the next poll sees them again; they count once.""" + """A read whose future raises is dropped and handed to take_failed_reads; + the read delivered before it counts once, the failed one not at all.""" producer, consumer = _mooncake("worker_0"), _mooncake("worker_1") consumer._async_reader = AsyncMooncakeReader(_FailSecondRead(), device="cpu") try: @@ -278,9 +278,11 @@ def test_a_failed_read_does_not_count_the_reads_before_it_again(): for future in [ep.future for ep in consumer.pending]: with contextlib.suppress(RuntimeError): future.result(timeout=5) - for _ in range(3): - with pytest.raises(RuntimeError, match="Mooncake read failed"): - consumer.get_ready_tensors() + assert [e.name for e in consumer.get_ready_tensors()["r1"]] == ["h"] + for _ in range(2): + assert consumer.get_ready_tensors() == {} + (failed,) = consumer.take_failed_reads() + assert failed.request_id == "r1" and [e.name for e in failed.graph_edges] == ["g"] assert _series(stats.snapshot())[("enc", "llm", "TCP")][:3] == (16, 1, 1) finally: producer._async_reader.shutdown() From c8173119f7243bb71e489717dd37b8883498f18a Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Tue, 20 Oct 2026 21:29:29 +0000 Subject: [PATCH 06/15] test: the serving-metrics stub tensor manager reports no failed reads --- test/modular/test_serving_metrics.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/modular/test_serving_metrics.py b/test/modular/test_serving_metrics.py index 29a670c2c4..4000d23185 100644 --- a/test/modular/test_serving_metrics.py +++ b/test/modular/test_serving_metrics.py @@ -686,6 +686,9 @@ def register_for_send(self, request_id, tensor_infos): def get_ready_tensors(self): return {} + def take_failed_reads(self): + return [] + def has_inflight_reads(self, request_id): return False From c4cbd401b44f81cf7c84b41e762bf20208055fc4 Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Wed, 21 Oct 2026 02:18:21 +0000 Subject: [PATCH 07/15] test: the fake Mooncake engine's report waits for the reader to submit the batch it targets On a loaded host the reader thread can submit after the test reports a status, which dropped the report and left the read waiting out its deadline. --- test/modular/test_failed_tensor_read.py | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/test/modular/test_failed_tensor_read.py b/test/modular/test_failed_tensor_read.py index 8d0ca07fcf..8b31eabe52 100644 --- a/test/modular/test_failed_tensor_read.py +++ b/test/modular/test_failed_tensor_read.py @@ -11,6 +11,7 @@ import itertools import json import threading +import time from concurrent.futures import Future from types import SimpleNamespace @@ -305,12 +306,23 @@ def transfer_check_status(self, batch_id): self._freed.add(batch_id) return status - def report(self, remote, status): - """From now on, the live batch reading ``remote`` reports ``status``.""" - with self._lock: - for batch_id, batch in self._batches.items(): - if batch[1] == remote and batch_id not in self._freed: + def report(self, remote, status, wait=True): + """From now on, the live batch reading ``remote`` reports ``status``. + ``wait`` waits for the reader thread to submit it, which a loaded host + delays.""" + deadline = time.monotonic() + 5 + while True: + with self._lock: + live = [ + batch for batch_id, batch in self._batches.items() + if batch[1] == remote and batch_id not in self._freed + ] + for batch in live: batch[3] = [status] + if live or not wait: + return + assert time.monotonic() < deadline, f"no live batch reads {remote:#x}" + time.sleep(0.001) def submits(self, remote): return self._submits[remote] @@ -337,7 +349,7 @@ def make(): yield engine # let any reader still polling a running batch finish for remote in list(engine.plan): - engine.report(remote, 1) + engine.report(remote, 1, wait=False) for mgr in managers: mgr._async_reader._executor.shutdown(wait=True) assert engine.after_free == [], "a freed batch was polled" From 6bf3a03188332e03404ca516baee078f11711f2b Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Wed, 21 Oct 2026 04:36:15 +0000 Subject: [PATCH 08/15] Count a read's retries under the device filter the Mooncake engine was initialized with get_local_topology builds its own engine, so without the TCP device it counted HCAs the engine filtered out, and retried a failed tensor more often than the sync read does. --- mstar/communication/tensors.py | 18 +++++++++++------- test/modular/test_failed_tensor_read.py | 15 +++++++++++++-- test/modular/test_mooncake_flat_layout.py | 14 ++++++++++++-- 3 files changed, 36 insertions(+), 11 deletions(-) diff --git a/mstar/communication/tensors.py b/mstar/communication/tensors.py index 7328baf601..5ea96539be 100644 --- a/mstar/communication/tensors.py +++ b/mstar/communication/tensors.py @@ -331,11 +331,12 @@ def _mooncake_transfer_timeout_s() -> int: return max(5, int(match.group())) if match else 5 -def _local_hca_count(engine) -> int: - """HCAs in the engine's local topology: Mooncake's sync read tries a - failed batch ``numContexts() + 1`` times.""" +def _local_hca_count(engine, device_filter: str) -> int: + """HCAs in the engine's local topology under the device filter it was + initialized with: Mooncake's sync read tries a failed batch + ``numContexts() + 1`` times.""" try: - topology = json.loads(engine.get_local_topology()) + topology = json.loads(engine.get_local_topology(device_filter)) return len({hca for tiers in topology.values() for tier in tiers for hca in tier}) except Exception: # noqa: BLE001 — a topology we can't read means no retries return 0 @@ -449,9 +450,9 @@ class TrackedMooncakeReader(AsyncMooncakeReader): Retries and the deadline are the sync read's. """ - def __init__(self, engine, device, max_workers: int = 3): + def __init__(self, engine, device, max_workers: int = 3, device_filter: str = ""): super().__init__(engine, device, max_workers=max_workers) - self._read_attempts = _local_hca_count(engine) + 1 + self._read_attempts = _local_hca_count(engine, device_filter) + 1 self._timeout_s = _mooncake_transfer_timeout_s() def _read(self, read_info: list["TransferReadInfo"]) -> None: @@ -2100,7 +2101,10 @@ def __init__( # the consumer has the engine copy ``nbytes`` from ``address + offset`` # into a buffer it views with the descriptor's strides self.needs_flat_layout = protocol in (CommProtocol.RDMA, CommProtocol.TCP) - self._async_reader = TrackedMooncakeReader(engine._engine, device=device) + self._async_reader = TrackedMooncakeReader( + engine._engine, device=device, + device_filter=tcp_transfer_device if protocol == CommProtocol.TCP else "", + ) self._transfer_status = self._async_reader.transfer_status @property diff --git a/test/modular/test_failed_tensor_read.py b/test/modular/test_failed_tensor_read.py index 8b31eabe52..73039c3d7b 100644 --- a/test/modular/test_failed_tensor_read.py +++ b/test/modular/test_failed_tensor_read.py @@ -270,8 +270,11 @@ def __init__(self, hcas=0): self._ids = itertools.count(1) self._lock = threading.Lock() - def get_local_topology(self): - return json.dumps({"cpu:0": [[f"mlx5_{i}" for i in range(self.hcas)], []]}) + def get_local_topology(self, device_name=""): + hcas = [f"mlx5_{i}" for i in range(self.hcas)] + if device_name: + hcas = [hca for hca in hcas if hca in device_name.split(",")] + return json.dumps({"cpu:0": [hcas, []]}) def register_memory(self, ptr, nbytes): return 0 @@ -603,6 +606,14 @@ def test_a_failed_tensor_is_retried_as_often_as_the_sync_read_retries(engine, hc assert "Status: -1" in str(error) +@pytest.mark.parametrize("device_filter, attempts", [("", 3), ("mlx5_1", 2)]) +def test_retries_count_only_the_hcas_the_engines_device_filter_keeps(engine, device_filter, attempts): + engine.hcas = 2 + reader = TrackedMooncakeReader(engine, "cpu", device_filter=device_filter) + assert reader._read_attempts == attempts + reader.shutdown() + + def test_a_read_past_its_deadline_hands_over_its_running_transfers(engine): mgr = engine.manager() mgr._async_reader._timeout_s = 0 diff --git a/test/modular/test_mooncake_flat_layout.py b/test/modular/test_mooncake_flat_layout.py index 75f2b3e3d6..8b63eec0c6 100644 --- a/test/modular/test_mooncake_flat_layout.py +++ b/test/modular/test_mooncake_flat_layout.py @@ -58,8 +58,8 @@ def __init__(self, hostname, **kwargs): class _FlatCopyReader: """The transfer engine's batch read: ``nbytes`` from ``remote_ptr``.""" - def __init__(self, engine, device, enable_prof=False): - pass + def __init__(self, engine, device, enable_prof=False, device_filter=""): + self.device_filter = device_filter def submit(self, reads): for r in reads: @@ -143,6 +143,16 @@ def test_batch_store_gapped_views_arrive_intact(bookkeeping): assert torch.equal(got, want) +@pytest.mark.parametrize("protocol, device_filter", [(CommProtocol.TCP, "eth1"), (CommProtocol.RDMA, "")]) +def test_the_reader_gets_the_device_filter_the_engine_was_initialized_with(protocol, device_filter): + """The engine filters its topology by the TCP device only, so the reader + counts retries under the same filter.""" + mgr = MooncakeCommunicationManager( + "e", "host", "cpu", _NullCommunicator(), protocol=protocol, tcp_transfer_device="eth1", + ) + assert mgr._async_reader.device_filter == device_filter + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") def test_flat_copy_waits_on_the_ready_event(): """The copy runs on its own stream, so a producer still writing on From a61c341676bf9af0ec5cc4d2a175f0d15d00d080 Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Wed, 21 Oct 2026 04:36:15 +0000 Subject: [PATCH 09/15] Log a failed output read in one line, as a failed input read is The tensor manager already logs the read's error; the data worker logged a full traceback per tensor slot on top. --- mstar/api_server/data_worker.py | 14 +++++++++++--- test/modular/test_failed_tensor_read.py | 11 ++++++++--- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/mstar/api_server/data_worker.py b/mstar/api_server/data_worker.py index 50308cffe7..773fe2e010 100644 --- a/mstar/api_server/data_worker.py +++ b/mstar/api_server/data_worker.py @@ -890,7 +890,7 @@ def _summarize_inputs(input: PreprocessInput) -> list[InputInfo]: def _fail_request( self, request_id: str, exc: BaseException, stage: str, count: int = 1, - sequence: int | None = None, abort: bool = True, + sequence: int | None = None, abort: bool = True, traceback: bool = True, ): """Report a per-request data-worker failure to the API server. @@ -905,8 +905,16 @@ def _fail_request( ``abort`` (False for a request the conductor has not seen) stops the request's walk, since its client already has its answer. + + ``traceback`` False logs one line, for an error already logged where + it happened. """ - logger.error("%s failed for request %s", stage, request_id, exc_info=exc) + if traceback: + logger.error("%s failed for request %s", stage, request_id, exc_info=exc) + else: + logger.error( + "%s failed for request %s: %s: %s", stage, request_id, type(exc).__name__, exc, + ) status = 400 if isinstance(exc, (ValueError, TypeError)) else 500 chunks = [ ResultChunk( @@ -1057,7 +1065,7 @@ def _fail_unread_outputs(self, failed: FailedRead) -> None: slot = state.order.pop(info.uuid, None) if state is not None else None self._fail_request( request_id, failed.error, "output tensor read", - sequence=slot[0] if slot is not None else None, + sequence=slot[0] if slot is not None else None, traceback=False, ) def _process_read_tensors(self): diff --git a/test/modular/test_failed_tensor_read.py b/test/modular/test_failed_tensor_read.py index 73039c3d7b..366593355a 100644 --- a/test/modular/test_failed_tensor_read.py +++ b/test/modular/test_failed_tensor_read.py @@ -220,10 +220,11 @@ def test_the_worker_fails_a_request_once_when_its_input_read_fails(tmp_path): assert w.scheduler.failed_rids == {"r1"} -def test_the_data_worker_answers_a_failed_output_read_in_its_slot(tmp_path): +def test_the_data_worker_answers_a_failed_output_read_in_its_slot(tmp_path, caplog): """The read of output 0 failed while output 1 waits in the reorder buffer: the error takes slot 0, so 1 follows it instead of waiting for the TTL, - and the API server gets one chunk per tensor it counted.""" + and the API server gets one chunk per tensor it counted. The failure is + logged in one line each, without a traceback, as an input read's is.""" wt = _preprocess_thread(model=None) held = ResultChunk(request_id="r1", modality="text", data=b"later", metadata={}) wt.request_output_state["r1"].pending[1] = held @@ -231,7 +232,11 @@ def test_the_data_worker_answers_a_failed_output_read_in_its_slot(tmp_path): wt.tensor_manager, _ = _shm_manager(tmp_path) wt.tensor_manager.pending = [_read("r1", _edge("u1", name="text_output"), _failed_future())] - assert wt._process_read_tensors() is True + with caplog.at_level("ERROR"): + assert wt._process_read_tensors() is True + errors = [r for r in caplog.records if r.levelname == "ERROR"] + assert [r.exc_info for r in errors] == [None, None] + assert "output tensor read failed for request r1: RuntimeError" in errors[1].getMessage() chunks = [wt.out_queue.get_nowait(), wt.out_queue.get_nowait()] assert [c.modality for c in chunks] == ["error", "text"] assert b"output tensor read failed: RuntimeError: Mooncake read failed" in chunks[0].data From 9ff9f8d5b9cc4f4f1358da73f0d5c682f96dbbfe Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Fri, 23 Oct 2026 22:29:31 +0000 Subject: [PATCH 10/15] Count the HCAs for read retries on the first failed tensor, not at startup get_local_topology starts a second Mooncake engine, about 1 s per manager, and the count is only used to retry a failed tensor. --- mstar/communication/tensors.py | 16 +++++++++++-- test/modular/test_failed_tensor_read.py | 31 +++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/mstar/communication/tensors.py b/mstar/communication/tensors.py index 5ea96539be..c696880d66 100644 --- a/mstar/communication/tensors.py +++ b/mstar/communication/tensors.py @@ -11,6 +11,7 @@ import shutil import stat import sys +import threading import time from abc import ABC, abstractmethod from collections.abc import Callable @@ -452,9 +453,20 @@ class TrackedMooncakeReader(AsyncMooncakeReader): def __init__(self, engine, device, max_workers: int = 3, device_filter: str = ""): super().__init__(engine, device, max_workers=max_workers) - self._read_attempts = _local_hca_count(engine, device_filter) + 1 + self._device_filter = device_filter + self._attempts: int | None = None + self._attempts_lock = threading.Lock() self._timeout_s = _mooncake_transfer_timeout_s() + @property + def _read_attempts(self) -> int: + """Counted on the first failed tensor, not at startup: + ``get_local_topology`` starts a second engine, about 1 s.""" + with self._attempts_lock: + if self._attempts is None: + self._attempts = _local_hca_count(self._engine, self._device_filter) + 1 + return self._attempts + def _read(self, read_info: list["TransferReadInfo"]) -> None: """Raises ``AsyncReadError``, after which this thread polls none of the batch ids it hands over.""" @@ -492,7 +504,7 @@ def submit(info: TransferReadInfo) -> None: if status != _TRANSFER_FAILED: # may or may not have freed the batch: never poll it again raise fail(f"Mooncake read returned unknown status {status}", {ptr}) - if attempts[ptr] >= self._read_attempts or time.monotonic() > deadline: + if time.monotonic() > deadline or attempts[ptr] >= self._read_attempts: raise fail(f"Mooncake read failed. Status: {status}") submit(info) if not running: diff --git a/test/modular/test_failed_tensor_read.py b/test/modular/test_failed_tensor_read.py index 366593355a..421899ef57 100644 --- a/test/modular/test_failed_tensor_read.py +++ b/test/modular/test_failed_tensor_read.py @@ -264,6 +264,8 @@ class _Engine: def __init__(self, hcas=0): self.plan: dict[int, list] = {} self.hcas = hcas + self.topology_calls = 0 + self.topology_delay = 0 self.submit_failure = 0 self.unregister_result = 0 self.unregistered: list[int] = [] @@ -276,6 +278,8 @@ def __init__(self, hcas=0): self._lock = threading.Lock() def get_local_topology(self, device_name=""): + self.topology_calls += 1 + time.sleep(self.topology_delay) hcas = [f"mlx5_{i}" for i in range(self.hcas)] if device_name: hcas = [hca for hca in hcas if hca in device_name.split(",")] @@ -619,6 +623,33 @@ def test_retries_count_only_the_hcas_the_engines_device_filter_keeps(engine, dev reader.shutdown() +def test_reads_that_succeed_never_count_the_hcas(engine): + """Counting starts a second Mooncake engine, so neither startup nor a + successful read pays for it.""" + mgr = engine.manager() + src = _sources(2) + engine.plan = {s.data_ptr(): [[0, 1]] for s in src} + future, _ = _start(mgr, "ok", src) + assert future.exception(timeout=5) is None + assert engine.topology_calls == 0 + + +def test_the_first_failed_tensor_counts_the_hcas_once(engine): + """Reads failing at once on the reader's three threads count once.""" + engine.hcas, engine.topology_delay = 1, 0.05 + mgr = engine.manager() + futures = [] + for i in range(3): + src = _sources(1) + engine.plan[src[0].data_ptr()] = [[-1], [1]] + (future,) = mgr.start_read_tensors(f"r{i}", [_remote_edge(src, first_uuid=100 + i)]) + futures.append((future, src)) + for future, src in futures: + assert future.exception(timeout=5) is None + assert engine.submits(src[0].data_ptr()) == 2 + assert engine.topology_calls == 1 + + def test_a_read_past_its_deadline_hands_over_its_running_transfers(engine): mgr = engine.manager() mgr._async_reader._timeout_s = 0 From 9e3fd1447e727dfe4a8624c4a97b61ba858fc199 Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Sat, 24 Oct 2026 06:46:33 +0000 Subject: [PATCH 11/15] test: over real Mooncake TCP, a failed read keeps its buffers until a transfer still running reports it stopped --- test/modular/test_failed_tensor_read.py | 54 +++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/test/modular/test_failed_tensor_read.py b/test/modular/test_failed_tensor_read.py index 421899ef57..eaf1a0fe6d 100644 --- a/test/modular/test_failed_tensor_read.py +++ b/test/modular/test_failed_tensor_read.py @@ -22,6 +22,7 @@ from mstar.api_server.request_types import ResultChunk from mstar.communication.communicator import BaseCommunicator, CommProtocol from mstar.communication.tensors import ( + _TRANSFER_COMPLETED, MOONCAKE_IMPORT_ERROR, AsyncMooncakeReader, AsyncReadError, @@ -745,3 +746,56 @@ def make(entity): finally: for mgr in (producer, consumer): mgr._async_reader._executor.shutdown(wait=True) + + +@pytest.mark.skipif(MOONCAKE_IMPORT_ERROR is not None, reason=f"mooncake unavailable: {MOONCAKE_IMPORT_ERROR}") +def test_real_mooncake_tcp_keeps_a_failed_reads_buffers_until_its_running_transfer_stops(): + """One read of two tensors over real Mooncake TCP: the first was never + registered on the producer and fails at once, the second (2 GiB) is still + moving when the read gives up. The request fails at once, but its buffers + stay registered and the request keeps a read in flight (READS_DONE waits) + until Mooncake reports the large transfer COMPLETED; then both buffers + are unregistered and freed.""" + def make(entity): + return create_tensor_communication_manager( + protocol=CommProtocol.TCP, my_entity_id=entity, hostname="localhost", + device="cpu", communicator=_Communicator(), tcp_transfer_device="lo", + ) + + producer, consumer = make("producer"), make("consumer") + recording = _Recording(consumer._async_reader._engine) + consumer._async_reader._engine = recording + unregistered = [] + unregister = consumer.transfer_engine.unregister_memory + consumer.transfer_engine.unregister_memory = lambda ptr: unregistered.append(ptr) or unregister(ptr) + try: + bad = torch.ones(4096, dtype=torch.int32) + big = torch.ones(2 << 30, dtype=torch.uint8) + bad_info = producer.store_and_return_tensor_info("r", {"bad": [bad]})["bad"] + big_info = producer.store_and_return_tensor_info("r", {"big": [big]})["big"] + producer.register_for_send("r", big_info) + edge = GraphEdge(next_node="n", name="x") + edge.tensor_info = bad_info + big_info + + (future,) = consumer.start_read_tensors("r", [edge]) + ptrs = [consumer.get_tensor(info.uuid).data_ptr() for info in edge.tensor_info] + error = future.exception(timeout=60) + assert isinstance(error, AsyncReadError), error + assert ptrs[1] in error.running, "the 2 GiB transfer finished before the read gave up" + + assert consumer.get_ready_tensors() == {} + assert [f.request_id for f in consumer.take_failed_reads()] == ["r"] + assert consumer.has_inflight_reads("r") + assert not set(ptrs) & set(unregistered) + deadline = time.monotonic() + 120 + while consumer.has_inflight_reads("r"): + assert not set(ptrs) & set(unregistered), "a buffer was freed while its transfer ran" + assert time.monotonic() < deadline, "the large transfer never reported that it stopped" + consumer.get_ready_tensors() + time.sleep(0.01) + assert _TRANSFER_COMPLETED in recording.statuses + assert set(ptrs) <= set(unregistered) + assert consumer._kept_reads == [] and consumer._unreleasable_buffers == [] + finally: + for mgr in (producer, consumer): + mgr._async_reader._executor.shutdown(wait=True) From 490b808847ef8cbb6e31bcb4c3b267048f06ce6f Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Sun, 25 Oct 2026 23:58:00 +0000 Subject: [PATCH 12/15] test: acceptance tests for a TP follower whose input read fails after rank 0 read its copy Two real ranks in their own processes over gloo, running the real Worker.run, scheduler, SHM tensor manager, Engine/StepRunner/KVManager and Qwen3.5's prepare_inputs. Checks behavior only: no hang and only the failed request fails, nothing of the failed row routed or sent, its pages never leasable with prefix_cache_tp, and the request beside it bit-identical to a run without the failure; for an image-grid row and a text-only row. --- .../test_tp_read_failure_acceptance.py | 983 ++++++++++++++++++ 1 file changed, 983 insertions(+) create mode 100644 test/modular/test_tp_read_failure_acceptance.py diff --git a/test/modular/test_tp_read_failure_acceptance.py b/test/modular/test_tp_read_failure_acceptance.py new file mode 100644 index 0000000000..8f1e6f1dd6 --- /dev/null +++ b/test/modular/test_tp_read_failure_acceptance.py @@ -0,0 +1,983 @@ +"""Acceptance tests: a TP=2 follower's async input read fails after rank 0 has +read its own copy of that input. + +Rank 0 schedules a TP step from its own copy of each row's inputs, so a +follower whose read of one of them fails can be left holding a step it cannot +run as rank 0 runs it. Whatever the fix, these must hold: + + A1 no hang, and only the request whose read failed fails; a request + running beside it completes + A2 nothing of the failed row leaves either rank: no routed output, no + tensor registered to be read elsewhere, no output sent on + A3 with ``prefix_cache_tp`` on, the failed row's pages never become + leasable, so a later request with the same prefix doesn't reuse them + A4 a request running beside the failed one gives bit-identical outputs + to the same run with no failure + A5 the same for a text-only row, which has no image grid + +Each run is two real ranks of one TP=2 instance, each in its own process +with a gloo process group, running the real ``Worker.run`` loop, the real +``MicroScheduler``, the real shared-memory tensor manager, the real +``Engine`` (prepare, admit, forward, commit through the real ``StepRunner`` +and ``KVManager``) and Qwen3.5's real ``prepare_inputs``, which reads +``image_grid_thw`` on the host to size the image span. The LLM forward is a +stand-in that all-reduces each rank's partial over the TP group, as a +row-parallel layer does, so a rank that never joins a step, or joins it with +another token count, shows up as a collective timeout or a gloo abort. + +Fakes: the graph runtime and request state (inputs land through +``ingest_inputs_batch``; what a rank routes and sends is recorded off +``complete_and_route_batch`` and ``send_outputs``), the conductor (played by +the test: FAIL_REQUESTS starts the drain, both ranks' READS_DONE sends the +REMOVE), and the reads, whose futures the test resolves. The follower's read +of r1 resolves when the follower first hears from rank 0 about a step, or a +couple of seconds in, whichever comes first, so it fails after rank 0 has +scheduled from its own copy under any design that lets it. + +The tests check behavior only, never a hook, so a design passes them however +it gets there (holding the row on rank 0, as the read gate in #754 does, or +zero-filling the follower's span and dropping the row at the end of the +step). Per-rank setup that a design adds (``Worker._setup_read_gate``) is +called when it exists; a design with its own setup has to be added to +``_DESIGN_SETUP``. +""" + +from __future__ import annotations + +import datetime +import multiprocessing as mp +import os +import queue +import sys +import time +import traceback +from concurrent.futures import Future +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from mstar.utils.ipc_format import WorkerMessageType + +HERE = Path(__file__).resolve().parent + +WORKERS = ("worker_0", "worker_1") +NODE = "LLM" +VISION_WALK, TEXT_WALK = "prefill_vision", "prefill" +KV = "kv" +HIDDEN = 16 +VOCAB = 512 +MERGE = 2 +PAGE_SIZE = 16 +ROOT = b"tp read failure" +SIGNALS = ("new_token", "hidden") +# A rank waiting this long in a collective for its partner is the hang. +COLLECTIVE_S = 8.0 +# The follower's late read resolves at its first SCHEDULE_TP, or after this. +TRIGGER_S = 2.0 +# The whole run; a run still going past this hung. +DEADLINE_S = 30.0 +# Per-rank setup a fix adds to Worker.__init__ (see the module docstring). +_DESIGN_SETUP = ("_setup_read_gate",) + +# What each request reads. A grid prompt is Qwen3.5's prefill_vision layout: +# text, image, text; an image of grid (t, h, w) takes t*(h/2)*(w/2) tokens. +_PROMPTS = { + "grid": { + "r0": dict(texts=(5, 7), grid=(1, 16, 24)), + "r1": dict(texts=(6, 4), grid=(1, 32, 32)), + }, + "text": { + "r0": dict(texts=(40,)), + "r1": dict(texts=(53,)), + }, +} +# the input of r1 whose follower read fails +_FAILING = {"grid": "image_grid_thw", "text": "text_inputs"} + + +def _ids(n: int, seed: int) -> list[int]: + return [3 + (seed * 7919 + i * 31) % (VOCAB - 3) for i in range(n)] + + +def _inputs(kind: str, rid: str) -> dict[str, list[torch.Tensor]]: + """The tensors each rank reads for ``rid``, by signal.""" + spec = _PROMPTS[kind][rid] + seed = int(rid[1:]) + 1 + texts = [torch.tensor(_ids(n, seed * 10 + i)) for i, n in enumerate(spec["texts"])] + if kind == "text": + return {"text_inputs": texts} + t, h, w = spec["grid"] + span = t * (h // MERGE) * (w // MERGE) + gen = torch.Generator().manual_seed(seed) + return { + "text_inputs": texts, + "image_grid_thw": [torch.tensor([[t, h, w]], dtype=torch.int64)], + "vision_embeds": [torch.randn(span, HIDDEN, generator=gen)], + } + + +def _walk(kind: str) -> str: + return VISION_WALK if kind == "grid" else TEXT_WALK + + +def _prompt_tokens(kind: str, rid: str) -> list[int]: + return [int(v) for t in _inputs(kind, rid)["text_inputs"] for v in t] + + +@dataclass +class Scenario: + kind: str = "grid" + # the follower's read of r1's _FAILING input fails + fail: bool = True + prefix_cache: bool = False + shm_root: str = "" + store: str = "" + + +# ── the rank (runs in its own process) ────────────────────────────────── + + +class _Stop(BaseException): + """Ends ``Worker.run`` from inside the loop; not an Exception, so the + loop's error handler lets it through.""" + + +class _Communicator: + """Messages to and from the parent process, which routes them between the + ranks and plays the conductor. Every message crosses the wire codec.""" + + def __init__(self, worker_id, inbox, outbox, on_message): + self.worker_id = worker_id + self.inbox, self.outbox = inbox, outbox + self.on_message = on_message + self.stopped = False + self._held = [] + + def send(self, dest, msg=None): + from mstar.communication.wire import encode + self.outbox.put(("send", self.worker_id, dest, encode(msg))) + + def _pull(self, timeout): + try: + item = self.inbox.get(timeout=timeout) if timeout else self.inbox.get_nowait() + except queue.Empty: + return False + if item == "stop": + self.stopped = True + else: + self._held.append(item) + return True + + def get_all_new_messages(self): + from mstar.communication.wire import decode + while self._pull(0): + pass + if self.stopped: + raise _Stop + held, self._held = self._held, [] + messages = [decode(b) for b in held] + self.on_message(messages) + return messages + + def wait_for_work(self, timeout_ms=10): + self.on_message([]) + if self._held or self.stopped: + return + self._pull(min(timeout_ms, 20) / 1000) + + +class _Rope: + """The position counter Qwen3.5's prepare_inputs reads its start from.""" + + def position(self, rid, label): + return 0 + + +def _submodule_class(): + from mstar.engine.resources.kv.config import KVStep + from mstar.engine.resources.step import Segment, SubmoduleStep + from mstar.model.qwen3_5.config import ROPE + from mstar.model.qwen3_5.submodules import LLMSubmodule + + class _LLM(LLMSubmodule): + """Qwen3.5's LLM node: everything but the model is the real submodule + (``prepare_inputs``, which reads the image grid on the host, and the + capture configs' dummy rows), so a fix that builds a row the way the + submodule does gets the real code. The forward is a stand-in: each + rank's partial of every token's embedding is summed over the TP group, + and a row's output is its mean, which every rank then holds whole. A + row's output depends on nothing but its own tokens.""" + + def __init__(self, group, rank): + torch.manual_seed(1234) + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.embed_tokens = torch.nn.Embedding(VOCAB, HIDDEN) + super().__init__( + model=model, + config=SimpleNamespace( + vision_start_token_id=1, vision_end_token_id=2, hidden_size=HIDDEN, + stop_token_ids=(), + ), + vision_config=SimpleNamespace(spatial_merge_size=MERGE), + ) + self.head = torch.randn(HIDDEN, VOCAB) + self.requires_grad_(False) + self.bind_node_resources({ROPE: _Rope()}) + self.group, self.rank = group, rank + + # the real model's step, stop and emit settings, which read state this + # stand-in doesn't have + def _device_loopback(self): + return False + + def uniform_row_inputs(self, graph_walk, has_loopback=False): + return None + + def device_loopback_signals(self, graph_walk): + return frozenset() + + def max_batch_size(self, graph_walk): + return None + + def max_rows_in_combined(self, graph_walk): + return None + + def postprocess(self, request_id, request_info, outputs, **kwargs): + return + + def check_stop(self, request_id, request_info, outputs): + return set() + + def check_stop_batched(self, request_ids, request_infos, host_rows, step_walk=None): + return None + + def step_is_input_free(self, graph_walk): + return False + + def inline_client_signals(self, graph_walk): + return {} + + def can_batch(self, batch, model_inputs): + return True + + def declare_step(self, graph_walk, request_ids, inputs, **kwargs): + return SubmoduleStep( + steps={KV: KVStep()}, + segments=[ + Segment(rid, "main", inp.input_seq_len) + for rid, inp in zip(request_ids, inputs, strict=True) + ], + ) + + def preprocess(self, graph_walk, engine_inputs, inputs): + return {"embeds": [ + inp.input_embeds if inp.input_embeds is not None + else self.model.model.embed_tokens(inp.input_ids) + for inp in inputs + ]} + + def forward(self, graph_walk, engine_inputs, **kwargs): + raise AssertionError("every step here runs batched") + + def forward_batched(self, graph_walk, engine_inputs, embeds): + packed = torch.cat(embeds).to(torch.float32) + # rank 0 holds a third of each value and rank 1 two thirds + partial = packed * ((self.rank + 1) / 3.0) + whole = self.group.all_reduce(partial) + out, at = {}, 0 + for rid, rows in zip(engine_inputs.request_ids, embeds, strict=True): + hidden = whole[at:at + rows.shape[0]].mean(dim=0) + at += rows.shape[0] + out[rid] = { + "new_token": [torch.argmax(hidden @ self.head).reshape(1)], + "hidden": [hidden.clone()], + } + return out + + return _LLM + + +def _kv_manager(rank, prefix_cache): + from mstar.engine.resources.base import KeyedWalks + from mstar.engine.resources.kv import manager as manager_mod + from mstar.engine.resources.kv.config import PagedKVConfig + from mstar.engine.resources.kv.manager import KVManager + + class _NoTransfer: + def __init__(self, *args, **kwargs): + pass + + def __getattr__(self, name): + return lambda *args, **kwargs: None + + manager_mod.KVTransferManager = _NoTransfer + kv = KVManager( + cfg=PagedKVConfig( + num_layers=1, num_kv_heads=1, head_dim=8, max_seq_len=4096, num_qo_heads=2, + max_num_pages=128, page_size=PAGE_SIZE, prefix_cache_tp=prefix_cache, + ), + name=KV, joint_comm_group=SimpleNamespace(rank=rank, world_size=2), + transfer_engine_info=None, device=torch.device("cpu"), dtype=torch.float32, + ) + if prefix_cache: + kv.enable_prefix_cache(ROOT, {"main": KeyedWalks((TEXT_WALK,), ("decode",))}) + return kv + + +def _kv_config(tokens): + from mstar.engine.resources.kv.config import KVReqConfig, PrefixSpan, StreamLayout + from mstar.engine.resources.kv.keys import chain + + whole = len(tokens) // PAGE_SIZE + return KVReqConfig( + prefix_keys={"main": chain([tokens[at:at + PAGE_SIZE] for at in range(0, len(tokens), PAGE_SIZE)])}, + prefix_tail={"main": tokens[whole * PAGE_SIZE:]}, + kv_layout={"main": StreamLayout( + spans=(PrefixSpan(len(tokens), len(tokens), TEXT_WALK),), keyed_spans=1, + )}, + ) + + +def _engine(rank, group, kv, prefix_cache, graph_runtime): + from mstar.engine.engine import Engine, SubmoduleManagement + from mstar.engine.resources import StepRunner + + sub = _submodule_class()(group, rank) + engine = Engine(graph_runtime=graph_runtime, autocast_dtype=None) + if rank == 0: + engine._tp_leader_nodes = {NODE} + else: + engine._tp_follower_nodes = {NODE} + engine._resources = {KV: kv} + engine._runner = StepRunner(engine._resources, node_resources={NODE: [KV]}) + engine._submodules = {NODE: SubmoduleManagement( + submodule=sub, forward=sub.forward, forward_batched=sub.forward_batched, + joint_comm_group=SimpleNamespace(rank=rank, world_size=2), resources={KV: kv}, + )} + engine._keyed_walks = {NODE: {TEXT_WALK}} if prefix_cache else {} + engine._prefix_model = "_LLM" + engine._device = torch.device("cpu") + engine._loopback_steps = {} + engine._autocast_dtype_for = lambda submodule: None + engine.reserve_replay_slot = lambda batch: None + return engine + + +class _Graph: + """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.""" + + def __init__(self, rank, kind, report): + self.rank, self.kind, self.report = rank, kind, report + self.walk = _walk(kind) + self.required = tuple(_inputs(kind, "r0")) + self.landed: dict[str, dict[str, list[int]]] = {} + self.ready: set[str] = set() + self.in_flight: set[str] = set() + self.done: set[str] = set() + self.per_request_info = {} + self.fwd = {} + self.tensor_manager = None + self.combined_walk_of = {} + # no client copy rides inline: every output is stored and routed + self.supports_inline_emit = False + + def admit(self, rid, info): + self.per_request_info[rid] = SimpleNamespace( + request_id=rid, stream_buffers={}, stream_buffers_by_consumer={}, ended_streams=set(), + ) + self.fwd[rid] = info + self.landed[rid] = {} + + # rid handles are the wire ids themselves + def get_rid_handle(self, wire_rid): + return wire_rid if wire_rid in self.per_request_info else None + + def get_rid_string(self, rid): + return rid + + def get_sharding_config(self, rid): + group = SimpleNamespace( + nodes={NODE}, tp_size=2, _tp_rank=self.rank, _workers=list(WORKERS), + ) + return SimpleNamespace(groups=[group], get_sharding_group=lambda node, walk: group) + + # inputs + def ingest_inputs_batch(self, block, can_buffer): + at = 0 + for i, rid in enumerate(block.rids): + n = block.tensors_per_edge[i] + signal = block.signal_names[block.signal_name_idxs[i]] + if rid in self.landed: + self.landed[rid][signal] = list(block.uuids[at:at + n]) + at += n + for rid, landed in self.landed.items(): + if rid not in self.done and rid not in self.in_flight and all(s in landed for s in self.required): + self.ready.add(rid) + + # the scheduler's view + def get_ready_nodes(self, exclude_rids, target=None, exclude_target=None): + from mstar.graph.runtime.base import ReadyNodeSpec + if exclude_target is not None and exclude_target == (NODE, self.walk): + return [] + if target is not None and target != (NODE, self.walk): + return [] + rids = [r for r in sorted(self.ready) if r not in exclude_rids and r in self.per_request_info] + return [ReadyNodeSpec(NODE, self.walk, rids)] if rids else [] + + def has_ready_excluding(self, exclude_rids, exclude_target=None): + return bool(self.get_ready_nodes(exclude_rids, exclude_target=exclude_target)) + + def get_worker_graph_id_for_node(self, node_name, graph_walk): + return 0 + + def pop_rids(self, node_name, graph_walk, request_ids, check_ready=False): + from mstar.graph.runtime.base import ColumnarEdgeSpecs, PopRidsOutput + from mstar.utils.containers import ParallelList + if check_ready and not all(r in self.ready for r in request_ids): + return None + block = ColumnarEdgeSpecs.empty() + rids = [] + for rid in request_ids: + if rid not in self.ready and check_ready: + continue + self.ready.discard(rid) + rids.append(rid) + for signal, uuids in self.landed.get(rid, {}).items(): + block.add(rid, signal, uuids, False) + return PopRidsOutput( + wg_ids=ParallelList(rids, [0] * len(rids)), input_edges=block, + output_signals=SIGNALS, rid_walk_idx=[0] * len(rids), walk_output_signals=(SIGNALS,), + ) + + def pop_walk_rids(self, node_name, rows, check_ready=False): + return self.pop_rids(node_name, rows.walks[0], list(rows.rids), check_ready) + + def get_nodes(self, node_name, rids, wg_ids): + return [SimpleNamespace(name=node_name) for _ in rids] + + def push_back_node(self, node_name, rids, wg_ids): + for rid in rids: + if rid in self.per_request_info and rid not in self.done: + self.ready.add(rid) + + def set_in_flight(self, node_name, rids, wg_ids, in_flight): + for rid in rids: + (self.in_flight.add if in_flight else self.in_flight.discard)(rid) + + def is_async_schedulable(self, node_name, walk): + return False + + def get_output_signals(self, node_name, walk): + return list(SIGNALS) + + def get_partition_for_node(self, node_name): + return "default" + + def get_graph_walk(self, rid, partition): + return self.walk + + def get_fwd_info(self, rid, partition): + return self.fwd[rid] + + def ingested_stream_chunks(self, rid, node_name): + return {} + + # post-processing + def cleanup_consumed_inputs(self, node_name, rids, wg_ids): + from mstar.graph.runtime.base import FreedTensors + return FreedTensors.none() + + def clear_pending_loop_stops(self): + pass + + def pending_loop_stop_rids(self, walks, loop_name): + return set() + + def stop_loops_batched(self, **kwargs): + return [] + + def update_dynamic_loop_iters(self, per_request_info, partition): + pass + + def mark_stream_partition_done(self, rid, partition): + pass + + def complete_and_route_batch(self, route_input): + """Every output of a routed row goes to a remote reader (the client), so + the worker registers each one to be read.""" + from mstar.graph.runtime.base import FreedTensors, RouteOutput + rids, uuids, reg_rids, at = route_input.wg_ids.keys, list(route_input.tensors), [], 0 + outputs = {} + for i, rid in enumerate(rids): + outputs[rid] = {} + for j, signal in enumerate(route_input.output_signals): + n = route_input.num_tensors[i * len(route_input.output_signals) + j] + tensors = [self.tensor_manager.get_tensor(u) for u in uuids[at:at + n]] + outputs[rid][signal] = [ + (str(t.dtype), tuple(t.shape), t.contiguous().numpy().tobytes()) for t in tensors + ] + reg_rids += [rid] * n + at += n + self.done.add(rid) + self.report(("routed", self.rank, outputs)) + return RouteOutput( + completion_id=0, register_uuids=uuids, register_rids=reg_rids, + new_token_output_idxs=[], local_streaming_by_signal={}, + freed_inputs=FreedTensors.none(), + ) + + def send_outputs(self, send_input): + self.report(("sent", self.rank, sorted(set(send_input.per_request_info.keys)))) + return [] + + def buffer_publish_info(self, *args): + pass + + def flush_publish_info(self, *args): + pass + + def get_pending_publish_info(self, rid, partition): + return {} + + def remove_request(self, rid): + self.per_request_info.pop(rid, None) + self.landed.pop(rid, None) + self.ready.discard(rid) + self.report(("removed", self.rank, rid)) + + def commit_speculation(self, *args, **kwargs): + pass + + +def _pointer(tensor, uuid): + from mstar.graph.base import TensorPointerInfo + return TensorPointerInfo( + list(tensor.shape), tensor.dtype, tensor.nbytes, 0, list(tensor.stride()), + uuid, "s", "api_server", + ) + + +class _Reads: + """This rank's reads of every request's inputs, through the real SHM + manager. A read's buffer is registered when it starts, holds the values + once it lands, and keeps whatever it held if it fails.""" + + def __init__(self, tm, kind, rank, scenario): + from mstar.communication.tensors import FutureAndPointers + from mstar.graph.base import GraphEdge + self.tm, self.late, self.fired = tm, [], False + uuid = 1000 * (rank + 1) + for rid in ("r0", "r1"): + for signal, tensors in _inputs(kind, rid).items(): + infos, buffers = [], [] + for t in tensors: + uuid += 1 + info = _pointer(t, uuid) + buf = torch.zeros_like(t) + tm.tensor_store.put_tensor(rid, uuid, buf, info) + infos.append(info) + buffers.append((buf, t)) + edge = GraphEdge(next_node=NODE, name=signal, tensor_info=infos) + future = Future() + tm.pending.append(FutureAndPointers(future=future, graph_edges=[edge], request_id=rid)) + if rank == 1 and rid == "r1": + fails = scenario.fail and signal == _FAILING[kind] + self.late.append((future, buffers, fails)) + else: + self._land(future, buffers) + + @staticmethod + def _land(future, buffers): + for buf, values in buffers: + buf.copy_(values) + future.set_result(None) + + def fire(self): + """The follower's reads of r1 resolve.""" + if self.fired: + return + self.fired = True + for future, buffers, fails in self.late: + if fails: + future.set_exception(RuntimeError("Mooncake read failed. Status: -1")) + else: + self._land(future, buffers) + + +def _tp_group(rank): + import torch.distributed as dist + + from mstar.distributed.all_reduce import AllReduce + from mstar.distributed.communication import CommGroup + + group = CommGroup(my_global_rank=rank, my_group_rank=rank, group_members=[0, 1]) + group.device_group = group.cpu_group = dist.group.WORLD + group.initialized = True + group._all_reduce = AllReduce(group.device_group) + return group + + +def _worker(rank, scenario, comm, graph, engine, tm, group): + import torch.distributed as dist + + from mstar.utils.containers import RecentSet + from mstar.worker.micro_scheduler import MicroScheduler + from mstar.worker.worker import _REMOVED_RIDS_KEEP, Worker + + w = Worker.__new__(Worker) + w.worker_id = WORKERS[rank] + w.device = torch.device("cpu") + w.enable_nvtx = w.enable_prof = False + w._launch_wait_s = 1.0 + w._max_consecutive_spec = 1 + w._device_loopback = {} + w._host_rows_slot = 0 + w._stream_buffer_rids = set() + w._phase_period = 0 + w._phase_buf = {} + w._gpu_spans = [] + w.profile_info = SimpleNamespace(pop_request=lambda rid: None) + w.parallel_groups = SimpleNamespace( + barrier_all=dist.barrier, + restore_default_timeout=lambda: None, + get_tp_config_for_node=lambda node: group, + get_sp_config_for_node=lambda node: SimpleNamespace(world_size=1), + get_instance_rank_for_node=lambda node: rank, + get_instance_world_size_for_node=lambda node: 2, + ) + w.communicator = comm + w.wakeup_event = SimpleNamespace(drain=lambda: None, register_future=lambda f: None) + w.tensor_manager = tm + w._graph_runtime = graph + w.engine_manager = SimpleNamespace( + get_engine=lambda node: engine, warmup_all=lambda: None, + remove_request=engine.remove_request, evictable_nodes=lambda: (), + ) + w.request_state = graph + w.parallel_nodes = {NODE} + w.parallel_leader_nodes = {NODE} if rank == 0 else set() + w.is_tp_follower = rank != 0 + w.tp_async_sched, w.tp_async_nodes = False, None + w._tp_broadcast_seq = 0 + w._tp_nospec = RecentSet(64) + w._tp_leader_gap_warned = False + w.scheduler = MicroScheduler(w.engine_manager, parallel_leader_nodes=w.parallel_leader_nodes) + w.scheduler.runtime = graph + w._unprocessed_messages = {} + w._removed_rids = RecentSet(_REMOVED_RIDS_KEEP) + w._last_active = {} + w._hold_logged = {} + w._in_flight_rids = set() + w._pending_removes = {} + w._removes_awaiting_step = {} + w._pending_drains, w._draining_rids, w._reads_done_sent = set(), set(), set() + w.streaming_buffers = {} + w.partition_topology = None + w._my_consumer_connections = [] + w._streaming_edge_names = set() + w._consumer_node_cache = {} + w._stream_partition = {} + w._innermost_loop_by_node = {} + w._metrics = None + for setup in _DESIGN_SETUP: + if hasattr(w, setup): + getattr(w, setup)() + return w + + +def _probe(engine, rid, tokens): + """The leases rank 0 would take for a new request with ``tokens`` if it + sent a step for it now, from this rank's own index.""" + from mstar.engine.resources.step import StepContext + engine.add_request(rid, {KV: _kv_config(tokens)}) + batch = SimpleNamespace( + node_name=NODE, request_ids=(rid,), + step_context=StepContext(request_ids=(rid,), graph_walk=TEXT_WALK, slot=0, capture=False), + ) + hits = engine.decide_cached_prefix(batch) + leases = hits.get(rid, {}) + return {label: list(pages) for label, pages in leases.items()} + + +def _rank_main(rank, scenario, inbox, outbox): + sys.path[:0] = [str(HERE), str(HERE.parents[1])] + os.environ.setdefault("MSTAR_PRE_PLAN_SPEC", "0") + import logging + logging.basicConfig(level=logging.WARNING, format=f"[rank {rank}] %(levelname)s %(name)s: %(message)s") + try: + import conftest # noqa: F401 (the triton stub) + import torch.distributed as dist + + import mstar.communication.wire_types # noqa: F401 (registers the tags) + from mstar.communication.tensors import SharedMemoryCommunicationManager + from mstar.conductor.request_info import CurrentForwardPassInfo + from mstar.model.qwen3_5.qwen3_5_model import IMAGE_PART, TEXT_PART + + dist.init_process_group( + "gloo", init_method=f"file://{scenario.store}", rank=rank, world_size=2, + timeout=datetime.timedelta(seconds=COLLECTIVE_S), + ) + group = _tp_group(rank) + kind = scenario.kind + report = outbox.put + graph = _Graph(rank, kind, report) + kv = _kv_manager(rank, scenario.prefix_cache) + engine = _engine(rank, group, kv, scenario.prefix_cache, graph) + tm = SharedMemoryCommunicationManager( + WORKERS[rank], "localhost", "cpu", None, shm_dir=os.path.join(scenario.shm_root, str(rank)), + ) + graph.tensor_manager = tm + order = [TEXT_PART, IMAGE_PART, TEXT_PART] if kind == "grid" else [TEXT_PART] + for rid in ("r0", "r1"): + graph.admit(rid, CurrentForwardPassInfo( + request_id=rid, graph_walk=_walk(kind), fwd_index=0, random_seed=0, max_tokens=1, + step_metadata={"prefill_order": order}, + )) + tokens = _prompt_tokens(kind, rid) + engine.add_request(rid, {KV: _kv_config(tokens)} if scenario.prefix_cache else None) + reads = _Reads(tm, kind, rank, scenario) + started = time.monotonic() + + def on_message(messages): + from mstar.utils.ipc_format import WorkerMessageType + if rank != 1 or reads.fired: + return + if any(m.message_type == WorkerMessageType.SCHEDULE_TP for m in messages) \ + or time.monotonic() - started > TRIGGER_S: + reads.fire() + + comm = _Communicator(WORKERS[rank], inbox, outbox, on_message) + tm.communicator = comm + w = _worker(rank, scenario, comm, graph, engine, tm, group) + try: + w.run() + except _Stop: + pass + final = {} + if scenario.prefix_cache: + failed_prefix = _prompt_tokens(kind, "r1")[:3 * PAGE_SIZE] + _ids(9, 77) + healthy_prefix = _prompt_tokens(kind, "r0")[:2 * PAGE_SIZE] + _ids(9, 78) + final["probe_failed_prefix"] = _probe(engine, "p1", failed_prefix) + final["probe_healthy_prefix"] = _probe(engine, "p0", healthy_prefix) + report(("final", rank, final)) + except BaseException: + report(("crash", rank, traceback.format_exc())) + raise + + +# ── the parent: routes messages and plays the conductor ───────────────── + + +@dataclass +class Run: + hung: bool = False + crashed: dict = field(default_factory=dict) + exitcodes: dict = field(default_factory=dict) + # rid -> the rank that first reported it failed, and why + failed: dict = field(default_factory=dict) + # rid -> ranks that sent READS_DONE for it + reads_done: dict = field(default_factory=dict) + removed: dict = field(default_factory=dict) + # rank -> rid -> signal -> [(dtype, shape, bytes)] + routed: dict = field(default_factory=lambda: {0: {}, 1: {}}) + sent: dict = field(default_factory=lambda: {0: set(), 1: set()}) + # messages a rank sent anywhere but the other rank and the conductor + other_sends: list = field(default_factory=list) + final: dict = field(default_factory=dict) + elapsed: float = 0.0 + + def completed(self, rid): + """Rank 0 routed the request's output; it never failed.""" + return rid in self.routed[0] and rid not in self.failed + + def torn_down(self, rid): + return rid in self.failed and self.reads_done.get(rid, set()) >= {0, 1} + + def summary(self): + return ( + f"hung={self.hung} crashed={sorted(self.crashed)} exit={self.exitcodes} " + f"failed={ {r: e[:120] for r, e in self.failed.items()} } " + f"reads_done={self.reads_done} routed={ {k: sorted(v) for k, v in self.routed.items()} } " + f"after {self.elapsed:.1f}s" + ) + + +def _run(tmp_path, scenario: Scenario) -> Run: + import mstar.communication.wire_types # noqa: F401 + from mstar.communication.wire import decode, encode + from mstar.utils.ipc_format import ( + ConductorMessageType, + DrainRequest, + MessageSource, + RemoveRequest, + WorkerMessage, + WorkerMessageType, + ) + + tmp_path.mkdir(parents=True, exist_ok=True) + scenario.shm_root = str(tmp_path / "shm") + scenario.store = str(tmp_path / "store") + ctx = mp.get_context("spawn") + inboxes = [ctx.Queue() for _ in WORKERS] + outbox = ctx.Queue() + procs = [ctx.Process(target=_rank_main, args=(r, scenario, inboxes[r], outbox), daemon=True) for r in (0, 1)] + for p in procs: + p.start() + run = Run() + start = time.monotonic() + stopping_since = None + + def to_both(message): + for box in inboxes: + box.put(encode(message)) + + while True: + now = time.monotonic() + if stopping_since is None: + settled = all(run.completed(r) or run.torn_down(r) for r in ("r0", "r1")) + if settled or now - start > DEADLINE_S: + run.hung = not settled + stopping_since = now + for box in inboxes: + box.put("stop") + elif len(run.final) + len(run.crashed) == 2 or now - stopping_since > COLLECTIVE_S + 5: + break + if stopping_since is None and not all(p.is_alive() for p in procs): + # a rank died: what it left is all there is + run.hung = not all(run.completed(r) or run.torn_down(r) for r in ("r0", "r1")) + stopping_since = now + for box in inboxes: + box.put("stop") + try: + item = outbox.get(timeout=0.05) + except queue.Empty: + continue + tag = item[0] + if tag == "send": + _, src, dest, data = item + if dest in WORKERS: + inboxes[WORKERS.index(dest)].put(data) + continue + msg = decode(data) + if dest != "conductor": + run.other_sends.append((src, dest, msg)) + continue + rank = WORKERS.index(src) + if msg.message_type == ConductorMessageType.FAIL_REQUESTS: + for rid, error in msg.body.errors.items(): + if rid in run.failed: + continue + run.failed[rid] = f"rank {rank}: {error}" + to_both(WorkerMessage( + message_type=WorkerMessageType.DRAIN_REQUEST, + body=DrainRequest(request_id=rid, source=MessageSource.CONDUCTOR), + )) + elif msg.message_type == ConductorMessageType.READS_DONE: + rid = msg.body.request_id + run.reads_done.setdefault(rid, set()).add(WORKERS.index(msg.body.entity_id)) + if run.reads_done[rid] >= {0, 1} and rid not in run.removed: + run.removed[rid] = True + to_both(WorkerMessage( + message_type=WorkerMessageType.REMOVE_REQUEST, + body=RemoveRequest(request_id=rid, source=MessageSource.CONDUCTOR), + )) + elif tag == "routed": + for rid, outputs in item[2].items(): + run.routed[item[1]].setdefault(rid, []).append(outputs) + elif tag == "sent": + run.sent[item[1]] |= set(item[2]) + elif tag == "final": + run.final[item[1]] = item[2] + elif tag == "crash": + run.crashed[item[1]] = item[2] + for p in procs: + p.join(timeout=2) + if p.is_alive(): + p.kill() + p.join(timeout=2) + run.exitcodes = {r: p.exitcode for r, p in enumerate(procs)} + for r, code in run.exitcodes.items(): + if code not in (0, None) and r not in run.crashed: + run.crashed[r] = f"exit code {code} (a gloo abort is a collective with another shape)" + run.elapsed = time.monotonic() - start + return run + + +# ── the tests ─────────────────────────────────────────────────────────── + +KINDS = [ + pytest.param("grid", id="grid"), + # A5: a text-only row has no grid, so its step has rank 0's shape anyway + pytest.param("text", id="text_only"), +] + + +def _healthy_ok(run: Run): + assert not run.crashed, run.summary() + "\n" + "\n".join(map(str, run.crashed.values())) + assert not run.hung, "the instance hung: " + run.summary() + assert run.completed("r0"), "r0, beside the failed request, did not complete: " + run.summary() + + +@pytest.mark.parametrize("kind", KINDS) +def test_a_followers_failed_read_fails_only_its_request_and_nothing_hangs(tmp_path, kind): + """A1 (and A5 for text_only). The follower's read of r1's image grid (or + text) fails after rank 0 read its own copy: the run settles before the + deadline, r1 fails and is torn down on both ranks, and r0 completes.""" + run = _run(tmp_path, Scenario(kind=kind)) + _healthy_ok(run) + assert set(run.failed) == {"r1"}, run.summary() + assert "input tensor read failed" in run.failed["r1"], run.summary() + assert run.torn_down("r1"), "r1's drain never finished on both ranks: " + run.summary() + + +@pytest.mark.parametrize("kind", KINDS) +def test_nothing_of_the_failed_row_leaves_either_rank(tmp_path, kind): + """A2 (and A5). No rank routes r1's outputs (so nothing is persisted, + registered for another reader or sent to the client), and no rank sends + anything for r1 past its runtime.""" + run = _run(tmp_path, Scenario(kind=kind)) + _healthy_ok(run) + for rank in (0, 1): + assert "r1" not in run.routed[rank], f"rank {rank} routed r1's outputs: " + run.summary() + assert "r1" not in run.sent[rank], f"rank {rank} sent r1's outputs on: " + run.summary() + out = [s for s in run.other_sends if getattr(s[2].body, "request_id", None) == "r1" + and s[2].message_type != WorkerMessageType.TENSOR_RECEIVED] # an input's ack + assert not out, out + + +def test_the_failed_rows_pages_never_become_leasable(tmp_path): + """A3. With prefix_cache_tp on, a text-only r1 whose follower read failed + leaves no pages in either rank's index: a new request sharing r1's first + three pages gets no lease, on rank 0 (whose lease the followers take) or + on the follower. A request sharing r0's prefix does, so the cache is on + and r0's commit indexed it.""" + run = _run(tmp_path, Scenario(kind="text", prefix_cache=True)) + _healthy_ok(run) + assert set(run.final) == {0, 1}, "a rank never reported its index: " + run.summary() + for rank in (0, 1): + assert run.final[rank]["probe_healthy_prefix"], f"rank {rank}: r0's prefix isn't cached at all" + assert not run.final[rank]["probe_failed_prefix"], ( + f"rank {rank} would lease the failed row's pages to a new request: " + f"{run.final[rank]['probe_failed_prefix']}" + ) + + +@pytest.mark.parametrize("kind", KINDS) +def test_a_request_beside_the_failed_one_is_bit_identical_to_a_run_without_failure(tmp_path, kind): + """A4 (and A5). r0's outputs (token and hidden state, by bytes) on each + rank equal the same run's with every read landing.""" + clean = _run(tmp_path / "clean", Scenario(kind=kind, fail=False)) + _healthy_ok(clean) + assert clean.completed("r1") and not clean.failed, "the run without failure failed: " + clean.summary() + run = _run(tmp_path / "failed", Scenario(kind=kind)) + _healthy_ok(run) + for rank in (0, 1): + assert run.routed[rank].get("r0") == clean.routed[rank].get("r0"), ( + f"rank {rank}: r0's outputs differ from the run without failure" + ) From 3748ab92ba2e83fe203a21f09d140e143078cd58 Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Mon, 26 Oct 2026 00:21:50 +0000 Subject: [PATCH 13/15] TP: take a row whose follower read failed out of the step on every rank, before preparing it Rank 0 schedules a lockstep step from its own copy of each input, so a follower whose read of one fails after rank 0 sent the step could never build the row, and the instance hung in its next collective. Rank 0 now marks a step whose rows take an input the followers read from elsewhere (ScheduleTPNode.read_check). A follower builds such a step as soon as each row's read landed or failed, with the rows it could not read in it and no inputs, so every rank has rank 0's rows. On the GPU thread, before prepare_inputs, one gloo all-reduce of a flag per row tells every rank which rows to take out; they go to failed_requests and fail through the conductor, as a per-row failure does. They are never prepared, admitted, committed or sent, so their outputs and pages never exist, and a prepare that runs a collective per row (Qwen3.5's vocab-parallel embedding) stays in step. A read-checked step is not pre-planned, so every rank admits it after the check. --- mstar/engine/engine.py | 4 + mstar/utils/ipc_format.py | 3 + mstar/worker/micro_scheduler.py | 73 ++- mstar/worker/worker.py | 166 ++++- test/modular/test_kv_prefix_tp.py | 4 +- test/modular/test_preplan_error_paths.py | 2 +- .../test_request_failure_propagation.py | 1 + test/modular/test_resident_delta.py | 5 + test/modular/test_tp_async_lockstep_sim.py | 43 ++ test/modular/test_tp_async_teardown_order.py | 2 +- test/modular/test_tp_read_check.py | 576 ++++++++++++++++++ .../test_tp_read_failure_acceptance.py | 2 +- test/modular/tp_read_check_sim.py | 134 ++++ 13 files changed, 1000 insertions(+), 15 deletions(-) create mode 100644 test/modular/test_tp_read_check.py create mode 100644 test/modular/tp_read_check_sim.py diff --git a/mstar/engine/engine.py b/mstar/engine/engine.py index 41c54e6b21..e4bfc206fc 100644 --- a/mstar/engine/engine.py +++ b/mstar/engine/engine.py @@ -313,6 +313,10 @@ class ExecutingBatch: # rid -> error, for per-rid stages that raised. The rid leaves the batch; # the rest of it runs. failed_requests: dict[int, str] = field(default_factory=dict) + # A lockstep step whose rows' input reads every rank checks before + # preparing it: the leader's rows in its order (see + # ``Worker._drop_rows_whose_read_failed``); None for no check + read_check_rids: tuple[int, ...] | None = None admit_error: AdmitFailedReason | None = None # the resource that ran out, so an eviction can be scoped to it diff --git a/mstar/utils/ipc_format.py b/mstar/utils/ipc_format.py index 92bd85e5ec..b0703db8cf 100644 --- a/mstar/utils/ipc_format.py +++ b/mstar/utils/ipc_format.py @@ -197,6 +197,9 @@ class ScheduleTPNode(MessageBody): # rank 0's prefix-cache leases, by wire request id and then resource prefix_hits: dict[str, dict[str, list[int]]] = field(default_factory=dict) page_trails: dict[str, tuple[int, int]] = field(default_factory=dict) + # A row of the step takes an input a follower reads from elsewhere: every + # rank agrees on whose read failed before preparing the step + read_check: bool = False @dataclass diff --git a/mstar/worker/micro_scheduler.py b/mstar/worker/micro_scheduler.py index c7e6e7a986..bc69acf91d 100644 --- a/mstar/worker/micro_scheduler.py +++ b/mstar/worker/micro_scheduler.py @@ -73,6 +73,9 @@ class ScheduledBatch: request_walks: dict[int, str] = field(default_factory=dict) # real walk -> its output edge names, under a combined walk walk_output_signals: dict[str, list[str]] = field(default_factory=dict) + # A TP follower's copy of ``ScheduleTPNode.read_check``: the leader's rows, + # in its order, including those this rank could not read; None for no check + read_check_rids: tuple[int, ...] | None = None # `rids_by_walk`, computed once; `merge` and `discard_rid` reset it _by_walk: dict[str, list[int]] | None = field(default=None, repr=False, compare=False) @@ -254,6 +257,12 @@ def __init__( # rid -> message, for requests a resource declared unservable during a # readiness check. Drained by the worker, which reports them onward. self.admit_errors: dict[int, str] = {} + # Rids an input read of which failed on this rank. A TP follower never + # gets those inputs, so it builds the leader's step without the row and + # the step's read check takes it out on every rank (see + # ``Worker._drop_rows_whose_read_failed``). + self.read_failed_rids: set[int] = set() + self._warned_read_check = False # lockstep-parallel (TP / SP instance) scheduling self.parallel_leader_nodes = parallel_leader_nodes @@ -522,14 +531,24 @@ def _try_schedule_tp_follow( # leader's graph walk. No step of the node is in flight on this path. if not self.settle_tp_follow_delta(caught_up=True): return # every pending move has to land before the step is built + rids = self.tp_rids(first_tp_node) + unread = self.unread_rows(first_tp_node, rids) popped = self.pop_ready_rids( - request_state, first_tp_node.node_name, - first_tp_node.graph_walk, self.tp_rids(first_tp_node), - walks=first_tp_node.walks, walk_idx=first_tp_node.walk_idx, + request_state, first_tp_node.node_name, first_tp_node.graph_walk, + [r for r in rids if r not in unread], + walks=first_tp_node.walks, + walk_idx=[ + w for r, w in zip(rids, first_tp_node.walk_idx, strict=True) if r not in unread + ] if first_tp_node.walks else None, ) if popped is None: + self.warn_unchecked_read_failure(first_tp_node, rids) return request_to_worker_graph, input_edges, output_signals, request_walks, walk_signals = popped + if unread: + request_to_worker_graph, request_walks = self.with_unread_rows( + first_tp_node, rids, unread, request_to_worker_graph, request_walks, + ) self.pop_tp_follow_head() @@ -542,6 +561,53 @@ def _try_schedule_tp_follow( tp_seq=first_tp_node.spec_seq, request_walks=request_walks, walk_output_signals=walk_signals, + read_check_rids=tuple(rids) if first_tp_node.read_check else None, + ) + + def unread_rows(self, head: ScheduleTPNode, rids: list[int]) -> set[int]: + """The rows of ``head`` whose input read failed on this rank, which a + follower puts in the step it builds with no inputs, for the step's read + check to take out on every rank. Only under a read check: without one no + rank takes the row out, so the follower pops it like any other (the + failed read may have been for another step).""" + failed = self.read_failed_rids + if not head.read_check or not failed: + return set() + return {r for r in rids if r in failed} + + def with_unread_rows( + self, head: ScheduleTPNode, rids: list[int], unread: set[int], + request_to_worker_graph: dict[int, int], request_walks: dict[int, str], + ) -> tuple[dict[int, int], dict[int, str]]: + """``request_to_worker_graph`` and ``request_walks`` with ``unread``'s + rows put back, in the leader's order: the step has rank 0's rows on + every rank until its read check, so it is built, submitted, voided and + settled alike. The unread rows were never popped.""" + walk_of = dict(zip(rids, (head.walks[w] for w in head.walk_idx), strict=True)) if head.walks else {} + request_walks = dict(request_walks) + rows = {} + for rid in rids: + if rid not in unread: + rows[rid] = request_to_worker_graph[rid] + continue + walk = walk_of.get(rid, head.graph_walk) + rows[rid] = self.runtime.get_worker_graph_id_for_node(head.node_name, walk) + if walk_of: + request_walks[rid] = walk + return rows, request_walks + + def warn_unchecked_read_failure(self, head: ScheduleTPNode, rids: list[int]) -> None: + """Once: a head with no read check waits on a row whose read failed + here, which it never gets. Only a read rank 0 did not mark gets here.""" + if self._warned_read_check or head.read_check or not self.read_failed_rids: + return + if self.read_failed_rids.isdisjoint(rids): + return + self._warned_read_check = True + logger.error( + "Node %s walk %s: a TP head with no read check names a request whose " + "input read failed on this rank; the instance waits on it", + head.node_name, head.graph_walk, ) @@ -1325,6 +1391,7 @@ def clear_rid(self, rid: int, rid_str: str) -> None: Takes both identities because ``pending_tp_follow_count`` is keyed by the wire string and everything else by the handle.""" self.failed_rids.discard(rid) + self.read_failed_rids.discard(rid) self.admit_errors.pop(rid, None) self.held_until.pop(rid, None) self._drop_backlogged_rid(rid) diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 164fea98ea..554469e938 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -7,16 +7,18 @@ import time import time as _time from collections import defaultdict, deque -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping from concurrent.futures import Future, ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, field from enum import Enum from functools import cached_property, partial from time import sleep +from types import MappingProxyType from typing import NamedTuple import torch +import torch.distributed as dist from mstar.communication.codec import WireCodec from mstar.communication.communicator import CommProtocol, make_communicator @@ -374,6 +376,12 @@ class Worker: # worker built without __init__ reads them as off too. _metrics: WorkerMetrics | None = None + # The lockstep node this worker leads, and per request the inputs of it + # that its followers read from elsewhere, not yet consumed by a step (see + # ``_setup_tp_read_check``). Empty on a bare instance. + _lockstep_led: frozenset[str] = frozenset() + _lockstep_unread: Mapping[int, set[str]] = MappingProxyType({}) + @cached_property def device_capability(self) -> DeviceCapability: return DeviceCapability.of(self.device) @@ -573,6 +581,7 @@ def __init__( ) self.is_tp_follower = len(self.parallel_nodes - self.parallel_leader_nodes) > 0 + self._setup_tp_read_check() # TP async scheduling: the leader speculates N+1 during forward N and # broadcasts it at once; followers rebuild it during their own N. @@ -1007,6 +1016,8 @@ def _remove_request(self, body: RemoveRequest) -> None: # once the api server finished reading the outputs. self._mark_removed(body.request_id) self.engine_manager.remove_request(request_id) + if self._lockstep_led: + self._lockstep_unread.pop(request_id, None) if self._stream_buffer_rids: self._stream_buffer_rids.discard(request_id) self.request_state.remove_request(request_id) @@ -1456,6 +1467,105 @@ def _poll_stream_buffers(self) -> None: self._mark_stream_ingested(rid, se) + # ------------------------------------------------------------------ + # TP read check: a follower's failed input read + # ------------------------------------------------------------------ + + def _setup_tp_read_check(self) -> None: + """Rank 0 of a lockstep instance schedules from its own copy of each + input, so a follower's read of one can still fail after rank 0 sent the + step. Rank 0 marks a step that takes such an input (``read_check``), and + before preparing it every rank takes out the rows whose read failed on + any rank (``_drop_rows_whose_read_failed``).""" + self._lockstep_led = frozenset(self.parallel_nodes & self.parallel_leader_nodes) + self._lockstep_unread = {} + + def _note_lockstep_reads(self, ready: dict[int, list[GraphEdge]]) -> None: + """Rank 0: the inputs of its lockstep node that landed here by a read, + which each follower reads too. This rank's own output of the node coming + back (a persisted token as the decode walk's input) is every rank's own + copy, with no read.""" + led, me = self._lockstep_led, self.tensor_manager.my_entity_id + for rid, edges in ready.items(): + for edge in edges: + if edge.next_node not in led or not edge.tensor_info: + continue + if all( + i._source_node_name in led and i.source_entity == me + for i in edge.tensor_info + ): + continue + self._lockstep_unread.setdefault(rid, set()).add(edge.name) + + def _note_local_lockstep_inputs(self, stored: StoredOutputs, signals: list[str]) -> None: + """Rank 0: outputs of a node outside the instance (an encoder beside + rank 0), which its lockstep node takes without a read and each follower + reads.""" + unread = self._lockstep_unread + for rid, sig in zip(stored.flat_rids, stored.signal_idxs, strict=True): + unread.setdefault(rid, set()).add(signals[sig]) + + def _step_reads(self, node_batch: ExecutingBatch) -> bool: + """Rank 0: whether a row of this step takes an input a follower reads, + so a follower's read of it can have failed.""" + unread, streamed = self._lockstep_unread, self._streaming_edge_names + if not unread and not streamed: + return False + tensors = node_batch.per_request_input_tensors + for rid in node_batch.request_ids: + names = tensors.get(rid) + if not names: + continue + # a stream chunk lands with no step to mark, so a chunk is always a read + if streamed and not streamed.isdisjoint(names): + return True + pending = unread.get(rid) + if pending and not pending.isdisjoint(names): + return True + return False + + def _consume_lockstep_reads(self, node_batch: ExecutingBatch) -> None: + """Rank 0, once a step of its lockstep node ran: the marked inputs its + rows took are spent. Not at the broadcast: a step that does not run (a + voided speculation, a refused admit) leaves its rows to a later one.""" + unread = self._lockstep_unread + tensors = node_batch.per_request_input_tensors + for rid in node_batch.request_ids: + pending = unread.get(rid) + if pending: + pending.difference_update(tensors.get(rid, ())) + if not pending: + del unread[rid] + + def _drop_rows_whose_read_failed(self, node_batch: ExecutingBatch) -> None: + """Before preparing a step rank 0 sent with ``read_check``: take out the + rows whose input read failed on any rank of the instance, on every rank. + A follower already left such a row out of what it built; the rank that + failed the read reported it, and the others fail it here, so its outputs + and its pages never exist on any rank. + + One gloo all-reduce of a flag per row (MAX over TP, then SP), in the + leader's row order, on the GPU thread where every rank prepares the + step. CPU tensors: it waits on the other ranks, never on the device.""" + rids = node_batch.read_check_rids + failed = self.scheduler.read_failed_rids + flags = torch.tensor([r in failed for r in rids], dtype=torch.int32) + node = node_batch.node_name + for group in ( + self.parallel_groups.get_tp_config_for_node(node), + self.parallel_groups.get_sp_config_for_node(node), + ): + if group.world_size > 1: + dist.all_reduce(flags, op=dist.ReduceOp.MAX, group=group.cpu_group) + flagged = {r for r, f in zip(rids, flags.tolist(), strict=True) if f} + if not flagged: + return + for rid in flagged: + node_batch.failed_requests.setdefault( + rid, "Error in worker: input tensor read failed on a rank of its TP group", + ) + node_batch.drop_rids(flagged) + def _check_ready_tensors(self) -> None: """Poll for completed RDMA transfers, feed ready graph edges to worker graph queues.""" self.wakeup_event.drain() @@ -1466,7 +1576,13 @@ def _check_ready_tensors(self) -> None: if f.request_id not in self.scheduler.failed_rids } if failed: + # only a TP step's read check asks (see ``unread_rows``) + self.scheduler.read_failed_rids.update( + rid for rid in failed if rid in self.request_state.per_request_info + ) self._fail_requests(failed) + if ready and self._lockstep_led: + self._note_lockstep_reads(ready) for request_id, edges in ready.items(): # Separate streaming edges from normal edges streaming = [e for e in edges if e.is_streaming] @@ -1643,6 +1759,7 @@ def _build_executing_batch(self, batch: ScheduledBatch) -> ExecutingBatch: per_request_input_metadata=per_request_input_metadata, request_walks=batch.request_walks, ) + node_batch.read_check_rids = batch.read_check_rids for rid in unresolved: # Reported the way a per-rid stage reports: left in # ``failed_requests`` for the run loop to drop and fail, so the @@ -1745,6 +1862,9 @@ def maybe_send_zmq_to_tp_followers( resident_delta = engine.take_resident_delta(node_batch.node_name) prefix_hits = engine.decide_cached_prefix(node_batch) page_trails = engine.page_trails(node_batch.node_name) + read_check = self._step_reads(node_batch) + if read_check: + node_batch.read_check_rids = tuple(node_batch.request_ids) # this worker is only a part of one TP group for this node, # so, we can just look at the sharding_config for the first # request to get the relevant workers @@ -1775,6 +1895,7 @@ def maybe_send_zmq_to_tp_followers( walk_idx=walk_idx, prefix_hits=prefix_hits, page_trails=page_trails, + read_check=read_check, ) ) ) @@ -2140,6 +2261,9 @@ def _execute_on_gpu_thread( ) try: try: + if node_batch.read_check_rids is not None: + with self._span("worker.gpu_thread.read_check"): + self._drop_rows_whose_read_failed(node_batch) with self._span( "worker.gpu_thread.prepare_inputs", cpu="worker.gpu_thread.prepare_inputs.cpu", @@ -2453,6 +2577,10 @@ def _plans_ahead(self, spec: Speculation) -> bool: do. A pre-plan refused on rank 0 alone puts its pages back at the end of rank 0's free list, and the ranks then hand out different page ids, which only the cache reads.""" node = spec.node_batch.node_name + if spec.node_batch.read_check_rids is not None: + # every rank admits a read-checked step inline, after the check: a follower + # that could not read a row pre-plans without it, rank 0 with it + return False return not (spec.is_yield_away and node in self.parallel_nodes and self._prefix_cache_open(node)) def _prefix_cache_open(self, node_name: str) -> bool: @@ -2961,18 +3089,23 @@ def _try_follow_speculation(self, pending: PendingBatch) -> Speculation | None: batch_N = pending.batch # The leader names its rids by wire string; these are this rank's handles. head_rids = self.scheduler.tp_rids(head) + # rows whose input read failed here: built without, as on the serial path + unread = self.scheduler.unread_rows(head, head_rids) head_walk_idx = dict(zip(head_rids, head.walk_idx, strict=True)) if head.walks else None in_batch = batch_N.request_to_worker_graph - continuing = [r for r in head_rids if r in in_batch] - fresh = [r for r in head_rids if r not in in_batch] - if not continuing: + # rows read here: prepped (continuing) or popped (fresh); every row, + # unread ones included, runs the leader's continuation of N + in_n = [r for r in head_rids if r in in_batch] + continuing = [r for r in in_n if r not in unread] + fresh = [r for r in head_rids if r not in in_batch and r not in unread] + if not in_n: return None - spec_walk = head.walks[head_walk_idx[continuing[0]]] if head.walks else head.graph_walk + spec_walk = head.walks[head_walk_idx[in_n[0]]] if head.walks else head.graph_walk # The leader already chose the target; this just reports its loop # context. speculate_node's eligibility filter cannot run here -- it # requires the node be a parallel LEADER node. spec_target_info = self._graph_runtime.get_spec_target( - batch_N.node_name, head.node_name, spec_walk, continuing[0], + batch_N.node_name, head.node_name, spec_walk, in_n[0], ) if spec_target_info is None: return None @@ -3003,6 +3136,7 @@ def _return_all_polled() -> None: # rid reads not-ready and waits for a replayed delta. The delta is # settled above, so a rid still off-device is a divergence, not a lag. _return_all_polled() + self.scheduler.warn_unchecked_read_failure(head, fresh) return None # output_signals are handled by _assemble_speculation @@ -3061,6 +3195,15 @@ def _return_all_polled() -> None: if rid in prep_inputs: new_request_to_worker_graph[rid] = prep_wg[rid] per_request_inputs[rid] = prep_inputs[rid] + elif rid in unread: + # In with no inputs, as on the serial path: the step has rank 0's + # rows until its read check takes this one out. + new_request_to_worker_graph[rid] = in_batch[rid] if rid in in_batch else \ + self._graph_runtime.get_worker_graph_id_for_node( + head.node_name, + head.walks[head_walk_idx[rid]] if head.walks else head.graph_walk, + ) + per_request_inputs[rid] = {} else: new_request_to_worker_graph[rid] = fresh_wg[rid] per_request_inputs[rid] = fresh_inputs[rid] @@ -3074,15 +3217,18 @@ def _return_all_polled() -> None: "Follow-speculating: %s %s (seq %d from %d)", head.node_name, head_rids, head.spec_seq, head.spec_from_seq, ) - return self._assemble_speculation( + speculation = self._assemble_speculation( pending, spec_target_info, new_request_to_worker_graph, per_request_inputs, - consumed_streaming_edges, continuing, + consumed_streaming_edges, in_n, is_same_node=True, spec_id=prep.spec_id, tp_seq=head.spec_seq, request_walks={ r: head.walks[w] for r, w in head_walk_idx.items() } if head.walks else None, ) + if head.read_check and speculation is not None: + speculation.node_batch.read_check_rids = tuple(head_rids) + return speculation # How many "no speculative head from step s" seqs a follower remembers. _TP_NOSPEC_KEEP = 1024 @@ -3246,6 +3392,8 @@ def _pp_stage(name: str) -> None: range_push("worker.postprocess.cleanup_inputs", synchronize=False) rids = list(batch_N.batch.request_to_worker_graph) + if self._lockstep_unread and batch_N.node_name in self._lockstep_led: + self._consume_lockstep_reads(batch_N.node_batch) # What the runtime dereferenced to zero and cannot reclaim itself: a # runtime behind the contract has the bookkeeper, not the shm files or # the registered memory. @@ -3580,6 +3728,8 @@ def store(rids): stored = stored._replace(num_tensors=counts) else: stored = store(rids) + if self._lockstep_led and batch_N.node_name not in self.parallel_nodes: + self._note_local_lockstep_inputs(stored, signals) flat_uuids = stored.flat_uuids flat_rids = stored.flat_rids signal_idxs = stored.signal_idxs diff --git a/test/modular/test_kv_prefix_tp.py b/test/modular/test_kv_prefix_tp.py index dd0ac37daa..5c28241776 100644 --- a/test/modular/test_kv_prefix_tp.py +++ b/test/modular/test_kv_prefix_tp.py @@ -354,7 +354,9 @@ def test_a_tp_leader_admits_a_yield_away_step_unplanned_as_its_followers_do(cach leader = SimpleNamespace(parallel_nodes={NODE}, _prefix_cache_open=lambda node: cache_open) def spec(node, yield_away): - return SimpleNamespace(is_yield_away=yield_away, node_batch=SimpleNamespace(node_name=node)) + return SimpleNamespace( + is_yield_away=yield_away, node_batch=SimpleNamespace(node_name=node, read_check_rids=None), + ) assert Worker._plans_ahead(leader, spec(NODE, True)) is not cache_open, "only the cache reads page ids" assert Worker._plans_ahead(leader, spec(NODE, False)), "the head is pre-planned on every rank" diff --git a/test/modular/test_preplan_error_paths.py b/test/modular/test_preplan_error_paths.py index 69a8d2d7ab..239e472d9a 100644 --- a/test/modular/test_preplan_error_paths.py +++ b/test/modular/test_preplan_error_paths.py @@ -131,7 +131,7 @@ def test_gpu_thread_drops_the_stage_when_prepare_inputs_raises() -> None: w.engine_manager = SimpleNamespace(get_engine=lambda name: engine) released = [] node_batch = SimpleNamespace( - node_name="dit", launch_started_event=None, preplanned_rids=("a",), + node_name="dit", launch_started_event=None, preplanned_rids=("a",), read_check_rids=None, request_ids=("a",), release_waiters=lambda: released.append(True), per_request_info={"a": SimpleNamespace(request_id="a")}, ) diff --git a/test/modular/test_request_failure_propagation.py b/test/modular/test_request_failure_propagation.py index 044a662dc4..49801bda32 100644 --- a/test/modular/test_request_failure_propagation.py +++ b/test/modular/test_request_failure_propagation.py @@ -74,6 +74,7 @@ def _worker(known_rids=("r1", "r2")): w._graph_runtime = _Runtime(known_rids) w.scheduler = MicroScheduler.__new__(MicroScheduler) w.scheduler.failed_rids = set() + w.scheduler.read_failed_rids = set() w.scheduler.held_until = {} w.scheduler.admit_errors = {} w.scheduler.backlog = {} diff --git a/test/modular/test_resident_delta.py b/test/modular/test_resident_delta.py index e557e837b2..6b14183654 100644 --- a/test/modular/test_resident_delta.py +++ b/test/modular/test_resident_delta.py @@ -598,6 +598,7 @@ class _Leader: maybe_send_zmq_to_tp_followers = Worker.maybe_send_zmq_to_tp_followers _rid_str = Worker._rid_str + _step_reads = Worker._step_reads def __init__(self, followers=("rank1", "rank2")): self.engine = _Engine("leader") @@ -622,6 +623,9 @@ def __init__(self, followers=("rank1", "rank2")): self.communicator = SimpleNamespace( send=lambda worker, msg: self.sent.append((worker, msg)) ) + # nothing the followers read: no step goes out read-checked + self._lockstep_unread: dict = {} + self._streaming_edge_names: set[str] = set() def _node_batch(): @@ -923,6 +927,7 @@ class _ForwardingLeader: def __init__(self, broadcast_seq: int): self.is_tp_follower = False + self._lockstep_led = frozenset() self._tp_broadcast_seq = broadcast_seq # handle-keyed state, per rid_table's list self._in_flight_rids: set[int] = set() diff --git a/test/modular/test_tp_async_lockstep_sim.py b/test/modular/test_tp_async_lockstep_sim.py index 7147cda74d..0ba95524f3 100644 --- a/test/modular/test_tp_async_lockstep_sim.py +++ b/test/modular/test_tp_async_lockstep_sim.py @@ -72,3 +72,46 @@ def test_b2_retract_is_refuted_on_structural_cancel(): def test_b2_waste_matches_design_prediction(): # expected run-ahead cost: one wasted forward per request end at B=1 assert _run("B2_VOID", "clean-3step").max_waste == 1 + + +# --- a follower's input read failing after rank 0 scheduled the row ---------- +# Model: tp_read_check_sim.py. It explores every order of the follower's read +# outcome, rank 0's scheduling and the lockstep step. + +try: + from .tp_read_check_sim import explore as explore_read_check +except ImportError: # collected without package context + from tp_read_check_sim import explore as explore_read_check + + +def test_the_read_check_holds_for_every_input_kind(): + """Every rank takes out the rows whose read failed before preparing the + step: no interleaving hangs, every step runs the same collectives on both + ranks, and only the request whose read failed fails.""" + for kind in ("data", "shape", "prepare"): + res = explore_read_check("CHECK", kind) + assert res.ok, (kind, res.violations[:1]) + + +def test_scheduling_from_rank_0s_reads_alone_hangs_the_instance(): + """A step rank 0 broadcast can need a read that then fails on the + follower, which can never run it.""" + for kind in ("data", "shape", "prepare"): + res = explore_read_check("NONE", kind) + assert any(v == "live" for v, _ in res.violations), kind + + +def test_zero_fill_desyncs_the_step_when_the_input_sets_its_length(): + """Zeros of the advertised shape keep a data input's step whole, but an + input the step derives its length from (image_grid_thw) gives the follower + a different step than rank 0's.""" + assert explore_read_check("ZERO", "data").ok + assert any(v == "lockstep" for v, _ in explore_read_check("ZERO", "shape").violations) + + +def test_a_capture_input_row_desyncs_the_step_when_prepare_runs_a_collective(): + """The walk's capture input at rank 0's length fixes the length, but rank 0 + prepared the real row, and a prepare that runs a collective for it (a text + part's vocab-parallel embedding) has no partner on the follower.""" + assert explore_read_check("CAPTURE", "shape").ok + assert any(v == "lockstep" for v, _ in explore_read_check("CAPTURE", "prepare").violations) diff --git a/test/modular/test_tp_async_teardown_order.py b/test/modular/test_tp_async_teardown_order.py index 9e8164caba..513d80b659 100644 --- a/test/modular/test_tp_async_teardown_order.py +++ b/test/modular/test_tp_async_teardown_order.py @@ -57,7 +57,7 @@ def _batch(seq: int, rids: list[int]): ) executing = SimpleNamespace( node_name=NODE, request_ids=list(rids), admit_error=None, failed_requests={}, - launch_started_event=None, + launch_started_event=None, read_check_rids=None, ) return scheduled, executing diff --git a/test/modular/test_tp_read_check.py b/test/modular/test_tp_read_check.py new file mode 100644 index 0000000000..4ea6921d8b --- /dev/null +++ b/test/modular/test_tp_read_check.py @@ -0,0 +1,576 @@ +"""A TP follower's input read fails after rank 0 scheduled the row from its own +copy (the TP side of #418). + +Rank 0 marks a step whose rows take an input the followers read from elsewhere +(``ScheduleTPNode.read_check``). A follower builds such a step once each row's +read landed or failed, leaving out the rows it could not read, and before any +rank prepares the step one all-reduce of a flag per row tells every rank which +rows to take out. Those rows never run, so nothing of them is computed, +committed or sent, and every rank runs the same collectives. + +- The two-rank runs reuse the acceptance harness + (``test_tp_read_failure_acceptance.py``) with a ``prepare_inputs`` that runs a + collective per text part, as Qwen3.5's vocab-parallel embedding does, and + with a step whose every row failed. +- TP4 and TP2 x SP2: four real gloo ranks run the check itself. +- The follower's scheduling (serial and ``MSTAR_TP_ASYNC_SCHED``) and rank 0's + marking run the real ``MicroScheduler`` and ``Worker`` methods on stubs. +""" + +from __future__ import annotations + +import datetime +import multiprocessing as mp +import os +import queue +import traceback +from concurrent.futures import Future +from types import SimpleNamespace + +import pytest +import test_tp_read_failure_acceptance as acc # the two-rank harness +import torch +import torch.nn.functional as F +from test_failed_tensor_read import _failed_future, _read, _shm_manager +from test_micro_scheduler import NODE, WALK, _Engine, _Manager + +import mstar.communication.wire_types # noqa: F401 (registers the tags) +from mstar.communication.tensors import StoredOutputs +from mstar.communication.wire import decode, encode +from mstar.engine.engine import ExecutingBatch +from mstar.engine.resources.step import StepContext +from mstar.graph.base import GraphEdge, TensorPointerInfo +from mstar.graph.runtime.base import ColumnarEdgeSpecs, SpeculationPrepOutput +from mstar.utils.ipc_format import OffloadDelta, ScheduleTPNode, WorkerMessageType +from mstar.worker.micro_scheduler import MicroScheduler, ScheduledBatch +from mstar.worker.worker import Speculation, Worker + +# ── two ranks through Worker.run (the acceptance harness) ───────────────── + + +def _vocab_parallel_submodule_class(harness_class): + """The harness's Qwen3.5 LLM submodule, with an embedding that all-reduces + over the TP group as ``VocabParallelEmbedding`` does. Qwen3.5's + ``prefill_vision`` embeds each text part (and its vision sentinels) inside + ``prepare_inputs``, so a rank that prepares a row another rank does not + runs a collective with no partner.""" + base = harness_class() + + class _Embedding(torch.nn.Module): + def __init__(self, weight, group, rank): + super().__init__() + self.weight, self.group, self.rank = weight, group, rank + + def forward(self, ids): + partial = F.embedding(ids, self.weight) * ((self.rank + 1) / 3.0) + return self.group.all_reduce(partial) + + class _LLM(base): + def __init__(self, group, rank): + super().__init__(group, rank) + self.model.model.embed_tokens = _Embedding( + self.model.model.embed_tokens.weight, group, rank, + ) + + return _LLM + + +def _rank_main_vocab_parallel(rank, scenario, inbox, outbox): + harness_class = acc._submodule_class + acc._submodule_class = lambda: _vocab_parallel_submodule_class(harness_class) + acc._rank_main(rank, scenario, inbox, outbox) + + +class _EveryLateReadFails(acc._Reads): + """The follower's reads of both requests resolve late, and its read of + each one's failing input fails: rank 0's step holds only rows this rank + could not read.""" + + def __init__(self, tm, kind, rank, scenario): + from mstar.communication.tensors import FutureAndPointers + from mstar.graph.base import GraphEdge as Edge + self.tm, self.late, self.fired = tm, [], False + uuid = 1000 * (rank + 1) + for rid in ("r0", "r1"): + for signal, tensors in acc._inputs(kind, rid).items(): + infos, buffers = [], [] + for t in tensors: + uuid += 1 + info = acc._pointer(t, uuid) + buf = torch.zeros_like(t) + tm.tensor_store.put_tensor(rid, uuid, buf, info) + infos.append(info) + buffers.append((buf, t)) + future = Future() + tm.pending.append(FutureAndPointers( + future=future, graph_edges=[Edge(next_node=acc.NODE, name=signal, tensor_info=infos)], + request_id=rid, + )) + if rank == 1: + self.late.append((future, buffers, signal == acc._FAILING[kind])) + else: + self._land(future, buffers) + + +def _rank_main_every_row_fails(rank, scenario, inbox, outbox): + acc._Reads = _EveryLateReadFails + acc._rank_main(rank, scenario, inbox, outbox) + + +@pytest.mark.parametrize("kind", acc.KINDS) +def test_a_prepare_that_runs_a_collective_per_row_stays_in_step(tmp_path, monkeypatch, kind): + """With the embedding all-reducing inside ``prepare_inputs``, r1's failed + follower read still fails only r1, and r0's outputs equal a run with no + failure: no rank prepares r1, so no rank runs r1's embedding collectives.""" + monkeypatch.setattr(acc, "_rank_main", _rank_main_vocab_parallel) + clean = acc._run(tmp_path / "clean", acc.Scenario(kind=kind, fail=False)) + acc._healthy_ok(clean) + assert clean.completed("r1") and not clean.failed, clean.summary() + run = acc._run(tmp_path / "failed", acc.Scenario(kind=kind)) + acc._healthy_ok(run) + assert set(run.failed) == {"r1"} and run.torn_down("r1"), run.summary() + for rank in (0, 1): + assert "r1" not in run.routed[rank], run.summary() + assert run.routed[rank]["r0"] == clean.routed[rank]["r0"], f"rank {rank}" + + +def test_a_step_whose_every_row_failed_runs_empty_on_every_rank(tmp_path, monkeypatch): + """Rank 0 sends a step of rows the follower could read none of: both ranks + run it empty, both requests fail and drain, and nothing is routed.""" + monkeypatch.setattr(acc, "_rank_main", _rank_main_every_row_fails) + run = acc._run(tmp_path, acc.Scenario(kind="text")) + assert not run.crashed, run.summary() + "\n" + "\n".join(map(str, run.crashed.values())) + assert not run.hung, run.summary() + assert set(run.failed) == {"r0", "r1"}, run.summary() + assert run.torn_down("r0") and run.torn_down("r1"), run.summary() + assert run.routed == {0: {}, 1: {}}, run.summary() + # no rank sat in a collective the other never joined + assert run.elapsed < acc.COLLECTIVE_S, run.summary() + + +# ── the check itself at TP4 and TP2 x SP2 (four gloo ranks) ─────────────── + +ROWS = (10, 11, 12, 13) + + +def _groups(rank, layout): + """This rank's TP and SP groups: ``tp4`` is one group of four; ``tp2sp2`` + puts ranks {0,1} and {2,3} in TP groups and {0,2} and {1,3} in SP groups.""" + import torch.distributed as dist + + from mstar.distributed.communication import CommGroup + + def group(members, pg): + g = CommGroup(my_global_rank=rank, my_group_rank=members.index(rank), group_members=members) + g.device_group = g.cpu_group = pg + return g + + if layout == "tp4": + return group([0, 1, 2, 3], dist.group.WORLD), CommGroup.trivial() + pgs = {m: dist.new_group(list(m)) for m in ((0, 1), (2, 3), (0, 2), (1, 3))} + tp = next(m for m in ((0, 1), (2, 3)) if rank in m) + sp = next(m for m in ((0, 2), (1, 3)) if rank in m) + return group(list(tp), pgs[tp]), group(list(sp), pgs[sp]) + + +def _check_rank(rank, layout, failed_here, store, out): + import torch.distributed as dist + try: + dist.init_process_group( + "gloo", init_method=f"file://{store}", rank=rank, world_size=4, + timeout=datetime.timedelta(seconds=10), + ) + tp, sp = _groups(rank, layout) + w = Worker.__new__(Worker) + w.scheduler = SimpleNamespace(read_failed_rids=set(failed_here)) + w.parallel_groups = SimpleNamespace( + get_tp_config_for_node=lambda node: tp, get_sp_config_for_node=lambda node: sp, + ) + # a follower that could not read a row left it out; the others hold it + rows = [r for r in ROWS if r not in failed_here] + batch = ExecutingBatch( + node_name=NODE, per_request_info=dict.fromkeys(rows, SimpleNamespace()), + step_context=StepContext(request_ids=tuple(rows), graph_walk=WALK, slot=0, capture=False), + ) + batch.read_check_rids = ROWS + w._drop_rows_whose_read_failed(batch) + out.put((rank, list(batch.request_ids), sorted(batch.failed_requests))) + dist.destroy_process_group() + except BaseException: + out.put((rank, "crash", traceback.format_exc())) + + +@pytest.mark.parametrize("layout", ["tp4", "tp2sp2"]) +def test_every_rank_takes_out_the_rows_any_rank_could_not_read(tmp_path, layout): + """Rank 2's read of row 11 and rank 3's of row 13 failed. Every rank, rank + 0 and rank 1 included, takes out both rows and fails both, and keeps the + other two in rank 0's order. Under TP2 x SP2 rank 1 shares no group with + rank 2, so it learns of row 11 only through the SP pass.""" + ctx = mp.get_context("spawn") + out = ctx.Queue() + failed = {2: {11}, 3: {13}} + procs = [ + ctx.Process(target=_check_rank, args=(r, layout, failed.get(r, set()), tmp_path / "store", out)) + for r in range(4) + ] + for p in procs: + p.start() + results = {} + try: + for _ in procs: + rank, *rest = out.get(timeout=60) + results[rank] = rest + except queue.Empty: + pass + for p in procs: + p.join(timeout=10) + if p.is_alive(): + p.kill() + assert sorted(results) == [0, 1, 2, 3], results + for rank, rest in results.items(): + assert rest[0] != "crash", rest[1] + assert rest == [[10, 12], [11, 13]], (rank, rest) + + +# ── a follower builds rank 0's step without the rows it could not read ──── + + +def _follower_scheduler(ready, read_failed=()): + manager = _Manager(["r0", "r1"]) + manager.queues["wg0"]._ready = {rid: {NODE} for rid in ready} + sched = MicroScheduler( + engine_manager=SimpleNamespace(get_engine=lambda n: _Engine()), + parallel_leader_nodes=set(), + ) + sched.runtime = manager.runtime + sched.read_failed_rids.update(read_failed) + return sched, manager + + +def _head(read_check, rids=("r0", "r1"), **kwargs): + return ScheduleTPNode( + node_name=NODE, graph_walk=WALK, request_ids=list(rids), spec_seq=0, + read_check=read_check, **kwargs, + ) + + +def test_a_follower_builds_a_read_checked_step_with_the_row_it_could_not_read_empty(): + """r1's read failed here, so r1 never readies. The follower takes the head + without waiting: it pops r0 alone and puts r1 in with no inputs, so the step + has rank 0's rows, in its order, until the check takes r1 out. The head no + longer holds r1's drain.""" + sched, manager = _follower_scheduler(ready=["r0"], read_failed=["r1"]) + sched.register_tp_follow(_head(read_check=True, rids=("r1", "r0"))) + + batch = sched.get_next_batch(manager) + + assert list(batch.request_to_worker_graph) == ["r1", "r0"] + assert batch.read_check_rids == ("r1", "r0") + assert manager.pops[-1][1] == ["r0"] + assert list(batch.input_edges.rids) == ["r0"] + assert sched.peek_tp_follow() is None + assert "r1" not in sched.pending_tp_follow_count + + +def test_a_follower_takes_a_step_whose_every_row_it_could_not_read(): + """The step still has rank 0's rows, so this rank submits it and joins its + check, as rank 0 does.""" + sched, manager = _follower_scheduler(ready=[], read_failed=["r0", "r1"]) + sched.register_tp_follow(_head(read_check=True)) + + batch = sched.get_next_batch(manager) + + assert list(batch.request_to_worker_graph) == ["r0", "r1"] + assert batch.read_check_rids == ("r0", "r1") + assert sched.peek_tp_follow() is None + + +def test_a_follower_waits_on_a_row_whose_read_is_still_running(): + sched, manager = _follower_scheduler(ready=["r0"]) + sched.register_tp_follow(_head(read_check=True)) + + assert sched.get_next_batch(manager) is None + assert sched.peek_tp_follow() is not None + + manager.ingest("r1") + batch = sched.get_next_batch(manager) + assert list(batch.request_to_worker_graph) == ["r0", "r1"] + + +def test_without_a_read_check_a_follower_pops_the_row_like_any_other(caplog): + """No rank would take the row out, so the follower does not either: a row + whose failed read was another step's input runs; one whose input never + lands waits, as before the check, and says why once.""" + sched, manager = _follower_scheduler(ready=["r0", "r1"], read_failed=["r1"]) + sched.register_tp_follow(_head(read_check=False)) + batch = sched.get_next_batch(manager) + assert list(batch.request_to_worker_graph) == ["r0", "r1"] + assert batch.read_check_rids is None + + sched, manager = _follower_scheduler(ready=["r0"], read_failed=["r1"]) + sched.register_tp_follow(_head(read_check=False)) + with caplog.at_level("ERROR"): + assert sched.get_next_batch(manager) is None + assert sched.get_next_batch(manager) is None + assert sum("no read check" in r.getMessage() for r in caplog.records) == 1 + + +def test_a_removed_requests_read_failure_is_forgotten(): + """The flag is keyed by the handle, which a later request can get.""" + sched, _ = _follower_scheduler(ready=[], read_failed=["r1"]) + sched.clear_rid("r1", "r1") + assert "r1" not in sched.read_failed_rids + + +def _spec_follower(read_failed, continuing, fresh, head_rids): + """An async follower with step N (``continuing``) in flight and rank 0's + speculative head for N+1 queued. The graph runtime's prep and the + speculation's assembly are recorded, not run.""" + # a row whose read failed never readies + sched, manager = _follower_scheduler( + ready=[r for r in fresh if r not in read_failed], read_failed=read_failed, + ) + w = Worker.__new__(Worker) + w.scheduler, w.request_state, w.enable_nvtx = sched, manager, False + w._in_flight_rids = set() + w.preps, w.assembled = [], [] + + def prep(prep_input): + w.preps.append(list(prep_input.rids)) + return SpeculationPrepOutput([], list(prep_input.rids), ["wg0"] * len(prep_input.rids), + ColumnarEdgeSpecs.empty()) + + w._graph_runtime = SimpleNamespace( + get_spec_target=lambda *a: SimpleNamespace(), + prep_follow_spec_rids=prep, + push_back_node=lambda node, rids, wg_ids: ( + w.pushed_back.extend(rids), manager.runtime.push_back_node(node, rids, wg_ids), + ), + commit_speculation=lambda *a, **k: None, + get_worker_graph_id_for_node=manager.runtime.get_worker_graph_id_for_node, + ) + w.pushed_back = [] + w.released = [] + w.engine_manager = SimpleNamespace(get_engine=lambda n: SimpleNamespace( + release_cached_prefix=lambda batch: w.released.extend(batch.request_ids), + )) + w.tensor_manager = SimpleNamespace(get_tensor=lambda uuid: torch.zeros(1)) + + w.continuing, w.inputs = [], [] + + def assemble(pending, info, rows, inputs, streamed, continuing, **kwargs): + w.assembled.append(list(rows)) + w.continuing.append(list(continuing)) + w.inputs.append(dict(inputs)) + rids = list(rows) + return Speculation( + scheduled_batch=ScheduledBatch(node_name=NODE, graph_walk=WALK, request_to_worker_graph=dict(rows)), + node_batch=ExecutingBatch( + node_name=NODE, per_request_info=dict.fromkeys(rids, SimpleNamespace()), + step_context=StepContext(request_ids=tuple(rids), graph_walk=WALK, slot=0, capture=False), + ), + consumed_edges=set(), continuing_rids=set(), partition="default", + is_new_iter=False, is_same_node=True, + ) + + w._assemble_speculation = assemble + pending = SimpleNamespace( + tp_seq=0, node_name=NODE, graph_walk=WALK, + batch=ScheduledBatch(node_name=NODE, graph_walk=WALK, + request_to_worker_graph=dict.fromkeys(continuing, "wg0")), + node_batch=SimpleNamespace(commit_done=SimpleNamespace(is_set=lambda: True)), + ) + sched.register_tp_follow(ScheduleTPNode( + node_name=NODE, graph_walk=WALK, request_ids=list(head_rids), spec_seq=1, + speculative=True, spec_from_seq=0, read_check=True, + )) + return w, pending + + +@pytest.mark.parametrize("unread", ["fresh", "continuing", "every"]) +def test_an_async_follower_builds_a_read_checked_head_with_the_row_empty(unread): + """``MSTAR_TP_ASYNC_SCHED``: the follower builds a speculative head's + continuing rows (prepped from N) and fresh rows (popped) apart. A row it + could not read is neither: it goes in with no inputs, so the speculation + has rank 0's rows (a continuing one still takes N's loop-back, and its + loop count), and is submitted, voided and settled as rank 0's is.""" + head = ["c0", "c1", "f0", "f1"] + failed = {"fresh": ["f1"], "continuing": ["c1"], "every": head}[unread] + w, pending = _spec_follower( + read_failed=failed, continuing=["c0", "c1"], fresh=["f0", "f1"], head_rids=head, + ) + + spec = Worker._try_follow_speculation(w, pending) + + assert spec is not None, "the head waits on a row this rank will never read" + assert w.preps == [[r for r in ("c0", "c1") if r not in failed]] + assert w.assembled == [head] + assert w.continuing == [["c0", "c1"]] + assert all(w.inputs[-1][r] == {} for r in failed) + assert spec.node_batch.read_check_rids == tuple(head) + assert w.scheduler.peek_tp_follow() is None + # a read-checked step is admitted and planned inline on every rank, after the check + assert not Worker._plans_ahead(w, spec) + + # N's verdict voids the head: rank 0 gives back the leases of every row it + # sent, and so does this rank, the unread rows' included + Worker._settle_speculation(w, spec, success=False) + assert sorted(w.released) == sorted(head) + + +def test_an_async_follower_that_cannot_prep_gives_back_only_the_rows_it_popped(): + """The continuing rows' prep fails (a chunk not landed yet): the follower + hands back the fresh row it popped, not the row it could not read, which it + never popped.""" + w, pending = _spec_follower( + read_failed=["f1"], continuing=["c0"], fresh=["f0", "f1"], head_rids=["c0", "f0", "f1"], + ) + w._graph_runtime.prep_follow_spec_rids = lambda prep_input: None + + assert Worker._try_follow_speculation(w, pending) is None + assert w.pushed_back == ["f0"] + assert w.scheduler.peek_tp_follow() is not None + + +# ── rank 0 marks the steps whose rows take an input a follower reads ────── + +FOLLOWERS = ["w0", "w1"] + + +def _leader(tmp_path, streamed=()): + tm, _ = _shm_manager(tmp_path) + w = Worker.__new__(Worker) + w.worker_id = "w0" + w.parallel_nodes, w.parallel_leader_nodes = {NODE}, {NODE} + w.is_tp_follower, w.enable_nvtx = False, False + w._tp_broadcast_seq = 0 + w.sent = [] + w.communicator = SimpleNamespace( + send=lambda dest, msg=None: w.sent.append((dest, decode(encode(msg)))), + ) + group = SimpleNamespace(nodes={NODE}, tp_size=2, _tp_rank=0, _workers=FOLLOWERS) + sharding = SimpleNamespace(groups=[group], get_sharding_group=lambda n, g: group) + w._graph_runtime = SimpleNamespace( + get_sharding_config=lambda r: sharding, get_rid_string=lambda h: h, + get_rid_handle=lambda r: r, ingest_inputs_batch=lambda block, can_buffer: None, + ) + w.engine_manager = SimpleNamespace(get_engine=lambda n: SimpleNamespace( + take_resident_delta=lambda n: OffloadDelta.new(), + decide_cached_prefix=lambda b: {}, page_trails=lambda n: {}, + )) + w.tensor_manager = tm + w.scheduler = MicroScheduler( + engine_manager=w.engine_manager, parallel_leader_nodes={NODE}, + ) + w.request_state = SimpleNamespace(per_request_info={"r0": None, "r1": None}) + w.wakeup_event = SimpleNamespace(drain=lambda: None) + w._streaming_edge_names = set(streamed) + w._setup_tp_read_check() + return w + + +def _input(uuid, name, source_entity="api_server", source_node=None, node=NODE): + t = torch.ones(3, dtype=torch.int64) + info = TensorPointerInfo( + list(t.shape), t.dtype, t.nbytes, 0, list(t.stride()), uuid, "s", source_entity, + _source_node_name=source_node, + ) + return GraphEdge(next_node=node, name=name, tensor_info=[info]) + + +def _land(w, rid, edge): + """``edge`` lands on ``w`` through its real read path.""" + done = Future() + done.set_result(None) + for info in edge.tensor_info: + w.tensor_manager.tensor_store.put_tensor(rid, info.uuid, torch.zeros(3, dtype=torch.int64), info) + w.tensor_manager.pending.append(_read(rid, edge, done)) + w._check_ready_tensors() + + +def _step(rows): + """A step whose rows take ``rows[rid]``'s input names.""" + rids = list(rows) + return ExecutingBatch( + node_name=NODE, per_request_info=dict.fromkeys(rids, SimpleNamespace()), + step_context=StepContext(request_ids=tuple(rids), graph_walk=WALK, slot=0, capture=False), + per_request_input_tensors={ + rid: {name: [torch.zeros(1)] for name in names} for rid, names in rows.items() + }, + ) + + +def _broadcast(w, batch): + w.maybe_send_zmq_to_tp_followers(batch) + (msg,) = [m for d, m in w.sent if m.message_type == WorkerMessageType.SCHEDULE_TP] + w.sent.clear() + return msg.body.read_check, batch.read_check_rids + + +def test_a_step_taking_an_input_read_from_elsewhere_is_checked_until_it_runs(tmp_path): + """r1's prompt landed on rank 0 by a read; the step that takes it goes out + checked, with its rows in order. A step that never ran (a voided + speculation, a refused admit) leaves the row to a later one, which is + checked too. Once one ran, the next step's input is the node's own output, + which no rank reads.""" + w = _leader(tmp_path) + _land(w, "r1", _input(11, "text_inputs")) + rows = {"r0": ["text_inputs"], "r1": ["text_inputs"]} + + assert _broadcast(w, _step(rows)) == (True, ("r0", "r1")) + ran = _step(rows) + assert _broadcast(w, ran) == (True, ("r0", "r1")) + w._consume_lockstep_reads(ran) + assert _broadcast(w, _step(rows)) == (False, None) + + +def test_the_nodes_own_output_coming_back_is_no_read(tmp_path): + """A persisted token returning as the decode walk's input is every rank's + own copy; an input for another node, or with no tensor, has no read either.""" + w = _leader(tmp_path) + me = w.tensor_manager.my_entity_id + _land(w, "r1", _input(12, "text_inputs", source_entity=me, source_node=NODE)) + _land(w, "r1", _input(13, "text_inputs", node="vision_encoder")) + _land(w, "r1", GraphEdge(next_node=NODE, name="text_inputs")) + + assert _broadcast(w, _step({"r1": ["text_inputs"]})) == (False, None) + + +def test_a_colocated_encoders_output_is_a_followers_read(tmp_path): + """An encoder beside rank 0 hands its output to rank 0's node without a + read; each follower reads it.""" + w = _leader(tmp_path) + w._note_local_lockstep_inputs( + StoredOutputs(flat_uuids=[21], flat_rids=["r1"], signal_idxs=[0], num_tensors=[1]), + ["vision_embeds"], + ) + + assert _broadcast(w, _step({"r0": [], "r1": ["vision_embeds"]})) == (True, ("r0", "r1")) + + +def test_a_streamed_input_is_checked_on_every_step(tmp_path): + w = _leader(tmp_path, streamed={"thinker_states"}) + for _ in range(2): + assert _broadcast(w, _step({"r0": ["thinker_states"]})) == (True, ("r0",)) + assert _broadcast(w, _step({"r0": ["text_inputs"]})) == (False, None) + + +def test_rank_0_records_its_own_failed_reads_for_the_check(tmp_path): + w = _leader(tmp_path) + w._fail_requests = lambda errors: None + w.tensor_manager.pending.append(_read("r1", _input(15, "text_inputs"), _failed_future())) + w._check_ready_tensors() + assert w.scheduler.read_failed_rids == {"r1"} + + +def test_single_rank_workers_mark_nothing(tmp_path): + w = _leader(tmp_path) + w.parallel_leader_nodes = set() + w._setup_tp_read_check() + _land(w, "r1", _input(16, "text_inputs")) + assert not w._lockstep_unread + + +if __name__ == "__main__": # pragma: no cover + os.environ.setdefault("OMP_NUM_THREADS", "2") + raise SystemExit(pytest.main([__file__, "-q"])) diff --git a/test/modular/test_tp_read_failure_acceptance.py b/test/modular/test_tp_read_failure_acceptance.py index 8f1e6f1dd6..8ca7dc1eba 100644 --- a/test/modular/test_tp_read_failure_acceptance.py +++ b/test/modular/test_tp_read_failure_acceptance.py @@ -80,7 +80,7 @@ # The whole run; a run still going past this hung. DEADLINE_S = 30.0 # Per-rank setup a fix adds to Worker.__init__ (see the module docstring). -_DESIGN_SETUP = ("_setup_read_gate",) +_DESIGN_SETUP = ("_setup_read_gate", "_setup_tp_read_check") # What each request reads. A grid prompt is Qwen3.5's prefill_vision layout: # text, image, text; an image of grid (t, h, w) takes t*(h/2)*(w/2) tokens. diff --git a/test/modular/tp_read_check_sim.py b/test/modular/tp_read_check_sim.py new file mode 100644 index 0000000000..6f06fc431f --- /dev/null +++ b/test/modular/tp_read_check_sim.py @@ -0,0 +1,134 @@ +"""Explicit-state model of a lockstep instance whose follower input read can fail. + +Two requests, r0 and r1, each need one input before their step. Rank 0 (the +leader) reads both successfully; the follower reads r0, and its read of r1 +either lands or fails. Every interleaving of {follower read outcome, leader +scheduling, follower execution, failure drain} is explored. + +Modes: + NONE -- rank 0 schedules from its own reads alone (main). + ZERO -- as NONE, but a follower whose read failed runs the row's + ``prepare_inputs`` on zeros of the advertised shape (#628). + CAPTURE -- as NONE, but a follower whose read failed puts the walk's + capture input at rank 0's prepared length in the row's place, + without running ``prepare_inputs`` on it; the row is dropped + after the forward. + CHECK -- a follower builds rank 0's step once every row's read landed or + failed, leaving out the rows it could not read; before preparing + the step every rank agrees on those rows (one all-reduce) and + takes them out (``Worker._drop_rows_whose_read_failed``). + +Input kinds: + data -- the values do not change the step (``vision_embeds``). + shape -- the step's length is derived from the values (``image_grid_thw``): + zeros give the follower a different step than the leader's. + prepare -- ``prepare_inputs`` issues a collective for the row (a text part's + vocab-parallel embedding in Qwen3.5's ``prefill_vision``), so a + rank that puts another row in its place skips that collective. + +Checked at every terminal state: + live -- no reachable state is stuck short of done (the instance hang). + lockstep -- every step ran the same collectives on both ranks. + isolated -- r0 completes; only r1 can fail. +""" +from __future__ import annotations + +from dataclasses import dataclass, field + +RIDS = ("r0", "r1") + + +@dataclass +class Result: + states: int = 0 + violations: list[tuple[str, str]] = field(default_factory=list) + + @property + def ok(self) -> bool: + return not self.violations + + +def _initial(): + return ( + # follower read state per rid: pending | landed | failed + (("r0", "pending"), ("r1", "pending")), + frozenset(), # rids rank 0 has scheduled + (), # follower FIFO: tuple of batches (tuples of rids) + frozenset(), # rids whose step ran on both ranks + False, # conductor drained r1 (follower reported failure) + False, # a step ran different collectives on the two ranks + ) + + +def _step_mismatch(mode, kind, head, rd): + """Whether the follower's step differs from rank 0's for ``head``, given + which of its rows the follower could not read.""" + failed = [r for r in head if rd[r] == "failed"] + if not failed: + return False + if mode == "ZERO": + # prepare runs on zeros: its collectives match, its length may not + return kind == "shape" + if mode == "CAPTURE": + # rank 0's length, but rank 0's prepare ran collectives the follower's didn't + return kind == "prepare" + # CHECK: every rank takes the row out before preparing + return False + + +def _successors(state, mode, kind): + reads, scheduled, fifo, done, drained, mismatch = state + rd = dict(reads) + # the follower's reads resolve: r0 always lands, r1 may land or fail + for rid in RIDS: + if rd[rid] == "pending": + outcomes = ("landed",) if rid == "r0" else ("landed", "failed") + for out in outcomes: + nr = dict(rd) + nr[rid] = out + yield (tuple(sorted(nr.items())), scheduled, fifo, done, + drained or out == "failed", mismatch) + # rank 0 schedules any non-empty set of its ready rows as one step + ready = [r for r in RIDS if r not in scheduled and not (r == "r1" and drained)] + for mask in range(1, 1 << len(ready)): + batch = tuple(r for i, r in enumerate(ready) if mask >> i & 1) + yield (reads, scheduled | set(batch), fifo + (batch,), done, drained, mismatch) + # the follower runs its FIFO head together with rank 0 (the collective) + if fifo: + head = fifo[0] + if all(rd[r] == "landed" for r in head): + yield (reads, scheduled, fifo[1:], done | set(head), drained, mismatch) + elif mode != "NONE" and all(rd[r] != "pending" for r in head): + # the failed row is dropped (before the forward for CHECK, after it + # for the others); the rest of the step runs + yield (reads, scheduled, fifo[1:], + done | {r for r in head if rd[r] == "landed"}, + drained, mismatch or _step_mismatch(mode, kind, head, rd)) + + +def _terminal_ok(state): + reads, _scheduled, fifo, done, drained, _mismatch = state + rd = dict(reads) + if fifo or "r0" not in done: + return False + # r1 either ran or failed and was drained + return "r1" in done or (rd["r1"] == "failed" and drained) + + +def explore(mode: str, kind: str) -> Result: + res = Result() + seen = set() + stack = [_initial()] + while stack: + state = stack.pop() + if state in seen: + continue + seen.add(state) + succ = list(_successors(state, mode, kind)) + if state[5]: + res.violations.append(("lockstep", repr(state))) + if not succ and not _terminal_ok(state): + res.violations.append(("live", repr(state))) + stack.extend(succ) + res.states = len(seen) + return res From 8087222be50562e1c958e3e72d4a5f8b515ff8ca Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Mon, 26 Oct 2026 18:16:12 +0000 Subject: [PATCH 14/15] TP read check: record every failed read for the check, and share the row vote - A follower's failed read of a request it already failed for another reason now goes to read_failed_rids too (the failed_rids filter stays on the report only). Before, a read-checked head naming it waited on the follower with no warning while rank 0 waited in the check's all-reduce. - The per-row vote moves to distributed/all_reduce.py as flagged_on_any_rank, next to _all_ranks_agree. - The async follower puts its unread rows back through with_unread_rows, as the serial path does; a comment and a test pin why pushing back an unread fresh row on a voided speculation stays harmless. - test_a_step_whose_every_row_failed_runs_empty_on_every_rank bounds its run time against a run with no failure instead of a fixed 8 s, which slow hosts went past on startup alone. --- mstar/distributed/all_reduce.py | 14 ++++++ mstar/worker/worker.py | 65 +++++++++++++++------------ test/modular/test_tp_read_check.py | 72 ++++++++++++++++++++++++++++-- 3 files changed, 119 insertions(+), 32 deletions(-) diff --git a/mstar/distributed/all_reduce.py b/mstar/distributed/all_reduce.py index 128143f01e..342eb2d0dd 100644 --- a/mstar/distributed/all_reduce.py +++ b/mstar/distributed/all_reduce.py @@ -9,6 +9,7 @@ import logging import math from abc import ABC, abstractmethod +from collections.abc import Iterable, Sequence from typing import TYPE_CHECKING import torch @@ -41,6 +42,19 @@ def _all_ranks_agree(flag: bool, cpu_group) -> bool: return bool(t.item()) +def flagged_on_any_rank(flags: Sequence[bool], groups: Iterable["CommGroup"]) -> list[bool]: + """Per entry, OR ``flags`` across the ranks of each group in ``groups`` in + turn: one gloo all-reduce (MAX) each, skipping a group of one rank. Over a + TP group and then an SP group, an entry ends up set on every rank of the + grid if any rank set it. Every rank passes the same number of entries, in + the same order. CPU tensors: it waits on the other ranks, never on a device.""" + t = torch.tensor(list(flags), dtype=torch.int32) + for group in groups: + if group.world_size > 1: + dist.all_reduce(t, op=dist.ReduceOp.MAX, group=group.cpu_group) + return [bool(f) for f in t.tolist()] + + class AllReduce: """NCCL for every message, and the fallback of the backends below.""" diff --git a/mstar/worker/worker.py b/mstar/worker/worker.py index 554469e938..bd0cace0df 100644 --- a/mstar/worker/worker.py +++ b/mstar/worker/worker.py @@ -18,7 +18,6 @@ from typing import NamedTuple import torch -import torch.distributed as dist from mstar.communication.codec import WireCodec from mstar.communication.communicator import CommProtocol, make_communicator @@ -30,6 +29,7 @@ create_tensor_communication_manager, ) from mstar.conductor.request_info import CurrentForwardPassInfo +from mstar.distributed.all_reduce import flagged_on_any_rank from mstar.distributed.base import ShardingConfig from mstar.distributed.communication import WorkerParallelGroups from mstar.engine import apply_dynamo_config @@ -1544,20 +1544,17 @@ def _drop_rows_whose_read_failed(self, node_batch: ExecutingBatch) -> None: failed the read reported it, and the others fail it here, so its outputs and its pages never exist on any rank. - One gloo all-reduce of a flag per row (MAX over TP, then SP), in the - leader's row order, on the GPU thread where every rank prepares the - step. CPU tensors: it waits on the other ranks, never on the device.""" + One vote of a flag per row over TP, then SP (``flagged_on_any_rank``), + in the leader's row order, on the GPU thread where every rank prepares + the step.""" rids = node_batch.read_check_rids failed = self.scheduler.read_failed_rids - flags = torch.tensor([r in failed for r in rids], dtype=torch.int32) node = node_batch.node_name - for group in ( + votes = flagged_on_any_rank([r in failed for r in rids], ( self.parallel_groups.get_tp_config_for_node(node), self.parallel_groups.get_sp_config_for_node(node), - ): - if group.world_size > 1: - dist.all_reduce(flags, op=dist.ReduceOp.MAX, group=group.cpu_group) - flagged = {r for r, f in zip(rids, flags.tolist(), strict=True) if f} + )) + flagged = {r for r, f in zip(rids, votes, strict=True) if f} if not flagged: return for rid in flagged: @@ -1570,17 +1567,23 @@ def _check_ready_tensors(self) -> None: """Poll for completed RDMA transfers, feed ready graph edges to worker graph queues.""" self.wakeup_event.drain() ready = self.tensor_manager.get_ready_tensors() - failed = { - f.request_id: f"Error in worker: input tensor read failed: {type(f.error).__name__}: {f.error}" - for f in self.tensor_manager.take_failed_reads() - if f.request_id not in self.scheduler.failed_rids - } - if failed: - # only a TP step's read check asks (see ``unread_rows``) + failed_reads = self.tensor_manager.take_failed_reads() + if failed_reads: + # Every failed read is recorded for a TP step's read check (see + # ``unread_rows``), a request this rank already failed for another + # reason included: until the conductor's removal lands, rank 0 can + # still send a read-checked step naming it. self.scheduler.read_failed_rids.update( - rid for rid in failed if rid in self.request_state.per_request_info + f.request_id for f in failed_reads if f.request_id in self.request_state.per_request_info ) - self._fail_requests(failed) + # reported once per request + failed = { + f.request_id: f"Error in worker: input tensor read failed: {type(f.error).__name__}: {f.error}" + for f in failed_reads + if f.request_id not in self.scheduler.failed_rids + } + if failed: + self._fail_requests(failed) if ready and self._lockstep_led: self._note_lockstep_reads(ready) for request_id, edges in ready.items(): @@ -3196,17 +3199,23 @@ def _return_all_polled() -> None: new_request_to_worker_graph[rid] = prep_wg[rid] per_request_inputs[rid] = prep_inputs[rid] elif rid in unread: - # In with no inputs, as on the serial path: the step has rank 0's - # rows until its read check takes this one out. - new_request_to_worker_graph[rid] = in_batch[rid] if rid in in_batch else \ - self._graph_runtime.get_worker_graph_id_for_node( - head.node_name, - head.walks[head_walk_idx[rid]] if head.walks else head.graph_walk, - ) per_request_inputs[rid] = {} else: new_request_to_worker_graph[rid] = fresh_wg[rid] per_request_inputs[rid] = fresh_inputs[rid] + request_walks = {r: head.walks[w] for r, w in head_walk_idx.items()} if head.walks else None + if unread: + # In with no inputs, as on the serial path: the step has rank 0's + # rows until its read check takes these out. An unread row not in + # N is no continuing row, so voiding the speculation + # (``_maybe_clear_spec``) pushes it back though it was never popped. + # Harmless only while every later head naming it is read-checked, + # so ``unread_rows`` keeps it out of the pop: rank 0 checks each + # step taking the input until one runs (``_consume_lockstep_reads``), + # and that step's check takes the row out on every rank. + new_request_to_worker_graph, _ = self.scheduler.with_unread_rows( + head, head_rids, unread, new_request_to_worker_graph, request_walks or {}, + ) # Committed: the serial path must not see the head any more. self.scheduler.pop_tp_follow_head() @@ -3222,9 +3231,7 @@ def _return_all_polled() -> None: new_request_to_worker_graph, per_request_inputs, consumed_streaming_edges, in_n, is_same_node=True, spec_id=prep.spec_id, tp_seq=head.spec_seq, - request_walks={ - r: head.walks[w] for r, w in head_walk_idx.items() - } if head.walks else None, + request_walks=request_walks, ) if head.read_check and speculation is not None: speculation.node_batch.read_check_rids = tuple(head_rids) diff --git a/test/modular/test_tp_read_check.py b/test/modular/test_tp_read_check.py index 4ea6921d8b..b0909200fe 100644 --- a/test/modular/test_tp_read_check.py +++ b/test/modular/test_tp_read_check.py @@ -137,15 +137,19 @@ def test_a_prepare_that_runs_a_collective_per_row_stays_in_step(tmp_path, monkey def test_a_step_whose_every_row_failed_runs_empty_on_every_rank(tmp_path, monkeypatch): """Rank 0 sends a step of rows the follower could read none of: both ranks run it empty, both requests fail and drain, and nothing is routed.""" + clean = acc._run(tmp_path / "clean", acc.Scenario(kind="text", fail=False)) + assert not clean.crashed and not clean.hung, clean.summary() monkeypatch.setattr(acc, "_rank_main", _rank_main_every_row_fails) - run = acc._run(tmp_path, acc.Scenario(kind="text")) + run = acc._run(tmp_path / "failed", acc.Scenario(kind="text")) assert not run.crashed, run.summary() + "\n" + "\n".join(map(str, run.crashed.values())) assert not run.hung, run.summary() assert set(run.failed) == {"r0", "r1"}, run.summary() assert run.torn_down("r0") and run.torn_down("r1"), run.summary() assert run.routed == {0: {}, 1: {}}, run.summary() - # no rank sat in a collective the other never joined - assert run.elapsed < acc.COLLECTIVE_S, run.summary() + # No rank sat in a collective the other never joined, which costs a + # COLLECTIVE_S wait. Measured against the run with no failure, not as a + # fixed bound: starting the two ranks alone takes seconds on a slow host. + assert run.elapsed < clean.elapsed + acc.COLLECTIVE_S, (run.summary(), clean.summary()) # ── the check itself at TP4 and TP2 x SP2 (four gloo ranks) ─────────────── @@ -315,6 +319,36 @@ def test_without_a_read_check_a_follower_pops_the_row_like_any_other(caplog): assert sum("no read check" in r.getMessage() for r in caplog.records) == 1 +def test_a_read_that_fails_after_its_request_failed_still_goes_to_the_check(tmp_path): + """This rank already failed r1 for another reason, and r1's read then fails + before the conductor's removal lands. The read is not reported again, but it + is recorded for the check: rank 0's read-checked step naming r1 is built + with r1 empty, so this rank joins the step's check and the check takes r1 + out, rather than this rank waiting on r1 with no warning while rank 0 waits + in the check.""" + sched, manager = _follower_scheduler(ready=["r0"]) + sched.fail_rids({"r1"}) + tm, _ = _shm_manager(tmp_path) + w = Worker.__new__(Worker) + w.scheduler, w.request_state, w.tensor_manager = sched, manager, tm + w.wakeup_event = SimpleNamespace(drain=lambda: None) + w.enable_nvtx = False + reported = [] + w._fail_requests = reported.append + tm.pending.append(_read("r1", _input(17, "text_inputs"), _failed_future())) + w._check_ready_tensors() + assert reported == [], "r1 was already reported failed" + assert sched.read_failed_rids == {"r1"} + + sched.register_tp_follow(_head(read_check=True)) + batch = sched.get_next_batch(manager) + + assert batch is not None, "the follower waits on r1, which never readies" + assert list(batch.request_to_worker_graph) == ["r0", "r1"] + assert batch.read_check_rids == ("r0", "r1") + assert manager.pops[-1][1] == ["r0"] + + def test_a_removed_requests_read_failure_is_forgotten(): """The flag is keyed by the handle, which a later request can get.""" sched, _ = _follower_scheduler(ready=[], read_failed=["r1"]) @@ -432,6 +466,38 @@ def test_an_async_follower_that_cannot_prep_gives_back_only_the_rows_it_popped() assert w.scheduler.peek_tp_follow() is not None +def test_a_voided_speculations_unread_row_stays_out_of_the_next_read_checked_head(): + """f1's read failed here and f1 is not in N, so it is no continuing row of + the speculation: voiding it (``_maybe_clear_spec``) pushes f1 back with the + fresh rows, though it was never popped, and the graph then holds f1 as + ready with no input. Rank 0 sends the step again, read-checked, since it + never ran: the follower still pops only the rows it read. Only the read + check keeps f1 out: a head without one would pop it.""" + w, pending = _spec_follower( + read_failed=["f1"], continuing=["c0"], fresh=["f0", "f1"], head_rids=["c0", "f0", "f1"], + ) + spec = Worker._try_follow_speculation(w, pending) + spec.continuing_rids = set(w.continuing[-1]) # as ``_assemble_speculation`` sets it + manager = w.request_state + + # what ``_maybe_clear_spec`` gives back + sb = spec.scheduled_batch + fresh = [rid for rid in sb.request_to_worker_graph if rid not in spec.continuing_rids] + assert fresh == ["f0", "f1"] + w._graph_runtime.push_back_node(NODE, fresh, [sb.request_to_worker_graph[r] for r in fresh]) + assert NODE in manager.queues["wg0"]._ready["f1"] + + manager.ingest("c0") # N's routing readies its continuing row + w.scheduler.register_tp_follow(_head(read_check=True, rids=("c0", "f0", "f1"))) + batch = w.scheduler.get_next_batch(manager) + assert list(batch.request_to_worker_graph) == ["c0", "f0", "f1"] + assert manager.pops[-1][1] == ["c0", "f0"] + + w.scheduler.register_tp_follow(_head(read_check=False, rids=("f1",))) + assert w.scheduler.get_next_batch(manager) is not None + assert manager.pops[-1][1] == ["f1"] + + # ── rank 0 marks the steps whose rows take an input a follower reads ────── FOLLOWERS = ["w0", "w1"] From 1628c8b8d2f7b4f6784d28a1ef15a550b1881f96 Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Mon, 26 Oct 2026 18:43:22 +0000 Subject: [PATCH 15/15] test: the in-process Mooncake engine of the slice tests reads async too The Mooncake reader now reads with batch_transfer_async_read and polls transfer_check_status, so the fake engine #623's slice tests use gets both: each batch is copied on submit and reports completed. --- test/modular/test_shm_tensor_comm.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/test/modular/test_shm_tensor_comm.py b/test/modular/test_shm_tensor_comm.py index c7ad2b3f71..b7a3355b28 100644 --- a/test/modular/test_shm_tensor_comm.py +++ b/test/modular/test_shm_tensor_comm.py @@ -692,6 +692,13 @@ def batch_transfer_sync_read(self, session_id, local_ptrs, remote_ptrs, lengths) ctypes.memmove(dst, src, n) return 0 + def batch_transfer_async_read(self, session_id, local_ptrs, remote_ptrs, lengths): + self.batch_transfer_sync_read(session_id, local_ptrs, remote_ptrs, lengths) + return 1 # any nonzero batch id; every batch is done once submitted + + def transfer_check_status(self, batch_id): + return 1 # COMPLETED + def _transport(protocol, entity_id, shm_dir): if protocol == CommProtocol.SHM: