diff --git a/experiments/bench_gpu_direct_control.py b/experiments/bench_gpu_direct_control.py new file mode 100644 index 000000000..0757264f1 --- /dev/null +++ b/experiments/bench_gpu_direct_control.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +"""Measure GPU-direct release-control round trips without tensor payloads.""" + +from __future__ import annotations + +import argparse +import json +import socketserver +import sys +import time +from pathlib import Path + + +class _ControlHandler(socketserver.StreamRequestHandler): + def handle(self) -> None: + while True: + payload = self.rfile.readline(1 << 20) + if not payload: + return + request = json.loads(payload) + op = request.get("op") + if op in {"release", "abort"}: + released = 1 + elif op in {"release_batch", "abort_batch"}: + released = len(request.get("items", [])) + else: + response = {"ok": False, "error": f"unsupported operation {op!r}"} + self.wfile.write(json.dumps(response).encode() + b"\n") + self.wfile.flush() + continue + self.server.request_count += 1 # type: ignore[attr-defined] + self.server.item_count += released # type: ignore[attr-defined] + self.wfile.write( + json.dumps({"ok": True, "released": released}).encode() + b"\n" + ) + self.wfile.flush() + + +class _ControlServer(socketserver.ThreadingTCPServer): + allow_reuse_address = True + daemon_threads = True + + def __init__(self, address: tuple[str, int]): + self.request_count = 0 + self.item_count = 0 + super().__init__(address, _ControlHandler) + + +def _serve(host: str, port: int) -> None: + with _ControlServer((host, port)) as server: + print(json.dumps({"event": "ready", "host": host, "port": port}), flush=True) + server.serve_forever() + + +def _client( + specforge_root: str, + endpoint: str, + *, + mode: str, + batch_size: int, + iterations: int, +) -> None: + sys.path.insert(0, str(Path(specforge_root).resolve())) + from specforge.runtime.data_plane.gpu_direct_store import _control_request + + items = [ + {"sample_id": f"sample-{index}", "generation": 1} + for index in range(batch_size) + ] + request_count = 0 + started = time.perf_counter() + for _ in range(iterations): + if mode == "legacy": + for item in items: + _control_request( + endpoint, + { + "op": "release", + "token": "benchmark", + **item, + "reason": "benchmark", + }, + ) + request_count += 1 + else: + _control_request( + endpoint, + { + "op": "release_batch", + "token": "benchmark", + "items": items, + "reason": "benchmark", + }, + ) + request_count += 1 + elapsed = time.perf_counter() - started + item_count = batch_size * iterations + print( + json.dumps( + { + "batch_size": batch_size, + "elapsed_s": elapsed, + "item_count": item_count, + "items_per_s": item_count / elapsed, + "iterations": iterations, + "mode": mode, + "request_count": request_count, + }, + sort_keys=True, + ), + flush=True, + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers(dest="command", required=True) + server = subparsers.add_parser("server") + server.add_argument("--host", default="0.0.0.0") + server.add_argument("--port", type=int, default=39100) + client = subparsers.add_parser("client") + client.add_argument("--specforge-root", required=True) + client.add_argument("--endpoint", required=True) + client.add_argument("--mode", choices=("legacy", "batch"), required=True) + client.add_argument("--batch-size", type=int, default=32) + client.add_argument("--iterations", type=int, default=100) + args = parser.parse_args() + if args.command == "server": + _serve(args.host, args.port) + else: + _client( + args.specforge_root, + args.endpoint, + mode=args.mode, + batch_size=args.batch_size, + iterations=args.iterations, + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/bench_live_capture_pipeline.py b/experiments/bench_live_capture_pipeline.py new file mode 100644 index 000000000..ae97102f9 --- /dev/null +++ b/experiments/bench_live_capture_pipeline.py @@ -0,0 +1,268 @@ +#!/usr/bin/env python3 +"""Run a finite DSpark capture, RDMA readback, and durable cleanup pipeline.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import sys +import tempfile +import threading +import time +import uuid +from pathlib import Path + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--specforge-root", required=True) + parser.add_argument("--server-urls", required=True) + parser.add_argument("--dataset-file", required=True) + parser.add_argument("--samples", type=int, default=64) + parser.add_argument("--local-hostname", required=True) + parser.add_argument("--rdma-devices", required=True) + parser.add_argument("--concurrency", type=int, default=32) + parser.add_argument("--lease", type=int, default=1) + parser.add_argument("--prefetch-batches", type=int, default=0) + parser.add_argument("--ingest-batch-size", type=int, default=32) + parser.add_argument( + "--prompt-routing", choices=("shared", "least_tokens"), default="shared" + ) + parser.add_argument( + "--prompt-batching", + choices=("shuffle", "length_bucketed"), + default="shuffle", + ) + parser.add_argument("--materialize", action="store_true") + parser.add_argument( + "--materialize-mode", + choices=("serial", "overlap"), + default="serial", + ) + args = parser.parse_args() + + sys.path.insert(0, str(Path(args.specforge_root).resolve())) + from datasets import Dataset + + from specforge.algorithms.builtin import builtin_algorithm_registry + from specforge.data.prompt_builder import _ProcessedPromptSequence + from specforge.inference.adapters.server_capture import ( + SGLangServerCaptureAdapter, + ServerCaptureSchema, + ) + from specforge.launch import build_disagg_online_producer + from specforge.runtime.data_plane.gpu_direct_store import ( + MooncakeGpuDirectFeatureStore, + ) + from specforge.runtime.data_plane.streaming_ref_channel import ( + StreamingRefChannel, + ) + + if args.samples < 1: + raise ValueError("samples must be positive") + dataset = Dataset.from_file(args.dataset_file) + if args.samples > len(dataset): + raise ValueError(f"requested {args.samples} rows from {len(dataset)}") + prompts = _ProcessedPromptSequence( + dataset.select(range(args.samples)), + max_length=8192, + min_loss_tokens=0, + loss_mask_filter=None, + ) + algorithm = builtin_algorithm_registry().resolve("dspark") + layout = algorithm.providers.server_streaming_for("text").layout + run_id = f"q38-live-capture-{uuid.uuid4().hex[:10]}" + store = MooncakeGpuDirectFeatureStore( + store_id=run_id, + local_hostname=args.local_hostname, + transport="rdma", + rdma_devices=args.rdma_devices, + retain_on_release=True, + ) + schema = ServerCaptureSchema( + aux_feature=layout.aux_feature, + last_hidden_feature=layout.last_hidden_feature, + passthrough=layout.passthrough, + attention_mask_feature=layout.attention_mask_feature, + ) + server_urls = [url.strip().rstrip("/") for url in args.server_urls.split(",")] + server_urls = [url for url in server_urls if url] + adapters = [ + SGLangServerCaptureAdapter( + url, + store, + run_id=run_id, + algorithm=algorithm.name, + schema=schema, + timeout_s=300, + ) + for url in server_urls + ] + workdir = tempfile.mkdtemp(prefix="q38-live-capture-") + channel = StreamingRefChannel(os.path.join(workdir, "refs.jsonl")) + channel.publish_consumer_quantum(1) + optional = {} + if args.prefetch_batches: + optional.update( + producer_prompt_prefetch_batches=args.prefetch_batches, + producer_reorder_buffer=args.concurrency, + ) + _workers, drive = build_disagg_online_producer( + algorithm=algorithm, + feature_source=adapters, + prompts=prompts, + feature_store=store, + channel=channel, + run_id=run_id, + target_hidden_size=5120, + target_repr=None, + aux_hidden_state_layer_ids=[5, 19, 33, 47, 61], + lease=args.lease, + producer_concurrency=args.concurrency, + producer_ordered_publish=True, + producer_prompt_routing=args.prompt_routing, + producer_prompt_batching=args.prompt_batching, + prompt_ingest_batch_size=args.ingest_batch_size, + in_flight_high_watermark=args.samples + 32, + in_flight_low_watermark=args.samples + 16, + backpressure_poll_s=0.01, + peer_wait_timeout_s=300, + prompt_epochs=1, + prompt_seed=42, + **optional, + ) + + refs = [] + materialize_result = { + "active_elapsed_s": 0.0, + "bytes": 0, + "elapsed_s": 0.0, + "error": None, + } + materializer_failed = threading.Event() + + def materialize_refs(ref_stream) -> None: + torch = __import__("torch") + torch.cuda.set_device(0) + started = time.perf_counter() + active_elapsed = 0.0 + try: + for ref in ref_stream: + get_started = time.perf_counter() + tensors, handle = store.get(ref, device="cuda:0") + active_elapsed += time.perf_counter() - get_started + materialize_result["bytes"] += sum( + tensor.numel() * tensor.element_size() + for tensor in tensors.values() + ) + store.release(handle, reason="live-pipeline-readback") + refs.append(ref) + del tensors + sync_started = time.perf_counter() + torch.cuda.synchronize() + active_elapsed += time.perf_counter() - sync_started + except BaseException as exc: + materialize_result["error"] = exc + materializer_failed.set() + try: + channel.mark_consumer_failed( + f"{type(exc).__name__}: {exc}" + ) + except Exception: + pass + finally: + materialize_result["active_elapsed_s"] = active_elapsed + materialize_result["elapsed_s"] = time.perf_counter() - started + + pipeline_started = time.perf_counter() + materializer = None + if args.materialize and args.materialize_mode == "overlap": + materializer = threading.Thread( + target=materialize_refs, + args=(channel.stream(poll_s=0.001, idle_timeout_s=300),), + name="streaming-rdma-materializer", + ) + materializer.start() + + capture_started = time.perf_counter() + capture_error = None + try: + produced = drive(should_stop=materializer_failed.is_set) + except BaseException as exc: + capture_error = exc + produced = channel.published + capture_elapsed = time.perf_counter() - capture_started + + if materializer is not None: + materializer.join() + else: + captured_refs = channel.poll() + if args.materialize: + materialize_refs(iter(captured_refs)) + else: + refs.extend(captured_refs) + + pipeline_elapsed = time.perf_counter() - pipeline_started + if materialize_result["error"] is not None: + raise RuntimeError("streaming RDMA materializer failed") from materialize_result[ + "error" + ] + if capture_error is not None: + raise capture_error + if produced != args.samples or len(refs) != args.samples: + raise RuntimeError( + f"capture count mismatch produced={produced} refs={len(refs)}" + ) + + cleanup_started = time.perf_counter() + sample_ids = [ref.sample_id for ref in refs] + abort_many = getattr(store, "abort_many", None) + if args.prefetch_batches and callable(abort_many): + removed = abort_many(sample_ids, reason="live-pipeline-durable-ack") + else: + for sample_id in sample_ids: + store.abort(sample_id, reason="live-pipeline-durable-ack") + removed = len(sample_ids) + store.drain_pending_removals(max_attempts=8, retry_interval_s=0.25) + cleanup_elapsed = time.perf_counter() - cleanup_started + health = store.health() + print( + json.dumps( + { + "capture_elapsed_s": capture_elapsed, + "capture_samples_per_s": produced / capture_elapsed, + "cleanup_elapsed_s": cleanup_elapsed, + "consumer_health": health, + "id_sha256": hashlib.sha256("\n".join(sample_ids).encode()).hexdigest(), + "materialize_active_elapsed_s": materialize_result[ + "active_elapsed_s" + ], + "materialize_bytes": materialize_result["bytes"], + "materialize_elapsed_s": materialize_result["elapsed_s"], + "materialize_mode": args.materialize_mode, + "materialize_samples_per_s": ( + len(refs) / materialize_result["active_elapsed_s"] + if materialize_result["active_elapsed_s"] + else 0.0 + ), + "pipeline_elapsed_s": pipeline_elapsed, + "pipeline_samples_per_s": produced / pipeline_elapsed, + "lease": args.lease, + "prefetch_batches": args.prefetch_batches, + "prompt_routing": args.prompt_routing, + "prompt_batching": args.prompt_batching, + "produced": produced, + "removed": removed, + "run_id": run_id, + "server_urls": server_urls, + }, + sort_keys=True, + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/bench_prompt_capture_overlap.py b/experiments/bench_prompt_capture_overlap.py new file mode 100644 index 000000000..07ff546d8 --- /dev/null +++ b/experiments/bench_prompt_capture_overlap.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +"""Measure prompt materialization overlap with a deterministic capture stub.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import sys +import tempfile +import time +from pathlib import Path + + +class _DelayedPrompts: + def __init__(self, count: int, delay_s: float) -> None: + self.count = count + self.delay_s = delay_s + + def __len__(self) -> int: + return self.count + + def __getitem__(self, index: int): + if index < 0 or index >= self.count: + raise IndexError(index) + time.sleep(self.delay_s) + length = 16 + index % 8 + return { + "task_id": f"prompt-{index}", + "payload": { + "input_ids": list(range(1, length + 1)), + "loss_mask": [1] * length, + }, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--specforge-root", required=True) + parser.add_argument("--prefetch-batches", type=int, required=True) + parser.add_argument("--prompts", type=int, default=256) + parser.add_argument("--ingest-batch-size", type=int, default=32) + parser.add_argument("--materialize-delay-ms", type=float, default=1.0) + parser.add_argument("--capture-delay-ms", type=float, default=30.0) + parser.add_argument("--concurrency", type=int, default=4) + args = parser.parse_args() + + sys.path.insert(0, str(Path(args.specforge_root).resolve())) + from specforge.algorithms.builtin import builtin_algorithm_registry + from specforge.inference.adapters.server_capture import ( + SGLangServerCaptureAdapter, + ServerCaptureSchema, + ) + from specforge.launch import build_disagg_online_producer + from specforge.runtime.data_plane.mooncake_store import MooncakeFeatureStore + from specforge.runtime.data_plane.streaming_ref_channel import ( + StreamingRefChannel, + ) + from tests.test_runtime.test_server_capture import ( + AUX_LAYERS, + HIDDEN, + _FakeMooncakeStore, + _StubCaptureServer, + ) + + algorithm = builtin_algorithm_registry().resolve("dflash") + layout = algorithm.providers.server_streaming_for("text").layout + backend = _FakeMooncakeStore() + sink = _StubCaptureServer(backend) + + def delayed_post(url, json_body, timeout): + time.sleep(args.capture_delay_ms / 1000.0) + return sink(url, json_body, timeout) + + store = MooncakeFeatureStore(store=backend, store_id="overlap-benchmark") + adapter = SGLangServerCaptureAdapter( + "http://capture-benchmark:30000", + store, + run_id="overlap-benchmark", + 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, + ), + post_fn=delayed_post, + ) + workdir = tempfile.mkdtemp(prefix="specforge-overlap-") + channel = StreamingRefChannel(os.path.join(workdir, "refs.jsonl")) + channel.publish_consumer_quantum(1) + kwargs = {} + if args.prefetch_batches: + kwargs.update( + producer_prompt_prefetch_batches=args.prefetch_batches, + producer_reorder_buffer=args.concurrency, + ) + prompts = _DelayedPrompts( + args.prompts, args.materialize_delay_ms / 1000.0 + ) + _workers, drive = build_disagg_online_producer( + algorithm=algorithm, + feature_source=adapter, + prompts=prompts, + feature_store=store, + channel=channel, + run_id="overlap-benchmark", + target_hidden_size=HIDDEN, + target_repr=None, + aux_hidden_state_layer_ids=AUX_LAYERS, + lease=8, + producer_concurrency=args.concurrency, + producer_ordered_publish=True, + prompt_ingest_batch_size=args.ingest_batch_size, + in_flight_high_watermark=args.prompts + 32, + in_flight_low_watermark=args.prompts + 16, + backpressure_poll_s=0.001, + sleep=lambda delay: time.sleep(min(delay, 0.001)), + **kwargs, + ) + started = time.perf_counter() + produced = drive() + elapsed = time.perf_counter() - started + sample_ids = [ref.sample_id for ref in channel.poll()] + print( + json.dumps( + { + "elapsed_s": elapsed, + "id_sha256": hashlib.sha256("\n".join(sample_ids).encode()).hexdigest(), + "prefetch_batches": args.prefetch_batches, + "produced": produced, + "refs_per_s": produced / elapsed, + }, + sort_keys=True, + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/bench_prompt_materialization.py b/experiments/bench_prompt_materialization.py new file mode 100644 index 000000000..00dd0b964 --- /dev/null +++ b/experiments/bench_prompt_materialization.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +"""Measure one deterministic online prompt-materialization batch.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import pickle +import sys +import time +from pathlib import Path + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--specforge-root", required=True) + parser.add_argument("--dataset-file", required=True) + parser.add_argument("--batch-size", type=int, default=4096) + parser.add_argument("--max-length", type=int, default=8192) + parser.add_argument("--min-loss-tokens", type=int, default=0) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + sys.path.insert(0, str(Path(args.specforge_root).resolve())) + from datasets import Dataset + + from specforge.data.prompt_builder import _ProcessedPromptSequence + from specforge.launch import _iter_epoch_online_prompt_batches + + dataset = Dataset.from_file(args.dataset_file) + prompts = _ProcessedPromptSequence( + dataset, + max_length=args.max_length, + min_loss_tokens=args.min_loss_tokens, + loss_mask_filter=None, + ) + started = time.perf_counter() + batch = next( + _iter_epoch_online_prompt_batches( + prompts, + 0, + 3, + seed=args.seed, + batch_size=args.batch_size, + ) + ) + elapsed = time.perf_counter() - started + encoded = pickle.dumps(batch, protocol=5) + payload_tokens = sum(len(item["payload"]["input_ids"]) for item in batch) + loss_tokens = sum(sum(item["payload"]["loss_mask"]) for item in batch) + print( + json.dumps( + { + "batch_size": len(batch), + "elapsed_s": elapsed, + "payload_tokens": payload_tokens, + "loss_tokens": loss_tokens, + "sha256": hashlib.sha256(encoded).hexdigest(), + }, + sort_keys=True, + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/specforge/application/planning.py b/specforge/application/planning.py index 44830a927..6197148e3 100644 --- a/specforge/application/planning.py +++ b/specforge/application/planning.py @@ -150,10 +150,14 @@ def _validate_training_topology( "requires model.target_backend=sglang" ) deployment = cfg.deployment.disaggregated - if deployment is None or deployment.backend != "mooncake": + if deployment is None or deployment.backend not in { + "mooncake", + "mooncake_gpu_direct", + }: raise ValueError( "online disaggregated training requires " - "deployment.disaggregated.backend=mooncake" + "deployment.disaggregated.backend=mooncake or " + "mooncake_gpu_direct" ) if cfg.model.shard_target_output: raise ValueError( diff --git a/specforge/config/schema.py b/specforge/config/schema.py index 27870d0ca..bf8b6d3f4 100644 --- a/specforge/config/schema.py +++ b/specforge/config/schema.py @@ -199,6 +199,12 @@ class RuntimeConfig(StrictConfigModel): producer_lease: int = Field(default=8, gt=0) producer_concurrency: int = Field(default=1, gt=0) + producer_ordered_publish: bool = False + producer_prompt_prefetch_batches: int = Field(default=1, ge=0) + producer_reorder_buffer: Optional[int] = Field(default=None, ge=0) + producer_prompt_routing: Literal["shared", "least_tokens"] = "shared" + producer_prompt_batching: Literal["shuffle", "length_bucketed"] = "shuffle" + producer_prompt_ingest_batch_size: int = Field(default=4096, gt=0) in_flight_high_watermark: int = Field(default=256, gt=0) in_flight_low_watermark: int = Field(default=192, ge=0) resident_high_watermark_bytes: Optional[int] = Field(default=None, gt=0) @@ -411,7 +417,7 @@ class DisaggregatedDeploymentConfig(StrictConfigModel): #: shared-filesystem requirement between trainer nodes while keeping the #: authority-owned source channel and SQLite/WAL on trainer node 0. inbox_server_url: Optional[str] = None - backend: Literal["shared_dir", "mooncake"] + backend: Literal["shared_dir", "mooncake", "mooncake_gpu_direct"] store_root: Optional[str] = None store_id: Optional[str] = None server_urls: List[str] = Field(default_factory=list) @@ -470,6 +476,16 @@ def _validate_store(self): raise ValueError( "deployment.disaggregated.store_root is required for shared_dir" ) + if ( + self.backend == "mooncake_gpu_direct" + and self.mooncake_protocol is not None + and self.mooncake_protocol + not in ("nvlink", "nvlink_intra", "mnnvl", "rdma") + ): + raise ValueError( + "mooncake_gpu_direct requires mooncake_protocol=nvlink, " + "nvlink_intra, or rdma" + ) if self.managed_local is not None: if self.backend != "mooncake": raise ValueError("managed_local requires backend=mooncake") @@ -788,11 +804,21 @@ def _validate_run_structure(self): mode == "online" and deployment == "disaggregated" and self.deployment.disaggregated is not None - and self.deployment.disaggregated.backend != "mooncake" + and self.deployment.disaggregated.backend + not in {"mooncake", "mooncake_gpu_direct"} ): raise ValueError( "online disaggregated training requires " - "deployment.disaggregated.backend=mooncake" + "deployment.disaggregated.backend=mooncake or mooncake_gpu_direct" + ) + if ( + mode != "online" + and deployment == "disaggregated" + and self.deployment.disaggregated is not None + and self.deployment.disaggregated.backend == "mooncake_gpu_direct" + ): + raise ValueError( + "deployment.disaggregated.backend=mooncake_gpu_direct requires online mode" ) managed_local = ( self.deployment.disaggregated.managed_local diff --git a/specforge/data/prompt_builder.py b/specforge/data/prompt_builder.py index a3461d6c2..7fb302b24 100644 --- a/specforge/data/prompt_builder.py +++ b/specforge/data/prompt_builder.py @@ -32,15 +32,17 @@ def prepare_prompt_tasks( max_prompts: int | None = None, loss_mask_filter: Callable[[Sequence[int]], bool] | None = None, ) -> Sequence[PromptTaskDict]: - """Prepare runtime prompt dictionaries from a JSONL file. + """Prepare runtime prompt dictionaries from JSONL or a disk HF dataset. Each returned item has the control-plane shape ``{"payload": {"input_ids": [...], "loss_mask": [...]}}`` and contains no tensors. Files whose first record contains ``input_ids`` and ``loss_mask`` - are treated as pre-tokenized. Other files are treated as raw conversation - data and processed through :func:`build_eagle3_dataset`, then exposed as a - lazy random-access sequence so large Arrow datasets are not expanded into - Python token lists before rollout starts. + are treated as pre-tokenized. A directory created by + :meth:`datasets.Dataset.save_to_disk` is exposed as a lazy random-access + sequence. Other files are treated as raw conversation data and processed + through :func:`build_eagle3_dataset`, then exposed as a lazy random-access + sequence so large Arrow datasets are not expanded into Python token lists + before rollout starts. ``max_prompts`` caps accepted prompts; ``None`` and ``0`` mean no cap. """ @@ -54,6 +56,31 @@ def prepare_prompt_tasks( loss_mask_filter=loss_mask_filter, ) path_string = os.fspath(path) + limit = None if max_prompts in (None, 0) else max_prompts + if os.path.isdir(path_string): + try: + from datasets import load_from_disk + except ImportError as exc: # pragma: no cover - production dependency + raise ImportError( + "loading a disk prompt dataset requires the datasets package" + ) from exc + dataset = load_from_disk(path_string) + required = {"input_ids", "loss_mask"} + missing = required.difference(dataset.column_names) + if missing: + raise ValueError( + f"disk prompt dataset {path_string!r} is missing columns " + f"{sorted(missing)}" + ) + if limit is not None and limit < len(dataset): + dataset = dataset.select(range(limit)) + return _ProcessedPromptSequence( + dataset, + max_length=max_length, + min_loss_tokens=min_loss_tokens, + loss_mask_filter=loss_mask_filter, + ) + first_record = next(_iter_records(path_string), None) if first_record is None: return [] @@ -68,7 +95,6 @@ def prepare_prompt_tasks( f"loss_mask; missing {missing} in {path_string!r}" ) - limit = None if max_prompts in (None, 0) else max_prompts if has_input_ids: rows = ( (record, f"{path_string}:{line_number}") diff --git a/specforge/inference/adapters/server_capture.py b/specforge/inference/adapters/server_capture.py index 6647c627f..1dfd5391a 100644 --- a/specforge/inference/adapters/server_capture.py +++ b/specforge/inference/adapters/server_capture.py @@ -6,15 +6,14 @@ # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 -"""Server-side spec-capture rollout source (zero-copy Mooncake transport). +"""Server-side spec-capture rollout source (Mooncake tensor transport). An external SGLang server patched with ``patches/sglang/v0.5.14/spec-capture.patch`` runs -the prefill and writes captured features straight into Mooncake in -:class:`MooncakeFeatureStore`'s key layout. Tensors never pass through this -process — the ``/generate`` response's ``meta_info["spec_capture"]`` carries -only key/shape/dtype, from which :meth:`SGLangServerCaptureAdapter.produce_refs` -builds committed-ready ``SampleRef``s. +the prefill and publishes captured features through Mooncake Store or the +GPU-direct TransferEngine path. Tensors never pass through this process; the +``/generate`` response carries object metadata or GPU buffer descriptors from +which :meth:`SGLangServerCaptureAdapter.produce_refs` builds ``SampleRef``s. The server knows only generic artifacts (``aux`` = capture layers concatenated, ``last_hidden`` = post-norm final hidden) plus passthrough @@ -25,6 +24,7 @@ from __future__ import annotations import logging +import threading import uuid from dataclasses import dataclass from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Union @@ -65,14 +65,6 @@ class ServerCaptureFailure: retryable: bool = True -def _default_post(url: str, json_body: Dict[str, Any], timeout: float): - import requests - - resp = requests.post(url, json=json_body, timeout=timeout) - resp.raise_for_status() - return resp.json() - - def _flatten_list_wrappers(value: Any) -> List[Any]: """Flatten list-only response wrappers while leaving row objects intact.""" if not isinstance(value, list): @@ -119,10 +111,9 @@ class SGLangServerCaptureAdapter: are verified against the :class:`CaptureConfig` from their FeatureSpecs alone — same loud extraction-boundary guarantee, no tensor fetch. - ``store`` must be the run's :class:`MooncakeFeatureStore` (its ``store_id`` - namespaces the keys and ``adopt()`` registers each ref so a later - ``abort()``/``gc()`` on the producer side can free server-written objects). - ``post_fn`` is injectable for tests. + ``store`` supplies the run's namespace and adopts server-owned references; + it can be the object-store backend or ``MooncakeGpuDirectFeatureStore``. + ``post_fn`` is injectable for callers that own the HTTP transport. """ def __init__( @@ -169,10 +160,31 @@ def __init__( ) self.request_input_adapter = request_input_adapter self.timeout_s = timeout_s - self.post_fn = post_fn or _default_post + self._post_local = threading.local() + self.post_fn = post_fn or self._pooled_post self.target_model_version = target_model_version self._healthy = True + def _pooled_post(self, url: str, json_body: Dict[str, Any], timeout: float): + """Reuse one keep-alive HTTP connection per capture executor thread.""" + import requests + + session = getattr(self._post_local, "session", None) + if session is None: + session = requests.Session() + adapter = requests.adapters.HTTPAdapter( + pool_connections=1, + pool_maxsize=1, + pool_block=True, + max_retries=0, + ) + session.mount("http://", adapter) + session.mount("https://", adapter) + self._post_local.session = session + response = session.post(url, json=json_body, timeout=timeout) + response.raise_for_status() + return response.json() + # -- request construction ------------------------------------------------- def _sample_id(self, task: PromptTask) -> str: return f"{self.run_id}:{task.task_id}" @@ -241,7 +253,7 @@ def _spec_capture_payload(self, task: PromptTask) -> Dict[str, Any]: "dtype": "int64", } ) - return { + payload = { "store_id": self.store.store_id, "sample_id": self._sample_id(task), # A task id is unique within a run, and retries happen only before @@ -254,6 +266,10 @@ def _spec_capture_payload(self, task: PromptTask) -> Dict[str, Any]: "features": features, "passthrough": passthrough, } + transport = getattr(self.store, "transport", None) + if transport is not None: + payload["transport"] = str(transport) + return payload # -- ref construction ------------------------------------------------------ def _ref_from_result( @@ -278,6 +294,24 @@ def _ref_from_result( } specs[name] = FeatureSpec(name=name, shape=shape, dtype=dtype, **extra) nbytes += _spec_nbytes(shape, dtype) + gpu_direct = "session_id" in result + if gpu_direct: + required = ( + "transport", + "session_id", + "control_endpoint", + "control_token", + ) + missing = [name for name in required if not result.get(name)] + if missing: + raise RuntimeError( + f"GPU-direct capture result is missing {missing}" + ) + for name, meta in feats.items(): + if int(meta.get("address", 0)) <= 0 or int(meta.get("nbytes", 0)) <= 0: + raise RuntimeError( + f"GPU-direct feature {name!r} has an invalid buffer descriptor" + ) num_tokens = int(task.metadata.get("num_tokens", 0)) or len( task.payload["input_ids"] ) @@ -285,7 +319,11 @@ def _ref_from_result( sample_id=sample_id, run_id=self.run_id, source_task_id=task.task_id, - feature_store_uri=f"mooncake://{result['store_id']}/{sample_id}", + feature_store_uri=( + f"mooncake+gpu://{result['session_id']}/{result['store_id']}/{sample_id}" + if gpu_direct + else f"mooncake://{result['store_id']}/{sample_id}" + ), feature_keys={n: f"{sample_id}/{n}" for n in specs}, feature_specs=specs, strategy=self.strategy, @@ -300,9 +338,32 @@ def _ref_from_result( "strategy": self.strategy, "target_repr": capture.target_repr, "vocab_map_version": capture.vocab_map_version, - "transport": "sglang_server_capture", + "transport": ( + "sglang_server_capture_gpu_direct" + if gpu_direct + else "sglang_server_capture" + ), "server": self.base_url, # which server captured it (provenance) "generation": gen, # the zero-copy get() locator + **( + { + "mooncake_gpu_direct": { + "transport": str(result["transport"]), + "session_id": str(result["session_id"]), + "control_endpoint": str(result["control_endpoint"]), + "control_token": str(result["control_token"]), + "features": { + name: { + "address": int(meta["address"]), + "nbytes": int(meta["nbytes"]), + } + for name, meta in feats.items() + }, + } + } + if gpu_direct + else {} + ), }, ) diff --git a/specforge/launch.py b/specforge/launch.py index 001c56696..9cd71fe93 100644 --- a/specforge/launch.py +++ b/specforge/launch.py @@ -455,9 +455,10 @@ def _iter_epoch_online_prompt_batches( """Yield a shuffled epoch while bounding expanded token-list residency.""" indices = _epoch_prompt_indices(prompts, epoch, seed=seed) for start in range(0, len(indices), batch_size): + batch_indices = indices[start : start + batch_size] yield [ _epoch_online_prompt(prompts[index], index, epoch, prompt_epochs) - for index in indices[start : start + batch_size] + for index in batch_indices ] @@ -807,6 +808,11 @@ def build_disagg_online_producer( feature_source=None, lease: int = 8, producer_concurrency: int = 1, + producer_ordered_publish: bool = False, + producer_prompt_prefetch_batches: int = 1, + producer_reorder_buffer: Optional[int] = None, + producer_prompt_routing: str = "shared", + producer_prompt_batching: str = "shuffle", in_flight_high_watermark: int = 256, in_flight_low_watermark: Optional[int] = None, resident_high_watermark_bytes: Optional[int] = None, @@ -846,6 +852,13 @@ def build_disagg_online_producer( keeping a reconstructed plan stable across restarts. Prompt payloads are normalized and ingested in ``prompt_ingest_batch_size`` chunks so a large memory-mapped dataset does not expand every token list before rollout. + ``producer_ordered_publish`` retains concurrent capture while publishing + completed request batches in submission order. ``producer_reorder_buffer`` + bounds completed out-of-order requests separately from active capture calls, + preventing one long request from collapsing target concurrency while + preserving request order across paired runs. A positive + ``producer_prompt_prefetch_batches`` keeps rollout workers alive and prepares + and ingests prompt chunks on a background feeder while capture is running. Failure semantics: a worker whose source raises (dead/unreachable server) has already failed its leases retryable — the surviving workers re-lease @@ -901,6 +914,23 @@ def elapsed(start: float) -> str: producer_concurrency = int(producer_concurrency) if producer_concurrency < 1: raise ValueError("producer_concurrency must be >= 1") + producer_prompt_prefetch_batches = int(producer_prompt_prefetch_batches) + if producer_prompt_prefetch_batches < 0: + raise ValueError("producer_prompt_prefetch_batches must be >= 0") + if producer_reorder_buffer is None: + producer_reorder_buffer = producer_concurrency + producer_reorder_buffer = int(producer_reorder_buffer) + if producer_reorder_buffer < 0: + raise ValueError("producer_reorder_buffer must be >= 0") + if producer_prompt_routing not in ("shared", "least_tokens"): + raise ValueError( + "producer_prompt_routing must be either 'shared' or 'least_tokens'" + ) + if producer_prompt_batching not in ("shuffle", "length_bucketed"): + raise ValueError( + "producer_prompt_batching must be either 'shuffle' or " + "'length_bucketed'" + ) prompt_ingest_batch_size = int(prompt_ingest_batch_size) if prompt_ingest_batch_size < 1: raise ValueError("prompt_ingest_batch_size must be >= 1") @@ -940,6 +970,10 @@ def elapsed(start: float) -> str: f"prompt_ingest_batch_size={prompt_ingest_batch_size} " f"lease={worker_lease} workers={num_rollout_workers} " f"concurrency={producer_concurrency} " + f"ordered_publish={producer_ordered_publish} " + f"prompt_prefetch_batches={producer_prompt_prefetch_batches} " + f"prompt_batching={producer_prompt_batching} " + f"reorder_buffer={producer_reorder_buffer} " f"watermarks={in_flight_high_watermark}/" f"{flow_control.limits.resolved_low_watermark_refs}" ) @@ -949,6 +983,7 @@ def elapsed(start: float) -> str: metadata_store=NoOpMetadataStore(), max_prompt_attempts=max_prompt_attempts, enable_sample_queue=False, + prompt_routing=producer_prompt_routing, ) producer_timing(f"DataFlowController created elapsed={elapsed(phase)}") @@ -993,6 +1028,10 @@ def drive_producer(max_rounds: int = 1_000_000, should_stop=None) -> int: "drive_producer enter " f"workers={len(workers)} lease={worker_lease} " f"concurrency={producer_concurrency} max_rounds={max_rounds} " + f"prompt_prefetch_batches={producer_prompt_prefetch_batches} " + f"reorder_buffer={producer_reorder_buffer} " + f"prompt_routing={producer_prompt_routing} " + f"prompt_batching={producer_prompt_batching} " f"watermarks={in_flight_high_watermark}/" f"{flow_control.limits.resolved_low_watermark_refs} " f"progress_interval={progress_interval}" @@ -1059,12 +1098,20 @@ def drive_producer(max_rounds: int = 1_000_000, should_stop=None) -> int: published_sizes = deque() last_publish_log = {"t": time.perf_counter()} dead: dict = {} # worker_id -> last failure reason + prompt_feed_done = threading.Event() + prompt_feed_abort = threading.Event() + if producer_prompt_prefetch_batches == 0: + prompt_feed_done.set() def pool_drained() -> bool: st = controller.status() # leased counts too: a peer's in-flight lease may fail retryable # and come back — leaving then would strand it. - return st["prompts_pending"] == 0 and st["prompts_leased"] == 0 + return ( + prompt_feed_done.is_set() + and st["prompts_pending"] == 0 + and st["prompts_leased"] == 0 + ) def reconcile_consumed_locked() -> int: consumed = channel.consumed_remote() @@ -1146,23 +1193,27 @@ def publish_refs(w, refs, run_once_start: float) -> None: state["first_ref_logged"] = True last_publish_log["t"] = now - def abort_unpublished(futures) -> None: + def abort_unpublished(futures, ready=()) -> None: """Drain active capture calls and free refs after a fatal driver error.""" + unpublished = [] for future in futures: try: - refs = future.result() + unpublished.extend(future.result()) except BaseException: continue - for ref in refs: - try: - feature_store.abort( - ref.sample_id, - reason="producer-driver-failed-before-publication", - ) - except Exception: - logger.exception( - "failed to abort unpublished ref %s", ref.sample_id - ) + for _request_start, refs, exc in ready: + if exc is None: + unpublished.extend(refs) + for ref in unpublished: + try: + feature_store.abort( + ref.sample_id, + reason="producer-driver-failed-before-publication", + ) + except Exception: + logger.exception( + "failed to abort unpublished ref %s", ref.sample_id + ) def run_worker(w) -> None: failures = 0 @@ -1171,6 +1222,35 @@ def run_worker(w) -> None: backpressure_started = None last_backpressure_log = 0.0 futures = {} + ready = {} + next_publish_sequence = 0 + + def finish_request(request_start, refs, exc) -> None: + nonlocal failures, accepting + if exc is not None: + # The worker already failed this call's leases retryable. + # Other active calls keep draining. + failures += 1 + logger.warning( + "rollout worker %s capture call failed (%d/%d): %s", + w.worker_id, + failures, + max_worker_failures, + exc, + ) + if failures >= max_worker_failures: + accepting = False + dead[w.worker_id] = str(exc) + logger.error( + "dropping rollout worker %s after %d consecutive " + "capture failures; health=%s", + w.worker_id, + failures, + w.health(), + ) + return + failures = 0 + publish_refs(w, refs, request_start) with concurrent.futures.ThreadPoolExecutor( max_workers=producer_concurrency, @@ -1207,7 +1287,7 @@ def run_worker(w) -> None: producer_timing( "backpressure wait " f"worker={w.worker_id} " - f"active_requests={len(futures)} " + f"active_requests={len(futures) + len(ready)} " f"produced={state['produced']} " f"in_flight={in_flight} " f"resident_bytes={current_resident_bytes} " @@ -1236,6 +1316,13 @@ def run_worker(w) -> None: and not paused and submitted < max_rounds and len(futures) < producer_concurrency + and len(futures) + len(ready) + < producer_concurrency + + ( + producer_reorder_buffer + if producer_ordered_publish + else 0 + ) and status["prompts_pending"] > 0 ): request_start = time.perf_counter() @@ -1243,11 +1330,11 @@ def run_worker(w) -> None: w.run_once, max_tasks=worker_lease, ) - futures[future] = request_start + futures[future] = (submitted, request_start) submitted += 1 status = controller.status() - if not futures: + if not futures and not ready: if pool_drained(): producer_timing( "pool drained " @@ -1267,38 +1354,32 @@ def run_worker(w) -> None: return_when=concurrent.futures.FIRST_COMPLETED, ) for future in done: - request_start = futures.pop(future) + sequence, request_start = futures.pop(future) try: refs = future.result() except Exception as exc: - # The worker already failed this call's leases - # retryable. Other active calls keep draining. - failures += 1 - logger.warning( - "rollout worker %s capture call failed (%d/%d): %s", - w.worker_id, - failures, - max_worker_failures, - exc, - ) - if failures >= max_worker_failures: - accepting = False - dead[w.worker_id] = str(exc) - logger.error( - "dropping rollout worker %s after %d " - "consecutive capture failures; health=%s", - w.worker_id, - failures, - w.health(), - ) - continue - failures = 0 - publish_refs(w, refs, request_start) + outcome = (request_start, (), exc) + else: + outcome = (request_start, refs, None) + if producer_ordered_publish: + ready[sequence] = outcome + else: + finish_request(*outcome) + if producer_ordered_publish: + while next_publish_sequence in ready: + outcome = ready.pop(next_publish_sequence) + next_publish_sequence += 1 + finish_request(*outcome) except BaseException: - abort_unpublished(futures) + abort_unpublished(futures, ready.values()) raise def ingest_prompt_batch(epoch: int, batch_index: int, epoch_prompts) -> None: + if producer_prompt_batching == "length_bucketed": + epoch_prompts.sort( + key=lambda prompt: len(prompt["payload"]["input_ids"]), + reverse=True, + ) phase = time.perf_counter() producer_timing( "controller.ingest_prompts start " @@ -1322,6 +1403,7 @@ def run_worker_guarded(w) -> None: run_worker(w) except BaseException as exc: # e.g. a channel publish failure fatal.append((w.worker_id, exc)) + prompt_feed_abort.set() if len(live_workers) == 1: run_worker(live_workers[0]) @@ -1342,51 +1424,127 @@ def run_worker_guarded(w) -> None: if fatal: raise fatal[0][1] + def run_prefetched_prompt_stream(live_workers) -> None: + """Feed prompt chunks concurrently while one persistent worker pool runs.""" + feeder_failures: list = [] + max_outstanding = ( + producer_prompt_prefetch_batches * prompt_ingest_batch_size + ) + + def feed_prompts() -> None: + try: + for epoch in range(prompt_epochs): + if prompt_feed_abort.is_set() or ( + should_stop is not None and should_stop() + ): + return + epoch_batches = _iter_epoch_online_prompt_batches( + prompts, + epoch, + prompt_epochs, + seed=prompt_seed, + batch_size=prompt_ingest_batch_size, + ) + queued = 0 + for batch_index, prompt_batch in enumerate(epoch_batches): + while True: + if prompt_feed_abort.is_set() or ( + should_stop is not None and should_stop() + ): + return + if len(dead) >= len(workers): + raise RuntimeError( + "all rollout workers were dropped while " + "the prompt feeder still had work" + ) + status = controller.status() + outstanding = ( + status["prompts_pending"] + + status["prompts_leased"] + ) + if outstanding <= max_outstanding: + break + sleep(backpressure_poll_s) + ingest_prompt_batch(epoch, batch_index, prompt_batch) + queued += len(prompt_batch) + producer_timing( + "epoch enqueued " + f"epoch={epoch + 1}/{prompt_epochs} " + f"prompts={queued} elapsed={elapsed(drive_start)}" + ) + except BaseException as exc: + feeder_failures.append(exc) + finally: + prompt_feed_done.set() + + feeder = threading.Thread( + target=feed_prompts, + name="prompt-feeder", + daemon=True, + ) + feeder.start() + try: + run_epoch_workers(live_workers) + except BaseException: + prompt_feed_abort.set() + feeder.join() + raise + feeder.join() + if feeder_failures: + raise feeder_failures[0] + try: live_workers = list(workers) - for epoch in range(prompt_epochs): - if should_stop is not None and should_stop(): - break - epoch_batches = _iter_epoch_online_prompt_batches( - prompts, - epoch, - prompt_epochs, - seed=prompt_seed, - batch_size=prompt_ingest_batch_size, - ) - stopped = False - for batch_index, prompt_batch in enumerate(epoch_batches): + if producer_prompt_prefetch_batches > 0: + run_prefetched_prompt_stream(live_workers) + else: + for epoch in range(prompt_epochs): if should_stop is not None and should_stop(): - stopped = True break - ingest_prompt_batch(epoch, batch_index, prompt_batch) - if not live_workers: - raise RuntimeError( - f"all rollout workers were already dropped before " - f"epoch {epoch + 1}/{prompt_epochs} batch " - f"{batch_index + 1} could run — dead workers: {dead}" - ) - run_epoch_workers(live_workers) - stopped = should_stop is not None and should_stop() - live_workers = [w for w in live_workers if w.worker_id not in dead] - if dead and not stopped and not pool_drained(): - raise RuntimeError( - f"all rollout workers exited with {len(dead)} dropped " - f"as dead and prompts remaining — dead workers: {dead}" - ) + epoch_batches = _iter_epoch_online_prompt_batches( + prompts, + epoch, + prompt_epochs, + seed=prompt_seed, + batch_size=prompt_ingest_batch_size, + ) + stopped = False + for batch_index, prompt_batch in enumerate(epoch_batches): + if should_stop is not None and should_stop(): + stopped = True + break + ingest_prompt_batch(epoch, batch_index, prompt_batch) + if not live_workers: + raise RuntimeError( + f"all rollout workers were already dropped before " + f"epoch {epoch + 1}/{prompt_epochs} batch " + f"{batch_index + 1} could run — dead workers: {dead}" + ) + run_epoch_workers(live_workers) + stopped = should_stop is not None and should_stop() + live_workers = [ + w for w in live_workers if w.worker_id not in dead + ] + if dead and not stopped and not pool_drained(): + raise RuntimeError( + f"all rollout workers exited with {len(dead)} " + f"dropped as dead and prompts remaining — " + f"dead workers: {dead}" + ) + if stopped: + break if stopped: break - if stopped: - break - st = controller.status() - producer_timing( - "epoch drained " - f"epoch={epoch + 1}/{prompt_epochs} " - f"produced={state['produced']} " - f"prompts_failed={st['prompts_failed']} " - f"pending={st['prompts_pending']} leased={st['prompts_leased']} " - f"elapsed={elapsed(drive_start)}" - ) + st = controller.status() + producer_timing( + "epoch drained " + f"epoch={epoch + 1}/{prompt_epochs} " + f"produced={state['produced']} " + f"prompts_failed={st['prompts_failed']} " + f"pending={st['prompts_pending']} " + f"leased={st['prompts_leased']} " + f"elapsed={elapsed(drive_start)}" + ) st = controller.status() stopped = should_stop is not None and should_stop() if st["prompts_failed"] and not stopped: diff --git a/specforge/launch_plan.py b/specforge/launch_plan.py index 5633c3af9..f7481e108 100644 --- a/specforge/launch_plan.py +++ b/specforge/launch_plan.py @@ -272,7 +272,7 @@ def _disaggregated_env( if deployment.inbox_server_url: values["DISAGG_INBOX_SERVER_URL"] = deployment.inbox_server_url if cfg.mode == "online": - if deployment.backend != "mooncake": + if deployment.backend not in {"mooncake", "mooncake_gpu_direct"}: raise ValueError("online disaggregated training requires Mooncake") # Online feature objects are allocated by the external capture server. # SpecForge roles only read or publish references to those objects. @@ -340,6 +340,15 @@ def _disaggregated_env( "Mooncake endpoints must be provided by deployment config or " f"environment: {missing}" ) + if deployment.backend == "mooncake_gpu_direct": + protocol = values.get("MOONCAKE_PROTOCOL") or base_env.get( + "MOONCAKE_PROTOCOL" + ) + if protocol not in {"nvlink", "nvlink_intra", "mnnvl", "rdma"}: + raise ValueError( + "mooncake_gpu_direct requires MOONCAKE_PROTOCOL=nvlink, " + "nvlink_intra, or rdma" + ) return values diff --git a/specforge/runtime/control_plane/controller.py b/specforge/runtime/control_plane/controller.py index 53a577f50..1bca0a739 100644 --- a/specforge/runtime/control_plane/controller.py +++ b/specforge/runtime/control_plane/controller.py @@ -44,8 +44,14 @@ def __init__( metadata_store: Optional[MetadataStore] = None, max_prompt_attempts: Optional[int] = None, enable_sample_queue: bool = True, + prompt_routing: str = "shared", ) -> None: + if prompt_routing not in ("shared", "least_tokens"): + raise ValueError( + "prompt_routing must be either 'shared' or 'least_tokens'" + ) self.run_id = run_id + self.prompt_routing = prompt_routing self.sample_queue = SampleRefQueue() if enable_sample_queue else None self.store = metadata_store or InMemoryMetadataStore() # Retryable-failure bound: a task failed this many attempts goes @@ -54,6 +60,10 @@ def __init__( self.max_prompt_attempts = max_prompt_attempts self._prompts: "OrderedDict[str, PromptTask]" = OrderedDict() self._prompt_pending: Deque[str] = deque() + self._prompt_pending_by_worker: Dict[str, Deque[str]] = {} + self._prompt_route: Dict[str, str] = {} + self._prompt_cost: Dict[str, int] = {} + self._worker_outstanding_cost: Dict[str, int] = {} self._prompt_leased: Dict[str, str] = {} # task_id -> worker_id self._prompt_failed: Dict[str, str] = {} # task_id -> terminal reason self._workers: Dict[str, Dict[str, Any]] = {} @@ -66,6 +76,8 @@ def register_rollout_worker(self, info: Dict[str, Any]) -> str: worker_id = info.get("worker_id") or f"rollout-{uuid.uuid4().hex[:8]}" with self._lock: self._workers[worker_id] = dict(info) + self._prompt_pending_by_worker.setdefault(worker_id, deque()) + self._worker_outstanding_cost.setdefault(worker_id, 0) return worker_id def register_trainer(self, info: Dict[str, Any]) -> str: @@ -104,22 +116,57 @@ def ingest_prompts(self, prompts: List[Dict[str, Any]]) -> List[str]: # Validation and object construction can be expensive for large prompt # batches, so keep them outside the controller's global state lock. with self._lock: + worker_ids = sorted(self._workers) + if self.prompt_routing == "least_tokens" and not worker_ids: + raise RuntimeError( + "least_tokens prompt routing requires a registered rollout worker" + ) for task in prepared: self._prompts[task.task_id] = task - self._prompt_pending.append(task.task_id) + if self.prompt_routing == "shared": + self._prompt_pending.append(task.task_id) + continue + input_ids = task.payload.get("input_ids", ()) + cost = max(1, len(input_ids)) + worker_id = min( + worker_ids, + key=lambda candidate: ( + self._worker_outstanding_cost[candidate], + candidate, + ), + ) + self._prompt_route[task.task_id] = worker_id + self._prompt_cost[task.task_id] = cost + self._worker_outstanding_cost[worker_id] += cost + self._prompt_pending_by_worker[worker_id].append(task.task_id) return task_ids def lease_prompt_tasks(self, worker_id: str, max_tasks: int) -> List[PromptTask]: out: List[PromptTask] = [] with self._lock: + pending = ( + self._prompt_pending_by_worker.setdefault(worker_id, deque()) + if self.prompt_routing == "least_tokens" + else self._prompt_pending + ) for _ in range(max_tasks): - if not self._prompt_pending: + if not pending: break - task_id = self._prompt_pending.popleft() + task_id = pending.popleft() self._prompt_leased[task_id] = worker_id out.append(self._prompts[task_id]) return out + def _retire_prompt_locked(self, task_id: str) -> None: + worker_id = self._prompt_route.pop(task_id, None) + cost = self._prompt_cost.pop(task_id, 0) + if worker_id is not None: + self._worker_outstanding_cost[worker_id] = max( + 0, self._worker_outstanding_cost.get(worker_id, 0) - cost + ) + self._prompt_leased.pop(task_id, None) + self._prompts.pop(task_id, None) + def complete_prompt_tasks(self, worker_id: str, task_ids: List[str]) -> None: """Retire successfully captured prompts that belong to another TP rank. @@ -133,8 +180,7 @@ def complete_prompt_tasks(self, worker_id: str, task_ids: List[str]) -> None: owner = self._prompt_leased.get(task_id) if owner is not None and owner != worker_id: continue - self._prompt_leased.pop(task_id, None) - self._prompts.pop(task_id, None) + self._retire_prompt_locked(task_id) def fail_prompt_tasks( self, worker_id: str, task_ids: List[str], reason: str, retryable: bool @@ -161,16 +207,21 @@ def fail_prompt_tasks( self._prompts[task_id] = dataclasses.replace( task, attempt=task.attempt + 1 ) - if task_id not in self._prompt_pending: + if self.prompt_routing == "least_tokens": + route = self._prompt_route[task_id] + pending = self._prompt_pending_by_worker[route] + if task_id not in pending: + pending.append(task_id) + elif task_id not in self._prompt_pending: self._prompt_pending.append(task_id) elif retryable: self._prompt_failed[task_id] = ( f"{reason} (attempts exhausted: {task.attempt + 1})" ) - self._prompts.pop(task_id, None) + self._retire_prompt_locked(task_id) else: self._prompt_failed[task_id] = reason - self._prompts.pop(task_id, None) + self._retire_prompt_locked(task_id) def commit_samples(self, worker_id: str, refs: List[SampleRef]) -> List[SampleRef]: """Commit refs and return the subset newly accepted by the ledger. @@ -194,8 +245,7 @@ def commit_samples(self, worker_id: str, refs: List[SampleRef]) -> List[SampleRe with self._lock: for ref in fresh: if ref.source_task_id is not None: - self._prompt_leased.pop(ref.source_task_id, None) - self._prompts.pop(ref.source_task_id, None) + self._retire_prompt_locked(ref.source_task_id) if fresh and self.sample_queue is not None: self.sample_queue.put(fresh) return fresh @@ -271,7 +321,11 @@ def reconcile_on_restart( def status(self) -> Dict[str, Any]: with self._lock: prompts = len(self._prompts) - pending = len(self._prompt_pending) + pending = ( + sum(len(queue) for queue in self._prompt_pending_by_worker.values()) + if self.prompt_routing == "least_tokens" + else len(self._prompt_pending) + ) leased = len(self._prompt_leased) failed = len(self._prompt_failed) workers = len(self._workers) diff --git a/specforge/runtime/control_plane/dp_ack.py b/specforge/runtime/control_plane/dp_ack.py index ae829b521..cb11de943 100644 --- a/specforge/runtime/control_plane/dp_ack.py +++ b/specforge/runtime/control_plane/dp_ack.py @@ -178,13 +178,29 @@ def ack_train_refs( self._cleanup_boundary += 1 boundary = self._cleanup_boundary failures = [] - for sample_id in local_ids: + abort_many = getattr(self.feature_store, "abort_many", None) + if callable(abort_many): try: - self.feature_store.abort( - sample_id, reason="optimizer-boundary-durable-ack" + abort_many( + local_ids, + reason="optimizer-boundary-durable-ack", ) except BaseException as exc: - failures.append(f"{sample_id}: {type(exc).__name__}: {exc}") + failures.append( + "batched optimizer-boundary cleanup: " + f"{type(exc).__name__}: {exc}" + ) + else: + for sample_id in local_ids: + try: + self.feature_store.abort( + sample_id, reason="optimizer-boundary-durable-ack" + ) + except BaseException as exc: + failures.append( + f"{sample_id}: {type(exc).__name__}: {exc}" + ) + for sample_id in local_ids: self._cleanup_pending.setdefault(sample_id, boundary) eligible_ids = [ diff --git a/specforge/runtime/data_plane/__init__.py b/specforge/runtime/data_plane/__init__.py index 90eaf9b5f..7a98679bd 100644 --- a/specforge/runtime/data_plane/__init__.py +++ b/specforge/runtime/data_plane/__init__.py @@ -21,6 +21,7 @@ "list_feature_files", "SharedDirFeatureStore", "MooncakeFeatureStore", + "MooncakeGpuDirectFeatureStore", "AuthPolicy", ] @@ -38,6 +39,7 @@ "list_feature_files": "offline_reader", "SharedDirFeatureStore": "disaggregated", "MooncakeFeatureStore": "mooncake_store", + "MooncakeGpuDirectFeatureStore": "gpu_direct_store", "AuthPolicy": "disaggregated", } diff --git a/specforge/runtime/data_plane/gpu_direct_store.py b/specforge/runtime/data_plane/gpu_direct_store.py new file mode 100644 index 000000000..33f0855f0 --- /dev/null +++ b/specforge/runtime/data_plane/gpu_direct_store.py @@ -0,0 +1,603 @@ +# coding=utf-8 +# Copyright 2024 The SpecForge team. All rights reserved. +# Licensed under the Apache License, Version 2.0 +"""Mooncake GPU-direct materialization for SGLang server captures. + +The tensor payload moves from the capture server's CUDA allocation directly +into trainer CUDA allocations. TCP is restricted to release acknowledgements. +""" + +from __future__ import annotations + +import concurrent.futures +import json +import os +import socket +import threading +import time +import uuid +from typing import Any, Callable, Dict, List, Optional, Tuple + +import torch + +from specforge.runtime.contracts import FeatureHandle, SampleRef +from specforge.runtime.data_plane.feature_store import FeatureStore + +_TORCH_DTYPES = { + "float32": torch.float32, + "float64": torch.float64, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "int64": torch.int64, + "int32": torch.int32, + "int16": torch.int16, + "int8": torch.int8, + "uint8": torch.uint8, + "bool": torch.bool, +} + + +def _nbytes(tensor: torch.Tensor) -> int: + return tensor.numel() * tensor.element_size() + + +def _control_request(endpoint: str, payload: Dict[str, Any]) -> Dict[str, Any]: + host, port_text = endpoint.rsplit(":", 1) + with socket.create_connection((host, int(port_text)), timeout=30.0) as conn: + conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + conn.sendall( + json.dumps(payload, separators=(",", ":")).encode("utf-8") + b"\n" + ) + response = conn.makefile("rb").readline(1 << 20) + if not response: + raise RuntimeError(f"empty GPU-direct control response from {endpoint}") + decoded = json.loads(response.decode("utf-8")) + if not decoded.get("ok"): + raise RuntimeError(str(decoded.get("error", "GPU-direct control failure"))) + return decoded + + +class MooncakeGpuDirectFeatureStore(FeatureStore): + """Consumer-side FeatureStore for Mooncake NVLink and RDMA descriptors.""" + + def __init__( + self, + *, + store_id: str, + local_hostname: str, + transport: str, + rdma_devices: str = "", + device: "torch.device | str | int" = "cuda", + retain_on_release: bool = False, + enable_transfers: bool = True, + ) -> None: + normalized = "nvlink" if transport == "mnnvl" else transport + if normalized not in {"nvlink", "nvlink_intra", "rdma"}: + raise ValueError( + "GPU-direct Mooncake transport must be nvlink, nvlink_intra, or rdma" + ) + self.store_id = str(store_id) + self.local_hostname = str(local_hostname) + self.transport = normalized + self.rdma_devices = str(rdma_devices) + self.retain_on_release = retain_on_release + self.enable_transfers = enable_transfers + self.device = ( + self._resolve_device(device) if enable_transfers else torch.device("cpu") + ) + self.materialize_device = self.device + self.clone_on_fetch = False + self._engine = None + self._session_id: Optional[str] = None + self._lock = threading.RLock() + self._active_leases: Dict[str, SampleRef] = {} + self._known_refs: Dict[str, SampleRef] = {} + self._external_attempts: set[Tuple[str, int]] = set() + self._release_pending: Dict[str, SampleRef] = {} + self._freed: set[Tuple[str, int]] = set() + self._stats = { + "gets": 0, + "releases": 0, + "aborts": 0, + "gpu_direct_bytes": 0, + "host_payload_bytes": 0, + "release_batch_requests": 0, + "release_batch_items": 0, + "release_batch_fallbacks": 0, + } + if enable_transfers: + self._ensure_engine() + + @staticmethod + def _resolve_device(device: "torch.device | str | int") -> torch.device: + if not torch.cuda.is_available(): + raise RuntimeError("Mooncake GPU-direct transport requires CUDA") + resolved = ( + torch.device("cuda", device) + if isinstance(device, int) + else torch.device(device) + ) + if resolved.type != "cuda": + raise ValueError("Mooncake GPU-direct destination must be CUDA") + if resolved.index is None: + resolved = torch.device("cuda", torch.cuda.current_device()) + return resolved + + def _ensure_engine(self): + if self._engine is not None: + return self._engine + with self._lock: + if self._engine is not None: + return self._engine + try: + from mooncake import engine as mooncake_engine + except ImportError as exc: + raise ImportError( + "mooncake-transfer-engine is required for GPU-direct capture" + ) from exc + if self.transport == "nvlink": + os.environ["MC_FORCE_MNNVL"] = "true" + os.environ.pop("MC_INTRA_NVLINK", None) + os.environ.pop("MC_FORCE_HCA", None) + if not bool(getattr(mooncake_engine, "SUPPORT_MNNVL", False)): + raise RuntimeError( + "installed Mooncake wheel has SUPPORT_MNNVL=False" + ) + elif self.transport == "nvlink_intra": + os.environ.pop("MC_FORCE_MNNVL", None) + os.environ["MC_INTRA_NVLINK"] = "true" + os.environ.pop("MC_FORCE_HCA", None) + if not bool( + getattr(mooncake_engine, "SUPPORT_INTRA_NVLINK", False) + ): + raise RuntimeError( + "installed Mooncake wheel has SUPPORT_INTRA_NVLINK=False" + ) + else: + os.environ.pop("MC_FORCE_MNNVL", None) + os.environ.pop("MC_INTRA_NVLINK", None) + os.environ["MC_FORCE_HCA"] = "true" + engine = mooncake_engine.TransferEngine() + status = engine.initialize( + self.local_hostname, + "P2PHANDSHAKE", + self.transport, + self.rdma_devices if self.transport == "rdma" else "", + ) + if int(status) != 0: + raise RuntimeError( + f"Mooncake {self.transport} initialization failed: {status}" + ) + self._engine = engine + self._session_id = ( + f"{self.local_hostname}:{int(engine.get_rpc_port())}" + ) + return engine + + @staticmethod + def _descriptor(ref: SampleRef) -> Dict[str, Any]: + descriptor = ref.metadata.get("mooncake_gpu_direct") + if not isinstance(descriptor, dict): + raise KeyError("SampleRef carries no mooncake_gpu_direct descriptor") + return descriptor + + def put( + self, + tensors: Dict[str, torch.Tensor], + *, + sample_id: str, + metadata: Dict[str, Any], + ) -> SampleRef: + raise RuntimeError( + "MooncakeGpuDirectFeatureStore receives server-owned captures; " + "publish through SGLang spec_capture" + ) + + def adopt(self, sample_ref: SampleRef) -> None: + descriptor = self._descriptor(sample_ref) + ref_transport = str(descriptor["transport"]) + if ref_transport != self.transport: + raise ValueError( + f"capture transport {ref_transport!r} differs from store " + f"transport {self.transport!r}" + ) + generation = int(sample_ref.metadata["generation"]) + with self._lock: + self._known_refs[sample_ref.sample_id] = sample_ref + self._external_attempts.discard((sample_ref.sample_id, generation)) + + def track_external_attempt( + self, + sample_id: str, + *, + generation: int, + feature_names: List[str], + ) -> None: + del feature_names + with self._lock: + self._external_attempts.add((str(sample_id), int(generation))) + + def discard_external_attempts( + self, *, reason: str = "unadopted-external-capture" + ) -> int: + del reason + with self._lock: + count = len(self._external_attempts) + self._external_attempts.clear() + return count + + def _allocate_outputs( + self, + sample_ref: SampleRef, + names: List[str], + device: torch.device, + ) -> Tuple[Dict[str, torch.Tensor], List[int], List[int], List[int]]: + descriptor = self._descriptor(sample_ref) + features = descriptor.get("features") + if not isinstance(features, dict): + raise KeyError("GPU-direct descriptor carries no feature buffers") + outputs: Dict[str, torch.Tensor] = {} + local_addresses: List[int] = [] + remote_addresses: List[int] = [] + lengths: List[int] = [] + for name in names: + spec = sample_ref.feature_specs.get(name) + remote = features.get(name) + if spec is None or not isinstance(remote, dict): + raise KeyError(f"missing GPU-direct descriptor for {name!r}") + dtype = _TORCH_DTYPES.get(spec.dtype) + if dtype is None: + raise TypeError(f"unsupported GPU-direct dtype {spec.dtype!r}") + output = torch.empty(tuple(spec.shape), dtype=dtype, device=device) + expected = _nbytes(output) + described = int(remote["nbytes"]) + if expected != described: + raise ValueError( + f"feature {name!r} descriptor has {described} bytes; " + f"FeatureSpec requires {expected}" + ) + outputs[name] = output + local_addresses.append(output.data_ptr()) + remote_addresses.append(int(remote["address"])) + lengths.append(expected) + return outputs, local_addresses, remote_addresses, lengths + + def get( + self, + sample_ref: SampleRef, + *, + device: "torch.device | str" = "cuda", + names: Optional[List[str]] = None, + ) -> Tuple[Dict[str, torch.Tensor], FeatureHandle]: + destination = torch.device(device) + if destination.type != "cuda": + raise ValueError("GPU-direct feature materialization requires CUDA") + if destination.index is None: + destination = self.device + generation = int(sample_ref.metadata["generation"]) + with self._lock: + if (sample_ref.sample_id, generation) in self._freed: + raise KeyError( + f"sample {sample_ref.sample_id} generation {generation} was freed" + ) + descriptor = self._descriptor(sample_ref) + if str(descriptor["transport"]) != self.transport: + raise ValueError("SampleRef transport does not match this store") + wanted = names or list(sample_ref.feature_keys) + with torch.cuda.device(destination): + outputs, local, remote, lengths = self._allocate_outputs( + sample_ref, wanted, destination + ) + engine = self._ensure_engine() + registered: List[int] = [] + try: + if self.transport == "rdma": + for address, length in zip(local, lengths): + status = engine.register_memory(address, length) + if status is not None and int(status) != 0: + raise RuntimeError( + f"Mooncake CUDA registration failed: {status}" + ) + registered.append(address) + batch_read = getattr(engine, "batch_transfer_sync_read", None) + # Mooncake documents an accuracy caveat for batch reads over + # multi-node NVLink. NVL72 therefore uses individual reads. + if self.transport != "nvlink" and callable(batch_read): + status = batch_read( + str(descriptor["session_id"]), local, remote, lengths + ) + if int(status) != 0: + raise RuntimeError( + f"Mooncake batch GPU read failed: {status}" + ) + else: + for local_address, remote_address, length in zip( + local, remote, lengths + ): + status = engine.transfer_sync_read( + str(descriptor["session_id"]), + local_address, + remote_address, + length, + ) + if int(status) != 0: + raise RuntimeError( + f"Mooncake GPU read failed: {status}" + ) + finally: + for address in registered: + engine.unregister_memory(address) + + handle = FeatureHandle( + sample_id=sample_ref.sample_id, + generation=generation, + lease_token=uuid.uuid4().hex, + ) + with self._lock: + self._known_refs[sample_ref.sample_id] = sample_ref + self._active_leases[handle.lease_token] = sample_ref + self._stats["gets"] += 1 + self._stats["gpu_direct_bytes"] += sum(lengths) + return outputs, handle + + @staticmethod + def _release_remote(ref: SampleRef, *, op: str, reason: str) -> None: + descriptor = MooncakeGpuDirectFeatureStore._descriptor(ref) + _control_request( + str(descriptor["control_endpoint"]), + { + "op": op, + "token": str(descriptor["control_token"]), + "sample_id": ref.sample_id, + "generation": int(ref.metadata["generation"]), + "reason": reason, + }, + ) + + @staticmethod + def _release_remote_batch( + refs: List[SampleRef], *, op: str, reason: str + ) -> None: + if not refs: + return + descriptors = [ + MooncakeGpuDirectFeatureStore._descriptor(ref) for ref in refs + ] + endpoint = str(descriptors[0]["control_endpoint"]) + token = str(descriptors[0]["control_token"]) + for descriptor in descriptors[1:]: + if str(descriptor["control_endpoint"]) != endpoint: + raise ValueError("GPU-direct release batch spans control endpoints") + if str(descriptor["control_token"]) != token: + raise ValueError("GPU-direct release batch spans control tokens") + _control_request( + endpoint, + { + "op": f"{op}_batch", + "token": token, + "items": [ + { + "sample_id": ref.sample_id, + "generation": int(ref.metadata["generation"]), + } + for ref in refs + ], + "reason": reason, + }, + ) + + def _finish_release_many( + self, refs: List[SampleRef], *, op: str, reason: str + ) -> int: + unique = {ref.sample_id: ref for ref in refs} + groups: Dict[Tuple[str, str], List[SampleRef]] = {} + for ref in unique.values(): + descriptor = self._descriptor(ref) + key = ( + str(descriptor["control_endpoint"]), + str(descriptor["control_token"]), + ) + groups.setdefault(key, []).append(ref) + + def release_group(group: List[SampleRef]): + if len(group) == 1: + ref = group[0] + try: + self._release_remote(ref, op=op, reason=reason) + except Exception: + return {ref.sample_id: False}, False + return {ref.sample_id: True}, False + try: + self._release_remote_batch(group, op=op, reason=reason) + except Exception as exc: + if "unsupported control operation" not in str(exc): + return {ref.sample_id: False for ref in group}, False + outcomes = {} + for ref in group: + try: + self._release_remote(ref, op=op, reason=reason) + except Exception: + outcomes[ref.sample_id] = False + else: + outcomes[ref.sample_id] = True + return outcomes, True + return {ref.sample_id: True for ref in group}, False + + outcomes: Dict[str, bool] = {} + fallbacks = 0 + grouped = list(groups.values()) + if len(grouped) == 1: + group_outcomes, fallback = release_group(grouped[0]) + outcomes.update(group_outcomes) + fallbacks += int(fallback) + elif grouped: + with concurrent.futures.ThreadPoolExecutor( + max_workers=len(grouped), + thread_name_prefix="gpudirect-release", + ) as executor: + futures = [executor.submit(release_group, group) for group in grouped] + for future in futures: + group_outcomes, fallback = future.result() + outcomes.update(group_outcomes) + fallbacks += int(fallback) + + released = 0 + with self._lock: + self._stats["release_batch_requests"] += sum( + len(group) > 1 for group in grouped + ) + self._stats["release_batch_items"] += sum( + len(group) for group in grouped if len(group) > 1 + ) + self._stats["release_batch_fallbacks"] += fallbacks + for sample_id, ref in unique.items(): + if not outcomes.get(sample_id, False): + self._release_pending[sample_id] = ref + continue + generation = int(ref.metadata["generation"]) + self._release_pending.pop(sample_id, None) + self._known_refs.pop(sample_id, None) + self._freed.add((sample_id, generation)) + released += 1 + return released + + def _finish_release(self, ref: SampleRef, *, op: str, reason: str) -> bool: + return bool(self._finish_release_many([ref], op=op, reason=reason)) + + def release(self, handle: FeatureHandle, *, reason: str = "consumed") -> None: + with self._lock: + ref = self._active_leases.pop(handle.lease_token, None) + if ref is None: + return + if self.retain_on_release: + return + if self._finish_release(ref, op="release", reason=reason): + with self._lock: + self._stats["releases"] += 1 + + def abort(self, sample_id: str, *, reason: str) -> None: + with self._lock: + ref = self._known_refs.get(sample_id) + if ref is None: + return + if self._finish_release(ref, op="abort", reason=reason): + with self._lock: + self._stats["aborts"] += 1 + + def abort_many(self, sample_ids: List[str], *, reason: str) -> int: + target = set(sample_ids) + with self._lock: + refs = [ + ref + for sample_id, ref in self._known_refs.items() + if sample_id in target + ] + removed = self._finish_release_many(refs, op="abort", reason=reason) + with self._lock: + self._stats["aborts"] += removed + return removed + + def retry_sample_removals(self, sample_ids: List[str]) -> Dict[str, Any]: + target = set(sample_ids) + with self._lock: + pending = [ + ref + for sample_id, ref in self._release_pending.items() + if sample_id in target + ] + removed = self._finish_release_many( + pending, op="abort", reason="release-retry" + ) + with self._lock: + remaining = [sid for sid in self._release_pending if sid in target] + return { + "removed": removed, + "removed_bytes": 0, + "release_pending": len(remaining), + "remaining_ids": remaining, + "attempts": 1 if pending else 0, + } + + def drain_sample_removals( + self, + sample_ids: List[str], + *, + max_attempts: int = 8, + retry_interval_s: float = 0.25, + sleep: Callable[[float], None] = time.sleep, + ) -> Dict[str, int]: + return self._drain( + sample_ids=sample_ids, + max_attempts=max_attempts, + retry_interval_s=retry_interval_s, + sleep=sleep, + ) + + def drain_pending_removals( + self, + *, + max_attempts: int = 40, + retry_interval_s: float = 0.5, + sleep: Callable[[float], None] = time.sleep, + ) -> Dict[str, int]: + return self._drain( + sample_ids=None, + max_attempts=max_attempts, + retry_interval_s=retry_interval_s, + sleep=sleep, + ) + + def _drain( + self, + *, + sample_ids: Optional[List[str]], + max_attempts: int, + retry_interval_s: float, + sleep: Callable[[float], None], + ) -> Dict[str, int]: + target = None if sample_ids is None else set(sample_ids) + removed = 0 + for attempt in range(max_attempts): + with self._lock: + pending = [ + ref + for sample_id, ref in self._release_pending.items() + if target is None or sample_id in target + ] + if not pending: + return { + "removed": removed, + "removed_bytes": 0, + "release_pending": 0, + "attempts": attempt, + } + removed += self._finish_release_many( + pending, op="abort", reason="lifecycle-drain" + ) + if attempt + 1 < max_attempts and retry_interval_s: + sleep(retry_interval_s) + with self._lock: + remaining = [ + sample_id + for sample_id in self._release_pending + if target is None or sample_id in target + ] + raise RuntimeError( + f"MooncakeGpuDirectFeatureStore could not release {remaining[:16]}" + ) + + def health(self) -> Dict[str, Any]: + with self._lock: + return { + "backend": "mooncake_gpu_direct", + "transport": self.transport, + "device": str(self.device), + "session_id": self._session_id, + "known_samples": len(self._known_refs), + "active_leases": len(self._active_leases), + "release_pending": len(self._release_pending), + **self._stats, + } + + +__all__ = ["MooncakeGpuDirectFeatureStore"] diff --git a/specforge/training/disaggregated.py b/specforge/training/disaggregated.py index 47d175e80..8efa7bddc 100644 --- a/specforge/training/disaggregated.py +++ b/specforge/training/disaggregated.py @@ -154,6 +154,39 @@ def _mooncake_store(cfg: Config, *, retain_on_release: bool = False): ) +def _mooncake_gpu_direct_store( + cfg: Config, *, retain_on_release: bool = False +): + from specforge.runtime.data_plane.gpu_direct_store import ( + MooncakeGpuDirectFeatureStore, + ) + + transport = os.environ.get("MOONCAKE_PROTOCOL", "rdma") + if transport == "mnnvl": + transport = "nvlink" + return MooncakeGpuDirectFeatureStore( + store_id=os.environ.get("DISAGG_STORE_ID", cfg.run_id), + local_hostname=os.environ.get("MOONCAKE_LOCAL_HOSTNAME", "127.0.0.1"), + transport=transport, + rdma_devices=os.environ.get("MOONCAKE_RDMA_DEVICES", ""), + retain_on_release=retain_on_release, + enable_transfers=cfg.training.role == "consumer", + ) + + +def _online_store(cfg: Config, *, retain_on_release: bool = False): + backend = os.environ.get("DISAGG_BACKEND", "mooncake") + if backend == "mooncake_gpu_direct": + return _mooncake_gpu_direct_store( + cfg, retain_on_release=retain_on_release + ) + if backend == "mooncake": + return _mooncake_store(cfg, retain_on_release=retain_on_release) + raise ValueError( + f"online disaggregated training does not support backend {backend!r}" + ) + + def _offline_store(cfg: Config, *, retain_on_release: bool = False): backend = os.environ.get("DISAGG_BACKEND", "shared_dir") if backend == "mooncake": @@ -581,7 +614,9 @@ def _build_online( # The producer owns capture and explicit attempt cleanup. The consumer must # 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") + store = _online_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 @@ -663,6 +698,13 @@ def _build_online( feature_source=adapters if len(adapters) > 1 else adapters[0], num_rollout_workers=len(adapters), producer_concurrency=cfg.runtime.producer_concurrency, + producer_ordered_publish=cfg.runtime.producer_ordered_publish, + producer_prompt_prefetch_batches=( + cfg.runtime.producer_prompt_prefetch_batches + ), + producer_reorder_buffer=cfg.runtime.producer_reorder_buffer, + producer_prompt_routing=cfg.runtime.producer_prompt_routing, + producer_prompt_batching=cfg.runtime.producer_prompt_batching, run_id=cfg.run_id, target_hidden_size=hidden_size, target_vocab_size=target_vocab, @@ -688,6 +730,9 @@ def _build_online( cfg.runtime.feature_store_max_resident_bytes ), peer_wait_timeout_s=peer_wait_timeout_s, + prompt_ingest_batch_size=( + cfg.runtime.producer_prompt_ingest_batch_size + ), ) def produce() -> int: diff --git a/specforge/training/trainer.py b/specforge/training/trainer.py index d465b947b..352dc87e2 100644 --- a/specforge/training/trainer.py +++ b/specforge/training/trainer.py @@ -150,6 +150,8 @@ def __init__( batch_size=batch_size, collate_fn=collate_fn, per_sample_transform=per_sample_transform, + device=getattr(store, "materialize_device", "cpu"), + clone_on_fetch=getattr(store, "clone_on_fetch", True), drop_last=True, strategy=algorithm_name, ack=not defer_queue_ack,