# SPDX-License-Identifier: Apache-2.0 """Reference implementation of the worker side of CONTRACT.md, with a fake drafter. Two uses: * offline tests of the Spark proxy (lane O3): every state rule of CONTRACT 7 is implemented here exactly; the fake drafter's output is a hash of the window rows the anchor may read, so any window error changes the drafts and the tests see it; * a reference for lane O2's server.py: ``WorkerCore`` is the whole protocol state machine with the drafter behind one function (``draft_fn``); O2 can replace ``fake_draft`` by the real model call and keep the rest, or port it. Transports: the contract's TCP stub (length-prefixed frames, newest connection wins) and lane O1's ofl server (``--ofl-dir``; TCP + RDMA backends, service "dspark-v1"). python -m remote_dspark.mock_worker --listen 127.0.0.1:18620 --identity <64hex> """ from __future__ import annotations import argparse import hashlib import json import logging import socket import threading import time from dataclasses import dataclass, field from typing import Callable import numpy as np from . import wire from .state import WindowRange from .transport import TransportError, recv_frame_blocking, send_frame_blocking log = logging.getLogger("dsv41.mock_worker") _LCG_A = 6364136223846793005 _LCG_C = 1442695040888963407 _M64 = (1 << 64) - 1 def fake_draft(window: list[tuple[int, bytes]], anchor_token: int, anchor_pos: int, num_spec: int, vocab: int) -> tuple[np.ndarray, np.ndarray]: """Deterministic stand-in for the drafter: ids/confidences from a hash of the window rows (position + bytes, ascending) and the anchor.""" h = hashlib.blake2b(digest_size=8) for pos, row in window: h.update(int(pos).to_bytes(8, "little", signed=True)) h.update(row) h.update(int(anchor_token).to_bytes(8, "little", signed=True)) h.update(int(anchor_pos).to_bytes(8, "little", signed=True)) x = int.from_bytes(h.digest(), "little") ids = np.empty(num_spec, dtype=np.int32) conf = np.empty(num_spec, dtype=np.float32) for j in range(num_spec): x = (x * _LCG_A + _LCG_C) & _M64 ids[j] = (x >> 17) % vocab conf[j] = ((x >> 40) & 0xFFFF) / 65535.0 return ids, conf @dataclass class SlotWindow: window: int row_bytes: int generation: int | None = None rng: WindowRange = field(init=False) rows: np.ndarray = field(init=False) def __post_init__(self) -> None: self.rng = WindowRange(self.window) self.rows = np.zeros((self.window, self.row_bytes), dtype=np.uint8) def reset(self, generation: int | None) -> None: self.generation = generation self.rng.clear() def window_rows(self, anchor_pos: int) -> list[tuple[int, bytes]]: return [(p, self.rows[p % self.window].tobytes()) for p in self.rng.positions(anchor_pos)] @dataclass class Inject: """One-shot / persistent fault injection for tests.""" delay_ms: float = 0.0 # persistent: sleep before replying to STEP frame_status_once: int | None = None # next STEP gets this frame status frame_status_always: int | None = None # persistent: every STEP gets this frame status req_status_once: dict[int, int] = field(default_factory=dict) # slot -> status req_status_always: int | None = None # persistent: every entry gets this status drop_next_step: bool = False # close the connection instead of replying # called with (StepReply, StepRequest) before the reply is packed; may edit the reply # (e.g. corrupt values); it stays armed until it returns True reply_hook_once: Callable[..., bool] | None = None reply_hook_fired: bool = False class WorkerCore: """Transport-agnostic contract state machine: handle(frame) -> reply frame.""" def __init__(self, expected: wire.Hello, *, worker_name: str = "mock-worker", draft_fn: Callable[..., tuple[np.ndarray, np.ndarray]] = fake_draft, capabilities: int = wire.CAP_TAPS_BF16 | wire.CAP_HAS_MARGIN): self.expected = expected self.worker_name = worker_name self.draft_fn = draft_fn self.capabilities = capabilities self.session: int | None = None self.last_seq = 0 self.slots = [SlotWindow(expected.window, expected.row_bytes) for _ in range(expected.max_slots)] self.counters = wire.Counters() self.inject = Inject() self.lock = threading.Lock() self.started = time.time() self.step_log: list[dict] = [] # tests: decoded entries per STEP self.keep_step_log = False # -------------------------------------------------------------- entry point def handle(self, buf) -> bytes | None: """Returns the reply frame, or None to drop the connection (fault injection).""" with self.lock: return self._handle(buf) def _handle(self, buf) -> bytes | None: try: h, body = wire.split_frame(buf) except wire.WireError as e: self.counters.inc("malformed_frames") log.warning("malformed frame: %s", e) return wire.frame(0, 0, 0, b"", wire.ST_MALFORMED) if h.op == wire.OP_PING: return wire.frame(h.op, h.seq, h.session) if h.op == wire.OP_HELLO: return self._hello(h, body) if h.op not in (wire.OP_STEP, wire.OP_RESET, wire.OP_STATS): return wire.frame(h.op, h.seq, h.session, b"", wire.ST_UNSUPPORTED_OP) if self.session is None or h.session != self.session: self.counters.inc("no_session") return wire.frame(h.op, h.seq, h.session, b"", wire.ST_NO_SESSION) if h.seq <= self.last_seq: self.counters.inc("stale_seq") return wire.frame(h.op, h.seq, h.session, b"", wire.ST_STALE_SEQ) self.last_seq = h.seq if h.op == wire.OP_STEP: return self._step(h, body) if h.op == wire.OP_RESET: if len(body) != wire.RESET_BODY.size: return wire.frame(h.op, h.seq, h.session, b"", wire.ST_MALFORMED) (mask,) = wire.RESET_BODY.unpack(bytes(body)) for i, s in enumerate(self.slots): if mask >> i & 1: s.reset(None) self.counters.inc("resets") return wire.frame(h.op, h.seq, h.session) stats = json.dumps(self.stats()).encode() return wire.frame(h.op, h.seq, h.session, stats) def stats(self) -> dict: return {"worker": self.worker_name, "session": self.session, "uptime_s": round(time.time() - self.started, 1), **self.counters.values} # -------------------------------------------------------------- HELLO def _hello(self, h: wire.AppHeader, body) -> bytes: try: req = wire.Hello.unpack(body) except wire.WireError: return wire.frame(h.op, h.seq, h.session, b"", wire.ST_MALFORMED) rep = wire.HelloReply( proto=1, max_slots=self.expected.max_slots, max_rows_per_step=self.expected.max_rows_per_step, capabilities=self.capabilities, window=self.expected.window, num_spec=self.expected.num_spec, drafter_identity=self.expected.drafter_identity, markov_identity=self.expected.markov_identity, worker_name=self.worker_name) if not req.proto_min <= 1 <= req.proto_max: return wire.frame(h.op, h.seq, h.session, rep.pack(), wire.ST_BAD_VERSION) if (req.drafter_identity != self.expected.drafter_identity or req.markov_identity != self.expected.markov_identity): self.counters.inc("identity_mismatch") return wire.frame(h.op, h.seq, h.session, rep.pack(), wire.ST_IDENTITY_MISMATCH) if (req.config_tuple() != self.expected.config_tuple() or req.max_slots > self.expected.max_slots or req.max_rows_per_step > self.expected.max_rows_per_step): self.counters.inc("config_mismatch") return wire.frame(h.op, h.seq, h.session, rep.pack(), wire.ST_CONFIG_MISMATCH) if h.session == 0: return wire.frame(h.op, h.seq, h.session, b"", wire.ST_MALFORMED) if h.session != self.session: for s in self.slots: s.reset(None) self.session = h.session self.last_seq = h.seq self.counters.inc("new_sessions") else: if h.seq <= self.last_seq: return wire.frame(h.op, h.seq, h.session, b"", wire.ST_STALE_SEQ) self.last_seq = h.seq self.counters.inc("rehellos") return wire.frame(h.op, h.seq, h.session, rep.pack()) # -------------------------------------------------------------- STEP def _step(self, h: wire.AppHeader, body) -> bytes | None: t0 = time.perf_counter() inj = self.inject self.counters.inc("step_frames") if inj.frame_status_once is not None: st, inj.frame_status_once = inj.frame_status_once, None return wire.frame(h.op, h.seq, h.session, b"", st) if inj.frame_status_always is not None: return wire.frame(h.op, h.seq, h.session, b"", inj.frame_status_always) exp = self.expected try: req = wire.decode_step(body, max_slots=exp.max_slots, max_rows=exp.max_rows_per_step, expected_row_bytes=exp.row_bytes, window=exp.window) except wire.WireError as e: self.counters.inc("malformed_steps") log.warning("malformed STEP: %s", e) return wire.frame(h.op, h.seq, h.session, b"", wire.ST_MALFORMED) n, N = req.n_req, exp.num_spec status = np.zeros(n, dtype=np.uint16) wrows = np.zeros(n, dtype=np.uint16) gens = np.zeros(n, dtype=np.uint32) ids = np.zeros((n, N), dtype=np.int32) conf = np.zeros((n, N), dtype=np.float32) want_margin = bool(req.step_flags & wire.SF_WANT_MARGIN) and bool( self.capabilities & wire.CAP_HAS_MARGIN) margin = np.zeros((n, N), dtype=np.float32) if want_margin else None row = 0 for i, e in enumerate(req.entries): gens[i] = e.generation & wire.U32_MASK rows = req.rows[row:row + e.n_rows] row += e.n_rows if self.keep_step_log: self.step_log.append({"step_id": req.step_id, "slot": e.slot, "flags": e.flags, "gen": e.generation, "first_pos": e.first_pos, "n_rows": e.n_rows, "anchor_pos": e.anchor_pos, "anchor_token": e.anchor_token}) s = self.slots[e.slot] # CONTRACT 7.2 step 1: an older generation is stale whatever its flags; NEW_GEN # resets the window only for the slot's own or a newer generation (resync) if s.generation is not None and wire.gen_newer(s.generation, e.generation): status[i] = wire.RS_STALE_GENERATION self.counters.inc("stale_generation") continue if s.generation is None or e.flags & wire.RF_NEW_GEN or e.generation != s.generation: s.reset(e.generation) try: s.rng.apply(e.first_pos, e.n_rows, e.anchor_pos) except ValueError as err: status[i] = wire.RS_BAD_ENTRY self.counters.inc("bad_entries") log.warning("slot %d: bad entry: %s", e.slot, err) s.rng.clear() continue for k in range(e.n_rows): s.rows[(e.first_pos + k) % exp.window] = rows[k] self.counters.inc("rows", e.n_rows) have = s.rng.rows_for_anchor(e.anchor_pos) wrows[i] = have forced = inj.req_status_once.pop(e.slot, None) if forced is None: forced = inj.req_status_always if forced is not None: status[i] = forced if forced not in wire.USABLE_REQ_STATUS: continue if e.flags & wire.RF_NEED_DRAFT: d, c = self.draft_fn(s.window_rows(e.anchor_pos), e.anchor_token, e.anchor_pos, N, exp.vocab_size) # CONTRACT 1.3: positions >= n_out are not consumed and come back as zeros; # the others are exactly the full-block draft k = N if e.n_out == 0 else min(int(e.n_out), N) ids[i, :k], conf[i, :k] = d[:k], c[:k] if margin is not None: margin[i, :k] = c[:k] self.counters.inc(f"n_out_{k}") self.counters.inc("drafted_requests") if forced is None: if have < s.rng.expected_rows(e.anchor_pos): status[i] = wire.RS_PARTIAL_WINDOW self.counters.inc("partial_windows") else: status[i] = wire.RS_OK self.counters.inc("steps") if inj.delay_ms > 0: time.sleep(inj.delay_ms / 1000.0) if inj.drop_next_step: inj.drop_next_step = False return None worker_us = int((time.perf_counter() - t0) * 1e6) rep = wire.StepReply(step_id=req.step_id, num_spec=N, status=status, window_rows=wrows, generation=gens, draft_ids=ids, confidence=conf, margin=margin, worker_us=worker_us, recv_to_send_us=worker_us) if inj.reply_hook_once is not None and inj.reply_hook_once(rep, req): inj.reply_hook_once = None inj.reply_hook_fired = True return wire.frame(h.op, h.seq, h.session, rep.pack()) # ------------------------------------------------------------------ TCP stub server class TcpStubServer: """Length-prefixed frames; the newest connection supersedes the previous one.""" def __init__(self, core: WorkerCore, host: str, port: int): if host in ("0.0.0.0", "", "::"): raise ValueError("bind a specific (direct-link) address, never 0.0.0.0") self.core = core self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) self.sock.bind((host, port)) self.sock.listen(4) self.port = self.sock.getsockname()[1] self.host = host self._conn: socket.socket | None = None self._stop = False self._lock = threading.Lock() self.connections = 0 self._thread = threading.Thread(target=self._accept_loop, daemon=True, name="mock-worker-accept") def start(self) -> "TcpStubServer": self._thread.start() return self def stop(self) -> None: self._stop = True try: self.sock.close() except OSError: pass with self._lock: if self._conn is not None: try: self._conn.close() except OSError: pass def drop_connection(self) -> None: with self._lock: if self._conn is not None: try: self._conn.shutdown(socket.SHUT_RDWR) except OSError: pass def _accept_loop(self) -> None: while not self._stop: try: conn, _ = self.sock.accept() except OSError: return conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) with self._lock: old, self._conn = self._conn, conn self.connections += 1 if old is not None: try: old.shutdown(socket.SHUT_RDWR) old.close() except OSError: pass threading.Thread(target=self._serve, args=(conn,), daemon=True, name="mock-worker-conn").start() def _serve(self, conn: socket.socket) -> None: buf = bytearray(1 << 20) try: while not self._stop: n = recv_frame_blocking(conn, buf) if n is None: return reply = self.core.handle(memoryview(buf)[:n]) if reply is None: conn.close() return send_frame_blocking(conn, reply) except (OSError, TransportError): return finally: try: conn.close() except OSError: pass def serve_ofl(core: WorkerCore, bind_ip: str, port: int, ofl_dir: str, rdma: bool = True): """Serve the contract on lane O1's ofl server (service dspark-v1).""" import os os.environ["VLLM_DSPARK_REMOTE_OFL_DIR"] = ofl_dir from .o1_adapter import DEFAULT_MAX_REQ, DEFAULT_MAX_RESP, SERVICE, load_ofl ofl = load_ofl() def handler(req) -> int: reply = core.handle(req.payload) if reply is None: return -1 req.respond(reply) return 0 backends = (ofl.BACKEND_TCP, ofl.BACKEND_RDMA) if rdma else (ofl.BACKEND_TCP,) srv = ofl.Server(bind_ip, port, service=SERVICE, handler=handler, backends=backends, max_req=DEFAULT_MAX_REQ, max_resp=DEFAULT_MAX_RESP, inflight=1) srv.start() return srv def main() -> None: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--listen", required=True, help="host:port (direct-link address)") ap.add_argument("--identity", required=True, help="64-hex drafter identity") ap.add_argument("--markov-identity", default="0" * 64) ap.add_argument("--ofl-dir", default="", help="serve over lane O1's ofl instead of TCP stub") ap.add_argument("--no-rdma", action="store_true") ap.add_argument("--delay-ms", type=float, default=0.0, help="simulated draft time") args = ap.parse_args() logging.basicConfig(level=logging.INFO) host, port = args.listen.rsplit(":", 1) expected = wire.Hello(drafter_identity=wire.parse_hex32(args.identity, "--identity"), markov_identity=wire.parse_hex32(args.markov_identity, "--markov-identity")) core = WorkerCore(expected) core.inject.delay_ms = args.delay_ms if args.ofl_dir: srv = serve_ofl(core, host, int(port), args.ofl_dir, rdma=not args.no_rdma) log.info("mock worker on ofl %s", srv.describe()) else: srv = TcpStubServer(core, host, int(port)).start() log.info("mock worker on tcp %s:%d", host, srv.port) try: while True: time.sleep(3600) except KeyboardInterrupt: srv.stop() if __name__ == "__main__": main()