| """The DSpark worker: CONTRACT.md v1.3 (lane O3) over lane O1's ofl transport, the drafter on one RTX 3090. |
| |
| State machine (DraftCore.handle(frame) -> reply frame) follows the contract and lane O3's reference worker |
| (remote_dspark/mock_worker.py WorkerCore) rule for rule: HELLO identity / config checks, sessions, strictly increasing |
| seq, per-slot generations (u32 serial arithmetic; older generations are stale even with NEW_GEN), the 7.2 window |
| rule (window.py), positions bounded by the session's max_model_len, PARTIAL_WINDOW, RESET / STATS / PING, |
| WANT_MARGIN, n_out (1.3). It exposes the same test surface (inject, slots[i].rng, step_log) so lane O3's tests can drive it |
| (tools/run_o3_conformance.py). The drafter sits behind a backend: |
| * GpuBackend: rows -> main_x -> window KV on the GPU (incremental: only the rows a STEP carries), batched drafts |
| through CUDA graphs per (batch size, k, n_out), with and without margins (WANT_MARGIN). If the CUDA context |
| dies, the worker answers INTERNAL and exits with status 3 so systemd restarts it. |
| * FakeBackend: lane O3's fake_draft over the raw rows (conformance tests only). |
| |
| Transports: ofl (service dspark-v1, TCP + RDMA backends, zero-copy payload) on the direct-link addresses only, and the |
| contract's length-prefixed TCP stub (optional). One lock serializes all frames; within a connection ofl answers in |
| order; frames of a superseded connection are refused by seq. |
| |
| Run (x86host): CUDA_VISIBLE_DEVICES=GPU-c4e7a371-d8a8-ebf0-918c-da3e38c87a93 python -m worker.server \\ |
| --root /opt/dsv41/x86host-data/dspark --ofl 10.202.28.1:18620 |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import ctypes |
| import json |
| import logging |
| import os |
| import signal |
| import socket |
| import struct |
| import sys |
| import threading |
| import time |
| from dataclasses import dataclass, field |
|
|
| import numpy as np |
|
|
| from . import protocol as P |
| from .window import WindowRange |
|
|
| log = logging.getLogger("dspark_worker") |
|
|
| GPU_UUID = "GPU-c4e7a371-d8a8-ebf0-918c-da3e38c87a93" |
| SERVICE = "dspark-v1" |
| MAX_REQ = 16 << 20 |
| MAX_RESP = 64 << 10 |
| DIRECT_LINK_IPS = ("10.202.28.1", "10.202.29.1") |
|
|
|
|
| def check_bind(ip: str) -> None: |
| """COMMON rule 9: listeners only on the direct-link addresses. DSPARK_WORKER_ALLOW_LOOPBACK=1 additionally allows |
| 127.0.0.0/8 (tests on one host); nothing else, never a wildcard.""" |
| if ip in DIRECT_LINK_IPS: |
| return |
| if os.environ.get("DSPARK_WORKER_ALLOW_LOOPBACK") == "1": |
| import ipaddress |
| try: |
| if ipaddress.ip_address(ip).version == 4 and ipaddress.ip_address(ip).is_loopback: |
| return |
| except ValueError: |
| pass |
| raise ValueError(f"bind {ip}: listeners only on the direct links {DIRECT_LINK_IPS}") |
|
|
|
|
| |
|
|
|
|
| @dataclass |
| class Inject: |
| """Fault injection (tests): the same knobs as lane O3's mock worker.""" |
| delay_ms: float = 0.0 |
| frame_status_once: int | None = None |
| req_status_once: dict[int, int] = field(default_factory=dict) |
| drop_next_step: bool = False |
|
|
|
|
| @dataclass |
| class SlotState: |
| window: int |
| generation: int | None = None |
| rng: WindowRange = field(init=False) |
|
|
| def __post_init__(self) -> None: |
| self.rng = WindowRange(self.window) |
|
|
|
|
| class Counters: |
| def __init__(self) -> None: |
| self.values: dict[str, int] = {} |
|
|
| def inc(self, key: str, n: int = 1) -> None: |
| self.values[key] = self.values.get(key, 0) + n |
|
|
|
|
| def expected_hello(drafter_identity: bytes, markov_identity: bytes, **kw) -> P.Hello: |
| """The live configuration (CONTRACT 6.2); kw overrides for tests.""" |
| base = dict(proto_min=1, proto_max=1, num_spec=5, window=128, hidden=5120, n_taps=3, tap_layer_ids=(37, 38, 39, 0), |
| payload_kind=P.PAYLOAD_TAPS_BF16, row_bytes=3 * 5120 * 2, noise_token_id=128799, vocab_size=129280, |
| max_model_len=327680, max_slots=4, max_rows_per_step=512, sample_mode=P.SAMPLE_GREEDY, tp_size=2, |
| drafter_identity=drafter_identity, markov_identity=markov_identity, client_name="") |
| base.update(kw) |
| return P.Hello(**base) |
|
|
|
|
| def _config_tuple(h: P.Hello) -> tuple: |
| return (h.num_spec, h.window, h.hidden, h.n_taps, tuple(h.tap_layer_ids[:h.n_taps]), h.payload_kind, |
| h.row_bytes, h.noise_token_id, h.vocab_size, h.sample_mode) |
|
|
|
|
| class DraftCore: |
| """Transport-independent contract state machine.""" |
|
|
| def __init__(self, expected: P.Hello, backend, *, worker_name: str = "o2-dspark-3090", |
| capabilities: int = P.CAP_TAPS_BF16 | P.CAP_HAS_MARGIN): |
| from .sample_policy import enabled |
| if expected.sample_mode == P.SAMPLE_TOP32: |
| if not enabled() or expected.num_spec not in (4, 5): |
| raise ValueError("sampled session requires switch and fixed depth4/5") |
| capabilities |= P.CAP_TOP32 |
| self.expected = expected |
| self.backend = backend |
| self.worker_name = worker_name |
| self.capabilities = capabilities |
| self.session: int | None = None |
| self.last_seq = 0 |
| self.max_model_len = expected.max_model_len |
| self.slots = [SlotState(expected.window) for _ in range(expected.max_slots)] |
| self.counters = Counters() |
| self.inject = Inject() |
| self.lock = threading.Lock() |
| self.fatal = threading.Event() |
| self.started = time.time() |
| self.step_log: list[dict] = [] |
| self.keep_step_log = False |
| self.worker_us: list[int] = [] |
|
|
| |
| def handle(self, buf) -> bytes | None: |
| t_in = time.perf_counter() |
| with self.lock: |
| try: |
| return self._handle(buf, t_in) |
| except Exception: |
| log.exception("frame failed") |
| self.counters.inc("internal_errors") |
| |
| |
| |
| self._classify_backend("frame") |
| try: |
| h = P.AppHeader.decode(buf) |
| return P.reply(h, P.INTERNAL) |
| except P.ProtocolError: |
| return P.AppHeader(P.MAGIC, P.VERSION, 0, 0, 0, 0, P.MALFORMED, 0).encode() |
|
|
| def _handle(self, buf, t_in: float) -> bytes | None: |
| mv = memoryview(buf) |
| try: |
| h = P.AppHeader.decode(mv) |
| if h.magic != P.MAGIC or h.version != P.VERSION or h.body_len != len(mv) - P.APP_HDR.size: |
| raise P.ProtocolError(f"magic {h.magic:#x} version {h.version} body_len {h.body_len} " |
| f"for {len(mv)} B") |
| except P.ProtocolError as e: |
| self.counters.inc("malformed_frames") |
| log.warning("malformed frame: %s", e) |
| return P.AppHeader(P.MAGIC, P.VERSION, 0, 0, 0, 0, P.MALFORMED, 0).encode() |
| body = mv[P.APP_HDR.size:] |
| if h.op == P.OP_PING: |
| return P.reply(h, P.OK) |
| if h.op == P.OP_HELLO: |
| return self._hello(h, body) |
| if h.op not in (P.OP_STEP, P.OP_RESET, P.OP_STATS): |
| return P.reply(h, P.UNSUPPORTED_OP) |
| if self.session is None or h.session != self.session: |
| self.counters.inc("no_session") |
| return P.reply(h, P.NO_SESSION) |
| if h.seq <= self.last_seq: |
| self.counters.inc("stale_seq") |
| return P.reply(h, P.STALE_SEQ) |
| self.last_seq = h.seq |
| if h.op == P.OP_STEP: |
| return self._step(h, body, t_in) |
| if h.op == P.OP_RESET: |
| if self.fatal.is_set(): |
| return P.reply(h, P.INTERNAL) |
| if len(body) != 4: |
| return P.reply(h, P.MALFORMED) |
| (mask,) = struct.unpack("<I", bytes(body)) |
| for i, s in enumerate(self.slots): |
| if mask >> i & 1: |
| self._reset_slot(i, None) |
| self.counters.inc("resets") |
| return P.reply(h, P.OK) |
| return P.reply(h, P.OK, json.dumps(self.stats()).encode()) |
|
|
| def _classify_backend(self, where: str) -> None: |
| """After a failure: if the backend is unusable (dead CUDA context), set fatal; a healthy backend keeps |
| serving (transient failure). A health probe that itself raises counts as unusable.""" |
| if self.fatal.is_set(): |
| return |
| try: |
| ok = bool(self.backend.healthy()) |
| except Exception: |
| log.exception("backend health probe raised after a %s failure", where) |
| ok = False |
| if not ok: |
| log.critical("the GPU backend is unusable after a %s failure: stopping the worker (non-zero exit; " |
| "systemd restarts it)", where) |
| self.fatal.set() |
|
|
| def _reset_slot(self, i: int, generation: int | None) -> None: |
| s = self.slots[i] |
| s.generation = generation |
| s.rng.clear() |
| self.backend.reset(i) |
|
|
| def stats(self) -> dict: |
| us = sorted(self.worker_us[-4096:]) |
| out = {"worker": self.worker_name, "session": self.session, "uptime_s": round(time.time() - self.started, 1), |
| **self.counters.values} |
| if us: |
| out["worker_us_p50"] = us[len(us) // 2] |
| out["worker_us_p99"] = us[min(len(us) - 1, int(len(us) * 0.99))] |
| out.update(self.backend.stats()) |
| return out |
|
|
| |
| def _hello(self, h: P.AppHeader, body) -> bytes: |
| if self.fatal.is_set(): |
| return P.reply(h, P.INTERNAL) |
| exp = self.expected |
| rep = P.hello_reply(1, exp.max_slots, exp.max_rows_per_step, self.capabilities, exp.window, exp.num_spec, |
| exp.drafter_identity, exp.markov_identity, self.worker_name) |
| try: |
| req = P.Hello.decode(body) |
| except P.ProtocolError: |
| return P.reply(h, P.MALFORMED) |
| if not req.proto_min <= 1 <= req.proto_max: |
| return P.reply(h, P.BAD_VERSION, rep) |
| if req.drafter_identity != exp.drafter_identity or req.markov_identity != exp.markov_identity: |
| self.counters.inc("identity_mismatch") |
| log.error("HELLO identity mismatch: drafter %s (mine %s), markov %s (mine %s)", |
| req.drafter_identity.hex(), exp.drafter_identity.hex(), req.markov_identity.hex(), |
| exp.markov_identity.hex()) |
| return P.reply(h, P.IDENTITY_MISMATCH, rep) |
| if (_config_tuple(req) != _config_tuple(exp) or not 1 <= req.max_slots <= exp.max_slots |
| or not 1 <= req.max_rows_per_step <= exp.max_rows_per_step |
| or not 1 <= req.max_model_len <= exp.max_model_len): |
| self.counters.inc("config_mismatch") |
| log.error("HELLO config mismatch: %s vs mine %s", req, exp) |
| return P.reply(h, P.CONFIG_MISMATCH, rep) |
| if h.session == 0: |
| return P.reply(h, P.MALFORMED) |
| if h.session != self.session: |
| for i in range(len(self.slots)): |
| self._reset_slot(i, None) |
| self.session = h.session |
| self.last_seq = h.seq |
| self.counters.inc("new_sessions") |
| log.info("HELLO new session %#x from %s (seq %d)", h.session, req.client_name, h.seq) |
| else: |
| if h.seq <= self.last_seq: |
| return P.reply(h, P.STALE_SEQ) |
| self.last_seq = h.seq |
| self.counters.inc("rehellos") |
| self.backend.set_max_model_len(req.max_model_len) |
| self.max_model_len = req.max_model_len |
| return P.reply(h, P.OK, rep) |
|
|
| |
| def _entry_problem(self, e: P.ReqEntry) -> str | None: |
| """CONTRACT 7.2 step 2 checks that need no slot state, plus the model range; None if fine.""" |
| exp = self.expected |
| if not 0 <= e.n_rows <= exp.window: |
| return f"n_rows {e.n_rows} outside 0..{exp.window}" |
| if not 0 <= e.anchor_pos <= self.max_model_len: |
| return f"anchor_pos {e.anchor_pos} outside 0..max_model_len {self.max_model_len}" |
| if e.n_rows and (e.first_pos < 0 or e.anchor_pos != e.first_pos + e.n_rows): |
| return f"rows [{e.first_pos}, +{e.n_rows}) do not end at anchor_pos {e.anchor_pos}" |
| if e.req_flags & P.REQ_NEED_DRAFT and not 0 <= e.anchor_token < exp.vocab_size: |
| return f"anchor_token {e.anchor_token} outside the vocabulary" |
| return None |
|
|
| def _step(self, h: P.AppHeader, body, t_in: float) -> bytes | None: |
| t0 = time.perf_counter() |
| if self.fatal.is_set(): |
| return P.reply(h, P.INTERNAL) |
| self.backend.discard() |
| inj = self.inject |
| if inj.frame_status_once is not None: |
| st, inj.frame_status_once = inj.frame_status_once, None |
| return P.reply(h, st) |
| exp = self.expected |
| try: |
| req = P.decode_step(body, exp.row_bytes, exp.max_slots, exp.max_rows_per_step) |
| if bool(req.step_flags & P.STEP_SAMPLED) != (exp.sample_mode == P.SAMPLE_TOP32): |
| raise P.ProtocolError("sampled flag/HELLO mismatch") |
| if exp.sample_mode == P.SAMPLE_TOP32 and (req.step_flags & P.STEP_WANT_MARGIN or |
| any(e.n_out not in (0, exp.num_spec) for e in req.entries)): |
| raise P.ProtocolError("sampled proposals require fixed depth and no margin admission") |
| for e in req.entries: |
| if e.slot >= exp.max_slots: |
| raise P.ProtocolError(f"slot {e.slot} >= max_slots {exp.max_slots}") |
| except P.CapacityError as e: |
| self.counters.inc("capacity_steps") |
| log.warning("STEP over capacity: %s", e) |
| return P.reply(h, P.CAPACITY) |
| except P.ProtocolError as e: |
| self.counters.inc("malformed_steps") |
| log.warning("malformed STEP: %s", e) |
| return P.reply(h, P.MALFORMED) |
| n, N = req.n_req, exp.num_spec |
| statuses = [[P.R_OK, 0, e.generation & 0xFFFFFFFF] for e in req.entries] |
| ids = np.zeros((n, N), dtype=np.int32) |
| conf = np.zeros((n, N), dtype=np.float32) |
| q_ids = np.zeros((n, N, 32), dtype=np.int32) if req.sample_meta else None |
| q_probs = np.zeros((n, N, 32), dtype=np.float32) if req.sample_meta else None |
| q_mask = np.zeros(n, dtype=np.int32) if req.sample_meta else None |
| want_margin = bool(req.step_flags & P.STEP_WANT_MARGIN) and bool(self.capabilities & P.CAP_HAS_MARGIN) |
| margin = np.zeros((n, N), dtype=np.float32) if want_margin else None |
| |
| drafts: list[tuple[int, int, int, int, WindowRange, int]] = [] |
| forced: dict[int, int | None] = {} |
| for i, e in enumerate(req.entries): |
| if self.keep_step_log: |
| self.step_log.append({"step_id": req.step_id, "slot": e.slot, "flags": e.req_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] |
| |
| if s.generation is not None and P.serial_newer(s.generation, e.generation): |
| statuses[i][0] = P.R_STALE_GENERATION |
| self.counters.inc("stale_generation") |
| continue |
| |
| |
| problem = self._entry_problem(e) |
| if problem is not None: |
| statuses[i][0] = P.R_BAD_ENTRY |
| self.counters.inc("bad_entries") |
| log.warning("slot %d: bad entry: %s", e.slot, problem) |
| continue |
| if s.generation is None or e.req_flags & P.REQ_NEW_GEN or e.generation != s.generation: |
| self._reset_slot(e.slot, e.generation) |
| try: |
| s.rng.apply(e.first_pos, e.n_rows, e.anchor_pos) |
| except ValueError as err: |
| statuses[i][0] = P.R_BAD_ENTRY |
| self.counters.inc("bad_entries") |
| log.warning("slot %d: bad entry: %s", e.slot, err) |
| self._reset_slot(e.slot, s.generation) |
| continue |
| if e.n_rows: |
| self.backend.stage(e.slot, e.first_pos, req.rows[e.row_start: e.row_start + e.n_rows]) |
| self.counters.inc("rows", e.n_rows) |
| have = s.rng.rows_for_anchor(e.anchor_pos) |
| statuses[i][1] = have |
| f = inj.req_status_once.pop(e.slot, None) |
| forced[i] = f |
| if f is not None: |
| statuses[i][0] = f |
| if f not in (P.R_OK, P.R_PARTIAL_WINDOW): |
| continue |
| if e.req_flags & P.REQ_NEED_DRAFT: |
| |
| drafts.append((i, e.slot, e.anchor_token, e.anchor_pos, s.rng, N if e.n_out == 0 else min(e.n_out, N))) |
| if f is None and have < s.rng.expected_rows(e.anchor_pos): |
| statuses[i][0] = P.R_PARTIAL_WINDOW |
| self.counters.inc("partial_windows") |
| try: |
| self.backend.flush() |
| if drafts: |
| calls = [(slot, tok, apos, rng, k) for _, slot, tok, apos, rng, k in drafts] |
| if req.sample_meta: |
| from .sample_policy import request_key |
| keys = [(request_key(h.session, slot, self.slots[slot].generation, |
| req.sample_meta[i][0]), req.step_id, apos) |
| for i, slot, _tok, apos, _rng, _k in drafts] |
| masks = [req.sample_meta[i][1] for i, *_rest in drafts] |
| d_ids, d_conf, d_qids, d_q = self.backend.draft_sampled(calls, N, keys, masks) |
| d_margin = None |
| for j, (i, *_rest) in enumerate(drafts): |
| q_ids[i], q_probs[i], q_mask[i] = d_qids[j], d_q[j], masks[j] |
| else: |
| d_ids, d_conf, d_margin = self.backend.draft(calls, N, want_margin) |
| for j, (i, *_rest, k) in enumerate(drafts): |
| ids[i, :k], conf[i, :k] = d_ids[j, :k], d_conf[j, :k] |
| if margin is not None: |
| margin[i, :k] = d_margin[j, :k] |
| self.counters.inc(f"n_out_{k}") |
| self.counters.inc("drafted_requests", len(drafts)) |
| except Exception: |
| log.exception("STEP %d: backend failed", req.step_id) |
| self.counters.inc("internal_errors") |
| self._classify_backend(f"STEP {req.step_id}") |
| self.backend.discard() |
| for i, e in enumerate(req.entries): |
| if statuses[i][0] in (P.R_OK, P.R_PARTIAL_WINDOW): |
| statuses[i][0] = P.R_INTERNAL |
| self._reset_slot(e.slot, self.slots[e.slot].generation) |
| ids[:] = 0 |
| conf[:] = 0.0 |
| if q_mask is not None: |
| q_mask[:] = 0 |
| q_probs[:] = 0 |
| 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) |
| self.worker_us.append(worker_us) |
| if len(self.worker_us) > 65536: |
| del self.worker_us[:32768] |
| rts = int((time.perf_counter() - t_in) * 1e6) |
| body_out = P.encode_step_reply(req.step_id, N, [tuple(x) for x in statuses], ids, conf, margin, |
| worker_us, rts) |
| if q_mask is not None: |
| import zlib |
| header = bytearray(body_out) |
| struct.pack_into("<I", header, 12, struct.unpack_from("<I", header, 12)[0] | P.REPLY_SAMPLED) |
| data = (struct.pack("<4sIII", b"PP3Q", 1, n, N) + q_mask.astype("<i4").tobytes() |
| + q_ids.astype("<i4").tobytes() + q_probs.astype("<f4").tobytes()) |
| body_out = bytes(header) + data |
| body_out += struct.pack("<I", zlib.crc32(body_out)) |
| return P.reply(h, P.OK, body_out) |
|
|
|
|
| |
|
|
|
|
| class GpuBackend: |
| """Rows -> window KV and drafts on the 3090 through model.DSparkDrafter.""" |
|
|
| def __init__(self, model, max_rows: int, row_bytes: int, graphs: bool = True, graph_reply: bool = False): |
| import torch |
|
|
| self.torch = torch |
| self.m = model |
| self.dev = model.dev |
| self.row_bytes = row_bytes |
| self.hidden3 = row_bytes // 2 |
| self.stage_buf = torch.empty((max_rows, row_bytes), dtype=torch.uint8, pin_memory=True) |
| self.pending: list[tuple[int, int, np.ndarray]] = [] |
| self.copy_done = torch.cuda.Event() |
| self.copy_done.record() |
| self.graphs = graphs |
| if graph_reply and not graphs: |
| raise ValueError("graph_reply requires CUDA graphs") |
| self.graph_reply = graph_reply |
| self.n_flush_rows = 0 |
| self.max_model_len = model.max_model_len |
|
|
| def set_max_model_len(self, n: int) -> None: |
| self.m.set_query_clamp(n) |
|
|
| def reset(self, slot: int) -> None: |
| self.pending = [p for p in self.pending if p[0] != slot] |
| self.m.reset_slots(self.torch.tensor([slot], device=self.dev)) |
|
|
| def stage(self, slot: int, first_pos: int, rows: np.ndarray) -> None: |
| self.pending.append((slot, first_pos, rows)) |
|
|
| def discard(self) -> None: |
| self.pending = [] |
|
|
| def healthy(self) -> bool: |
| """A trivial kernel and a sync: False once the CUDA context is dead (illegal address, Xid, ...).""" |
| torch = self.torch |
| try: |
| torch.cuda.synchronize(self.dev) |
| return float((torch.ones(1, device=self.dev) + 1).item()) == 2.0 |
| except Exception: |
| log.exception("GPU health probe failed") |
| return False |
|
|
| def flush(self) -> None: |
| torch = self.torch |
| if not self.pending: |
| return |
| n = sum(r.shape[0] for _, _, r in self.pending) |
| self.copy_done.synchronize() |
| sb = self.stage_buf[:n].numpy() |
| slot_rows = np.empty(n, dtype=np.int64) |
| pos_rows = np.empty(n, dtype=np.int64) |
| k = 0 |
| for slot, p0, rows in self.pending: |
| c = rows.shape[0] |
| sb[k:k + c] = rows |
| slot_rows[k:k + c] = slot |
| pos_rows[k:k + c] = np.arange(p0, p0 + c) |
| k += c |
| self.pending = [] |
| taps = self.stage_buf[:n].to(self.dev, non_blocking=True).view(torch.bfloat16).view(n, self.hidden3) |
| self.copy_done.record() |
| sr = torch.from_numpy(slot_rows).to(self.dev, non_blocking=True) |
| pr = torch.from_numpy(pos_rows).to(self.dev, non_blocking=True) |
| if self.graphs and n <= max(self.m.insert_graphs or {0: 0}): |
| self.m.insert_taps_graphed(sr, pr, taps) |
| else: |
| for a in range(0, n, 128): |
| b = min(n, a + 128) |
| self.m.insert(sr[a:b], pr[a:b], self.m.main_x(taps[a:b])) |
| self.n_flush_rows += n |
|
|
| def draft(self, reqs, N: int, want_margin: bool): |
| """reqs: (slot, anchor token, anchor pos, window range, k = positions consumed). The head runs for the batch's |
| largest k (the caller zeroes each request's positions >= its own k).""" |
| torch = self.torch |
| B = len(reqs) |
| n = max(k for *_r, k in reqs) |
| host = np.array([[slot, tok, apos - 1, rng.read_lo(apos)] for slot, tok, apos, rng, _k in reqs], |
| dtype=np.int64) |
| t = torch.from_numpy(host).to(self.dev, non_blocking=True) |
| slots, anchors, last, lo = t[:, 0], t[:, 1], t[:, 2], t[:, 3] |
| mg = None |
| packed = None |
| if self.graphs and self.m.has_graph(B, N, n, margin=want_margin): |
| out = self.m.draft_graphed(slots, anchors, last, N, ctx_lo=lo, n_out=n, margin=want_margin, |
| packed_reply=self.graph_reply) |
| if self.graph_reply: |
| packed = out[0] |
| else: |
| ids, conf = out[0][:B], out[1][:B] |
| mg = out[2][:B] if want_margin else None |
| elif want_margin: |
| ids, conf, mg = self.m.draft_with_margin(slots, anchors, last, N, ctx_lo=lo, n_out=n) |
| else: |
| ids, conf = self.m.draft(slots, anchors, last, N, ctx_lo=lo, n_out=n) |
| if packed is None: |
| parts = [ids.to(torch.float64), conf.to(torch.float64)] + ([mg.to(torch.float64)] if want_margin else []) |
| packed = torch.cat(parts, 1) |
| packed = packed.cpu().numpy() |
| margin = packed[:, 2 * N:].astype(np.float32) if want_margin else None |
| return packed[:, :N].astype(np.int32), packed[:, N:2 * N].astype(np.float32), margin |
|
|
| def stats(self) -> dict: |
| torch = self.torch |
| return {"gpu_mem_bytes": torch.cuda.memory_reserved(self.dev), "rows_inserted": self.n_flush_rows} |
|
|
|
|
| class FakeBackend: |
| """Raw rows per slot + a draft function of the window rows (lane O3's fake_draft); conformance tests only.""" |
|
|
| def __init__(self, window: int, row_bytes: int, max_slots: int, draft_fn, vocab: int): |
| self.window, self.row_bytes, self.draft_fn, self.vocab = window, row_bytes, draft_fn, vocab |
| self.rows = [np.zeros((window, row_bytes), np.uint8) for _ in range(max_slots)] |
|
|
| def set_max_model_len(self, n: int) -> None: |
| pass |
|
|
| def reset(self, slot: int) -> None: |
| pass |
|
|
| def stage(self, slot: int, first_pos: int, rows: np.ndarray) -> None: |
| for k in range(rows.shape[0]): |
| self.rows[slot][(first_pos + k) % self.window] = rows[k] |
|
|
| def discard(self) -> None: |
| pass |
|
|
| def healthy(self) -> bool: |
| return True |
|
|
| def flush(self) -> None: |
| pass |
|
|
| def draft(self, reqs, N: int, want_margin: bool): |
| ids = np.zeros((len(reqs), N), np.int32) |
| conf = np.zeros((len(reqs), N), np.float32) |
| for j, (slot, tok, apos, rng, _k) in enumerate(reqs): |
| win = [(p, self.rows[slot][p % self.window].tobytes()) for p in rng.positions(apos)] |
| ids[j], conf[j] = self.draft_fn(win, tok, apos, N, self.vocab) |
| return ids, conf, (conf.copy() if want_margin else None) |
|
|
| def stats(self) -> dict: |
| return {} |
|
|
|
|
| |
|
|
|
|
| def _load_ofl(): |
| """Lane O1's ofl.py: $OFL_PY_DIR, else offload/transport next to offload/dspark (merged branch), else the copy |
| pinned by run/pin_ofl.sh, else O1's x86host tree.""" |
| here = os.path.dirname(os.path.abspath(__file__)) |
| for d in (os.environ.get("OFL_PY_DIR"), os.path.normpath(os.path.join(here, "..", "..", "transport")), |
| os.path.expanduser("~/dsv41-offload/o2/lib"), |
| os.path.expanduser("~/dsv41-offload/o1/src/offload/transport")): |
| if d and os.path.exists(os.path.join(d, "ofl.py")): |
| if d not in sys.path: |
| sys.path.insert(0, d) |
| import ofl |
| return ofl |
| raise ImportError("ofl.py not found (set OFL_PY_DIR)") |
|
|
|
|
| class OflServer: |
| """core.handle on an ofl server; the request payload is read in place (no copy) and the reply memmoved into the |
| registered response buffer.""" |
|
|
| def __init__(self, core: DraftCore, bind_ip: str, port: int, rdma: bool = True): |
| check_bind(bind_ip) |
| ofl = _load_ofl() |
| L = ofl.lib() |
| self.core = core |
|
|
| def _cb(user, req): |
| try: |
| n = ctypes.c_uint32() |
| p = L.ofl_req_payload(req, ctypes.byref(n)) |
| buf = (ctypes.c_char * n.value).from_address(p) if n.value else bytearray() |
| rep = core.handle(memoryview(buf)) |
| if rep is None: |
| return -1 |
| cap = ctypes.c_uint32() |
| out = L.ofl_req_resp_buf(req, ctypes.byref(cap)) |
| if len(rep) > cap.value: |
| log.error("reply %d B > response capacity %d B", len(rep), cap.value) |
| return -2 |
| ctypes.memmove(out, rep, len(rep)) |
| L.ofl_req_set_len(req, len(rep)) |
| return 0 |
| except Exception: |
| log.exception("ofl handler") |
| return -3 |
|
|
| self._cb = ofl.HANDLER(_cb) |
| mask = (1 << ofl.BACKEND_TCP) | ((1 << ofl.BACKEND_RDMA) if rdma else 0) |
| cfg = ofl._ServerConfig(bind_ip.encode(), port, mask, None, 1, -1, SERVICE.encode(), MAX_REQ, MAX_RESP, |
| 1, 4, 2_000_000) |
| err = ctypes.create_string_buffer(512) |
| self.h = L.ofl_server_create(ctypes.byref(cfg), self._cb, None, err, 512) |
| if not self.h: |
| raise RuntimeError(f"ofl_server_create {bind_ip}:{port}: {err.value.decode()}") |
| self.L = L |
| L.ofl_server_start(self.h) |
| self.desc = L.ofl_server_describe(self.h).decode() |
|
|
| def stop(self) -> None: |
| if self.h: |
| self.L.ofl_server_destroy(self.h) |
| self.h = None |
|
|
|
|
| class TcpStubServer: |
| """The contract's TCP stub (CONTRACT 5): u32 length + frame; the newest connection supersedes the previous one.""" |
|
|
| def __init__(self, core: DraftCore, host: str, port: int): |
| if host in ("0.0.0.0", "", "::"): |
| raise ValueError("bind a specific (direct-link) address, never 0.0.0.0") |
| check_bind(host) |
| 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._conn: socket.socket | None = None |
| self._stop = False |
| self._lock = threading.Lock() |
| self.connections = 0 |
| threading.Thread(target=self._accept, daemon=True, name="dspark-tcp-accept").start() |
|
|
| def _accept(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="dspark-tcp-conn").start() |
|
|
| @staticmethod |
| def _recv_exact(conn: socket.socket, view: memoryview) -> bool: |
| got = 0 |
| while got < len(view): |
| k = conn.recv_into(view[got:]) |
| if k == 0: |
| return False |
| got += k |
| return True |
|
|
| def _serve(self, conn: socket.socket) -> None: |
| buf = bytearray(MAX_REQ + 64) |
| mv = memoryview(buf) |
| try: |
| while not self._stop: |
| if not self._recv_exact(conn, mv[:4]): |
| return |
| (n,) = struct.unpack("<I", mv[:4]) |
| if n > MAX_REQ + 64 or not self._recv_exact(conn, mv[:n]): |
| return |
| rep = self.core.handle(mv[:n]) |
| if rep is None: |
| return |
| conn.sendall(struct.pack("<I", len(rep)) + rep) |
| except OSError: |
| return |
| finally: |
| try: |
| conn.close() |
| except OSError: |
| pass |
|
|
| 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 build_core(root: str, graphs: bool = True, tp: int = 1, max_model_len: int = 327680) -> tuple[DraftCore, object]: |
| from .arm_selection import selected_arm |
| arm = selected_arm() |
| if tp == 2: |
| return build_core_tp2(root, graphs, max_model_len=max_model_len) |
| if arm != "k2-tp2": |
| raise ValueError(f"{arm} requires --tp 2; no one-card fallback") |
| if tp != 1: |
| raise SystemExit(f"--tp {tp}: 1 (one 3090) or 2 (the NVLink pair)") |
| import torch |
|
|
| from .checkpoint import Checkpoint |
| from .identity import compute as compute_identity |
| from .model import DSparkDrafter |
|
|
| vis = os.environ.get("CUDA_VISIBLE_DEVICES", "") |
| if vis != GPU_UUID: |
| raise SystemExit(f"CUDA_VISIBLE_DEVICES={vis!r}: this worker runs on {GPU_UUID} only (COMMON rule 3)") |
| t0 = time.time() |
| ck = Checkpoint(root) |
| ident = compute_identity(root, ck.manifest) |
| dev = torch.device("cuda", 0) |
| |
| |
| model = DSparkDrafter(ck, dev, slots=8, max_model_len=max_model_len) |
| exp = expected_hello(bytes.fromhex(ident["drafter_identity"]), bytes.fromhex(ident["markov_identity"]), |
| max_model_len=max_model_len) |
| graph_reply = os.environ.get("DSPARK_GRAPH_REPLY", "0") == "1" |
| if graphs: |
| model.capture(batch_sizes=tuple(range(1, exp.max_slots + 1)), ks=(exp.num_spec,), |
| insert_rows=(8, 16, 32, 64, 128), n_outs=tuple(range(1, exp.num_spec + 1)), |
| pack_reply=graph_reply) |
| backend = GpuBackend(model, exp.max_rows_per_step, exp.row_bytes, graphs=graphs, graph_reply=graph_reply) |
| core = DraftCore(exp, backend) |
| log.info("worker ready in %.1f s: drafter_identity %s, markov_identity %s, %.2f GiB reserved", time.time() - t0, |
| ident["drafter_identity"], ident["markov_identity"], torch.cuda.memory_reserved(dev) / 2**30) |
| return core, model |
|
|
|
|
| def build_core_tp2(root: str, graphs: bool = True, max_model_len: int = 327680) -> tuple[DraftCore, object]: |
| """The TP = 2 worker (tp2.py, lane E) on the NVLink pair: CUDA_VISIBLE_DEVICES = c4e7a371,7744b44b (in that order; |
| rank 0 = c4e7a371, PHB to the ConnectX-5). DSPARK_WORKER_TP_EMULATE=1 runs both ranks on GPU-7744b44b alone |
| (tests only: one card's speed, the same contract and drafts). max_model_len as in build_core (each rank holds the |
| replicated RoPE table: (max_model_len + 2048) x 64 fp32, 269 MB at 1,048,576).""" |
| from .arm_selection import selected_arm |
|
|
| arm = selected_arm() |
| import torch |
|
|
| from .checkpoint import Checkpoint |
| from .identity import compute as compute_identity |
| from .tp2 import TP2Backend, TPDrafter, check_visible |
|
|
| emulate = os.environ.get("DSPARK_WORKER_TP_EMULATE") == "1" |
| check_visible(pair=not emulate) |
| t0 = time.time() |
| ck = Checkpoint(root) |
| ident = compute_identity(root, ck.manifest) |
| devs = [torch.device("cuda", 0), torch.device("cuda", 0 if emulate else 1)] |
| if arm == "k2-tp2": |
| model = TPDrafter(ck, devs, slots=8, max_model_len=max_model_len) |
| else: |
| from offload.dspark.native.tp2_native import NATIVE_ROOT, build |
| from offload.dspark.native.identity import compute as compute_native_identity |
| ident = compute_native_identity(arm, ident, native_root=NATIVE_ROOT) |
| try: |
| model = build(arm, devs, ck, slots=8, max_model_len=max_model_len, native_root=NATIVE_ROOT) |
| except Exception as error: |
| raise RuntimeError(f"{arm}: native load/prepack failed; startup refused: {error}") from error |
| exp = expected_hello(bytes.fromhex(ident["drafter_identity"]), bytes.fromhex(ident["markov_identity"]), |
| max_model_len=max_model_len) |
| from .sample_policy import enabled |
| if enabled(): |
| if arm != "k2-tp2": |
| raise ValueError("PP3 supports only the production K2 arm") |
| exp.sample_mode = P.SAMPLE_TOP32 |
| if os.environ.get("DSPARK_V1_DUMP_DIR"): |
| raise ValueError("PP3 cannot combine sampled proposals and decode dumps") |
| dump_dir = os.environ.get("DSPARK_V1_DUMP_DIR", "") |
| sink = None |
| if dump_dir: |
| from .v1_dump import V1Dump |
| sink = V1Dump(model, ident, dump_dir, os.environ.get("DSPARK_V1_ROUNDS", "1")) |
| if graphs: |
| model.capture(batch_sizes=tuple(range(1, exp.max_slots + 1)), ks=(exp.num_spec,), |
| insert_rows=(8, 16, 32, 64, 128), n_outs=tuple(range(1, exp.num_spec + 1))) |
| if sink is None: |
| backend = TP2Backend(model, exp.max_rows_per_step, exp.row_bytes, graphs=graphs, max_slots=exp.max_slots) |
| core_type = DraftCore |
| else: |
| from .v1_dump import DumpDraftCore, DumpTP2Backend |
| backend = DumpTP2Backend(model, exp.max_rows_per_step, exp.row_bytes, graphs=graphs, |
| max_slots=exp.max_slots, sink=sink) |
| core_type = DumpDraftCore |
| core = core_type(exp, backend, worker_name="e-dspark-3090x2" + ("-emulated" if emulate else "")) |
| if arm != "k2-tp2": |
| core.worker_name = f"{arm}-K6" + ("-emulated" if emulate else "") |
| backend.arm_identity = ident |
| log.info("TP=2 worker ready in %.1f s: drafter_identity %s, markov_identity %s, %s GiB reserved per device, graphs=%s", |
| time.time() - t0, ident["drafter_identity"], ident["markov_identity"], |
| [round(torch.cuda.memory_reserved(d) / 2**30, 2) for d in dict.fromkeys(devs)], graphs) |
| return core, model |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--root", default="/opt/dsv41/x86host-data/dspark") |
| ap.add_argument("--ofl", action="append", default=[], help="ip:port for an ofl listener (TCP + RDMA)") |
| ap.add_argument("--tcp-stub", action="append", default=[], help="ip:port for the contract's TCP stub") |
| ap.add_argument("--no-rdma", action="store_true") |
| ap.add_argument("--no-graphs", action="store_true") |
| ap.add_argument("--tp", type=int, default=int(os.environ.get("DSPARK_WORKER_TP", "1")), |
| help="1: one 3090 (default); 2: the NVLink pair (tp2.py; or DSPARK_WORKER_TP=2)") |
| ap.add_argument("--max-model-len", type=int, default=327680, |
| help="largest max_model_len a session may announce (engine MAX_LEN); sizes the RoPE table") |
| ap.add_argument("--status-file", default=os.path.expanduser("~/dsv41-offload/o2/run/worker.status.json")) |
| args = ap.parse_args() |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") |
| if not args.ofl and not args.tcp_stub: |
| raise SystemExit("no listener: give --ofl and/or --tcp-stub") |
| if not 1 <= args.max_model_len <= 1048576: |
| raise SystemExit(f"--max-model-len {args.max_model_len} outside 1..1048576 (the model's max_position_embeddings)") |
| core, model = build_core(args.root, graphs=not args.no_graphs, tp=args.tp, max_model_len=args.max_model_len) |
| servers = [] |
| for spec in args.ofl: |
| ip, port = spec.rsplit(":", 1) |
| s = OflServer(core, ip, int(port), rdma=not args.no_rdma) |
| log.info("ofl listener %s", s.desc) |
| servers.append(s) |
| for spec in args.tcp_stub: |
| ip, port = spec.rsplit(":", 1) |
| s = TcpStubServer(core, ip, int(port)) |
| log.info("TCP stub listener %s:%d", ip, s.port) |
| servers.append(s) |
| stop = threading.Event() |
| signal.signal(signal.SIGTERM, lambda *_: stop.set()) |
| signal.signal(signal.SIGINT, lambda *_: stop.set()) |
| os.makedirs(os.path.dirname(args.status_file), exist_ok=True) |
| tick = 0 |
| while not stop.is_set() and not core.fatal.is_set(): |
| if tick % 10: |
| tick += 1 |
| stop.wait(0.5) |
| continue |
| tick += 1 |
| try: |
| with core.lock: |
| st = core.stats() |
| st["listeners"] = args.ofl + args.tcp_stub |
| st["pid"] = os.getpid() |
| st["time"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) |
| tmp = args.status_file + ".tmp" |
| with open(tmp, "w") as f: |
| json.dump(st, f) |
| os.replace(tmp, args.status_file) |
| except Exception: |
| log.exception("status file") |
| log.info("stopping") |
| for s in servers: |
| s.stop() |
| if core.fatal.is_set(): |
| log.critical("exit 3: GPU backend unusable") |
| sys.exit(3) |
|
|
|
|
|
|
| def build_adapted_evaluation_core(root: str, adapter: str, native_root: str, *, |
| emulated: bool = True, graphs: bool = True, |
| max_model_len: int = 327680) -> tuple[DraftCore, object]: |
| """Explicit isolated D2 entry point; never selected by the service CLI.""" |
| from offload.dspark.d2.train.worker_eval_build import build |
| return build(root, adapter, native_root, emulated=emulated, graphs=graphs, max_model_len=max_model_len) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
