# SPDX-License-Identifier: Apache-2.0 """The rank-0 round after the gather, torch-free: ring bookkeeping, circuit breaker, STEP, reply validation and the packed reply, plus the async proxy's per-round driver (wait for the gather, run the round, publish). Moved out of round.py without behaviour change so that one implementation runs in all three places: * the engine thread (synchronous rounds; round.RemoteRound), * the in-process proxy thread (async, VLLM_DSPARK_REMOTE_ASYNC_PROXY=thread; round.RemoteRound), * the out-of-process proxy (async, =process; proxy_proc.ProcessRound), which has its own interpreter and GIL, so the success path (publish) never waits for the engine's GIL. RoundCore holds no torch object and makes no torch or CUDA call; a subclass sets the attributes listed in its docstring and may override the hooks (_bcast_writable, _request_resync, _clear_resync). """ from __future__ import annotations import logging import time from dataclasses import dataclass import numpy as np from . import wire from .client import StepOutcome from .hostflag import FAR_AHEAD log = logging.getLogger("vllm.remote_dspark") STATUS_OK = 0 # every request that needed drafts got usable remote drafts STATUS_SOME_FAILED = 1 # some requests did not (see ok_mask) STATUS_ROUND_FAILED = 2 # no remote drafts this round STATUS_NAMES = {STATUS_OK: "ok", STATUS_SOME_FAILED: "some_failed", STATUS_ROUND_FAILED: "round_failed"} # async: the gather metadata of round t lives in slot t % META_SLOTS, so a proxy that reaches round t after the # watchdog completed it (and the stream went on) still reads round t's own metadata, never a later round's META_SLOTS = 64 WATCHDOG_EVENTS_KEPT = 64 class ReplyLayout: """Offsets of the int32 reply: [status, n, ok_mask[M], drafts[M*N], conf_bits[M*N]].""" def __init__(self, max_reqs: int, num_spec: int, sampled: bool = False): self.M, self.N = max_reqs, num_spec self.o_ok = 2 self.o_ids = 2 + max_reqs self.o_conf = self.o_ids + max_reqs * num_spec self.o_sample_mask = self.o_conf + max_reqs * num_spec self.o_qids = self.o_sample_mask + max_reqs self.o_q = self.o_qids + max_reqs * num_spec * 32 self.size = self.o_q + max_reqs * num_spec * 32 if sampled else self.o_sample_mask self.sampled = sampled @dataclass class BatchView: """CPU-side facts of a batch, from InputBatch (batch order).""" req_ids: list[str] slots: np.ndarray # [n] request-state index (idx_mapping_np) qsl: np.ndarray # [n+1] query_start_loc_np need_draft: np.ndarray # [n] bool k_window: np.ndarray # [n] drafts verified this step (0 for prefill) max_q: int | None = None # longest query span; None = from qsl (the process proxy gets it in the descriptor) @property def n(self) -> int: return len(self.req_ids) @property def longest_query(self) -> int: if self.max_q is not None: return int(self.max_q) return int(np.diff(self.qsl).max()) if len(self.qsl) > 1 else 0 @dataclass class GatherResult: valid: np.ndarray # [n] valid rows per request (rows sent) first_pos: np.ndarray # [n] position of the first candidate (= first valid) row anchor: np.ndarray # [n] anchor token sampled: np.ndarray # [n] bool: num_sampled > 0 n_candidates: int sample_meta: list[tuple[int, int]] | None = None @property def n_valid(self) -> int: return int(self.valid.sum()) def collect_meta(meta_np: np.ndarray, n: int, n_candidates: int) -> GatherResult: """The gather's metadata rows [valid, first_pos, anchor, sampled] x max_slots (int64, host) of one enqueue (its metadata slot), read after its copy completed.""" mh = meta_np return GatherResult(valid=mh[0, :n].copy(), first_pos=mh[1, :n].copy(), anchor=mh[2, :n].copy(), sampled=mh[3, :n].astype(bool), n_candidates=n_candidates, sample_meta=[(int(mh[5, i]), int(mh[4, i])) for i in range(n)] if mh.shape[0] == 6 else None) class RoundReply: """What an async round keeps of its result: the accepted reply (nothing reads it after the round).""" reply: wire.StepReply | None = None class RoundCore: """Rank-0 round logic shared by RemoteRound and the process proxy. A subclass sets: cfg, stats, N, W, M, layout (ReplyLayout), bcast_np (int32 reply being built), n_out_by_reqs, rounds, _last_epoch, _last_log, _err_logged, client, slots, breaker, gatherer (msg_np, max_rows, row_bytes, collect(n, R); gather_from_ring for synchronous resyncs), is_sender, async_mode; async only: sync (HostFlagCore: wait_gathered, taken, publish, fail_all, values, watchdog_stats), _stop (is_set()), _meta_slots (metadata slots of the gatherer: round t uses t % _meta_slots), _staged, _dead, resync_np, and either _resync_lock + _resync_req or its own _request_resync / _clear_resync. The round only READS the gather staging (gatherer.msg_np rows, resync_np windows): the STEP headers and entries go into a private buffer (_step_prefix) and resync rows into another (_resync_rows), and the message is sent as [headers, rows] (client.step rows=). In async mode the watchdog can complete a round while the proxy is inside it, and the engine stream then refills the same staging with the next round's rows; a proxy that wrote there would corrupt that later round's rows, silently and persistently (S11 P-1). A late round reads torn rows at worst: its reply is discarded (lost claim) and every slot resends its window. """ # ------------------------------------------------------------------ hooks def _bcast_writable(self) -> None: """Wait until bcast_np may be rewritten (RemoteRound: the last pinned copy out of it ran).""" def _request_resync(self, slots: set[int]) -> None: """Async: ask the engine side to stage these slots' ring windows with a later round's gather.""" with self._resync_lock: self._resync_req |= slots def _clear_resync(self, staged: frozenset[int]) -> None: """Async: these slots' windows were staged with the round being run.""" with self._resync_lock: self._resync_req -= staged # ------------------------------------------------------------------ round def n_out_for(self, n: int) -> int: """ReqEntry.n_out for a round of n requests (0 = all N positions).""" t = self.n_out_by_reqs if not t: return 0 c = int(t[min(n, len(t) - 1)]) return 0 if c >= self.N else max(1, c) def _pack_failed(self, n: int) -> None: self._bcast_writable() b = self.bcast_np b[:] = 0 b[0] = STATUS_ROUND_FAILED b[1] = n def _after_gather(self, bv: BatchView, g, res: RoundReply, step_id: int) -> int: assert self.gatherer is not None and self.slots is not None and self.client is not None n = bv.n # rank-0 bookkeeping of the ring (CONTRACT 7.1, 7.3). ready + epoch are read once, together: # every entry below is built for this connection, and step() refuses any newer one epoch = self.client.connection_epoch() if epoch is not None and epoch != self._last_epoch: self._last_epoch = epoch self.slots.mark_all_dirty() self.stats.inc(f"rounds_{self._phase(bv)}") slots = self.slots anchors_pos = self._apply_ring(bv, g) assert self.breaker is not None if not self.breaker.allow(): # open breaker: no network this round (CONTRACT 10); the ring above stays exact if self.breaker.state == "half_open": self.stats.set("breaker", self.breaker.snapshot()) return self._fail_round(bv, "cooldown") if epoch is None: reason = "disabled" if self.client.disabled else "not_ready" return self._fail_round(bv, reason) entries: list[wire.ReqEntry] = [] spans: list[tuple[int, int, int]] = [] resync = False n_out = self.n_out_for(n) for i in range(n): s = int(bv.slots[i]) st = slots[s] v = int(g.valid[i]) anchor_pos = int(anchors_pos[i]) flags = wire.RF_NEED_DRAFT if bv.need_draft[i] else 0 first, k = int(g.first_pos[i]), v if st.dirty or not st.announced: # the worker's window is unknown or new: (re)send the whole ring window, # which already holds this round's rows (CONTRACT 7.4) flags |= wire.RF_NEW_GEN if not st.rng.empty: pos = st.rng.positions(anchor_pos) if (pos.start, len(pos)) != (first, k): first, k = pos.start, len(pos) flags |= wire.RF_RESYNC resync = True if k == 0: # nothing to send for this request (cannot happen for a scheduled # request); an empty NEW_GEN entry would only fail BAD_ENTRY return self._fail_round(bv, "no_rows") entries.append(wire.ReqEntry(slot=s, generation=st.generation, first_pos=first, n_rows=k, anchor_token=int(g.anchor[i]), anchor_pos=anchor_pos, flags=flags, k_window=int(min(bv.k_window[i], 255)), n_out=n_out)) spans.append((s, first, k)) n_rows = sum(e.n_rows for e in entries) if resync: if self.async_mode: # the whole message is rebuilt from the ring: every slot's window must have been # staged by the engine stream this round; otherwise ask for it and fail this round need = {sl for sl, _, _ in spans} if not need <= self._staged: self._request_resync(need) return self._fail_round(bv, "resync_pending") rows = self._rows_from_staged(spans) got = rows.shape[0] else: got = self.gatherer.gather_from_ring(n, spans) rows = self._staged_rows(n, got) if got != n_rows: raise RuntimeError(f"resync gathered {got} rows, expected {n_rows}") self.stats.inc("resync_rounds") self.stats.inc("resync_rows", n_rows) elif n_rows != g.n_valid: raise RuntimeError(f"{n_rows} rows to send but {g.n_valid} valid rows staged") else: rows = self._staged_rows(n, n_rows) if self.async_mode and self.sync.taken(step_id): # completed by the watchdog (or a fail-all) while this proxy built it: the stream may already be # refilling the staging for a later round, so the rows are not sent (the publish below loses too) return self._fail_round(bv, "preempted") step_flags = wire.SF_WANT_MARGIN if self.cfg.want_margin else 0 if self.cfg.sampled_top32: if g.sample_meta is None: raise RuntimeError("sampled mode needs same-round seed and target-temperature metadata") step_flags |= wire.SF_SAMPLED out: StepOutcome = self.client.step(self._step_prefix(), entries, n_rows, step_id=step_id, step_flags=step_flags, epoch=epoch, rows=rows, sample_meta=g.sample_meta if self.cfg.sampled_top32 else None) if not out.ok: if out.reason not in ("not_ready", "disabled", "reconnected"): # sent, and failed self._breaker_record(False, out.reason or "unknown") return self._fail_round(bv, out.reason or "unknown") self.stats.observe("rtt_us", out.rtt_us) self.stats.observe(f"rtt_us_{self._phase(bv)}", out.rtt_us) self.stats.observe("worker_us", out.reply.worker_us) self.stats.inc("rows_sent", n_rows) self.stats.inc(f"rows_sent_{self._phase(bv)}", n_rows) if n_out: self.stats.inc("rounds_n_out_cut") status = self._accept_reply(bv, entries, out.reply, res) needed = bv.need_draft[:n].astype(bool) if needed.any(): # failed for the breaker when every request that needed drafts failed on the # worker. A round with nothing to draft (prefill appends only) never exercises # drafting: neutral, so it neither resets the count nor closes a probe usable = np.isin(out.reply.status, list(wire.USABLE_REQ_STATUS)) self._breaker_record(bool(usable[needed].any()), "every request failed") return status def _apply_ring(self, bv: BatchView, g) -> np.ndarray: """The round's rows in rank 0's model of the device ring (owner, generation, [lo, hi]); returns the anchor positions.""" slots = self.slots anchors_pos = g.first_pos + g.valid for i in range(bv.n): s = int(bv.slots[i]) st = slots[s] slots.begin_request(s, bv.req_ids[i]) v = int(g.valid[i]) if v <= 0: st.rng.clear() st.dirty = True continue try: st.rng.apply(int(g.first_pos[i]), v, int(anchors_pos[i])) except ValueError as e: log.warning("slot %d ring bookkeeping: %s", s, e) st.rng.clear() st.dirty = True return anchors_pos def _phase(self, bv: BatchView | None) -> str: """decode: every request's query is a verify span (<= N + 1 tokens); prefill: the step carried a longer chunk.""" if bv is None: return "unknown" return "prefill" if bv.longest_query > self.N + 1 else "decode" def _breaker_record(self, ok: bool, why: str) -> None: assert self.breaker is not None change = self.breaker.record(ok) if change is None: return self.stats.set("breaker", self.breaker.snapshot()) if change == "tripped": self.stats.inc("breaker_trips") log.warning("remote DSpark circuit breaker open (%s; trip %d): no network for the " "next %d rounds, then one probe", why, self.breaker.trips, self.breaker.cooldown) else: self.stats.inc("breaker_recoveries") log.info("remote DSpark circuit breaker closed: the probe round succeeded") def _fail_round(self, bv: BatchView, reason: str) -> int: assert self.slots is not None for s in bv.slots: self.slots[int(s)].dirty = True self.stats.inc(f"fallback_{reason}") self.stats.inc("rounds_failed") self._pack_failed(bv.n) return STATUS_ROUND_FAILED def _accept_reply(self, bv: BatchView, entries: list[wire.ReqEntry], reply: wire.StepReply, res: RoundReply) -> int: assert self.slots is not None n, N, L = bv.n, self.N, self.layout self._bcast_writable() b = self.bcast_np b[:] = 0 b[1] = n okm = b[L.o_ok:L.o_ok + n] ids = b[L.o_ids:L.o_ids + n * N].reshape(n, N) conf = b[L.o_conf:L.o_conf + n * N].view(np.float32).reshape(n, N) all_ok = True for i, e in enumerate(entries): rs = int(reply.status[i]) st = self.slots[e.slot] self.stats.inc(f"req_status_{wire.REQ_STATUS_NAMES.get(rs, rs)}".lower()) usable = rs in wire.USABLE_REQ_STATUS if usable: st.announced = True # a partial window is resent only if the ring holds more than the worker # has; otherwise it is genuine (rows never existed) and only counted st.dirty = (rs == wire.RS_PARTIAL_WINDOW and int(reply.window_rows[i]) < st.rng.rows_for_anchor( e.anchor_pos)) else: st.dirty = True if not bv.need_draft[i]: continue if usable: okm[i] = 1 ids[i] = reply.draft_ids[i] conf[i] = reply.confidence[i] if L.sampled: if reply.support_ids is None or reply.proposal_q is None or reply.sampled_mask is None: raise RuntimeError("sampled q missing before publish") b[L.o_sample_mask+i] = reply.sampled_mask[i] b[L.o_qids:L.o_qids+n*N*32].reshape(n,N,32)[i] = reply.support_ids[i] b[L.o_q:L.o_q+n*N*32].view(np.float32).reshape(n,N,32)[i] = reply.proposal_q[i] else: all_ok = False res.reply = reply stale = [e.slot for i, e in enumerate(entries) if int(reply.status[i]) == wire.RS_STALE_GENERATION] if stale: # rank 0's generations only grow (CONTRACT 7.1), so the worker's slot state does # not match ours; a resync under the same session would stay stale self.stats.inc("stale_generation_new_sessions") self.client.new_session(f"worker reported STALE_GENERATION for slots {stale}") self.stats.inc("rounds_ok") status = STATUS_OK if all_ok else STATUS_SOME_FAILED if not all_ok: self.stats.inc("rounds_some_failed") b[0] = status return status # ------------------------------------------------------------------ async proxy (thread or process) def _wait_gathered(self, t: int) -> bool: """Proxy thread: True once the engine stream finished round t's gather; False when stopping. A prefill chunk's forward can take seconds, so there is no give-up.""" assert self.sync is not None t0 = time.monotonic() next_log = 30.0 while not self._stop.is_set(): if self.sync.wait_gathered(t, 0.1): return True waited = time.monotonic() - t0 if waited >= next_log: next_log += 30.0 log.warning("remote DSpark async: the engine stream has not finished round %d's " "gather after %.0f s (words: %s)", t, waited, self.sync.values()) return False def _async_round(self, item) -> bool: """One round on the proxy thread. False when close() stopped it before round t was gathered.""" assert self.sync is not None t, bv, R, staged = item self._clear_resync(staged) # a request made below persists to the next round t0 = time.perf_counter() # outside the per-round handler: a failing wait means the hand-off itself is broken (dies) if not self._wait_gathered(t): return False # from here the reply buffer may be written: the stream finished round t-1's reply copy waited_us = (time.perf_counter() - t0) * 1e6 self.stats.observe("gather_us", waited_us) if self.sync.taken(t): self._skip_taken_round(t, bv, R, waited_us) self.stats.observe("rank0_round_us", (time.perf_counter() - t0) * 1e6) self._maybe_log() return True try: if bv is None: raise RuntimeError("enqueue failed on the engine thread") g = self.gatherer.collect(bv.n, R, t % self._meta_slots) res = RoundReply() self._staged = staged self._after_gather(bv, g, res, step_id=t) except Exception as e: # the round fails closed; the thread goes on n = bv.n if bv is not None else 0 self._pack_failed(n) if bv is not None: self.stats.inc("fallback_internal_error") if self.slots is not None: self.slots.mark_all_dirty() now = time.monotonic() if now - self._err_logged > 10.0: self._err_logged = now log.exception("remote DSpark async round %d failed: %s", t, e) finally: self._staged = frozenset() if not self.sync.reply_consumed(t - 1): # stream order makes gathered >= t imply that round t-1's reply was copied; a missing acknowledgment # means the hand-off's ordering is broken, and writing the page now could tear that copy raise RuntimeError(f"round {t}: the engine stream has not acknowledged round {t - 1}'s reply " f"(words {self.sync.values()})") # claim t, copy the reply into the page, release t: one GIL-free C call (S9 F5) if not self.sync.publish(t, self.bcast_np): # the watchdog failed round t first: the worker may have taken rows we now treat as # unsent, so every slot re-announces its window if self.slots is not None: self.slots.mark_all_dirty() self.stats.inc("watchdog_preempted") self.stats.inc(f"watchdog_preempted_{self._phase(bv)}") self._note_watchdog_event(t, bv, R, "late_publish", waited_us) done_us = (time.perf_counter() - t0) * 1e6 self.stats.observe("rank0_round_us", done_us) # gathered seen -> published: how long the engine stream stays parked on this round, by phase self.stats.observe(f"proxy_parked_us_{self._phase(bv)}", done_us - waited_us) self._maybe_log() return True def _skip_taken_round(self, t: int, bv: BatchView | None, R: int, waited_us: float) -> None: """Round t was completed (failed) by the watchdog or a fail-all before the proxy reached it. The stream has gone on and may already be writing a later round's rows into the staging, so no row is read and nothing is sent. The ring bookkeeping still runs, from round t's own metadata slot, and the round's slots resend their windows (NEW_GEN, resync from the ring) in a later round. A proxy more than META_SLOTS rounds behind cannot trust that slot either: those slots start their windows afresh.""" self.stats.inc("watchdog_preempted") self.stats.inc(f"watchdog_preempted_{self._phase(bv)}") self.stats.inc("watchdog_skipped_rounds") self._note_watchdog_event(t, bv, R, "skipped", waited_us) if bv is None or self.slots is None: return if self.rounds - t >= self._meta_slots: for i in range(bv.n): s = int(bv.slots[i]) self.slots.begin_request(s, bv.req_ids[i]) self.slots[s].rng.clear() self.slots[s].dirty = True self.stats.inc("watchdog_skipped_unknown_ring") return self._apply_ring(bv, self.gatherer.collect(bv.n, R, t % self._meta_slots)) for s in bv.slots: self.slots[int(s)].dirty = True def _note_watchdog_event(self, t: int, bv: BatchView | None, R: int, how: str, waited_us: float) -> None: """The last WATCHDOG_EVENTS_KEPT rounds the watchdog completed, with wall time and batch shape, so a run can attribute them to its stages (stats key watchdog_events).""" ev = getattr(self, "_wd_events", None) if ev is None: ev = self._wd_events = [] ev.append({"round": t, "time": round(time.time(), 3), "how": how, "phase": self._phase(bv), "n": bv.n if bv is not None else 0, "rows": int(R), "longest_query": bv.longest_query if bv is not None else 0, "proxy_wait_ms": round(waited_us / 1e3, 3)}) del ev[:-WATCHDOG_EVENTS_KEPT] self.stats.set("watchdog_events", list(ev)) def _die(self, reason: str) -> None: """The proxy thread is ending abnormally: fail every pending and later round first (one C call: failed reply, monotonic far-ahead release), then record it.""" assert self.sync is not None self.sync.fail_all(self.rounds + FAR_AHEAD) if self._dead is None: self._dead = reason try: self.stats.set("async_dead", reason) log.error("remote DSpark async proxy thread stopped (%s): every later round is " "target-only (zero drafts)", reason) except Exception: # noqa: BLE001 -- logging must not undo the release above pass def _rows_from_staged(self, spans: list[tuple[int, int, int]]) -> np.ndarray: """Async resync: the STEP rows for spans [(slot, first_pos, n_rows)] from the windows the engine stream staged (= gatherer.gather_from_ring on the CPU, same ring indexing), built in the round's private buffer: never in the shared STEP staging, which a later round may already own (S11 P-1). Returns [total, row_bytes] uint8.""" assert self.gatherer is not None g, W = self.gatherer, self.W total = sum(k for _, _, k in spans) if total > g.max_rows: raise ValueError(f"{total} resync rows > max_rows {g.max_rows}") buf = getattr(self, "_resync_rows", None) if buf is None or buf.shape[0] < total: buf = self._resync_rows = np.empty((total, g.row_bytes), dtype=np.uint8) rows = buf[:total] o = 0 for sl, p, k in spans: rows[o:o + k] = self.resync_np[sl][(p + np.arange(k, dtype=np.int64)) % W] o += k return rows def _staged_rows(self, n_req: int, n_rows: int) -> np.ndarray: """The first n_rows rows the gather staged for a round of n_req requests: a read-only use of the staging (uint8 view, no copy).""" g = self.gatherer off = wire.APP_HEADER_SIZE + wire.rows_offset_for(n_req, self.cfg.sampled_top32) return g.msg_np[off:off + n_rows * g.row_bytes] def _step_prefix(self) -> bytearray: """The round's private buffer for the STEP's AppHeader, StepHeader and entries (client.step rows=).""" b = getattr(self, "_prefix", None) if b is None: b = self._prefix = bytearray(wire.APP_HEADER_SIZE + wire.rows_offset_for(self.M, self.cfg.sampled_top32)) return b # ------------------------------------------------------------------ logging def _maybe_log(self) -> None: every = self.cfg.log_every if every <= 0 or self.rounds % every: return if self.breaker is not None: self.stats.set("breaker", self.breaker.snapshot()) if self.async_mode: self.stats.set("async_watchdog", self.sync.watchdog_stats()) snap = self.stats.snapshot() snap["rounds"] = self.rounds c = snap["counters"] lat = snap["latency"] log.info( "remote DSpark %s: %d rounds, ok %d, failed %d, resync %d; rtt p50/p99 %s/%s us, " "worker p50 %s us, gather p50 %s us, rank-0 round p50 %s us", self.cfg.policy, self.rounds, c.get("rounds_ok", 0), c.get("rounds_failed", 0), c.get("resync_rounds", 0), lat.get("rtt_us", {}).get("p50_us"), lat.get("rtt_us", {}).get("p99_us"), lat.get("worker_us", {}).get("p50_us"), lat.get("gather_us", {}).get("p50_us"), lat.get("rank0_round_us", {}).get("p50_us")) self.write_stats() def write_stats(self) -> None: if not self.is_sender or not self.cfg.stats_file: return self.stats.set("rounds", self.rounds) if self.breaker is not None: self.stats.set("breaker", self.breaker.snapshot()) if self.async_mode: self.stats.set("async_watchdog", self.sync.watchdog_stats()) try: self.stats.write_json(self.cfg.stats_file) except OSError as e: log.warning("remote DSpark stats file %s: %s", self.cfg.stats_file, e)