| |
| """Environment configuration of the remote DSpark proxy (Spark side). |
| |
| The overlay is inactive unless VLLM_DSPARK_REMOTE is set (then the image's own code |
| path runs byte-identically). Every value is validated at startup; a bad value raises. |
| |
| VLLM_DSPARK_REMOTE endpoint for rank 0: o1rdma://10.202.28.1:18620 (lane O1's ofl, RDMA), |
| o1tcp://... (ofl's TCP backend) or tcp://... (the contract's TCP stub) |
| VLLM_DSPARK_REMOTE_BIND this side's direct-link IPv4 (10.202.28.2); same as ?bind= in the |
| endpoint, without a '?' that a launcher shell could glob |
| VLLM_DSPARK_REMOTE_OFL_DIR directory holding libofl.so + ofl.py (o1tcp/o1rdma endpoints) |
| VLLM_DSPARK_REMOTE_POLICY shadow (default) | resilience | capacity |
| VLLM_DSPARK_REMOTE_IDENTITY 64-hex drafter identity (CONTRACT 6.7); required |
| VLLM_DSPARK_REMOTE_MARKOV_IDENTITY 64-hex; default: sha256 of the VLLM_DSPARK_MARKOV_CACHE |
| file, zeros if that variable is unset |
| VLLM_DSPARK_REMOTE_TIMEOUT_MS per-round deadline (default 50) |
| VLLM_DSPARK_REMOTE_CONNECT_TIMEOUT_MS connect + HELLO deadline (default 1000) |
| VLLM_DSPARK_REMOTE_BACKOFF_MS first reconnect delay after a failure (default 1000, |
| doubles up to VLLM_DSPARK_REMOTE_BACKOFF_MAX_MS, default 30000) |
| VLLM_DSPARK_REMOTE_ON_FAIL target_only (default) | raise (capacity policy) |
| VLLM_DSPARK_REMOTE_TRIP circuit breaker (CONTRACT 10, breaker.py): consecutive failed attempts |
| before the remote is skipped (default 3; 0 = off) |
| VLLM_DSPARK_REMOTE_COOLDOWN rounds without network after a trip, then one probe (default 256) |
| VLLM_DSPARK_REMOTE_STRICT_START 1: raise at startup if connect + HELLO fails (default 0) |
| VLLM_DSPARK_REMOTE_SHADOW_USE local (default) | remote (shadow policy) |
| VLLM_DSPARK_REMOTE_WANT_MARGIN 1: ask the worker for top1-top2 margins (default 1 in shadow) |
| VLLM_DSPARK_REMOTE_DRAFT_MS remote draft round cost for the depth policy (ms, finite >= 0); |
| replaces VLLM_DSPARK_DRAFT_MS in capacity/resilience |
| VLLM_DSPARK_REMOTE_ROUND_MS verify-only forward table for remote mode: ms for depth 0..k |
| (k + 1 finite values > 0, non-decreasing); replaces |
| VLLM_DSPARK_ROUND_MS (whose live value includes the local |
| drafter) in capacity/resilience. With VLLM_DSPARK_DEPTH_POLICY |
| on, capacity/resilience need BOTH remote costs (no mixed cost |
| model) and refuse "confr:<table>" (its own table includes the |
| local drafter). Shadow keeps the local measurements. |
| VLLM_DSPARK_REMOTE_ASYNC 1: capacity policy only - rank 0 does not block its CPU on the round; |
| the engine stream and the round's proxy thread hand each round over |
| through two words in registered host memory (streamsync.HostFlagSync); |
| the proxy thread makes no CUDA call. Self-tested at startup; falls back |
| to the synchronous round with a warning. Default 0. |
| VLLM_DSPARK_REMOTE_ASYNC_PROXY async mode: where the round's proxy runs. thread (default): a Python thread of the |
| engine process (needs the engine's GIL to publish). process: a torch-free child |
| process with its own interpreter (proxy_proc.py); the gather staging, the descriptor |
| ring and the hand-off words share one registered memfd region (shmregion.py), so the |
| success path never waits for the engine's GIL; the switch interval is not changed |
| VLLM_DSPARK_REMOTE_HOSTFLAG_LIB async mode: path of libdsv41_hostflag.so (built from |
| remote_dspark/hostflag.c). Default: VLLM_DSPARK_REMOTE_OFL_DIR's copy, |
| else hostflag.c compiled with cc at startup. |
| VLLM_DSPARK_REMOTE_ASYNC_NAP_US async mode: the proxy thread spins on the gathered word; > 0 = sleep this |
| many us between checks after VLLM_DSPARK_REMOTE_ASYNC_SPIN_US of |
| spinning (default 0: spin; one rank-0 core while a round is pending) |
| VLLM_DSPARK_REMOTE_ASYNC_SPIN_US see ASYNC_NAP_US (default 1000) |
| VLLM_DSPARK_REMOTE_ASYNC_WATCHDOG_MS async mode: a native watchdog fails and releases a round the engine stream |
| gathered but nobody released within this many ms (zero drafts for that round), so a |
| proxy thread that cannot run never parks the GPU for good. Default 0 = auto: |
| max(4 x VLLM_DSPARK_REMOTE_TIMEOUT_MS, 200) |
| VLLM_DSPARK_REMOTE_ASYNC_RESPAWNS async process mode: new proxy processes the engine may start after its proxy died |
| or stalled (default 3; 0 = none: target-only rounds for the rest of the run). Between |
| the loss and the new proxy's start every round is target-only (zero drafts); the new |
| proxy starts with a new session, so the worker drops its windows and every slot is |
| re-announced (NEW_GEN). Spent: target-only rounds, one error log line |
| VLLM_DSPARK_REMOTE_ASYNC_RESPAWN_RESET_S a proxy process that served this long restores the full respawn budget |
| (default 3600; 0 = never) |
| VLLM_DSPARK_REMOTE_SWITCH_INTERVAL_US async thread mode: Python GIL switch interval while the proxy thread |
| runs (default 500; the interpreter default is 5000; 0 = unchanged); ignored with |
| ASYNC_PROXY=process |
| VLLM_DSPARK_REMOTE_FUSED_GATHER auto (default: on CUDA, after a startup self-check against the torch |
| path) | 0 | 1: the one-launch Triton row gather (fused_gather.py) |
| VLLM_DSPARK_REMOTE_N_OUT auto (default) | 0: ReqEntry.n_out (CONTRACT 6.3). auto = cap[n] of the dynamic |
| speculation schedule when the depth policy masks positions >= cap |
| (conf/confb/confr, no VLLM_DSPARK_DEPTH_FILE) in capacity/resilience; |
| otherwise, and in shadow, always all N positions |
| VLLM_DSPARK_REMOTE_LOG_EVERY rounds between stats log lines (default 1024, 0 = off) |
| VLLM_DSPARK_REMOTE_STATS_FILE JSON stats written by rank 0 (default |
| /tmp/dsv41-remote-dspark.json, inside the container) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import math |
| import os |
| from dataclasses import dataclass |
|
|
| from . import wire |
| from .transport import Endpoint, parse_endpoint |
|
|
| POLICIES = ("shadow", "resilience", "capacity") |
| ENV_ENDPOINT = "VLLM_DSPARK_REMOTE" |
|
|
|
|
| def remote_requested(environ=None) -> bool: |
| env = os.environ if environ is None else environ |
| return bool((env.get(ENV_ENDPOINT) or "").strip()) |
|
|
|
|
| def _int(env, key: str, default: int, lo: int = 0) -> int: |
| raw = env.get(key) |
| if raw is None or raw.strip() == "": |
| return default |
| try: |
| v = int(raw) |
| except ValueError as e: |
| raise ValueError(f"{key}={raw!r}: not an integer") from e |
| if v < lo: |
| raise ValueError(f"{key}={v}: must be >= {lo}") |
| return v |
|
|
|
|
| def _choice(env, key: str, default: str, choices: tuple[str, ...]) -> str: |
| v = (env.get(key) or default).strip().lower() |
| if v not in choices: |
| raise ValueError(f"{key}={v!r}: expected one of {choices}") |
| return v |
|
|
|
|
| def _bool(env, key: str, default: bool) -> bool: |
| raw = env.get(key) |
| if raw is None or raw.strip() == "": |
| return default |
| v = raw.strip().lower() |
| if v in ("1", "true", "on", "yes"): |
| return True |
| if v in ("0", "false", "off", "no"): |
| return False |
| raise ValueError(f"{key}={raw!r}: expected 0 or 1") |
|
|
|
|
| def _float_list(env, key: str) -> list[float] | None: |
| raw = env.get(key) |
| if raw is None or raw.strip() == "": |
| return None |
| try: |
| vals = [float(x) for x in raw.split(",") if x.strip()] |
| except ValueError as e: |
| raise ValueError(f"{key}={raw!r}: expected comma-separated numbers") from e |
| if not vals or not all(math.isfinite(v) for v in vals): |
| raise ValueError(f"{key}={raw!r}: expected finite numbers") |
| if any(v < 0 for v in vals): |
| raise ValueError(f"{key}={raw!r}: expected non-negative numbers") |
| return vals |
|
|
|
|
| @dataclass |
| class RemoteConfig: |
| endpoint: Endpoint |
| policy: str |
| drafter_identity: bytes |
| markov_identity: bytes |
| timeout_ms: int |
| connect_timeout_ms: int |
| backoff_ms: int |
| backoff_max_ms: int |
| on_fail: str |
| strict_start: bool |
| shadow_use: str |
| want_margin: bool |
| remote_draft_ms: float | None |
| remote_round_ms: list[float] | None |
| log_every: int |
| stats_file: str |
| async_mode: bool = False |
| switch_interval_us: int = 500 |
| hostflag_lib: str = "" |
| async_nap_us: int = 0 |
| async_spin_us: int = 1000 |
| async_watchdog_ms: int = 0 |
| async_proxy: str = "thread" |
| async_respawns: int = 3 |
| async_respawn_reset_s: int = 3600 |
| fused_gather: bool | None = None |
| n_out: str = "auto" |
| trip: int = 3 |
| cooldown: int = 256 |
| sampled_top32: bool = False |
|
|
| @property |
| def watchdog_ms(self) -> int: |
| """The async watchdog deadline actually used.""" |
| return self.async_watchdog_ms or max(4 * self.timeout_ms, 200) |
|
|
| @property |
| def keeps_local_drafter(self) -> bool: |
| return self.policy in ("shadow", "resilience") |
|
|
| @classmethod |
| def from_env(cls, environ=None) -> "RemoteConfig": |
| env = os.environ if environ is None else environ |
| if not remote_requested(env): |
| raise ValueError(f"{ENV_ENDPOINT} is not set") |
| endpoint = parse_endpoint(env[ENV_ENDPOINT]) |
| bind = (env.get("VLLM_DSPARK_REMOTE_BIND") or "").strip() |
| if bind: |
| if endpoint.options.get("bind", bind) != bind: |
| raise ValueError(f"VLLM_DSPARK_REMOTE_BIND={bind} conflicts with the endpoint's " |
| f"?bind={endpoint.options['bind']}") |
| endpoint.options["bind"] = bind |
| policy = _choice(env, "VLLM_DSPARK_REMOTE_POLICY", "shadow", POLICIES) |
| ident_raw = env.get("VLLM_DSPARK_REMOTE_IDENTITY", "") |
| if not ident_raw.strip(): |
| raise ValueError( |
| "VLLM_DSPARK_REMOTE_IDENTITY is required: the 64-hex drafter identity " |
| "(offload/dspark/spark/tools/identity.py over lane S2's DSPARK_COPY_DONE.json)") |
| drafter_identity = wire.parse_hex32(ident_raw, "VLLM_DSPARK_REMOTE_IDENTITY") |
| markov_raw = env.get("VLLM_DSPARK_REMOTE_MARKOV_IDENTITY", "") |
| if markov_raw.strip(): |
| markov_identity = wire.parse_hex32(markov_raw, "VLLM_DSPARK_REMOTE_MARKOV_IDENTITY") |
| else: |
| markov_identity = markov_identity_from_env(env) |
| draft_ms_list = _float_list(env, "VLLM_DSPARK_REMOTE_DRAFT_MS") |
| if draft_ms_list is not None and len(draft_ms_list) != 1: |
| raise ValueError("VLLM_DSPARK_REMOTE_DRAFT_MS: one number (ms)") |
| cfg = cls( |
| endpoint=endpoint, |
| policy=policy, |
| drafter_identity=drafter_identity, |
| markov_identity=markov_identity, |
| timeout_ms=_int(env, "VLLM_DSPARK_REMOTE_TIMEOUT_MS", 50, lo=1), |
| connect_timeout_ms=_int(env, "VLLM_DSPARK_REMOTE_CONNECT_TIMEOUT_MS", 1000, lo=1), |
| backoff_ms=_int(env, "VLLM_DSPARK_REMOTE_BACKOFF_MS", 1000, lo=1), |
| backoff_max_ms=_int(env, "VLLM_DSPARK_REMOTE_BACKOFF_MAX_MS", 30000, lo=1), |
| on_fail=_choice(env, "VLLM_DSPARK_REMOTE_ON_FAIL", "target_only", |
| ("target_only", "raise")), |
| strict_start=_bool(env, "VLLM_DSPARK_REMOTE_STRICT_START", False), |
| shadow_use=_choice(env, "VLLM_DSPARK_REMOTE_SHADOW_USE", "local", |
| ("local", "remote")), |
| want_margin=_bool(env, "VLLM_DSPARK_REMOTE_WANT_MARGIN", policy == "shadow"), |
| remote_draft_ms=None if draft_ms_list is None else draft_ms_list[0], |
| remote_round_ms=_float_list(env, "VLLM_DSPARK_REMOTE_ROUND_MS"), |
| log_every=_int(env, "VLLM_DSPARK_REMOTE_LOG_EVERY", 1024, lo=0), |
| stats_file=(env.get("VLLM_DSPARK_REMOTE_STATS_FILE") |
| or "/tmp/dsv41-remote-dspark.json"), |
| async_mode=_bool(env, "VLLM_DSPARK_REMOTE_ASYNC", False), |
| switch_interval_us=_int(env, "VLLM_DSPARK_REMOTE_SWITCH_INTERVAL_US", 500, lo=0), |
| hostflag_lib=(env.get("VLLM_DSPARK_REMOTE_HOSTFLAG_LIB") or "").strip(), |
| async_nap_us=_int(env, "VLLM_DSPARK_REMOTE_ASYNC_NAP_US", 0, lo=0), |
| async_spin_us=_int(env, "VLLM_DSPARK_REMOTE_ASYNC_SPIN_US", 1000, lo=0), |
| async_watchdog_ms=_int(env, "VLLM_DSPARK_REMOTE_ASYNC_WATCHDOG_MS", 0, lo=0), |
| async_proxy=_choice(env, "VLLM_DSPARK_REMOTE_ASYNC_PROXY", "thread", ("thread", "process")), |
| async_respawns=_int(env, "VLLM_DSPARK_REMOTE_ASYNC_RESPAWNS", 3, lo=0), |
| async_respawn_reset_s=_int(env, "VLLM_DSPARK_REMOTE_ASYNC_RESPAWN_RESET_S", 3600, lo=0), |
| fused_gather={"auto": None, "1": True, "0": False}[ |
| _choice(env, "VLLM_DSPARK_REMOTE_FUSED_GATHER", "auto", ("auto", "0", "1"))], |
| n_out=_choice(env, "VLLM_DSPARK_REMOTE_N_OUT", "auto", ("auto", "0")), |
| sampled_top32=_bool(env, "VLLM_DSPARK_SAMPLED_TOP32", False), |
| trip=_int(env, "VLLM_DSPARK_REMOTE_TRIP", 3, lo=0), |
| cooldown=_int(env, "VLLM_DSPARK_REMOTE_COOLDOWN", 256, lo=1), |
| ) |
| if cfg.async_mode and cfg.policy != "capacity": |
| raise ValueError("VLLM_DSPARK_REMOTE_ASYNC=1 needs VLLM_DSPARK_REMOTE_POLICY=capacity " |
| "(resilience/shadow decide on the host every round)") |
| if cfg.async_mode and cfg.on_fail == "raise": |
| raise ValueError("VLLM_DSPARK_REMOTE_ASYNC=1 cannot raise on a failed round " |
| "(the host never waits for it); use ON_FAIL=target_only") |
| if cfg.backoff_max_ms < cfg.backoff_ms: |
| raise ValueError("VLLM_DSPARK_REMOTE_BACKOFF_MAX_MS < VLLM_DSPARK_REMOTE_BACKOFF_MS") |
| return cfg |
|
|
|
|
| def markov_identity_from_env(env) -> bytes: |
| path = (env.get("VLLM_DSPARK_MARKOV_CACHE") or "").strip() |
| if not path: |
| return bytes(32) |
| try: |
| with open(path, "rb") as f: |
| return hashlib.sha256(f.read()).digest() |
| except OSError as e: |
| raise ValueError( |
| f"VLLM_DSPARK_MARKOV_CACHE={path!r} unreadable ({e}); set " |
| "VLLM_DSPARK_REMOTE_MARKOV_IDENTITY explicitly") from e |
|
|
|
|
| def apply_cost_overrides(cfg: RemoteConfig, environ=None, depth_policy_on: bool = False, |
| depth_mode: str | None = None, num_spec: int = 5, |
| depth_file: str | None = None) -> dict[str, str]: |
| """Make the remote round cost an explicit input of the depth policy (CONTRACT 10, |
| brief "recalibration"). Writes VLLM_DSPARK_DRAFT_MS / VLLM_DSPARK_ROUND_MS in the |
| process environment before the model runner builds DSparkDepthPolicy (it reads them |
| in initialize_kv_cache, after the speculator exists; it is built only when |
| VLLM_DSPARK_DEPTH_POLICY is on). Same env on both ranks, so the cost tables stay |
| identical across TP ranks. Returns what was changed; raises before writing anything. |
| |
| The policy prices depth d as ROUND_MS[d] + DRAFT_MS (depth_policy.round_cost_ms). The |
| live ROUND_MS includes the local drafter, so with the policy on, capacity/resilience |
| need both remote inputs: a remote DRAFT_MS over the live table counts the drafter |
| twice, a verify-only table under the live DRAFT_MS counts the wrong drafter. confr's |
| table is part of the policy string and cannot take them at all. fixed:<d> never prices |
| a depth, so it needs them only when a depth file can switch the mode at runtime.""" |
| env = os.environ if environ is None else environ |
| changed: dict[str, str] = {} |
| if cfg.policy == "shadow": |
| return changed |
| if depth_policy_on and not (depth_mode == "fixed" and not depth_file): |
| missing = [k for k, v in (("VLLM_DSPARK_REMOTE_DRAFT_MS", cfg.remote_draft_ms), |
| ("VLLM_DSPARK_REMOTE_ROUND_MS", cfg.remote_round_ms)) |
| if v is None] |
| if missing: |
| raise ValueError( |
| f"remote DSpark ({cfg.policy}) with VLLM_DSPARK_DEPTH_POLICY needs both " |
| "VLLM_DSPARK_REMOTE_DRAFT_MS (measured remote round, ms) and " |
| "VLLM_DSPARK_REMOTE_ROUND_MS (verify-only table, ms by depth 0..k): missing " |
| + " and ".join(missing) + ". The live VLLM_DSPARK_ROUND_MS includes the local " |
| "drafter, so a partial override mixes two cost models") |
| if depth_mode == "confr": |
| raise ValueError( |
| f"remote DSpark ({cfg.policy}): VLLM_DSPARK_DEPTH_POLICY=confr:<rows ms> carries " |
| "its own step table (measured with the local drafter), which the remote cost " |
| "inputs cannot replace; use conf or confb with VLLM_DSPARK_REMOTE_DRAFT_MS / " |
| "VLLM_DSPARK_REMOTE_ROUND_MS") |
| if cfg.remote_round_ms is not None: |
| t = cfg.remote_round_ms |
| if len(t) != num_spec + 1: |
| raise ValueError(f"VLLM_DSPARK_REMOTE_ROUND_MS: {len(t)} values, expected " |
| f"{num_spec + 1} (verify ms for depth 0..{num_spec})") |
| if any(v <= 0 for v in t): |
| raise ValueError("VLLM_DSPARK_REMOTE_ROUND_MS: verify times must be > 0 ms") |
| if any(b < a for a, b in zip(t, t[1:])): |
| raise ValueError("VLLM_DSPARK_REMOTE_ROUND_MS: must be non-decreasing with the " |
| f"depth (more live rows never verify faster): {t}") |
| changed["VLLM_DSPARK_ROUND_MS"] = ",".join(f"{v:g}" for v in t) |
| if cfg.remote_draft_ms is not None: |
| changed["VLLM_DSPARK_DRAFT_MS"] = f"{cfg.remote_draft_ms:g}" |
| for k, v in changed.items(): |
| env[k] = v |
| return changed |
|
|
|
|
| |
| |
| CAP_MASKING_DEPTH_MODES = ("conf", "confb", "confr") |
|
|
|
|
| def n_out_table(cfg: RemoteConfig, num_spec: int, depth_mode: str, depth_file: str | None, |
| dense_cap: list[int] | None) -> tuple[list[int] | None, str]: |
| """ReqEntry.n_out by request count (CONTRACT 6.3), or None (always all N positions), |
| with the reason. dense_cap is build_dynamic_sd_schedule_lookup's table (index = number |
| of requests, index 0 unused), as the model runner passes it to DSparkDepthPolicy. |
| |
| Positions >= cap[n] are only safe to drop when nothing reads them: the depth policy |
| must zero them (a cap-masking mode, fixed for the life of the serve, so no depth file |
| that could switch it to "fixed"/"off"), and shadow compares every position.""" |
| if cfg.policy == "shadow": |
| return None, "shadow compares every position" |
| if cfg.n_out == "0": |
| return None, "VLLM_DSPARK_REMOTE_N_OUT=0" |
| if depth_mode not in CAP_MASKING_DEPTH_MODES: |
| return None, f"depth policy {depth_mode!r} does not mask positions >= cap" |
| if depth_file: |
| return None, "VLLM_DSPARK_DEPTH_FILE can switch the depth policy at runtime" |
| if not dense_cap: |
| return None, "no dynamic speculation schedule" |
| table = [0 if i == 0 else min(num_spec, int(c)) for i, c in enumerate(dense_cap)] |
| if all(c >= num_spec for c in table[1:]): |
| return None, "schedule never caps below N" |
| return table, "cap by request count" |
|
|
|
|
