diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 10e7f888d..1af37f747 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -33,7 +33,7 @@ jobs: python-version: "3.12" - name: Install libzmq run: sudo apt-get update && sudo apt-get install -y libzmq3-dev - - name: Cargo tests (transport semantics) + - name: Cargo tests (transport + arena semantics) working-directory: rust run: cargo test --release - name: Build and install the extension @@ -42,10 +42,15 @@ jobs: python -m pip install --upgrade pip maturin maturin build --release pip install target/wheels/*.whl - - name: Interop tests (pyzmq <-> Rust) + - name: Interop tests (transport + arena) run: | - pip install pytest pyzmq - # the package itself, without its (heavy) deps — the transport - # test chain needs only pyzmq + stdlib + pip install pytest pyzmq numpy pyyaml + # cpu torch + triton: the arena test imports the tensor-transport + # stack, which reaches mstar.utils.sampling (triton at module level) + pip install torch --index-url https://download.pytorch.org/whl/cpu + pip install triton + # the package itself, without its remaining (heavy) deps pip install --no-deps -e . - pytest test/rust/test_rust_communicator.py -v + pytest test/rust/ -v + - name: Walk-layer A/B (informational) + run: python test/rust/bench_walk_ab.py diff --git a/docs/environment_variables.rst b/docs/environment_variables.rst index 5957fab13..3b6f1e0d2 100644 --- a/docs/environment_variables.rst +++ b/docs/environment_variables.rst @@ -35,3 +35,136 @@ Communication - ``19000`` - Base of the deterministic entity-id → TCP port map (``api_server`` = base, ``conductor`` = base+1, ``worker_`` = base+100+rank). + * - ``MSTAR_SHM_ARENA`` + - ``0`` + - SHM tensor-transport implementation. ``0``: per-uuid files. + ``1``: the Rust shared-memory arena (requires the ``rust/`` + extension; raises if missing). ``AUTO``: the arena when the + extension imports, files otherwise. Must match across the + deployment — arena locations ride in the tensor descriptors. + * - ``MSTAR_SHM_ARENA_SEGMENT_MB`` + - ``256`` + - Size of each arena segment. The arena grows segment by segment; + existing segments never move (registrations stay valid). + * - ``MSTAR_SHM_ARENA_MAX_SEGMENTS`` + - ``32`` + - Growth cap PER ENTITY. Every entity (workers + the api-server data + worker) creates its own arena, so node-wide /dev/shm demand can + reach ``MAX_SEGMENTS x SEGMENT_MB x num_entities`` — size against + ``df -h /dev/shm`` (tmpfs defaults to ~50% of RAM). Construction + fails fast if one entity's ceiling exceeds /dev/shm. At the cap, + sends spill (see ``MSTAR_SHM_ARENA_SPILL``). + * - ``MSTAR_SHM_ARENA_FULL_TIMEOUT_S`` + - ``30`` + - Strict mode only (``MSTAR_SHM_ARENA_SPILL=0``): how long a send + backpressures on a full arena before failing. + * - ``MSTAR_SHM_ARENA_SPILL`` + - ``1`` + - Degrade gracefully at the segment cap: stage the tensor through the + per-uuid file protocol instead — slower, never fails, matching the + file transport's saturation behavior. ``0`` restores strict + backpressure + timeout — only meaningful where ANOTHER thread + drains consumer ACKs (the threaded api-server); on a worker the + ACKs arrive on the very thread that would be waiting. + * - ``MSTAR_SHM_ARENA_SPILL_AFTER_S`` + - ``0`` + - Optional grace before spilling, for deployments where another + thread frees slots concurrently. Default 0: spill immediately + (a worker cannot receive ACKs while it waits). + * - ``MSTAR_SHM_ARENA_PIN`` + - ``1`` + - ``cudaHostRegister`` each mapped segment (both sides) so D2H/H2D + copies through the side streams run at page-locked bandwidth and + stay asynchronous. ``0`` disables (pageable copies). + * - ``MSTAR_SHM_ARENA_PIN_MAX_MB`` + - ``4096`` + - Budget for TOTAL pinned host memory PER PROCESS, distinct from the + segment cap (pinned pages come out of the OS's pageable pool + system-wide). Node-wide pinned demand is approx + ``PIN_MAX_MB x num_entities`` — a consumer pins peer segments too, + so one process can pin more than its own arena holds. Segments + past the budget stay unpinned: copies work, without async overlap. + * - ``MSTAR_SHM_ARENA_SLOT_TTL_S`` + - ``0`` + - TTL backstop for abort-orphaned slots (a request aborted after + staging but before all consumer ACKs defers reclaim forever). + A slot older than the request timeout cannot have a legitimate + reader, so a bound safely above it (recommend >= 2x the request + timeout) cannot race a real consumer. ``0`` disables (default, + pending review discussion); reclaims run under capacity pressure + and with the periodic stats sweep, logging loudly. + * - ``MSTAR_SHM_ARENA_STATS_INTERVAL_S`` + - ``60`` + - Under ``--log-stats``: how often the arena logs its occupancy / + fragmentation snapshot (segments, free bytes, largest contiguous + free block, pinned bytes). + +Graph / scheduler core +---------------------- + +.. list-table:: + :header-rows: 1 + :widths: 28 14 58 + + * - Variable + - Default + - Meaning + * - ``MSTAR_RUST_WALK`` + - ``0`` + - ``shadow``: run the Rust walk core in lockstep with every + per-request ``WorkerGraphIO`` on real traffic — Python stays + authoritative; ready-set / doneness / loop-counter divergence is + logged as an error (events the core does not model yet suspend + comparison for that request with a logged reason). ``1``: Rust + decisions (ready set, doneness, loop indices) are authoritative — + Python keeps executing values, comparison stays on, and divergence + or an unmodeled event falls the request back to Python with an + error logged. ``0``: off. + * - ``MSTAR_RUST_WALK_STRICT`` + - ``0`` + - With shadow mode, raise on divergence instead of logging (CI / + debugging). + +Serving (Rust frontend) +----------------------- + +Read by the ``mstar-server`` binary and its bridge +(``mstar-serve --rust-frontend``; see :doc:`installation`). + +.. list-table:: + :header-rows: 1 + :widths: 28 14 58 + + * - Variable + - Default + - Meaning + * - ``MSTAR_SERVER_BIN`` + - unset + - Path to the ``mstar-server`` binary. Fallback order: + ``--rust-frontend-bin``, this variable, ``$PATH``, then the in-repo + ``rust/server/target/release`` build. + * - ``MSTAR_REQUEST_TIMEOUT_S`` + - ``600`` + - Per-request budget in the Rust frontend; on expiry the client gets + an error and the request is aborted in the backend. + * - ``MSTAR_SAMPLE_RATE`` + - ``24000`` + - Sample rate stamped on ``/v1/audio/speech`` WAV output. + * - ``MSTAR_ALLOW_REMOTE`` + - ``0`` + - Allow ``http(s)`` media URLs in requests (fetched server-side, + 30 s timeout). Off by default. + * - ``MSTAR_MAX_CONCURRENT_REQUESTS`` + - ``256`` + - Admission cap on in-flight generation requests; beyond it clients + get an immediate 503 instead of queueing into the request timeout. + ``/health`` and ``/v1/models`` bypass the cap. + * - ``MSTAR_MAX_BODY_MB`` + - ``128`` + - Request body limit (multipart uploads included). + * - ``MSTAR_TOKENIZER`` + - unset + - Path to a HuggingFace ``tokenizer.json`` enabling frontend + tokenization. Leave unset with the Python backend — its preprocess + worker owns tokenization, and the bridge rejects pre-tokenized + ingest. diff --git a/docs/installation.rst b/docs/installation.rst index 4009c0022..e0f57def1 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -248,13 +248,20 @@ Verify the install mstar --help mstar-serve --help -Optional: the Rust ZMQ transport --------------------------------- - -The ZeroMQ control mesh can run over a Rust transport (vendored in ``rust/``) -instead of pyzmq — same endpoints, same wire format, selectable per process -with ``MSTAR_RUST_ZMQ`` (see :doc:`environment_variables`). It is optional: -without it, everything runs on pyzmq as before. +Optional: Rust support +---------------------- + +Parts of the runtime can run on Rust components (vendored in ``rust/``), +each optional and selectable per process — without the extension everything +runs on the pure-Python paths as before. Migrated so far (see +:doc:`environment_variables` for the flags): + +* **ZeroMQ control mesh** — ``MSTAR_RUST_ZMQ``: same endpoints, same wire + format as pyzmq. +* **SHM tensor arena** — ``MSTAR_SHM_ARENA``: replaces the per-tensor file + transport; requires the extension, interoperates on the same descriptor + wire (and depends on the transport above only in the sense that both ship + in the same extension). Build the extension into your environment with `maturin `_ (needs a Rust toolchain; ``rustup`` works): diff --git a/mstar/communication/arena.py b/mstar/communication/arena.py new file mode 100644 index 000000000..22a453899 --- /dev/null +++ b/mstar/communication/arena.py @@ -0,0 +1,809 @@ +"""Tensor transport over a shared-memory arena. + +``ArenaShmCommunicationManager`` replaces the file transport's per-tensor +open/write/read/unlink with a Rust segmented ``/dev/shm`` arena (persistent +mmaps + first-fit coalescing allocator, vendored in ``rust/``). Producer: +``register_for_send`` reserves a slot and D2H-copies into the segment on the +copy stream; the ``(segment, offset)`` rides the existing descriptors. +Consumer: ``start_read_tensors`` maps the named segment once, reads +zero-copy, H2D-copies on the copy stream; the producer reclaims on ACK, +gated by a CUDA-event future so an edge is never ACKed before its copies +land. Segments are mapped once and never move, so the one-time +``cudaHostRegister`` per segment (within ``MSTAR_SHM_ARENA_PIN_MAX_MB``) +holds for its lifetime and keeps the side-stream copies truly async. + +Capacity degrades in layers: grow by segments up to +``MSTAR_SHM_ARENA_MAX_SEGMENTS``; at the cap, briefly backpressure for +consumer ACKs; then spill the tensor to the per-uuid file protocol +(``MSTAR_SHM_ARENA_SPILL``, default on) — slower, never fails, like the old +transport at saturation. ``stats_summary()`` exposes occupancy and the +fragmentation gauge (largest contiguous free block); ``--log-stats`` logs it +periodically. + +Ceilings are PER-ENTITY and multiply across a node: with E entities +(workers + the api-server data worker), /dev/shm demand can reach +``MAX_SEGMENTS x SEGMENT_MB x E`` and pinned host RAM approx +``PIN_MAX_MB x E`` (consumers pin peer segments too, so one process can pin +more than its own arena holds). Construction fails fast when one entity's +ceiling already exceeds /dev/shm, and warns when it exceeds current free +space or when the pin budget is an outsized share of physical RAM. + +Selection: ``create_tensor_communication_manager`` picks this manager for +the SHM protocol when ``MSTAR_SHM_ARENA`` is ``1`` (require) or ``AUTO`` +(use if the ``mstar_rust`` extension imports); default ``0`` keeps the file +transport. See :doc:`environment_variables` for all knobs. +""" + +from __future__ import annotations + +import logging +import os +import queue +import threading +import time +import weakref +from concurrent.futures import Future + +import torch + +from mstar.communication.communicator import BaseCommunicator +from mstar.communication.tensors import ( + FutureAndPointers, + SharedMemoryCommunicationManager, + _deserialize_tensor, + _nullcontext, + _serialize_tensor, +) +from mstar.graph.base import GraphEdge, TensorPointerInfo + +logger = logging.getLogger(__name__) + +_CUDA_HOST_ALREADY_REGISTERED = 712 + + +class _CudaEventFuture: + """Future-shaped CUDA event. An edge whose H2D copies ride behind this + event reports ready only once the copies have completed on the device — + so the ACK that lets the producer reclaim the arena slot is deferred by + ``get_ready_tensors``'s existing future polling instead of a blocking + host synchronize.""" + + def __init__(self, stream): + # blocking=True: `synchronize` SLEEPS until the event fires. The + # default (False) busy-waits — the watcher thread that turns these + # into wake futures would burn a full core per wait. + self._event = torch.cuda.Event(blocking=True) + self._event.record(stream) + + def done(self) -> bool: + return self._event.query() + + def result(self) -> None: + self._event.synchronize() + + +_CUDART = None + + +def _cudart(): + """libcudart via ctypes: unlike the torch binding, a ctypes call + RELEASES the GIL, so registering a 256 MiB segment (tens of ms) on the + send path cannot stall the process's other Python threads (serve loop, + stream relays, the wake watcher).""" + global _CUDART + if _CUDART is None: + import ctypes + import ctypes.util + + name = ctypes.util.find_library("cudart") or "libcudart.so" + lib = ctypes.CDLL(name) + lib.cudaHostRegister.argtypes = [ + ctypes.c_void_p, ctypes.c_size_t, ctypes.c_uint] + lib.cudaHostRegister.restype = ctypes.c_int + _CUDART = lib + return _CUDART + + +def _unlink_paths(paths: list) -> None: + """Exit-time segment unlink (weakref.finalize target — must not + reference the manager). Workers exit with the manager still alive, so + the interpreter never garbage-collects it and the Rust Drop (which + unlinks) never runs; without this, every worker run leaks its full + segments in /dev/shm until swept by a later start.""" + for path in paths: + try: + os.unlink(path) + except OSError: + pass + + +def _unpin(ptr: int) -> bool: + """cudaHostUnregister a previously pinned mapping (GIL released via + ctypes). Returns success.""" + if not torch.cuda.is_available(): + return False + try: + lib = _cudart() + import ctypes + + lib.cudaHostUnregister.argtypes = [ctypes.c_void_p] + lib.cudaHostUnregister.restype = ctypes.c_int + rc = lib.cudaHostUnregister(ptr) + except OSError: + rc = torch.cuda.cudart().cudaHostUnregister(ptr) + return rc == 0 + + +def _pin(ptr: int, nbytes: int) -> bool: + """cudaHostRegister a mapped segment (idempotent). Returns success. + GIL released for the duration (see ``_cudart``).""" + if not torch.cuda.is_available(): + return False + try: + rc = _cudart().cudaHostRegister(ptr, nbytes, 0) + except OSError: # no loadable libcudart: fall back to the torch binding + rc = torch.cuda.cudart().cudaHostRegister(ptr, nbytes, 0) + ok = rc in (0, _CUDA_HOST_ALREADY_REGISTERED) + if not ok: + logger.warning("cudaHostRegister(%#x, %d) failed rc=%d — copies fall " + "back to pageable bandwidth", ptr, nbytes, rc) + return ok + + +class ArenaShmCommunicationManager(SharedMemoryCommunicationManager): + """Tensor transport via the Rust shared-memory arena (``mstar_rust``).""" + + def __init__( + self, + my_entity_id: str, + hostname: str, + device: str, + communicator: BaseCommunicator, + shm_dir: str | None = None, + enable_prof: bool = False, + ): + super().__init__( + my_entity_id=my_entity_id, hostname=hostname, device=device, + communicator=communicator, shm_dir=shm_dir, + enable_prof=enable_prof, + ) + from mstar_rust import SegmentedShmArena, ShmArena + + self._ShmArena = ShmArena + segment_mb = int(os.getenv("MSTAR_SHM_ARENA_SEGMENT_MB", "256")) + max_segments = int(os.getenv("MSTAR_SHM_ARENA_MAX_SEGMENTS", "32")) + self._full_timeout_s = float( + os.getenv("MSTAR_SHM_ARENA_FULL_TIMEOUT_S", "30")) + self._pin_segments = ( + os.getenv("MSTAR_SHM_ARENA_PIN", "1") == "1" + and torch.cuda.is_available() and str(device) != "cpu" + ) + # Pinned host memory is a system-wide resource (pages come out of the + # OS's pageable pool), so it gets its own budget, distinct from the + # segment cap. Segments past the budget stay unpinned — copies still + # work, they just lose async overlap. Oversized dedicated segments + # (single allocations larger than a segment) are never pinned: a + # one-shot transfer doesn't amortize the registration cost. + self._pin_budget = int( + os.getenv("MSTAR_SHM_ARENA_PIN_MAX_MB", "4096")) << 20 + self._pinned_bytes = 0 + self._pin_budget_warned = False + self._segment_bytes = segment_mb << 20 + # Arena saturation spills to the per-uuid file transport (the old + # protocol) instead of failing: bursty or oversized workloads degrade + # to file-copy speed, matching the prior manager's "slower, never + # fails" behavior. MSTAR_SHM_ARENA_SPILL=0 restores strict + # backpressure + timeout. + self._spill = os.getenv("MSTAR_SHM_ARENA_SPILL", "1") == "1" + # Default 0: on a worker, TENSOR_RECEIVED ACKs (which free slots) + # are processed by the SAME thread that would sit in this grace + # wait, so waiting is pure dead time there — spill immediately. + # Deployments where another thread drains the communicator (the + # threaded api_server) can set a small grace to ride out bursts. + self._spill_after_s = float( + os.getenv("MSTAR_SHM_ARENA_SPILL_AFTER_S", "0")) + self._frag_warned = False + # h2d completion watcher: turns a CUDA event into a real Future so + # the worker's eventfd wakes the moment reads finish (no 10 ms tick). + self._wake_q: queue.Queue | None = None + self._wake_lock = threading.Lock() + # Serializes pin accounting and peer-map insertion: _pin releases + # the GIL (ctypes), so unlocked read-modify-write of _pinned_bytes + # tears, and check-then-insert on _peer_segments can map+pin the + # same segment twice (the loser leaking forever). + self._pin_lock = threading.Lock() + # Concurrent start_read_tensors calls (threaded api-server): used to + # keep eviction from unmapping a segment another thread is copying + # from before its future lands in `pending`. + self._reads_active = 0 + # Periodic occupancy/fragmentation logging, tied to --log-stats + # (enable_prof) and time-gated. + self._stats_interval_s = float( + os.getenv("MSTAR_SHM_ARENA_STATS_INTERVAL_S", "60")) + self._stats_last = 0.0 + # CEILINGS ARE PER-ENTITY and multiply across a node: every entity + # (workers + the api-server data worker) creates its own arena, so + # node /dev/shm demand can reach + # MAX_SEGMENTS x SEGMENT_MB x num_entities + # and node pinned RAM approx PIN_MAX_MB x num_entities (a consumer + # pins peer segments too, so one process's pinned bytes can exceed + # its own arena). Check the static tmpfs ceiling NOW instead of + # surfacing as ENOSPC on a growth mid-run. + per_entity_max = (segment_mb << 20) * max_segments + try: + st = os.statvfs("/dev/shm") + shm_total = st.f_frsize * st.f_blocks + shm_avail = st.f_frsize * st.f_bavail + except OSError: + shm_total = shm_avail = None + self._shm_total = shm_total + if shm_total is not None: + if per_entity_max > shm_total: + raise RuntimeError( + f"SHM arena ceiling for ONE entity " + f"({per_entity_max >> 20} MiB = " + f"MSTAR_SHM_ARENA_MAX_SEGMENTS x _SEGMENT_MB) exceeds " + f"/dev/shm ({shm_total >> 20} MiB) — and every entity " + f"multiplies this. Lower the knobs or grow tmpfs.") + if per_entity_max > shm_avail: + logger.warning( + "ARENA: this entity's ceiling (%d MiB) exceeds current " + "/dev/shm free space (%d MiB); ceilings are per-entity " + "and multiply across workers — growth may hit ENOSPC " + "under load", per_entity_max >> 20, shm_avail >> 20) + try: + phys = os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") + except (ValueError, OSError): + phys = None + if (self._pin_segments and phys is not None + and self._pin_budget > phys // 4): + logger.warning( + "ARENA: MSTAR_SHM_ARENA_PIN_MAX_MB (%d MiB) exceeds a " + "quarter of physical RAM (%d MiB) FOR ONE ENTITY; pinned " + "budgets multiply across entities and come out of the OS's " + "pageable pool", self._pin_budget >> 20, phys >> 20) + # INSTANCE-UNIQUE base name: a fixed per-entity name in the global + # /dev/shm namespace collides across servers (a second server's + # create() truncates the first's live segments — silent corruption) + # and across users (permission-denied at startup). pid + random + # token makes each instance's names unique; wire-compatibility is + # free because consumers open whatever segment name the descriptor + # carries. The pid embedded in the name also enables the orphan + # sweep below. + import secrets + + self._sweep_orphans() + base = (f"mstar_arena_{my_entity_id}_{os.getpid()}_" + f"{secrets.token_hex(4)}") + self._arena = SegmentedShmArena.create( + base, segment_mb << 20, max_segments) + # Guaranteed exit-time unlink: weakref.finalize runs at interpreter + # shutdown even when the manager is still referenced (the worker + # case — no explicit cleanup path runs there, so the Rust Drop + # never fires). The callback captures only the mutable name list, + # which _sync_segments extends as the arena grows. + self._own_segment_paths: list[str] = [] + self._finalizer = weakref.finalize( + self, _unlink_paths, self._own_segment_paths) + # Producer-side segment views (memoryviews are stable: segments + # never move or resize) + how many segments are already pinned. + self._seg_views: list[memoryview] = [] + self._pinned_segments = 0 + self._sync_segments() + # uuid -> (segment_idx, offset) for sender-side reclaim. + # (register_for_send receives the TensorPointerInfos directly and + # stamps them in place — no side-table needed.) + self._arena_locs: dict[str, tuple[int, int]] = {} + # uuid -> stage time, for the TTL backstop: a request aborted after + # staging but before every consumer ACKs defers reclaim forever + # (cleanup_request waits for ACKs that will never come). A slot + # older than the REQUEST timeout cannot have a legitimate reader — + # the request is dead by contract — so freeing past a bound safely + # above it cannot race a real consumer. Default OFF pending review + # discussion; enable with MSTAR_SHM_ARENA_SLOT_TTL_S (recommend + # >= 2x the request timeout). + self._arena_ts: dict[str, float] = {} + self._slot_ttl_s = float( + os.getenv("MSTAR_SHM_ARENA_SLOT_TTL_S", "0")) + self._ttl_reclaimed_total = 0 + # Consumer-side: peer segment name -> (arena, memoryview, + # pinned_nbytes) — pinned_nbytes 0 when the segment wasn't pinned. + # Entries are EVICTED (unpin + unmap) once the backing file is gone + # (producer finished or restarted): with instance-unique names a + # restarting producer mints new names every generation, so a + # never-evicting cache would grow mappings and pinned bytes without + # bound on any long-lived consumer. + self._peer_segments: dict[str, tuple[object, memoryview, int]] = {} + self._peer_evict_last = 0.0 + + # -- segments -------------------------------------------------------- + + def _sweep_orphans(self) -> None: + """A SIGKILLed server never runs Drop, orphaning up to its full + arena in /dev/shm until reboot. Names embed the owning pid, so a + startup sweep can reclaim any segment whose owner is gone. Files we + cannot judge (foreign naming) or cannot remove (another user's) + are left with a debug note.""" + try: + names = os.listdir("/dev/shm") + except OSError: + return + for name in names: + if not name.startswith("mstar_arena_"): + continue + parts = name.split(".")[0].rsplit("_", 2) + if len(parts) != 3 or not parts[1].isdigit(): + continue # pre-uniquification or foreign naming: skip + if os.path.exists(f"/proc/{parts[1]}"): + continue # owner alive + try: + os.unlink(f"/dev/shm/{name}") + logger.info("ARENA: swept orphaned segment %s " + "(owner pid %s is gone)", name, parts[1]) + except OSError as e: + logger.debug("ARENA: cannot sweep %s: %s", name, e) + + def _maybe_pin(self, ptr: int, nbytes: int) -> int: + """Pin within budget; returns the bytes actually pinned (0 if + skipped/failed). Runs under _pin_lock: _pin releases the GIL, so + budget check + register + accounting must be one atomic unit.""" + # Oversized dedicated segments ARE pinned (within budget): freed + # segments are reused for later large tensors, so the registration + # amortizes over the segment's lifetime, not one transfer. + if not self._pin_segments: + return 0 + with self._pin_lock: + return self._pin_locked(ptr, nbytes) + + def _pin_locked(self, ptr: int, nbytes: int) -> int: + if self._pinned_bytes + nbytes > self._pin_budget: + if not self._pin_budget_warned: + logger.warning( + "ARENA: pinned-memory budget reached (%d MiB, " + "MSTAR_SHM_ARENA_PIN_MAX_MB); further segments stay " + "unpinned — copies work but lose async overlap", + self._pin_budget >> 20) + self._pin_budget_warned = True + return 0 + if _pin(ptr, nbytes): + self._pinned_bytes += nbytes + self._pinned_segments += 1 + return nbytes + return 0 + + def _sync_segments(self) -> None: + grew = False + while len(self._seg_views) < self._arena.num_segments: + i = len(self._seg_views) + seg = self._arena.segment(i) + self._own_segment_paths.append( + f"/dev/shm/{self._arena.segment_name(i)}") + self._seg_views.append(memoryview(seg)) + self._maybe_pin(*seg.ptr_len()) + grew = True + if grew: + total, free, largest = self._arena.stats() + if (self._shm_total is not None + and total > self._shm_total * 0.8): + logger.warning( + "ARENA: this entity's segments now total %d MiB — over " + "80%% of /dev/shm (%d MiB). Oversized dedicated " + "segments grow past the static ceiling; other entities " + "multiply this further.", + total >> 20, self._shm_total >> 20) + logger.info( + "ARENA: grew to %d segments (%d MiB total, %d MiB free, " + "largest free block %d MiB, %d MiB pinned)", + self._arena.num_segments, total >> 20, free >> 20, + largest >> 20, self._pinned_bytes >> 20) + + def stats_summary(self) -> dict: + """Occupancy/fragmentation snapshot (named apart from the raw + ``SegmentedShmArena.stats`` tuple). The fragmentation signature is + `largest_free_block` collapsing while `free_bytes` stays high.""" + total, free, largest = self._arena.stats() + return { + "segments": self._arena.num_segments, + "total_bytes": total, + "free_bytes": free, + "largest_free_block": largest, + "pinned_bytes": self._pinned_bytes, + } + + def _reclaim_expired(self) -> int: + """TTL backstop for abort-orphaned slots (see _arena_ts). Returns + the number of slots/files reclaimed.""" + if not self._slot_ttl_s: + return 0 + now = time.monotonic() + n = 0 + for uuid, ts in list(self._arena_ts.items()): + if now - ts < self._slot_ttl_s: + continue + self._arena_ts.pop(uuid, None) + if (loc := self._arena_locs.pop(uuid, None)) is not None: + self._arena.free(*loc) + n += 1 + if (path := self._shm_files.pop(uuid, None)) is not None: + try: + os.unlink(path) + n += 1 + except FileNotFoundError: + pass + if n: + self._ttl_reclaimed_total += n + logger.warning( + "ARENA: TTL-reclaimed %d slot(s)/file(s) older than %.0fs " + "(%d total) — requests aborted without consumer ACKs; " + "the ACK path is leaking", + n, self._slot_ttl_s, self._ttl_reclaimed_total) + return n + + def _maybe_log_stats(self) -> None: + """Under ``--log-stats`` (enable_prof), log the snapshot at most + once per MSTAR_SHM_ARENA_STATS_INTERVAL_S so a long soak leaves an + occupancy/fragmentation time series in the logs.""" + if not self.enable_prof: + return + now = time.monotonic() + if now - self._stats_last < self._stats_interval_s: + return + self._stats_last = now + self._reclaim_expired() + st = self.stats_summary() + # live_slots/spill_files climbing while requests finish = reclaim + # being deferred (e.g. aborts without ACKs) — the soak's leak canary. + st["live_slots"] = len(self._arena_locs) + st["spill_files"] = len(self._shm_files) + logger.info("ARENA stats: %s", st) + + def _peer_view(self, segment_name: str) -> memoryview: + entry = self._peer_segments.get(segment_name) + if entry is None: + with self._pin_lock: + entry = self._peer_segments.get(segment_name) + if entry is None: # lost the race: another thread mapped it + arena = self._ShmArena.open(segment_name) + pinned = (self._pin_locked(*arena.ptr_len()) + if self._pin_segments else 0) + entry = self._peer_segments[segment_name] = ( + arena, memoryview(arena), pinned) + return entry[1] + + def _evict_dead_peers(self) -> None: + """Drop cached peer segments whose backing file is gone. Gated on + `self.pending` being empty: an mmap must outlive any in-flight h2d + copy that reads from it, and pending futures are exactly those + copies. Time-gated to keep the existence checks off the hot path.""" + now = time.monotonic() + # _reads_active > 1: another thread may have queued an async H2D + # copy whose future hasn't reached `pending` yet — unmapping under + # it would be a use-after-free. Only the sole active reader evicts. + if (self.pending or self._reads_active > 1 + or now - self._peer_evict_last < 10.0): + return + self._peer_evict_last = now + for name in list(self._peer_segments): + if os.path.exists(f"/dev/shm/{name}"): + continue + with self._pin_lock: + arena, view, pinned = self._peer_segments[name] + if pinned: + # Unregister BEFORE unmapping: tearing down a mapping + # CUDA still holds registered leaves a dangling pinned + # range. On failure keep the entry (retry next sweep) + # so accounting stays truthful. + ptr, _len = arena.ptr_len() + if not _unpin(ptr): + logger.warning( + "ARENA: cudaHostUnregister failed for dead peer " + "%s; keeping mapping until it succeeds", name) + continue + self._pinned_bytes -= pinned + del self._peer_segments[name] + view.release() + logger.debug("ARENA: evicted dead peer segment %s " + "(%d MiB pinned released)", name, pinned >> 20) + + # -- producer --------------------------------------------------------- + + def _reserve(self, nbytes: int) -> tuple[int, int] | None: + """Reserve with layered degradation. At the segment cap: wait + briefly for consumer ACKs to free space; then (default) return None + so the caller SPILLS this tensor to the per-uuid file transport — + slower, never fails, exactly the old manager's saturation behavior. + MSTAR_SHM_ARENA_SPILL=0 keeps strict backpressure for the full + timeout, then raises.""" + grace = self._spill_after_s if self._spill else self._full_timeout_s + deadline = time.monotonic() + grace + warned = False + while True: + try: + seg, off = self._arena.reserve(max(nbytes, 1)) + except RuntimeError: + if self._reclaim_expired(): + continue # expired slots freed: retry the reserve + total, free, largest = self._arena.stats() + if free >= nbytes and not self._frag_warned: + # The fragmentation signature: enough TOTAL free space, + # but no contiguous block large enough. + logger.warning( + "ARENA: fragmentation — need %d bytes with %d free " + "but largest free block is %d (of %d total). " + "Consider larger MSTAR_SHM_ARENA_SEGMENT_MB.", + nbytes, free, largest, total) + self._frag_warned = True + if time.monotonic() > deadline: + if self._spill: + return None + raise RuntimeError( + f"SHM arena full for >{self._full_timeout_s}s " + f"({self._arena.num_segments} segments); raise " + "MSTAR_SHM_ARENA_MAX_SEGMENTS / _SEGMENT_MB or check " + "for consumers not ACKing") from None + if not warned: + logger.warning( + "SHM arena at capacity; backpressuring sends until " + "consumers ACK%s", + " (then spilling to file)" if self._spill else "") + warned = True + time.sleep(0.002) + continue + self._sync_segments() + return seg, off + + def register_for_send( + self, request_id: str, tensor_infos: list[TensorPointerInfo], + skip_cuda_sync: bool = False, + ): + if not skip_cuda_sync and torch.cuda.is_available(): + torch.cuda.default_stream().synchronize() + ctx = ( + torch.cuda.stream(self._d2h_stream) + if self._d2h_stream is not None + else _nullcontext() + ) + queued = False + self._maybe_log_stats() + with ctx: + for info_arg in tensor_infos: + uuid = info_arg.uuid + if self.tensor_store.is_registered(request_id, uuid): + continue + tensor = self.tensor_store.get_tensor(request_id, uuid) + t0 = time.perf_counter() + t = tensor.detach().contiguous() + nbytes = t.numel() * t.element_size() + loc = self._reserve(nbytes) + if loc is not None and self.tensor_store.is_registered( + request_id, uuid): + # Lost a concurrent-duplicate race: another thread + # registered this uuid while our reserve released the + # GIL. Return our slot instead of orphaning it. + self._arena.free(*loc) + continue + if loc is None: + # Arena saturated: spill THIS tensor to the per-uuid file + # protocol (infos keep shm_segment=None — the consumer + # falls back to the file read for exactly those). + data = _serialize_tensor(t) + path = self._shm_path(self.my_entity_id, uuid) + with open(path, "wb") as f: + f.write(data) + self._shm_files[uuid] = path + self._arena_ts[uuid] = time.monotonic() + self.tensor_store.set_metadata( + request_id, uuid, mem_registered=True) + if self.enable_prof: + self._record_tx(request_id, uuid, len(data), + time.perf_counter() - t0) + logger.debug("ARENA: spilled %s to %s (%d bytes)", + uuid, path, len(data)) + continue + seg, off = loc + try: + if nbytes: + host = torch.frombuffer( + self._seg_views[seg][off:off + nbytes], + dtype=torch.uint8, + ).view(t.dtype).reshape(t.shape) + # Async D2H into the pinned segment when a copy + # stream exists (one sync below covers the batch); + # blocking otherwise, so the descriptor can never + # ship ahead of the bytes. + host.copy_( + t, non_blocking=self._d2h_stream is not None) + queued = True + self._arena_locs[uuid] = (seg, off) + self._arena_ts[uuid] = time.monotonic() + except BaseException: + # Anything that unwinds between reserve and the + # _arena_locs record would orphan the slot forever + # (cleanup can only free what is recorded). + self._arena.free(seg, off) + raise + seg_name = self._arena.segment_name(seg) + info_arg.shm_segment = seg_name + info_arg.shm_offset = off + self.tensor_store.set_metadata( + request_id, uuid, mem_registered=True) + if self.enable_prof: + self._record_tx( + request_id, uuid, nbytes, time.perf_counter() - t0) + logger.debug("ARENA: staged %s at %s+%d (%d bytes)", + uuid, seg_name, off, nbytes) + if queued and self._d2h_stream is not None: + # The control message referencing these bytes is sent after we + # return; the consumer must never observe a partial copy. + self._d2h_stream.synchronize() + + # -- consumer --------------------------------------------------------- + + def start_read_tensors( + self, request_id: str, graph_edges: list[GraphEdge], + graph_walk: str | None = None, + ): + # Increment races are benign here: a torn count can only make + # eviction OVER-cautious (skip a sweep), never unsafe. + self._reads_active += 1 + try: + return self._start_read_tensors(request_id, graph_edges, + graph_walk) + finally: + self._reads_active -= 1 + + def _start_read_tensors(self, request_id, graph_edges, graph_walk): + self._evict_dead_peers() + h2d_did_work = False + read_edges: list[tuple[GraphEdge, float]] = [] + ctx = ( + torch.cuda.stream(self._h2d_stream) + if self._h2d_stream is not None + else _nullcontext() + ) + with ctx: + for graph_edge in graph_edges: + if len(graph_edge.tensor_info) == 0: + continue + rx_t0 = time.perf_counter() + for info in graph_edge.tensor_info: + if info.source_entity == self.my_entity_id: + self._slice_existing_tensor( + request_id=request_id, name=graph_edge.name, + next_node=graph_edge.next_node, + graph_walk=graph_walk, info=info, + ) + self.tensor_store.increment_ref( + request_id, info.uuid, 1) + continue + if self.tensor_store.check_uuid_presence( + request_id, info.uuid): + self.tensor_store.increment_ref( + request_id, info.uuid, 1) + continue + if info.shm_segment is None: + # Spilled at the producer (arena saturated): read the + # per-uuid file, the old protocol's path. + path = self._shm_path(info.source_entity, info.uuid) + with open(path, "rb") as f: + f.seek(info.offset) + data = f.read(info.nbytes) + tensor = _deserialize_tensor( + data, self.device, tensor_info=info) + else: + tensor = self._read_from_arena(info) + h2d_did_work = h2d_did_work or tensor.numel() > 0 + self.tensor_store.put_tensor( + request_id, info.uuid, tensor) + self.tensor_store.set_metadata( + request_id, info.uuid, mem_registered=False) + # +1 transit (released by get_ready_tensors), +1 usage + # (released by _cleanup_consumed_inputs). + self.tensor_store.increment_ref(request_id, info.uuid, 2) + read_edges.append( + (graph_edge, time.perf_counter() - rx_t0)) + future = None + if h2d_did_work and self._h2d_stream is not None: + # Downstream kernels see the data (device-side ordering only). + torch.cuda.default_stream(self.device).wait_stream( + self._h2d_stream) + # The producer reclaims the slot when an edge is ACKed, and the + # source is its live mapping — so the edge must not report ready + # until its copies have completed (the file path's f.read() made + # this implicit). One event covers the batch: all copies were + # queued on the h2d stream in program order. get_ready_tensors + # polls it — no host block here. + future = _CudaEventFuture(self._h2d_stream) + for graph_edge, rx_time in read_edges: + self.pending.append( + FutureAndPointers( + future=future, graph_edges=[graph_edge], + request_id=request_id, rx_time=rx_time, + ) + ) + if future is None: + return [] + # A real Future for the worker's eventfd (EventWakeup.register_ + # futures needs add_done_callback): completed by the watcher thread + # the moment the h2d copies finish, so an otherwise-idle worker + # re-checks get_ready_tensors immediately instead of on its next + # poll tick. + return [self._wake_when_done(future)] + + def _wake_when_done(self, cuda_future) -> Future: + with self._wake_lock: # two callers must not race the create + if self._wake_q is None: + self._wake_q = queue.Queue() + # staticmethod target closing over ONLY the queue: a bound + # method would pin the whole manager (and its arena — so + # segments never unlink) for the daemon thread's lifetime. + threading.Thread( + target=self._watch_wakes, args=(self._wake_q,), + daemon=True, + name=f"arena-h2d-wake-{self.my_entity_id}").start() + fut: Future = Future() + self._wake_q.put((cuda_future, fut)) + return fut + + @staticmethod + def _watch_wakes(q: queue.Queue) -> None: + # Events are queued in stream order, so sequential waits are exact. + # None is the close() sentinel. + while True: + item = q.get() + if item is None: + return + cuda_future, fut = item + try: + cuda_future.result() # event.synchronize (GIL released) + fut.set_result(None) + except Exception: # noqa: BLE001 — wake best-effort; + # a cancelled/already-resolved future must not kill the + # watcher (its death silently downgrades every later wake + # to the poll tick and grows the queue unboundedly). + if not fut.done(): + try: + fut.set_result(None) + except Exception: # noqa: BLE001 + pass + + def _read_from_arena(self, info: TensorPointerInfo) -> torch.Tensor: + if info.nbytes == 0: + t = torch.empty( + info.dims, dtype=info.dtype) + return t.to(self.device) if self.device != "cpu" else t + view = self._peer_view(info.shm_segment) + start = info.shm_offset + info.offset # offset = TP-shard read offset + flat = torch.frombuffer( + view[start:start + info.nbytes], dtype=torch.uint8) + t = flat.view(info.dtype).reshape(info.dims) + if self.device != "cpu": + return t.to(self.device, non_blocking=True) + return t.clone() # CPU consumer: own the bytes past reclaim + + # -- reclaim ---------------------------------------------------------- + + def close(self) -> None: + """Stop the wake watcher (segments/pins release with the arena's + Drop once the manager is garbage collected).""" + with self._wake_lock: + if self._wake_q is not None: + self._wake_q.put(None) + self._wake_q = None + + def _cleanup_by_uuid(self, request_id: str, uuid: str): + # Grandparent cleanup (refcounts): skip the file manager's unlink. + super(SharedMemoryCommunicationManager, self)._cleanup_by_uuid( + request_id, uuid) + self._arena_ts.pop(uuid, None) + if (loc := self._arena_locs.pop(uuid, None)) is not None: + self._arena.free(*loc) + logger.debug("ARENA: freed %s at %s", uuid, loc) + if (path := self._shm_files.pop(uuid, None)) is not None: + try: + os.unlink(path) # spilled tensor: reclaim the file + except FileNotFoundError: + pass + if not self.tensor_store.check_uuid_presence(request_id, uuid): + return + self.tensor_store.remove_tensor(request_id, uuid) diff --git a/mstar/communication/tensors.py b/mstar/communication/tensors.py index 92eae3379..c92d0229f 100644 --- a/mstar/communication/tensors.py +++ b/mstar/communication/tensors.py @@ -1132,6 +1132,12 @@ def start_read_tensors( if self.tensor_store.check_uuid_presence(request_id, info.uuid): self.tensor_store.increment_ref(request_id, info.uuid, 1) continue + if info.shm_segment is not None: + raise RuntimeError( + f"tensor {info.uuid} from {info.source_entity} " + "was staged in an SHM arena but this consumer " + "runs the file transport — MSTAR_SHM_ARENA must " + "match across the deployment") path = self._shm_path(info.source_entity, info.uuid) with open(path, "rb") as f: f.seek(info.offset) @@ -1186,8 +1192,40 @@ def create_tensor_communication_manager( shm_dir: str | None = None, enable_prof: bool=False ) -> TensorCommunicationManager: - """Select tensor transport backend based on protocol.""" + """Select tensor transport backend based on protocol. + + For the SHM protocol, ``MSTAR_SHM_ARENA`` selects the implementation + (see ``docs/environment_variables.rst``): ``0`` (default) — per-uuid + files; ``1`` — the Rust shared-memory arena (raises if the + ``mstar_rust`` extension is missing); ``AUTO`` — the arena when the + extension imports, files otherwise. The flag must match across the + deployment: the arena location rides in the tensor descriptors, so a + file-transport consumer cannot read an arena producer. + """ if protocol == CommProtocol.SHM: + choice = os.getenv("MSTAR_SHM_ARENA", "0").upper() + if choice not in ("0", "1", "AUTO"): + raise ValueError( + f"MSTAR_SHM_ARENA must be 0, 1, or AUTO; got {choice!r}") + if choice != "0": + try: + from mstar.communication.arena import ( + ArenaShmCommunicationManager, + ) + except ImportError: + if choice == "1": + raise + logger.debug("MSTAR_SHM_ARENA=AUTO: mstar_rust not " + "installed, using the file transport") + else: + return ArenaShmCommunicationManager( + my_entity_id=my_entity_id, + hostname=hostname, + device=device, + communicator=communicator, + shm_dir=shm_dir, + enable_prof=enable_prof + ) return SharedMemoryCommunicationManager( my_entity_id=my_entity_id, hostname=hostname, diff --git a/mstar/conductor/conductor.py b/mstar/conductor/conductor.py index e45570b30..c3b81cfb6 100644 --- a/mstar/conductor/conductor.py +++ b/mstar/conductor/conductor.py @@ -3,6 +3,7 @@ import logging import multiprocessing as mp import os +import signal import socket import time from collections import defaultdict @@ -100,6 +101,17 @@ def _worker_process_target( tcp_transfer_device="", ): """Top-level target for spawned worker processes. Must be module-level for picklability.""" + # SIGTERM (the conductor's p.terminate()) defaults to immediate death: + # no unwinding, no atexit, no finalizers — so a worker's shared-memory + # segments were never unlinked and leaked into /dev/shm for the life of + # the box. Turning it into SystemExit unwinds the interpreter normally, + # which runs the transport's cleanup. The main process gets this for + # free from SIGINT -> KeyboardInterrupt, which is why only the workers + # leaked. + def _graceful_exit(_signum, _frame): + raise SystemExit(0) + + signal.signal(signal.SIGTERM, _graceful_exit) logging.basicConfig( level=getattr(logging, log_level), format=f"%(asctime)s %(levelname)s [{worker_id}] %(name)s: %(message)s", @@ -439,13 +451,24 @@ def _launch_workers(self): atexit.register(self.shutdown) def shutdown(self): - logger.info("Shutting down conductor...") """Terminate and join all worker processes.""" + logger.info("Shutting down conductor...") + # SIGTERM is handled in _worker_process_target as a graceful exit so + # the transports' cleanup runs (shm segments unlinked). A worker + # blocked in a C call cannot service it in time, so escalate to + # SIGKILL after the join window rather than hang the shutdown — + # the leftover segments are then reclaimed by the next arena + # start's orphan sweep. for p in self._worker_processes: if p.is_alive(): p.terminate() for p in self._worker_processes: p.join(timeout=5) + if p.is_alive(): + logger.warning( + "Worker pid %s did not exit on SIGTERM; killing", p.pid) + p.kill() + p.join(timeout=5) self._worker_processes.clear() def _assign_worker_graphs_to_workers(self) -> dict[str, list[str]]: diff --git a/mstar/graph/base.py b/mstar/graph/base.py index 72d57f695..17c38de34 100644 --- a/mstar/graph/base.py +++ b/mstar/graph/base.py @@ -29,6 +29,12 @@ class TensorPointerInfo: source_tp_size: int = 1 source_tp_rank: int = 0 + # SHM-arena transport (MSTAR_SHM_ARENA): the producer wrote this tensor's + # bytes at `shm_offset` inside the named arena segment. None for the + # other transports (per-uuid files, Mooncake). + shm_segment: str | None = None + shm_offset: int = 0 + _source_node_name: str | None = None _source_graph_walk: str | None = None @@ -45,6 +51,8 @@ def clone(self): offset=self.offset, source_tp_size=self.source_tp_size, source_tp_rank=self.source_tp_rank, + shm_segment=self.shm_segment, + shm_offset=self.shm_offset, _source_node_name=self._source_node_name, _source_graph_walk=self._source_graph_walk ) diff --git a/mstar/graph/rust_core.py b/mstar/graph/rust_core.py new file mode 100644 index 000000000..06ea92ee6 --- /dev/null +++ b/mstar/graph/rust_core.py @@ -0,0 +1,586 @@ +"""The graph-core translation seam: translate ``GraphSection`` trees into the Rust +walk core's spec. + +``walks_to_json`` produces the JSON that ``mstar_rust.WalkSet.from_json`` +compiles; a ``WalkSet.state(walk)`` is then the Rust counterpart of a +per-request ``WorkerGraphIO`` — same readiness, completion-routing, loop +iteration, and termination semantics (asserted by +``test/rust/test_walk_parity.py``, which drives both implementations with +identical event sequences). + +Scope: the walk state machine only. Streaming buffer semantics and the +worker/conductor wiring stay in Python for now — this module is the +translation seam those later steps build on. +""" + +from __future__ import annotations + +import json +import logging +import os +from itertools import count as _count + +from mstar.graph.base import ( + GraphEdge, + GraphNode, + GraphSection, + Loop, + NodeCompletionOutput, + Parallel, + Sequential, +) +from mstar.graph.special_destinations import EMIT_TO_CLIENT, EMPTY_DESTINATION + +# mstar's sentinel destinations -> the Rust core's. +_DEST = {EMIT_TO_CLIENT: "EMIT_TO_CLIENT", EMPTY_DESTINATION: "EMPTY_DESTINATION"} + + +def edge_to_spec(edge: GraphEdge) -> dict: + return { + "next_node": _DEST.get(edge.next_node, edge.next_node), + "name": edge.name, + "persist": bool(edge.persist), + "output_modality": edge.output_modality or None, + } + + +def section_to_spec(section: GraphSection) -> dict: + if isinstance(section, GraphNode): + return { + "kind": "node", + "name": section.name, + "input_names": sorted(section.input_names), + "outputs": [edge_to_spec(e) for e in section.outputs], + } + if isinstance(section, Loop): + return { + "kind": "loop", + "name": section.name, + "body": section_to_spec(section.section), + "max_iters": int(section.max_iters), + "outputs": [edge_to_spec(e) for e in section.outputs], + "accumulated_outputs": [ + edge_to_spec(e) for e in section.accumulated_outputs + ], + } + if isinstance(section, (Sequential, Parallel)): + kind = "sequential" if isinstance(section, Sequential) else "parallel" + return { + "kind": kind, + "sections": [section_to_spec(s) for s in section.sections], + } + raise TypeError(f"unknown GraphSection type: {type(section).__name__}") + + +def walks_to_json(walks: dict[str, GraphSection]) -> str: + """``{walk_name: GraphSection}`` -> the WalkSet JSON spec.""" + return json.dumps({name: section_to_spec(s) for name, s in walks.items()}) + + +# --------------------------------------------------------------------------- +# Shadow adoption (MSTAR_RUST_WALK=shadow): the Rust walk state runs in +# lockstep with the Python WorkerGraphIO on real traffic — Python stays +# authoritative, every event is mirrored, and ready-set / doneness / loop +# divergence is reported loudly. This is the pre-authority adoption step: +# it exercises the Rust core against every real model's walks in situ. +# --------------------------------------------------------------------------- + +logger = logging.getLogger(__name__) + +_WALKSET_CACHE: dict[int, tuple] = {} # id(section) -> (section ref, WalkSet) + + +def rust_walk_mode() -> str: + """``MSTAR_RUST_WALK``: ``0`` (default, off), ``shadow`` (mirror + + compare, Python authoritative), or ``1`` (Rust decisions authoritative; + Python keeps executing values; divergence falls back per-request).""" + mode = os.getenv("MSTAR_RUST_WALK", "0").lower() + if mode not in ("0", "shadow", "1", "pure"): + raise ValueError( + f"MSTAR_RUST_WALK must be 0, shadow, 1, or pure; got {mode!r}") + return mode + + +class _RustReadyView(set): + """The ready set served from the Rust state in authority mode. Callers + mutate it (``pop_ready_nodes`` discards; OOM push-back adds) — discard + marks the node scheduled in Rust; add cannot be modeled (no unschedule), + so it falls the request back to Python.""" + + def __init__(self, names, owner): + super().__init__(names) + self._owner = owner + + def discard(self, name): + if name in self: + self._owner._mark_scheduled(name) + super().discard(name) + + def add(self, name): + self._owner._suspend("ready push-back (OOM hold) not modeled") + super().add(name) + + +class _RegistryShim: + """`wg_state_registry.is_done` view honoring the current authority.""" + + def __init__(self, owner): + self._owner = owner + + @property + def is_done(self): + return self._owner._done() + + def __getattr__(self, name): + return getattr(self._owner._io.wg_state_registry, name) + + +class ShadowedWorkerGraphIO: + """A ``WorkerGraphIO`` with a Rust ``WalkState`` shadowing every event. + + Delegates everything to the Python io (authoritative). Mutating calls are + mirrored into the Rust state; after each, ready sets, doneness, and loop + indices are compared. A divergence logs an error (or raises with + ``MSTAR_RUST_WALK_STRICT=1``). Events the Rust core does not model yet + (streaming-buffer edges, speculative buffers) suspend comparison for the + request with a single logged reason rather than false-positive. + """ + + def __init__(self, io, section, wg_id, authority: bool = False): + from mstar_rust import WalkSet + + self._io = io + self._authority = authority + self._scheduled: set[str] = set() + if authority: + self.wg_state_registry = _RegistryShim(self) + entry = _WALKSET_CACHE.get(id(section)) + if entry is None: + walkset = WalkSet.from_json(walks_to_json({"walk": section})) + _WALKSET_CACHE[id(section)] = (section, walkset) + else: + walkset = entry[1] + self._walkset_key = id(section) + self._rs = walkset.state("walk") + self._wg_id = wg_id + # (name, next_node) pairs produced INSIDE this graph: the Rust core + # routes them itself on complete(); the worker re-ingests them into + # the Python io, and mirroring that would double-ingest. + self._internal: set[tuple[str, str]] = set() + for node in io.nodes.values(): + for edge in node.outputs: + self._internal.add((edge.name, edge.next_node)) + for loop in io.loops.values(): + for edge in list(loop.outputs) + list(loop.accumulated_outputs): + self._internal.add((edge.name, edge.next_node)) + self._suspended: str | None = None + # Internal edges from the last completion the worker has yet to + # re-ingest: Rust routes them inside complete(), Python's worker + # re-ingests them after — compare only once both have settled. A + # multiset of (name, next_node): the same pair can also arrive as an + # EXTERNAL seed (e.g. iteration 0 of a loop-back input), which must + # be mirrored, so membership here is what says "already routed". + self._pending_edges: list[tuple[str, str]] = [] + self._strict = os.getenv("MSTAR_RUST_WALK_STRICT") == "1" + + def __getattr__(self, name): + return getattr(self._io, name) + + # -- authority views ------------------------------------------------------ + + def _rust_drives(self) -> bool: + return self._authority and self._suspended is None + + @property + def ready_node_names(self): + if self._rust_drives(): + return _RustReadyView(self._rs.ready_nodes(), self) + return self._io.ready_node_names + + def _done(self) -> bool: + if self._rust_drives(): + return self._rs.is_done() + return self._io.wg_state_registry.is_done + + def get_loop_indices(self): + if self._rust_drives(): + return dict(self._rs.loop_iters()) + return self._io.get_loop_indices() + + def _mark_scheduled(self, name: str) -> None: + if self._suspended is None and name not in self._scheduled: + try: + self._rs.schedule(name) + self._scheduled.add(name) + except Exception as e: # noqa: BLE001 + self._suspend(f"schedule rejected: {e!r}") + # keep the Python io in step (it stays the value store) + self._io.ready_node_names.discard(name) + + # -- mirroring ----------------------------------------------------------- + + def _suspend(self, reason: str) -> None: + if self._suspended is None: + self._suspended = reason + logger.info("rust-walk shadow suspended (%s): %s", + self._wg_id, reason) + + def _check(self, event: str) -> None: + if self._suspended: + return + py_ready = sorted(self._io.ready_node_names) + rs_ready = sorted(self._rs.ready_nodes()) + py_done = self._io.wg_state_registry.is_done + rs_done = self._rs.is_done() + rs_iters = dict(self._rs.loop_iters()) + py_iters = self._io.get_loop_indices() + if py_ready != rs_ready or py_done != rs_done or py_iters != rs_iters: + msg = (f"rust-walk shadow divergence after {event} ({self._wg_id}): " + f"ready py={py_ready} rs={rs_ready}; " + f"done py={py_done} rs={rs_done}; " + f"iters py={py_iters} rs={rs_iters}") + if self._strict: + raise AssertionError(msg) + logger.error(msg) + self._suspend("diverged; comparison stopped for this request") + + def ingest_input(self, graph_edge, can_buffer: bool = True) -> bool: + claimed = self._io.ingest_input(graph_edge, can_buffer) + if claimed: + if getattr(graph_edge, "is_streaming", False): + self._suspend("streaming edge (buffer semantics stay Python)") + elif (pair := (graph_edge.name, graph_edge.next_node)) in \ + self._pending_edges: + # Rust already routed this inside complete(); Python is + # catching up now. + self._pending_edges.remove(pair) + else: + try: + self._rs.seed([(graph_edge.next_node, graph_edge.name)]) + except Exception as e: # noqa: BLE001 + self._suspend(f"seed rejected: {e!r}") + if not self._pending_edges: + self._check( + f"ingest {graph_edge.name}->{graph_edge.next_node}") + return claimed + + def mark_node_complete(self, node_name: str): + completion = self._io.mark_node_complete(node_name) + if self._suspended is None: + try: + if node_name not in self._scheduled: + self._rs.schedule(node_name) + self._scheduled.discard(node_name) + self._rs.complete(node_name) + except Exception as e: # noqa: BLE001 + self._suspend(f"complete rejected: {e!r}") + # EVERYTHING a completion hands back for local re-ingest is + # already accounted for inside the Rust complete(): routed + # internal edges, loop-back promotions, AND the re-injected + # external inputs mstar re-emits on a loop advance. + self._pending_edges = [ + (e.name, e.next_node) for e in completion.output_edges + if e.next_node in self._io.nodes] + if not self._pending_edges: + self._check(f"complete {node_name}") + return completion + + def register_loop_finish_signal(self, loop_name: str): + self._io.register_loop_finish_signal(loop_name) + if self._suspended is None and loop_name in self._io.loops: + try: + self._rs.signal_loop_finish(loop_name) + except Exception as e: # noqa: BLE001 + self._suspend(f"loop signal rejected: {e!r}") + + def clear(self): + # End-of-pass reset for multi-forward-pass requests: fresh Rust + # state (the worker re-ingests the next pass's inputs), comparison + # re-armed. + self._io.clear() + walkset = _WALKSET_CACHE[self._walkset_key][1] + self._rs = walkset.state("walk") + self._suspended = None + + def ingest_for_speculation(self, edges, source_node): + # Speculative buffers are a Python-side concern (not modeled yet); + # they don't mutate real readiness, so no mirror and no suspend. + return self._io.ingest_for_speculation(edges, source_node) + + +def wrap_worker_graph_io(io, section, wg_id): + """The adoption seam: wrap a fresh per-request WorkerGraphIO according to + ``MSTAR_RUST_WALK``. ``0`` returns it untouched; ``shadow`` mirrors with + Python authoritative; ``1`` serves decisions (ready set, doneness, loop + indices) from the Rust state — Python keeps executing values, comparison + stays on, and any divergence or unmodeled event falls the request back + to Python with an error logged.""" + mode = rust_walk_mode() + if mode == "pure": + # No Python registries at all. Streaming graphs and speculation are + # not modeled yet: graphs with stream consumers fall back to + # authority mode (Rust decisions, Python values) with a logged note. + if any(getattr(n, "consumes_stream", False) + for n in io.nodes.values()): + logger.info("rust-walk pure: %s consumes streams; using " + "authority mode for it", wg_id) + return ShadowedWorkerGraphIO(io, section, wg_id, authority=True) + return PureRustWorkerGraphIO(section, wg_id) + if mode in ("shadow", "1"): + return ShadowedWorkerGraphIO(io, section, wg_id, + authority=(mode == "1")) + return io + + +# --------------------------------------------------------------------------- +# Pure mode (MSTAR_RUST_WALK=pure): no Python WorkerGraphIO at all — the Rust +# WalkState owns readiness/loops/termination, and this adapter keeps only the +# value store (uuid -> GraphEdge, the mstar-rs pattern) plus the node views +# the worker's execution path reads. Streaming graphs and speculation fall +# back / disable (documented; those ports are staged separately). +# --------------------------------------------------------------------------- + +class _Slot: + """ready_signals-shaped view: just the fields execution reads.""" + + def __init__(self): + self.ready_inputs: dict = {} + self.ready_names: set = set() + + def clear(self): + self.ready_inputs.clear() + self.ready_names.clear() + + +class _NodeView: + """Per-request node facade satisfying worker execution reads.""" + + def __init__(self, node): + self.name = node.name + self.input_names = node.input_names + self.outputs = node.outputs + self.enable_async_scheduling = getattr( + node, "enable_async_scheduling", True) + self.consumes_stream = getattr(node, "consumes_stream", False) + self.ready_signals = _Slot() + self.ready_next_iter = _Slot() + + +class _PureRegistryShim: + __slots__ = ("_io",) + + def __init__(self, io): + self._io = io + + @property + def is_done(self): + return self._io._rs.is_done() + + +class PureRustWorkerGraphIO: + """`WorkerGraphIO` with the Rust core as the ONLY state machine. + + Values: every ingested/produced edge gets a uuid; the Rust state carries + uuids, this adapter maps them back to `GraphEdge`s (whose `tensor_info` + the worker fills after execution — mutation is visible at loop + termination when events return the captured uuids). Per (node, input) + the adapter keeps a FIFO of pending edges; scheduling a node moves the + oldest edge per input into its `ready_signals` view — FIFO order matches + iteration order because completion re-emits loop externals exactly as + mstar's registries do. + """ + + def __init__(self, section, wg_id): + from mstar_rust import WalkSet + + entry = _WALKSET_CACHE.get(id(section)) + if entry is None: + walkset = WalkSet.from_json(walks_to_json({"walk": section})) + _WALKSET_CACHE[id(section)] = (section, walkset) + else: + walkset = entry[1] + self._rs = walkset.state("walk") + self.graph = section + self.wg_id = wg_id + self.num_times_run = 0 + self._graph_nodes = section.get_nodes() + self.loops = section.get_loops() + self.nodes = {n: _NodeView(g) for n, g in self._graph_nodes.items()} + self._uuid = _count(1) + self._store: dict[int, GraphEdge] = {} + # (node, input) -> FIFO of pending ingested edges + self._pending: dict[tuple[str, str], list[GraphEdge]] = {} + self._scheduled: set[str] = set() + # externals per loop, keyed by (node, input): re-emitted on each + # advance (mstar's contract). Keyed — not appended — because the + # re-emitted edges come back through ingest_input; appending would + # double the list every iteration. + self._loop_externals: dict[str, dict[tuple, GraphEdge]] = {} + self._loop_members = { + ln: set(lp.section.get_nodes().keys()) + for ln, lp in self.loops.items() + } + self._internal = { + (e.name, e.next_node) + for g in self._graph_nodes.values() for e in g.outputs + } + self.ready_for_streaming: set = set() + self._registry_shim = _PureRegistryShim(self) + # complete_full return values cached between completions; ingest and + # schedule dirty the ready cache (recomputed lazily). + self._loop_cache = {n: (i, t) for n, i, t in self._rs.loop_states()} + self._ready_cache: list | None = [] + self._done_cache = False + # ids of edges WE re-emitted on the last loop advance: Rust already + # re-injected those values internally, so their re-ingest must only + # refill the Python FIFO views — re-seeding Rust would append to the + # loop's external_inputs every iteration (quadratic complete cost). + self._reemitted_ids: set[int] = set() + self._tm = None + self._rid = None + + # -- registry-shaped surface ---------------------------------------------- + + @property + def ready_node_names(self): + if self._ready_cache is None: + self._ready_cache = self._rs.ready_nodes() + return _RustReadyView( + (n for n in self._ready_cache if n not in self._scheduled), + self) + + @property + def wg_state_registry(self): + return self._registry_shim + + def get_loop_indices(self): + return dict(self._rs.loop_iters()) + + def register_communication_info(self, tm, rid): + self._tm, self._rid = tm, rid + + def get_node(self, name): + return self.nodes[name] + + # -- events ---------------------------------------------------------------- + + def _mark_scheduled(self, name): + if name in self._scheduled: + return + self._rs.schedule(name) + self._scheduled.add(name) + view = self.nodes[name] + view.ready_signals.clear() + for inp in view.input_names: + fifo = self._pending.get((name, inp)) + if fifo: + edge = fifo.pop(0) + view.ready_signals.ready_inputs[inp] = edge + view.ready_signals.ready_names.add(inp) + + def ingest_input(self, edge, can_buffer=True): + if edge.next_node not in self.nodes: + return False + if edge.name not in self.nodes[edge.next_node].input_names: + return False + self._pending.setdefault((edge.next_node, edge.name), []).append(edge) + if id(edge) in self._reemitted_ids: + self._reemitted_ids.discard(id(edge)) + else: + uuid = next(self._uuid) + self._store[uuid] = edge + self._rs.seed_with([(edge.next_node, edge.name, uuid)]) + self._ready_cache = None # readiness may have changed + # track loop externals for re-emission on advance + pair = (edge.name, edge.next_node) + if pair not in self._internal: + for ln, members in self._loop_members.items(): + if edge.next_node in members: + self._loop_externals.setdefault(ln, {})[ + (edge.next_node, edge.name)] = edge + return True + + def mark_node_complete(self, node_name): + before = self._loop_cache + self._mark_scheduled(node_name) # idempotent + self._scheduled.discard(node_name) + view = self.nodes[node_name] + # Return the declared edge objects themselves — mstar's registries do + # the same (the worker overwrites tensor_info each pass); fresh uuids + # keep the Rust-side value identity per iteration. COPY the list + # (sharing the objects): the loop-advance path extends the + # completion's list, and sharing it would grow the node's declared + # outputs every iteration. + out_edges = list(view.outputs) + outputs = [] + for e in out_edges: + uuid = next(self._uuid) + self._store[uuid] = e + outputs.append((e.name, [uuid])) + events, done, ready, loop_states = self._rs.complete_full( + node_name, outputs) + self._ready_cache = ready + self._done_cache = done + after = {n: (i, t) for n, i, t in loop_states} + self._loop_cache = after + + completion = NodeCompletionOutput(output_edges=out_edges) + for ln in self.loops: + b_i, b_t = before.get(ln, (0, False)) + a_i, a_t = after.get(ln, (0, False)) + if a_t and not b_t: # terminated now: filter its loop-backs + completion.filtered_signals |= self.loops[ln]._loop_back_inputs + # A dead loop's nodes must show no stale readiness (the EOS + # contract: termination clears ready signals): drop their + # buffered inputs and blank the views. + for member in self._loop_members.get(ln, ()): + view = self.nodes.get(member) + if view is None: + continue + view.ready_signals.clear() + view.ready_next_iter.clear() + for inp in view.input_names: + self._pending.pop((member, inp), None) + # loop outputs with captured values (uuids -> stored edges) + for kind, name, _tgt, uuids in events: + del kind, name, uuids + for le in self.loops[ln].outputs: + src = [e for e in out_edges if e.name == le.name] + lo = le.clone() + if src: + lo.tensor_info = src[0].tensor_info + lo._src_edge = src[0] + completion.output_edges.append(lo) + elif a_i > b_i: # advanced: re-emit external inputs + reemit = list(self._loop_externals.get(ln, {}).values()) + completion.output_edges.extend(reemit) + # EVERY locally-destined edge in this completion was already routed + # (or re-injected) inside the Rust complete — their re-ingest by the + # worker must only refill the Python FIFO views, never re-seed Rust + # (a seed of a loop-member input appends to the loop's + # external_inputs, growing every iteration). + self._reemitted_ids.update( + id(e) for e in completion.output_edges + if e.next_node in self.nodes) + return completion + + def register_loop_finish_signal(self, loop_name): + if loop_name in self.loops: + self._rs.signal_loop_finish(loop_name) + + def ingest_for_speculation(self, edges, source_node): + return [] # speculation disabled in pure mode (staged separately) + + def clear_speculative_inputs(self): + pass + + def clear(self): + walkset = _WALKSET_CACHE[id(self.graph)][1] + self._rs = walkset.state("walk") + self._scheduled.clear() + self._pending.clear() + self._loop_externals.clear() + for v in self.nodes.values(): + v.ready_signals.clear() + v.ready_next_iter.clear() + self.num_times_run += 1 diff --git a/mstar/worker/micro_scheduler.py b/mstar/worker/micro_scheduler.py index 75ab61815..13b175f10 100644 --- a/mstar/worker/micro_scheduler.py +++ b/mstar/worker/micro_scheduler.py @@ -102,7 +102,7 @@ def _select_node_priority( walk_counts[e.graph_walk] = walk_counts.get(e.graph_walk, 0) + 1 graph_walk = max(walk_counts, key=walk_counts.get) - return node_name, graph_walk + return best_node_name, graph_walk def _select_node_rr( self, node_name_to_requests: dict[str, list[ReadyNodeEntry]] diff --git a/mstar/worker/node_manager_utils.py b/mstar/worker/node_manager_utils.py index ea350eebc..759f456c0 100644 --- a/mstar/worker/node_manager_utils.py +++ b/mstar/worker/node_manager_utils.py @@ -101,6 +101,11 @@ def add_request(self, request_id: str): queue.register_communication_info( self.tensor_manager, request_id ) + # MSTAR_RUST_WALK=shadow: run the Rust walk core in lockstep with + # this io on real traffic (Python authoritative, divergence loud). + from mstar.graph.rust_core import wrap_worker_graph_io + queue = wrap_worker_graph_io( + queue, self.worker_graph.section, self.worker_graph_id) self.per_request_queues[request_id] = queue def remove_request(self, request_id: str): diff --git a/rust/Cargo.lock b/rust/Cargo.lock index f41870dab..0835867ce 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -169,6 +169,12 @@ dependencies = [ "rustversion", ] +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + [[package]] name = "jobserver" version = "0.1.35" @@ -207,6 +213,15 @@ version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[package]] +name = "memmap2" +version = "0.9.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" +dependencies = [ + "libc", +] + [[package]] name = "memoffset" version = "0.9.1" @@ -221,9 +236,11 @@ name = "mstar-rust" version = "0.1.0" dependencies = [ "libc", + "memmap2", "pyo3", "rmp-serde", "serde", + "serde_json", "thiserror", "zmq", ] @@ -426,6 +443,19 @@ dependencies = [ "syn", ] +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + [[package]] name = "serde_spanned" version = "0.6.9" @@ -602,6 +632,12 @@ dependencies = [ "dircpy", ] +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" + [[package]] name = "zmq" version = "0.10.0" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 721746a2d..260e6889e 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -6,14 +6,17 @@ license = "Apache-2.0" [lib] name = "mstar_rust" -# cdylib: the Python extension module. rlib: the same crate as a Rust -# library, so the migration's later steps (Rust conductor / API server) -# consume the typed `ZmqCommunicator` layer directly. +# cdylib: the Python extension module. rlib: the same code as a plain Rust +# library, so other Rust components (servers, services, tools) can consume +# any of this crate's modules directly — typed messaging today, and +# whatever the crate grows to hold — without going through Python. crate-type = ["cdylib", "rlib"] [dependencies] +memmap2 = "0.9" serde = { version = "1", features = ["derive"] } rmp-serde = "1.3" +serde_json = "1" thiserror = "2" zmq = "0.10" pyo3 = { version = "0.23", features = ["extension-module"] } diff --git a/rust/src/core/error.rs b/rust/src/core/error.rs new file mode 100644 index 000000000..437cd091a --- /dev/null +++ b/rust/src/core/error.rs @@ -0,0 +1,21 @@ +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum CoreError { + #[error("unknown node '{0}' in walk '{1}'")] + UnknownNode(String, String), + #[error("unknown loop '{0}' in walk '{1}'")] + UnknownLoop(String, String), + #[error("unknown walk '{0}'")] + UnknownWalk(String), + #[error("node '{node}' is not ready (missing inputs: {missing:?})")] + NodeNotReady { node: String, missing: Vec }, + #[error("node '{0}' was not scheduled; complete_node without take_node_inputs")] + NodeNotScheduled(String), + #[error("duplicate node name '{0}' in walk '{1}'")] + DuplicateNode(String, String), + #[error("invalid walk spec: {0}")] + InvalidSpec(String), +} + +pub type Result = std::result::Result; diff --git a/rust/src/core/graph.rs b/rust/src/core/graph.rs new file mode 100644 index 000000000..d15554732 --- /dev/null +++ b/rust/src/core/graph.rs @@ -0,0 +1,230 @@ +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; + +use serde::{Deserialize, Serialize}; + +use super::error::{CoreError, Result}; + +/// Special edge destination: route the edge's tensors to the client. +pub const EMIT_TO_CLIENT: &str = "EMIT_TO_CLIENT"; +/// Special edge destination: drop the value (side-effect-only outputs). +pub const EMPTY_DESTINATION: &str = "EMPTY_DESTINATION"; + +/// A directed dataflow edge. `name` is both the producing node's output name +/// and the input name at the destination (as in `mstar/graph/base.py`). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct EdgeSpec { + pub next_node: String, + pub name: String, + /// Persisted edges become walk outputs handed to the policy on WalkDone + /// (mstar's "persist signals" flowing back to the conductor). + #[serde(default)] + pub persist: bool, + #[serde(default)] + pub output_modality: Option, + /// Streaming edge (mstar's `StreamingGraphEdge`): the value goes into + /// the stream buffer of the connection targeting this partition instead + /// of a node in this walk. The producer is otherwise unaware. + #[serde(default)] + pub target_partition: Option, +} + +/// A unit of computation. Executed on the data plane (Python/torch); the +/// control plane only tracks its readiness and routes its outputs. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct NodeSpec { + pub name: String, + pub input_names: BTreeSet, + #[serde(default)] + pub outputs: Vec, +} + +/// A loop over a body section. `outputs` snapshot the last iteration's +/// values; `accumulated_outputs` collect a value per iteration and emit the +/// whole sequence at termination. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LoopSpec { + pub name: String, + pub body: Box
, + pub max_iters: u32, + #[serde(default)] + pub outputs: Vec, + #[serde(default)] + pub accumulated_outputs: Vec, +} + +/// The composable graph structure a model declares per walk. Sequential and +/// Parallel exist for construction ergonomics and (later) worker-graph +/// splitting; runtime readiness is pure dataflow, so compilation flattens +/// them. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum Section { + Node(NodeSpec), + Sequential { sections: Vec
}, + Parallel { sections: Vec
}, + Loop(LoopSpec), +} + +/// A loop after compilation: membership plus termination/emission spec. +#[derive(Debug, Clone)] +pub struct CompiledLoop { + pub name: String, + pub max_iters: u32, + /// Direct member nodes — nodes whose INNERMOST enclosing loop is this one. + pub members: BTreeSet, + pub outputs: Vec, + pub accumulated_outputs: Vec, + /// Enclosing loop, for nesting (mstar's loop-inside-loop): an inner loop + /// is an *entity* of its parent's iteration — the parent's iteration + /// completes only once the inner loop has fully terminated. + pub parent: Option, + /// Direct child loops (indices into `CompiledWalk::loops`). + pub children: Vec, +} + +/// A walk's graph, flattened for execution: nodes by name, loops, and the +/// node -> loop membership index. +#[derive(Debug, Clone)] +pub struct CompiledWalk { + pub name: String, + pub nodes: BTreeMap, + pub loops: Vec, + pub node_loop: BTreeMap, +} + +impl CompiledWalk { + pub fn compile(walk_name: &str, section: &Section) -> Result { + let mut walk = CompiledWalk { + name: walk_name.to_string(), + nodes: BTreeMap::new(), + loops: Vec::new(), + node_loop: BTreeMap::new(), + }; + walk.collect(section, None)?; + walk.validate()?; + Ok(walk) + } + + fn collect(&mut self, section: &Section, enclosing_loop: Option) -> Result<()> { + match section { + Section::Node(node) => { + if self.nodes.contains_key(&node.name) { + return Err(CoreError::DuplicateNode( + node.name.clone(), + self.name.clone(), + )); + } + self.nodes.insert(node.name.clone(), node.clone()); + if let Some(loop_idx) = enclosing_loop { + self.loops[loop_idx].members.insert(node.name.clone()); + self.node_loop.insert(node.name.clone(), loop_idx); + } + Ok(()) + } + Section::Sequential { sections } | Section::Parallel { sections } => { + for s in sections { + self.collect(s, enclosing_loop)?; + } + Ok(()) + } + Section::Loop(spec) => { + if spec.max_iters == 0 { + return Err(CoreError::InvalidSpec(format!( + "loop '{}' has max_iters == 0", + spec.name + ))); + } + let loop_idx = self.loops.len(); + self.loops.push(CompiledLoop { + name: spec.name.clone(), + max_iters: spec.max_iters, + members: BTreeSet::new(), + outputs: spec.outputs.clone(), + accumulated_outputs: spec.accumulated_outputs.clone(), + parent: enclosing_loop, + children: Vec::new(), + }); + if let Some(outer) = enclosing_loop { + self.loops[outer].children.push(loop_idx); + } + self.collect(&spec.body, Some(loop_idx))?; + if self.loops[loop_idx].members.is_empty() + && self.loops[loop_idx].children.is_empty() + { + return Err(CoreError::InvalidSpec(format!( + "loop '{}' has an empty body", + spec.name + ))); + } + Ok(()) + } + } + } + + fn validate(&self) -> Result<()> { + if self.nodes.is_empty() { + return Err(CoreError::InvalidSpec(format!( + "walk '{}' has no nodes", + self.name + ))); + } + // Every internal edge destination must exist. + let edge_iter = self + .nodes + .values() + .flat_map(|n| n.outputs.iter()) + .chain(self.loops.iter().flat_map(|l| l.outputs.iter())) + .chain(self.loops.iter().flat_map(|l| l.accumulated_outputs.iter())); + for edge in edge_iter { + // Streaming edges leave this walk (their next_node lives in the + // target partition's walk); everything else must resolve here. + if edge.target_partition.is_none() + && edge.next_node != EMIT_TO_CLIENT + && edge.next_node != EMPTY_DESTINATION + && !self.nodes.contains_key(&edge.next_node) + { + return Err(CoreError::UnknownNode( + edge.next_node.clone(), + self.name.clone(), + )); + } + } + Ok(()) + } + + pub fn loop_index(&self, loop_name: &str) -> Result { + self.loops + .iter() + .position(|l| l.name == loop_name) + .ok_or_else(|| CoreError::UnknownLoop(loop_name.to_string(), self.name.clone())) + } +} + +/// A model's full set of named walks (`prefill`, `decode`, ...), compiled. +/// This is what `get_graph_walk_graphs()` returns in Python mstar. +#[derive(Debug, Clone)] +pub struct WalkSet { + pub walks: BTreeMap>, +} + +impl WalkSet { + /// Build from the JSON spec the Python side sends: + /// `{"walk_name":
, ...}`. + pub fn from_json(json: &str) -> Result { + let raw: BTreeMap = serde_json::from_str(json) + .map_err(|e| CoreError::InvalidSpec(format!("bad walk-set JSON: {e}")))?; + let mut walks = BTreeMap::new(); + for (name, section) in &raw { + walks.insert(name.clone(), Arc::new(CompiledWalk::compile(name, section)?)); + } + Ok(Self { walks }) + } + + pub fn get(&self, walk: &str) -> Result> { + self.walks + .get(walk) + .cloned() + .ok_or_else(|| CoreError::UnknownWalk(walk.to_string())) + } +} diff --git a/rust/src/core/mod.rs b/rust/src/core/mod.rs new file mode 100644 index 000000000..0d73c7128 --- /dev/null +++ b/rust/src/core/mod.rs @@ -0,0 +1,17 @@ +//! The graph/walk core compiled walk graphs and the +//! per-request walk state machine — the Rust port of the runtime behavior of +//! `mstar/graph/base.py`'s `GraphNode`/`Loop` registries and +//! `WorkerGraphIO`. Parity is asserted by `test/rust/test_walk_parity.py`, +//! which drives both implementations with identical event sequences. + +pub mod error; +pub mod graph; +pub mod sched; +pub mod tensor; +pub mod walk; + +pub use error::{CoreError, Result}; +pub use graph::{CompiledWalk, Section, WalkSet, EMIT_TO_CLIENT, EMPTY_DESTINATION}; +pub use tensor::{TensorRef, Uuid}; +pub use sched::{BatchFilter, MicroScheduler, ReadyEntry, ScheduledBatch, SchedulingType}; +pub use walk::{CompletionResult, IncomingInput, RouteEvent, WalkState}; diff --git a/rust/src/core/sched.rs b/rust/src/core/sched.rs new file mode 100644 index 000000000..d87858c52 --- /dev/null +++ b/rust/src/core/sched.rs @@ -0,0 +1,439 @@ +//! The micro-scheduler the Rust port of +//! `mstar/worker/micro_scheduler.py`'s decision logic. +//! +//! The seam is decide-vs-mutate: the caller hands a SNAPSHOT of ready work +//! (one [`ReadyEntry`] per ready (node, request) pair, with the engine-level +//! readiness and priority already evaluated — those live with the engines on +//! the Python side), plus the current monotonic time; the scheduler returns +//! which (node, walk, requests) to batch. Popping the ready queues and +//! executing stay with the worker. Time is always passed in, never read — +//! decisions are deterministic and testable. +//! +//! Ported semantics, asserted by `test/rust/test_sched_parity.py`: +//! round-robin by least-recent (node, walk) batch number; priority mode +//! (lowest engine priority, then the walk with the most requests); OOM +//! hold-with-backoff; deferred removes; leader-node filtering; target / +//! exclude filters; max-batch truncation; TP-follower batches first, with +//! the consecutive-batch cap yielding to other ready work (fairness). +//! +//! On exact ties (equal round-robin recency, equal walk counts) the Python +//! implementation's choice follows set/dict iteration order; here the first +//! entry in snapshot order wins. Callers that need reproducibility across +//! the two must not depend on tie order (mstar's does not). + +use std::collections::{BTreeMap, BTreeSet, VecDeque}; + +/// One ready (node, request) pair, with engine state pre-evaluated. +#[derive(Debug, Clone)] +pub struct ReadyEntry { + pub node: String, + pub walk: String, + pub request_id: String, + pub worker_graph_id: String, + /// `engine.check_ready(node, rid, fwd_info)` — e.g. KV cache read in. + pub engine_ready: bool, + /// Engine priority (lower schedules first in priority mode; mstar: + /// KV_CACHE = 0, STATELESS = 2, unknown = 99). + pub priority: u32, + /// Whether this node may INITIATE a batch on this rank + /// (mstar's `parallel_leader_nodes`; follower ranks replay instead). + pub leader: bool, +} + +/// The scheduling decision. The caller pops `request_ids` (in order) for +/// `node` from their worker-graph queues and executes. +#[derive(Debug, Clone, PartialEq)] +pub struct ScheduledBatch { + pub node: String, + pub walk: String, + pub request_ids: Vec, + pub worker_graph_ids: Vec, + /// True when this batch replays a TP leader's decision. + pub tp_follow: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SchedulingType { + Priority, + RoundRobin, +} + +/// Optional constraints for one `get_next_batch` call. +#[derive(Debug, Default, Clone)] +pub struct BatchFilter { + pub max_batch_size: Option, + pub target_node: Option, + pub target_walk: Option, + /// Skip this (node, walk) pair entirely. + pub exclude_target: Option<(String, String)>, +} + +#[derive(Debug, Clone)] +struct TpFollow { + node: String, + walk: String, + request_ids: Vec, +} + +#[derive(Debug)] +pub struct MicroScheduler { + sched_type: SchedulingType, + batch_number: u64, + /// (node, walk) -> batch number of its last scheduled batch (round-robin). + last_batch_num: BTreeMap<(String, String), u64>, + /// request -> monotonic ms until which it is held (OOM backoff). + held_until: BTreeMap, + /// Requests with a deferred remove: stop initiating new work. + pending_removes: BTreeSet, + /// Leader decisions to replay, in arrival order. + tp_pending: VecDeque, + consec_tp_follower_batches: u32, + max_consec_tp_follower_batches: u32, + /// mstar's `HOLD_BACKOFF_SECONDS` (50 ms), in ms. + hold_backoff_ms: u64, +} + +impl MicroScheduler { + pub fn new(sched_type: SchedulingType, max_consec_tp_follower_batches: u32) -> Self { + Self { + sched_type, + batch_number: 0, + last_batch_num: BTreeMap::new(), + held_until: BTreeMap::new(), + pending_removes: BTreeSet::new(), + tp_pending: VecDeque::new(), + consec_tp_follower_batches: 0, + max_consec_tp_follower_batches, + hold_backoff_ms: 50, + } + } + + /// OOM backoff: hold these requests for `hold_backoff_ms` from `now_ms`. + pub fn hold_requests(&mut self, request_ids: &[String], now_ms: u64) { + for rid in request_ids { + self.held_until + .insert(rid.clone(), now_ms + self.hold_backoff_ms); + } + } + + pub fn add_pending_remove(&mut self, request_id: &str) { + self.pending_removes.insert(request_id.to_string()); + } + + pub fn clear_pending_remove(&mut self, request_id: &str) { + self.pending_removes.remove(request_id); + } + + /// A TP leader's batch decision to replay on this follower rank. + pub fn register_tp_follow(&mut self, node: String, walk: String, request_ids: Vec) { + self.tp_pending.push_back(TpFollow { + node, + walk, + request_ids, + }); + } + + /// Admissible = not held, not pending-remove, engine-ready. + fn admissible(&self, e: &ReadyEntry, now_ms: u64) -> bool { + if !e.engine_ready || self.pending_removes.contains(&e.request_id) { + return false; + } + match self.held_until.get(&e.request_id) { + Some(&t) => t <= now_ms, + None => true, + } + } + + /// Any admissible ready work other than `exclude`? (mstar's + /// `has_ready_excluding` — the spec-chain fairness peek. Note: like + /// mstar's, this does NOT apply the leader filter.) + pub fn has_ready_excluding( + &self, + ready: &[ReadyEntry], + exclude: Option<(&str, &str)>, + now_ms: u64, + ) -> bool { + ready.iter().any(|e| { + self.admissible(e, now_ms) + && exclude != Some((e.node.as_str(), e.walk.as_str())) + }) + } + + fn try_schedule_tp_follow( + &mut self, + ready: &[ReadyEntry], + now_ms: u64, + ) -> Option { + let head = self.tp_pending.front()?.clone(); + // Fairness: after N consecutive follower batches, yield if anything + // else is ready (identical to the leader's spec-chain cap). + if self.consec_tp_follower_batches >= self.max_consec_tp_follower_batches + && self.has_ready_excluding( + ready, + Some((head.node.as_str(), head.walk.as_str())), + now_ms, + ) + { + return None; + } + // Every rid of the leader's decision must be ready here (graph-level + // AND engine-level); otherwise wait — replay order is fixed. + let mut worker_graph_ids = Vec::with_capacity(head.request_ids.len()); + for rid in &head.request_ids { + let entry = ready.iter().find(|e| { + e.node == head.node && e.request_id == *rid + })?; + if !entry.engine_ready { + return None; + } + worker_graph_ids.push(entry.worker_graph_id.clone()); + } + self.batch_number += 1; + self.last_batch_num + .insert((head.node.clone(), head.walk.clone()), self.batch_number); + self.tp_pending.pop_front(); + Some(ScheduledBatch { + node: head.node, + walk: head.walk, + request_ids: head.request_ids, + worker_graph_ids, + tp_follow: true, + }) + } + + /// The scheduling decision (mstar's `get_next_batch`). + pub fn get_next_batch( + &mut self, + ready: &[ReadyEntry], + filter: &BatchFilter, + now_ms: u64, + ) -> Option { + // Expire stale holds (mstar does this each call). + self.held_until.retain(|_, &mut t| t > now_ms); + + match self.try_schedule_tp_follow(ready, now_ms) { + Some(batch) => { + self.consec_tp_follower_batches += 1; + return Some(batch); + } + None => self.consec_tp_follower_batches = 0, + } + + // Group admissible leader entries by node, preserving snapshot order. + let mut by_node: Vec<(String, Vec<&ReadyEntry>)> = Vec::new(); + for e in ready { + if !e.leader || !self.admissible(e, now_ms) { + continue; + } + if let Some(t) = &filter.target_node { + if e.node != *t { + continue; + } + } + if let Some(t) = &filter.target_walk { + if e.walk != *t { + continue; + } + } + if let Some((xn, xw)) = &filter.exclude_target { + if e.node == *xn && e.walk == *xw { + continue; + } + } + match by_node.iter_mut().find(|(n, _)| *n == e.node) { + Some((_, v)) => v.push(e), + None => by_node.push((e.node.clone(), vec![e])), + } + } + if by_node.is_empty() { + return None; + } + + let (node, walk) = match self.sched_type { + SchedulingType::Priority => Self::select_priority(&by_node)?, + SchedulingType::RoundRobin => self.select_round_robin(&by_node)?, + }; + + let mut request_ids = Vec::new(); + let mut worker_graph_ids = Vec::new(); + for (n, entries) in &by_node { + if *n != node { + continue; + } + for e in entries { + if e.walk == walk { + request_ids.push(e.request_id.clone()); + worker_graph_ids.push(e.worker_graph_id.clone()); + } + } + } + if let Some(cap) = filter.max_batch_size { + request_ids.truncate(cap); + worker_graph_ids.truncate(cap); + } + if request_ids.is_empty() { + return None; + } + + self.batch_number += 1; + self.last_batch_num + .insert((node.clone(), walk.clone()), self.batch_number); + Some(ScheduledBatch { + node, + walk, + request_ids, + worker_graph_ids, + tp_follow: false, + }) + } + + /// Lowest engine priority wins; within it, the walk with the most + /// requests (mstar maximizes batch size; the rest wait a cycle). + fn select_priority(by_node: &[(String, Vec<&ReadyEntry>)]) -> Option<(String, String)> { + let (node, entries) = by_node + .iter() + .min_by_key(|(_, entries)| entries.first().map(|e| e.priority).unwrap_or(99))?; + let mut walk_counts: Vec<(String, usize)> = Vec::new(); + for e in entries { + match walk_counts.iter_mut().find(|(w, _)| *w == e.walk) { + Some((_, c)) => *c += 1, + None => walk_counts.push((e.walk.clone(), 1)), + } + } + let (walk, _) = walk_counts.into_iter().max_by_key(|&(_, c)| c)?; + Some((node.clone(), walk)) + } + + /// Least-recently-batched (node, walk) wins (mstar's round-robin). + fn select_round_robin( + &self, + by_node: &[(String, Vec<&ReadyEntry>)], + ) -> Option<(String, String)> { + let mut best: Option<(u64, String, String)> = None; + for (node, entries) in by_node { + for e in entries { + let step = self + .last_batch_num + .get(&(node.clone(), e.walk.clone())) + .copied() + .unwrap_or(0); + if best.as_ref().map(|(s, _, _)| step < *s).unwrap_or(true) { + best = Some((step, node.clone(), e.walk.clone())); + } + } + } + best.map(|(_, n, w)| (n, w)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn entry(node: &str, walk: &str, rid: &str) -> ReadyEntry { + ReadyEntry { + node: node.into(), + walk: walk.into(), + request_id: rid.into(), + worker_graph_id: "wg0".into(), + engine_ready: true, + priority: 0, + leader: true, + } + } + + #[test] + fn round_robin_rotates_across_node_walks() { + let mut s = MicroScheduler::new(SchedulingType::RoundRobin, 1); + let ready = vec![entry("A", "w", "r1"), entry("B", "w", "r2")]; + let f = BatchFilter::default(); + let b1 = s.get_next_batch(&ready, &f, 0).unwrap(); + let b2 = s.get_next_batch(&ready, &f, 0).unwrap(); + assert_ne!(b1.node, b2.node); + } + + #[test] + fn priority_prefers_low_then_biggest_walk() { + let mut s = MicroScheduler::new(SchedulingType::Priority, 1); + let mut kv1 = entry("KV", "decode", "r1"); + kv1.priority = 0; + let mut kv2 = entry("KV", "prefill", "r2"); + kv2.priority = 0; + let mut kv3 = entry("KV", "prefill", "r3"); + kv3.priority = 0; + let mut st = entry("VOC", "decode", "r4"); + st.priority = 2; + let b = s + .get_next_batch(&[kv1, kv2, kv3, st], &BatchFilter::default(), 0) + .unwrap(); + assert_eq!((b.node.as_str(), b.walk.as_str()), ("KV", "prefill")); + assert_eq!(b.request_ids, vec!["r2", "r3"]); + } + + #[test] + fn holds_expire_after_backoff() { + let mut s = MicroScheduler::new(SchedulingType::RoundRobin, 1); + s.hold_requests(&["r1".to_string()], 1000); + let ready = vec![entry("A", "w", "r1")]; + assert!(s.get_next_batch(&ready, &BatchFilter::default(), 1010).is_none()); + assert!(s.get_next_batch(&ready, &BatchFilter::default(), 1051).is_some()); + } + + #[test] + fn tp_follow_first_with_fairness_cap() { + let mut s = MicroScheduler::new(SchedulingType::RoundRobin, 1); + s.register_tp_follow("A".into(), "w".into(), vec!["r1".into()]); + s.register_tp_follow("A".into(), "w".into(), vec!["r1".into()]); + let ready = vec![entry("A", "w", "r1"), entry("B", "w", "r2")]; + let f = BatchFilter::default(); + let b1 = s.get_next_batch(&ready, &f, 0).unwrap(); + assert!(b1.tp_follow); + // consec cap = 1 and B is ready: the second follow yields to B. + let b2 = s.get_next_batch(&ready, &f, 0).unwrap(); + assert!(!b2.tp_follow); + assert_eq!(b2.node, "B"); + // then the queued follow goes through. + let b3 = s.get_next_batch(&ready, &f, 0).unwrap(); + assert!(b3.tp_follow); + } + + #[test] + fn tp_follow_waits_for_all_rids_ready() { + let mut s = MicroScheduler::new(SchedulingType::RoundRobin, 1); + s.register_tp_follow("A".into(), "w".into(), vec!["r1".into(), "r2".into()]); + // On a follower rank the followed node is not leader-initiable + // (mstar's parallel_leader_nodes filter) — model that here. + let follower = |rid: &str| { + let mut e = entry("A", "w", rid); + e.leader = false; + e + }; + let only_r1 = vec![follower("r1")]; + assert!(s.get_next_batch(&only_r1, &BatchFilter::default(), 0).is_none()); + let both = vec![follower("r1"), follower("r2")]; + let b = s.get_next_batch(&both, &BatchFilter::default(), 0).unwrap(); + assert_eq!(b.request_ids, vec!["r1", "r2"]); + } + + #[test] + fn filters_and_truncation() { + let mut s = MicroScheduler::new(SchedulingType::RoundRobin, 1); + let ready = vec![ + entry("A", "w", "r1"), + entry("A", "w", "r2"), + entry("A", "w", "r3"), + ]; + let f = BatchFilter { + max_batch_size: Some(2), + ..Default::default() + }; + let b = s.get_next_batch(&ready, &f, 0).unwrap(); + assert_eq!(b.request_ids.len(), 2); + + let f = BatchFilter { + exclude_target: Some(("A".into(), "w".into())), + ..Default::default() + }; + assert!(s.get_next_batch(&ready, &f, 0).is_none()); + } +} diff --git a/rust/src/core/tensor.rs b/rust/src/core/tensor.rs new file mode 100644 index 000000000..4e72552af --- /dev/null +++ b/rust/src/core/tensor.rs @@ -0,0 +1,26 @@ +use serde::{Deserialize, Serialize}; + +/// Process-wide unique tensor identity. Assigned by the runtime; the Python +/// side keys its `uuid -> torch.Tensor` object store on this. +pub type Uuid = u64; + +/// Descriptor for a tensor living on the data plane. The control plane never +/// sees tensor bytes — only this. Mirrors mstar's `TensorPointerInfo`, minus +/// the transport-specific fields (address/offset/session) which belong to the +/// multi-process transport tier (T4). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct TensorRef { + pub uuid: Uuid, + pub dims: Vec, + pub dtype: String, +} + +impl TensorRef { + pub fn new(uuid: Uuid, dims: Vec, dtype: impl Into) -> Self { + Self { + uuid, + dims, + dtype: dtype.into(), + } + } +} diff --git a/rust/src/core/walk.rs b/rust/src/core/walk.rs new file mode 100644 index 000000000..a53882e67 --- /dev/null +++ b/rust/src/core/walk.rs @@ -0,0 +1,1086 @@ +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; + +use super::error::{CoreError, Result}; +use super::graph::{CompiledWalk, EdgeSpec, EMIT_TO_CLIENT, EMPTY_DESTINATION}; +use super::tensor::TensorRef; + +/// An external input injected into a walk (from the policy's forward-pass +/// args, or — in a future multi-worker world — from a peer worker). +#[derive(Debug, Clone)] +pub struct IncomingInput { + pub node: String, + pub name: String, + pub tensors: Vec, +} + +/// Something the walk state machine routed outward while completing a node. +#[derive(Debug, Clone, PartialEq)] +pub enum RouteEvent { + /// Edge destined for the client (`EMIT_TO_CLIENT`). + Emission { + name: String, + modality: Option, + tensors: Vec, + }, + /// Edge marked `persist: true` — a walk output for the policy. + Persist { + name: String, + tensors: Vec, + }, + /// Streaming edge: goes to the stream buffer of the connection into + /// `target_partition` (the runtime owns buffers; walks don't see them). + Stream { + name: String, + target_partition: String, + tensors: Vec, + }, +} + +#[derive(Debug, Default)] +pub struct CompletionResult { + pub events: Vec, + pub walk_done: bool, +} + +#[derive(Debug, Default)] +struct NodeState { + /// Inputs received for the current iteration. + current: BTreeMap>, + /// Inputs buffered for the next loop iteration (mstar's `ready_next_iter`): + /// filled when an input arrives at a node that already has that input or + /// already ran this iteration — i.e. loop-back edges. + next_iter: BTreeMap>, + completed: bool, + scheduled: bool, +} + +#[derive(Debug, Default)] +struct LoopState { + curr_iter: u32, + finish_signal: bool, + terminated: bool, + /// Latest value per output name (snapshot for `outputs`). + last_values: BTreeMap>, + /// Per-iteration values per output name (for `accumulated_outputs`). + accumulated: BTreeMap>, + /// External inputs into loop members, re-injected on every advance + /// (mstar's `_ingested_external_inputs`). + external_inputs: Vec<(String, String, Vec)>, + /// Child loops that terminated within THIS iteration of this loop. An + /// inner loop is an entity of its parent's iteration (mstar's nested + /// loops): the parent's iteration completes only when its member nodes + /// completed AND every child loop terminated. + children_done: BTreeSet, +} + +/// Per-request state machine over one compiled walk graph. Ports the runtime +/// behavior of `GraphNode`/`Loop`/`WorkerGraphIO` from `mstar/graph/`. +#[derive(Debug)] +pub struct WalkState { + graph: Arc, + nodes: BTreeMap, + loops: Vec, +} + +impl WalkState { + pub fn new(graph: Arc) -> Self { + let nodes = graph + .nodes + .keys() + .map(|name| (name.clone(), NodeState::default())) + .collect(); + let loops = graph.loops.iter().map(|_| LoopState::default()).collect(); + Self { + graph, + nodes, + loops, + } + } + + pub fn walk_name(&self) -> &str { + &self.graph.name + } + + /// Inject external inputs (walk seeding). External inputs into loop + /// members are recorded for re-injection on each loop advance. + pub fn seed(&mut self, inputs: Vec) -> Result<()> { + for input in inputs { + if !self.graph.nodes.contains_key(&input.node) { + return Err(CoreError::UnknownNode( + input.node.clone(), + self.graph.name.clone(), + )); + } + if let Some(&loop_idx) = self.graph.node_loop.get(&input.node) { + self.loops[loop_idx].external_inputs.push(( + input.node.clone(), + input.name.clone(), + input.tensors.clone(), + )); + } + self.ingest(&input.node, &input.name, input.tensors); + } + Ok(()) + } + + fn ingest(&mut self, node: &str, name: &str, tensors: Vec) { + let state = self.nodes.get_mut(node).expect("node validated"); + if state.completed || state.scheduled || state.current.contains_key(name) { + state.next_iter.insert(name.to_string(), tensors); + } else { + state.current.insert(name.to_string(), tensors); + } + } + + /// Nodes whose inputs are all present and are not running/finished. + pub fn ready_nodes(&self) -> Vec { + self.graph + .nodes + .iter() + .filter(|(name, spec)| { + let st = &self.nodes[*name]; + !st.completed + && !st.scheduled + && spec.input_names.iter().all(|i| st.current.contains_key(i)) + }) + .map(|(name, _)| name.clone()) + .collect() + } + + /// Mark a ready node as scheduled and hand back its input tensors for + /// the data plane to execute with. + pub fn take_node_inputs(&mut self, node: &str) -> Result>> { + let spec = self + .graph + .nodes + .get(node) + .ok_or_else(|| CoreError::UnknownNode(node.to_string(), self.graph.name.clone()))?; + let st = self.nodes.get_mut(node).expect("node validated"); + let missing: Vec = spec + .input_names + .iter() + .filter(|i| !st.current.contains_key(*i)) + .cloned() + .collect(); + if st.completed || st.scheduled || !missing.is_empty() { + return Err(CoreError::NodeNotReady { + node: node.to_string(), + missing, + }); + } + st.scheduled = true; + Ok(st.current.clone()) + } + + fn stream_target(&self, input_name: &str) -> Option<&str> { + self.graph.nodes.iter().find_map(|(name, spec)| { + let st = &self.nodes[name]; + (spec.input_names.contains(input_name) + && !st.completed + && !st.scheduled + && !st.current.contains_key(input_name)) + .then_some(name.as_str()) + }) + } + + /// Can a stream chunk for `input_name` be ingested right now? (mstar's + /// `process_new_streaming_inputs` with `can_buffer=False` — chunks are + /// only delivered into a node that is waiting for exactly this input.) + pub fn can_accept_input(&self, input_name: &str) -> bool { + self.stream_target(input_name).is_some() + } + + /// Ingest a stream chunk (the window, in order). Caller must have + /// checked `can_accept_input`. + pub fn ingest_stream_input(&mut self, input_name: &str, tensors: Vec) { + let node = self + .stream_target(input_name) + .expect("caller checked can_accept_input") + .to_string(); + self.ingest(&node, input_name, tensors); + } + + /// Whether `node` is a loop body that will run again after the current + /// iteration — the speculation predicate (mstar's `_can_speculate`): + /// the node's innermost loop is live, has iterations left, and no finish + /// signal is pending. Conservative: any ancestor loop being on its final + /// iteration doesn't matter (the INNER loop advancing is what re-runs the + /// node), but a pending finish signal anywhere in the chain vetoes. + pub fn node_continues_in_loop(&self, node: &str) -> bool { + let Some(&idx) = self.graph.node_loop.get(node) else { + return false; // not in a loop: runs once + }; + let mut i = Some(idx); + while let Some(li) = i { + let lst = &self.loops[li]; + if lst.terminated || lst.finish_signal { + return false; + } + i = self.graph.loops[li].parent; + } + let lst = &self.loops[idx]; + lst.curr_iter + 1 < self.graph.loops[idx].max_iters + } + + /// Record a finish signal for a loop (mstar's `check_stop -> STOP_LOOPS`). + /// Takes effect when the current iteration completes. + pub fn signal_loop_finish(&mut self, loop_name: &str) -> Result<()> { + let idx = self.graph.loop_index(loop_name)?; + self.loops[idx].finish_signal = true; + Ok(()) + } + + /// The data plane finished executing `node`; route its named outputs. + pub fn complete_node( + &mut self, + node: &str, + outputs: BTreeMap>, + ) -> Result { + // Arc-clone the compiled graph (refcount bump) so the spec can be + // borrowed while node/loop state is mutated — no per-complete clone + // of the NodeSpec and its edge strings. + let graph = self.graph.clone(); + let spec = graph + .nodes + .get(node) + .ok_or_else(|| CoreError::UnknownNode(node.to_string(), self.graph.name.clone()))?; + { + let st = self.nodes.get_mut(node).expect("node validated"); + if !st.scheduled { + return Err(CoreError::NodeNotScheduled(node.to_string())); + } + st.scheduled = false; + st.completed = true; + } + + let mut result = CompletionResult::default(); + let loop_idx = self.graph.node_loop.get(node).copied(); + + // Capture loop output values produced by this node — along the whole + // ancestor chain, since an inner node's output may be an OUTER loop's + // declared output (mstar's nested loops). + let mut ancestor = loop_idx; + while let Some(idx) = ancestor { + let loop_spec = &self.graph.loops[idx]; + let capture: Vec<(String, bool)> = loop_spec + .outputs + .iter() + .map(|e| (e.name.clone(), false)) + .chain( + loop_spec + .accumulated_outputs + .iter() + .map(|e| (e.name.clone(), true)), + ) + .collect(); + let lst = &mut self.loops[idx]; + for (name, accumulate) in capture { + if let Some(tensors) = outputs.get(&name) { + if accumulate { + lst.accumulated + .entry(name.clone()) + .or_default() + .extend(tensors.iter().cloned()); + } + lst.last_values.insert(name, tensors.clone()); + } + } + ancestor = self.graph.loops[idx].parent; + } + + // Route the node's own edges. + for edge in &spec.outputs { + let tensors = outputs.get(&edge.name).cloned().unwrap_or_default(); + Self::route_edge(&mut self.nodes, &mut result.events, edge, tensors); + } + + // Loop iteration bookkeeping. + if let Some(idx) = loop_idx { + self.try_complete_loop(idx, &mut result.events); + } else { + // A completed non-loop node will never re-read its inputs (a + // re-arrival routes into `next_iter`, not `current`). Drop them so + // the runtime's reclaim sweep can free any tensor now unreferenced + // — otherwise every intermediate lingers until request finish. + self.nodes + .get_mut(node) + .expect("node validated") + .current + .clear(); + } + + result.walk_done = self.is_done(); + Ok(result) + } + + /// Collect the uuids of every tensor this walk still holds in a node input + /// slot (`current`/`next_iter`) or a loop-output capture. Used by the + /// runtime's per-request reachability sweep to reclaim shared-memory + /// buffers as soon as a tensor is no longer referenced (rather than at + /// request finish). + pub fn collect_live_uuids(&self, out: &mut BTreeSet) { + for st in self.nodes.values() { + for tensors in st.current.values().chain(st.next_iter.values()) { + out.extend(tensors.iter().map(|t| t.uuid)); + } + } + for lst in &self.loops { + for tensors in lst.last_values.values().chain(lst.accumulated.values()) { + out.extend(tensors.iter().map(|t| t.uuid)); + } + for (_, _, tensors) in &lst.external_inputs { + out.extend(tensors.iter().map(|t| t.uuid)); + } + } + } + + /// This loop and every loop nested under it (depth-first). + fn subtree_loops(&self, idx: usize) -> Vec { + let mut out = vec![idx]; + let mut stack = self.graph.loops[idx].children.clone(); + while let Some(i) = stack.pop() { + out.push(i); + stack.extend(self.graph.loops[i].children.iter().copied()); + } + out + } + + /// Complete this loop's iteration if all its entities are done: every + /// direct member node completed AND every child loop terminated this + /// iteration (an inner loop is an entity of its parent's iteration). + fn try_complete_loop(&mut self, idx: usize, events: &mut Vec) { + if self.loops[idx].terminated { + return; + } + let members_done = self.graph.loops[idx] + .members + .iter() + .all(|m| self.nodes[m].completed); + let children_done = self.graph.loops[idx] + .children + .iter() + .all(|c| self.loops[idx].children_done.contains(c)); + if members_done && children_done { + self.complete_loop_iter(idx, events); + } + } + + /// mstar `Loop.complete_iter`: terminate on max_iters/finish signal, + /// otherwise advance one iteration. + fn complete_loop_iter(&mut self, idx: usize, events: &mut Vec) { + let finished = { + let lst = &self.loops[idx]; + lst.finish_signal || lst.curr_iter + 1 >= self.graph.loops[idx].max_iters + }; + if finished { + let loop_spec = self.graph.loops[idx].clone(); + let (last_values, accumulated) = { + let lst = &mut self.loops[idx]; + lst.terminated = true; + ( + std::mem::take(&mut lst.last_values), + std::mem::take(&mut lst.accumulated), + ) + }; + for edge in &loop_spec.outputs { + let tensors = last_values.get(&edge.name).cloned().unwrap_or_default(); + Self::route_edge(&mut self.nodes, events, edge, tensors); + } + for edge in &loop_spec.accumulated_outputs { + let tensors = accumulated.get(&edge.name).cloned().unwrap_or_default(); + Self::route_edge(&mut self.nodes, events, edge, tensors); + } + // mstar clears the terminated loop's BODY registries (descendant + // loops read iteration 0 again); the loop's own counter persists + // until a parent advance resets it. Matched for `loop_iters` + // parity — descendant state is otherwise reset on the parent's + // next advance anyway. + for li in self.subtree_loops(idx) { + if li != idx { + self.loops[li].curr_iter = 0; + } + } + // The loop has finished: its buffered loop-back signals must NOT + // propagate (mstar's `filtered_signals`) — without this, a parent + // loop's advance would promote the final loop-back value over the + // re-injected external input. + for li in self.subtree_loops(idx) { + let members: Vec = + self.graph.loops[li].members.iter().cloned().collect(); + for member in members { + self.nodes + .get_mut(&member) + .expect("member validated") + .next_iter + .clear(); + } + } + // Nested: this loop is an entity of its parent's iteration — the + // parent may now be complete too. + if let Some(parent) = loop_spec.parent { + self.loops[parent].children_done.insert(idx); + self.try_complete_loop(parent, events); + } + } else { + // Advance: reset the whole SUBTREE. Every transitive member node + // promotes its buffered next-iter inputs; descendant loops reset + // fully (curr_iter, termination, finish signal, output caches) so + // they run afresh inside the new iteration; then external inputs + // of this loop AND descendants re-inject (never overwriting a + // fresher promoted loop-back value). + let subtree = self.subtree_loops(idx); + self.loops[idx].curr_iter += 1; + self.loops[idx].children_done.clear(); + for &li in &subtree { + let members: Vec = + self.graph.loops[li].members.iter().cloned().collect(); + for member in &members { + let st = self.nodes.get_mut(member).expect("member validated"); + st.completed = false; + st.current = std::mem::take(&mut st.next_iter); + } + if li != idx { + let lst = &mut self.loops[li]; + lst.curr_iter = 0; + lst.terminated = false; + lst.finish_signal = false; // a stop applies to one run + lst.children_done.clear(); + lst.last_values.clear(); + lst.accumulated.clear(); + } + } + for &li in &subtree { + let externals = self.loops[li].external_inputs.clone(); + for (node, name, tensors) in externals { + let st = self.nodes.get_mut(&node).expect("member validated"); + st.current.entry(name).or_insert(tensors); + } + } + } + } + + fn route_edge( + nodes: &mut BTreeMap, + events: &mut Vec, + edge: &EdgeSpec, + tensors: Vec, + ) { + if edge.persist { + events.push(RouteEvent::Persist { + name: edge.name.clone(), + tensors: tensors.clone(), + }); + } + if let Some(partition) = &edge.target_partition { + events.push(RouteEvent::Stream { + name: edge.name.clone(), + target_partition: partition.clone(), + tensors, + }); + return; + } + match edge.next_node.as_str() { + EMIT_TO_CLIENT => events.push(RouteEvent::Emission { + name: edge.name.clone(), + modality: edge.output_modality.clone(), + tensors, + }), + EMPTY_DESTINATION => {} + dest => { + let state = nodes.get_mut(dest).expect("edges validated at compile"); + if state.completed || state.scheduled || state.current.contains_key(&edge.name) { + state.next_iter.insert(edge.name.clone(), tensors); + } else { + state.current.insert(edge.name.clone(), tensors); + } + } + } + } + + /// A walk is done when every node has completed and every loop has + /// terminated. One walk completion == one forward pass, after which the + /// policy picks the next walk (or finishes the request). + /// Current iteration per loop, by name (mstar's `get_loop_indices`). + pub fn loop_iters(&self) -> Vec<(String, u32)> { + self.graph + .loops + .iter() + .zip(self.loops.iter()) + .map(|(spec, st)| (spec.name.clone(), st.curr_iter)) + .collect() + } + + /// (name, curr_iter, terminated) per loop — the adapter uses the + /// terminated flag flips to compute filtered loop-back signals and the + /// curr_iter bumps to re-emit external inputs, mirroring mstar's + /// completion contract without duplicating loop mechanics. + pub fn loop_states(&self) -> Vec<(String, u32, bool)> { + self.graph + .loops + .iter() + .zip(self.loops.iter()) + .map(|(spec, st)| (spec.name.clone(), st.curr_iter, st.terminated)) + .collect() + } + + pub fn is_done(&self) -> bool { + self.graph + .loops + .iter() + .enumerate() + .all(|(i, _)| self.loops[i].terminated) + && self + .nodes + .iter() + .all(|(name, st)| match self.graph.node_loop.get(name) { + Some(&idx) => self.loops[idx].terminated, + None => st.completed, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::graph::{LoopSpec, NodeSpec, Section, WalkSet}; + + fn tref(uuid: u64) -> TensorRef { + TensorRef::new(uuid, vec![1], "float32") + } + + fn node(name: &str, inputs: &[&str], outputs: Vec) -> Section { + Section::Node(NodeSpec { + name: name.to_string(), + input_names: inputs.iter().map(|s| s.to_string()).collect(), + outputs, + }) + } + + fn edge(next: &str, name: &str) -> EdgeSpec { + EdgeSpec { + next_node: next.to_string(), + name: name.to_string(), + persist: false, + output_modality: None, + target_partition: None, + } + } + + fn emit(name: &str, modality: &str, persist: bool) -> EdgeSpec { + EdgeSpec { + next_node: EMIT_TO_CLIENT.to_string(), + name: name.to_string(), + persist, + output_modality: Some(modality.to_string()), + target_partition: None, + } + } + + /// vjepa2 `prefill_video`: video_frames -> [encoder] -> encoder_hidden + /// -> [predictor] -> predicted_hidden -> EMIT_TO_CLIENT. + fn vjepa2_like_walk() -> WalkState { + let section = Section::Sequential { + sections: vec![ + node( + "video_encoder", + &["video_frames"], + vec![edge("predictor", "encoder_hidden")], + ), + node( + "predictor", + &["encoder_hidden"], + vec![emit("predicted_hidden", "video", true)], + ), + ], + }; + WalkState::new(Arc::new( + CompiledWalk::compile("prefill_video", §ion).unwrap(), + )) + } + + #[test] + fn sequential_two_node_walk_end_to_end() { + let mut walk = vjepa2_like_walk(); + walk.seed(vec![IncomingInput { + node: "video_encoder".into(), + name: "video_frames".into(), + tensors: vec![tref(1)], + }]) + .unwrap(); + + assert_eq!(walk.ready_nodes(), vec!["video_encoder".to_string()]); + let inputs = walk.take_node_inputs("video_encoder").unwrap(); + assert_eq!(inputs["video_frames"], vec![tref(1)]); + assert!(walk.ready_nodes().is_empty(), "scheduled node not ready"); + + let res = walk + .complete_node( + "video_encoder", + BTreeMap::from([("encoder_hidden".to_string(), vec![tref(2)])]), + ) + .unwrap(); + assert!(res.events.is_empty()); + assert!(!res.walk_done); + + assert_eq!(walk.ready_nodes(), vec!["predictor".to_string()]); + walk.take_node_inputs("predictor").unwrap(); + let res = walk + .complete_node( + "predictor", + BTreeMap::from([("predicted_hidden".to_string(), vec![tref(3)])]), + ) + .unwrap(); + assert!(res.walk_done); + assert_eq!(res.events.len(), 2); // persist + emission + assert!(res.events.contains(&RouteEvent::Persist { + name: "predicted_hidden".into(), + tensors: vec![tref(3)], + })); + assert!(res.events.contains(&RouteEvent::Emission { + name: "predicted_hidden".into(), + modality: Some("video".into()), + tensors: vec![tref(3)], + })); + } + + #[test] + fn take_requires_ready_and_complete_requires_scheduled() { + let mut walk = vjepa2_like_walk(); + assert!(matches!( + walk.take_node_inputs("video_encoder"), + Err(CoreError::NodeNotReady { .. }) + )); + assert!(matches!( + walk.complete_node("video_encoder", BTreeMap::new()), + Err(CoreError::NodeNotScheduled(_)) + )); + } + + /// pi05-style flow-matching loop: a single LLM node feeding itself + /// noisy_actions/timestep for max_iters, then emitting the final actions. + fn flow_loop_walk(max_iters: u32) -> WalkState { + let body = node( + "LLM", + &["noisy_actions", "timestep_index"], + vec![ + edge("LLM", "noisy_actions"), + edge("LLM", "timestep_index"), + ], + ); + let section = Section::Loop(LoopSpec { + name: "flow".into(), + body: Box::new(body), + max_iters, + outputs: vec![emit("noisy_actions", "action", true)], + accumulated_outputs: vec![], + }); + WalkState::new(Arc::new(CompiledWalk::compile("action_gen", §ion).unwrap())) + } + + fn run_loop_iter(walk: &mut WalkState, out_uuid: u64) -> CompletionResult { + walk.take_node_inputs("LLM").unwrap(); + walk.complete_node( + "LLM", + BTreeMap::from([ + ("noisy_actions".to_string(), vec![tref(out_uuid)]), + ("timestep_index".to_string(), vec![tref(out_uuid + 1000)]), + ]), + ) + .unwrap() + } + + #[test] + fn loop_runs_max_iters_and_emits_last_value() { + let mut walk = flow_loop_walk(3); + walk.seed(vec![ + IncomingInput { + node: "LLM".into(), + name: "noisy_actions".into(), + tensors: vec![tref(10)], + }, + IncomingInput { + node: "LLM".into(), + name: "timestep_index".into(), + tensors: vec![tref(11)], + }, + ]) + .unwrap(); + + for iter in 0..3u64 { + assert_eq!(walk.ready_nodes(), vec!["LLM".to_string()], "iter {iter}"); + let res = run_loop_iter(&mut walk, 100 + iter); + if iter < 2 { + assert!(!res.walk_done, "iter {iter} should continue"); + assert!(res.events.is_empty()); + } else { + assert!(res.walk_done); + assert!(res.events.contains(&RouteEvent::Emission { + name: "noisy_actions".into(), + modality: Some("action".into()), + tensors: vec![tref(102)], // last iteration's value + })); + } + } + assert!(walk.ready_nodes().is_empty()); + } + + #[test] + fn loop_back_values_flow_between_iterations() { + let mut walk = flow_loop_walk(2); + walk.seed(vec![ + IncomingInput { + node: "LLM".into(), + name: "noisy_actions".into(), + tensors: vec![tref(10)], + }, + IncomingInput { + node: "LLM".into(), + name: "timestep_index".into(), + tensors: vec![tref(11)], + }, + ]) + .unwrap(); + + run_loop_iter(&mut walk, 100); + // Second iteration must see iteration 1's loop-back outputs, not the + // re-injected seeds (externals must not overwrite loop-back values). + let inputs = walk.take_node_inputs("LLM").unwrap(); + assert_eq!(inputs["noisy_actions"], vec![tref(100)]); + assert_eq!(inputs["timestep_index"], vec![tref(1100)]); + } + + #[test] + fn finish_signal_terminates_loop_early() { + let mut walk = flow_loop_walk(10); + walk.seed(vec![ + IncomingInput { + node: "LLM".into(), + name: "noisy_actions".into(), + tensors: vec![tref(10)], + }, + IncomingInput { + node: "LLM".into(), + name: "timestep_index".into(), + tensors: vec![tref(11)], + }, + ]) + .unwrap(); + + run_loop_iter(&mut walk, 100); + walk.signal_loop_finish("flow").unwrap(); + let res = run_loop_iter(&mut walk, 200); + assert!(res.walk_done, "finish signal should stop before max_iters"); + assert!(res.events.contains(&RouteEvent::Emission { + name: "noisy_actions".into(), + modality: Some("action".into()), + tensors: vec![tref(200)], + })); + } + + #[test] + fn accumulated_outputs_collect_every_iteration() { + let body = node( + "gen", + &["state"], + vec![edge("gen", "state")], + ); + let section = Section::Loop(LoopSpec { + name: "rollout".into(), + body: Box::new(body), + max_iters: 3, + outputs: vec![], + accumulated_outputs: vec![emit("state", "video", false)], + }); + let mut walk = WalkState::new(Arc::new( + CompiledWalk::compile("rollout_walk", §ion).unwrap(), + )); + walk.seed(vec![IncomingInput { + node: "gen".into(), + name: "state".into(), + tensors: vec![tref(1)], + }]) + .unwrap(); + + let mut last = CompletionResult::default(); + for i in 0..3u64 { + walk.take_node_inputs("gen").unwrap(); + last = walk + .complete_node( + "gen", + BTreeMap::from([("state".to_string(), vec![tref(100 + i)])]), + ) + .unwrap(); + } + assert!(last.walk_done); + assert_eq!( + last.events, + vec![RouteEvent::Emission { + name: "state".into(), + modality: Some("video".into()), + tensors: vec![tref(100), tref(101), tref(102)], + }] + ); + } + + #[test] + fn parallel_branches_join() { + // a -> c, b -> c (fan-in), c emits. + let section = Section::Sequential { + sections: vec![ + Section::Parallel { + sections: vec![ + node("a", &["xa"], vec![edge("c", "ya")]), + node("b", &["xb"], vec![edge("c", "yb")]), + ], + }, + node("c", &["ya", "yb"], vec![emit("z", "text", false)]), + ], + }; + let mut walk = WalkState::new(Arc::new(CompiledWalk::compile("par", §ion).unwrap())); + walk.seed(vec![ + IncomingInput { + node: "a".into(), + name: "xa".into(), + tensors: vec![tref(1)], + }, + IncomingInput { + node: "b".into(), + name: "xb".into(), + tensors: vec![tref(2)], + }, + ]) + .unwrap(); + + let mut ready = walk.ready_nodes(); + ready.sort(); + assert_eq!(ready, vec!["a".to_string(), "b".to_string()]); + + walk.take_node_inputs("a").unwrap(); + walk.complete_node("a", BTreeMap::from([("ya".to_string(), vec![tref(3)])])) + .unwrap(); + assert_eq!( + walk.ready_nodes(), + vec!["b".to_string()], + "c waits for both branches; b is still ready" + ); + + walk.take_node_inputs("b").unwrap(); + walk.complete_node("b", BTreeMap::from([("yb".to_string(), vec![tref(4)])])) + .unwrap(); + assert_eq!(walk.ready_nodes(), vec!["c".to_string()]); + + walk.take_node_inputs("c").unwrap(); + let res = walk + .complete_node("c", BTreeMap::from([("z".to_string(), vec![tref(5)])])) + .unwrap(); + assert!(res.walk_done); + } + + #[test] + fn walkset_parses_json() { + let json = r#"{ + "prefill_video": { + "kind": "sequential", + "sections": [ + {"kind": "node", "name": "video_encoder", + "input_names": ["video_frames"], + "outputs": [{"next_node": "predictor", "name": "encoder_hidden"}]}, + {"kind": "node", "name": "predictor", + "input_names": ["encoder_hidden"], + "outputs": [{"next_node": "EMIT_TO_CLIENT", "name": "predicted_hidden", + "persist": true, "output_modality": "video"}]} + ] + } + }"#; + let set = WalkSet::from_json(json).unwrap(); + let walk = set.get("prefill_video").unwrap(); + assert_eq!(walk.nodes.len(), 2); + assert!(set.get("nope").is_err()); + } + + #[test] + fn compile_rejects_bad_specs() { + // Unknown edge destination. + let bad = node("a", &["x"], vec![edge("ghost", "y")]); + assert!(CompiledWalk::compile("w", &bad).is_err()); + + // A loop with an empty body is meaningless. + let empty = Section::Loop(LoopSpec { + name: "empty".into(), + body: Box::new(Section::Sequential { sections: vec![] }), + max_iters: 2, + outputs: vec![], + accumulated_outputs: vec![], + }); + assert!(CompiledWalk::compile("w", &empty).is_err()); + } + + // ---- nested loops (mstar's loop-inside-loop) --------------------------- + + /// outer(max=O) { inner(max=I) { step } } — `step` self-feeds `x`; the + /// inner loop emits its last x on each termination; the outer re-runs the + /// inner each iteration and emits the final value at the end. + fn nested_walk(outer_iters: u32, inner_iters: u32) -> WalkState { + let step = node("step", &["x"], vec![edge("step", "x")]); + let inner = Section::Loop(LoopSpec { + name: "inner".into(), + body: Box::new(step), + max_iters: inner_iters, + outputs: vec![emit("x", "text", true)], + accumulated_outputs: vec![], + }); + let outer = Section::Loop(LoopSpec { + name: "outer".into(), + body: Box::new(inner), + max_iters: outer_iters, + outputs: vec![], + accumulated_outputs: vec![], + }); + WalkState::new(Arc::new(CompiledWalk::compile("nested", &outer).unwrap())) + } + + fn run_step(walk: &mut WalkState, out_uuid: u64) -> CompletionResult { + walk.take_node_inputs("step").unwrap(); + walk.complete_node( + "step", + BTreeMap::from([("x".to_string(), vec![tref(out_uuid)])]), + ) + .unwrap() + } + + #[test] + fn nested_loop_runs_outer_times_inner_iterations() { + // 2 outer x 3 inner = 6 step executions; the walk is done only after + // the LAST inner termination completes the LAST outer iteration. + let mut walk = nested_walk(2, 3); + walk.seed(vec![IncomingInput { + node: "step".into(), + name: "x".into(), + tensors: vec![tref(1)], + }]) + .unwrap(); + + let mut executions = 0u64; + let mut emissions = Vec::new(); + while !walk.is_done() { + assert_eq!(walk.ready_nodes(), vec!["step".to_string()], + "step must be ready each iteration"); + executions += 1; + let result = run_step(&mut walk, 100 + executions); + for ev in result.events { + if let RouteEvent::Emission { tensors, .. } = ev { + emissions.push(tensors[0].uuid); + } + } + assert!(executions <= 6, "ran more than outer*inner iterations"); + } + assert_eq!(executions, 6); + // The inner loop terminated twice (once per outer iteration), emitting + // its last value each time: after executions 3 and 6. + assert_eq!(emissions, vec![103, 106]); + } + + #[test] + fn nested_inner_reruns_from_external_input_each_outer_iteration() { + // The seed (an external input into the inner loop) must re-inject at + // the start of EVERY outer iteration — the inner loop's second run + // starts from the seed again, not from the first run's final value. + let mut walk = nested_walk(2, 2); + walk.seed(vec![IncomingInput { + node: "step".into(), + name: "x".into(), + tensors: vec![tref(7)], + }]) + .unwrap(); + + // Outer iteration 1: inner runs twice (input 7, then loop-back 101). + let inp = walk.take_node_inputs("step").unwrap(); + assert_eq!(inp["x"][0].uuid, 7); + walk.complete_node("step", BTreeMap::from([("x".to_string(), vec![tref(101)])])) + .unwrap(); + let inp = walk.take_node_inputs("step").unwrap(); + assert_eq!(inp["x"][0].uuid, 101); + walk.complete_node("step", BTreeMap::from([("x".to_string(), vec![tref(102)])])) + .unwrap(); + + // Outer iteration 2 begins: the inner loop reset; the SEED re-injected. + assert!(!walk.is_done()); + let inp = walk.take_node_inputs("step").unwrap(); + assert_eq!(inp["x"][0].uuid, 7, "seed must re-inject on outer advance"); + walk.complete_node("step", BTreeMap::from([("x".to_string(), vec![tref(201)])])) + .unwrap(); + let inp = walk.take_node_inputs("step").unwrap(); + assert_eq!(inp["x"][0].uuid, 201); + walk.complete_node("step", BTreeMap::from([("x".to_string(), vec![tref(202)])])) + .unwrap(); + assert!(walk.is_done()); + } + + #[test] + fn nested_inner_finish_signal_applies_to_one_run() { + // Stopping the inner loop ends its CURRENT run; the outer's next + // iteration runs the inner afresh (the signal does not persist). + let mut walk = nested_walk(2, 5); + walk.seed(vec![IncomingInput { + node: "step".into(), + name: "x".into(), + tensors: vec![tref(1)], + }]) + .unwrap(); + + walk.signal_loop_finish("inner").unwrap(); + run_step(&mut walk, 100); // inner terminates after 1 iter (signal) + assert!(!walk.is_done(), "outer has another iteration to run"); + + // Second outer iteration: the inner runs its full 5 iterations. + for i in 0..5 { + assert!(!walk.is_done()); + run_step(&mut walk, 200 + i); + } + assert!(walk.is_done()); + } + + #[test] + fn outer_loop_output_captured_from_inner_node() { + // An inner node's output declared as the OUTER loop's output must be + // captured across the nesting boundary and emitted at outer end. + let step = node("step", &["x"], vec![edge("step", "x")]); + let inner = Section::Loop(LoopSpec { + name: "inner".into(), + body: Box::new(step), + max_iters: 2, + outputs: vec![], + accumulated_outputs: vec![], + }); + let outer = Section::Loop(LoopSpec { + name: "outer".into(), + body: Box::new(inner), + max_iters: 2, + outputs: vec![emit("x", "text", true)], + accumulated_outputs: vec![], + }); + let mut walk = + WalkState::new(Arc::new(CompiledWalk::compile("nested_out", &outer).unwrap())); + walk.seed(vec![IncomingInput { + node: "step".into(), + name: "x".into(), + tensors: vec![tref(1)], + }]) + .unwrap(); + + let mut last_emission = None; + let mut n = 0; + while !walk.is_done() { + n += 1; + let result = run_step(&mut walk, 100 + n); + for ev in result.events { + if let RouteEvent::Emission { tensors, .. } = ev { + last_emission = Some(tensors[0].uuid); + } + } + } + assert_eq!(n, 4); + assert_eq!(last_emission, Some(104), "outer emits the final inner value"); + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 9847321af..8721f7172 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -1,17 +1,31 @@ -//! mstar's Rust transport: the ZMQ PUSH/PULL control mesh. -//! `communicator.rs` is the transport + codec split; this file is the PyO3 -//! surface (`mstar_rust.ZmqCommunicator`) the Python `RustZMQCommunicator` -//! wrapper drives. Build: `maturin develop --release` in rust/. +//! mstar's Rust components, one crate: each module is an independent +//! capability behind its own opt-in flag, and this file is the PyO3 +//! surface the Python wrappers drive. Currently: `communicator.rs` (the +//! ZMQ PUSH/PULL control mesh, as a transport + codec split) and `shm.rs` +//! (the shared-memory tensor arena). The crate is not limited to these — +//! new components land as new modules with their own bindings. Also +//! usable as a plain Rust library (rlib) by Rust-side consumers. +//! Build: `maturin develop --release` in rust/. pub mod communicator; +pub mod core; +pub mod shm; +use std::os::raw::{c_int, c_void}; use std::time::Duration; -use pyo3::exceptions::PyRuntimeError; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::ffi; use pyo3::prelude::*; use pyo3::types::PyBytes; use communicator::{RawZmqCommunicator, RecvEvent}; +use core::graph::CompiledWalk; +use core::sched::{BatchFilter, MicroScheduler, ReadyEntry, SchedulingType}; +use core::tensor::TensorRef; +use core::walk::{IncomingInput, RouteEvent, WalkState}; +use core::WalkSet; +use shm::{SegmentedShmArena, ShmArena}; /// Opaque byte frames over ZMQ PUSH/PULL: ipc or tcp endpoints, lazily-cached /// peers, and wakeup-fd polling (an eventfd wakes `recv_or_wake` instantly). @@ -104,9 +118,524 @@ impl PyZmqCommunicator { } } +/// Shared-memory tensor arena for cross-process transport. Producer: +/// `create(name, size)` -> `reserve(nbytes)` -> `torch.frombuffer( +/// memoryview(arena)[off:off+n], dtype=..).copy_(cpu_tensor)`; send the +/// offset descriptor; `free(off)` on reclaim. Consumer: `open(name)` and +/// `torch.frombuffer(memoryview(arena)[off:off+n], ..)` (then H2D). +#[pyclass(name = "ShmArena")] +struct PyShmArena { + arena: ShmArena, +} + +#[pymethods] +impl PyShmArena { + #[staticmethod] + fn create(name: &str, size: usize) -> PyResult { + Ok(Self { + arena: ShmArena::create(name, size) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?, + }) + } + + #[staticmethod] + fn open(name: &str) -> PyResult { + Ok(Self { + arena: ShmArena::open(name).map_err(|e| PyRuntimeError::new_err(e.to_string()))?, + }) + } + + fn reserve(&self, nbytes: usize) -> PyResult { + self.arena + .reserve(nbytes) + .map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + fn free(&self, offset: usize) -> bool { + self.arena.free(offset) + } + + #[getter] + fn size(&self) -> usize { + self.arena.size() + } + + #[getter] + fn bytes_free(&self) -> usize { + self.arena.bytes_free() + } + + /// `(base_ptr, len)` for the pinning hook — `cudaHostRegister` the + /// mapping once (e.g. `torch.cuda.cudart().cudaHostRegister(ptr, len, 0)`). + fn ptr_len(&self) -> (usize, usize) { + (self.arena.as_mut_ptr() as usize, self.arena.size()) + } + + fn close(&mut self) { + self.arena.close(); + } + + /// Whole arena as a writable memoryview -> zero-copy `torch.frombuffer`. + unsafe fn __getbuffer__( + slf: Bound<'_, Self>, + view: *mut ffi::Py_buffer, + flags: c_int, + ) -> PyResult<()> { + if view.is_null() { + return Err(PyValueError::new_err("null buffer view")); + } + let borrow = slf.borrow(); + let ptr = borrow.arena.as_mut_ptr() as *mut c_void; + let len = borrow.arena.size() as ffi::Py_ssize_t; + let ret = ffi::PyBuffer_FillInfo(view, slf.as_ptr(), ptr, len, 0, flags); + if ret != 0 { + Err(PyErr::fetch(slf.py())) + } else { + Ok(()) + } + } + + unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) {} +} + +/// One segment of a `SegmentedShmArena`, exposed with the buffer protocol so +/// staging stays zero-copy per segment. The mapping never moves, so a +/// memoryview (and a CUDA host-registration of the segment) stays valid for +/// the segment's lifetime. +#[pyclass(name = "ShmSegment")] +struct PyShmSegment { + seg: std::sync::Arc, +} + +#[pymethods] +impl PyShmSegment { + #[getter] + fn size(&self) -> usize { + self.seg.size() + } + + /// `(base_ptr, len)` for the pinning hook (see `ShmArena::ptr_len`). + fn ptr_len(&self) -> (usize, usize) { + (self.seg.as_mut_ptr() as usize, self.seg.size()) + } + + unsafe fn __getbuffer__( + slf: Bound<'_, Self>, + view: *mut ffi::Py_buffer, + flags: c_int, + ) -> PyResult<()> { + if view.is_null() { + return Err(PyValueError::new_err("null buffer view")); + } + let borrow = slf.borrow(); + let ptr = borrow.seg.as_mut_ptr() as *mut c_void; + let len = borrow.seg.size() as ffi::Py_ssize_t; + let ret = ffi::PyBuffer_FillInfo(view, slf.as_ptr(), ptr, len, 0, flags); + if ret != 0 { + Err(PyErr::fetch(slf.py())) + } else { + Ok(()) + } + } + + unsafe fn __releasebuffer__(&self, _view: *mut ffi::Py_buffer) {} +} + +/// Grow-by-segments producer arena with uuid-grouped reclaim. +/// `reserve(n) -> (segment_idx, offset)`; descriptors carry +/// `segment_name(idx)` so consumers keep opening plain `ShmArena`s by name. +/// Segments are created once and never move — registration-friendly. +#[pyclass(name = "SegmentedShmArena")] +struct PySegmentedShmArena { + arena: SegmentedShmArena, +} + +#[pymethods] +impl PySegmentedShmArena { + #[staticmethod] + fn create(base: &str, segment_size: usize, max_segments: usize) -> PyResult { + Ok(Self { + arena: SegmentedShmArena::create(base, segment_size, max_segments) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?, + }) + } + + /// -> (segment_idx, offset); grows by one segment when full (dedicated + /// segment for oversized allocations), errors at the max_segments cap. + /// GIL released: the growth path creates + maps a new shm segment + /// (file create, ftruncate, mmap — milliseconds), which must not stall + /// other Python threads (the serve loop, stream relays). + fn reserve(&self, py: Python<'_>, nbytes: usize) -> PyResult<(usize, usize)> { + py.allow_threads(|| self.arena.reserve(nbytes)) + .map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + /// GIL released: growth (in `reserve`) holds the segments mutex + /// across an mmap for milliseconds; blocking on that mutex with the + /// GIL held would freeze every Python thread for the duration. + fn free(&self, py: Python<'_>, segment: usize, offset: usize) -> bool { + py.allow_threads(|| self.arena.free(segment, offset)) + } + + #[getter] + fn num_segments(&self, py: Python<'_>) -> usize { + py.allow_threads(|| self.arena.num_segments()) + } + + /// `(total_bytes, free_bytes, largest_free_block)` across all segments. + /// `largest_free_block` collapsing while `free_bytes` stays high is the + /// fragmentation signature (allocations fail / segments grow despite + /// healthy total free space). + fn stats(&self, py: Python<'_>) -> (usize, usize, usize) { + py.allow_threads(|| self.arena.stats()) + } + + fn segment_name(&self, i: usize) -> String { + self.arena.segment_name(i) + } + + /// Shared buffer-protocol view of segment `i`. + fn segment(&self, i: usize) -> PyResult { + self.arena + .segment(i) + .map(|seg| PyShmSegment { seg }) + .ok_or_else(|| PyValueError::new_err(format!("no segment {i}"))) + } +} + +/// A model's compiled walk graphs, built from the JSON spec the Python +/// translator (`mstar/graph/rust_core.py`) produces from `GraphSection`s. +#[pyclass(name = "WalkSet")] +struct PyWalkSet { + inner: WalkSet, +} + +#[pymethods] +impl PyWalkSet { + #[staticmethod] + fn from_json(spec: &str) -> PyResult { + Ok(Self { + inner: WalkSet::from_json(spec).map_err(|e| PyValueError::new_err(e.to_string()))?, + }) + } + + #[getter] + fn walk_names(&self) -> Vec { + self.inner.walks.keys().cloned().collect() + } + + /// Fresh per-request walk state over the named walk. + fn state(&self, walk: &str) -> PyResult { + let graph = self + .inner + .get(walk) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + Ok(PyWalkState { + inner: WalkState::new(graph.clone()), + graph, + next_uuid: 1, + }) + } +} + +/// Per-request walk state machine — the Rust `WorkerGraphIO`. Tensor payloads +/// stay on the Python data plane; here every value is an opaque uuid. +#[pyclass(name = "WalkState")] +struct PyWalkState { + inner: WalkState, + graph: std::sync::Arc, + next_uuid: u64, +} + +impl PyWalkState { + fn fresh_ref(&mut self) -> TensorRef { + let u = self.next_uuid; + self.next_uuid += 1; + TensorRef::new(u, vec![], "opaque") + } +} + +#[pymethods] +impl PyWalkState { + /// Inject external inputs: [(node, input_name), ...]. + fn seed(&mut self, inputs: Vec<(String, String)>) -> PyResult<()> { + let seeded = inputs + .into_iter() + .map(|(node, name)| { + let t = self.fresh_ref(); + IncomingInput { + node, + name, + tensors: vec![t], + } + }) + .collect(); + self.inner + .seed(seeded) + .map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + fn ready_nodes(&self) -> Vec { + self.inner.ready_nodes() + } + + /// Claim a ready node for execution (mstar's pop from the ready queue). + fn schedule(&mut self, node: &str) -> PyResult<()> { + self.inner + .take_node_inputs(node) + .map(|_| ()) + .map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + /// Complete a scheduled node, producing a value for each of its declared + /// output edge names. Returns (route_events, walk_done); route_events are + /// (kind, name, target) with kind in {"emission", "persist", "stream"} — + /// internal edges route inside the state machine. + fn complete(&mut self, node: &str) -> PyResult<(Vec<(String, String, String)>, bool)> { + let names: Vec = self + .graph + .nodes + .get(node) + .ok_or_else(|| PyRuntimeError::new_err(format!("unknown node {node:?}")))? + .outputs + .iter() + .map(|e| e.name.clone()) + .collect(); + let mut outputs = std::collections::BTreeMap::new(); + for name in names { + let t = self.fresh_ref(); + outputs.insert(name, vec![t]); + } + let result = self + .inner + .complete_node(node, outputs) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; + let events = result + .events + .into_iter() + .map(|ev| match ev { + RouteEvent::Emission { name, modality, .. } => ( + "emission".to_string(), + name, + modality.unwrap_or_default(), + ), + RouteEvent::Persist { name, .. } => ("persist".to_string(), name, String::new()), + RouteEvent::Stream { + name, + target_partition, + .. + } => ("stream".to_string(), name, target_partition), + }) + .collect(); + Ok((events, result.walk_done)) + } + + /// Pure-mode seed: caller-supplied uuids so the Python side can key its + /// own uuid -> tensor-descriptor store (the mstar value-map pattern). + fn seed_with(&mut self, inputs: Vec<(String, String, u64)>) -> PyResult<()> { + let seeded = inputs + .into_iter() + .map(|(node, name, uuid)| IncomingInput { + node, + name, + tensors: vec![TensorRef::new(uuid, vec![], "opaque")], + }) + .collect(); + self.inner + .seed(seeded) + .map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + /// Pure-mode complete: caller-supplied uuids per output name; route + /// events return (kind, name, target, uuids) so loop outputs can be + /// reconstructed with real tensor descriptors at termination. + fn complete_with( + &mut self, + node: &str, + outputs: Vec<(String, Vec)>, + ) -> PyResult<(Vec<(String, String, String, Vec)>, bool)> { + let mut map = std::collections::BTreeMap::new(); + for (name, uuids) in outputs { + let refs = uuids + .into_iter() + .map(|u| TensorRef::new(u, vec![], "opaque")) + .collect(); + map.insert(name, refs); + } + let result = self + .inner + .complete_node(node, map) + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; + let uu = |ts: Vec| ts.into_iter().map(|t| t.uuid).collect::>(); + let events = result + .events + .into_iter() + .map(|ev| match ev { + RouteEvent::Emission { name, modality, tensors } => ( + "emission".to_string(), name, + modality.unwrap_or_default(), uu(tensors)), + RouteEvent::Persist { name, tensors } => ( + "persist".to_string(), name, String::new(), uu(tensors)), + RouteEvent::Stream { name, target_partition, tensors } => ( + "stream".to_string(), name, target_partition, uu(tensors)), + }) + .collect(); + Ok((events, result.walk_done)) + } + + /// One-crossing completion for the pure adapter: events + walk-done + + /// the post-completion ready set + loop states, replacing four separate + /// calls per step. + #[allow(clippy::type_complexity)] + fn complete_full( + &mut self, + node: &str, + outputs: Vec<(String, Vec)>, + ) -> PyResult<( + Vec<(String, String, String, Vec)>, + bool, + Vec, + Vec<(String, u32, bool)>, + )> { + let (events, done) = self.complete_with(node, outputs)?; + Ok((events, done, self.inner.ready_nodes(), self.inner.loop_states())) + } + + /// External loop-termination signal (mstar's stop_loops / EOS). + fn signal_loop_finish(&mut self, loop_name: &str) -> PyResult<()> { + self.inner + .signal_loop_finish(loop_name) + .map_err(|e| PyRuntimeError::new_err(e.to_string())) + } + + fn loop_iters(&self) -> Vec<(String, u32)> { + self.inner.loop_iters() + } + + /// (name, curr_iter, terminated) per loop. + fn loop_states(&self) -> Vec<(String, u32, bool)> { + self.inner.loop_states() + } + + fn is_done(&self) -> bool { + self.inner.is_done() + } +} + +/// The worker's batch-decision engine (`mstar/worker/micro_scheduler.py` in +/// Rust). Decide-vs-mutate seam: takes a snapshot of ready work as tuples +/// `(node, walk, rid, worker_graph_id, engine_ready, priority, leader)` plus +/// an explicit monotonic time; returns the batch to run. Queue pops and +/// execution stay with the caller. +#[pyclass(name = "MicroScheduler")] +struct PyMicroScheduler { + inner: MicroScheduler, +} + +type PyReady = (String, String, String, String, bool, u32, bool); + +fn to_entries(ready: Vec) -> Vec { + ready + .into_iter() + .map(|(node, walk, request_id, worker_graph_id, engine_ready, priority, leader)| { + ReadyEntry { + node, + walk, + request_id, + worker_graph_id, + engine_ready, + priority, + leader, + } + }) + .collect() +} + +#[pymethods] +impl PyMicroScheduler { + #[new] + #[pyo3(signature = (sched_type = "round_robin", max_consec_tp_follower_batches = 1))] + fn new(sched_type: &str, max_consec_tp_follower_batches: u32) -> PyResult { + let st = match sched_type { + "priority" => SchedulingType::Priority, + "round_robin" => SchedulingType::RoundRobin, + other => { + return Err(PyValueError::new_err(format!( + "sched_type must be 'priority' or 'round_robin', got {other:?}" + ))) + } + }; + Ok(Self { + inner: MicroScheduler::new(st, max_consec_tp_follower_batches), + }) + } + + fn hold_requests(&mut self, request_ids: Vec, now_ms: u64) { + self.inner.hold_requests(&request_ids, now_ms); + } + + fn add_pending_remove(&mut self, request_id: &str) { + self.inner.add_pending_remove(request_id); + } + + fn clear_pending_remove(&mut self, request_id: &str) { + self.inner.clear_pending_remove(request_id); + } + + fn register_tp_follow(&mut self, node: String, walk: String, request_ids: Vec) { + self.inner.register_tp_follow(node, walk, request_ids); + } + + #[pyo3(signature = (ready, exclude, now_ms))] + fn has_ready_excluding( + &self, + ready: Vec, + exclude: Option<(String, String)>, + now_ms: u64, + ) -> bool { + let entries = to_entries(ready); + self.inner.has_ready_excluding( + &entries, + exclude.as_ref().map(|(n, w)| (n.as_str(), w.as_str())), + now_ms, + ) + } + + /// -> None or (node, walk, request_ids, worker_graph_ids, tp_follow). + #[pyo3(signature = (ready, now_ms, max_batch_size = None, target_node = None, + target_walk = None, exclude_target = None))] + #[allow(clippy::too_many_arguments)] + fn get_next_batch( + &mut self, + ready: Vec, + now_ms: u64, + max_batch_size: Option, + target_node: Option, + target_walk: Option, + exclude_target: Option<(String, String)>, + ) -> Option<(String, String, Vec, Vec, bool)> { + let entries = to_entries(ready); + let filter = BatchFilter { + max_batch_size, + target_node, + target_walk, + exclude_target, + }; + self.inner + .get_next_batch(&entries, &filter, now_ms) + .map(|b| (b.node, b.walk, b.request_ids, b.worker_graph_ids, b.tp_follow)) + } +} + #[pymodule] fn mstar_rust(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add("__version__", env!("CARGO_PKG_VERSION"))?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/rust/src/shm.rs b/rust/src/shm.rs new file mode 100644 index 000000000..a13a8392e --- /dev/null +++ b/rust/src/shm.rs @@ -0,0 +1,478 @@ +//! Shared-memory tensor arena — the data plane for cross-process transport. +//! +//! One persistent `/dev/shm` mmap per producer entity plus an in-process +//! first-fit free-list allocator, replacing the per-tensor file open/write/ +//! read/unlink dance in `SharedMemoryCommunicationManager` (the Python +//! buffer-protocol view lives in `lib.rs`). +//! +//! - Producer: [`ShmArena::create`], [`ShmArena::reserve`] -> offset, copy +//! D2H bytes into `arena[off..off+n]`, send the offset as a descriptor, +//! [`ShmArena::free`] the offset once every consumer has ACKed the +//! transfer (the embedder tracks which offsets belong to which tensors). +//! - Consumer: [`ShmArena::open`] the same name, read `arena[off..off+n]` +//! (H2D). Zero syscalls per tensor, one memcpy each way. + +use std::collections::HashMap; +use std::fs::OpenOptions; +use std::path::Path; +use std::sync::Mutex; + +use memmap2::{MmapMut, MmapOptions}; +use thiserror::Error; + +/// Cache-line-friendly alignment; keeps neighbouring tensors off shared lines. +pub const ALIGN: usize = 256; + +#[inline] +fn align_up(n: usize) -> usize { + n.div_ceil(ALIGN) * ALIGN +} + +#[derive(Debug, Error)] +pub enum ShmError { + #[error("size must be > 0")] + ZeroSize, + #[error("arena full: need {need} B, {free} B free of {total} B")] + Full { + need: usize, + free: usize, + total: usize, + }, + #[error("io on {path}: {source}")] + Io { + path: String, + source: std::io::Error, + }, +} + +/// First-fit free-list allocator over the arena. Coalesces on free. +struct Allocator { + free: Vec<(usize, usize)>, // (offset, len), sorted by offset, disjoint + live: HashMap, // offset -> aligned len (free() needs only offset) +} + +impl Allocator { + fn new(size: usize) -> Self { + Self { + free: vec![(0, size)], + live: HashMap::new(), + } + } + + fn alloc(&mut self, n: usize) -> Option { + let n = align_up(n.max(1)); + for i in 0..self.free.len() { + let (off, len) = self.free[i]; + if len >= n { + if len == n { + self.free.remove(i); + } else { + self.free[i] = (off + n, len - n); + } + self.live.insert(off, n); + return Some(off); + } + } + None + } + + fn dealloc(&mut self, off: usize) -> bool { + let Some(n) = self.live.remove(&off) else { + return false; // double-free / unknown -> no-op + }; + self.free.push((off, n)); + self.free.sort_by_key(|b| b.0); + let mut merged: Vec<(usize, usize)> = Vec::with_capacity(self.free.len()); + for &(o, l) in &self.free { + if let Some(last) = merged.last_mut() { + if last.0 + last.1 == o { + last.1 += l; + continue; + } + } + merged.push((o, l)); + } + self.free = merged; + true + } + + fn bytes_free(&self) -> usize { + self.free.iter().map(|b| b.1).sum() + } + + /// The fragmentation gauge: the biggest single allocation that can + /// currently succeed. Collapsing toward zero while `bytes_free` stays + /// high is the telltale of fragmentation. + fn largest_free_block(&self) -> usize { + self.free.iter().map(|&(_, l)| l).max().unwrap_or(0) + } +} + +/// A named shared-memory arena. The producer owns it (unlinks on drop); a +/// consumer opens the same name read/write without ownership. +pub struct ShmArena { + mmap: MmapMut, + len: usize, + path: String, + owner: bool, + alloc: Mutex, +} + +// The mmap intentionally aliases across processes; in-process access is +// serialised by the allocator Mutex (and, from Python, the GIL). +unsafe impl Send for ShmArena {} +unsafe impl Sync for ShmArena {} + +fn shm_path(name: &str) -> String { + if Path::new("/dev/shm").is_dir() { + format!("/dev/shm/{name}") + } else { + format!("/tmp/{name}") + } +} + +impl ShmArena { + /// Producer: create (or replace) the arena and own it (unlink on drop). + pub fn create(name: &str, size: usize) -> Result { + if size == 0 { + return Err(ShmError::ZeroSize); + } + let path = shm_path(name); + let io = |source| ShmError::Io { + path: path.clone(), + source, + }; + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open(&path) + .map_err(io)?; + file.set_len(size as u64).map_err(io)?; + let mmap = unsafe { MmapOptions::new().len(size).map_mut(&file) }.map_err(io)?; + Ok(Self { + mmap, + len: size, + path, + owner: true, + alloc: Mutex::new(Allocator::new(size)), + }) + } + + /// Consumer: open an existing arena read/write; never unlinks. + pub fn open(name: &str) -> Result { + let path = shm_path(name); + let io = |source| ShmError::Io { + path: path.clone(), + source, + }; + let file = OpenOptions::new() + .read(true) + .write(true) + .open(&path) + .map_err(io)?; + let size = file.metadata().map_err(io)?.len() as usize; + let mmap = unsafe { MmapOptions::new().len(size).map_mut(&file) }.map_err(io)?; + Ok(Self { + mmap, + len: size, + path, + owner: false, + alloc: Mutex::new(Allocator::new(size)), // unused on a consumer + }) + } + + /// Reserve `nbytes`; returns the byte offset into the arena. Producer only. + pub fn reserve(&self, nbytes: usize) -> Result { + let mut a = self.alloc.lock().expect("alloc lock"); + a.alloc(nbytes).ok_or_else(|| ShmError::Full { + need: nbytes, + free: a.bytes_free(), + total: self.len, + }) + } + + /// Release a reserved offset. Producer only. Idempotent (false = unknown). + pub fn free(&self, offset: usize) -> bool { + self.alloc.lock().expect("alloc lock").dealloc(offset) + } + + pub fn size(&self) -> usize { + self.len + } + + pub fn bytes_free(&self) -> usize { + self.alloc.lock().expect("alloc lock").bytes_free() + } + + /// Largest single contiguous free block (see `Allocator::largest_free_block`). + pub fn largest_free_block(&self) -> usize { + self.alloc.lock().expect("alloc lock").largest_free_block() + } + + /// Base pointer of the arena — used by the `mstar-py` buffer-protocol view + /// so torch can `frombuffer(arena[off:off+n])` with no intermediate copy. + pub fn as_mut_ptr(&self) -> *mut u8 { + self.mmap.as_ptr() as *mut u8 + } + + /// Byte slice into the arena (bounds-checked; for tests / Rust-side copies). + pub fn bytes(&self, offset: usize, len: usize) -> &[u8] { + &self.mmap[offset..offset + len] + } + + /// Copy `src` into the arena at `offset` (producer D2H staging). + pub fn write_at(&mut self, offset: usize, src: &[u8]) { + self.mmap[offset..offset + src.len()].copy_from_slice(src); + } + + pub fn close(&mut self) { + if self.owner && !self.path.is_empty() { + let _ = std::fs::remove_file(&self.path); + self.path.clear(); + } + } +} + +impl Drop for ShmArena { + fn drop(&mut self) { + self.close(); + } +} + +// --------------------------------------------------------------------------- +// Segmented arena: grow-by-segments + uuid-grouped reclaim +// --------------------------------------------------------------------------- + +/// A producer arena that grows by adding fixed-size segments instead of +/// erroring (or resizing) when full. +/// +/// Why segments, not `mremap`: each segment's mapping is created once and +/// never moves, so a CUDA host-registration (`cudaHostRegister`) of a segment +/// stays valid for its lifetime — the property that keeps D2H/H2D copies on +/// side streams truly asynchronous. A dynamically-resized arena would +/// invalidate the registration on every growth. Registration itself is the +/// embedder's job (torch/cuda): watch `num_segments()` and register each new +/// segment's `(ptr, len)` once. +/// +/// Descriptors are `(segment_name, offset)`: every segment is an ordinary +/// named [`ShmArena`] (`{base}.seg{i}`), so consumers keep opening arenas +/// lazily by name — a consumer never needs to know segmentation exists. +/// +/// Reclaim is offset-based: the embedder tracks each allocation's +/// `(segment, offset)` and calls [`Self::free`] when the consumer has +/// ACKed the transfer. +pub struct SegmentedShmArena { + base: String, + segment_size: usize, + max_segments: usize, + // Behind a Mutex so every method takes `&self` (interior mutability): + // the Python binding releases the GIL across reserve's growth path, and + // a `&mut` PyO3 borrow held there would make any concurrent `&self` + // call (stats, free) raise "Already mutably borrowed". + // Growth serializes on this lock instead of the PyCell borrow. + segments: Mutex>>, +} + +impl SegmentedShmArena { + /// Create with one initial segment. `max_segments` caps total growth + /// (the arena-full backpressure boundary): `segment_size * max_segments` + /// bytes, after which `reserve` reports `Full`. + pub fn create( + base: &str, + segment_size: usize, + max_segments: usize, + ) -> Result { + if segment_size == 0 || max_segments == 0 { + return Err(ShmError::ZeroSize); + } + let first = ShmArena::create(&format!("{base}.seg0"), segment_size)?; + Ok(Self { + base: base.to_string(), + segment_size, + max_segments, + segments: Mutex::new(vec![std::sync::Arc::new(first)]), + }) + } + + /// Reserve `nbytes`; returns `(segment_index, offset)`. Tries existing + /// segments first-fit, then grows by one segment (a dedicated one when + /// `nbytes` exceeds the segment size), up to `max_segments`. + pub fn reserve(&self, nbytes: usize) -> Result<(usize, usize), ShmError> { + let mut segments = self.segments.lock().expect("segments lock"); + for (i, seg) in segments.iter().enumerate() { + if let Ok(off) = seg.reserve(nbytes) { + return Ok((i, off)); + } + } + if segments.len() >= self.max_segments { + return Err(ShmError::Full { + need: nbytes, + free: segments.iter().map(|s| s.bytes_free()).sum(), + total: segments.iter().map(|s| s.size()).sum(), + }); + } + let idx = segments.len(); + let size = self.segment_size.max(align_up(nbytes.max(1))); + let seg = ShmArena::create(&format!("{}.seg{idx}", self.base), size)?; + let off = seg.reserve(nbytes)?; // fresh segment sized to fit: infallible + segments.push(std::sync::Arc::new(seg)); + Ok((idx, off)) + } + + /// Release one offset (embedder-tracked descriptors). Idempotent. + pub fn free(&self, segment: usize, offset: usize) -> bool { + self.segments + .lock() + .expect("segments lock") + .get(segment) + .is_some_and(|s| s.free(offset)) + } + + pub fn num_segments(&self) -> usize { + self.segments.lock().expect("segments lock").len() + } + + /// Occupancy + fragmentation snapshot across all segments: + /// `(total_bytes, free_bytes, largest_free_block)`. The fragmentation + /// signature is `largest_free_block` collapsing while `free_bytes` stays + /// high — allocations then fail (or force segment growth) even though + /// total free space looks healthy. + pub fn stats(&self) -> (usize, usize, usize) { + let segments = self.segments.lock().expect("segments lock"); + let total = segments.iter().map(|s| s.size()).sum(); + let free = segments.iter().map(|s| s.bytes_free()).sum(); + let largest = segments + .iter() + .map(|s| s.largest_free_block()) + .max() + .unwrap_or(0); + (total, free, largest) + } + + /// The `{base}.seg{i}` arena name — what descriptors carry, and what a + /// consumer passes to [`ShmArena::open`]. + pub fn segment_name(&self, i: usize) -> String { + format!("{}.seg{i}", self.base) + } + + /// A shared handle to segment `i` (buffer views, registration hooks). + pub fn segment(&self, i: usize) -> Option> { + self.segments.lock().expect("segments lock").get(i).cloned() + } + + /// `(base_ptr, len)` of segment `i` — the registration hook: the embedder + /// `cudaHostRegister`s each new segment ONCE; the mapping never moves. + pub fn segment_ptr_len(&self, i: usize) -> Option<(*mut u8, usize)> { + self.segments + .lock() + .expect("segments lock") + .get(i) + .map(|s| (s.as_mut_ptr(), s.size())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn name(tag: &str) -> String { + format!("mstar_shm_test_{tag}_{:?}", std::thread::current().id()) + } + + #[test] + fn alloc_free_coalesce() { + let mut a = Allocator::new(1024); + let x = a.alloc(200).unwrap(); // -> 256 aligned + let y = a.alloc(200).unwrap(); + assert_eq!(x, 0); + assert_eq!(y, 256); + assert_eq!(a.bytes_free(), 1024 - 512); + a.dealloc(x); + a.dealloc(y); + // Fully coalesced back to one free block. + assert_eq!(a.free, vec![(0, 1024)]); + assert!(!a.dealloc(9999)); // unknown offset + } + + #[test] + fn producer_consumer_roundtrip_same_process() { + let n = name("rt"); + let mut prod = ShmArena::create(&n, 4096).unwrap(); + let off = prod.reserve(1500).unwrap(); + let payload: Vec = (0..1500).map(|i| (i % 251) as u8).collect(); + prod.write_at(off, &payload); + + // Consumer opens the same name and reads at the descriptor offset. + let cons = ShmArena::open(&n).unwrap(); + assert_eq!(cons.bytes(off, 1500), &payload[..]); + + prod.free(off); + assert_eq!(prod.bytes_free(), 4096); + } + + #[test] + fn reserve_reports_full() { + let n = name("full"); + let arena = ShmArena::create(&n, 512).unwrap(); + arena.reserve(300).unwrap(); // -> 512 aligned, arena full + assert!(matches!(arena.reserve(1), Err(ShmError::Full { .. }))); + } + + #[test] + fn zero_size_rejected() { + assert!(matches!( + ShmArena::create(&name("z"), 0), + Err(ShmError::ZeroSize) + )); + } + + // ---- segmented arena (grow-by-segments + uuid reclaim) ---------------- + + #[test] + fn segments_grow_under_pressure_and_names_are_openable() { + let base = name("seg"); + let mut a = SegmentedShmArena::create(&base, 1024, 4).unwrap(); + assert_eq!(a.num_segments(), 1); + + // Fill segment 0, forcing growth into segment 1. + let (s0, _) = a.reserve(900).unwrap(); + let (s1, o1) = a.reserve(900).unwrap(); + assert_eq!((s0, s1), (0, 1)); + assert_eq!(a.num_segments(), 2); + + // A late segment is an ordinary named arena a consumer can open — + // exactly what a (segment_name, offset) descriptor promises. + let cons = ShmArena::open(&a.segment_name(1)).unwrap(); + assert_eq!(cons.size(), 1024); + let _ = (cons.bytes(o1, 4), ()); + + // The registration hook surface: stable (ptr, len) per segment. + let (p0, l0) = a.segment_ptr_len(0).unwrap(); + assert!(!p0.is_null()); + assert_eq!(l0, 1024); + } + + #[test] + fn oversized_allocation_gets_dedicated_segment() { + let base = name("big"); + let mut a = SegmentedShmArena::create(&base, 1024, 4).unwrap(); + let (seg, off) = a.reserve(5000).unwrap(); // > segment_size + assert_eq!((seg, off), (1, 0)); + assert!(a.segment(1).unwrap().size() >= 5000); + } + + #[test] + fn growth_cap_reports_full() { + let base = name("cap"); + let mut a = SegmentedShmArena::create(&base, 512, 2).unwrap(); + a.reserve(500).unwrap(); // fills seg 0 + a.reserve(500).unwrap(); // grows + fills seg 1 + assert!(matches!(a.reserve(500), Err(ShmError::Full { .. }))); + // Freeing makes room again (backpressure boundary, not a dead end). + assert!(a.free(0, 0)); + assert!(matches!(a.reserve(500), Ok((0, 0)))); + } + +} diff --git a/test/rust/bench_walk_ab.py b/test/rust/bench_walk_ab.py new file mode 100644 index 000000000..bd17d775d --- /dev/null +++ b/test/rust/bench_walk_ab.py @@ -0,0 +1,71 @@ +"""Walk-layer A/B: Python WorkerGraphIO vs PureRustWorkerGraphIO vs the +minimal uuid protocol (the target boundary). Interleaved per request so +background load hits all columns equally. Informational — run by CI after +the test suite; prints a table, never fails.""" +import sys +import time +from copy import deepcopy +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) + +from test_walk_parity import _loop_graph # noqa: E402 + +from mstar.graph.base import GraphEdge # noqa: E402 +from mstar.graph.graph_io import WorkerGraphIO # noqa: E402 +from mstar.graph.rust_core import ( # noqa: E402 + PureRustWorkerGraphIO, + walks_to_json, +) +from mstar_rust import WalkSet # noqa: E402 + +ITERS, REQS = 64, 60 +SECTION = _loop_graph(ITERS) +STEPS = ITERS + 1 +WS = WalkSet.from_json(walks_to_json({"walk": SECTION})) + + +def drive_io(io): + io.ingest_input(GraphEdge(next_node="L", name="seed")) + io.ingest_input(GraphEdge(next_node="L", name="fb")) + while not io.wg_state_registry.is_done: + name = sorted(io.ready_node_names)[0] + io.ready_node_names.discard(name) + comp = io.mark_node_complete(name) + for e in comp.output_edges: + if e.next_node in io.nodes: + io.ingest_input(e) + + +def drive_minimal(): + st = WS.state("walk") + st.seed_with([("L", "seed", 1), ("L", "fb", 2)]) + u = 3 + while not st.is_done(): + n = st.ready_nodes()[0] + st.schedule(n) + names = ("fb", "tok", "final") if n == "L" else ("out",) + st.complete_full(n, [(e, [u + i]) for i, e in enumerate(names)]) + u += 3 + + +RUNS = { + "python registries": lambda: drive_io(WorkerGraphIO(deepcopy(SECTION), wg_id="w")), + "pure rust adapter": lambda: drive_io(PureRustWorkerGraphIO(SECTION, "w")), + "minimal uuid protocol": drive_minimal, +} +for fn in RUNS.values(): + fn() # warm + +totals = dict.fromkeys(RUNS, 0.0) +for _ in range(REQS): # interleaved + for name, fn in RUNS.items(): + t0 = time.perf_counter() + fn() + totals[name] += time.perf_counter() - t0 + +print(f"\nwalk A/B ({REQS} reqs x {STEPS} steps, interleaved)") +base = totals["python registries"] +for name, tot in totals.items(): + per = tot / REQS / STEPS * 1e6 + print(f" {name:24} {per:8.2f} us/step ({base / tot:4.2f}x vs python)") diff --git a/test/rust/test_arena_transport.py b/test/rust/test_arena_transport.py new file mode 100644 index 000000000..780f7078c --- /dev/null +++ b/test/rust/test_arena_transport.py @@ -0,0 +1,364 @@ +"""ArenaShmCommunicationManager: producer stages tensors +into the Rust shared-memory arena, the location rides the TensorPointerInfo, +a separate consumer manager reads them zero-copy, and reclaim frees the +arena slots. Skipped unless the ``mstar_rust`` extension is installed.""" +import os + +import pytest +import torch + +pytest.importorskip("mstar_rust") + +from mstar.communication.arena import ArenaShmCommunicationManager +from mstar.communication.communicator import CommProtocol +from mstar.communication.tensors import ( + SharedMemoryCommunicationManager, + create_tensor_communication_manager, +) +from mstar.graph.base import GraphEdge + + +class _NullCommunicator: + """The store/register/read path under test never touches the mesh.""" + + def send(self, entity_id, msg): + raise AssertionError("unexpected control-mesh send") + + def get_all_new_messages(self): + return [] + + +def _manager(entity, tmp_path): + os.environ["MSTAR_SHM_ARENA_SEGMENT_MB"] = "1" + os.environ["MSTAR_SHM_ARENA_MAX_SEGMENTS"] = "4" + return ArenaShmCommunicationManager( + my_entity_id=entity, hostname="localhost", device="cpu", + communicator=_NullCommunicator(), shm_dir=str(tmp_path), + ) + + +def test_producer_to_consumer_roundtrip(tmp_path): + prod = _manager("w0", tmp_path) + cons = _manager("w1", tmp_path) + + tensors = { + "hidden": [torch.randn(4, 8), torch.arange(32, dtype=torch.int64)], + "empty": [torch.empty(0, 3)], + } + infos = prod.store_and_return_tensor_info("r1", tensors) + prod.register_for_send("r1", [i for il in infos.values() for i in il]) + + # The location was stamped onto the shipped descriptors. + for il in infos.values(): + for info in il: + assert info.shm_segment is not None and info.shm_segment.startswith( + "mstar_arena_w0") + + edges = [ + GraphEdge(next_node="B", name=name, tensor_info=il) + for name, il in infos.items() + ] + cons.start_read_tensors("r1", edges) + for name, originals in tensors.items(): + for original, info in zip(originals, infos[name], strict=True): + got = cons.tensor_store.get_tensor("r1", info.uuid) + assert torch.equal(got, original), name + + +def test_reclaim_frees_arena_slots(tmp_path): + prod = _manager("w2", tmp_path) + infos = prod.store_and_return_tensor_info( + "r2", {"x": [torch.randn(16)]}) + (info,) = infos["x"] + prod.register_for_send("r2", [info]) + assert prod._arena_locs + prod._cleanup_by_uuid("r2", info.uuid) + assert not prod._arena_locs + + +def test_arena_grows_then_spills(tmp_path): + """Past the segment cap the producer SPILLS to the per-uuid file + protocol (slower, never fails — the old manager's saturation behavior); + the consumer reads spilled tensors through the file fallback; reclaim + unlinks the files. Stats expose the fragmentation gauge.""" + os.environ["MSTAR_SHM_ARENA_SPILL_AFTER_S"] = "0.05" + try: + prod = _manager("w3", tmp_path) + # 1 MiB segments, cap 4: filling ~3.5 MiB grows the arena... + infos = prod.store_and_return_tensor_info( + "r3", {"big": [torch.zeros(300_000, dtype=torch.uint8) + for _ in range(12)]}) + prod.register_for_send("r3", list(infos["big"])) + assert prod._arena.num_segments > 1 + st = prod.stats_summary() + assert st["segments"] == prod._arena.num_segments + assert 0 < st["largest_free_block"] <= st["free_bytes"] + # ...and past the cap, further tensors spill to files instead of + # failing: shm_segment stays None and a per-uuid file appears. + vals = [torch.full((300_000,), i, dtype=torch.uint8) + for i in range(6)] + more = prod.store_and_return_tensor_info("r3", {"more": vals}) + prod.register_for_send("r3", list(more["more"])) + spilled = [i for i in more["more"] if i.shm_segment is None] + assert spilled, "expected at least one spill past the cap" + assert all(u in prod._shm_files for u in + (i.uuid for i in spilled)) + # The consumer round-trips spilled tensors via the file fallback. + cons = _manager("w4", tmp_path) + edge = GraphEdge(next_node="B", name="more", + tensor_info=more["more"]) + cons.start_read_tensors("r3", [edge]) + for val, info in zip(vals, more["more"], strict=True): + got = cons.tensor_store.get_tensor("r3", info.uuid) + assert torch.equal(got, val) + # Reclaim unlinks the spilled files. + for info in spilled: + path = prod._shm_files[info.uuid] + prod._cleanup_by_uuid("r3", info.uuid) + assert not os.path.exists(path) + finally: + del os.environ["MSTAR_SHM_ARENA_SPILL_AFTER_S"] + + +def test_mixed_edge_and_fragmentation_signature(tmp_path, caplog): + """One edge can mix arena-staged and spilled tensors (the consumer + dispatches per descriptor), and a reserve that fails while TOTAL free + space covers it logs the fragmentation signature (largest free block + collapsed) before spilling.""" + import logging + + os.environ["MSTAR_SHM_ARENA_SPILL_AFTER_S"] = "0.05" + try: + prod = _manager("w6", tmp_path) + # Fill to the 4-segment cap with 12 x 300 KB... + fill = prod.store_and_return_tensor_info( + "r6", {"fill": [torch.zeros(300_000, dtype=torch.uint8) + for _ in range(12)]}) + prod.register_for_send("r6", list(fill["fill"])) + # ...then free ALTERNATE allocations: ~1.2 MB total free, but no + # contiguous block larger than ~300 KB. + for info in fill["fill"][::2]: + prod._cleanup_by_uuid("r6", info.uuid) + st = prod.stats_summary() + assert st["free_bytes"] > 500_000 > st["largest_free_block"] + + # A small tensor fits a hole (arena); a 500 KB one has the total + # free space but no block -> fragmentation warning, then spill. + small = torch.arange(1000, dtype=torch.uint8) + big = torch.full((500_000,), 7, dtype=torch.uint8) + mixed = prod.store_and_return_tensor_info( + "r6", {"mixed": [small, big]}) + with caplog.at_level(logging.WARNING, + logger="mstar.communication.arena"): + prod.register_for_send("r6", list(mixed["mixed"])) + s_info, b_info = mixed["mixed"] + assert s_info.shm_segment is not None # staged in a hole + assert b_info.shm_segment is None # spilled + assert any("fragmentation" in r.message for r in caplog.records) + + # The consumer reads the MIXED edge: one from the arena, one from + # the spill file, in a single start_read_tensors call. + cons = _manager("w7", tmp_path) + edge = GraphEdge(next_node="B", name="mixed", + tensor_info=mixed["mixed"]) + cons.start_read_tensors("r6", [edge]) + assert torch.equal( + cons.tensor_store.get_tensor("r6", s_info.uuid), small) + assert torch.equal( + cons.tensor_store.get_tensor("r6", b_info.uuid), big) + finally: + del os.environ["MSTAR_SHM_ARENA_SPILL_AFTER_S"] + + +def test_strict_mode_backpressures_then_fails(tmp_path): + """MSTAR_SHM_ARENA_SPILL=0 restores the strict contract: backpressure + at the cap, then a loud arena-full error.""" + os.environ["MSTAR_SHM_ARENA_SPILL"] = "0" + os.environ["MSTAR_SHM_ARENA_FULL_TIMEOUT_S"] = "0.2" + try: + prod = _manager("w5", tmp_path) + infos = prod.store_and_return_tensor_info( + "r5", {"big": [torch.zeros(300_000, dtype=torch.uint8) + for _ in range(12)]}) + prod.register_for_send("r5", list(infos["big"])) + more = prod.store_and_return_tensor_info( + "r5", {"more": [torch.zeros(300_000, dtype=torch.uint8) + for _ in range(6)]}) + with pytest.raises(RuntimeError, match="arena full"): + prod.register_for_send("r5", list(more["more"])) + finally: + del os.environ["MSTAR_SHM_ARENA_SPILL"] + del os.environ["MSTAR_SHM_ARENA_FULL_TIMEOUT_S"] + + +def test_transport_mismatch_fails_loudly(tmp_path): + """A mixed deployment (arena producer + file consumer, or the reverse) + fails with an explicit MSTAR_SHM_ARENA message where data would be + unreachable (arena producer -> file consumer); the reverse direction + interops via the arena consumer's file fallback.""" + arena_prod = _manager("mx0", tmp_path) + file_cons = SharedMemoryCommunicationManager( + my_entity_id="mx1", hostname="localhost", device="cpu", + communicator=_NullCommunicator(), shm_dir=str(tmp_path)) + infos = arena_prod.store_and_return_tensor_info( + "rm", {"x": [torch.randn(4)]}) + arena_prod.register_for_send("rm", [infos["x"][0]]) + edge = GraphEdge(next_node="B", name="x", tensor_info=infos["x"]) + with pytest.raises(RuntimeError, match="MSTAR_SHM_ARENA"): + file_cons.start_read_tensors("rm", [edge]) + + # The reverse direction now INTEROPS: a file-producer's tensors carry + # no arena location, which is exactly the spill wire shape — the arena + # consumer reads them through its file fallback. + file_prod = SharedMemoryCommunicationManager( + my_entity_id="mx2", hostname="localhost", device="cpu", + communicator=_NullCommunicator(), shm_dir=str(tmp_path)) + arena_cons = _manager("mx3", tmp_path) + y = torch.randn(4) + infos = file_prod.store_and_return_tensor_info("rm2", {"y": [y]}) + file_prod.register_for_send("rm2", [infos["y"][0]]) + edge = GraphEdge(next_node="B", name="y", tensor_info=infos["y"]) + arena_cons.start_read_tensors("rm2", [edge]) + got = arena_cons.tensor_store.get_tensor("rm2", infos["y"][0].uuid) + assert torch.equal(got, y) + + +def test_factory_flag(tmp_path, monkeypatch): + def make(value): + monkeypatch.setenv("MSTAR_SHM_ARENA", value) + return create_tensor_communication_manager( + protocol=CommProtocol.SHM, my_entity_id=f"f_{value}", + hostname="localhost", device="cpu", + communicator=_NullCommunicator(), shm_dir=str(tmp_path), + ) + + assert type(make("0")) is SharedMemoryCommunicationManager + assert type(make("1")) is ArenaShmCommunicationManager + assert type(make("AUTO")) is ArenaShmCommunicationManager + with pytest.raises(ValueError): + make("yes") + + +def test_instance_unique_names_no_collision(tmp_path): + """Two servers with the SAME entity id must not share /dev/shm names — + a fixed name would let the second create() truncate the first's live + segments (silent corruption, observed on a shared cluster).""" + a = _manager("dup", tmp_path) + x = torch.arange(64, dtype=torch.uint8) + infos = a.store_and_return_tensor_info("r", {"x": [x]}) + a.register_for_send("r", list(infos["x"])) + b = _manager("dup", tmp_path) # same entity id, second instance + assert b._arena.segment_name(0) != a._arena.segment_name(0) + # a's staged data survives b's creation; the descriptor still resolves. + cons = _manager("dupc", tmp_path) + edge = GraphEdge(next_node="B", name="x", tensor_info=infos["x"]) + cons.start_read_tensors("r", [edge]) + assert torch.equal( + cons.tensor_store.get_tensor("r", infos["x"][0].uuid), x) + + +def test_orphan_sweep(tmp_path): + """A SIGKILLed server's segments (owner pid gone) are reclaimed by the + next construction's sweep; live owners' files are left alone.""" + dead = "/dev/shm/mstar_arena_zombie_999999999_deadbeef.seg0" + keep = f"/dev/shm/mstar_arena_alive_{os.getpid()}_cafebabe.seg0" + with open(dead, "wb") as f: + f.write(b"x" * 64) + with open(keep, "wb") as f: + f.write(b"x") + try: + _manager("sweeper", tmp_path) + assert not os.path.exists(dead), "dead-owner orphan not swept" + assert os.path.exists(keep), "live-owner file wrongly swept" + finally: + for f in (dead, keep): + try: + os.unlink(f) + except FileNotFoundError: + pass + + +def test_dead_peer_segments_evicted(tmp_path): + """A consumer must not accumulate mappings for peer segments whose + backing file is gone (instance-unique names mean every producer restart + mints NEW names — a never-evicting cache leaks a generation of mappings + per restart).""" + prod = _manager("evp", tmp_path) + x = torch.arange(64, dtype=torch.uint8) + infos = prod.store_and_return_tensor_info("re", {"x": [x]}) + prod.register_for_send("re", list(infos["x"])) + cons = _manager("evc", tmp_path) + edge = GraphEdge(next_node="B", name="x", tensor_info=infos["x"]) + cons.start_read_tensors("re", [edge]) + seg = infos["x"][0].shm_segment + assert seg in cons._peer_segments + cons.pending.clear() # no in-flight reads + # Producer goes away gracefully: Drop unlinks its segments. + prod._cleanup_by_uuid("re", infos["x"][0].uuid) + del prod + import gc + + gc.collect() + assert not os.path.exists(f"/dev/shm/{seg}") + cons._peer_evict_last = 0.0 # bypass the time gate + cons.start_read_tensors("re", []) # triggers the eviction sweep + assert seg not in cons._peer_segments + + +def test_ttl_backstop_reclaims_abort_orphans(tmp_path): + """A slot staged but never ACKed (abort) is force-freed once older + than MSTAR_SHM_ARENA_SLOT_TTL_S, letting a full arena recover instead + of spilling forever. Off by default.""" + os.environ["MSTAR_SHM_ARENA_SLOT_TTL_S"] = "0.05" + try: + prod = _manager("wt", tmp_path) + infos = prod.store_and_return_tensor_info( + "rt", {"x": [torch.zeros(300_000, dtype=torch.uint8)]}) + prod.register_for_send("rt", list(infos["x"])) + assert prod._arena_locs + import time as _t + + _t.sleep(0.06) + assert prod._reclaim_expired() >= 1 + assert not prod._arena_locs + finally: + del os.environ["MSTAR_SHM_ARENA_SLOT_TTL_S"] + + +def test_segments_unlinked_at_interpreter_exit(tmp_path): + """A worker exits with its manager still referenced (no explicit + cleanup path), so the Rust Drop never runs — the exit finalizer must + unlink the segments anyway.""" + import subprocess + import sys + from pathlib import Path + + repo_root = str(Path(__file__).resolve().parents[2]) + code = f""" +import sys +sys.path.insert(0, {repo_root!r}) +import os +os.environ["MSTAR_SHM_ARENA_SEGMENT_MB"] = "1" +os.environ["MSTAR_SHM_ARENA_MAX_SEGMENTS"] = "2" +import torch +from mstar.communication.arena import ArenaShmCommunicationManager + +class _C: + def send(self, *a): pass + def get_all_new_messages(self): return [] + +m = ArenaShmCommunicationManager( + my_entity_id="exitcase", hostname="localhost", device="cpu", + communicator=_C(), shm_dir={repr(str(tmp_path))}) +infos = m.store_and_return_tensor_info( + "r", {{"x": [torch.arange(64, dtype=torch.uint8)]}}) +m.register_for_send("r", [infos["x"][0]]) +print(m._own_segment_paths[0]) +KEEP_ALIVE = m # global reference survives to interpreter exit +""" + out = subprocess.run([sys.executable, "-c", code], check=False, + capture_output=True, text=True, timeout=120) + assert out.returncode == 0, out.stderr[-500:] + seg_path = out.stdout.strip().splitlines()[-1] + assert seg_path.startswith("/dev/shm/mstar_arena_exitcase_") + assert not os.path.exists(seg_path), "segment survived interpreter exit" diff --git a/test/rust/test_graph_modes.py b/test/rust/test_graph_modes.py new file mode 100644 index 000000000..6ce933b22 --- /dev/null +++ b/test/rust/test_graph_modes.py @@ -0,0 +1,56 @@ +"""The existing graph test suite (test/modular/test_graph.py), re-run with +the Rust walk core engaged — the validation the graph layer's adoption is +gated on. Each scenario (pipeline, fixed-iteration loop, dynamic finish, +nested loops, EOS ready-signal clearing) runs under every ``MSTAR_RUST_WALK`` +mode by wrapping the ``WorkerGraphIO`` the tests construct: + +* ``shadow`` — Python stays authoritative; the test additionally FAILS on + any logged divergence, so a silent Rust-side mismatch cannot pass. +* ``1`` — Rust decisions (ready set, doneness, loop indices) drive the + scenarios; the suite's own assertions validate the behavior. + +Skipped unless the ``mstar_rust`` extension is installed.""" + +import importlib.util +import logging +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("mstar_rust") + +from mstar.graph.graph_io import WorkerGraphIO +from mstar.graph.rust_core import wrap_worker_graph_io + +_SPEC = importlib.util.spec_from_file_location( + "modular_test_graph", + Path(__file__).resolve().parents[1] / "modular" / "test_graph.py") +_tg = importlib.util.module_from_spec(_SPEC) +sys.modules["modular_test_graph"] = _tg +_SPEC.loader.exec_module(_tg) + +SCENARIOS = [ + _tg.test_simple_pipeline, + _tg.test_diffusion_loop_fixed_iters, + _tg.test_ar_generation_with_dynamic_finish, + _tg.test_nested_loops, + _tg.test_eos_clears_ready_signals, +] +MODES = ["shadow", "1", "pure"] + + +@pytest.mark.parametrize("scenario", SCENARIOS, ids=lambda f: f.__name__) +@pytest.mark.parametrize("mode", MODES) +def test_graph_suite_under_rust_walk(mode, scenario, monkeypatch, caplog): + monkeypatch.setenv("MSTAR_RUST_WALK", mode) + + def wrapped(graph): + return wrap_worker_graph_io(WorkerGraphIO(graph), graph, wg_id=0) + + monkeypatch.setattr(_tg, "WorkerGraphIO", wrapped) + with caplog.at_level(logging.ERROR, logger="mstar.graph.rust_core"): + scenario() + diverged = [r.message for r in caplog.records + if "diverg" in r.message.lower() or "fell back" in r.message.lower()] + assert not diverged, diverged diff --git a/test/rust/test_sched_parity.py b/test/rust/test_sched_parity.py new file mode 100644 index 000000000..e12c67477 --- /dev/null +++ b/test/rust/test_sched_parity.py @@ -0,0 +1,211 @@ +"""Scheduler parity: the Rust MicroScheduler vs mstar's, driven with +identical ready-work and event sequences. mstar's scheduler reads live +manager/engine objects and pops queues itself; the Rust one takes a snapshot +and leaves pops to the caller — the harness bridges the two and asserts the +DECISIONS (node, walk, request set, order of batches) match at every call. +Skipped unless the ``mstar_rust`` extension is installed.""" +import time +from types import SimpleNamespace + +import pytest + +pytest.importorskip("mstar_rust") +pytest.importorskip("torch") + +from mstar_rust import MicroScheduler as RustScheduler + +from mstar.engine.base import EngineType +from mstar.utils.ipc_format import ScheduleTPNode +from mstar.worker.micro_scheduler import MicroScheduler, SchedulingType + +PRIORITY = {EngineType.KV_CACHE: 0, EngineType.STATELESS: 2} + + +class _Engine: + def __init__(self, etype): + self._etype = etype + self.not_ready: set = set() # (node, rid) pairs failing check_ready + + def engine_type(self): + return self._etype + + def check_ready(self, node, rid, fwd_info): + return (node, rid) not in self.not_ready + + +class _World: + """One mutable ready-state driving BOTH schedulers.""" + + def __init__(self, node_engines, leaders): + self.engines = {n: _Engine(t) for n, t in node_engines.items()} + self.leaders = set(leaders) + # rid -> (node, walk) currently ready (one ready node per rid here) + self.ready: dict[str, tuple[str, str]] = {} + + # ---- mstar-side stubs ---- + world = self + + class _Queue: + def get_ready_node_names(self): + return {rid: {nw[0]} for rid, nw in world.ready.items()} + + @property + def per_request_queues(self): + return { + rid: SimpleNamespace(ready_node_names={nw[0]}) + for rid, nw in world.ready.items() + } + + def pop_ready_nodes(self, rid, names): + nw = world.ready.get(rid) + if nw and nw[0] in names: + del world.ready[rid] + return [SimpleNamespace(name=nw[0])] + return [] + + queue = _Queue() + self.wgm = SimpleNamespace( + queues={"wg0": queue}, + per_request_info={}, # filled per rid below + get_partition_for_node=lambda n: "p0", + get_graph_walk=lambda rid, part: self.ready[rid][1], + get_fwd_info=lambda rid, part: None, + get_worker_graph_id_for_node=lambda rid, node, graph_walk=None: "wg0", + ) + em = SimpleNamespace( + node_to_engine=self.engines, + get_engine=lambda name: self.engines[name], + ) + self.py = MicroScheduler( + em, sched_type=SchedulingType.ROUND_ROBIN, + parallel_leader_nodes=self.leaders, + max_consec_tp_follower_batches=1, + ) + self.rs = RustScheduler("round_robin", 1) + + def use_priority(self): + self.py.sched_type = SchedulingType.PRIORITY + self.rs = RustScheduler("priority", 1) + return self + + def add(self, rid, node, walk): + self.ready[rid] = (node, walk) + self.wgm.per_request_info[rid] = True + + def _snapshot(self): + return [ + ( + node, walk, rid, "wg0", + self.engines[node].check_ready(node, rid, None), + PRIORITY.get(self.engines[node].engine_type(), 99), + node in self.leaders, + ) + for rid, (node, walk) in self.ready.items() + ] + + def step(self, **kw): + """One get_next_batch on both; assert identical decisions; apply the + Rust decision's pops to the shared state (Python popped its own).""" + now_ms = int(time.monotonic() * 1000) + snapshot = self._snapshot() + py_batch = self.py.get_next_batch( + self.wgm, + max_batch_size=kw.get("max_batch_size"), + target_node_name=kw.get("target_node"), + target_graph_walk=kw.get("target_walk"), + exclude_target=kw.get("exclude_target"), + ) + rs_batch = self.rs.get_next_batch( + snapshot, now_ms, + max_batch_size=kw.get("max_batch_size"), + target_node=kw.get("target_node"), + target_walk=kw.get("target_walk"), + exclude_target=kw.get("exclude_target"), + ) + if py_batch is None: + assert rs_batch is None, f"rust scheduled {rs_batch}, python None" + return None + assert rs_batch is not None, f"python scheduled {py_batch}, rust None" + node, walk, rids, _wgids, _tpf = rs_batch + assert node == py_batch.node_name + assert walk == py_batch.graph_walk + assert sorted(rids) == sorted(py_batch.node_objects.keys()) + # Python's stub pop already removed its rids; mirror for Rust's view. + for rid in rids: + self.ready.pop(rid, None) + return node, walk, sorted(rids) + + def hold(self, rids): + self.py.hold_requests(rids) + self.rs.hold_requests(rids, int(time.monotonic() * 1000)) + + def tp_follow(self, node, walk, rids): + self.py.register_tp_follow( + ScheduleTPNode(node_name=node, graph_walk=walk, request_ids=rids)) + self.rs.register_tp_follow(node, walk, rids) + + +def test_round_robin_rotates(): + w = _World({"A": EngineType.KV_CACHE, "B": EngineType.KV_CACHE}, + leaders={"A", "B"}) + w.add("r1", "A", "w") + w.add("r2", "B", "w") + first = w.step() + w.add(first[2][0], first[0], "w") # re-ready the scheduled one + second = w.step() + assert first[0] != second[0], "round robin must rotate" + + +def test_priority_and_biggest_walk(): + w = _World({"KV": EngineType.KV_CACHE, "VOC": EngineType.STATELESS}, + leaders={"KV", "VOC"}).use_priority() + w.add("r1", "KV", "decode") + w.add("r2", "KV", "prefill") + w.add("r3", "KV", "prefill") + w.add("r4", "VOC", "decode") + node, walk, rids = w.step() + assert (node, walk) == ("KV", "prefill") and rids == ["r2", "r3"] + + +def test_engine_not_ready_is_skipped(): + w = _World({"A": EngineType.KV_CACHE}, leaders={"A"}) + w.add("r1", "A", "w") + w.engines["A"].not_ready.add(("A", "r1")) + assert w.step() is None + w.engines["A"].not_ready.clear() + assert w.step() is not None + + +def test_hold_backoff_expires(): + w = _World({"A": EngineType.KV_CACHE}, leaders={"A"}) + w.add("r1", "A", "w") + w.hold(["r1"]) + assert w.step() is None + time.sleep(0.06) # HOLD_BACKOFF_SECONDS = 0.05 + assert w.step() is not None + + +def test_max_batch_and_exclude(): + w = _World({"A": EngineType.KV_CACHE}, leaders={"A"}) + for i in range(4): + w.add(f"r{i}", "A", "w") + node, walk, rids = w.step(max_batch_size=2) + assert len(rids) == 2 + assert w.step(exclude_target=("A", "w")) is None + + +def test_tp_follow_order_and_fairness(): + w = _World({"T": EngineType.KV_CACHE, "B": EngineType.STATELESS}, + leaders={"B"}) # follower rank for T + w.add("r1", "T", "w") + w.add("r2", "B", "w") + w.tp_follow("T", "w", ["r1"]) + node, _, _ = w.step() + assert node == "T" # the follow replays first + w.add("r1", "T", "w") + w.tp_follow("T", "w", ["r1"]) + node, _, _ = w.step() + # consec cap hit and B is ready: fairness yields to B in both. + assert node == "B" + node, _, _ = w.step() + assert node == "T" diff --git a/test/rust/test_walk_parity.py b/test/rust/test_walk_parity.py new file mode 100644 index 000000000..a0b7f1ee1 --- /dev/null +++ b/test/rust/test_walk_parity.py @@ -0,0 +1,283 @@ +"""Walk-core parity: the Rust walk core vs mstar's ``WorkerGraphIO``, driven +with identical event sequences over the same graphs. At every step the ready +sets must match; at the end, loop iteration counts and doneness must match. +Skipped unless the ``mstar_rust`` extension is installed.""" +from copy import deepcopy + +import pytest + +pytest.importorskip("mstar_rust") + +from mstar_rust import WalkSet + +from mstar.graph.base import GraphEdge, GraphNode, Loop, Parallel, Sequential +from mstar.graph.graph_io import WorkerGraphIO +from mstar.graph.rust_core import walks_to_json +from mstar.graph.special_destinations import EMIT_TO_CLIENT + + +def node(name, inputs, outputs): + return GraphNode( + name=name, input_names=set(inputs), + outputs=[GraphEdge(next_node=d, name=n) for n, d in outputs], + ) + + +class Harness: + """Drives both implementations in lockstep and asserts parity.""" + + def __init__(self, section, seeds): + self.py = WorkerGraphIO(deepcopy(section), wg_id="parity") + rust_set = WalkSet.from_json(walks_to_json({"walk": section})) + self.rs = rust_set.state("walk") + for node_name, input_name in seeds: + assert self.py.ingest_input( + GraphEdge(next_node=node_name, name=input_name)) + self.rs.seed(seeds) + self.assert_parity() + + def assert_parity(self): + assert sorted(self.py.ready_node_names) == sorted( + self.rs.ready_nodes()), "ready sets diverged" + assert self.py.wg_state_registry.is_done == self.rs.is_done(), \ + "doneness diverged" + assert dict(self.rs.loop_iters()) == self.py.get_loop_indices(), \ + "loop iteration counts diverged" + + def run_node(self, name): + """Pop + execute + complete `name` in both; route Python's internal + edges back (the Rust core routes internally).""" + assert name in self.py.ready_node_names + self.py.ready_node_names.discard(name) + completion = self.py.mark_node_complete(name) + for edge in completion.output_edges: + if edge.next_node in self.py.nodes: + self.py.ingest_input(edge) + + self.rs.schedule(name) + self.rs.complete(name) + self.assert_parity() + + def stop_loop(self, loop_name): + self.py.register_loop_finish_signal(loop_name) + self.rs.signal_loop_finish(loop_name) + + def run_until_done(self, max_steps=200): + steps = 0 + while not self.rs.is_done(): + ready = sorted(self.rs.ready_nodes()) + assert ready, "stuck: nothing ready but not done" + self.run_node(ready[0]) + steps += 1 + assert steps < max_steps, "walk did not terminate" + + def loop_parity(self): + assert dict(self.rs.loop_iters()) == self.py.get_loop_indices(), \ + "loop iteration counts diverged" + + +def test_sequential_chain(): + section = Sequential([ + node("A", ["x"], [("h", "B")]), + node("B", ["h"], [("h2", "C")]), + node("C", ["h2"], [("out", EMIT_TO_CLIENT)]), + ]) + h = Harness(section, [("A", "x")]) + h.run_until_done() + + +def test_parallel_fan_out_and_join(): + section = Sequential([ + node("A", ["x"], [("l", "B"), ("r", "C")]), + Parallel([ + node("B", ["l"], [("lb", "D")]), + node("C", ["r"], [("rc", "D")]), + ]), + node("D", ["lb", "rc"], [("out", EMIT_TO_CLIENT)]), + ]) + h = Harness(section, [("A", "x")]) + # After A completes, both branches must light up in both implementations. + h.run_node("A") + assert sorted(h.rs.ready_nodes()) == ["B", "C"] + h.run_until_done() + + +def _loop_graph(max_iters): + # Loop `outputs` snapshot the body's values BY NAME, so the body node + # declares a "final" edge (destination empty = value only captured by + # the loop). Both loop inputs are seeded: "fb" is otherwise fed only by + # the loop-back edge, which does not exist on iteration 0. + from mstar.graph.special_destinations import EMPTY_DESTINATION + body = node("L", ["seed", "fb"], + [("fb", "L"), ("tok", EMIT_TO_CLIENT), + ("final", EMPTY_DESTINATION)]) + loop = Loop(section=body, max_iters=max_iters, outputs=[ + GraphEdge(next_node="post", name="final")], name="dec") + return Sequential([loop, node("post", ["final"], + [("out", EMIT_TO_CLIENT)])]) + + +def test_loop_runs_to_max_iters(): + h = Harness(_loop_graph(4), [("L", "seed"), ("L", "fb")]) + h.run_until_done() + h.loop_parity() + + +def test_loop_early_stop_signal(): + h = Harness(_loop_graph(50), [("L", "seed"), ("L", "fb")]) + h.run_node("L") + h.run_node("L") + h.stop_loop("dec") + h.run_until_done() + h.loop_parity() + + +def test_nested_loops(): + from mstar.graph.special_destinations import EMPTY_DESTINATION + # inner: I runs 2 iters, its "chunk" snapshot feeds O; O restarts the + # inner loop for the next outer iteration and produces "done", which the + # outer loop snapshots for post. + inner = Loop( + section=node("I", ["iseed", "ifb"], + [("ifb", "I"), ("chunk", EMPTY_DESTINATION)]), + max_iters=2, outputs=[GraphEdge(next_node="O", name="chunk")], + name="inner") + outer_body = Sequential([ + inner, + node("O", ["chunk"], + [("iseed", "I"), ("ifb", "I"), ("done", EMPTY_DESTINATION)]), + ]) + outer = Loop( + section=outer_body, max_iters=3, + outputs=[GraphEdge(next_node="post", name="done")], name="outer") + section = Sequential([ + outer, + node("post", ["done"], [("out", EMIT_TO_CLIENT)]), + ]) + h = Harness(section, [("I", "iseed"), ("I", "ifb")]) + h.run_until_done() + h.loop_parity() + + +def test_shadow_mode_mirrors_and_detects(monkeypatch, caplog): + """MSTAR_RUST_WALK=shadow: the wrapper mirrors real WorkerGraphIO events + with no divergence on a healthy run, and STRICT mode raises when the + states are forced apart.""" + import logging + + from mstar.graph.rust_core import wrap_worker_graph_io + + monkeypatch.setenv("MSTAR_RUST_WALK", "shadow") + section = _loop_graph(3) + io = wrap_worker_graph_io( + WorkerGraphIO(deepcopy(section), wg_id="wg"), section, "wg") + + def drive(io): + io.ingest_input(GraphEdge(next_node="L", name="seed")) + io.ingest_input(GraphEdge(next_node="L", name="fb")) + while not io.wg_state_registry.is_done: + name = sorted(io.ready_node_names)[0] + io.ready_node_names.discard(name) + completion = io.mark_node_complete(name) + for edge in completion.output_edges: + if edge.next_node in io.nodes: + io.ingest_input(edge) + + with caplog.at_level(logging.ERROR): + drive(io) + assert io._suspended is None, io._suspended + assert not [r for r in caplog.records if "divergence" in r.message] + + # Fault injection: desync the Rust state -> strict mode must raise. + monkeypatch.setenv("MSTAR_RUST_WALK_STRICT", "1") + io2 = wrap_worker_graph_io( + WorkerGraphIO(deepcopy(section), wg_id="wg2"), section, "wg2") + io2.ingest_input(GraphEdge(next_node="L", name="seed")) + io2.ingest_input(GraphEdge(next_node="L", name="fb")) + io2._rs.signal_loop_finish("dec") # rust-only event = forced divergence + io2.ready_node_names.discard("L") + with pytest.raises(AssertionError, match="divergence"): + # The check fires at the settle point (after the completion's + # locally-destined edges are re-ingested), so drive a full step. + completion = io2.mark_node_complete("L") + for edge in completion.output_edges: + if edge.next_node in io2.nodes: + io2.ingest_input(edge) + + +def test_authority_mode_rust_drives(monkeypatch): + """MSTAR_RUST_WALK=1: ready set / doneness / loop indices come from the + Rust state; a full loop walk completes, and a forced divergence falls + back to Python instead of breaking the request.""" + from mstar.graph.rust_core import wrap_worker_graph_io + + monkeypatch.setenv("MSTAR_RUST_WALK", "1") + section = _loop_graph(3) + io = wrap_worker_graph_io( + WorkerGraphIO(deepcopy(section), wg_id="wg"), section, "wg") + io.ingest_input(GraphEdge(next_node="L", name="seed")) + io.ingest_input(GraphEdge(next_node="L", name="fb")) + steps = 0 + while not io.wg_state_registry.is_done: + name = sorted(io.ready_node_names)[0] + io.ready_node_names.discard(name) # -> rust schedule + completion = io.mark_node_complete(name) + for edge in completion.output_edges: + if edge.next_node in io.nodes: + io.ingest_input(edge) + steps += 1 + assert steps < 20 + assert io._suspended is None + assert io.get_loop_indices() == {"dec": 2} + + # Divergence -> per-request fallback to Python, request still completes. + io2 = wrap_worker_graph_io( + WorkerGraphIO(deepcopy(section), wg_id="wg2"), section, "wg2") + io2.ingest_input(GraphEdge(next_node="L", name="seed")) + io2.ingest_input(GraphEdge(next_node="L", name="fb")) + io2._rs.signal_loop_finish("dec") # force desync + steps = 0 + while not io2.wg_state_registry.is_done: + name = sorted(io2.ready_node_names)[0] + io2.ready_node_names.discard(name) + completion = io2.mark_node_complete(name) + for edge in completion.output_edges: + if edge.next_node in io2.nodes: + io2.ingest_input(edge) + steps += 1 + assert steps < 20 + assert io2._suspended is not None # fell back, and the walk finished + + +def test_pure_mode_no_python_registries(monkeypatch): + """MSTAR_RUST_WALK=pure: the Rust state is the only walk machine; the + adapter serves node views whose ready_inputs carry the REAL ingested + edges (the worker's tensor path), completion re-emits loop externals and + emits loop outputs at termination, and the walk finishes.""" + from mstar.graph.rust_core import PureRustWorkerGraphIO + + monkeypatch.setenv("MSTAR_RUST_WALK", "pure") + section = _loop_graph(3) + io = PureRustWorkerGraphIO(section, "wg-pure") + + seed = GraphEdge(next_node="L", name="seed") + fb0 = GraphEdge(next_node="L", name="fb") + assert io.ingest_input(seed) and io.ingest_input(fb0) + + completions = 0 + while not io.wg_state_registry.is_done: + name = sorted(io.ready_node_names)[0] + io.ready_node_names.discard(name) # -> schedule + slot fill + view = io.nodes[name] + # the worker's execution read: real edges, right names + assert set(view.ready_signals.ready_inputs) == view.input_names + if completions == 0 and name == "L": + assert view.ready_signals.ready_inputs["seed"] is seed + completion = io.mark_node_complete(name) + for edge in completion.output_edges: + if edge.next_node in io.nodes: + io.ingest_input(edge) + completions += 1 + assert completions < 20 + assert io.get_loop_indices() == {"dec": 2} + assert completions == 4 # 3 loop iters + post