Download code/offload/dspark/spark/remote_dspark/roundcore.py from ProprietaryLegal/dsv41-sampled-drafting: direct link, hf CLI and curl.
- Browser
- Download file 28.2 kB
-
https://huggingface.co/ProprietaryLegal/dsv41-sampled-drafting/resolve/main/code/offload/dspark/spark/remote_dspark/roundcore.py
- Command line
-
hf download hf://ProprietaryLegal/dsv41-sampled-drafting/code/offload/dspark/spark/remote_dspark/roundcore.py
-
curl -L -o roundcore.py https://huggingface.co/ProprietaryLegal/dsv41-sampled-drafting/resolve/main/code/offload/dspark/spark/remote_dspark/roundcore.py
28.2 kB
| # 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 | |
| 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) | |
| def n(self) -> int: | |
| return len(self.req_ids) | |
| 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 | |
| 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 | |
| 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) | |