diff --git a/examples/configs/README.md b/examples/configs/README.md index 936efa226..22f0e8034 100644 --- a/examples/configs/README.md +++ b/examples/configs/README.md @@ -47,6 +47,13 @@ The `qwen3-8b-dflash-1server-dp7-disaggregated.yaml`, full-stack examples. Their typed `managed_local` blocks own Mooncake, one or two patched SGLang capture servers, and the trainer GPU allocation; the same `specforge train -c ...` command starts and cleans up each complete stack. +The `qwen3-4b-dspark-live.yaml` recipe is the online-live example: external +serving traffic drives capture (`deployment.disaggregated.live`). With +`live.mooncake` set, one `specforge train` command supervises Mooncake + +producer + consumer; `scripts/online_live/launch_capture_server.py` launches +the live-patched SGLang server with flags derived from the same config, and +`scripts/online_live/live_traffic_mimic.py` stands in for real users. + Disaggregated recipes without `managed_local` keep Mooncake and SGLang external for scheduler- or service-managed deployments. @@ -298,6 +305,11 @@ For `deployment.mode: disaggregated`, also write: | `deployment.disaggregated.producer_hold_s` | `null` | Optional positive offline producer retention timeout. Unset is unbounded; expiration fails the attempt. | | `deployment.disaggregated.shutdown_grace_s` | `30.0` | SIGTERM-to-SIGKILL window for a plain supervisor teardown; must cover worker cleanup (Mooncake drains, checkpoint flush, failure sentinels). `managed_local` stacks use `managed_local.shutdown_grace_s`. | | `deployment.disaggregated.managed_local` | `null` | Optional owned single-node Mooncake + capture-server stack described below. | +| `deployment.disaggregated.live` | `null` | Online-live mode: external serving-traffic capture servers push records to a producer-hosted intake endpoint. Mutually exclusive with `managed_local` and `server_urls`; requires `training.max_steps` (or `total_steps`) and no data source. | +| `deployment.disaggregated.live.host` | `0.0.0.0` | Intake bind host on the producer. | +| `deployment.disaggregated.live.port` | required | Intake bind port; capture servers point `--spec-capture-intake-url` at it. | +| `deployment.disaggregated.live.mooncake` | `null` | Optional supervisor-owned loopback Mooncake master (same shape as `managed_local.mooncake`): one `specforge train` command launches Mooncake + producer + consumer. Set together with `trainer_cuda_visible_devices`. | +| `deployment.disaggregated.live.trainer_cuda_visible_devices` | `null` | Consumer CUDA devices for the supervised live launch; count must equal `deployment.trainer.nproc_per_node`. | The four path fields have different ownership: diff --git a/examples/configs/qwen3-4b-dspark-disaggregated.yaml b/examples/configs/qwen3-4b-dspark-disaggregated.yaml index 962d9c41f..909859217 100644 --- a/examples/configs/qwen3-4b-dspark-disaggregated.yaml +++ b/examples/configs/qwen3-4b-dspark-disaggregated.yaml @@ -12,6 +12,9 @@ training: strategy: dspark num_epochs: 6 max_steps: 10000 + # Data-parallel over 7 ranks (global batch 7, accumulation 1). max_steps + # binds long before the epoch cap, so runtime = 10000 x per-step time; the + # extra ranks cut per-step time rather than adding sequential work. batch_size: 1 learning_rate: 0.0006 warmup_ratio: 0.04 @@ -20,8 +23,14 @@ training: loss_decay_gamma: 4.0 objective_chunk_blocks: 128 save_interval: 1000 + log_interval: 10 dist_timeout: 30 seed: 42 +tracking: + # WANDB_API_KEY is supplied via the process environment, never YAML. + report_to: wandb + wandb_project: specforge + wandb_name: qwen3-4b-dspark-disaggregated run_id: qwen3-4b-dspark-disaggregated output_dir: outputs/qwen3-4b-dspark-disaggregated @@ -29,12 +38,31 @@ deployment: mode: disaggregated trainer: nnodes: 1 - nproc_per_node: 1 + nproc_per_node: 7 disaggregated: control_dir: outputs/qwen3-4b-dspark-disaggregated/control backend: mooncake - server_urls: - - http://127.0.0.1:30000 - mooncake_metadata_server: http://127.0.0.1:35880/metadata - mooncake_master_server_addr: 127.0.0.1:35551 - mooncake_protocol: tcp + managed_local: + trainer_cuda_visible_devices: + - "1" + - "2" + - "3" + - "4" + - "5" + - "6" + - "7" + mooncake: + # Loopback TCP; RDMA loopback is broken in this KVM guest (transport + # retry counter exceeded even intra-HCA) and NVLink does not apply to + # the host-RAM feature buffer. + protocol: tcp + global_segment_size_bytes: 34359738368 + local_buffer_size_bytes: 1073741824 + capture_servers: + # A 4B target generates samples far faster than the trainer consumes + # them, so one capture GPU is enough; the other 7 GPUs go to training. + - port: 30000 + cuda_visible_devices: + - "0" + tp_size: 1 + mem_fraction_static: 0.7 diff --git a/examples/configs/qwen3-4b-dspark-live.yaml b/examples/configs/qwen3-4b-dspark-live.yaml new file mode 100644 index 000000000..d45e40729 --- /dev/null +++ b/examples/configs/qwen3-4b-dspark-live.yaml @@ -0,0 +1,72 @@ +# Online-live DSpark training from real serving traffic. +# +# Unlike the driven qwen3-4b-dspark-disaggregated recipe, the serving engine +# and its traffic stay external: users (or the mimic script) send ordinary +# /generate requests, the live-patched server captures prefill + generated +# tokens, and SpecForge trains on them. live.mooncake makes the launcher a +# supervisor that owns Mooncake master + producer + consumer, so the whole +# workflow is three commands (one terminal each): +# +# 1. specforge train --config examples/configs/qwen3-4b-dspark-live.yaml +# 2. python scripts/online_live/launch_capture_server.py \ +# --config examples/configs/qwen3-4b-dspark-live.yaml --cuda 0 +# 3. python scripts/online_live/live_traffic_mimic.py \ +# --server-url http://127.0.0.1:30000 \ +# --config examples/configs/qwen3-4b-dspark-live.yaml +model: + target_model_path: Qwen/Qwen3-4B + draft_model_config: configs/qwen3-4b-dspark.json + target_backend: sglang +data: + # No training data source: samples come from serving traffic. max_length is + # the intake token cap (prompt + generated tokens); longer captures are shed. + max_length: 3072 + chat_template: qwen + cache_dir: cache +training: + strategy: dspark + max_steps: 10000 + batch_size: 1 + learning_rate: 0.0006 + warmup_ratio: 0.04 + max_grad_norm: 1.0 + num_anchors: 512 + loss_decay_gamma: 4.0 + objective_chunk_blocks: 128 + save_interval: 1000 + log_interval: 10 + dist_timeout: 30 + seed: 42 +tracking: + # WANDB_API_KEY is supplied via the process environment, never YAML. + report_to: wandb + wandb_project: specforge + wandb_name: qwen3-4b-dspark-live +run_id: qwen3-4b-dspark-live +output_dir: outputs/qwen3-4b-dspark-live + +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 7 + disaggregated: + control_dir: outputs/qwen3-4b-dspark-live/control + backend: mooncake + live: + host: 0.0.0.0 + port: 8600 + # Supervisor-owned loopback Mooncake master; the pool size is the live + # capture buffer (write-until-full, trainer acks free space). + mooncake: + protocol: tcp + global_segment_size_bytes: 34359738368 + local_buffer_size_bytes: 1073741824 + trainer_cuda_visible_devices: + - "1" + - "2" + - "3" + - "4" + - "5" + - "6" + - "7" diff --git a/examples/configs/qwen3.6-27b-dspark-disaggregated.yaml b/examples/configs/qwen3.6-27b-dspark-disaggregated.yaml index 962e898ad..b8795f112 100644 --- a/examples/configs/qwen3.6-27b-dspark-disaggregated.yaml +++ b/examples/configs/qwen3.6-27b-dspark-disaggregated.yaml @@ -15,9 +15,15 @@ data: training: strategy: dspark num_epochs: 10 - # Portable one-rank equivalent of the source recipe's global batch 512. + # Global batch ~512 as in the source recipe: 6 trainer ranks x batch 1 x + # accumulation 85 = 510. Measured balance: one capture server sustains + # ~7 samples/s while 7 trainer ranks consumed ~23 samples/s of compute, + # so a GPU moved from the trainer to a second capture server. batch_size: 1 - accumulation_steps: 512 + accumulation_steps: 85 + # One optimizer step is ~minutes here; log every step so WandB shows + # progress immediately (default 50 would mean hours to the first point). + log_interval: 1 learning_rate: 0.0006 warmup_ratio: 0.04 max_grad_norm: 1.0 @@ -27,6 +33,17 @@ training: save_interval: 125 dist_timeout: 30 seed: 42 +runtime: + # The consumer dispatches one full optimizer window at a time + # (nnodes * nproc_per_node * batch_size * accumulation_steps = 510), + # so both watermarks must be at least that quantum. + in_flight_high_watermark: 576 + in_flight_low_watermark: 512 +tracking: + # WANDB_API_KEY is supplied via the process environment, never YAML. + report_to: wandb + wandb_project: specforge + wandb_name: qwen3.6-27b-dspark-disaggregated run_id: qwen3.6-27b-dspark-disaggregated output_dir: outputs/qwen3.6-27b-dspark-disaggregated @@ -34,14 +51,24 @@ deployment: mode: disaggregated trainer: nnodes: 1 - nproc_per_node: 1 + nproc_per_node: 6 disaggregated: control_dir: outputs/qwen3.6-27b-dspark-disaggregated/control backend: mooncake managed_local: trainer_cuda_visible_devices: - - "1" + - "2" + - "3" + - "4" + - "5" + - "6" + - "7" mooncake: + # RDMA was tried here (mlx5_4/mlx5_5) but this KVM guest's IB fabric + # rejects same-host loopback: every slice dies with "transport retry + # counter exceeded", even intra-HCA. Loopback TCP (pinned via + # MC_TCP_BIND_ADDRESS in launch_plan) is the working transport for + # single-node runs; use RDMA only for true multi-node deployments. protocol: tcp global_segment_size_bytes: 68719476736 local_buffer_size_bytes: 1073741824 @@ -51,3 +78,8 @@ deployment: - "0" tp_size: 1 mem_fraction_static: 0.7 + - port: 30001 + cuda_visible_devices: + - "1" + tp_size: 1 + mem_fraction_static: 0.7 diff --git a/patches/sglang/online-live/spec-capture-live.patch b/patches/sglang/online-live/spec-capture-live.patch new file mode 100644 index 000000000..7cd17d428 --- /dev/null +++ b/patches/sglang/online-live/spec-capture-live.patch @@ -0,0 +1,394 @@ +--- a/python/sglang/srt/server_args.py ++++ b/python/sglang/srt/server_args.py +@@ -2147,6 +2147,19 @@ + "match the draft strategy being trained; they wire capture onto " + "different submodules (VL targets only populate the dflash path).", + ] = "eagle3" ++ spec_capture_intake_url: A[ ++ Optional[str], ++ "Online-live capture: SpecForge intake origin (http://host:port). The " ++ "sink fetches its capture config from this endpoint, captures ordinary " ++ "serving requests (prefill + generated tokens), writes tensors into " ++ "Mooncake, and pushes tensor-free records back. Requires " ++ "--enable-spec-capture.", ++ ] = None ++ spec_capture_sample_rate: A[ ++ float, ++ "Fraction of live serving requests to capture (deterministic per " ++ "request id). Only used with --spec-capture-intake-url.", ++ ] = 1.0 + enable_return_routed_experts: A[ + bool, + "Enable returning routed experts of each layer with responses.", +--- a/python/sglang/srt/spec_capture_sink.py ++++ b/python/sglang/srt/spec_capture_sink.py +@@ -26,14 +26,30 @@ + + Mooncake connection uses the standard ``MOONCAKE_*`` env vars (see + ``MooncakeFeatureStore``). ++ ++Online-live mode (``--spec-capture-intake-url``): ordinary serving requests ++are captured without any client ``spec_capture`` field. The sink GETs its ++capture config (store id, feature names, passthrough synthesis rules, token ++cap) from the SpecForge intake, mints a deterministic spec per sampled ++request, and — on a background writer thread, never the scheduler loop — ++writes prefill+decode rows to Mooncake and POSTs the tensor-free record back. ++Every failure (queue full, Mooncake down/full, intake down or non-2xx) drops ++just that capture, removes any keys already written, and never touches the ++user's response. + """ + + from __future__ import annotations + ++import hashlib ++import json + import logging + import os ++import queue + import threading ++import time + from typing import Any, Dict, List, Optional ++from urllib.error import HTTPError ++from urllib.request import Request, urlopen + + import torch + +@@ -66,7 +82,13 @@ + class SpecCaptureSink: + """Writes captured per-request tensors into Mooncake in SpecForge layout.""" + +- def __init__(self, aux_layer_ids: Optional[List[int]] = None) -> None: ++ def __init__( ++ self, ++ aux_layer_ids: Optional[List[int]] = None, ++ *, ++ intake_url: Optional[str] = None, ++ sample_rate: float = 1.0, ++ ) -> None: + self.aux_layer_ids = list(aux_layer_ids) if aux_layer_ids else None + self._store = None + self._put_config = None +@@ -74,6 +96,26 @@ + # Retried HTTP requests reuse deterministic keys. Striped locks keep + # replacement atomic per key without retaining one lock per sample. + self._write_locks = [threading.Lock() for _ in range(256)] ++ self.intake_url = intake_url.rstrip("/") if intake_url else None ++ self.sample_rate = float(sample_rate) ++ self._live_config: Optional[Dict[str, Any]] = None ++ self._live_config_next_fetch = 0.0 ++ self._live_stats = { ++ "captured": 0, ++ "dropped_queue_full": 0, ++ "dropped_put": 0, ++ "dropped_post": 0, ++ "dropped_bad_rows": 0, ++ } ++ self._live_stats_last_log = 0.0 ++ self._live_queue: Optional[queue.Queue] = None ++ if self.intake_url: ++ self._live_queue = queue.Queue( ++ maxsize=int(os.environ.get("SPEC_CAPTURE_LIVE_QUEUE", 64)) ++ ) ++ threading.Thread( ++ target=self._live_writer, name="spec-capture-live", daemon=True ++ ).start() + + # -- connection --------------------------------------------------------- + def _connect(self): +@@ -235,6 +277,146 @@ + "features": result_feats, + } + ++ # -- online-live mode ----------------------------------------------------- ++ def _fetch_live_config(self) -> Optional[Dict[str, Any]]: ++ """Lazy capture-config handshake; re-attempted at most every 5s.""" ++ if self._live_config is not None: ++ return self._live_config ++ now = time.monotonic() ++ if now < self._live_config_next_fetch: ++ return None ++ self._live_config_next_fetch = now + 5.0 ++ try: ++ with urlopen( ++ f"{self.intake_url}/v1/spec-capture/config", timeout=5.0 ++ ) as response: ++ self._live_config = json.load(response) ++ logger.info("spec-capture live config: %s", self._live_config) ++ except Exception as e: ++ logger.warning("spec-capture live config fetch failed: %s", e) ++ return self._live_config ++ ++ def maybe_live_spec(self, rid: str) -> Optional[Dict[str, Any]]: ++ """Server-minted capture spec for one serving request (None = skip). ++ ++ Both the sampling decision and the sample id derive deterministically ++ from ``rid`` so every TP rank makes the same capture decision. ++ """ ++ if self.intake_url is None: ++ return None ++ config = self._fetch_live_config() ++ if config is None: ++ return None ++ if self.sample_rate < 1.0: ++ digest = hashlib.sha1(str(rid).encode("utf-8")).digest() ++ if int.from_bytes(digest[:8], "big") / 2**64 >= self.sample_rate: ++ return None ++ return { ++ "store_id": config["store_id"], ++ "sample_id": f"live-{rid}", ++ "gen": int(config.get("gen", 1)), ++ "replace": False, ++ "features": dict(config.get("features") or {}), ++ "live": True, ++ } ++ ++ def put_sample_live( ++ self, ++ spec: Dict[str, Any], ++ *, ++ aux: Optional[torch.Tensor], ++ last_hidden: Optional[torch.Tensor], ++ tokens: List[int], ++ ) -> None: ++ """Fire-and-forget: hand one finished capture to the writer thread.""" ++ try: ++ self._live_queue.put_nowait((spec, aux, last_hidden, list(tokens))) ++ except queue.Full: ++ self._live_stats["dropped_queue_full"] += 1 ++ self._log_live_stats() ++ ++ def _live_writer(self) -> None: ++ while True: ++ item = self._live_queue.get() ++ try: ++ self._write_live_sample(*item) ++ except Exception as e: ++ logger.warning("spec-capture live write failed: %s", e) ++ finally: ++ self._live_queue.task_done() ++ ++ def _write_live_sample(self, spec, aux, last_hidden, tokens) -> None: ++ config = self._live_config or {} ++ min_tokens = int(config.get("min_num_tokens") or 2) ++ max_tokens = int(config.get("max_num_tokens") or 0) ++ rows = aux.shape[0] if aux is not None else -1 ++ if ( ++ rows != len(tokens) ++ or len(tokens) < min_tokens # warmup/probe one-token requests ++ or (max_tokens and len(tokens) > max_tokens) ++ ): ++ self._live_stats["dropped_bad_rows"] += 1 ++ self._log_live_stats() ++ return ++ passthrough = [] ++ for item in config.get("passthrough") or []: ++ data = tokens if item["source"] == "tokens" else [1] * len(tokens) ++ passthrough.append( ++ { ++ "name": item["name"], ++ "data": data, ++ "shape": [1, len(tokens)], ++ "dtype": "int64", ++ } ++ ) ++ spec = dict(spec, passthrough=passthrough) ++ try: ++ result = self.put_sample(spec, aux=aux, last_hidden=last_hidden) ++ except Exception as e: ++ # Pool full or Mooncake down: partial keys were already removed. ++ logger.warning("spec-capture live put failed for %s: %s", spec["sample_id"], e) ++ self._live_stats["dropped_put"] += 1 ++ self._log_live_stats() ++ return ++ result["num_tokens"] = len(tokens) ++ if self._post_live_record(result): ++ self._live_stats["captured"] += 1 ++ else: ++ # A record SpecForge never accepted must not stay hard-pinned. ++ for name in result["features"]: ++ self._remove_quiet( ++ self._tkey( ++ spec["store_id"], spec["sample_id"], int(spec["gen"]), name ++ ) ++ ) ++ self._live_stats["dropped_post"] += 1 ++ self._log_live_stats() ++ ++ def _post_live_record(self, result: Dict[str, Any]) -> bool: ++ body = json.dumps(result).encode("utf-8") ++ for _attempt in range(2): # the intake dedups sample_id; retry is safe ++ request = Request( ++ f"{self.intake_url}/v1/spec-capture/records", ++ data=body, ++ headers={"Content-Type": "application/json"}, ++ method="POST", ++ ) ++ try: ++ with urlopen(request, timeout=10.0) as response: ++ return 200 <= response.status < 300 ++ except HTTPError: ++ return False # shed/rejected: no point retrying ++ except Exception: ++ continue ++ return False ++ ++ def _log_live_stats(self) -> None: ++ now = time.monotonic() ++ if now - self._live_stats_last_log < 60.0: ++ return ++ self._live_stats_last_log = now ++ logger.info("spec-capture live stats: %s", self._live_stats) ++ + + _SINK: Optional[SpecCaptureSink] = None + +@@ -248,9 +430,17 @@ + global _SINK + if getattr(server_args, "enable_spec_capture", False) and _SINK is None: + _SINK = SpecCaptureSink( +- aux_layer_ids=getattr(server_args, "spec_capture_aux_layer_ids", None) ++ aux_layer_ids=getattr(server_args, "spec_capture_aux_layer_ids", None), ++ intake_url=getattr(server_args, "spec_capture_intake_url", None), ++ sample_rate=getattr(server_args, "spec_capture_sample_rate", 1.0), + ) + + + def get_sink() -> Optional[SpecCaptureSink]: + return _SINK ++ ++ ++def maybe_live_spec(rid: str) -> Optional[Dict[str, Any]]: ++ """Scheduler one-liner: a live capture spec for ``rid``, or None.""" ++ sink = _SINK ++ return sink.maybe_live_spec(rid) if sink is not None else None +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -153,6 +153,7 @@ + decide_needs_cpu_seq_lens, + resolve_forward_inputs, + ) ++from sglang.srt import spec_capture_sink + from sglang.srt.managers.prefill_delayer import ( + PrefillDelayer, + PrefillDelayerSinglePassExecutor, +@@ -585,6 +586,14 @@ + "(single-pass prefill) so captured hidden states cover the " + "whole sequence" + ) ++ if ( ++ server_args.spec_capture_intake_url ++ and server_args.speculative_algorithm is not None ++ ): ++ raise ValueError( ++ "--spec-capture-intake-url (online-live capture) does not " ++ "support serving with speculative decoding yet" ++ ) + from sglang.srt import spec_capture_sink + + spec_capture_sink.maybe_init_sink(server_args) +@@ -2067,7 +2076,10 @@ + dllm_config=self.dllm_config, + time_stats=recv_req.time_stats, + multi_item_delimiter_indices=recv_req.multi_item_delimiter_indices, +- spec_capture=recv_req.spec_capture, ++ # An explicit client spec (driven capture) wins; otherwise the ++ # sink may mint one for this live serving request. ++ spec_capture=recv_req.spec_capture ++ or spec_capture_sink.maybe_live_spec(recv_req.rid), + ) + req.tokenizer = self.tokenizer + +--- a/python/sglang/srt/managers/schedule_batch.py ++++ b/python/sglang/srt/managers/schedule_batch.py +@@ -1489,6 +1489,16 @@ + if self.input_embeds is not None: + self.output_ids = array("q") + ++ if self.spec_capture is not None and self.spec_capture.get("live"): ++ # The re-prefill would double-append captured rows; drop this ++ # request's live capture instead. return_hidden_states was forced ++ # on by the live spec, so reset it or the re-prefill would stream ++ # hidden states into the user's response. ++ self.spec_capture = None ++ self.spec_capture_aux = [] ++ self.spec_capture_last_hidden = [] ++ self.return_hidden_states = False ++ + def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator): + token_indices = req_to_token_pool.req_to_token[ + self.req_pool_idx, : self.seqlen - 1 +--- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py ++++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +@@ -498,6 +498,23 @@ + ) + return end + ++ def _append_spec_capture_decode_states( ++ self, ++ *, ++ req: Req, ++ logits_output: LogitsProcessorOutput, ++ start: int, ++ count: int, ++ ) -> None: ++ """Accumulate captured decode rows (aligned with the fed tokens).""" ++ req.spec_capture_aux.append( ++ logits_output.hidden_states[start : start + count].cpu().clone() ++ ) ++ if logits_output.last_hidden_states is not None: ++ req.spec_capture_last_hidden.append( ++ logits_output.last_hidden_states[start : start + count].cpu().clone() ++ ) ++ + def _sink_spec_capture(self, req: Req) -> None: + """Write a finished capture request's tensors to the Mooncake sink. + +@@ -505,6 +522,9 @@ + (or an ``{"error": ...}`` marker) is set on ``req.spec_capture_result``, + returned to the client via the dedicated ``spec_capture`` output field + (a per-request channel, unlike per-token ``customized_info``). ++ ++ Live capture instead hands off to the sink's background writer: the ++ record goes to the SpecForge intake, nothing enters the user response. + """ + from sglang.srt import spec_capture_sink + +@@ -517,6 +537,16 @@ + if req.spec_capture_last_hidden + else None + ) ++ if req.spec_capture.get("live"): ++ # Rows cover origin_input_ids plus every fed decode token; the ++ # final sampled token has no hidden row and is dropped. ++ tokens = list(req.origin_input_ids) + list(req.output_ids)[:-1] ++ sink.put_sample_live( ++ req.spec_capture, aux=aux, last_hidden=last_hidden, tokens=tokens ++ ) ++ req.spec_capture_aux = [] ++ req.spec_capture_last_hidden = [] ++ return + try: + req.spec_capture_result = sink.put_sample( + req.spec_capture, aux=aux, last_hidden=last_hidden +@@ -781,7 +811,20 @@ + logits_output=logits_output, + ) + +- if req.return_hidden_states and logits_output.hidden_states is not None: ++ if req.spec_capture is not None and logits_output.hidden_states is not None: ++ # Spec-training capture: one row per fed token (the previously ++ # sampled one), covering prompt + generated tokens by finish. ++ if not is_spec: ++ start, count = i, 1 ++ else: ++ start = i * result.speculative_num_draft_tokens ++ count = len(next_token_id) ++ self._append_spec_capture_decode_states( ++ req=req, logits_output=logits_output, start=start, count=count ++ ) ++ if req.finished(): ++ self._sink_spec_capture(req) ++ elif req.return_hidden_states and logits_output.hidden_states is not None: + if not is_spec: + req.hidden_states.append( + logits_output.hidden_states[i].cpu().clone().tolist() diff --git a/patches/sglang/v0.5.14/spec-capture.patch b/patches/sglang/v0.5.14/spec-capture.patch index a1081ba3d..1f946e5df 100644 --- a/patches/sglang/v0.5.14/spec-capture.patch +++ b/patches/sglang/v0.5.14/spec-capture.patch @@ -302,7 +302,7 @@ index 1cff5c983..a5935fa3c 100644 + # Aux capture without a draft worker, routed to the strategy's own + # capture method (they wire different submodules — e.g. VL models + # only populate layers_to_capture via the dflash path). -+ if getattr(server_args, "spec_capture_method", "eagle3") == "dflash": ++ if getattr(server_args, "spec_capture_method", "eagle3") in ("dflash", "dspark"): + self.dflash_use_aux_hidden_state = True + self.dflash_target_layer_ids = server_args.spec_capture_aux_layer_ids + if hasattr(self, "spec_aux_config"): @@ -332,7 +332,7 @@ diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index c7162c16d..d515b3c69 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py -@@ -2127,6 +2127,25 @@ class ServerArgs: +@@ -2127,6 +2127,26 @@ class ServerArgs: bool, "Enable returning hidden states with responses.", ] = False @@ -351,7 +351,8 @@ index c7162c16d..d515b3c69 100644 + ] = None + spec_capture_method: A[ + str, -+ "Capture method for --enable-spec-capture: 'eagle3' or 'dflash'. Must " ++ "Capture method for --enable-spec-capture: 'eagle3', 'dflash', or " ++ "'dspark' (dspark reuses the dflash aux wiring on stock SGLang). Must " + "match the draft strategy being trained; they wire capture onto " + "different submodules (VL targets only populate the dflash path).", + ] = "eagle3" diff --git a/scripts/apply_sglang_spec_capture_patch.sh b/scripts/apply_sglang_spec_capture_patch.sh index 1f2594d6b..c19b60997 100755 --- a/scripts/apply_sglang_spec_capture_patch.sh +++ b/scripts/apply_sglang_spec_capture_patch.sh @@ -12,15 +12,20 @@ # when a reverse dry-run proves it matches the current patch byte-for-byte; # anything else fails loudly rather than testing against unknown server code. # +# The optional --live layer (patches/sglang/online-live/) adds capture of +# real serving traffic on top of the base v0.5.14 patch. It has its own +# applied-copy record and is applied after / reversed before the base patch. +# # Usage: scripts/apply_sglang_spec_capture_patch.sh # [--target v0.5.14|kimi-k3-ee560a2|kimi-k3-9acd9cb|kimi-k3-f8493a4] -# [--reverse] +# [--live] [--reverse] set -euo pipefail HERE="$(cd "$(dirname "$0")/.." && pwd)" TARGET="v0.5.14" PATCH_TARGET="" REVERSE=0 +LIVE=0 while [[ $# -gt 0 ]]; do case "$1" in --target) @@ -31,6 +36,10 @@ while [[ $# -gt 0 ]]; do TARGET="$2" shift 2 ;; + --live) + LIVE=1 + shift + ;; --reverse) REVERSE=1 shift @@ -42,6 +51,11 @@ while [[ $# -gt 0 ]]; do esac done +if [[ "$LIVE" == 1 && "$TARGET" != "v0.5.14" ]]; then + echo "ERROR: --live is only supported for the v0.5.14 target" >&2 + exit 2 +fi + case "$TARGET" in v0.5.14) EXPECTED_VERSION_PREFIX="0.5.14" @@ -69,24 +83,59 @@ SGL_PARENT="$(python -c 'import sglang, os; print(os.path.dirname(os.path.dirnam SGL_VERSION="$(python -c 'import sglang; print(sglang.__version__)')" APPLIED_COPY="$SGL_PARENT/sglang/.spec_capture_patch.applied" SINK="$SGL_PARENT/sglang/srt/spec_capture_sink.py" +LIVE_PATCH="$HERE/patches/sglang/online-live/spec-capture-live.patch" +LIVE_APPLIED_COPY="$SGL_PARENT/sglang/.spec_capture_live_patch.applied" if [[ -n "$EXPECTED_VERSION_PREFIX" && "$SGL_VERSION" != "$EXPECTED_VERSION_PREFIX"* ]]; then echo "WARNING: installed sglang is $SGL_VERSION; the patch targets $TARGET" >&2 fi +reverse_live_if_applied() { + if [[ -f "$LIVE_APPLIED_COPY" ]]; then + patch --reverse -p2 --batch -d "$SGL_PARENT" < "$LIVE_APPLIED_COPY" + rm -f "$LIVE_APPLIED_COPY" + echo "online-live patch --reverse at $SGL_PARENT/sglang" + fi +} + if [[ "$REVERSE" == 1 ]]; then + # The live layer sits on top of the base patch: reverse it first. + reverse_live_if_applied + if [[ "$LIVE" == 1 ]]; then + exit 0 + fi patch --reverse -p2 --batch -N -d "$SGL_PARENT" < "$PATCH" rm -f "$APPLIED_COPY" echo "spec-capture patch $TARGET --reverse at $SGL_PARENT/sglang (sglang $SGL_VERSION)" exit 0 fi +apply_live_if_requested() { + if [[ "$LIVE" != 1 ]]; then + return 0 + fi + if [[ -f "$LIVE_APPLIED_COPY" ]]; then + if cmp -s "$LIVE_APPLIED_COPY" "$LIVE_PATCH"; then + echo "online-live patch already applied at $SGL_PARENT/sglang" + return 0 + fi + echo "online-live patch changed; reversing the recorded version first" + patch --reverse -p2 --batch -d "$SGL_PARENT" < "$LIVE_APPLIED_COPY" + fi + patch -p2 --batch -N -d "$SGL_PARENT" < "$LIVE_PATCH" + cp "$LIVE_PATCH" "$LIVE_APPLIED_COPY" + echo "online-live patch applied at $SGL_PARENT/sglang" +} + if [[ -f "$APPLIED_COPY" ]]; then if cmp -s "$APPLIED_COPY" "$PATCH"; then echo "spec-capture patch $TARGET already applied at $SGL_PARENT/sglang" + apply_live_if_requested exit 0 fi echo "spec-capture patch changed; reversing the recorded version first" + # The live layer sits on top of the base patch: it must come off first. + reverse_live_if_applied patch --reverse -p2 --batch -d "$SGL_PARENT" < "$APPLIED_COPY" elif [[ -f "$SINK" ]]; then # Patched before the applied-copy record existed. Adopt only a tree that @@ -110,3 +159,4 @@ fi patch -p2 --batch -N -d "$SGL_PARENT" < "$PATCH" cp "$PATCH" "$APPLIED_COPY" echo "spec-capture patch $TARGET applied at $SGL_PARENT/sglang (sglang $SGL_VERSION)" +apply_live_if_requested diff --git a/scripts/online_live/launch_capture_server.py b/scripts/online_live/launch_capture_server.py new file mode 100755 index 000000000..d63ecec8c --- /dev/null +++ b/scripts/online_live/launch_capture_server.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +# coding=utf-8 +"""One-command live capture server: resolve flags from the run config and exec. + +Reads an online-live SpecForge config, resolves the capture contract (method, +aux layer ids, context length) exactly like the producer does, applies the +online-live SGLang patch if needed, and execs ``sglang.launch_server`` with +the derived Mooncake environment and capture flags. + +Example: + python scripts/online_live/launch_capture_server.py \\ + --config examples/configs/qwen3-4b-dspark-live.yaml --cuda 0 + +Extra SGLang flags pass through after ``--``: + ... --cuda 0 -- --mem-fraction-static 0.8 +""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + + +def parse_args() -> tuple[argparse.Namespace, list[str]]: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", required=True, help="online-live run YAML") + parser.add_argument("--port", type=int, default=30000) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--cuda", default=None, help="CUDA_VISIBLE_DEVICES value") + parser.add_argument("--tp-size", type=int, default=1) + parser.add_argument( + "--intake-url", default=None, help="default: http://127.0.0.1:" + ) + parser.add_argument("--sample-rate", type=float, default=1.0) + parser.add_argument( + "--skip-patch", action="store_true", help="do not (re)apply the live patch" + ) + return parser.parse_known_args() + + +def main() -> None: + from specforge.application import resolve_capture_contract + from specforge.config import SGLANG_CAPTURE_CONTEXT_HEADROOM, Config + from specforge.launch_plan import _sglang_argv + + args, extra = parse_args() + if extra and extra[0] == "--": + extra = extra[1:] + cfg = Config.from_file(args.config) + deployment = cfg.deployment.disaggregated + live = deployment.live if deployment is not None else None + if live is None: + raise SystemExit(f"{args.config} is not an online-live config") + contract = resolve_capture_contract(cfg) + + if not args.skip_patch: + subprocess.check_call( + [str(REPO_ROOT / "scripts/apply_sglang_spec_capture_patch.sh"), "--live"] + ) + + env = dict(os.environ) + if live.mooncake is not None: + mooncake = live.mooncake + env.setdefault( + "MOONCAKE_METADATA_SERVER", + f"http://127.0.0.1:{mooncake.metadata_port}/metadata", + ) + env.setdefault( + "MOONCAKE_MASTER_SERVER_ADDR", f"127.0.0.1:{mooncake.rpc_port}" + ) + env.setdefault("MOONCAKE_LOCAL_HOSTNAME", mooncake.local_hostname) + env.setdefault("MOONCAKE_PROTOCOL", mooncake.protocol) + env.setdefault( + "MOONCAKE_GLOBAL_SEGMENT_SIZE", str(mooncake.global_segment_size_bytes) + ) + env.setdefault( + "MOONCAKE_LOCAL_BUFFER_SIZE", str(mooncake.local_buffer_size_bytes) + ) + env.setdefault("MC_TRANSFER_TIMEOUT", "300") + env.setdefault("MC_TCP_BIND_ADDRESS", mooncake.local_hostname) + else: + for name, value in ( + ("MOONCAKE_METADATA_SERVER", deployment.mooncake_metadata_server), + ("MOONCAKE_MASTER_SERVER_ADDR", deployment.mooncake_master_server_addr), + ): + if value: + env.setdefault(name, value) + if args.cuda is not None: + env["CUDA_VISIBLE_DEVICES"] = args.cuda + env.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1") + + intake_url = args.intake_url or f"http://127.0.0.1:{live.port}" + context_length = cfg.model.sglang_context_length or ( + cfg.data.max_length + SGLANG_CAPTURE_CONTEXT_HEADROOM + ) + argv = [ + sys.executable, + "-m", + "sglang.launch_server", + "--model-path", + cfg.model.target_model_path, + "--dtype", + cfg.model.torch_dtype, + ] + if cfg.model.trust_remote_code: + argv.append("--trust-remote-code") + if cfg.model.cache_dir: + argv.extend(("--download-dir", cfg.model.cache_dir)) + argv.extend( + [ + "--tp-size", + str(args.tp_size), + "--chunked-prefill-size", + "-1", + "--enable-spec-capture", + "--spec-capture-method", + contract.method, + "--spec-capture-aux-layer-ids", + *[str(layer) for layer in contract.aux_layer_ids], + "--spec-capture-intake-url", + intake_url, + "--spec-capture-sample-rate", + str(args.sample_rate), + "--host", + args.host, + "--port", + str(args.port), + ] + ) + argv.extend(_sglang_argv(cfg.model, overrides={"sglang_context_length": context_length})) + argv.extend(extra) + print(f"exec: {' '.join(argv)}", flush=True) + os.execvpe(argv[0], argv, env) + + +if __name__ == "__main__": + main() diff --git a/scripts/online_live/live_traffic_mimic.py b/scripts/online_live/live_traffic_mimic.py new file mode 100755 index 000000000..bd71f8a9c --- /dev/null +++ b/scripts/online_live/live_traffic_mimic.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python3 +# coding=utf-8 +"""Mimic user traffic against a live-capture SGLang server. + +Reads ShareGPT-style conversations (``{"id", "conversations": [{"role", +"content"}, ...]}`` JSONL — the same file the driven disaggregated recipes use +as ``data.train_data_path``), builds a chat prompt up to the last assistant +turn, and POSTs plain ``/generate`` requests exactly like a real user: no +``spec_capture`` field, sampled decoding. With the server launched with +``--spec-capture-intake-url``, every request feeds drafter training. + +Example: + python scripts/online_live/live_traffic_mimic.py \\ + --server-url http://127.0.0.1:30000 \\ + --config examples/configs/qwen3-4b-dspark-live.yaml +""" + +from __future__ import annotations + +import argparse +import itertools +import json +import random +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from urllib.request import Request, urlopen + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--server-url", required=True) + parser.add_argument( + "--config", default=None, help="run YAML; supplies the tokenizer default" + ) + parser.add_argument( + "--data", + default="./cache/dataset/sharegpt_train.jsonl", + help="ShareGPT JSONL path", + ) + parser.add_argument( + "--tokenizer", default=None, help="HF tokenizer id/path (or use --config)" + ) + parser.add_argument("--concurrency", type=int, default=8) + parser.add_argument("--qps", type=float, default=0.0, help="0 = unthrottled") + parser.add_argument("--max-requests", type=int, default=0, help="0 = all") + # Matches data.max_length in the disagg recipes (the intake token cap). + parser.add_argument("--max-prompt-tokens", type=int, default=3072) + # Temporarily 1 so live runs compare directly against the driven + # disaggregated recipe (prefill-only capture). + parser.add_argument("--max-new-tokens", type=int, default=1) + parser.add_argument("--temperature", type=float, default=0.7) + parser.add_argument("--timeout-s", type=float, default=300.0) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--loop", action="store_true", help="repeat the dataset") + parser.add_argument("--log-interval-s", type=float, default=10.0) + return parser.parse_args() + + +def iter_prompts(args, tokenizer): + """Yield token-id prompts: the full conversation through the last assistant turn. + + Matches the driven disagg-online producer, which prefills whole + conversations (assistant replies included) with max_new_tokens=1. + """ + rng = random.Random(args.seed) + with open(args.data, encoding="utf-8") as stream: + records = [json.loads(line) for line in stream if line.strip()] + rng.shuffle(records) + for record in itertools.cycle(records) if args.loop else records: + turns = [ + {"role": turn["role"], "content": turn["content"]} + for turn in record.get("conversations", []) + if turn.get("role") in ("system", "user", "assistant") + ] + while turns and turns[-1]["role"] != "assistant": + turns.pop() + if not turns: + continue + text = tokenizer.apply_chat_template( + turns, add_generation_prompt=False, tokenize=False + ) + input_ids = tokenizer(text, add_special_tokens=False).input_ids + if len(input_ids) > args.max_prompt_tokens: + continue + yield input_ids + + +def main() -> None: + from transformers import AutoTokenizer + + args = parse_args() + if args.tokenizer is None: + if args.config is None: + raise SystemExit("pass --tokenizer or --config") + import yaml + + with open(args.config, encoding="utf-8") as stream: + args.tokenizer = yaml.safe_load(stream)["model"]["target_model_path"] + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) + url = f"{args.server_url.rstrip('/')}/generate" + stats = {"sent": 0, "ok": 0, "error": 0} + lock = threading.Lock() + start = time.monotonic() + + def send(input_ids) -> None: + try: + body = json.dumps( + { + "input_ids": list(input_ids), + "sampling_params": { + "temperature": args.temperature, + "max_new_tokens": args.max_new_tokens, + }, + } + ).encode("utf-8") + request = Request( + url, data=body, headers={"Content-Type": "application/json"} + ) + with urlopen(request, timeout=args.timeout_s): + outcome = "ok" + except Exception as exc: # noqa: BLE001 — a user request may just fail + outcome = "error" + print(f"request failed: {exc}", flush=True) + with lock: + stats[outcome] += 1 + + last_log = time.monotonic() + with ThreadPoolExecutor(max_workers=args.concurrency) as executor: + for input_ids in iter_prompts(args, tokenizer): + if args.max_requests and stats["sent"] >= args.max_requests: + break + if args.qps > 0: + target = start + stats["sent"] / args.qps + delay = target - time.monotonic() + if delay > 0: + time.sleep(delay) + executor.submit(send, input_ids) + stats["sent"] += 1 + now = time.monotonic() + if now - last_log >= args.log_interval_s: + rate = stats["sent"] / (now - start) + print(f"{stats} ({rate:.2f} req/s)", flush=True) + last_log = now + print(f"done: {stats} in {time.monotonic() - start:.1f}s", flush=True) + + +if __name__ == "__main__": + main() diff --git a/specforge/application/__init__.py b/specforge/application/__init__.py index 4fce78ead..153829460 100644 --- a/specforge/application/__init__.py +++ b/specforge/application/__init__.py @@ -5,6 +5,7 @@ ResolvedRun, bind_run, build_application_run, + resolve_capture_contract, resolve_offline_capture, resolve_run, ) @@ -14,6 +15,7 @@ "ResolvedRun", "bind_run", "build_application_run", + "resolve_capture_contract", "resolve_offline_capture", "resolve_run", ] diff --git a/specforge/application/composition.py b/specforge/application/composition.py index 0eb271c17..1ffb30f86 100644 --- a/specforge/application/composition.py +++ b/specforge/application/composition.py @@ -148,11 +148,25 @@ def build_application_run( ) +def resolve_capture_contract(cfg: Config, registry: AlgorithmRegistry | None = None): + """Resolve the server capture contract (method, aux layers, dims). + + Composition-root accessor for external launchers (e.g. the online-live + capture-server script), which must not import ``specforge.training``. + """ + from specforge.training.capture_contract import resolve_server_capture_contract + + return resolve_server_capture_contract( + cfg, algorithm=resolve_run(cfg, registry).algorithm + ) + + __all__ = [ "ResolvedOfflineCapture", "ResolvedRun", "bind_run", "build_application_run", + "resolve_capture_contract", "resolve_offline_capture", "resolve_run", ] diff --git a/specforge/cli.py b/specforge/cli.py index bd057dd53..8d27d7ece 100644 --- a/specforge/cli.py +++ b/specforge/cli.py @@ -161,6 +161,11 @@ def _config_for_role(cfg: Config, role: str) -> Config: # child consumes the already-derived environment and must not attempt to # validate or own that stack again. disaggregated["managed_local"] = None + if disaggregated is not None and disaggregated.get("live") is not None: + # Same ownership rule for the supervised live stack; children keep the + # intake host/port but not the supervisor-owned Mooncake topology. + disaggregated["live"]["mooncake"] = None + disaggregated["live"]["trainer_cuda_visible_devices"] = None if role == "producer": raw["profiling"]["enabled"] = False return Config.model_validate(raw) diff --git a/specforge/config/__init__.py b/specforge/config/__init__.py index 70a885798..bdeb8ff21 100644 --- a/specforge/config/__init__.py +++ b/specforge/config/__init__.py @@ -7,6 +7,7 @@ DataConfig, DeploymentConfig, DisaggregatedDeploymentConfig, + LiveIntakeDeploymentConfig, ManagedLocalCaptureServerConfig, ManagedLocalMooncakeConfig, ManagedLocalStackConfig, @@ -27,6 +28,7 @@ "DataConfig", "DeploymentConfig", "DisaggregatedDeploymentConfig", + "LiveIntakeDeploymentConfig", "ManagedLocalMooncakeConfig", "ManagedLocalCaptureServerConfig", "ManagedLocalStackConfig", diff --git a/specforge/config/schema.py b/specforge/config/schema.py index 05686eda6..40a2222cb 100644 --- a/specforge/config/schema.py +++ b/specforge/config/schema.py @@ -150,20 +150,31 @@ class DataConfig(StrictConfigModel): max_prompts: Optional[int] = Field(default=None, ge=0) @model_validator(mode="after") - def _exactly_one_source(self): + def _at_most_one_source(self): sources = [ bool(self.train_data_path), bool(self.prompts_path), bool(self.hidden_states_path), ] - if sum(sources) != 1: + if sum(sources) > 1: raise ValueError( - "set exactly one of data.train_data_path (raw online data), " + "set at most one of data.train_data_path (raw online data), " "data.prompts_path (pre-tokenized online data), or " "data.hidden_states_path (offline features)" ) return self + @property + def source_count(self) -> int: + return sum( + bool(value) + for value in ( + self.train_data_path, + self.prompts_path, + self.hidden_states_path, + ) + ) + class TrackingConfig(StrictConfigModel): """Optional experiment tracking behind the trainer's logger seam.""" @@ -381,6 +392,47 @@ def _validate_local_resources(self): return self +class LiveIntakeDeploymentConfig(StrictConfigModel): + """Intake endpoint where external live capture servers push records. + + Setting ``mooncake`` opts into the one-command launch: the supervisor owns + a loopback Mooncake master and spawns producer + consumer, leaving only + the serving engine and traffic external. + """ + + host: str = "0.0.0.0" + port: int = Field(gt=0, le=65535) + mooncake: Optional[ManagedLocalMooncakeConfig] = None + trainer_cuda_visible_devices: Optional[List[str]] = None + + @model_validator(mode="after") + def _validate_live(self): + if not self.host or self.host.strip() != self.host: + raise ValueError( + "deployment.disaggregated.live.host must be non-empty and must " + "not contain surrounding whitespace" + ) + if (self.mooncake is None) != (self.trainer_cuda_visible_devices is None): + raise ValueError( + "live.mooncake and live.trainer_cuda_visible_devices opt into " + "the supervised launch together" + ) + if self.trainer_cuda_visible_devices is not None: + _validate_cuda_devices( + self.trainer_cuda_visible_devices, + field_name="live.trainer_cuda_visible_devices", + ) + if self.mooncake is not None: + ports = { + self.mooncake.rpc_port, + self.mooncake.metadata_port, + self.mooncake.metrics_port, + } + if self.port in ports: + raise ValueError("live.port must not collide with Mooncake ports") + return self + + class DisaggregatedDeploymentConfig(StrictConfigModel): """Shared, non-secret topology for producer/consumer launch planning.""" @@ -418,6 +470,9 @@ class DisaggregatedDeploymentConfig(StrictConfigModel): #: managed_local supervisors use managed_local.shutdown_grace_s instead. shutdown_grace_s: float = Field(default=30.0, gt=0) managed_local: Optional[ManagedLocalStackConfig] = None + #: Online-live mode: hidden states are captured from production serving + #: traffic and pushed to this producer-hosted intake endpoint. + live: Optional[LiveIntakeDeploymentConfig] = None @model_validator(mode="after") def _validate_store(self): @@ -480,6 +535,38 @@ def _validate_store(self): "managed_local online capture is server-owned; do not set " "producer_segment_size" ) + if self.live is not None: + if self.backend != "mooncake": + raise ValueError("live capture requires backend=mooncake") + if self.managed_local is not None: + raise ValueError( + "live capture uses external serving; do not set managed_local" + ) + if self.server_urls: + raise ValueError( + "live capture servers push to the intake; do not set " + "deployment.disaggregated.server_urls" + ) + if self.producer_segment_size is not None: + raise ValueError( + "live capture is server-owned; do not set producer_segment_size" + ) + if self.live.mooncake is not None: + explicit = [ + name + for name in ( + "mooncake_metadata_server", + "mooncake_master_server_addr", + "mooncake_local_hostname", + "mooncake_protocol", + "mooncake_rdma_devices", + ) + if getattr(self, name) + ] + if explicit: + raise ValueError( + f"live.mooncake derives Mooncake endpoints; do not set {explicit}" + ) return self @@ -702,6 +789,52 @@ def _validate_run_structure(self): mode = self.mode deployment = self.deployment.mode role = self.training.role + live = ( + self.deployment.disaggregated.live + if self.deployment.disaggregated is not None + else None + ) + + if live is not None: + if self.data.source_count != 0: + raise ValueError( + "live capture trains from serving traffic; do not set a " + "data source" + ) + if self.training.max_steps is None and self.training.total_steps is None: + raise ValueError( + "live capture is an unbounded stream; set training.max_steps " + "or training.total_steps" + ) + if live.mooncake is not None: + if self.deployment.trainer.nnodes != 1: + raise ValueError( + "live.mooncake requires deployment.trainer.nnodes=1" + ) + if self.training.role != "auto": + raise ValueError( + "live.mooncake requires the persisted training role to " + "be auto" + ) + if self.training.resume_from is not None: + raise ValueError( + "live.mooncake requires a fresh control_dir and does " + "not support resume" + ) + if ( + len(live.trainer_cuda_visible_devices) + != self.deployment.trainer.nproc_per_node + ): + raise ValueError( + "live.trainer_cuda_visible_devices count must equal " + "deployment.trainer.nproc_per_node" + ) + elif self.data.source_count != 1: + raise ValueError( + "set exactly one of data.train_data_path (raw online data), " + "data.prompts_path (pre-tokenized online data), or " + "data.hidden_states_path (offline features)" + ) if mode == "online" and deployment != "disaggregated": raise ValueError( diff --git a/specforge/inference/adapters/live_intake.py b/specforge/inference/adapters/live_intake.py new file mode 100644 index 000000000..31fc1d5d2 --- /dev/null +++ b/specforge/inference/adapters/live_intake.py @@ -0,0 +1,156 @@ +# coding=utf-8 +# Copyright 2024 The SpecForge team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Online-live ref source: capture records pushed by external serving engines. + +The driven transport (:class:`SGLangServerCaptureAdapter`) originates capture +requests and reads results from ``/generate`` responses. In live mode the +serving engine originates samples from real user traffic, writes tensors into +Mooncake itself, and pushes the same result record to the producer's +:class:`CaptureIntakeServer`. This source is the receiving half: it publishes +the capture-config handshake document and turns each validated record into a +committed-ready ``SampleRef``. +""" + +from __future__ import annotations + +from typing import Any, Dict, List + +from specforge.inference.adapters.server_capture import ( + ServerCaptureSchema, + build_server_capture_ref, + validate_server_capture_ref, +) +from specforge.inference.capture import CaptureConfig, CaptureMismatchError + +_LIVE_TRANSPORT = "sglang_live_capture" + + +class LiveIntakeRefSource: + """Validate pushed capture records against the run's capture contract.""" + + def __init__( + self, + store, + *, + run_id: str, + algorithm: str, + schema: ServerCaptureSchema, + capture: CaptureConfig, + max_num_tokens: int, + target_model_version: str = "unknown", + ) -> None: + if not hasattr(store, "adopt") or not hasattr(store, "store_id"): + raise TypeError( + "LiveIntakeRefSource needs a MooncakeFeatureStore-like store" + ) + if not algorithm: + raise ValueError("algorithm must be non-empty") + if max_num_tokens < 1: + raise ValueError("max_num_tokens must be >= 1") + self.store = store + self.run_id = run_id + self.strategy = algorithm + self.schema = schema + self.capture = capture + self.max_num_tokens = int(max_num_tokens) + self.target_model_version = target_model_version + self._expected_features = frozenset(self._feature_names()) + + def _feature_names(self) -> List[str]: + names = [ + name + for name in (self.schema.aux_feature, self.schema.last_hidden_feature) + if name is not None + ] + names.extend(name for name, _key, _trailing in self.schema.passthrough) + if self.schema.attention_mask_feature is not None: + names.append(self.schema.attention_mask_feature) + return names + + def config_payload(self) -> Dict[str, Any]: + """The handshake document a live capture server GETs at startup.""" + features: Dict[str, str] = {} + if self.schema.aux_feature is not None: + features["aux"] = self.schema.aux_feature + if self.schema.last_hidden_feature is not None: + features["last_hidden"] = self.schema.last_hidden_feature + passthrough: List[Dict[str, str]] = [] + for feature_name, payload_key, trailing in self.schema.passthrough: + if trailing: + raise ValueError( + f"live capture cannot synthesize passthrough {feature_name!r} " + f"with trailing shape {trailing}" + ) + source = "tokens" if payload_key == "input_ids" else "ones" + passthrough.append({"name": feature_name, "source": source}) + if self.schema.attention_mask_feature is not None: + passthrough.append( + {"name": self.schema.attention_mask_feature, "source": "ones"} + ) + return { + "store_id": str(self.store.store_id), + "run_id": self.run_id, + "gen": 1, + "features": features, + "passthrough": passthrough, + # Floor of 2 keeps warmup/probe one-token requests out of training + # (DFlash-family objectives need two consecutive supervised tokens). + "min_num_tokens": 2, + "max_num_tokens": self.max_num_tokens, + } + + def ref_from_record(self, record: Dict[str, Any]): + """Build and verify a SampleRef; raises loudly on any contract breach.""" + store_id = str(record.get("store_id")) + if store_id != str(self.store.store_id): + raise ValueError( + f"record store_id {store_id!r} != run store " + f"{self.store.store_id!r}" + ) + if int(record.get("gen", -1)) != 1: + raise ValueError(f"live capture requires gen=1, got {record.get('gen')}") + num_tokens = int(record.get("num_tokens", 0)) + if num_tokens < 2: + raise ValueError("record requires num_tokens >= 2") + if num_tokens > self.max_num_tokens: + raise ValueError( + f"record num_tokens {num_tokens} exceeds the run cap " + f"{self.max_num_tokens}" + ) + features = record.get("features") + if not isinstance(features, dict) or not features: + raise ValueError("record requires a features mapping") + if frozenset(features) != self._expected_features: + raise CaptureMismatchError( + f"record features {sorted(features)} != expected " + f"{sorted(self._expected_features)}" + ) + ref = build_server_capture_ref( + record, + schema=self.schema, + capture=self.capture, + run_id=self.run_id, + strategy=self.strategy, + source_task_id=str(record["sample_id"]), + target_model_version=self.target_model_version, + num_tokens=num_tokens, + transport=_LIVE_TRANSPORT, + origin=str(record.get("origin", "live")), + ) + validate_server_capture_ref( + ref, + record, + schema=self.schema, + capture=self.capture, + expected_len=num_tokens, + ) + return ref + + +__all__ = ["LiveIntakeRefSource"] diff --git a/specforge/inference/adapters/server_capture.py b/specforge/inference/adapters/server_capture.py index f30a4d292..4cee833a3 100644 --- a/specforge/inference/adapters/server_capture.py +++ b/specforge/inference/adapters/server_capture.py @@ -65,6 +65,103 @@ class ServerCaptureFailure: retryable: bool = True +def build_server_capture_ref( + result: Dict[str, Any], + *, + schema: ServerCaptureSchema, + capture: CaptureConfig, + run_id: str, + strategy: str, + source_task_id: str, + target_model_version: str, + num_tokens: int, + transport: str, + origin: str, + tokenizer_version: str = "unknown", +): + """Build a SampleRef from a server capture result dict (no validation).""" + from specforge.runtime.contracts import SampleRef + + sample_id = str(result["sample_id"]) + gen = int(result["gen"]) + feats: Dict[str, Dict[str, Any]] = result["features"] + specs: Dict[str, FeatureSpec] = {} + nbytes = 0 + for name, meta in feats.items(): + shape = tuple(int(d) for d in meta["shape"]) + dtype = str(meta["dtype"]) + extra: Dict[str, Any] = {} + if name == schema.last_hidden_feature: + extra["target_repr"] = capture.target_repr + if capture.vocab_map_version: + extra["target_meta"] = {"vocab_map_version": capture.vocab_map_version} + specs[name] = FeatureSpec(name=name, shape=shape, dtype=dtype, **extra) + nbytes += _spec_nbytes(shape, dtype) + return SampleRef( + sample_id=sample_id, + run_id=run_id, + source_task_id=source_task_id, + feature_store_uri=f"mooncake://{result['store_id']}/{sample_id}", + feature_keys={n: f"{sample_id}/{n}" for n in specs}, + feature_specs=specs, + strategy=strategy, + schema_version=SCHEMA_VERSION, + target_model_version=target_model_version, + tokenizer_version=tokenizer_version, + num_tokens=num_tokens, + estimated_bytes=nbytes, + metadata={ + "run_id": run_id, + "source_task_id": source_task_id, + "strategy": strategy, + "target_repr": capture.target_repr, + "vocab_map_version": capture.vocab_map_version, + "transport": transport, + "server": origin, # which server captured it (provenance) + "generation": gen, # the zero-copy get() locator + }, + ) + + +def validate_server_capture_ref( + ref, + result: Dict[str, Any], + *, + schema: ServerCaptureSchema, + capture: CaptureConfig, + expected_len: int, +) -> None: + """Verify a server-built ref against the capture contract; raise loudly.""" + short = { + name: spec.shape + for name, spec in ref.feature_specs.items() + if len(spec.shape) >= 2 and spec.shape[1] != expected_len + } + if short: + raise CaptureMismatchError( + f"[{ref.sample_id}] captured seq len != expected {expected_len} " + f"for {short}" + ) + recorded_aux_layer_ids = result.get("aux_layer_ids") + if capture.aux_hidden_state_layer_ids and recorded_aux_layer_ids is None: + raise CaptureMismatchError( + f"[{ref.sample_id}] capture omitted aux-layer ids; cannot verify " + f"requested layers {capture.aux_hidden_state_layer_ids}" + ) + verify_capture_specs( + ref.feature_specs, + capture, + sample_id=ref.sample_id, + recorded_aux_layer_ids=( + tuple(recorded_aux_layer_ids) + if recorded_aux_layer_ids is not None + else None + ), + aux_feature_name=schema.aux_feature or "hidden_state", + target_feature_name=schema.last_hidden_feature or "target", + ) + + def _default_post(url: str, json_body: Dict[str, Any], timeout: float): import requests @@ -259,51 +356,21 @@ def _spec_capture_payload(self, task: PromptTask) -> Dict[str, Any]: def _ref_from_result( self, task: PromptTask, result: Dict[str, Any], capture: CaptureConfig ): - from specforge.runtime.contracts import SampleRef - - sample_id = str(result["sample_id"]) - gen = int(result["gen"]) - feats: Dict[str, Dict[str, Any]] = result["features"] - specs: Dict[str, FeatureSpec] = {} - nbytes = 0 - for name, meta in feats.items(): - shape = tuple(int(d) for d in meta["shape"]) - dtype = str(meta["dtype"]) - extra: Dict[str, Any] = {} - if name == self.schema.last_hidden_feature: - extra["target_repr"] = capture.target_repr - if capture.vocab_map_version: - extra["target_meta"] = { - "vocab_map_version": capture.vocab_map_version - } - specs[name] = FeatureSpec(name=name, shape=shape, dtype=dtype, **extra) - nbytes += _spec_nbytes(shape, dtype) num_tokens = int(task.metadata.get("num_tokens", 0)) or len( task.payload["input_ids"] ) - return SampleRef( - sample_id=sample_id, + return build_server_capture_ref( + result, + schema=self.schema, + capture=capture, run_id=self.run_id, - source_task_id=task.task_id, - feature_store_uri=f"mooncake://{result['store_id']}/{sample_id}", - feature_keys={n: f"{sample_id}/{n}" for n in specs}, - feature_specs=specs, strategy=self.strategy, - schema_version=SCHEMA_VERSION, + source_task_id=task.task_id, target_model_version=self.target_model_version, - tokenizer_version=str(task.metadata.get("tokenizer_version", "unknown")), num_tokens=num_tokens, - estimated_bytes=nbytes, - metadata={ - "run_id": self.run_id, - "source_task_id": task.task_id, - "strategy": self.strategy, - "target_repr": capture.target_repr, - "vocab_map_version": capture.vocab_map_version, - "transport": "sglang_server_capture", - "server": self.base_url, # which server captured it (provenance) - "generation": gen, # the zero-copy get() locator - }, + transport="sglang_server_capture", + origin=self.base_url, + tokenizer_version=str(task.metadata.get("tokenizer_version", "unknown")), ) # -- the RefSource entry point ---------------------------------------------- @@ -336,10 +403,32 @@ def produce_refs( generation=int(payload["gen"]), feature_names=feature_names, ) + from specforge.runtime.workflow_log import wlog + + wlog( + "producer", + "POST /generate (prompts -> SGLang; capture spec asks the server " + "to write hidden states into Mooncake)", + url=f"{self.base_url}/generate", + n_prompts=len(tasks), + sample_ids=[self._sample_id(t) for t in tasks[:2]] + ["..."], + requested_features=list(capture_payloads[0]["features"].values()) + if capture_payloads + else [], + ) rows = self.post_fn( f"{self.base_url}/generate", json_body=body, timeout=self.timeout_s ) rows = _flatten_list_wrappers(rows) + wlog( + "producer", + "SGLang /generate response (tensors are NOT here; only Mooncake " + "keys + shapes come back in meta_info.spec_capture)", + n_rows=len(rows), + first_spec_capture=(rows[0].get("meta_info") or {}).get("spec_capture") + if rows and isinstance(rows[0], dict) + else None, + ) if len(rows) != len(tasks): raise RuntimeError( f"spec-capture server returned {len(rows)} rows for " @@ -400,52 +489,16 @@ def produce_refs( f"{expected_identity}" ) ref = self._ref_from_result(task, result, capture) - # A capture shorter than the prompt is corrupt: it means cache - # isolation or another full-prefill invariant failed despite the - # request's fresh namespace. - expected_len = len(task.payload["input_ids"]) - short = { - name: spec.shape - for name, spec in ref.feature_specs.items() - if len(spec.shape) >= 2 and spec.shape[1] != expected_len - } - if short: - self.store.adopt(ref) - self.store.abort(ref.sample_id, reason="seq-len-mismatch") - out.append( - ServerCaptureFailure( - task_id=task.task_id, - reason=( - f"server_capture: captured seq len != prompt len " - f"{expected_len} for {short}; the capture request " - "did not execute a complete prefill" - ), - retryable=False, - ) - ) - continue try: - recorded_aux_layer_ids = result.get("aux_layer_ids") - if ( - capture.aux_hidden_state_layer_ids - and recorded_aux_layer_ids is None - ): - raise CaptureMismatchError( - f"[{ref.sample_id}] capture omitted aux-layer ids; " - "cannot verify requested layers " - f"{capture.aux_hidden_state_layer_ids}" - ) - verify_capture_specs( - ref.feature_specs, - capture, - sample_id=ref.sample_id, - recorded_aux_layer_ids=( - tuple(recorded_aux_layer_ids) - if recorded_aux_layer_ids is not None - else None - ), - aux_feature_name=self.schema.aux_feature or "hidden_state", - target_feature_name=self.schema.last_hidden_feature or "target", + # A capture shorter than the prompt is corrupt: it means cache + # isolation or another full-prefill invariant failed despite + # the request's fresh namespace. + validate_server_capture_ref( + ref, + result, + schema=self.schema, + capture=capture, + expected_len=len(task.payload["input_ids"]), ) except CaptureMismatchError as exc: # Loud boundary failure; free the server-written keys so a @@ -503,4 +556,6 @@ def _spec_nbytes(shape: Tuple[int, ...], dtype: str) -> int: "ServerCaptureSchema", "ServerCaptureFailure", "SGLangServerCaptureAdapter", + "build_server_capture_ref", + "validate_server_capture_ref", ] diff --git a/specforge/inference/sglang_patch_inventory.md b/specforge/inference/sglang_patch_inventory.md index cb92085e6..1e9e32fe5 100644 --- a/specforge/inference/sglang_patch_inventory.md +++ b/specforge/inference/sglang_patch_inventory.md @@ -12,7 +12,7 @@ Online training uses one of these source-specific patches: | Target | Patch | Capture methods | |---|---|---| -| SGLang v0.5.14 / `inkling-support` | [`patches/sglang/v0.5.14/spec-capture.patch`](../../patches/sglang/v0.5.14/spec-capture.patch) | EAGLE3, DFlash | +| SGLang v0.5.14 / `inkling-support` | [`patches/sglang/v0.5.14/spec-capture.patch`](../../patches/sglang/v0.5.14/spec-capture.patch) | EAGLE3, DFlash, DSpark (DSpark reuses the DFlash aux wiring; stock models expose only `set_dflash_layers_to_capture`) | | Kimi K3 SGLang `9acd9cb` (`f8493a4` compatible) | [`patches/sglang/kimi-k3-f8493a4/spec-capture.patch`](../../patches/sglang/kimi-k3-f8493a4/spec-capture.patch) | EAGLE3, DFlash, DSpark | The patch adds `--enable-spec-capture` and a server-side sink that: @@ -51,6 +51,39 @@ Marlin reduction fallback when the token dimension exceeds CUDA grid.y's 65,535 limit. The server-capture unit and GPU gates must pass before updating either supported source revision. +## Online-live: capture from production serving traffic + +[`patches/sglang/online-live/spec-capture-live.patch`](../../patches/sglang/online-live/spec-capture-live.patch) +is a separate layer on top of the v0.5.14 base patch (apply with +`scripts/apply_sglang_spec_capture_patch.sh --live`, or let +`scripts/online_live/launch_capture_server.py` apply it and launch the server +in one command). It captures ordinary +user requests — no client `spec_capture` field — and pushes the resulting +records to a SpecForge producer, closing the serve→train loop: + +- Two new flags: `--spec-capture-intake-url http://:` and + `--spec-capture-sample-rate` (deterministic per request id, so all TP ranks + make the same capture decision). +- At startup the sink GETs `/v1/spec-capture/config` from the intake + ([`runtime/data_plane/intake_server.py`](../runtime/data_plane/intake_server.py)): + store id, feature names, passthrough synthesis rules (`input_ids` from the + request tokens, `loss_mask`/`attention_mask` all-ones), and the token cap. +- Capture covers prefill **and** generated tokens: one row per fed token, so a + finished request yields `prompt + output[:-1]` rows aligned with the + `input_ids` passthrough. +- Tensors still go straight into Mooncake via the base sink; the tensor-free + record then POSTs to `/v1/spec-capture/records` from a bounded background + writer queue (never the scheduler loop). Any failure — queue full, Mooncake + down or pool full, intake down, 429 shed, 422 reject — drops just that + capture, removes any keys already written, and never touches the user's + response. Flow control is therefore shed-based; the Mooncake pool size is + the final buffer bound and trainer acks free it. + +v1 limitations: `--chunked-prefill-size -1` is still required; the live +capture server cannot itself serve with speculative decoding; a retracted +request drops its capture; a sink crash between the Mooncake write and the +key removal leaks hard-pinned orphans until the store is restarted. + ## Offline: dedicated local capture [`../offline_capture`](../offline_capture) is used exclusively by diff --git a/specforge/launch.py b/specforge/launch.py index 001c56696..7d980e7fa 100644 --- a/specforge/launch.py +++ b/specforge/launch.py @@ -789,6 +789,59 @@ def refs_for_epoch(epoch): # --------------------------------------------------------------------------- +def _await_consumer_quantum( + channel, + *, + in_flight_high_watermark: int, + resolved_low_watermark: int, + peer_wait_timeout_s: Optional[float], + sleep, + poll_s: float, +) -> int: + """Block until the consumer publishes its optimizer window; validate it.""" + import time + + start = time.monotonic() + while True: + quantum = channel.consumer_quantum() + if quantum is not None: + break + consumer_failure = channel.consumer_failure() + if consumer_failure is not None: + raise RuntimeError( + "consumer failed before publishing its optimizer window: " + f"{consumer_failure}" + ) + if ( + peer_wait_timeout_s is not None + and time.monotonic() - start > peer_wait_timeout_s + ): + raise TimeoutError( + "producer timed out waiting for the consumer optimizer " + f"window after {peer_wait_timeout_s:.0f}s" + ) + sleep(poll_s) + if in_flight_high_watermark < quantum: + raise ValueError( + "producer in-flight high watermark " + f"{in_flight_high_watermark} is smaller than the consumer's " + f"global optimizer-step quantum {quantum}; set " + "DISAGG_IN_FLIGHT_HIGH_WATERMARK to at least that value" + ) + if resolved_low_watermark < quantum: + # The consumer dispatches only whole optimizer windows, so a paused + # producer must be resumable while the consumer still needs up to one + # full window: a low watermark below the quantum could leave both + # sides waiting on each other. + raise ValueError( + "producer in-flight low watermark " + f"{resolved_low_watermark} is smaller than the consumer's " + f"global optimizer-step quantum {quantum}; set " + "DISAGG_IN_FLIGHT_LOW_WATERMARK to at least that value" + ) + return quantum + + def build_disagg_online_producer( *, algorithm: AlgorithmRegistration, @@ -997,46 +1050,17 @@ def drive_producer(max_rounds: int = 1_000_000, should_stop=None) -> int: f"{flow_control.limits.resolved_low_watermark_refs} " f"progress_interval={progress_interval}" ) - quantum_wait_start = time.monotonic() try: - while True: - consumer_quantum = channel.consumer_quantum() - if consumer_quantum is not None: - break - consumer_failure = channel.consumer_failure() - if consumer_failure is not None: - raise RuntimeError( - "consumer failed before publishing its optimizer window: " - f"{consumer_failure}" - ) - if ( - peer_wait_timeout_s is not None - and time.monotonic() - quantum_wait_start > peer_wait_timeout_s - ): - raise TimeoutError( - "producer timed out waiting for the consumer optimizer " - f"window after {peer_wait_timeout_s:.0f}s" - ) - sleep(backpressure_poll_s) - if in_flight_high_watermark < consumer_quantum: - raise ValueError( - "producer in-flight high watermark " - f"{in_flight_high_watermark} is smaller than the consumer's " - f"global optimizer-step quantum {consumer_quantum}; set " - "DISAGG_IN_FLIGHT_HIGH_WATERMARK to at least that value" - ) - resolved_low_watermark = flow_control.limits.resolved_low_watermark_refs - if resolved_low_watermark < consumer_quantum: - # The consumer dispatches only whole optimizer windows, so a - # paused producer must be resumable while the consumer still - # needs up to one full window: a low watermark below the - # quantum could leave both sides waiting on each other. - raise ValueError( - "producer in-flight low watermark " - f"{resolved_low_watermark} is smaller than the consumer's " - f"global optimizer-step quantum {consumer_quantum}; set " - "DISAGG_IN_FLIGHT_LOW_WATERMARK to at least that value" - ) + consumer_quantum = _await_consumer_quantum( + channel, + in_flight_high_watermark=in_flight_high_watermark, + resolved_low_watermark=( + flow_control.limits.resolved_low_watermark_refs + ), + peer_wait_timeout_s=peer_wait_timeout_s, + sleep=sleep, + poll_s=backpressure_poll_s, + ) producer_timing( f"consumer optimizer window ready quantum={consumer_quantum}" ) @@ -1300,6 +1324,21 @@ def run_worker(w) -> None: def ingest_prompt_batch(epoch: int, batch_index: int, epoch_prompts) -> None: phase = time.perf_counter() + if epoch_prompts: + from specforge.runtime.workflow_log import wlog + + first = epoch_prompts[0] + wlog( + "producer", + "pull data: tokenized prompts from the dataset cache enter " + "the flow controller (leased to rollout workers next)", + n_prompts=len(epoch_prompts), + epoch=epoch + 1, + first_prompt_keys=list(first)[:8], + first_prompt_len=len(first.get("input_ids", ())) + if isinstance(first, dict) + else "?", + ) producer_timing( "controller.ingest_prompts start " f"epoch={epoch + 1}/{prompt_epochs} batch={batch_index + 1} " @@ -1429,6 +1468,226 @@ def run_worker_guarded(w) -> None: return workers, drive_producer +def build_disagg_live_producer( + *, + feature_store: FeatureStore, + channel, + ref_source, + intake_host: str, + intake_port: int, + in_flight_high_watermark: int = 256, + in_flight_low_watermark: Optional[int] = None, + resident_high_watermark_bytes: Optional[int] = None, + resident_low_watermark_bytes: Optional[int] = None, + feature_store_max_resident_bytes: Optional[int] = None, + gc_interval_s: float = 30.0, + backpressure_poll_s: float = 0.2, + peer_wait_timeout_s: Optional[float] = None, + progress_interval_s: float = 30.0, + sleep=None, +): + """Producer side of an ONLINE-LIVE run (intake pump, no rollout pool). + + External serving engines capture hidden states from real user traffic, + write tensors straight into Mooncake, and push tensor-free capture records + to the hosted :class:`CaptureIntakeServer`. Each fresh record is + validated by ``ref_source``, adopted into ``feature_store``, and published + on ``channel``; the consumer side is identical to the driven online mode. + + Flow control is shed-based because user traffic cannot be paused: above + the in-flight/resident watermarks (or the hard byte cap) a record answers + 429 and the serving sink drops that capture. The Mooncake pool itself is + the final bound — when it fills the server's write fails and the capture + is dropped before a record is ever pushed. + + Returns ``(intake_server, drive_producer)``. ``drive_producer`` blocks + until ``should_stop`` (normally the consumer's stop sentinel) and closes + the channel; a failure publishes the failure sentinel instead. + """ + import threading + import time + from collections import deque + + from specforge.runtime.control_plane.flow_control import ( + FlowControlLimits, + ProducerFlowControl, + ) + from specforge.runtime.data_plane.intake_server import CaptureIntakeServer + + sleep = sleep or time.sleep + flow_control = ProducerFlowControl( + FlowControlLimits( + high_watermark_refs=in_flight_high_watermark, + low_watermark_refs=in_flight_low_watermark, + high_watermark_bytes=resident_high_watermark_bytes, + low_watermark_bytes=resident_low_watermark_bytes, + ) + ) + if ( + feature_store_max_resident_bytes is not None + and feature_store_max_resident_bytes < 1 + ): + raise ValueError("feature_store_max_resident_bytes must be >= 1") + if ( + feature_store_max_resident_bytes is not None + and resident_high_watermark_bytes is not None + and feature_store_max_resident_bytes < resident_high_watermark_bytes + ): + raise ValueError( + "feature_store_max_resident_bytes must be >= resident_high_watermark_bytes" + ) + + publish_lock = threading.Lock() + published_sizes = deque() + state = { + "produced": 0, + "shed": 0, + "rejected": 0, + "quantum": None, + "accounted_consumed": 0, + "resident_bytes": 0, + "stopped": False, + "fatal": None, + } + + def reconcile_consumed_locked() -> int: + # Mooncake health is process-local and cannot observe deletes made by + # a remote consumer; track resident bytes from estimated_bytes and the + # channel's durable consumed counter. + consumed = channel.consumed_remote() + delta = consumed - state["accounted_consumed"] + if delta < 0 or delta > len(published_sizes): + raise RuntimeError( + "producer byte accounting does not match the channel: " + f"consumed advanced by {delta} with " + f"{len(published_sizes)} published refs tracked" + ) + for _ in range(delta): + state["resident_bytes"] -= published_sizes.popleft() + state["accounted_consumed"] = consumed + return state["resident_bytes"] + + def on_record(record): + with publish_lock: + if state["fatal"] is not None: + return "rejected", "producer failed" + if state["stopped"] or state["quantum"] is None: + return "rejected", "producer is not accepting records" + in_flight = channel.in_flight_remote() + current_bytes = reconcile_consumed_locked() + # Never shed below one optimizer window or the consumer starves. + if in_flight >= state["quantum"] and flow_control.should_pause( + in_flight_refs=in_flight, resident_bytes=current_bytes + ): + state["shed"] += 1 + return "shed", f"in_flight={in_flight} bytes={current_bytes}" + try: + ref = ref_source.ref_from_record(record) + except Exception as exc: + state["rejected"] += 1 + return "rejected", f"{type(exc).__name__}: {exc}" + ref_size = max(0, int(ref.estimated_bytes or 0)) + projected_bytes = current_bytes + ref_size + if ( + feature_store_max_resident_bytes is not None + and projected_bytes > feature_store_max_resident_bytes + ): + state["shed"] += 1 + return "shed", ( + f"resident byte hard cap {feature_store_max_resident_bytes}" + ) + feature_store.adopt(ref) + try: + _publish_refs_with_cleanup( + channel=channel, feature_store=feature_store, refs=[ref] + ) + except BaseException as exc: + state["fatal"] = exc + raise + published_sizes.append(ref_size) + state["resident_bytes"] = projected_bytes + state["produced"] += 1 + return "accepted", ref.sample_id + + intake_server = CaptureIntakeServer( + intake_host, + intake_port, + config_payload=ref_source.config_payload(), + on_record=on_record, + ) + + def drive_producer(should_stop=None) -> int: + try: + quantum = _await_consumer_quantum( + channel, + in_flight_high_watermark=in_flight_high_watermark, + resolved_low_watermark=( + flow_control.limits.resolved_low_watermark_refs + ), + peer_wait_timeout_s=peer_wait_timeout_s, + sleep=sleep, + poll_s=backpressure_poll_s, + ) + except BaseException as exc: + try: + channel.fail(f"{type(exc).__name__}: {exc}") + except Exception: + logger.exception("failed to publish producer setup failure") + raise + with publish_lock: + state["accounted_consumed"] = channel.consumed_remote() + state["quantum"] = quantum + intake_server.start() + logger.info( + "live capture intake listening on %s:%d (quantum=%d)", + intake_host, + intake_server.port, + quantum, + ) + last_gc = last_log = time.monotonic() + try: + while True: + if state["fatal"] is not None: + raise state["fatal"] + if should_stop is not None and should_stop(): + break + now = time.monotonic() + if gc_interval_s and now - last_gc >= gc_interval_s: + feature_store.gc() + last_gc = now + if progress_interval_s and now - last_log >= progress_interval_s: + logger.info( + "live producer produced=%d shed=%d rejected=%d " + "in_flight=%d resident_bytes=%d intake=%s", + state["produced"], + state["shed"], + state["rejected"], + channel.in_flight_remote(), + state["resident_bytes"], + intake_server.stats(), + ) + last_log = now + sleep(backpressure_poll_s) + except BaseException as exc: + with publish_lock: + state["stopped"] = True + intake_server.stop() + try: + channel.fail(f"{type(exc).__name__}: {exc}") + except Exception: + logger.exception("failed to publish producer failure sentinel") + raise + with publish_lock: + state["stopped"] = True + intake_server.stop() + channel.close() # successful EOF; a failure uses channel.fail() + return state["produced"] + + drive_producer.flow_control = flow_control + drive_producer.stats = lambda: dict(state) + return intake_server, drive_producer + + def build_disagg_online_consumer( *, algorithm: AlgorithmRegistration, diff --git a/specforge/launch_plan.py b/specforge/launch_plan.py index c9444ad78..cc6da68fa 100644 --- a/specforge/launch_plan.py +++ b/specforge/launch_plan.py @@ -355,6 +355,18 @@ def _managed_local_environment(cfg: Config) -> dict[str, str]: "MOONCAKE_LOCAL_HOSTNAME": mooncake.local_hostname, "MOONCAKE_PROTOCOL": mooncake.protocol, "DISAGG_SERVER_URLS": ",".join(server_urls), + # Mooncake's default 60s transfer deadline is too tight for large + # hidden-state tensors when the transfer engine stalls under load + # (observed TRANSFER_FAIL on loopback TCP); match the recipe docs. + "MC_TRANSFER_TIMEOUT": os.environ.get("MC_TRANSFER_TIMEOUT", "300"), + # Without an explicit bind, the transfer engine advertises data + # endpoints on an auto-discovered external interface even for this + # single-node topology, routing tensor traffic through the host NIC + # stack (observed stalls/EFAULT under cilium). Pin TCP data to the + # same loopback endpoint as the master. + "MC_TCP_BIND_ADDRESS": os.environ.get( + "MC_TCP_BIND_ADDRESS", mooncake.local_hostname + ), } if mooncake.rdma_devices: values["MOONCAKE_RDMA_DEVICES"] = mooncake.rdma_devices @@ -388,48 +400,91 @@ def _sglang_argv( return argv -def _managed_local_services( - cfg: Config, - *, - algorithm: "AlgorithmRegistration", -) -> tuple[ServiceSpec, ...]: - from specforge.training.capture_contract import resolve_server_capture_contract - +def _live_managed_environment(cfg: Config) -> dict[str, str]: deployment = cfg.deployment.disaggregated - assert deployment is not None and deployment.managed_local is not None - managed = deployment.managed_local - mooncake = managed.mooncake - control_dir = Path(deployment.control_dir) - log_dir = control_dir / "logs" - shared_env = _managed_local_environment(cfg) - capture_context_length = cfg.model.sglang_context_length or ( - cfg.data.max_length + SGLANG_CAPTURE_CONTEXT_HEADROOM - ) + assert deployment is not None and deployment.live is not None + mooncake = deployment.live.mooncake + assert mooncake is not None + values = { + "MOONCAKE_METADATA_SERVER": ( + f"http://127.0.0.1:{mooncake.metadata_port}/metadata" + ), + "MOONCAKE_MASTER_SERVER_ADDR": f"127.0.0.1:{mooncake.rpc_port}", + "MOONCAKE_LOCAL_HOSTNAME": mooncake.local_hostname, + "MOONCAKE_PROTOCOL": mooncake.protocol, + "MC_TRANSFER_TIMEOUT": os.environ.get("MC_TRANSFER_TIMEOUT", "300"), + "MC_TCP_BIND_ADDRESS": os.environ.get( + "MC_TCP_BIND_ADDRESS", mooncake.local_hostname + ), + } + if mooncake.rdma_devices: + values["MOONCAKE_RDMA_DEVICES"] = mooncake.rdma_devices + return values + - mooncake_service = ServiceSpec( +def _mooncake_master_service( + control_dir: Path, mooncake, metadata_server_url: str +) -> ServiceSpec: + # Prefer the interpreter's own mooncake_master: a system-wide binary of a + # different version speaks an incompatible RPC schema, and every client + # mount then fails with "invalid rpc arg" / RPC_FAIL. + _venv_master = Path(sys.executable).parent / "mooncake_master" + mooncake_master_bin = ( + str(_venv_master) if _venv_master.exists() else "mooncake_master" + ) + return ServiceSpec( command=CommandSpec( "mooncake", ( - "mooncake_master", + mooncake_master_bin, "--enable_http_metadata_server=true", "--http_metadata_server_host=127.0.0.1", f"--rpc_port={mooncake.rpc_port}", f"--http_metadata_server_port={mooncake.metadata_port}", f"--metrics_port={mooncake.metrics_port}", + # The default 10s KV lease can expire mid-transfer when large + # hidden-state tensors move slowly (observed LEASE_EXPIRED on + # loopback TCP). Do not raise this too far: a live read lease + # blocks the store's non-forced remove-on-consume, so a long + # TTL parks freed samples and can exhaust the global segment. + "--default_kv_lease_ttl=60000", ), {"CUDA_VISIBLE_DEVICES": ""}, ), readiness=ReadinessSpec( "mooncake", - shared_env["MOONCAKE_METADATA_SERVER"] + "?key=specforge-health-check", + metadata_server_url + "?key=specforge-health-check", mooncake.startup_timeout_s, tcp_host="127.0.0.1", tcp_port=mooncake.rpc_port, ), - log_path=str(log_dir / "mooncake.log"), + log_path=str(control_dir / "logs" / "mooncake.log"), phase=0, ) + +def _managed_local_services( + cfg: Config, + *, + algorithm: "AlgorithmRegistration", +) -> tuple[ServiceSpec, ...]: + from specforge.training.capture_contract import resolve_server_capture_contract + + deployment = cfg.deployment.disaggregated + assert deployment is not None and deployment.managed_local is not None + managed = deployment.managed_local + mooncake = managed.mooncake + control_dir = Path(deployment.control_dir) + log_dir = control_dir / "logs" + shared_env = _managed_local_environment(cfg) + capture_context_length = cfg.model.sglang_context_length or ( + cfg.data.max_length + SGLANG_CAPTURE_CONTEXT_HEADROOM + ) + + mooncake_service = _mooncake_master_service( + control_dir, mooncake, shared_env["MOONCAKE_METADATA_SERVER"] + ) + contract = resolve_server_capture_contract(cfg, algorithm=algorithm) capture_services = [] for index, server in enumerate(managed.capture_servers): @@ -624,6 +679,9 @@ def _validate_capture_urls( deployment = cfg.deployment.disaggregated if deployment is not None and deployment.managed_local is not None: return + if deployment is not None and deployment.live is not None: + # Live capture servers are external and push to the producer's intake. + return if deployment is not None and deployment.server_urls: return if base_env.get("DISAGG_SERVER_URLS") or base_env.get("DISAGG_SERVER_URL"): @@ -665,32 +723,39 @@ def build_launch_plan( raise ValueError( "managed_local launch planning requires a resolved algorithm registration" ) + live = deployment.live if deployment is not None else None + live_managed = live is not None and live.mooncake is not None + managed_stack = "managed_local" if managed_local is not None else ( + "live.mooncake" if live_managed else None + ) managed_child = base_env.get(_MANAGED_CHILD_ENV) == "1" - if managed_local is not None: + if managed_stack is not None: if managed_child: if requested_role not in ("producer", "consumer"): raise ValueError( - "managed_local child workers require an explicit producer " + f"{managed_stack} child workers require an explicit producer " "or consumer role" ) if not Path(deployment.control_dir, "logs").is_dir(): raise ValueError( - "managed_local child worker requires an active supervisor " + f"{managed_stack} child worker requires an active supervisor " f"control_dir: {deployment.control_dir}" ) else: if distributed: - raise ValueError("managed_local cannot run inside an existing torchrun") + raise ValueError( + f"{managed_stack} cannot run inside an existing torchrun" + ) if requested_role not in ("auto", "both"): raise ValueError( - "managed_local supports only --role auto or --role both" + f"{managed_stack} supports only --role auto or --role both" ) if node_rank is not None: - raise ValueError("managed_local does not accept --node-rank") + raise ValueError(f"{managed_stack} does not accept --node-rank") if os.path.exists(deployment.control_dir): raise ValueError( - "managed_local requires a fresh control_dir, but it already " - f"exists: {deployment.control_dir}" + f"{managed_stack} requires a fresh control_dir, but it " + f"already exists: {deployment.control_dir}" ) role = _resolve_role(cfg, requested_role, distributed=distributed) topology = cfg.deployment.trainer @@ -711,19 +776,25 @@ def build_launch_plan( managed_environment: dict[str, str] = {} if managed_local is not None: managed_environment = _managed_local_environment(cfg) + elif live_managed: + managed_environment = _live_managed_environment(cfg) + if managed_environment: role_base_env = {**base_env, **managed_environment} if role in ("producer", "both"): producer_env = _disaggregated_env(cfg, role_base_env, role="producer") if role in ("consumer", "both"): consumer_env = _disaggregated_env(cfg, role_base_env, role="consumer") - if managed_local is not None: + if managed_stack is not None: + trainer_devices = ( + managed_local.trainer_cuda_visible_devices + if managed_local is not None + else live.trainer_cuda_visible_devices + ) producer_env.update(managed_environment) producer_env["CUDA_VISIBLE_DEVICES"] = "" producer_env[_MANAGED_CHILD_ENV] = "1" consumer_env.update(managed_environment) - consumer_env["CUDA_VISIBLE_DEVICES"] = ",".join( - managed_local.trainer_cuda_visible_devices - ) + consumer_env["CUDA_VISIBLE_DEVICES"] = ",".join(trainer_devices) consumer_env[_MANAGED_CHILD_ENV] = "1" _validate_capture_urls(cfg, role=role, base_env=base_env) @@ -820,6 +891,28 @@ def build_launch_plan( ), shutdown_grace_s=managed_local.shutdown_grace_s, ) + if live_managed: + mooncake = live.mooncake + return LaunchPlan( + "managed_supervisor", + "both", + commands=(producer, consumer), + services=( + _mooncake_master_service( + Path(deployment.control_dir), + mooncake, + f"http://127.0.0.1:{mooncake.metadata_port}/metadata", + ), + ), + managed_root=deployment.control_dir, + managed_ports=( + mooncake.rpc_port, + mooncake.metadata_port, + mooncake.metrics_port, + live.port, + ), + shutdown_grace_s=deployment.shutdown_grace_s, + ) return LaunchPlan( "supervisor", "both", @@ -906,17 +999,21 @@ def _managed_preflight(plan: LaunchPlan) -> None: mooncake_available = False if not mooncake_available: raise RuntimeError("managed_local requires the mooncake Python package") - try: - patched_sglang = ( - importlib.util.find_spec("sglang.srt.spec_capture_sink") is not None - ) - except ModuleNotFoundError: - patched_sglang = False - if not patched_sglang: - raise RuntimeError( - "managed_local requires patched SGLang spec capture; run " - "scripts/apply_sglang_spec_capture_patch.sh" - ) + owns_capture_servers = any( + service.command.label.startswith("capture-server") for service in plan.services + ) + if owns_capture_servers: + try: + patched_sglang = ( + importlib.util.find_spec("sglang.srt.spec_capture_sink") is not None + ) + except ModuleNotFoundError: + patched_sglang = False + if not patched_sglang: + raise RuntimeError( + "managed_local requires patched SGLang spec capture; run " + "scripts/apply_sglang_spec_capture_patch.sh" + ) for port in plan.managed_ports: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: @@ -991,6 +1088,14 @@ def _spawn_command( stdout=None, stderr=None, ) -> subprocess.Popen: + from specforge.runtime.workflow_log import wlog + + wlog( + "launcher", + f"spawn {command.label!r}", + argv=" ".join(str(a) for a in command.argv), + gpus=command.env.get("CUDA_VISIBLE_DEVICES", ""), + ) child_env = os.environ.copy() child_env.update(command.env) kwargs = {"env": child_env, "start_new_session": True} diff --git a/specforge/modeling/draft/dflash.py b/specforge/modeling/draft/dflash.py index 4fea605bd..459f7b580 100644 --- a/specforge/modeling/draft/dflash.py +++ b/specforge/modeling/draft/dflash.py @@ -184,6 +184,13 @@ def forward( ) else: attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + if self.config._attn_implementation == "flex_attention": + # Draft blocks issue short (<128-token) GQA queries, which routes + # inductor to its flex-decoding kernel. That path computes + # BLOCK_M = next_pow2(q_len_hint * gqa_groups) and rejects every + # config once BLOCK_M exceeds the 128 sparse block size, raising + # NoValidChoicesError. Force the main flex kernel instead. + kwargs.setdefault("kernel_options", {"FORCE_USE_FLEX_ATTENTION": True}) attn_output, attn_weights = attn_fn( self, q, diff --git a/specforge/runtime/data_plane/intake_server.py b/specforge/runtime/data_plane/intake_server.py new file mode 100644 index 000000000..aa5078a01 --- /dev/null +++ b/specforge/runtime/data_plane/intake_server.py @@ -0,0 +1,170 @@ +# Copyright 2024 The SpecForge team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Producer-hosted HTTP intake for online-live capture records. + +Live capture servers write feature tensors straight into Mooncake and push +only tensor-free capture records here. The endpoint pair is the whole live +protocol: ``GET /v1/spec-capture/config`` hands the server its capture +document (store id, feature names, passthrough synthesis rules, token cap) and +``POST /v1/spec-capture/records`` submits one record per captured request. + +Responses drive the server's cleanup contract: any non-2xx status tells the +sink to remove the keys it just wrote. A duplicate ``sample_id`` (a sink +retry after a lost response) is acknowledged with 200 without re-dispatching. +""" + +from __future__ import annotations + +import json +import threading +from collections import OrderedDict +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Callable, Dict, Tuple +from urllib.parse import urlparse + +_CONFIG_PATH = "/v1/spec-capture/config" +_RECORDS_PATH = "/v1/spec-capture/records" +_MAX_RECORD_BYTES = 1 << 20 + +_STATUS_CODES = { + "accepted": HTTPStatus.OK, + "shed": HTTPStatus.TOO_MANY_REQUESTS, + "rejected": HTTPStatus.UNPROCESSABLE_ENTITY, +} + + +class _IntakeThreadingHTTPServer(ThreadingHTTPServer): + request_queue_size = 256 + daemon_threads = True + + +class CaptureIntakeServer: + """Serve the live capture config and accept pushed capture records. + + ``on_record`` is called with each fresh record under one server-wide lock + and returns ``(status, detail)`` with status in ``accepted``/``shed``/ + ``rejected``. Only accepted sample ids enter the dedup window. + """ + + def __init__( + self, + host: str, + port: int, + *, + config_payload: Dict, + on_record: Callable[[Dict], Tuple[str, str]], + dedup_capacity: int = 65536, + ) -> None: + self.config_payload = dict(config_payload) + self.on_record = on_record + self._dedup: OrderedDict[str, None] = OrderedDict() + self._dedup_capacity = int(dedup_capacity) + self._lock = threading.Lock() + self._stats = {"accepted": 0, "shed": 0, "rejected": 0, "duplicate": 0} + self._httpd = _IntakeThreadingHTTPServer((host, port), self._handler_type()) + self._thread: threading.Thread | None = None + + @property + def port(self) -> int: + return self._httpd.server_address[1] + + def stats(self) -> Dict[str, int]: + with self._lock: + return dict(self._stats) + + def _dispatch(self, record: Dict) -> Tuple[HTTPStatus, Dict]: + sample_id = record.get("sample_id") + if not isinstance(sample_id, str) or not sample_id: + return HTTPStatus.BAD_REQUEST, {"error": "record requires a sample_id"} + with self._lock: + if sample_id in self._dedup: + self._stats["duplicate"] += 1 + return HTTPStatus.OK, {"status": "duplicate"} + status, detail = self.on_record(record) + code = _STATUS_CODES.get(status) + if code is None: + raise RuntimeError(f"on_record returned unknown status {status!r}") + self._stats[status] += 1 + if status == "accepted": + self._dedup[sample_id] = None + while len(self._dedup) > self._dedup_capacity: + self._dedup.popitem(last=False) + return code, {"status": status, "detail": detail} + + def _handler_type(self): + owner = self + + class Handler(BaseHTTPRequestHandler): + server_version = "SpecForgeIntake/1" + + def log_message(self, _format, *_args): + return + + def _json(self, status: HTTPStatus, payload) -> None: + body = json.dumps(payload, separators=(",", ":")).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if urlparse(self.path).path != _CONFIG_PATH: + self._json(HTTPStatus.NOT_FOUND, {"error": "unknown path"}) + return + self._json(HTTPStatus.OK, owner.config_payload) + + def do_POST(self): + if urlparse(self.path).path != _RECORDS_PATH: + self._json(HTTPStatus.NOT_FOUND, {"error": "unknown path"}) + return + try: + length = int(self.headers.get("Content-Length", "0")) + if length < 1 or length > _MAX_RECORD_BYTES: + raise ValueError("invalid body size") + record = json.loads(self.rfile.read(length)) + if not isinstance(record, dict): + raise ValueError("record must be a JSON object") + except (TypeError, ValueError, json.JSONDecodeError) as exc: + self._json(HTTPStatus.BAD_REQUEST, {"error": str(exc)}) + return + record.setdefault("origin", self.client_address[0]) + try: + status, payload = owner._dispatch(record) + except Exception as exc: # noqa: BLE001 — sink must see a 5xx + self._json( + HTTPStatus.INTERNAL_SERVER_ERROR, + {"error": f"{type(exc).__name__}: {exc}"}, + ) + return + self._json(status, payload) + + return Handler + + def start(self) -> "CaptureIntakeServer": + if self._thread is not None: + return self + self._thread = threading.Thread( + target=self._httpd.serve_forever, + name="specforge-capture-intake", + daemon=True, + ) + self._thread.start() + return self + + def stop(self) -> None: + if self._thread is None: + return + self._httpd.shutdown() + self._httpd.server_close() + self._thread.join(timeout=5.0) + self._thread = None + + +__all__ = ["CaptureIntakeServer"] diff --git a/specforge/runtime/data_plane/mooncake_store.py b/specforge/runtime/data_plane/mooncake_store.py index ac3e9cdcc..70f6bc078 100644 --- a/specforge/runtime/data_plane/mooncake_store.py +++ b/specforge/runtime/data_plane/mooncake_store.py @@ -244,6 +244,11 @@ def __init__( # frees within a run (empty in retain_on_release/offline mode); a durable # shared index would own this in the online multi-node follow-up. self._freed: set = set() + # Receive buffers whose get_into failed transiently: the transfer + # engine may still hold in-flight slices targeting them, so they stay + # registered and alive for the life of the process (see + # _store_get_tensor). + self._quarantined_buffers: List[torch.Tensor] = [] self._lock = threading.RLock() self._counter = 0 self._gen_counter = 0 @@ -285,31 +290,85 @@ def _store_put_tensor(self, key: str, t: torch.Tensor) -> None: if rc is not None and int(rc) < 0: raise RuntimeError(f"mooncake put_from failed (status {rc}) for {key}") - def _store_get_tensor(self, key: str, out: torch.Tensor) -> None: - """Zero-copy fetch into a pre-allocated tensor. Raises KeyError if absent. - - The receive buffer is registered with the transfer engine for the get_into - (required by the raw-buffer path), then unregistered. + # Transport-level statuses that do not indicate a missing object: + # REPLICA_IS_NOT_READY (-703), LEASE_EXPIRED (-707, "lease expired before + # data transfer completed"), TRANSFER_FAIL (-800), CHECKSUM_MISMATCH + # (-801), RPC_FAIL (-900), RPC_TIMEOUT (-901). A single stalled TCP + # transfer must not kill a multi-day run, so these are retried on a fresh + # attempt (a retry acquires a fresh lease). OBJECT_NOT_FOUND (-704) stays + # fatal. + _TRANSIENT_GET_STATUSES = frozenset({-703, -707, -800, -801, -900, -901}) + _GET_ATTEMPTS = 3 + + def _attempt_get_into(self, key: str, buf: torch.Tensor, nb: int) -> Optional[int]: + """One registered get_into attempt. Returns rc, or None on a binding error. + + If the transfer fails transiently (e.g. a stalled TCP batch hits the + store's fixed 60s deadline -> TRANSFER_FAIL), the engine may still hold + in-flight slices targeting ``buf``, so it must NOT be unregistered: + doing so corrupted glibc's heap in practice (``munmap_chunk(): invalid + pointer``). Such buffers stay registered and pinned for the life of + the process instead (bounded leak: one receive buffer per rare + transport failure). """ - nb = _nbytes(out) + registered = False try: - self._store.register_buffer(out.data_ptr(), nb) + self._store.register_buffer(buf.data_ptr(), nb) + registered = True except Exception: # pragma: no cover - some builds auto-register pass + rc: Optional[int] = None try: - rc = self._store.get_into(key, out.data_ptr(), nb) - finally: - try: - self._store.unregister_buffer(out.data_ptr()) - except Exception: # pragma: no cover - pass - if rc is None or int(rc) < 0: + rc = int(self._store.get_into(key, buf.data_ptr(), nb)) + except Exception: # pragma: no cover - binding-level failure + rc = None + if registered: + if rc is None or (rc < 0 and rc in self._TRANSIENT_GET_STATUSES): + self._quarantined_buffers.append(buf) + else: + try: + self._store.unregister_buffer(buf.data_ptr()) + except Exception: # pragma: no cover + pass + return rc + + def _store_get_tensor(self, key: str, out: torch.Tensor) -> None: + """Zero-copy fetch into a pre-allocated tensor. Raises KeyError if absent. + + Transient transport failures are retried into a FRESH receive buffer + each time — a buffer whose transfer timed out may still receive late + slices and is quarantined by _attempt_get_into, never reused. A late + slice replaying into ``out`` after a successful retry is harmless: + tensor keys are per-generation and never rewritten, so it can only + write identical bytes. + """ + nb = _nbytes(out) + rc = self._attempt_get_into(key, out, nb) + for attempt in range(1, self._GET_ATTEMPTS): + if rc is not None and rc >= 0: + break + if rc is not None and rc not in self._TRANSIENT_GET_STATUSES: + raise KeyError(f"mooncake get_into failed (status {rc}) for {key}") + logger.warning( + "mooncake get_into transient failure (status %s) for %s; " + "retrying into a fresh buffer (attempt %d/%d)", + rc, + key, + attempt, + self._GET_ATTEMPTS - 1, + ) + time.sleep(2.0 * attempt) + tmp = torch.empty_like(out) + rc = self._attempt_get_into(key, tmp, nb) + if rc is not None and rc == nb: + out.copy_(tmp) + if rc is None or rc < 0: raise KeyError(f"mooncake get_into failed (status {rc}) for {key}") # get_into returns the number of bytes read; a full read returns exactly # nb. A short read (0 <= rc < nb) would leave the tail of this freshly # allocated buffer as uninitialized garbage. Reject it rather than hand # the trainer silently-corrupt data (B5: never serve wrong bytes). - if int(rc) != nb: + if rc != nb: raise KeyError( f"mooncake get_into short read for {key}: got {rc} of {nb} bytes" ) @@ -570,6 +629,16 @@ def _get_tensors( ) out[n] = _alloc_from_spec(spec) # fresh -> clone-on-fetch for free (B5) self._store_get_tensor(key, out[n]) + from specforge.runtime.workflow_log import wlog + + wlog( + "trainer", + "mooncake get_into: fetched one sample's tensors from the store " + "(zero-copy into pre-allocated CPU buffers)", + sample_id=sid, + gen=gen, + **out, + ) return out, gen # -- lifetime ---------------------------------------------------------- diff --git a/specforge/runtime/workflow_log.py b/specforge/runtime/workflow_log.py new file mode 100644 index 000000000..d8586be5b --- /dev/null +++ b/specforge/runtime/workflow_log.py @@ -0,0 +1,51 @@ +"""Env-gated tracing of the disaggregated training workflow, for studying it. + +Set ``SPECFORGE_WORKFLOW_LOG=1`` to print one line per pipeline event, tagged +with the process role and the source location that emitted it:: + + [workflow] producer server_capture.py:342 POST /generate url=... n_prompts=8 + [workflow] trainer mooncake_store.py:601 fetched sample ... hidden_states=Tensor(1, 812, 7680):bfloat16 + +Silent (single boolean check) when the variable is unset, so instrumentation +can stay in place permanently. +""" + +from __future__ import annotations + +import os +import sys + +ENABLED = os.environ.get("SPECFORGE_WORKFLOW_LOG", "0") == "1" + + +def _fmt(v) -> str: + """Compact one-line rendering; tensors become shape:dtype summaries.""" + try: + import torch + + if isinstance(v, torch.Tensor): + dtype = str(v.dtype).replace("torch.", "") + return f"Tensor{tuple(v.shape)}:{dtype}@{v.device}" + except Exception: # pragma: no cover - torch absent in some processes + pass + if isinstance(v, dict): + return "{" + ", ".join(f"{k}={_fmt(x)}" for k, x in v.items()) + "}" + if isinstance(v, (list, tuple)): + if len(v) > 6: + head = ", ".join(_fmt(x) for x in v[:6]) + return f"[{head}, ... +{len(v) - 6} more]" + return "[" + ", ".join(_fmt(x) for x in v) + "]" + if isinstance(v, float): + return f"{v:.4g}" + return str(v) + + +def wlog(role: str, event: str, **fields) -> None: + """Print one workflow event. ``role`` names the process (producer, trainer, + launcher, capture-server); caller file:line is added automatically.""" + if not ENABLED: + return + frame = sys._getframe(1) + where = f"{os.path.basename(frame.f_code.co_filename)}:{frame.f_lineno}" + kv = " ".join(f"{k}={_fmt(v)}" for k, v in fields.items()) + print(f"[workflow] {role:<10s} {where:<28s} {event} {kv}".rstrip(), flush=True) diff --git a/specforge/training/controller.py b/specforge/training/controller.py index aed40978d..30e3ffb1f 100644 --- a/specforge/training/controller.py +++ b/specforge/training/controller.py @@ -657,6 +657,20 @@ def _fit(self, data: Iterable[TrainBatch], progress: Optional[Any]) -> int: if self.ack_fn is not None: pending_ack.extend(batch.sample_ids) self._step_profiler.before_micro_step(self.global_step) + if self.micro_step <= 3: + from specforge.runtime.workflow_log import wlog + + wlog( + "trainer", + "train_step input (one micro-batch of fetched features " + "entering forward/backward; shown for micro-steps 1-3)", + micro_step=self.micro_step, + **{ + k: v + for k, v in vars(batch).items() + if not k.startswith("_") + }, + ) train_compute_started = time.perf_counter() result = self.core.train_step( batch, @@ -665,6 +679,16 @@ def _fit(self, data: Iterable[TrainBatch], progress: Optional[Any]) -> int: ), ) perf_train_compute_s += time.perf_counter() - train_compute_started + if self.micro_step <= 3: + from specforge.runtime.workflow_log import wlog + + wlog( + "trainer", + "train_step output", + micro_step=self.micro_step, + optimizer_stepped=result.optimizer_stepped, + metrics=result.metrics, + ) self.last_metrics = result.metrics # grad accumulated but optimizer has not stepped yet; everything # keyed on optimizer steps fires only at the boundary. diff --git a/specforge/training/disaggregated.py b/specforge/training/disaggregated.py index 47d175e80..3f672155b 100644 --- a/specforge/training/disaggregated.py +++ b/specforge/training/disaggregated.py @@ -545,6 +545,61 @@ def mark_consumer_failed(exc: BaseException) -> None: ) +def _finalize_online_producer( + store, + channel_path: str, + *, + discard_external: bool, + primary_exc: Optional[BaseException], + produced: int, +) -> int: + """Terminal producer cleanup shared by the driven and live online modes.""" + from specforge.runtime.data_plane.feature_store import drain_feature_store_removals + from specforge.runtime.data_plane.streaming_ref_channel import StreamingRefChannel + + cleanup_errors = [] + try: + reader = StreamingRefChannel(channel_path) + while True: + refs = reader.poll(max_n=1024) + if not refs: + break + for ref in refs: + try: + store.abort(ref.sample_id, reason="online-attempt-finished") + except Exception as exc: + cleanup_errors.append( + f"{ref.sample_id}: {type(exc).__name__}: {exc}" + ) + except Exception as exc: + cleanup_errors.append(f"published-ref scan: {type(exc).__name__}: {exc}") + if discard_external: + try: + store.discard_external_attempts(reason="online-attempt-unadopted-capture") + except Exception as exc: + cleanup_errors.append( + f"unadopted-capture cleanup: {type(exc).__name__}: {exc}" + ) + try: + drain_feature_store_removals(store) + except Exception as exc: + cleanup_errors.append(f"pending-remove drain: {type(exc).__name__}: {exc}") + if primary_exc is not None and cleanup_errors: + raise RuntimeError( + f"producer failed ({type(primary_exc).__name__}: " + f"{primary_exc}) and Mooncake cleanup also failed: " + f"{cleanup_errors}" + ) from primary_exc + if primary_exc is not None: + raise primary_exc + if cleanup_errors: + raise RuntimeError( + "producer could not clean all published Mooncake features: " + f"{cleanup_errors}" + ) + return produced + + def _producer_capture_metadata(cfg: Config, algorithm: AlgorithmRegistration): from specforge.training.capture_contract import resolve_server_capture_contract @@ -557,6 +612,121 @@ def _producer_capture_metadata(cfg: Config, algorithm: AlgorithmRegistration): ) +def _build_live_producer( + cfg: Config, + *, + algorithm: AlgorithmRegistration, + streaming, + store, + channel, + channel_path: str, + live, +): + """Producer role of the online-live mode: intake pump over pushed records.""" + from specforge.inference.adapters.live_intake import LiveIntakeRefSource + from specforge.inference.adapters.server_capture import ServerCaptureSchema + from specforge.inference.capture import CaptureConfig + from specforge.launch import build_disagg_live_producer + from specforge.training.assembly import TrainingRun + from specforge.training.capture_contract import resolve_server_capture_contract + + in_flight_high_watermark, in_flight_low_watermark = _online_flow_window(cfg) + contract = resolve_server_capture_contract(cfg, algorithm=algorithm) + feature_contract = algorithm.spec.feature_contract( + "streaming", cfg.model.input_modality + ) + target_repr = streaming.target_representation + allowed = feature_contract.allowed_target_representations + if allowed and target_repr not in allowed: + raise ValueError( + f"target representation {target_repr!r} is not supported by " + f"algorithm {algorithm.name!r}; expected one of {sorted(allowed)}" + ) + layout = streaming.layout + ref_source = LiveIntakeRefSource( + store, + run_id=cfg.run_id, + algorithm=algorithm.name, + schema=ServerCaptureSchema( + aux_feature=layout.aux_feature, + last_hidden_feature=layout.last_hidden_feature, + passthrough=layout.passthrough, + attention_mask_feature=layout.attention_mask_feature, + ), + capture=CaptureConfig.from_strategy( + required_features=feature_contract.required_tensors, + aux_hidden_state_layer_ids=contract.aux_layer_ids, + target_repr=target_repr, + target_hidden_size=contract.target_hidden_size, + target_vocab_size=contract.target_vocab_size, + draft_vocab_size=contract.draft_vocab_size, + ), + max_num_tokens=cfg.data.max_length, + target_model_version=cfg.model.target_model_path, + ) + _intake, drive = build_disagg_live_producer( + feature_store=store, + channel=channel, + ref_source=ref_source, + intake_host=live.host, + intake_port=live.port, + in_flight_high_watermark=in_flight_high_watermark, + in_flight_low_watermark=in_flight_low_watermark, + resident_high_watermark_bytes=( + int(os.environ["DISAGG_RESIDENT_HIGH_WATERMARK_BYTES"]) + if os.environ.get("DISAGG_RESIDENT_HIGH_WATERMARK_BYTES") + else cfg.runtime.resident_high_watermark_bytes + ), + resident_low_watermark_bytes=( + int(os.environ["DISAGG_RESIDENT_LOW_WATERMARK_BYTES"]) + if os.environ.get("DISAGG_RESIDENT_LOW_WATERMARK_BYTES") + else cfg.runtime.resident_low_watermark_bytes + ), + feature_store_max_resident_bytes=( + cfg.runtime.feature_store_max_resident_bytes + ), + peer_wait_timeout_s=_optional_timeout_s("DISAGG_PEER_WAIT_TIMEOUT"), + ) + aux_layer_args = " ".join(str(layer) for layer in contract.aux_layer_ids) + print( + "[live-producer] launch capture servers with: --enable-spec-capture " + f"--spec-capture-method {contract.method} " + f"--spec-capture-aux-layer-ids {aux_layer_args} " + f"--spec-capture-intake-url http://:{live.port} " + "--chunked-prefill-size -1", + flush=True, + ) + + def produce() -> int: + produced = 0 + primary_exc = None + try: + produced = drive(should_stop=channel.consumer_stopped) + consumer_failure = channel.consumer_failure() + if consumer_failure is not None: + raise RuntimeError(f"consumer failed: {consumer_failure}") + except BaseException as exc: + primary_exc = exc + try: + channel.fail(f"{type(exc).__name__}: {exc}") + except Exception as signal_exc: + print( + f"failed to publish producer failure: {signal_exc}", + flush=True, + ) + # Live capture pre-registers nothing, so there are no external + # provisional attempts to discard. + return _finalize_online_producer( + store, + channel_path, + discard_external=False, + primary_exc=primary_exc, + produced=produced, + ) + + return TrainingRun(execute=produce) + + def _build_online( cfg: Config, *, @@ -582,11 +752,26 @@ def _build_online( # retain materialized features until DPAckController commits the optimizer # boundary and explicitly aborts the acknowledged ids. store = _mooncake_store(cfg, retain_on_release=cfg.training.role == "consumer") - from specforge.runtime.data_plane.feature_store import drain_feature_store_removals from specforge.runtime.data_plane.streaming_ref_channel import StreamingRefChannel channel = StreamingRefChannel(channel_path) + live = ( + cfg.deployment.disaggregated.live + if cfg.deployment.disaggregated is not None + else None + ) + if cfg.training.role == "producer" and live is not None: + return _build_live_producer( + cfg, + algorithm=algorithm, + streaming=streaming, + store=store, + channel=channel, + channel_path=channel_path, + live=live, + ) + if cfg.training.role == "producer": from specforge.inference.adapters.server_capture import ( ServerCaptureSchema, @@ -719,52 +904,13 @@ def produce() -> int: f"failed to publish producer failure: {signal_exc}", flush=True, ) - cleanup_errors = [] - try: - reader = StreamingRefChannel(channel_path) - while True: - refs = reader.poll(max_n=1024) - if not refs: - break - for ref in refs: - try: - store.abort(ref.sample_id, reason="online-attempt-finished") - except Exception as exc: - cleanup_errors.append( - f"{ref.sample_id}: {type(exc).__name__}: {exc}" - ) - except Exception as exc: - cleanup_errors.append( - f"published-ref scan: {type(exc).__name__}: {exc}" - ) - try: - store.discard_external_attempts( - reason="online-attempt-unadopted-capture" - ) - except Exception as exc: - cleanup_errors.append( - f"unadopted-capture cleanup: {type(exc).__name__}: {exc}" - ) - try: - drain_feature_store_removals(store) - except Exception as exc: - cleanup_errors.append( - f"pending-remove drain: {type(exc).__name__}: {exc}" - ) - if primary_exc is not None and cleanup_errors: - raise RuntimeError( - f"producer failed ({type(primary_exc).__name__}: " - f"{primary_exc}) and Mooncake cleanup also failed: " - f"{cleanup_errors}" - ) from primary_exc - if primary_exc is not None: - raise primary_exc - if cleanup_errors: - raise RuntimeError( - "producer could not clean all published Mooncake features: " - f"{cleanup_errors}" - ) - return produced + return _finalize_online_producer( + store, + channel_path, + discard_external=True, + primary_exc=primary_exc, + produced=produced, + ) return TrainingRun(execute=produce) diff --git a/tests/test_config/test_launch_topology.py b/tests/test_config/test_launch_topology.py index 0e77be1c8..e0fdb94eb 100644 --- a/tests/test_config/test_launch_topology.py +++ b/tests/test_config/test_launch_topology.py @@ -40,7 +40,8 @@ "qwen3-32b-eagle3-online.yaml": 4, "qwen3-4b-dflash-online.yaml": 8, "qwen3-4b-dspark-offline.yaml": 1, - "qwen3-4b-dspark-disaggregated.yaml": 1, + "qwen3-4b-dspark-disaggregated.yaml": 7, + "qwen3-4b-dspark-live.yaml": 7, "qwen3-4b-eagle3-online.yaml": 1, "qwen3-8b-dspark-offline.yaml": 1, "qwen3-8b-dflash-disaggregated.yaml": 4, @@ -73,7 +74,7 @@ "qwen3.6-27b-dflash-multiserver-disaggregated.yaml": 2, "qwen3.6-27b-dflash-online.yaml": 8, "qwen3.6-27b-domino-online.yaml": 8, - "qwen3.6-27b-dspark-disaggregated.yaml": 1, + "qwen3.6-27b-dspark-disaggregated.yaml": 6, "qwq-32b-eagle3-online.yaml": 4, } @@ -104,8 +105,36 @@ "qwen3-4b-dspark-disaggregated.yaml": { "control_dir": "outputs/qwen3-4b-dspark-disaggregated/control", "backend": "mooncake", - "server_urls": ["http://127.0.0.1:30000"], - **LOCAL_MOONCAKE_ENDPOINTS, + "managed_local": { + "trainer_cuda_visible_devices": ["1", "2", "3", "4", "5", "6", "7"], + "mooncake": { + "protocol": "tcp", + "global_segment_size_bytes": 34359738368, + "local_buffer_size_bytes": 1073741824, + }, + "capture_servers": [ + { + "port": 30000, + "cuda_visible_devices": ["0"], + "tp_size": 1, + "mem_fraction_static": 0.7, + } + ], + }, + }, + "qwen3-4b-dspark-live.yaml": { + "control_dir": "outputs/qwen3-4b-dspark-live/control", + "backend": "mooncake", + "live": { + "host": "0.0.0.0", + "port": 8600, + "mooncake": { + "protocol": "tcp", + "global_segment_size_bytes": 34359738368, + "local_buffer_size_bytes": 1073741824, + }, + "trainer_cuda_visible_devices": ["1", "2", "3", "4", "5", "6", "7"], + }, }, "qwen3-8b-dflash-disaggregated.yaml": { "control_dir": "outputs/qwen3-8b-dflash-disaggregated/control", @@ -283,7 +312,7 @@ "control_dir": "outputs/qwen3.6-27b-dspark-disaggregated/control", "backend": "mooncake", "managed_local": { - "trainer_cuda_visible_devices": ["1"], + "trainer_cuda_visible_devices": ["2", "3", "4", "5", "6", "7"], "mooncake": { "protocol": "tcp", "global_segment_size_bytes": 68719476736, @@ -295,7 +324,13 @@ "cuda_visible_devices": ["0"], "tp_size": 1, "mem_fraction_static": 0.7, - } + }, + { + "port": 30001, + "cuda_visible_devices": ["1"], + "tp_size": 1, + "mem_fraction_static": 0.7, + }, ], }, }, @@ -313,7 +348,7 @@ def _recipes() -> dict[str, Path]: class ExampleLaunchTopologyTest(unittest.TestCase): def test_every_recipe_has_the_explicit_golden_topology(self): recipes = _recipes() - self.assertEqual(len(EXPECTED_NPROC_PER_NODE), 64) + self.assertEqual(len(EXPECTED_NPROC_PER_NODE), 65) self.assertEqual(set(recipes), set(EXPECTED_NPROC_PER_NODE)) for filename, nproc_per_node in EXPECTED_NPROC_PER_NODE.items(): diff --git a/tests/test_config/test_recipe_readme.py b/tests/test_config/test_recipe_readme.py index 97e1eac6c..cb401df93 100644 --- a/tests/test_config/test_recipe_readme.py +++ b/tests/test_config/test_recipe_readme.py @@ -8,6 +8,7 @@ DataConfig, DeploymentConfig, DisaggregatedDeploymentConfig, + LiveIntakeDeploymentConfig, ManagedLocalCaptureServerConfig, ManagedLocalMooncakeConfig, ManagedLocalStackConfig, @@ -46,6 +47,7 @@ def test_recipe_readme_names_every_typed_config_field(self): "deployment.disaggregated.managed_local.capture_servers[]", ManagedLocalCaptureServerConfig, ), + ("deployment.disaggregated.live", LiveIntakeDeploymentConfig), ) missing = [] diff --git a/tests/test_config/test_schema.py b/tests/test_config/test_schema.py index 03edbbd6a..289cbb9e6 100644 --- a/tests/test_config/test_schema.py +++ b/tests/test_config/test_schema.py @@ -63,6 +63,17 @@ def _managed_local_payload(*, ep_size: int) -> dict: return payload +def _live_payload() -> dict: + payload = _online_payload("dspark") + payload["data"] = {"max_length": 2048} + payload["deployment"]["disaggregated"] = { + "control_dir": "/control", + "backend": "mooncake", + "live": {"port": 8600}, + } + return payload + + def _write(payload: dict, suffix: str) -> str: fd, path = tempfile.mkstemp(suffix=suffix) with os.fdopen(fd, "w") as f: @@ -193,6 +204,94 @@ def test_exactly_one_data_source(self): raw = Config.model_validate(raw_payload) self.assertEqual(raw.mode, "online") + def test_live_capture_trains_from_serving_traffic_without_a_dataset(self): + config = Config.model_validate(_live_payload()) + self.assertEqual(config.mode, "online") + self.assertEqual(config.deployment.disaggregated.live.host, "0.0.0.0") + self.assertEqual(config.deployment.disaggregated.live.port, 8600) + + def test_live_capture_rejects_a_data_source(self): + payload = _live_payload() + payload["data"]["train_data_path"] = "/conversations.jsonl" + with self.assertRaisesRegex(ValidationError, "do not set a"): + Config.model_validate(payload) + + def test_live_capture_requires_a_bounded_schedule(self): + payload = _live_payload() + payload["training"].pop("max_steps") + with self.assertRaisesRegex(ValidationError, "unbounded stream"): + Config.model_validate(payload) + payload["training"]["total_steps"] = 100 + config = Config.model_validate(payload) + self.assertEqual(config.training.total_steps, 100) + + def test_live_capture_rejects_driven_capture_topologies(self): + payload = _live_payload() + payload["deployment"]["disaggregated"]["server_urls"] = [ + "http://127.0.0.1:30000" + ] + with self.assertRaisesRegex(ValidationError, "server_urls"): + Config.model_validate(payload) + + payload = _live_payload() + payload["deployment"]["disaggregated"]["managed_local"] = { + "trainer_cuda_visible_devices": ["1"], + "capture_servers": [{"port": 30000, "cuda_visible_devices": ["0"]}], + } + with self.assertRaisesRegex(ValidationError, "managed_local"): + Config.model_validate(payload) + + payload = _live_payload() + payload["deployment"]["disaggregated"]["producer_segment_size"] = 1024 + with self.assertRaisesRegex(ValidationError, "producer_segment_size"): + Config.model_validate(payload) + + payload = _live_payload() + payload["deployment"]["disaggregated"]["backend"] = "shared_dir" + payload["deployment"]["disaggregated"]["store_root"] = "/features" + with self.assertRaisesRegex(ValidationError, "backend=mooncake"): + Config.model_validate(payload) + + def test_live_supervised_launch_pairs_mooncake_with_trainer_devices(self): + payload = _live_payload() + payload["deployment"]["trainer"] = {"nnodes": 1, "nproc_per_node": 2} + payload["deployment"]["disaggregated"]["live"].update( + { + "mooncake": {"global_segment_size_bytes": 4096}, + "trainer_cuda_visible_devices": ["1", "2"], + } + ) + config = Config.model_validate(payload) + self.assertEqual( + config.deployment.disaggregated.live.mooncake.rpc_port, 35551 + ) + + missing_devices = _live_payload() + missing_devices["deployment"]["disaggregated"]["live"]["mooncake"] = {} + with self.assertRaisesRegex(ValidationError, "together"): + Config.model_validate(missing_devices) + + wrong_count = _live_payload() + wrong_count["deployment"]["disaggregated"]["live"].update( + {"mooncake": {}, "trainer_cuda_visible_devices": ["1", "2"]} + ) + with self.assertRaisesRegex(ValidationError, "nproc_per_node"): + Config.model_validate(wrong_count) + + explicit_endpoint = _live_payload() + explicit_endpoint["deployment"]["trainer"] = { + "nnodes": 1, + "nproc_per_node": 2, + } + explicit_endpoint["deployment"]["disaggregated"].update( + {"mooncake_metadata_server": "http://metadata:8080/metadata"} + ) + explicit_endpoint["deployment"]["disaggregated"]["live"].update( + {"mooncake": {}, "trainer_cuda_visible_devices": ["1", "2"]} + ) + with self.assertRaisesRegex(ValidationError, "derives Mooncake endpoints"): + Config.model_validate(explicit_endpoint) + def test_offline_eval_source_and_interval_form_one_pair(self): offline = Config.model_validate( { diff --git a/tests/test_config/test_unified_feature_reachability.py b/tests/test_config/test_unified_feature_reachability.py index 2421071ce..eba71be36 100644 --- a/tests/test_config/test_unified_feature_reachability.py +++ b/tests/test_config/test_unified_feature_reachability.py @@ -146,7 +146,7 @@ def test_all_example_configs_validate_through_the_typed_entry(self): for path in EXAMPLE_CONFIG_DIR.glob("*.yaml") if not path.name.startswith(".") ) - self.assertEqual(len(paths), 64) + self.assertEqual(len(paths), 65) resolved_runs = { path.name: resolve_run(Config.from_file(str(path))) for path in paths diff --git a/tests/test_runtime/test_disagg_live_producer.py b/tests/test_runtime/test_disagg_live_producer.py new file mode 100644 index 000000000..3d16d1948 --- /dev/null +++ b/tests/test_runtime/test_disagg_live_producer.py @@ -0,0 +1,257 @@ +# coding=utf-8 +"""End-to-end unit tests for the online-live producer (no GPU, no SGLang). + +A stub "sink" plays the live-patched serving engine: it POSTs capture records +over real HTTP to the producer-hosted :class:`CaptureIntakeServer` while the +REAL live stack runs underneath — ``LiveIntakeRefSource`` validation, +``adopt``, watermark shedding, byte accounting, ``StreamingRefChannel`` +publication, and terminal cleanup via ``_finalize_online_producer``.""" + +import json +import os +import tempfile +import threading +import time +import unittest +from typing import Any, Dict +from urllib.error import HTTPError +from urllib.request import Request, urlopen + +from specforge.inference.adapters.live_intake import LiveIntakeRefSource +from specforge.inference.adapters.server_capture import ServerCaptureSchema +from specforge.inference.capture import CaptureConfig +from specforge.launch import build_disagg_live_producer +from specforge.runtime.data_plane.streaming_ref_channel import StreamingRefChannel + +HIDDEN = 8 +AUX_LAYERS = (2, 5, 8) + + +class _FakeStore: + store_id = "store0" + + def __init__(self): + self.adopted = [] + self.aborted = [] + self.gc_calls = 0 + + def adopt(self, ref): + self.adopted.append(ref.sample_id) + + def abort(self, sample_id, *, reason="aborted"): + self.aborted.append((sample_id, reason)) + + def gc(self): + self.gc_calls += 1 + return {} + + +def _source(store) -> LiveIntakeRefSource: + return LiveIntakeRefSource( + store, + run_id="run0", + algorithm="dspark", + schema=ServerCaptureSchema( + aux_feature="hidden_states", + last_hidden_feature="target_last_hidden_states", + passthrough=( + ("input_ids", "input_ids", ()), + ("loss_mask", "loss_mask", ()), + ), + ), + capture=CaptureConfig.from_strategy( + required_features={ + "input_ids", + "hidden_states", + "loss_mask", + "target_last_hidden_states", + }, + aux_hidden_state_layer_ids=AUX_LAYERS, + target_repr="hidden_state", + target_hidden_size=HIDDEN, + ), + max_num_tokens=64, + ) + + +def _record(sample_id: str, length: int = 5) -> Dict[str, Any]: + aux_width = len(AUX_LAYERS) * HIDDEN + return { + "sample_id": sample_id, + "store_id": "store0", + "gen": 1, + "num_tokens": length, + "aux_layer_ids": list(AUX_LAYERS), + "features": { + "hidden_states": {"shape": [1, length, aux_width], "dtype": "bfloat16"}, + "target_last_hidden_states": { + "shape": [1, length, HIDDEN], + "dtype": "bfloat16", + }, + "input_ids": {"shape": [1, length], "dtype": "int64"}, + "loss_mask": {"shape": [1, length], "dtype": "int64"}, + }, + } + + +class LiveProducerTest(unittest.TestCase): + def setUp(self): + self.work = tempfile.mkdtemp(prefix="live_producer_") + self.channel_path = os.path.join(self.work, "refs.jsonl") + self.channel = StreamingRefChannel(self.channel_path) + self.store = _FakeStore() + + def _build(self, **overrides): + kwargs = dict( + feature_store=self.store, + channel=self.channel, + ref_source=_source(self.store), + intake_host="127.0.0.1", + intake_port=0, + in_flight_high_watermark=4, + in_flight_low_watermark=2, + backpressure_poll_s=0.01, + gc_interval_s=0.05, + ) + kwargs.update(overrides) + return build_disagg_live_producer(**kwargs) + + def _drive_async(self, drive): + result: Dict[str, Any] = {} + + def run(): + try: + result["produced"] = drive(should_stop=self.channel.consumer_stopped) + except BaseException as exc: # noqa: BLE001 — surfaced by the test + result["error"] = exc + + thread = threading.Thread(target=run, daemon=True) + thread.start() + return thread, result + + def _post(self, server, record): + body = json.dumps(record).encode("utf-8") + request = Request( + f"http://127.0.0.1:{server.port}/v1/spec-capture/records", + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with urlopen(request, timeout=5.0) as response: + return response.status, json.load(response) + except HTTPError as exc: + return exc.code, json.load(exc) + + def _await(self, predicate, timeout_s: float = 5.0): + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(0.01) + raise AssertionError("condition not met before timeout") + + def test_pushed_records_are_adopted_published_and_deduped(self): + server, drive = self._build() + self.channel.publish_consumer_quantum(2) + thread, result = self._drive_async(drive) + self._await(lambda: server.port != 0 and drive.stats()["quantum"] == 2) + + status, payload = self._post(server, _record("s0")) + self.assertEqual((status, payload["status"]), (200, "accepted")) + status, payload = self._post(server, _record("s0")) + self.assertEqual((status, payload["status"]), (200, "duplicate")) + status, payload = self._post(server, _record("s1", length=7)) + + self.assertEqual(self.store.adopted, ["s0", "s1"]) + refs = self.channel.poll() + self.assertEqual([ref.sample_id for ref in refs], ["s0", "s1"]) + self.assertEqual(refs[0].metadata["transport"], "sglang_live_capture") + self.assertEqual(refs[1].num_tokens, 7) + + self._await(lambda: self.store.gc_calls > 0) + self.channel.mark_consumer_done() + thread.join(timeout=5.0) + self.assertEqual(result.get("produced"), 2) + self.assertTrue(self.channel.is_closed()) + + def test_invalid_records_are_rejected_without_adoption(self): + server, drive = self._build() + self.channel.publish_consumer_quantum(1) + thread, _result = self._drive_async(drive) + self._await(lambda: drive.stats()["quantum"] == 1) + + wrong_store = _record("bad0") + wrong_store["store_id"] = "other" + status, payload = self._post(server, wrong_store) + self.assertEqual(status, 422) + self.assertIn("store_id", payload["detail"]) + + too_long = _record("bad1", length=65) + status, _ = self._post(server, too_long) + self.assertEqual(status, 422) + + self.assertEqual(self.store.adopted, []) + self.assertEqual(self.channel.poll(), []) + self.channel.mark_consumer_done() + thread.join(timeout=5.0) + + def test_watermark_shed_and_resume_hysteresis(self): + server, drive = self._build( + in_flight_high_watermark=2, in_flight_low_watermark=1 + ) + self.channel.publish_consumer_quantum(1) + thread, _result = self._drive_async(drive) + self._await(lambda: drive.stats()["quantum"] == 1) + + self.assertEqual(self._post(server, _record("s0"))[0], 200) + self.assertEqual(self._post(server, _record("s1"))[0], 200) + # in_flight == high watermark == 2 -> shed until consumption drops it. + status, payload = self._post(server, _record("s2")) + self.assertEqual((status, payload["status"]), (429, "shed")) + + self.channel.mark_consumed(1) # in_flight 1 == low watermark -> resume + self.assertEqual(self._post(server, _record("s3"))[0], 200) + self.assertEqual(drive.stats()["shed"], 1) + + self.channel.mark_consumer_done() + thread.join(timeout=5.0) + + def test_resident_byte_hard_cap_sheds_instead_of_failing(self): + record_bytes = 5 * (len(AUX_LAYERS) * HIDDEN + HIDDEN) * 2 + 2 * 5 * 8 + server, drive = self._build( + feature_store_max_resident_bytes=record_bytes + 1 + ) + self.channel.publish_consumer_quantum(1) + thread, _result = self._drive_async(drive) + self._await(lambda: drive.stats()["quantum"] == 1) + + self.assertEqual(self._post(server, _record("s0"))[0], 200) + status, payload = self._post(server, _record("s1")) + self.assertEqual((status, payload["status"]), (429, "shed")) + self.assertIn("hard cap", payload["detail"]) + + self.channel.mark_consumed(1) # frees the resident bytes + self.assertEqual(self._post(server, _record("s2"))[0], 200) + + self.channel.mark_consumer_done() + thread.join(timeout=5.0) + self.assertEqual(self.store.aborted, []) + + def test_setup_failure_publishes_the_failure_sentinel(self): + _server, drive = self._build(peer_wait_timeout_s=0.05) + with self.assertRaises(TimeoutError): + drive() + self.assertIsNotNone(self.channel.failure()) + + def test_undersized_watermarks_fail_against_the_consumer_quantum(self): + _server, drive = self._build( + in_flight_high_watermark=4, in_flight_low_watermark=2 + ) + self.channel.publish_consumer_quantum(3) + with self.assertRaisesRegex(ValueError, "low watermark"): + drive() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_runtime/test_launch_plan.py b/tests/test_runtime/test_launch_plan.py index e6637a8df..2166baa1c 100644 --- a/tests/test_runtime/test_launch_plan.py +++ b/tests/test_runtime/test_launch_plan.py @@ -658,15 +658,18 @@ def test_managed_local_plan_owns_mooncake_and_multiple_capture_servers(self): self.assertEqual([service.phase for service in plan.services], [0, 1, 1]) self.assertEqual(plan.managed_ports, (35551, 35880, 35903, 30000, 30001)) mooncake = plan.services[0] + venv_master = Path(sys.executable).parent / "mooncake_master" + expected_master = str(venv_master) if venv_master.exists() else "mooncake_master" self.assertEqual( mooncake.command.argv, ( - "mooncake_master", + expected_master, "--enable_http_metadata_server=true", "--http_metadata_server_host=127.0.0.1", "--rpc_port=35551", "--http_metadata_server_port=35880", "--metrics_port=35903", + "--default_kv_lease_ttl=60000", ), ) self.assertEqual(mooncake.readiness.kind, "mooncake") @@ -894,6 +897,62 @@ def test_capture_server_urls_are_required_only_for_producers(self): ) self.assertEqual(consumer.role, "consumer") + def test_live_managed_plan_owns_only_the_mooncake_master(self): + raw = _config(mode="disaggregated", nproc=2).model_dump() + raw["data"] = {} + raw["deployment"]["disaggregated"]["server_urls"] = [] + raw["deployment"]["disaggregated"]["control_dir"] = os.path.join( + tempfile.mkdtemp(prefix="live_managed_"), "attempt-1" + ) + raw["deployment"]["disaggregated"]["live"] = { + "port": 8600, + "mooncake": {"global_segment_size_bytes": 4096}, + "trainer_cuda_visible_devices": ["1", "2"], + } + cfg = Config.model_validate(raw) + plan = build_launch_plan(cfg, config_path="run.yaml", env={}) + self.assertEqual(plan.kind, "managed_supervisor") + self.assertEqual( + [service.command.label for service in plan.services], ["mooncake"] + ) + self.assertEqual(plan.managed_ports, (35551, 35880, 35903, 8600)) + producer, consumer = plan.commands + self.assertEqual(producer.env["CUDA_VISIBLE_DEVICES"], "") + self.assertEqual(consumer.env["CUDA_VISIBLE_DEVICES"], "1,2") + for command in plan.commands: + self.assertEqual( + command.env["MOONCAKE_MASTER_SERVER_ADDR"], "127.0.0.1:35551" + ) + self.assertEqual(command.env["SPECFORGE_MANAGED_LOCAL_CHILD"], "1") + with self.assertRaisesRegex(ValueError, "explicit producer or consumer"): + build_launch_plan( + cfg, + config_path="run.yaml", + env={"SPECFORGE_MANAGED_LOCAL_CHILD": "1"}, + ) + + def test_live_producer_needs_no_capture_server_urls(self): + raw = _config(mode="disaggregated").model_dump() + raw["data"] = {} + raw["deployment"]["disaggregated"]["server_urls"] = [] + raw["deployment"]["disaggregated"]["live"] = {"port": 8600} + cfg = Config.model_validate(raw) + producer = build_launch_plan( + cfg, + config_path="run.yaml", + requested_role="producer", + env=MOONCAKE_ENV, + ) + self.assertEqual(producer.role, "producer") + self.assertEqual(producer.worker_env["DISAGG_CLIENT_SEGMENT_SIZE"], "0") + consumer = build_launch_plan( + cfg, + config_path="run.yaml", + requested_role="consumer", + env=MOONCAKE_ENV, + ) + self.assertEqual(consumer.role, "consumer") + def test_multi_node_auto_fails_instead_of_remote_spawning(self): cfg = _config(mode="disaggregated", nproc=2, nnodes=2) with self.assertRaisesRegex(ValueError, "one trainer node only"): diff --git a/tests/test_runtime/test_live_intake.py b/tests/test_runtime/test_live_intake.py new file mode 100644 index 000000000..b51b9c8e8 --- /dev/null +++ b/tests/test_runtime/test_live_intake.py @@ -0,0 +1,271 @@ +# coding=utf-8 +"""Unit tests for the online-live intake: pushed records -> verified refs. + +No GPU and no real server: records are the JSON dicts a live-patched SGLang +sink would POST after writing tensors into Mooncake. These tests cover the +receiving half — the handshake ``config_payload`` document, record validation +in :class:`LiveIntakeRefSource`, and the :class:`CaptureIntakeServer` HTTP +protocol (dedup, shed, reject).""" + +import json +import unittest +from typing import Any, Dict +from urllib.error import HTTPError +from urllib.request import Request, urlopen + +from specforge.inference.adapters.live_intake import LiveIntakeRefSource +from specforge.inference.adapters.server_capture import ServerCaptureSchema +from specforge.inference.capture import CaptureConfig, CaptureMismatchError +from specforge.runtime.contracts import SampleRef +from specforge.runtime.data_plane.intake_server import CaptureIntakeServer + +HIDDEN = 8 +AUX_LAYERS = (2, 5, 8) +MAX_TOKENS = 64 + + +class _StubStore: + store_id = "store0" + + def __init__(self): + self.adopted = [] + + def adopt(self, ref): + self.adopted.append(ref) + + +def _dspark_schema() -> ServerCaptureSchema: + return ServerCaptureSchema( + aux_feature="hidden_states", + last_hidden_feature="target_last_hidden_states", + passthrough=(("input_ids", "input_ids", ()), ("loss_mask", "loss_mask", ())), + ) + + +def _dspark_contract() -> CaptureConfig: + return CaptureConfig.from_strategy( + required_features={ + "input_ids", + "hidden_states", + "loss_mask", + "target_last_hidden_states", + }, + aux_hidden_state_layer_ids=AUX_LAYERS, + target_repr="hidden_state", + target_hidden_size=HIDDEN, + ) + + +def _source(store=None) -> LiveIntakeRefSource: + return LiveIntakeRefSource( + store or _StubStore(), + run_id="run0", + algorithm="dspark", + schema=_dspark_schema(), + capture=_dspark_contract(), + max_num_tokens=MAX_TOKENS, + target_model_version="target-v1", + ) + + +def _record(length: int = 5, **overrides) -> Dict[str, Any]: + aux_width = len(AUX_LAYERS) * HIDDEN + record = { + "sample_id": "live-abc123", + "store_id": "store0", + "gen": 1, + "num_tokens": length, + "aux_layer_ids": list(AUX_LAYERS), + "features": { + "hidden_states": {"shape": [1, length, aux_width], "dtype": "bfloat16"}, + "target_last_hidden_states": { + "shape": [1, length, HIDDEN], + "dtype": "bfloat16", + }, + "input_ids": {"shape": [1, length], "dtype": "int64"}, + "loss_mask": {"shape": [1, length], "dtype": "int64"}, + }, + } + record.update(overrides) + return record + + +class LiveIntakeRefSourceTest(unittest.TestCase): + def test_config_payload_carries_the_capture_document(self): + payload = _source().config_payload() + self.assertEqual(payload["store_id"], "store0") + self.assertEqual(payload["gen"], 1) + self.assertEqual( + payload["features"], + {"aux": "hidden_states", "last_hidden": "target_last_hidden_states"}, + ) + self.assertEqual( + payload["passthrough"], + [ + {"name": "input_ids", "source": "tokens"}, + {"name": "loss_mask", "source": "ones"}, + ], + ) + self.assertEqual(payload["min_num_tokens"], 2) + self.assertEqual(payload["max_num_tokens"], MAX_TOKENS) + + def test_config_payload_rejects_trailing_passthrough_shapes(self): + schema = ServerCaptureSchema( + aux_feature="hidden_states", + last_hidden_feature=None, + passthrough=(("depths", "depths", (4,)),), + ) + source = LiveIntakeRefSource( + _StubStore(), + run_id="run0", + algorithm="peagle", + schema=schema, + capture=_dspark_contract(), + max_num_tokens=MAX_TOKENS, + ) + with self.assertRaisesRegex(ValueError, "trailing shape"): + source.config_payload() + + def test_record_becomes_a_committed_ready_ref(self): + ref = _source().ref_from_record(_record(length=6, origin="10.0.0.9")) + self.assertIsInstance(ref, SampleRef) + self.assertEqual(ref.sample_id, "live-abc123") + self.assertEqual(ref.num_tokens, 6) + self.assertEqual(ref.feature_store_uri, "mooncake://store0/live-abc123") + self.assertEqual(ref.feature_keys["input_ids"], "live-abc123/input_ids") + self.assertEqual(ref.metadata["transport"], "sglang_live_capture") + self.assertEqual(ref.metadata["server"], "10.0.0.9") + self.assertEqual(ref.metadata["generation"], 1) + self.assertEqual( + ref.feature_specs["target_last_hidden_states"].target_repr, + "hidden_state", + ) + + def test_identity_and_bounds_violations_are_rejected(self): + source = _source() + cases = { + "store_id": _record(store_id="other-store"), + "gen": _record(gen=2), + "num_tokens": _record(num_tokens=0), + "one_token_warmup": _record(num_tokens=1), + "cap": _record(num_tokens=MAX_TOKENS + 1), + "features": _record(features={}), + } + for name, record in cases.items(): + with self.subTest(case=name), self.assertRaises(ValueError): + source.ref_from_record(record) + + def test_contract_violations_are_rejected(self): + source = _source() + extra = _record() + extra["features"]["surprise"] = {"shape": [1, 5], "dtype": "int64"} + with self.assertRaises(CaptureMismatchError): + source.ref_from_record(extra) + + short = _record() + short["features"]["input_ids"]["shape"] = [1, 3] + with self.assertRaisesRegex(CaptureMismatchError, "seq len"): + source.ref_from_record(short) + + wrong_width = _record() + wrong_width["features"]["hidden_states"]["shape"] = [1, 5, HIDDEN] + with self.assertRaisesRegex(CaptureMismatchError, "aux width"): + source.ref_from_record(wrong_width) + + wrong_layers = _record(aux_layer_ids=[1, 2, 3]) + with self.assertRaisesRegex(CaptureMismatchError, "aux-layer id"): + source.ref_from_record(wrong_layers) + + missing_layers = _record(aux_layer_ids=None) + with self.assertRaisesRegex(CaptureMismatchError, "omitted aux-layer"): + source.ref_from_record(missing_layers) + + +def _get(url: str): + with urlopen(url, timeout=5.0) as response: + return response.status, json.load(response) + + +def _post(url: str, payload) -> tuple: + body = json.dumps(payload).encode("utf-8") + request = Request( + url, data=body, headers={"Content-Type": "application/json"}, method="POST" + ) + try: + with urlopen(request, timeout=5.0) as response: + return response.status, json.load(response) + except HTTPError as exc: + return exc.code, json.load(exc) + + +class CaptureIntakeServerTest(unittest.TestCase): + def setUp(self): + self.dispatched = [] + self.reply = ("accepted", "") + + def on_record(record): + self.dispatched.append(record) + return self.reply + + self.server = CaptureIntakeServer( + "127.0.0.1", + 0, + config_payload={"store_id": "store0", "gen": 1}, + on_record=on_record, + ).start() + self.origin = f"http://127.0.0.1:{self.server.port}" + self.addCleanup(self.server.stop) + + def test_config_endpoint_serves_the_handshake_document(self): + status, payload = _get(f"{self.origin}/v1/spec-capture/config") + self.assertEqual(status, 200) + self.assertEqual(payload, {"store_id": "store0", "gen": 1}) + + def test_accepted_record_is_dispatched_once_and_deduped_on_retry(self): + record = {"sample_id": "s0", "num_tokens": 4} + status, payload = _post(f"{self.origin}/v1/spec-capture/records", record) + self.assertEqual((status, payload["status"]), (200, "accepted")) + status, payload = _post(f"{self.origin}/v1/spec-capture/records", record) + self.assertEqual((status, payload["status"]), (200, "duplicate")) + self.assertEqual(len(self.dispatched), 1) + self.assertEqual(self.dispatched[0]["origin"], "127.0.0.1") + self.assertEqual( + self.server.stats(), + {"accepted": 1, "shed": 0, "rejected": 0, "duplicate": 1}, + ) + + def test_shed_and_rejected_records_return_non_2xx_and_are_not_deduped(self): + self.reply = ("shed", "watermark") + status, _ = _post( + f"{self.origin}/v1/spec-capture/records", {"sample_id": "s1"} + ) + self.assertEqual(status, 429) + + self.reply = ("rejected", "bad record") + status, payload = _post( + f"{self.origin}/v1/spec-capture/records", {"sample_id": "s1"} + ) + self.assertEqual(status, 422) + self.assertEqual(payload["detail"], "bad record") + self.assertEqual(len(self.dispatched), 2) + + def test_malformed_requests_never_reach_the_dispatcher(self): + status, _ = _post(f"{self.origin}/v1/spec-capture/records", ["not", "a", "dict"]) + self.assertEqual(status, 400) + status, _ = _post(f"{self.origin}/v1/spec-capture/records", {"sample_id": ""}) + self.assertEqual(status, 400) + status, _ = _post(f"{self.origin}/v1/other", {"sample_id": "s2"}) + self.assertEqual(status, 404) + self.assertEqual(self.dispatched, []) + + def test_dispatcher_crash_maps_to_500(self): + self.server.on_record = lambda record: 1 / 0 + status, payload = _post( + f"{self.origin}/v1/spec-capture/records", {"sample_id": "s3"} + ) + self.assertEqual(status, 500) + self.assertIn("ZeroDivisionError", payload["error"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_runtime/test_sglang_live_capture_patch.py b/tests/test_runtime/test_sglang_live_capture_patch.py new file mode 100644 index 000000000..4795b320b --- /dev/null +++ b/tests/test_runtime/test_sglang_live_capture_patch.py @@ -0,0 +1,171 @@ +# coding=utf-8 +"""Seam tests for the online-live SGLang patch layer (no GPU, no server). + +Runs against the installed sglang tree with the ``patches/sglang/online-live`` +layer applied (``scripts/apply_sglang_spec_capture_patch.sh --live``); skipped +otherwise. Exercises the sink's live logic in isolation: deterministic +per-request sampling, passthrough synthesis, and the remove-on-failed-POST +cleanup contract.""" + +import queue +import unittest +from unittest import mock + +import torch + +try: + from sglang.srt import spec_capture_sink +except ImportError: # pragma: no cover - sglang not installed + spec_capture_sink = None + +LIVE_PATCHED = spec_capture_sink is not None and hasattr( + spec_capture_sink, "maybe_live_spec" +) + +CONFIG = { + "store_id": "store0", + "run_id": "run0", + "gen": 1, + "features": {"aux": "hidden_states", "last_hidden": "target_last_hidden_states"}, + "passthrough": [ + {"name": "input_ids", "source": "tokens"}, + {"name": "loss_mask", "source": "ones"}, + ], + "max_num_tokens": 8, +} + + +def _live_sink(sample_rate: float = 1.0): + sink = spec_capture_sink.SpecCaptureSink( + aux_layer_ids=[2, 5, 8], + intake_url="http://intake:8600", + sample_rate=sample_rate, + ) + sink._live_config = dict(CONFIG) + return sink + + +@unittest.skipUnless(LIVE_PATCHED, "installed sglang lacks the online-live patch") +class LiveSinkSeamTest(unittest.TestCase): + def test_module_hook_is_none_without_a_sink(self): + with mock.patch.object(spec_capture_sink, "_SINK", None): + self.assertIsNone(spec_capture_sink.maybe_live_spec("rid0")) + + def test_live_spec_is_minted_deterministically_from_the_rid(self): + sink = _live_sink() + spec = sink.maybe_live_spec("rid-42") + self.assertEqual( + spec, + { + "store_id": "store0", + "sample_id": "live-rid-42", + "gen": 1, + "replace": False, + "features": dict(CONFIG["features"]), + "live": True, + }, + ) + # Sampling must agree across TP ranks: same rid -> same decision. + low, high = _live_sink(sample_rate=0.5), _live_sink(sample_rate=0.5) + for rid in (f"rid-{i}" for i in range(32)): + self.assertEqual( + low.maybe_live_spec(rid) is None, high.maybe_live_spec(rid) is None + ) + self.assertIsNone(_live_sink(sample_rate=0.0).maybe_live_spec("rid-0")) + + def test_no_intake_url_or_config_means_no_capture(self): + plain = spec_capture_sink.SpecCaptureSink(aux_layer_ids=[2]) + self.assertIsNone(plain.maybe_live_spec("rid-0")) + unfetched = _live_sink() + unfetched._live_config = None + unfetched._live_config_next_fetch = float("inf") # block a real fetch + self.assertIsNone(unfetched.maybe_live_spec("rid-0")) + + def test_write_synthesizes_passthrough_and_posts_the_record(self): + sink = _live_sink() + tokens = [11, 12, 13] + aux = torch.randn(3, 4) + seen = {} + + def fake_put_sample(spec, *, aux, last_hidden): + seen["spec"] = spec + return { + "sample_id": spec["sample_id"], + "store_id": spec["store_id"], + "gen": spec["gen"], + "aux_layer_ids": sink.aux_layer_ids, + "features": {"hidden_states": {"shape": [1, 3, 4], "dtype": "float32"}}, + } + + with ( + mock.patch.object(sink, "put_sample", side_effect=fake_put_sample), + mock.patch.object( + sink, "_post_live_record", return_value=True + ) as post, + ): + sink._write_live_sample( + sink.maybe_live_spec("rid-0"), aux, None, tokens + ) + passthrough = { + item["name"]: item for item in seen["spec"]["passthrough"] + } + self.assertEqual(passthrough["input_ids"]["data"], tokens) + self.assertEqual(passthrough["loss_mask"]["data"], [1, 1, 1]) + self.assertEqual(passthrough["input_ids"]["shape"], [1, 3]) + self.assertEqual(post.call_args.args[0]["num_tokens"], 3) + self.assertEqual(sink._live_stats["captured"], 1) + + def test_failed_post_removes_every_written_key(self): + sink = _live_sink() + removed = [] + result = { + "sample_id": "live-rid-0", + "store_id": "store0", + "gen": 1, + "aux_layer_ids": [2, 5, 8], + "features": { + "hidden_states": {"shape": [1, 2, 4], "dtype": "float32"}, + "input_ids": {"shape": [1, 2], "dtype": "int64"}, + }, + } + with ( + mock.patch.object(sink, "put_sample", return_value=dict(result)), + mock.patch.object(sink, "_post_live_record", return_value=False), + mock.patch.object(sink, "_remove_quiet", side_effect=removed.append), + ): + sink._write_live_sample( + sink.maybe_live_spec("rid-0"), torch.randn(2, 4), None, [1, 2] + ) + self.assertEqual( + sorted(removed), + [ + "store0/live-rid-0/g1/hidden_states", + "store0/live-rid-0/g1/input_ids", + ], + ) + self.assertEqual(sink._live_stats["dropped_post"], 1) + + def test_inconsistent_rows_and_overlong_requests_are_dropped_before_put(self): + sink = _live_sink() + with mock.patch.object(sink, "put_sample") as put: + spec = sink.maybe_live_spec("rid-0") + sink._write_live_sample(spec, torch.randn(3, 4), None, [1, 2]) # 3 != 2 + sink._write_live_sample( + spec, torch.randn(9, 4), None, list(range(9)) + ) # > max_num_tokens + sink._write_live_sample(spec, torch.randn(1, 4), None, [7]) # warmup + put.assert_not_called() + self.assertEqual(sink._live_stats["dropped_bad_rows"], 3) + + def test_full_queue_drops_instead_of_blocking(self): + sink = _live_sink() + sink._live_queue = queue.Queue(maxsize=1) + sink._live_queue.put_nowait(("occupied",)) + sink.put_sample_live( + {"sample_id": "s"}, aux=None, last_hidden=None, tokens=[1] + ) + self.assertEqual(sink._live_stats["dropped_queue_full"], 1) + + +if __name__ == "__main__": + unittest.main()