Blame for code/offload/dspark/spark/remote_dspark/client.py

# SPDX-License-Identifier: Apache-2.0
"""Rank-0 client of the remote DSpark worker: session, sequence numbers, HELLO,
background reconnect with backoff, and the STEP exchange (CONTRACT 6, 7.5, 9).

torch-free. The engine thread only calls ``step()`` (never blocks on connects); a
daemon thread owns connect + HELLO and hands a ready connection over.
"""

from __future__ import annotations

import logging
import secrets
import threading
import time
from dataclasses import dataclass
from typing import Callable

import numpy as np

from . import wire
from .config import RemoteConfig
from .stats import Stats
from .transport import RpcClient, TransportError, TransportTimeout, make_client

log = logging.getLogger("vllm.remote_dspark")

# While the worker stays unreachable the background thread retries every backoff interval (5 s in production), so only the
# first failure of a run and then one line per CONNECT_FAIL_LOG_S are logged; the counter connect_failures counts them all.
CONNECT_FAIL_LOG_S = 60.0


def check_reply_values(reply: wire.StepReply, entries: list[wire.ReqEntry], *,
                       vocab_size: int, window: int) -> None:
    """Refuse a STEP reply whose values the target must not see (CONTRACT 6.3, 7.2, 9.2).

    Shape and the step / generation echoes are checked before this. Draft ids become the
    target's input ids on both ranks and confidences feed the depth policy, so per entry:
      * status is a known per-request status;
      * window_rows <= min(W, anchor_pos) (no more rows than the anchor can read); for a
        drafted entry OK means exactly that many and PARTIAL_WINDOW fewer (7.2 step 5);
      * consumed positions (NEED_DRAFT, usable status, below n_out): draft ids in
        [0, vocab_size), confidences finite in [0, 1], margins finite and >= 0;
      * every other position (past n_out, entries without NEED_DRAFT, failed entries):
        draft 0, confidence 0.0, margin 0.0.
    Raises wire.WireError naming the first violation. Hot path (every round): scalar checks
    per entry, then one masked pass over the [n, N] arrays; the per-position search for the
    message runs only when that pass fails."""
    N = reply.num_spec
    status, wrows = reply.status.tolist(), reply.window_rows.tolist()
    ks = []
    for i, e in enumerate(entries):
        rs, wr = status[i], wrows[i]
        if rs not in wire.REQ_STATUS_NAMES:
            raise wire.WireError(f"slot {e.slot}: unknown request status {rs}")
        expected = min(window, max(e.anchor_pos, 0))
        if wr > expected:
            raise wire.WireError(f"slot {e.slot}: window_rows {wr} > {expected} readable rows "
                                 f"(anchor_pos {e.anchor_pos})")
        consumed = bool(e.flags & wire.RF_NEED_DRAFT) and rs in wire.USABLE_REQ_STATUS
        if consumed:
            if rs == wire.RS_OK and wr != expected:
                raise wire.WireError(f"slot {e.slot}: status OK with window_rows {wr} "
                                     f"!= {expected}")
            if rs == wire.RS_PARTIAL_WINDOW and wr >= expected:
                raise wire.WireError(f"slot {e.slot}: PARTIAL_WINDOW with window_rows {wr} "
                                     f">= {expected}")
        ks.append((N if e.n_out == 0 else min(e.n_out, N)) if consumed else 0)
    live = np.arange(N)[None, :] < np.array(ks)[:, None]          # [n, N] consumed positions
    ids, conf, m = reply.draft_ids, reply.confidence, reply.margin
    # NaN fails every comparison below, so a NaN anywhere fails the pass
    ok = (np.where(live, (ids >= 0) & (ids < vocab_size), ids == 0).all()
          and np.where(live, (conf >= 0.0) & (conf <= 1.0), conf == 0.0).all()
          and (m is None or np.where(live, np.isfinite(m) & (m >= 0.0), m == 0.0).all()))
    if not ok:
        _raise_first_value_violation(reply, entries, ks, vocab_size)


def _raise_first_value_violation(reply: wire.StepReply, entries: list[wire.ReqEntry],
                                 ks: list[int], vocab_size: int) -> None:
    for i, (e, k) in enumerate(zip(entries, ks)):
        ids, conf = reply.draft_ids[i], reply.confidence[i]
        bad = np.flatnonzero((ids[:k] < 0) | (ids[:k] >= vocab_size))
        if bad.size:
            j = int(bad[0])
            raise wire.WireError(f"slot {e.slot} position {j}: draft id {int(ids[j])} "
                                 f"outside [0, {vocab_size})")
        bad = np.flatnonzero(~((conf[:k] >= 0.0) & (conf[:k] <= 1.0)))
        if bad.size:
            j = int(bad[0])
            raise wire.WireError(f"slot {e.slot} position {j}: confidence {float(conf[j])!r} "
                                 "is not a probability")
        tails = [("draft id", ids[k:]), ("confidence", conf[k:])]
        if reply.margin is not None:
            m = reply.margin[i]
            bad = np.flatnonzero(~(np.isfinite(m[:k]) & (m[:k] >= 0.0)))
            if bad.size:
                j = int(bad[0])
                raise wire.WireError(f"slot {e.slot} position {j}: margin {float(m[j])!r} "
                                     "is not a finite top1 - top2")
            tails.append(("margin", m[k:]))
        for what, tail in tails:
            bad = np.flatnonzero(tail != 0)
            if bad.size:
                j = k + int(bad[0])
                raise wire.WireError(
                    f"slot {e.slot} position {j}: {what} {tail[j - k].item()!r} where the "
                    f"contract requires 0 (n_out {e.n_out}, status {int(reply.status[i])}, "
                    f"need_draft {bool(e.flags & wire.RF_NEED_DRAFT)})")
    raise wire.WireError("reply values violate CONTRACT 6.3")   # unreachable: pass and search agree


@dataclass
class StepOutcome:
    ok: bool
    reply: wire.StepReply | None = None
    reason: str | None = None        # short counter key when not ok
    rtt_us: float = 0.0
    # True when the frame may have been applied by the worker but its result is unknown
    # (timeout / transport error after sending): every slot in the STEP must resync.
    unknown_apply: bool = False


class RemoteDraftClient:
    def __init__(
        self,
        cfg: RemoteConfig,
        hello: wire.Hello,
        *,
        stats: Stats | None = None,
        transport_factory: Callable[..., RpcClient] | None = None,
        background: bool = True,
        session: int | None = None,
    ):
        """session: the session id of every HELLO and STEP (default: random). A respawned process proxy gets one its
        predecessor never used, so its first HELLO makes the worker drop every window (CONTRACT 6.2)."""
        self.cfg = cfg
        self.hello = hello
        self.stats = stats or Stats()
        self._factory = transport_factory or make_client
        if session is not None and not 0 < session <= 0xFFFFFFFF:
            raise ValueError(f"session {session} outside 1..2**32-1")
        self.session = session or secrets.randbits(32) or 1
        self._lock = threading.Lock()
        self._seq = 0
        self._client: RpcClient | None = None
        self._ready = False
        self.epoch = 0                       # successful HELLOs so far
        self.disabled_reason: str | None = None
        self.worker_info: wire.HelloReply | None = None
        self.last_error: str | None = None
        self._backoff_ms = cfg.backoff_ms
        self._next_attempt = 0.0
        self._fail_run = 0                   # connect/HELLO failures since the last successful HELLO
        self._fail_logged = 0.0              # monotonic time of the last connect-failure log line
        self._wake = threading.Event()
        self._stop = False
        self._resp = bytearray(1 << 16)
        self._thread: threading.Thread | None = None
        if background:
            self.start_background()

    def start_background(self) -> None:
        """Start the daemon thread that (re)connects whenever the connection is down."""
        if self._thread is None:
            self._thread = threading.Thread(
                target=self._run, name="dsv41-remote-dspark-connect", daemon=True)
            self._thread.start()

    # ------------------------------------------------------------------ state

    @property
    def ready(self) -> bool:
        return self._ready and self.disabled_reason is None

    @property
    def disabled(self) -> bool:
        return self.disabled_reason is not None

    def _next_seq(self) -> int:
        with self._lock:
            self._seq += 1
            return self._seq

    def close(self) -> None:
        self._stop = True
        self._wake.set()
        with self._lock:
            if self._client is not None:
                self._client.close()
            self._client = None
            self._ready = False

    def _disable(self, reason: str) -> None:
        self.disabled_reason = reason
        self.stats.inc("remote_disabled")
        self.stats.set("disabled_reason", reason)
        log.error("remote DSpark disabled for this engine run: %s", reason)
        with self._lock:
            if self._client is not None:
                self._client.close()
            self._client = None
            self._ready = False
        self._wake.set()

    def new_session(self, reason: str) -> None:
        """Start a new session and reconnect at once. The HELLO with the new session id
        makes the worker drop every window (CONTRACT 6.2); the epoch change then resyncs
        every live slot from rank 0's ring. For a worker state that no longer matches rank
        0's (STALE_GENERATION: rank 0 never sends an older generation, CONTRACT 7.1)."""
        with self._lock:     # one critical section: no reader sees the new id on the old connection
            old = self.session
            new = old
            while new == old or new == 0:
                new = secrets.randbits(32)
            self.session = new
            self._drop_locked(backoff=False)
        self.last_error = f"new session: {reason}"
        self.stats.inc("new_sessions")
        self.stats.inc("disconnects")
        log.error("remote DSpark: new session %#010x (was %#010x): %s", new, old, reason)
        self._wake.set()

    def connection_epoch(self) -> int | None:
        """The epoch of the ready connection, or None when there is none (one read under the
        lock, so ready and epoch belong to the same connection)."""
        with self._lock:
            return self.epoch if self._ready and self.disabled_reason is None else None

    def mark_failed(self, reason: str, *, backoff: bool = True) -> None:
        """Drop the connection; the connector thread reconnects (after the backoff)."""
        self.last_error = reason
        with self._lock:
            self._drop_locked(backoff=backoff)
        self.stats.inc("disconnects")
        self._wake.set()

    def _drop_locked(self, *, backoff: bool) -> None:
        if self._client is not None:
            self._client.close()
        self._client = None
        self._ready = False
        now = time.monotonic()
        if backoff:
            self._next_attempt = now + self._backoff_ms / 1000.0
            self._backoff_ms = min(self._backoff_ms * 2, self.cfg.backoff_max_ms)
        else:
            self._next_attempt = now

    # ------------------------------------------------------------------ connect

    def connect_now(self) -> bool:
        """Connect + HELLO synchronously. True when a ready connection is installed."""
        if self.disabled:
            return False
        timeout = self.cfg.connect_timeout_ms / 1000.0
        cli = None
        session = self.session
        try:
            cli = self._factory(self.cfg.endpoint)
            cli.connect(timeout)
            seq = self._next_seq()
            msg = bytearray(wire.frame(wire.OP_HELLO, seq, session, self.hello.pack()))
            buf = bytearray(1 << 16)  # >= the O1 transport's registered reply size
            n = cli.request(msg, buf, timeout)
            h, body = wire.split_frame(memoryview(buf)[:n])
            if h.op != wire.OP_HELLO or h.seq != seq or h.session != session:
                raise wire.WireError(
                    f"HELLO reply echoes op={h.op} seq={h.seq} session={h.session:#x}")
            if h.status != wire.ST_OK:
                name = wire.FRAME_STATUS_NAMES.get(h.status, str(h.status))
                if h.status in wire.FATAL_FRAME_STATUS:
                    detail = ""
                    if len(body) == wire.HELLO_REP.size:
                        rep = wire.HelloReply.unpack(body)
                        detail = (f" (worker identity {rep.drafter_identity.hex()}, "
                                  f"markov {rep.markov_identity.hex()}; ours "
                                  f"{self.hello.drafter_identity.hex()}, "
                                  f"{self.hello.markov_identity.hex()})")
                    cli.close()
                    self._disable(f"HELLO refused: {name}{detail}")
                    return False
                raise wire.WireError(f"HELLO refused: {name}")
            rep = wire.HelloReply.unpack(body)
            problems = []
            if rep.drafter_identity != self.hello.drafter_identity:
                problems.append("drafter identity")
            if rep.markov_identity != self.hello.markov_identity:
                problems.append("markov identity")
            if rep.window != self.hello.window or rep.num_spec != self.hello.num_spec:
                problems.append(f"window/num_spec {rep.window}/{rep.num_spec}")
            if self.hello.sample_mode == wire.SAMPLE_TOP32 and not rep.capabilities & wire.CAP_TOP32:
                problems.append("missing dyadic24 top32 capability")
            if not rep.capabilities & wire.CAP_TAPS_BF16:
                problems.append("no TAPS_BF16 capability")
            if rep.max_slots < self.hello.max_slots:
                problems.append(f"max_slots {rep.max_slots} < {self.hello.max_slots}")
            if rep.max_rows_per_step < self.hello.max_rows_per_step:
                problems.append(f"max_rows_per_step {rep.max_rows_per_step}")
            if problems:
                cli.close()
                self._disable("HELLO reply mismatch: " + ", ".join(problems))
                return False
        except (TransportError, wire.WireError, OSError) as e:
            if cli is not None:
                cli.close()
            self.last_error = f"connect: {e}"
            self.stats.inc("connect_failures")
            self._fail_run += 1
            now = time.monotonic()
            if self._fail_run == 1:
                self._fail_logged = now
                log.warning("remote DSpark connect/HELLO to %s failed: %s (target-only rounds; retrying, "
                            "further failures logged once per %.0f s)", self.cfg.endpoint, e, CONNECT_FAIL_LOG_S)
            elif now - self._fail_logged >= CONNECT_FAIL_LOG_S:
                self._fail_logged = now
                log.warning("remote DSpark connect/HELLO to %s still failing after %d attempts: %s",
                            self.cfg.endpoint, self._fail_run, e)
            return False
        with self._lock:
            stale = session != self.session     # new_session() ran meanwhile
            if stale:
                cli.close()
        if stale:                                # HELLO again at once under the new id
            return self.connect_now()
        with self._lock:
            if self._client is not None:
                self._client.close()
            self._client = cli
            self._ready = True
            self.epoch += 1
            self._backoff_ms = self.cfg.backoff_ms
        self.worker_info = rep
        self.stats.inc("hellos")
        self.stats.set("worker", {"name": rep.worker_name, "max_slots": rep.max_slots,
                                  "max_rows_per_step": rep.max_rows_per_step,
                                  "capabilities": rep.capabilities})
        failed, self._fail_run = self._fail_run, 0
        log.info("remote DSpark worker %r ready at %s (session %#010x, epoch %d)%s",
                 rep.worker_name, self.cfg.endpoint, session, self.epoch,
                 f" after {failed} failed connect attempts" if failed else "")
        return True

    def _run(self) -> None:
        while not self._stop:
            if self.disabled:
                return
            if self._ready:
                self._wake.wait()
                self._wake.clear()
                continue
            delay = self._next_attempt - time.monotonic()
            if delay > 0:
                self._wake.wait(delay)
                self._wake.clear()
                continue
            if not self.connect_now() and not self.disabled:
                with self._lock:
                    self._next_attempt = time.monotonic() + self._backoff_ms / 1000.0
                    self._backoff_ms = min(self._backoff_ms * 2, self.cfg.backoff_max_ms)

    # ------------------------------------------------------------------ STEP

    def step(
        self,
        buf,
        entries: list[wire.ReqEntry],
        n_rows: int,
        step_id: int,
        step_flags: int = 0,
        epoch: int | None = None,
        rows=None,
        sample_meta=None,
    ) -> StepOutcome:
        """Send one STEP. Without ``rows``: ``buf`` is a writable buffer whose rows already
        sit at ``wire.APP_HEADER_SIZE + wire.rows_offset_for(len(entries))``; the headers are
        written in front of them here, so the message goes out without a copy. With ``rows``
        (n_rows x row_bytes, contiguous; RoundCore always): ``buf`` is the caller's private
        buffer and takes only the headers and entries, and the message goes out as
        [headers, rows] (RpcClient.request_parts), so nothing is written into the memory that
        holds the rows: the gather staging the engine stream may already be refilling for a
        later round (S11 P-1). ``epoch``: the connection epoch the entries were built for; a
        newer connection is refused (reason "reconnected", nothing sent), since its slots
        were never resynced."""
        if self.disabled:
            return StepOutcome(False, reason="disabled")
        with self._lock:
            cli = self._client if self._ready else None
            session = self.session
            moved = epoch is not None and self.epoch != epoch
        if cli is None:
            return StepOutcome(False, reason="not_ready")
        if moved:
            return StepOutcome(False, reason="reconnected")
        row_bytes = self.hello.row_bytes
        seq = self._next_seq()
        ro = wire.pack_step_prefix(buf, wire.APP_HEADER_SIZE, step_id, entries, n_rows,
                                   row_bytes, step_flags, sample_meta=sample_meta)
        body_len = ro + n_rows * row_bytes
        wire.AppHeader(op=wire.OP_STEP, seq=seq, session=session,
                       body_len=body_len).pack_into(buf, 0)
        if rows is None:
            parts = (memoryview(buf).cast("B")[:wire.APP_HEADER_SIZE + body_len],)
        else:
            if rows.nbytes != n_rows * row_bytes:
                raise ValueError(f"rows hold {rows.nbytes} bytes, {n_rows} rows need {n_rows * row_bytes}")
            parts = (memoryview(buf).cast("B")[:wire.APP_HEADER_SIZE + ro], rows)
        t0 = time.perf_counter()
        try:
            if len(parts) == 1:
                n = cli.request(parts[0], self._resp, self.cfg.timeout_ms / 1000.0)
            else:
                n = cli.request_parts(parts, self._resp, self.cfg.timeout_ms / 1000.0)
        except TransportTimeout:
            self.mark_failed("timeout")
            return StepOutcome(False, reason="timeout", unknown_apply=True)
        except TransportError as e:
            self.mark_failed(f"transport: {e}")
            return StepOutcome(False, reason="transport_error", unknown_apply=True)
        rtt_us = (time.perf_counter() - t0) * 1e6
        self.stats.inc("bytes_sent", wire.APP_HEADER_SIZE + body_len)
        self.stats.inc("bytes_received", n)
        try:
            h, body = wire.split_frame(memoryview(self._resp)[:n])
            if h.op != wire.OP_STEP or h.seq != seq or h.session != session:
                raise wire.WireError(
                    f"reply echoes op={h.op} seq={h.seq} (want {seq}) "
                    f"session={h.session:#x}")
        except wire.WireError as e:
            self.mark_failed(f"bad reply: {e}")
            return StepOutcome(False, reason="bad_reply", rtt_us=rtt_us, unknown_apply=True)
        if h.status != wire.ST_OK:
            name = wire.FRAME_STATUS_NAMES.get(h.status, str(h.status))
            if h.status in wire.FATAL_FRAME_STATUS:
                self._disable(f"STEP refused: {name}")
            elif h.status in (wire.ST_NO_SESSION, wire.ST_STALE_SEQ):
                # The worker did not apply the frame; a fresh HELLO fixes the session.
                self.mark_failed(f"STEP refused: {name}", backoff=False)
            return StepOutcome(False, reason=f"frame_{name.lower()}", rtt_us=rtt_us,
                               unknown_apply=False)
        try:
            reply = wire.StepReply.unpack(body, expect_n_req=len(entries),
                                          expect_num_spec=self.hello.num_spec)
            if reply.step_id != step_id:
                raise wire.WireError(f"reply step_id {reply.step_id} != {step_id}")
            for e, g in zip(entries, reply.generation.tolist()):
                if g != (e.generation & wire.U32_MASK):
                    raise wire.WireError(f"slot {e.slot}: generation echo {g} != {e.generation}")
            if self.hello.sample_mode == wire.SAMPLE_TOP32:
                from .sampled_wire import validate
                if reply.support_ids is None or reply.proposal_q is None or reply.sampled_mask is None:
                    raise wire.WireError("missing sampled q; refusing sampled IDs")
                try:
                    needed = np.array([bool(e.flags & wire.RF_NEED_DRAFT) and
                                       int(reply.status[i]) in wire.USABLE_REQ_STATUS
                                       for i, e in enumerate(entries)])
                    validate(reply.sampled_mask, reply.support_ids, reply.proposal_q, reply.draft_ids,
                             needed, self.hello.vocab_size,
                             np.array([int(m[1]) for m in sample_meta], dtype=np.int32))
                except ValueError as exc:
                    raise wire.WireError(str(exc)) from exc
            elif reply.support_ids is not None:
                raise wire.WireError("unexpected q on greedy session")
            check_reply_values(reply, entries, vocab_size=self.hello.vocab_size,
                               window=self.hello.window)
        except wire.WireError as e:
            self.stats.set("last_bad_reply", str(e))
            log.warning("remote DSpark: STEP reply refused (round fails, slots resync): %s", e)
            self.mark_failed(f"bad reply: {e}")
            return StepOutcome(False, reason="bad_reply", rtt_us=rtt_us, unknown_apply=True)
        return StepOutcome(True, reply=reply, rtt_us=rtt_us)

    # ------------------------------------------------------------------ tools

    def simple_request(self, op: int, body: bytes = b"") -> tuple[int, bytes]:
        """PING / STATS / RESET on the engine's connection (not on the hot path)."""
        with self._lock:
            cli = self._client if self._ready else None
        if cli is None:
            raise TransportError("not connected")
        seq = self._next_seq()
        msg = bytearray(wire.frame(op, seq, self.session, body))
        n = cli.request(msg, self._resp, self.cfg.timeout_ms / 1000.0)
        h, rb = wire.split_frame(memoryview(self._resp)[:n])
        if h.op != op or h.seq != seq:
            raise wire.WireError("reply does not echo the request")
        return h.status, bytes(rb)