From 7a84f35449c5806257fdc04d0aeb2b30f04bd029 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Fri, 17 Jul 2026 13:25:49 +0800 Subject: [PATCH 01/20] bench: add independent consumer request accounting Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- .../independent_consumer_fanout/README.md | 42 + .../config.example.json | 218 ++++ scripts/benchmark_independent_consumers.py | 41 + src/speculators/benchmarks/__init__.py | 1 + .../benchmarks/independent_consumers.py | 1107 +++++++++++++++++ tests/unit/benchmarks/__init__.py | 1 + .../benchmarks/test_independent_consumers.py | 300 +++++ 7 files changed, 1710 insertions(+) create mode 100644 benchmarks/independent_consumer_fanout/README.md create mode 100644 benchmarks/independent_consumer_fanout/config.example.json create mode 100644 scripts/benchmark_independent_consumers.py create mode 100644 src/speculators/benchmarks/__init__.py create mode 100644 src/speculators/benchmarks/independent_consumers.py create mode 100644 tests/unit/benchmarks/__init__.py create mode 100644 tests/unit/benchmarks/test_independent_consumers.py diff --git a/benchmarks/independent_consumer_fanout/README.md b/benchmarks/independent_consumer_fanout/README.md new file mode 100644 index 000000000..939900af5 --- /dev/null +++ b/benchmarks/independent_consumer_fanout/README.md @@ -0,0 +1,42 @@ +# Independent consumer fan-out benchmark + +This fixture runs two fresh-server scenarios in order: + +1. one vLLM hidden-state producer and one single-process trainer (`1p1c`); +2. one fresh producer and three independently launched single-process trainers + (`1p3c`). + +The launcher rejects distributed launchers and any DP, TP, PP, SP, or process-count +option other than one. It also requires distinct physical GPU indices for every role +that overlaps in time. `CUDA_VISIBLE_DEVICES` and distributed rank variables are owned +by the launcher and cannot be supplied by a role config. + +Each trainer receives its own accounting endpoint through the `{endpoint}` command +placeholder. Other placeholders are `{consumer_id}`, `{output_dir}`, and `{scenario}`. +The proxy forwards non-streaming OpenAI requests to vLLM and records a completion only +when the response is successful and contains a hidden-state artifact path. It stores +only a digest of the request identity, never the returned artifact path. + +Start from `config.example.json`, replace the model and preprocessed-data placeholders, +and keep each consumer command as a direct, single-process `scripts/train.py` launch. +Run the fixture from the repository root: + +```bash +python scripts/benchmark_independent_consumers.py \ + benchmarks/independent_consumer_fanout/config.example.json \ + --run-directory /tmp/speculators-fanout-run \ + --report /tmp/speculators-fanout-report.json +``` + +The run directory must not already exist. Role logs remain there and the compact report +contains the exact command/configuration, package versions, request and valid-completion +counts, shared-sample multiplicity, a common post-warmup throughput window, native +per-consumer `profile/step_ms` summaries, makespan, and sampled GPU memory. Environment +values are omitted from the report. The command exits nonzero +if a role fails, a GPU is shared or already occupied, a completion is malformed, sample +multiplicity is ambiguous, or the common steady-state window is too small. + +For the unshared baseline, +`expected_service_completions_per_shared_sample` is one for `1p1c` and three for +`1p3c`. A publish-once implementation changes the latter to one; the logical consumer +commands and all other workload settings must remain equivalent. diff --git a/benchmarks/independent_consumer_fanout/config.example.json b/benchmarks/independent_consumer_fanout/config.example.json new file mode 100644 index 000000000..163e79459 --- /dev/null +++ b/benchmarks/independent_consumer_fanout/config.example.json @@ -0,0 +1,218 @@ +{ + "producer": { + "gpu": 0, + "endpoint": "http://127.0.0.1:8000/v1", + "startup_timeout_seconds": 900, + "command": [ + "python", + "scripts/launch_vllm.py", + "Qwen/Qwen3-8B", + "--hidden-states-path", + "{output_dir}/hidden_states", + "--target-layer-ids", + "2", + "18", + "33", + "--", + "--port", + "8000", + "--tensor-parallel-size", + "1" + ] + }, + "scenarios": [ + { + "kind": "1p1c", + "consumers": [ + { + "consumer_id": "consumer-0", + "gpu": 1, + "command": [ + "python", + "scripts/train.py", + "--verifier-name-or-path", + "Qwen/Qwen3-8B", + "--data-path", + "PREPROCESSED_DATA_PATH", + "--vllm-endpoint", + "{endpoint}", + "--save-path", + "{output_dir}/checkpoints", + "--epochs", + "1", + "--total-seq-len", + "3072", + "--speculator-type", + "dflash", + "--draft-arch", + "qwen3", + "--block-size", + "8", + "--max-anchors", + "256", + "--num-layers", + "5", + "--target-layer-ids", + "2", + "18", + "33", + "--mask-token-id", + "151669", + "--on-missing", + "generate", + "--on-generate", + "delete", + "--num-workers", + "1" + ] + } + ], + "warmup_completions_per_consumer": 10, + "minimum_steady_completions_per_consumer": 50, + "warmup_steps_per_consumer": 10, + "minimum_steady_steps_per_consumer": 50, + "minimum_shared_samples": 50, + "expected_service_completions_per_shared_sample": 1 + }, + { + "kind": "1p3c", + "consumers": [ + { + "consumer_id": "consumer-0", + "gpu": 1, + "command": [ + "python", + "scripts/train.py", + "--verifier-name-or-path", + "Qwen/Qwen3-8B", + "--data-path", + "PREPROCESSED_DATA_PATH", + "--vllm-endpoint", + "{endpoint}", + "--save-path", + "{output_dir}/checkpoints", + "--epochs", + "1", + "--total-seq-len", + "3072", + "--speculator-type", + "dflash", + "--draft-arch", + "qwen3", + "--block-size", + "4", + "--max-anchors", + "256", + "--num-layers", + "5", + "--target-layer-ids", + "2", + "18", + "33", + "--mask-token-id", + "151669", + "--on-missing", + "generate", + "--on-generate", + "delete", + "--num-workers", + "1" + ] + }, + { + "consumer_id": "consumer-1", + "gpu": 2, + "command": [ + "python", + "scripts/train.py", + "--verifier-name-or-path", + "Qwen/Qwen3-8B", + "--data-path", + "PREPROCESSED_DATA_PATH", + "--vllm-endpoint", + "{endpoint}", + "--save-path", + "{output_dir}/checkpoints", + "--epochs", + "1", + "--total-seq-len", + "3072", + "--speculator-type", + "dflash", + "--draft-arch", + "qwen3", + "--block-size", + "8", + "--max-anchors", + "256", + "--num-layers", + "5", + "--target-layer-ids", + "2", + "18", + "33", + "--mask-token-id", + "151669", + "--on-missing", + "generate", + "--on-generate", + "delete", + "--num-workers", + "1" + ] + }, + { + "consumer_id": "consumer-2", + "gpu": 3, + "command": [ + "python", + "scripts/train.py", + "--verifier-name-or-path", + "Qwen/Qwen3-8B", + "--data-path", + "PREPROCESSED_DATA_PATH", + "--vllm-endpoint", + "{endpoint}", + "--save-path", + "{output_dir}/checkpoints", + "--epochs", + "1", + "--total-seq-len", + "3072", + "--speculator-type", + "dflash", + "--draft-arch", + "qwen3", + "--block-size", + "16", + "--max-anchors", + "256", + "--num-layers", + "5", + "--target-layer-ids", + "2", + "18", + "33", + "--mask-token-id", + "151669", + "--on-missing", + "generate", + "--on-generate", + "delete", + "--num-workers", + "1" + ] + } + ], + "warmup_completions_per_consumer": 10, + "minimum_steady_completions_per_consumer": 50, + "warmup_steps_per_consumer": 10, + "minimum_steady_steps_per_consumer": 50, + "minimum_shared_samples": 50, + "expected_service_completions_per_shared_sample": 3 + } + ], + "allowed_gpus": [0, 1, 2, 3], + "proxy_timeout_seconds": 180, + "memory_sample_interval_seconds": 0.25 +} diff --git a/scripts/benchmark_independent_consumers.py b/scripts/benchmark_independent_consumers.py new file mode 100644 index 000000000..45c10e759 --- /dev/null +++ b/scripts/benchmark_independent_consumers.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 + +import argparse +from pathlib import Path + +from speculators.benchmarks.independent_consumers import ( + load_config, + run_benchmark, + write_report, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run serial 1P1C and independent 1P3C hidden-state benchmarks." + ) + parser.add_argument("config", type=Path, help="Strict benchmark JSON config") + parser.add_argument( + "--run-directory", type=Path, required=True, help="New directory for role logs" + ) + parser.add_argument( + "--report", type=Path, required=True, help="Path for the compact JSON report" + ) + parser.add_argument( + "--validate-only", action="store_true", help="Validate config without launching" + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + config = load_config(args.config) + if args.validate_only: + return 0 + report = run_benchmark(config, args.run_directory) + write_report(report, args.report) + return 0 if report["valid"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/speculators/benchmarks/__init__.py b/src/speculators/benchmarks/__init__.py new file mode 100644 index 000000000..ec62a489f --- /dev/null +++ b/src/speculators/benchmarks/__init__.py @@ -0,0 +1 @@ +"""Benchmark helpers for reproducible Speculators performance evidence.""" diff --git a/src/speculators/benchmarks/independent_consumers.py b/src/speculators/benchmarks/independent_consumers.py new file mode 100644 index 000000000..2ac296618 --- /dev/null +++ b/src/speculators/benchmarks/independent_consumers.py @@ -0,0 +1,1107 @@ +from __future__ import annotations + +import hashlib +import http.client +import importlib.metadata +import json +import os +import re +import signal +import subprocess +import threading +import time +import urllib.error +import urllib.request +from collections import Counter, defaultdict +from contextlib import suppress +from dataclasses import asdict, dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any, Literal +from urllib.parse import urlsplit + +import psutil +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + + +class EvidenceError(RuntimeError): + """Raised when a benchmark cannot produce trustworthy evidence.""" + + +_DISTRIBUTED_ENV = { + "LOCAL_RANK", + "RANK", + "WORLD_SIZE", + "LOCAL_WORLD_SIZE", + "GROUP_RANK", + "ROLE_RANK", + "MASTER_ADDR", + "MASTER_PORT", +} +_GPU_ENV = "CUDA_VISIBLE_DEVICES" +_DISTRIBUTED_LAUNCHERS = { + "accelerate", + "deepspeed", + "mpiexec", + "mpirun", + "srun", + "torchrun", +} +_PARALLEL_SIZE_OPTIONS = { + "-dp", + "-pp", + "-tp", + "--data-parallel-size", + "--data-parallel-size-local", + "--data_parallel_size", + "--nnodes", + "--nproc-per-node", + "--nproc_per_node", + "--pipeline-parallel-size", + "--pipeline_parallel_size", + "--sequence-parallel-size", + "--sequence_parallel_size", + "--tensor-parallel-size", + "--tensor_parallel_size", +} + + +def _option_value(command: list[str], index: int) -> tuple[str, int]: + token = command[index] + if "=" in token: + return token.split("=", 1)[1], index + if index + 1 >= len(command): + raise ValueError(f"Missing value for {token}") + return command[index + 1], index + 1 + + +def _validate_single_process_command(command: list[str]) -> list[str]: + executable = Path(command[0]).name + if executable in _DISTRIBUTED_LAUNCHERS: + raise ValueError( + f"Distributed launcher {executable!r} is not an independent consumer" + ) + if command[1:3] == ["-m", "torch.distributed.run"]: + raise ValueError("torch.distributed.run is not an independent consumer") + if not executable.startswith("python"): + raise ValueError("Benchmark roles must be direct Python commands") + if "--fsdp-shard" in command: + raise ValueError("--fsdp-shard requires one distributed trainer") + + index = 0 + while index < len(command): + option = command[index].split("=", 1)[0] + if option in _PARALLEL_SIZE_OPTIONS: + raw_value, value_index = _option_value(command, index) + try: + value = int(raw_value) + except ValueError as error: + raise ValueError(f"{option} must have an integer value") from error + if value != 1: + raise ValueError( + f"{option}={value} creates one distributed/parallel job, not an " + "independent single-GPU consumer" + ) + index = value_index + index += 1 + return command + + +class _StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class CommandSpec(_StrictModel): + command: list[str] = Field(min_length=1) + env: dict[str, str] = Field(default_factory=dict) + + @field_validator("command") + @classmethod + def validate_command(cls, value: list[str]) -> list[str]: + if any(not token for token in value): + raise ValueError("Command arguments must not be empty") + return _validate_single_process_command(value) + + @field_validator("env") + @classmethod + def validate_env(cls, value: dict[str, str]) -> dict[str, str]: + owned = (_DISTRIBUTED_ENV | {_GPU_ENV}) & value.keys() + if owned: + names = ", ".join(sorted(owned)) + raise ValueError(f"The harness owns these environment variables: {names}") + return value + + +class ProducerSpec(CommandSpec): + gpu: int = Field(ge=0) + endpoint: str + startup_timeout_seconds: float = Field(default=900.0, gt=0) + + @field_validator("endpoint") + @classmethod + def validate_endpoint(cls, value: str) -> str: + parsed = urlsplit(value) + if parsed.scheme not in {"http", "https"} or not parsed.hostname: + raise ValueError("endpoint must be an absolute HTTP(S) URL") + return value.rstrip("/") + + +class ConsumerSpec(CommandSpec): + consumer_id: str = Field(pattern=r"^[A-Za-z0-9_.-]+$") + gpu: int = Field(ge=0) + + +class ScenarioSpec(_StrictModel): + kind: Literal["1p1c", "1p3c"] + consumers: list[ConsumerSpec] + timeout_seconds: float = Field(default=3600.0, gt=0) + warmup_completions_per_consumer: int = Field(default=2, ge=0) + minimum_steady_completions_per_consumer: int = Field(default=5, ge=1) + warmup_steps_per_consumer: int = Field(default=10, ge=0) + minimum_steady_steps_per_consumer: int = Field(default=10, ge=1) + minimum_shared_samples: int = Field(default=5, ge=1) + expected_service_completions_per_shared_sample: int = Field(ge=1) + + @model_validator(mode="after") + def validate_consumers(self) -> ScenarioSpec: + expected_consumers = 1 if self.kind == "1p1c" else 3 + if len(self.consumers) != expected_consumers: + raise ValueError( + f"{self.kind} requires exactly {expected_consumers} independent " + "consumer process(es)" + ) + ids = [consumer.consumer_id for consumer in self.consumers] + if len(ids) != len(set(ids)): + raise ValueError(f"{self.kind} consumer IDs must be unique") + gpus = [consumer.gpu for consumer in self.consumers] + if len(gpus) != len(set(gpus)): + raise ValueError(f"{self.kind} assigns more than one consumer to a GPU") + allowed_multiplicities = {1, expected_consumers} + if self.expected_service_completions_per_shared_sample not in ( + allowed_multiplicities + ): + raise ValueError( + "Expected service multiplicity must be either publish-once (1) or " + f"one completion per consumer ({expected_consumers})" + ) + return self + + +class BenchmarkConfig(_StrictModel): + producer: ProducerSpec + scenarios: list[ScenarioSpec] = Field(min_length=2, max_length=2) + allowed_gpus: list[int] = Field(min_length=2) + proxy_timeout_seconds: float = Field(default=180.0, gt=0) + memory_sample_interval_seconds: float = Field(default=0.25, gt=0) + + @field_validator("allowed_gpus") + @classmethod + def validate_allowed_gpus(cls, value: list[int]) -> list[int]: + if any(gpu < 0 for gpu in value): + raise ValueError("GPU indices must be non-negative") + if len(value) != len(set(value)): + raise ValueError("allowed_gpus must not contain duplicates") + return value + + @model_validator(mode="after") + def validate_layout(self) -> BenchmarkConfig: + if [scenario.kind for scenario in self.scenarios] != ["1p1c", "1p3c"]: + raise ValueError("Scenarios must run serially in 1p1c, then 1p3c order") + allowed = set(self.allowed_gpus) + if self.producer.gpu not in allowed: + raise ValueError("Producer GPU is outside allowed_gpus") + for scenario in self.scenarios: + consumer_gpus = {consumer.gpu for consumer in scenario.consumers} + if self.producer.gpu in consumer_gpus: + raise ValueError( + f"{scenario.kind} places producer and consumer on the same GPU" + ) + outside = consumer_gpus - allowed + if outside: + raise ValueError( + f"{scenario.kind} uses GPUs outside allowed_gpus: {sorted(outside)}" + ) + return self + + +@dataclass(frozen=True) +class RequestEvent: + consumer_id: str + request_key: str | None + requested_at: float + completed_at: float + status_code: int | None + valid_request: bool + valid_completion: bool + error: str | None = None + + +class AccountingLedger: + def __init__(self) -> None: + self._events: list[RequestEvent] = [] + self._lock = threading.Lock() + + def add(self, event: RequestEvent) -> None: + with self._lock: + self._events.append(event) + + def snapshot(self) -> list[RequestEvent]: + with self._lock: + return list(self._events) + + +def canonical_request_key(path: str, body: bytes) -> str: + """Return a stable key for one hidden-state completion request.""" + try: + payload = json.loads(body) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise ValueError("Request body is not valid JSON") from error + if not isinstance(payload, dict): + raise ValueError("Request JSON must be an object") + model = payload.get("model") + if not isinstance(model, str) or not model: + raise ValueError("Request is missing a model") + if payload.get("max_tokens") != 1: + raise ValueError("Hidden-state requests must use max_tokens=1") + if payload.get("stream", False): + raise ValueError("Streaming responses cannot be accounted atomically") + + normalized_path = urlsplit(path).path.rstrip("/") + if normalized_path.endswith("/chat/completions"): + messages = payload.get("messages") + if not isinstance(messages, list) or not messages: + raise ValueError("Chat request is missing messages") + identity: dict[str, Any] = { + "api": "chat.completions", + "messages": messages, + "model": model, + } + elif normalized_path.endswith("/completions"): + prompt = payload.get("prompt") + if ( + not isinstance(prompt, list) + or not prompt + or not all( + isinstance(token, int) and not isinstance(token, bool) + for token in prompt + ) + ): + raise ValueError("Completion prompt must be one non-empty token-ID list") + identity = {"api": "completions", "model": model, "prompt": prompt} + else: + raise ValueError("Request is not a completions endpoint") + + encoded = json.dumps( + identity, ensure_ascii=True, separators=(",", ":"), sort_keys=True + ).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _is_valid_completion(body: bytes) -> bool: + try: + payload = json.loads(body) + except (UnicodeDecodeError, json.JSONDecodeError): + return False + if not isinstance(payload, dict): + return False + choices = payload.get("choices") + transfer = payload.get("kv_transfer_params") + return bool( + isinstance(choices, list) + and choices + and isinstance(transfer, dict) + and isinstance(transfer.get("hidden_states_path"), str) + and transfer["hidden_states_path"] + ) + + +_HOP_BY_HOP_HEADERS = { + "connection", + "content-length", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailers", + "transfer-encoding", + "upgrade", +} + + +class AccountingProxy: + """A non-streaming reverse proxy that counts validated service completions.""" + + def __init__( + self, + target_endpoint: str, + consumer_id: str, + ledger: AccountingLedger, + timeout_seconds: float, + ) -> None: + target = urlsplit(target_endpoint) + if target.scheme not in {"http", "https"} or not target.hostname: + raise ValueError("target_endpoint must be an absolute HTTP(S) URL") + self._target = target + self._consumer_id = consumer_id + self._ledger = ledger + self._timeout = timeout_seconds + self._server = ThreadingHTTPServer(("127.0.0.1", 0), self._handler_type()) + self._server.daemon_threads = True + self._thread = threading.Thread( + target=self._server.serve_forever, + name=f"accounting-proxy-{consumer_id}", + daemon=True, + ) + + @property + def endpoint(self) -> str: + port = self._server.server_address[1] + base_path = self._target.path.rstrip("/") + return f"http://127.0.0.1:{port}{base_path}" + + def _handler_type(self) -> type[BaseHTTPRequestHandler]: # noqa: C901 + proxy = self + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, _format: str, *_args: Any) -> None: + return + + def do_GET(self) -> None: + self._forward() + + def do_POST(self) -> None: + self._forward() + + def _forward(self) -> None: + content_length = int(self.headers.get("Content-Length", "0")) + body = self.rfile.read(content_length) if content_length else b"" + is_completion = self.command == "POST" and urlsplit( + self.path + ).path.rstrip("/").endswith("/completions") + requested_at = time.monotonic() + request_key = None + request_error = None + if is_completion: + try: + request_key = canonical_request_key(self.path, body) + except ValueError as error: + request_error = str(error) + + response_status = 502 + response_reason: str | None = None + response_headers: list[tuple[str, str]] = [ + ("Content-Type", "application/json") + ] + response_body = b'{"error":"accounting proxy upstream failure"}' + valid_completion = False + try: + connection_type = ( + http.client.HTTPSConnection + if proxy._target.scheme == "https" + else http.client.HTTPConnection + ) + port = proxy._target.port or ( + 443 if proxy._target.scheme == "https" else 80 + ) + connection = connection_type( + proxy._target.hostname, port, timeout=proxy._timeout + ) + headers = { + name: value + for name, value in self.headers.items() + if name.lower() not in _HOP_BY_HOP_HEADERS + and name.lower() not in {"host", "accept-encoding"} + } + headers["Accept-Encoding"] = "identity" + connection.request( + self.command, self.path, body=body, headers=headers + ) + response = connection.getresponse() + response_status = response.status + response_reason = response.reason + response_body = response.read() + response_headers = response.getheaders() + connection.close() + + valid_completion = bool( + is_completion + and request_key + and response.status in range(200, 300) + and _is_valid_completion(response_body) + ) + except Exception as error: # noqa: BLE001 + request_error = request_error or type(error).__name__ + + if is_completion: + proxy._ledger.add( + RequestEvent( + consumer_id=proxy._consumer_id, + request_key=request_key, + requested_at=requested_at, + completed_at=time.monotonic(), + status_code=response_status, + valid_request=request_key is not None, + valid_completion=valid_completion, + error=request_error, + ) + ) + + self.send_response(response_status, response_reason) + for name, value in response_headers: + if name.lower() not in _HOP_BY_HOP_HEADERS: + self.send_header(name, value) + self.send_header("Content-Length", str(len(response_body))) + self.end_headers() + with suppress(BrokenPipeError, ConnectionResetError): + self.wfile.write(response_body) + + return Handler + + def start(self) -> None: + self._thread.start() + + def close(self) -> None: + self._server.shutdown() + self._server.server_close() + self._thread.join(timeout=5) + + +def analyze_scenario( # noqa: C901 + scenario: ScenarioSpec, + events: list[RequestEvent], + started_at: float, + finished_at: float, +) -> dict[str, Any]: + """Analyze one scenario and reject incomplete or ambiguous evidence.""" + reasons: list[str] = [] + expected_ids = [consumer.consumer_id for consumer in scenario.consumers] + by_consumer: dict[str, list[RequestEvent]] = defaultdict(list) + for event in events: + by_consumer[event.consumer_id].append(event) + + unexpected_ids = set(by_consumer) - set(expected_ids) + if unexpected_ids: + reasons.append(f"events from unknown consumers: {sorted(unexpected_ids)}") + invalid_requests = sum(not event.valid_request for event in events) + invalid_completions = sum(not event.valid_completion for event in events) + if invalid_requests: + reasons.append( + f"{invalid_requests} request(s) were not canonical hidden-state calls" + ) + if invalid_completions: + reasons.append(f"{invalid_completions} request(s) lacked a valid completion") + + measured: list[RequestEvent] = [] + steady_start = started_at + steady_end = finished_at + per_consumer_total: dict[str, int] = {} + for consumer_id in expected_ids: + completed = sorted( + ( + event + for event in by_consumer.get(consumer_id, []) + if event.valid_completion + ), + key=lambda event: event.completed_at, + ) + per_consumer_total[consumer_id] = len(completed) + needed = ( + scenario.warmup_completions_per_consumer + + scenario.minimum_steady_completions_per_consumer + ) + if len(completed) < needed: + reasons.append( + f"{consumer_id} completed {len(completed)} request(s), " + f"fewer than {needed}" + ) + continue + warmup = scenario.warmup_completions_per_consumer + if warmup: + steady_start = max(steady_start, completed[warmup - 1].completed_at) + steady_end = min(steady_end, completed[-1].completed_at) + measured.extend(completed[warmup:]) + + duration = steady_end - steady_start + steady_events = [ + event for event in measured if steady_start <= event.completed_at <= steady_end + ] + per_consumer_steady = Counter(event.consumer_id for event in steady_events) + if duration <= 0: + reasons.append("common steady-state window is empty") + for consumer_id in expected_ids: + count = per_consumer_steady[consumer_id] + if count < scenario.minimum_steady_completions_per_consumer: + reasons.append( + f"{consumer_id} has only {count} completion(s) in the common " + "steady window" + ) + + key_counts = Counter( + event.request_key for event in measured if event.request_key is not None + ) + key_consumers: dict[str, set[str]] = defaultdict(set) + for event in measured: + if event.request_key is not None: + key_consumers[event.request_key].add(event.consumer_id) + + multiplicity = scenario.expected_service_completions_per_shared_sample + if multiplicity == len(expected_ids): + qualifying = [ + key + for key, count in key_counts.items() + if count == multiplicity and key_consumers[key] == set(expected_ids) + ] + malformed = [ + key + for key, count in key_counts.items() + if count != multiplicity or key_consumers[key] != set(expected_ids) + ] + else: + qualifying = [key for key, count in key_counts.items() if count == 1] + malformed = [key for key, count in key_counts.items() if count != 1] + if len(qualifying) < scenario.minimum_shared_samples: + reasons.append( + f"only {len(qualifying)} sample key(s) have expected multiplicity " + f"{multiplicity}; need {scenario.minimum_shared_samples}" + ) + if malformed: + reasons.append( + f"{len(malformed)} measured sample key(s) have ambiguous multiplicity" + ) + + safe_key_counts = { + key[:16]: count for key, count in sorted(key_counts.items()) if key is not None + } + return { + "valid": not reasons, + "invalid_reasons": reasons, + "request_accounting": { + "requests": len(events), + "valid_completions": len(events) - invalid_completions, + "invalid_requests": invalid_requests, + "invalid_completions": invalid_completions, + "per_consumer_completions": per_consumer_total, + "expected_service_completions_per_shared_sample": multiplicity, + "qualifying_shared_samples": len(qualifying), + "sample_completion_counts": safe_key_counts, + }, + "steady_state": { + "warmup_completions_per_consumer": ( + scenario.warmup_completions_per_consumer + ), + "started_at_monotonic": steady_start, + "finished_at_monotonic": steady_end, + "duration_seconds": max(duration, 0.0), + "completions": len(steady_events), + "completions_per_consumer": dict(per_consumer_steady), + "completions_per_second": ( + len(steady_events) / duration if duration > 0 else None + ), + }, + } + + +_STEP_TIME_PATTERN = re.compile( + r"profile/step_ms=(?P[0-9]+(?:\.[0-9]+)?(?:e[+-]?[0-9]+)?)", + re.IGNORECASE, +) + + +def _percentile(values: list[float], percentile: float) -> float: + index = max(0, min(len(values) - 1, int(len(values) * percentile + 0.999999) - 1)) + return sorted(values)[index] + + +def analyze_consumer_steps( + log_path: Path, warmup_steps: int, minimum_steady_steps: int +) -> dict[str, Any]: + """Extract native trainer step timings and separate warmup from steady state.""" + if not log_path.is_file(): + return { + "valid": False, + "invalid_reasons": ["consumer log is missing"], + "observed_steps": 0, + "warmup_steps": 0, + "steady_steps": 0, + "step_ms_mean": None, + "step_ms_p50": None, + "step_ms_p95": None, + } + values = [] + with log_path.open(errors="replace") as log_file: + for line in log_file: + values.extend( + float(match.group("value")) + for match in _STEP_TIME_PATTERN.finditer(line) + ) + steady = values[warmup_steps:] + reasons = [] + if len(steady) < minimum_steady_steps: + reasons.append( + f"only {len(steady)} steady step timing(s); need {minimum_steady_steps}" + ) + return { + "valid": not reasons, + "invalid_reasons": reasons, + "observed_steps": len(values), + "warmup_steps": min(warmup_steps, len(values)), + "steady_steps": len(steady), + "step_ms_mean": sum(steady) / len(steady) if steady else None, + "step_ms_p50": _percentile(steady, 0.50) if steady else None, + "step_ms_p95": _percentile(steady, 0.95) if steady else None, + } + + +@dataclass(frozen=True) +class _GpuSample: + captured_at: float + total_memory_mib: dict[int, int] + role_memory_mib: dict[int, int] + compute_pids: dict[int, list[int]] + + +def _run_nvidia_smi(query: str) -> list[str]: + result = subprocess.run( # noqa: S603 + [ # noqa: S607 + "nvidia-smi", + f"--query-{query}", + "--format=csv,noheader,nounits", + ], + check=True, + capture_output=True, + text=True, + timeout=15, + ) + return [line.strip() for line in result.stdout.splitlines() if line.strip()] + + +def _gpu_snapshot(target_gpus: set[int]) -> _GpuSample: + gpu_rows = _run_nvidia_smi("gpu=index,uuid,memory.used") + uuid_to_index: dict[str, int] = {} + total_memory: dict[int, int] = {} + for row in gpu_rows: + raw_index, uuid, raw_memory = (part.strip() for part in row.split(",", 2)) + index = int(raw_index) + if index in target_gpus: + uuid_to_index[uuid] = index + total_memory[index] = int(raw_memory) + if set(total_memory) != target_gpus: + missing = sorted(target_gpus - set(total_memory)) + raise EvidenceError(f"nvidia-smi did not report target GPUs {missing}") + + role_memory = dict.fromkeys(target_gpus, 0) + compute_pids: dict[int, list[int]] = defaultdict(list) + for row in _run_nvidia_smi("compute-apps=gpu_uuid,pid,used_gpu_memory"): + uuid, raw_pid, raw_memory = (part.strip() for part in row.split(",", 2)) + index = uuid_to_index.get(uuid) + if index is None: + continue + try: + pid = int(raw_pid) + memory = int(raw_memory) + except ValueError as error: + raise EvidenceError(f"Unparseable nvidia-smi compute row: {row}") from error + compute_pids[index].append(pid) + role_memory[index] += memory + return _GpuSample( + captured_at=time.monotonic(), + total_memory_mib=total_memory, + role_memory_mib=role_memory, + compute_pids=dict(compute_pids), + ) + + +class GpuMonitor: + def __init__(self, target_gpus: set[int], interval_seconds: float) -> None: + self._target_gpus = target_gpus + self._interval = interval_seconds + self._roots: dict[int, int] = {} + self._known_pids: dict[int, set[int]] = defaultdict(set) + self._samples: list[_GpuSample] = [] + self._errors: list[str] = [] + self._lock = threading.Lock() + self._stop = threading.Event() + self._started = False + self._thread = threading.Thread( + target=self._run, name="benchmark-gpu-monitor", daemon=True + ) + + def require_idle(self) -> dict[int, int]: + sample = _gpu_snapshot(self._target_gpus) + occupied = {gpu: pids for gpu, pids in sample.compute_pids.items() if pids} + if occupied: + raise EvidenceError( + f"Target GPUs already have compute processes: {occupied}" + ) + return sample.total_memory_mib + + def set_role_process(self, gpu: int, pid: int) -> None: + with self._lock: + self._roots[gpu] = pid + self._known_pids[gpu].add(pid) + + def start(self) -> None: + self._started = True + self._thread.start() + + def stop(self) -> None: + if not self._started: + return + self._stop.set() + self._thread.join(timeout=max(5.0, self._interval * 4)) + + def _allowed_pids(self, gpu: int) -> set[int]: + with self._lock: + root = self._roots.get(gpu) + known = set(self._known_pids[gpu]) + if root is None: + return known + try: + descendants = {child.pid for child in psutil.Process(root).children(True)} + except (psutil.NoSuchProcess, psutil.AccessDenied): + descendants = set() + allowed = known | {root} | descendants + with self._lock: + self._known_pids[gpu].update(allowed) + return allowed + + def _run(self) -> None: + while not self._stop.wait(self._interval): + try: + sample = _gpu_snapshot(self._target_gpus) + for gpu, pids in sample.compute_pids.items(): + if len(pids) > 1: + self._errors.append( + f"GPU {gpu} has {len(pids)} CUDA compute processes" + ) + foreign = set(pids) - self._allowed_pids(gpu) + if foreign: + self._errors.append( + f"GPU {gpu} has foreign compute PIDs {sorted(foreign)}" + ) + self._samples.append(sample) + except Exception as error: # noqa: BLE001 + self._errors.append(f"{type(error).__name__}: {error}") + + def summarize( + self, started_at: float, finished_at: float, baseline: dict[int, int] + ) -> dict[str, Any]: + samples = [ + sample + for sample in self._samples + if started_at <= sample.captured_at <= finished_at + ] + per_gpu: dict[str, Any] = {} + for gpu in sorted(self._target_gpus): + role_values = [sample.role_memory_mib[gpu] for sample in samples] + total_values = [sample.total_memory_mib[gpu] for sample in samples] + observed = any(sample.compute_pids.get(gpu) for sample in samples) + per_gpu[str(gpu)] = { + "baseline_memory_mib": baseline.get(gpu), + "peak_role_memory_mib": max(role_values) if role_values else None, + "peak_total_memory_mib": max(total_values) if total_values else None, + "max_compute_processes": max( + (len(sample.compute_pids.get(gpu, [])) for sample in samples), + default=0, + ), + "compute_process_observed": observed, + } + errors = list(dict.fromkeys(self._errors)) + if not samples: + errors.append("No memory samples fall inside the steady-state window") + missing = [ + gpu + for gpu, value in per_gpu.items() + if not value["compute_process_observed"] + ] + if missing: + errors.append(f"No compute process observed on GPU(s) {missing}") + return { + "reliable": not errors, + "invalid_reasons": errors, + "sample_count": len(samples), + "per_gpu": per_gpu, + } + + +def _render(values: list[str], replacements: dict[str, str]) -> list[str]: + rendered: list[str] = [] + for original in values: + value = original + for name, replacement in replacements.items(): + value = value.replace("{" + name + "}", replacement) + rendered.append(value) + return rendered + + +class _ManagedProcess: + def __init__( + self, + command: list[str], + env: dict[str, str], + gpu: int, + log_path: Path, + ) -> None: + process_env = os.environ.copy() + for name in _DISTRIBUTED_ENV: + process_env.pop(name, None) + process_env.update(env) + process_env[_GPU_ENV] = str(gpu) + self._log = log_path.open("w", encoding="utf-8") + self.started_at = time.monotonic() + self.finished_at: float | None = None + self.process = subprocess.Popen( # noqa: S603 + command, + env=process_env, + stdout=self._log, + stderr=subprocess.STDOUT, + start_new_session=True, + text=True, + ) + + def close_log(self) -> None: + self._log.close() + + def terminate(self, grace_seconds: float = 20.0) -> None: + if self.process.poll() is not None: + self.finished_at = self.finished_at or time.monotonic() + self.close_log() + return + os.killpg(self.process.pid, signal.SIGTERM) + try: + self.process.wait(timeout=grace_seconds) + except subprocess.TimeoutExpired: + os.killpg(self.process.pid, signal.SIGKILL) + self.process.wait(timeout=10) + self.finished_at = time.monotonic() + self.close_log() + + +def _wait_for_producer( + endpoint: str, process: subprocess.Popen[str], timeout: float +) -> None: + deadline = time.monotonic() + timeout + url = endpoint.rstrip("/") + "/models" + while time.monotonic() < deadline: + return_code = process.poll() + if return_code is not None: + raise EvidenceError( + f"Producer exited during startup with code {return_code}" + ) + try: + with urllib.request.urlopen(url, timeout=5) as response: # noqa: S310 + if response.status in range(200, 300): + return + except (urllib.error.URLError, TimeoutError): + pass + time.sleep(1) + raise EvidenceError(f"Producer was not ready after {timeout:.1f} seconds") + + +def _run_scenario( # noqa: C901 + config: BenchmarkConfig, scenario: ScenarioSpec, output_dir: Path +) -> dict[str, Any]: + scenario_dir = output_dir / scenario.kind + scenario_dir.mkdir(parents=True, exist_ok=False) + target_gpus = { + config.producer.gpu, + *(consumer.gpu for consumer in scenario.consumers), + } + monitor = GpuMonitor(target_gpus, config.memory_sample_interval_seconds) + baseline: dict[int, int] = {} + ledger = AccountingLedger() + producer: _ManagedProcess | None = None + consumers: list[_ManagedProcess] = [] + proxies: list[AccountingProxy] = [] + runtime_errors: list[str] = [] + started_at = time.monotonic() + finished_at = started_at + + producer_replacements = { + "output_dir": str(scenario_dir / "producer"), + "scenario": scenario.kind, + } + Path(producer_replacements["output_dir"]).mkdir() + try: + baseline = monitor.require_idle() + producer = _ManagedProcess( + _render(config.producer.command, producer_replacements), + { + key: _render([value], producer_replacements)[0] + for key, value in config.producer.env.items() + }, + config.producer.gpu, + scenario_dir / "producer.log", + ) + monitor.set_role_process(config.producer.gpu, producer.process.pid) + _wait_for_producer( + config.producer.endpoint, + producer.process, + config.producer.startup_timeout_seconds, + ) + + for consumer in scenario.consumers: + proxy = AccountingProxy( + config.producer.endpoint, + consumer.consumer_id, + ledger, + config.proxy_timeout_seconds, + ) + proxy.start() + proxies.append(proxy) + + monitor.start() + started_at = time.monotonic() + for consumer, proxy in zip(scenario.consumers, proxies, strict=True): + consumer_dir = scenario_dir / consumer.consumer_id + consumer_dir.mkdir() + replacements = { + "consumer_id": consumer.consumer_id, + "endpoint": proxy.endpoint, + "output_dir": str(consumer_dir), + "scenario": scenario.kind, + } + process = _ManagedProcess( + _render(consumer.command, replacements), + { + **{ + key: _render([value], replacements)[0] + for key, value in consumer.env.items() + }, + "SPECULATORS_CONSUMER_ID": consumer.consumer_id, + }, + consumer.gpu, + scenario_dir / f"{consumer.consumer_id}.log", + ) + consumers.append(process) + monitor.set_role_process(consumer.gpu, process.process.pid) + + deadline = started_at + scenario.timeout_seconds + for consumer, process in zip(scenario.consumers, consumers, strict=True): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise EvidenceError(f"{scenario.kind} exceeded its timeout") + try: + return_code = process.process.wait(timeout=remaining) + except subprocess.TimeoutExpired as error: + raise EvidenceError( + f"{consumer.consumer_id} exceeded the scenario timeout" + ) from error + if return_code != 0: + raise EvidenceError( + f"{consumer.consumer_id} exited with code {return_code}" + ) + process.finished_at = time.monotonic() + producer_return_code = producer.process.poll() + if producer_return_code is not None: + raise EvidenceError( + f"Producer exited unexpectedly with code {producer_return_code}" + ) + finished_at = time.monotonic() + except Exception as error: # noqa: BLE001 + runtime_errors.append(f"{type(error).__name__}: {error}") + finished_at = time.monotonic() + finally: + monitor.stop() + for process in consumers: + process.terminate() + for proxy in proxies: + proxy.close() + if producer is not None: + producer.terminate() + + analysis = analyze_scenario(scenario, ledger.snapshot(), started_at, finished_at) + consumer_steps = {} + for consumer in scenario.consumers: + step_result = analyze_consumer_steps( + scenario_dir / f"{consumer.consumer_id}.log", + scenario.warmup_steps_per_consumer, + scenario.minimum_steady_steps_per_consumer, + ) + consumer_steps[consumer.consumer_id] = step_result + runtime_errors.extend( + f"{consumer.consumer_id}: {reason}" + for reason in step_result["invalid_reasons"] + ) + steady = analysis["steady_state"] + memory = monitor.summarize( + steady["started_at_monotonic"], + steady["finished_at_monotonic"], + baseline, + ) + invalid_reasons = [ + *runtime_errors, + *analysis["invalid_reasons"], + *memory["invalid_reasons"], + ] + return { + "kind": scenario.kind, + "valid": not invalid_reasons, + "invalid_reasons": invalid_reasons, + "consumer_processes": [ + { + "consumer_id": consumer.consumer_id, + "gpu": consumer.gpu, + "return_code": process.process.returncode, + "runtime_seconds": ( + process.finished_at - process.started_at + if process.finished_at is not None + else None + ), + } + for consumer, process in zip(scenario.consumers, consumers, strict=False) + ], + "request_accounting": analysis["request_accounting"], + "steady_state": steady, + "consumer_step_times": consumer_steps, + "makespan_seconds": max(finished_at - started_at, 0.0), + "memory": memory, + } + + +def _redacted_config(config: BenchmarkConfig) -> dict[str, Any]: + value = config.model_dump(mode="json") + value["producer"]["env"] = sorted(value["producer"]["env"]) + for scenario in value["scenarios"]: + for consumer in scenario["consumers"]: + consumer["env"] = sorted(consumer["env"]) + return value + + +def run_benchmark(config: BenchmarkConfig, output_dir: Path) -> dict[str, Any]: + """Run fresh-producer 1P1C and 1P3C scenarios and return one JSON report.""" + output_dir.mkdir(parents=True, exist_ok=False) + scenarios = [ + _run_scenario(config, scenario, output_dir) for scenario in config.scenarios + ] + versions = {} + for package in ("speculators", "torch", "vllm"): + try: + versions[package] = importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + versions[package] = None + return { + "schema_version": 1, + "valid": all(scenario["valid"] for scenario in scenarios), + "config": _redacted_config(config), + "versions": versions, + "scenarios": scenarios, + } + + +def load_config(path: Path) -> BenchmarkConfig: + return BenchmarkConfig.model_validate_json(path.read_text(encoding="utf-8")) + + +def write_report(report: dict[str, Any], path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + temporary.replace(path) + + +def event_as_dict(event: RequestEvent) -> dict[str, Any]: + """Expose a stable serialization helper for tests and external tooling.""" + return asdict(event) diff --git a/tests/unit/benchmarks/__init__.py b/tests/unit/benchmarks/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/tests/unit/benchmarks/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/unit/benchmarks/test_independent_consumers.py b/tests/unit/benchmarks/test_independent_consumers.py new file mode 100644 index 000000000..98c72b3f4 --- /dev/null +++ b/tests/unit/benchmarks/test_independent_consumers.py @@ -0,0 +1,300 @@ +from __future__ import annotations + +import json +import threading +import urllib.request +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest +from pydantic import ValidationError + +from speculators.benchmarks.independent_consumers import ( + AccountingLedger, + AccountingProxy, + BenchmarkConfig, + ConsumerSpec, + ProducerSpec, + RequestEvent, + ScenarioSpec, + analyze_consumer_steps, + analyze_scenario, + canonical_request_key, +) + + +def _consumer(consumer_id: str, gpu: int, command: list[str] | None = None): + return ConsumerSpec( + consumer_id=consumer_id, + gpu=gpu, + command=command or ["python", "trainer.py", "--endpoint", "{endpoint}"], + ) + + +def _scenario( + kind: str = "1p3c", + *, + multiplicity: int | None = None, + warmup: int = 1, + minimum_steady: int = 2, + minimum_shared: int = 2, +) -> ScenarioSpec: + consumer_count = 1 if kind == "1p1c" else 3 + return ScenarioSpec( + kind=kind, + consumers=[ + _consumer(f"c{index}", index + 1) for index in range(consumer_count) + ], + warmup_completions_per_consumer=warmup, + minimum_steady_completions_per_consumer=minimum_steady, + minimum_shared_samples=minimum_shared, + expected_service_completions_per_shared_sample=( + multiplicity if multiplicity is not None else consumer_count + ), + ) + + +def _config() -> BenchmarkConfig: + return BenchmarkConfig( + producer=ProducerSpec( + gpu=0, + endpoint="http://127.0.0.1:8000/v1", + command=["python", "producer.py", "--tensor-parallel-size", "1"], + ), + scenarios=[_scenario("1p1c"), _scenario("1p3c")], + allowed_gpus=[0, 1, 2, 3], + ) + + +def _event(consumer: str, key: str, completed_at: float, *, valid: bool = True): + return RequestEvent( + consumer_id=consumer, + request_key=key, + requested_at=completed_at - 0.05, + completed_at=completed_at, + status_code=200 if valid else 500, + valid_request=True, + valid_completion=valid, + ) + + +def _valid_events() -> list[RequestEvent]: + events = [] + times = { + "c0": [0.10, 0.20, 0.30, 0.40], + "c1": [0.15, 0.25, 0.35, 0.45], + "c2": [0.18, 0.28, 0.38, 0.48], + } + for consumer, completed_times in times.items(): + for key, completed_at in zip( + ["warmup", "sample-a", "sample-b", "sample-c"], + completed_times, + strict=True, + ): + events.append(_event(consumer, key, completed_at)) + return events + + +def test_config_requires_serial_one_then_three_consumers(): + config = _config() + + assert [scenario.kind for scenario in config.scenarios] == ["1p1c", "1p3c"] + assert [len(scenario.consumers) for scenario in config.scenarios] == [1, 3] + + with pytest.raises(ValidationError, match="Scenarios must run serially"): + BenchmarkConfig( + producer=config.producer, + scenarios=list(reversed(config.scenarios)), + allowed_gpus=config.allowed_gpus, + ) + + +@pytest.mark.parametrize( + "command", + [ + ["torchrun", "--nproc-per-node", "3", "trainer.py"], + ["python", "-m", "torch.distributed.run", "trainer.py"], + ["python", "trainer.py", "--tensor-parallel-size=3"], + ["python", "trainer.py", "-tp", "3"], + ["python", "trainer.py", "--data-parallel-size", "3"], + ["python", "trainer.py", "--fsdp-shard"], + ["bash", "trainer.sh"], + ], +) +def test_config_rejects_distributed_consumer_commands(command): + with pytest.raises( + ValidationError, + match="not an independent|creates one|direct Python|distributed trainer", + ): + _consumer("c0", 1, command) + + +def test_config_rejects_gpu_sharing_and_owned_environment(): + with pytest.raises(ValidationError, match="more than one consumer"): + ScenarioSpec( + kind="1p3c", + consumers=[_consumer("c0", 1), _consumer("c1", 1), _consumer("c2", 2)], + expected_service_completions_per_shared_sample=3, + ) + + with pytest.raises(ValidationError, match="harness owns"): + ConsumerSpec( + consumer_id="c0", + gpu=1, + command=["python", "trainer.py"], + env={"WORLD_SIZE": "3"}, + ) + + +def test_canonical_request_key_is_stable_and_semantic(): + first = json.dumps( + {"prompt": [1, 2, 3], "max_tokens": 1, "model": "model", "temperature": 0} + ).encode() + reordered = json.dumps( + {"model": "model", "temperature": 1, "max_tokens": 1, "prompt": [1, 2, 3]} + ).encode() + changed = json.dumps( + {"model": "model", "max_tokens": 1, "prompt": [1, 2, 4]} + ).encode() + + assert canonical_request_key("/v1/completions", first) == canonical_request_key( + "/v1/completions", reordered + ) + assert canonical_request_key("/v1/completions", first) != canonical_request_key( + "/v1/completions", changed + ) + + +@pytest.mark.parametrize( + "payload", + [ + {"model": "model", "max_tokens": 2, "prompt": [1]}, + {"model": "model", "max_tokens": 1, "prompt": [[1]]}, + {"model": "model", "max_tokens": 1, "prompt": [1], "stream": True}, + ], +) +def test_canonical_request_key_rejects_non_fixture_calls(payload): + with pytest.raises(ValueError): + canonical_request_key("/v1/completions", json.dumps(payload).encode()) + + +class _UpstreamHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, _format, *_args): + return + + def do_POST(self): + content_length = int(self.headers["Content-Length"]) + self.rfile.read(content_length) + response = json.dumps( + { + "choices": [{"text": ""}], + "kv_transfer_params": {"hidden_states_path": "/not/reported"}, + } + ).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(response))) + self.end_headers() + self.wfile.write(response) + + +def test_accounting_proxy_counts_validated_hidden_state_completion(): + upstream = ThreadingHTTPServer(("127.0.0.1", 0), _UpstreamHandler) + upstream_thread = threading.Thread(target=upstream.serve_forever, daemon=True) + upstream_thread.start() + ledger = AccountingLedger() + target = f"http://127.0.0.1:{upstream.server_address[1]}/v1" + proxy = AccountingProxy(target, "c0", ledger, timeout_seconds=2) + proxy.start() + try: + body = json.dumps( + {"model": "model", "prompt": [1, 2, 3], "max_tokens": 1} + ).encode() + request = urllib.request.Request( # noqa: S310 + proxy.endpoint + "/completions", + data=body, + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=2) as response: # noqa: S310 + assert response.status == 200 + + [event] = ledger.snapshot() + assert event.consumer_id == "c0" + assert event.valid_request + assert event.valid_completion + assert event.status_code == 200 + finally: + proxy.close() + upstream.shutdown() + upstream.server_close() + upstream_thread.join(timeout=2) + + +def test_analysis_separates_warmup_and_requires_exact_multiplicity(): + result = analyze_scenario(_scenario(), _valid_events(), 0.0, 1.0) + + assert result["valid"] + assert result["request_accounting"]["requests"] == 12 + assert result["request_accounting"]["qualifying_shared_samples"] == 3 + assert result["steady_state"]["duration_seconds"] == pytest.approx(0.22) + assert result["steady_state"]["completions"] == 7 + + +def test_analysis_fails_closed_on_missing_consumer_evidence(): + events = [event for event in _valid_events() if event.consumer_id != "c2"] + + result = analyze_scenario(_scenario(), events, 0.0, 1.0) + + assert not result["valid"] + assert any("c2 completed 0" in reason for reason in result["invalid_reasons"]) + assert any( + "expected multiplicity" in reason for reason in result["invalid_reasons"] + ) + + +def test_analysis_fails_closed_on_failed_or_duplicate_completion(): + events = _valid_events() + events.append(_event("c0", "sample-a", 0.5)) + events[-2] = _event("c2", "sample-c", 0.48, valid=False) + + result = analyze_scenario(_scenario(), events, 0.0, 1.0) + + assert not result["valid"] + assert result["request_accounting"]["invalid_completions"] == 1 + assert any( + "ambiguous multiplicity" in reason for reason in result["invalid_reasons"] + ) + + +def test_consumer_step_analysis_excludes_warmup_and_reports_percentiles(tmp_path): + log_path = tmp_path / "consumer.log" + log_path.write_text( + "\n".join( + [ + "profile/step_ms=9.9e+03, global_step=1", + "profile/step_ms=10.0, global_step=2", + "profile/step_ms=20.0, global_step=3", + "profile/step_ms=30.0, global_step=4", + "profile/step_ms=40.0, global_step=5", + ] + ) + ) + + result = analyze_consumer_steps(log_path, warmup_steps=1, minimum_steady_steps=4) + + assert result["valid"] + assert result["observed_steps"] == 5 + assert result["steady_steps"] == 4 + assert result["step_ms_mean"] == 25.0 + assert result["step_ms_p50"] == 20.0 + assert result["step_ms_p95"] == 40.0 + + +def test_consumer_step_analysis_fails_closed_on_missing_log(tmp_path): + result = analyze_consumer_steps( + tmp_path / "missing.log", warmup_steps=1, minimum_steady_steps=2 + ) + + assert not result["valid"] + assert result["invalid_reasons"] == ["consumer log is missing"] From bcd3f9a88b227389d24d6484dcfc88a8401b71f5 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Fri, 17 Jul 2026 14:12:17 +0800 Subject: [PATCH 02/20] feat: coalesce shared hidden-state artifacts Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- .../independent_consumer_fanout/README.md | 18 +- docs/cli/train.md | 8 + scripts/train.py | 44 ++ .../benchmarks/independent_consumers.py | 188 +++++++- .../data_generation/artifact_cache.py | 418 ++++++++++++++++++ src/speculators/train/data.py | 80 ++++ src/speculators/train/dataloader.py | 16 + .../benchmarks/test_independent_consumers.py | 100 +++++ .../data_generation/test_artifact_cache.py | 345 +++++++++++++++ tests/unit/train/test_cli_args.py | 30 ++ tests/unit/train/test_shared_artifacts.py | 254 +++++++++++ 11 files changed, 1479 insertions(+), 22 deletions(-) create mode 100644 src/speculators/data_generation/artifact_cache.py create mode 100644 tests/unit/data_generation/test_artifact_cache.py create mode 100644 tests/unit/train/test_shared_artifacts.py diff --git a/benchmarks/independent_consumer_fanout/README.md b/benchmarks/independent_consumer_fanout/README.md index 939900af5..85c31e5b6 100644 --- a/benchmarks/independent_consumer_fanout/README.md +++ b/benchmarks/independent_consumer_fanout/README.md @@ -12,7 +12,8 @@ that overlaps in time. `CUDA_VISIBLE_DEVICES` and distributed rank variables are by the launcher and cannot be supplied by a role config. Each trainer receives its own accounting endpoint through the `{endpoint}` command -placeholder. Other placeholders are `{consumer_id}`, `{output_dir}`, and `{scenario}`. +placeholder. Other placeholders are `{consumer_id}`, `{output_dir}`, `{scenario}`, and +the scenario-local `{shared_artifacts_dir}`. The proxy forwards non-streaming OpenAI requests to vLLM and records a completion only when the response is successful and contains a hidden-state artifact path. It stores only a digest of the request identity, never the returned artifact path. @@ -40,3 +41,18 @@ For the unshared baseline, `expected_service_completions_per_shared_sample` is one for `1p1c` and three for `1p3c`. A publish-once implementation changes the latter to one; the logical consumer commands and all other workload settings must remain equivalent. + +To measure publish-once fan-out, pass the same cache to every consumer with +`--shared-hidden-states-path {shared_artifacts_dir}` and set the `1p3c` expected service +multiplicity to one. The report then includes aggregate logical request, hit, miss, +coalesced-waiter, retry, publish, failure, cleanup, and timeout counters under +`shared_artifact_cache`. The run fails closed unless three logical requests correspond +to every service completion, exactly one miss is published, the other two requests hit, +and all failure, retry, cleanup, and timeout counters are zero. Baseline scenarios that +do not use the shared-cache placeholder remain valid without cache accounting. + +In publish-once mode, `per_consumer_completions` and the steady-state per-consumer +completion map identify which consumer owned each service miss. They do not represent +logical trainer progress: cache logical-request totals and each independent consumer's +`consumer_step_times` provide that evidence. The report labels both maps with +`service_request_owner` to make this distinction explicit. diff --git a/docs/cli/train.md b/docs/cli/train.md index 9b004218e..d9167808c 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -84,6 +84,14 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--max-retries`** (int, default: `3`) Maximum number of retry attempts per vLLM request on failure. +- **`--shared-hidden-states-path`** (str, default: `None`) Optional filesystem cache that coalesces identical online hidden-state requests across independent trainers. All participating trainers must use the same path. This does not change `--hidden-states-path`, which remains the per-dataset indexed cache used by `--on-generate cache`. + +- **`--shared-hidden-states-namespace`** (str, default: `None`) Optional request-identity namespace for the producer's extraction configuration. Use the same value for trainers sharing artifacts and a different value when target layer IDs or other extraction semantics differ. + +- **`--shared-hidden-states-ttl`** (float, default: `3600.0`) Seconds to retain shared artifacts before regenerating them. Set to `0` to disable expiration. + +- **`--shared-hidden-states-lock-timeout`** (float, default: `300.0`) Maximum seconds to wait while another trainer generates and atomically publishes the same artifact. + - **`--legacy-data`** (flag) **DEPRECATED.** Use the old data format which stores hidden states alongside token_ids. - **`--total-seq-len`** (int, default: `8192`) Maximum total sequence length for training batches. Note: samples will be packed into batches with total combined sequence length `{total-seq-len}`. diff --git a/scripts/train.py b/scripts/train.py index 0de39860e..65f513ba2 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -643,6 +643,14 @@ def main(args: argparse.Namespace): # noqa: C901 verifier_name_or_path=args.verifier_name_or_path, request_timeout=args.request_timeout, max_retries=args.max_retries, + shared_artifacts_path=args.shared_hidden_states_path, + shared_artifacts_namespace=args.shared_hidden_states_namespace, + shared_artifacts_ttl_seconds=( + None + if args.shared_hidden_states_ttl == 0 + else args.shared_hidden_states_ttl + ), + shared_artifacts_lock_timeout_seconds=(args.shared_hidden_states_lock_timeout), hidden_size=hidden_size, num_target_layers=num_target_layers, num_workers=args.num_workers, @@ -894,6 +902,42 @@ def parse_args(): "Only applies if --on-missing=generate." ), ) + parser.add_argument( + "--shared-hidden-states-path", + type=str, + default=None, + help=( + "Optional shared artifact cache used to coalesce identical hidden-state " + "requests across independent trainers. Disabled by default." + ), + ) + parser.add_argument( + "--shared-hidden-states-namespace", + type=str, + default=None, + help=( + "Optional identity namespace for the producer's extraction configuration. " + "Trainers sharing artifacts must use the same value." + ), + ) + parser.add_argument( + "--shared-hidden-states-ttl", + type=float, + default=3600.0, + help=( + "Seconds to retain shared artifacts before regeneration; 0 disables " + "expiration. Only applies when --shared-hidden-states-path is set." + ), + ) + parser.add_argument( + "--shared-hidden-states-lock-timeout", + type=float, + default=300.0, + help=( + "Maximum seconds to wait for another trainer publishing the same shared " + "artifact. Only applies when --shared-hidden-states-path is set." + ), + ) parser.add_argument( "--legacy-data", action="store_true", diff --git a/src/speculators/benchmarks/independent_consumers.py b/src/speculators/benchmarks/independent_consumers.py index 2ac296618..1730ff44b 100644 --- a/src/speculators/benchmarks/independent_consumers.py +++ b/src/speculators/benchmarks/independent_consumers.py @@ -23,6 +23,8 @@ import psutil from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from speculators.data_generation.artifact_cache import HiddenStateArtifactCache + class EvidenceError(RuntimeError): """Raised when a benchmark cannot produce trustworthy evidence.""" @@ -493,35 +495,61 @@ def analyze_scenario( # noqa: C901 if invalid_completions: reasons.append(f"{invalid_completions} request(s) lacked a valid completion") + multiplicity = scenario.expected_service_completions_per_shared_sample + publish_once = multiplicity == 1 and len(expected_ids) > 1 measured: list[RequestEvent] = [] steady_start = started_at steady_end = finished_at - per_consumer_total: dict[str, int] = {} - for consumer_id in expected_ids: + per_consumer_total = { + consumer_id: sum( + event.valid_completion for event in by_consumer.get(consumer_id, []) + ) + for consumer_id in expected_ids + } + if publish_once: completed = sorted( - ( - event - for event in by_consumer.get(consumer_id, []) - if event.valid_completion - ), + (event for event in events if event.valid_completion), key=lambda event: event.completed_at, ) - per_consumer_total[consumer_id] = len(completed) needed = ( scenario.warmup_completions_per_consumer + scenario.minimum_steady_completions_per_consumer ) if len(completed) < needed: reasons.append( - f"{consumer_id} completed {len(completed)} request(s), " - f"fewer than {needed}" + f"service completed {len(completed)} request(s), fewer than {needed}" ) - continue - warmup = scenario.warmup_completions_per_consumer - if warmup: - steady_start = max(steady_start, completed[warmup - 1].completed_at) - steady_end = min(steady_end, completed[-1].completed_at) - measured.extend(completed[warmup:]) + elif completed: + warmup = scenario.warmup_completions_per_consumer + if warmup: + steady_start = max(steady_start, completed[warmup - 1].completed_at) + steady_end = min(steady_end, completed[-1].completed_at) + measured.extend(completed[warmup:]) + else: + for consumer_id in expected_ids: + completed = sorted( + ( + event + for event in by_consumer.get(consumer_id, []) + if event.valid_completion + ), + key=lambda event: event.completed_at, + ) + needed = ( + scenario.warmup_completions_per_consumer + + scenario.minimum_steady_completions_per_consumer + ) + if len(completed) < needed: + reasons.append( + f"{consumer_id} completed {len(completed)} request(s), " + f"fewer than {needed}" + ) + continue + warmup = scenario.warmup_completions_per_consumer + if warmup: + steady_start = max(steady_start, completed[warmup - 1].completed_at) + steady_end = min(steady_end, completed[-1].completed_at) + measured.extend(completed[warmup:]) duration = steady_end - steady_start steady_events = [ @@ -530,13 +558,20 @@ def analyze_scenario( # noqa: C901 per_consumer_steady = Counter(event.consumer_id for event in steady_events) if duration <= 0: reasons.append("common steady-state window is empty") - for consumer_id in expected_ids: - count = per_consumer_steady[consumer_id] - if count < scenario.minimum_steady_completions_per_consumer: + if publish_once: + if len(steady_events) < scenario.minimum_steady_completions_per_consumer: reasons.append( - f"{consumer_id} has only {count} completion(s) in the common " + f"service has only {len(steady_events)} completion(s) in the " "steady window" ) + else: + for consumer_id in expected_ids: + count = per_consumer_steady[consumer_id] + if count < scenario.minimum_steady_completions_per_consumer: + reasons.append( + f"{consumer_id} has only {count} completion(s) in the common " + "steady window" + ) key_counts = Counter( event.request_key for event in measured if event.request_key is not None @@ -546,7 +581,6 @@ def analyze_scenario( # noqa: C901 if event.request_key is not None: key_consumers[event.request_key].add(event.consumer_id) - multiplicity = scenario.expected_service_completions_per_shared_sample if multiplicity == len(expected_ids): qualifying = [ key @@ -583,11 +617,15 @@ def analyze_scenario( # noqa: C901 "invalid_requests": invalid_requests, "invalid_completions": invalid_completions, "per_consumer_completions": per_consumer_total, + "per_consumer_completions_semantics": ( + "service_request_owner" if publish_once else "logical_consumer" + ), "expected_service_completions_per_shared_sample": multiplicity, "qualifying_shared_samples": len(qualifying), "sample_completion_counts": safe_key_counts, }, "steady_state": { + "mode": "publish_once_service" if publish_once else "per_consumer_service", "warmup_completions_per_consumer": ( scenario.warmup_completions_per_consumer ), @@ -596,6 +634,9 @@ def analyze_scenario( # noqa: C901 "duration_seconds": max(duration, 0.0), "completions": len(steady_events), "completions_per_consumer": dict(per_consumer_steady), + "completions_per_consumer_semantics": ( + "service_request_owner" if publish_once else "logical_consumer" + ), "completions_per_second": ( len(steady_events) / duration if duration > 0 else None ), @@ -603,6 +644,87 @@ def analyze_scenario( # noqa: C901 } +def analyze_cache_accounting( + scenario: ScenarioSpec, + stats: dict[str, Any] | None, + service_completions: int, +) -> dict[str, Any]: + """Validate cache counters against service-level request accounting.""" + consumer_count = len(scenario.consumers) + publish_once = ( + scenario.expected_service_completions_per_shared_sample == 1 + and consumer_count > 1 + ) + if stats is None: + reasons = ( + ["publish-once scenario is missing cache accounting"] + if publish_once + else [] + ) + return {"valid": not reasons, "invalid_reasons": reasons, "stats": None} + + required = { + "schema_version", + "logical_requests", + "hits", + "misses", + "coalesced_waiters", + "retry_generations", + "publishes", + "generation_failures", + "publish_failures", + "invalid_artifacts_removed", + "expired_artifacts_removed", + "stale_temps_removed", + "lock_timeouts", + } + reasons = [] + if required - stats.keys(): + reasons.append(f"cache accounting is missing {sorted(required - stats.keys())}") + return {"valid": False, "invalid_reasons": reasons, "stats": stats} + if any( + not isinstance(stats[name], int) + or isinstance(stats[name], bool) + or stats[name] < 0 + for name in required + ): + reasons.append("cache accounting contains invalid counters") + return {"valid": False, "invalid_reasons": reasons, "stats": stats} + if stats["schema_version"] != 1: + reasons.append( + f"cache accounting schema_version={stats['schema_version']}, expected 1" + ) + + expected_logical_requests = service_completions * consumer_count + expected_hits = service_completions * (consumer_count - 1) + expected = { + "logical_requests": expected_logical_requests, + "hits": expected_hits, + "misses": service_completions, + "publishes": service_completions, + } + for name, value in expected.items(): + if stats[name] != value: + reasons.append(f"cache {name}={stats[name]}, expected {value}") + for name in ( + "retry_generations", + "generation_failures", + "publish_failures", + "invalid_artifacts_removed", + "expired_artifacts_removed", + "stale_temps_removed", + "lock_timeouts", + ): + if stats[name]: + reasons.append(f"cache {name} must be zero, got {stats[name]}") + if stats["coalesced_waiters"] > stats["hits"]: + reasons.append( + "cache coalesced_waiters cannot exceed cache hits, got " + f"{stats['coalesced_waiters']} > {stats['hits']}" + ) + return {"valid": not reasons, "invalid_reasons": reasons, "stats": stats} + + _STEP_TIME_PATTERN = re.compile( r"profile/step_ms=(?P[0-9]+(?:\.[0-9]+)?(?:e[+-]?[0-9]+)?)", re.IGNORECASE, @@ -916,6 +1038,12 @@ def _run_scenario( # noqa: C901 consumers: list[_ManagedProcess] = [] proxies: list[AccountingProxy] = [] runtime_errors: list[str] = [] + shared_artifacts_dir = scenario_dir / "shared-artifacts" + cache_accounting_enabled = any( + "{shared_artifacts_dir}" in value + for consumer in scenario.consumers + for value in (*consumer.command, *consumer.env.values()) + ) started_at = time.monotonic() finished_at = started_at @@ -962,6 +1090,7 @@ def _run_scenario( # noqa: C901 "endpoint": proxy.endpoint, "output_dir": str(consumer_dir), "scenario": scenario.kind, + "shared_artifacts_dir": str(shared_artifacts_dir), } process = _ManagedProcess( _render(consumer.command, replacements), @@ -1013,6 +1142,21 @@ def _run_scenario( # noqa: C901 producer.terminate() analysis = analyze_scenario(scenario, ledger.snapshot(), started_at, finished_at) + cache_stats = None + if cache_accounting_enabled: + try: + cache_stats = HiddenStateArtifactCache( + shared_artifacts_dir, artifact_ttl_seconds=None + ).snapshot_stats() + except Exception as error: # noqa: BLE001 + runtime_errors.append( + f"cache accounting unavailable: {type(error).__name__}: {error}" + ) + cache_accounting = analyze_cache_accounting( + scenario, + cache_stats, + analysis["request_accounting"]["valid_completions"], + ) consumer_steps = {} for consumer in scenario.consumers: step_result = analyze_consumer_steps( @@ -1034,6 +1178,7 @@ def _run_scenario( # noqa: C901 invalid_reasons = [ *runtime_errors, *analysis["invalid_reasons"], + *cache_accounting["invalid_reasons"], *memory["invalid_reasons"], ] return { @@ -1054,6 +1199,7 @@ def _run_scenario( # noqa: C901 for consumer, process in zip(scenario.consumers, consumers, strict=False) ], "request_accounting": analysis["request_accounting"], + "shared_artifact_cache": cache_accounting, "steady_state": steady, "consumer_step_times": consumer_steps, "makespan_seconds": max(finished_at - started_at, 0.0), diff --git a/src/speculators/data_generation/artifact_cache.py b/src/speculators/data_generation/artifact_cache.py new file mode 100644 index 000000000..2d6084cc3 --- /dev/null +++ b/src/speculators/data_generation/artifact_cache.py @@ -0,0 +1,418 @@ +from __future__ import annotations + +import fcntl +import hashlib +import json +import os +import re +import time +import uuid +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +from safetensors.torch import load_file, save_file + +from speculators.data_generation.vllm_client import ClientItem + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator, Mapping + +_REQUEST_ID_PATTERN = re.compile(r"^[0-9a-f]{64}$") +_STATS_VERSION = 1 +_COUNTERS = ( + "logical_requests", + "hits", + "misses", + "coalesced_waiters", + "retry_generations", + "publishes", + "generation_failures", + "publish_failures", + "invalid_artifacts_removed", + "expired_artifacts_removed", + "stale_temps_removed", + "lock_timeouts", +) + + +class ArtifactCacheError(RuntimeError): + """Base error raised by shared hidden-state artifact caching.""" + + +class ArtifactLockTimeoutError(ArtifactCacheError, TimeoutError): + """Raised when a request key remains owned beyond the configured timeout.""" + + +@dataclass(frozen=True) +class ArtifactResult: + data: dict[str, torch.Tensor] + request_id: str + path: Path + cache_hit: bool + coalesced: bool + + +def canonical_hidden_state_request_id( + model: str, + client_item: ClientItem, + *, + namespace: str | None = None, +) -> str: + """Build a stable identity for the hidden states produced by one request.""" + if not model: + raise ValueError("model must be non-empty") + token_ids = client_item.get("input_ids") + if ( + not isinstance(token_ids, list) + or not token_ids + or not all( + isinstance(token, int) and not isinstance(token, bool) + for token in token_ids + ) + ): + raise ValueError("input_ids must be one non-empty list of integer token IDs") + + identity: dict[str, Any] = { + "input_ids": token_ids, + "model": model, + "schema_version": 1, + } + messages = client_item.get("messages") + if messages is not None: + identity["messages"] = messages + if namespace is not None: + identity["namespace"] = namespace + try: + encoded = json.dumps( + identity, ensure_ascii=True, separators=(",", ":"), sort_keys=True + ).encode() + except (TypeError, ValueError) as error: + raise ValueError("request identity must be JSON serializable") from error + return hashlib.sha256(encoded).hexdigest() + + +def _empty_stats() -> dict[str, int]: + return {"schema_version": _STATS_VERSION, **dict.fromkeys(_COUNTERS, 0)} + + +def _fsync_directory(path: Path) -> None: + descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +class HiddenStateArtifactCache: + """Cross-process, publish-once cache for immutable hidden-state tensors.""" + + def __init__( + self, + root: str | os.PathLike[str], + *, + artifact_ttl_seconds: float | None = 3600.0, + stale_temp_seconds: float = 300.0, + lock_timeout_seconds: float = 300.0, + lock_poll_seconds: float = 0.05, + ) -> None: + if artifact_ttl_seconds is not None and artifact_ttl_seconds <= 0: + raise ValueError("artifact_ttl_seconds must be positive or None") + if stale_temp_seconds <= 0: + raise ValueError("stale_temp_seconds must be positive") + if lock_timeout_seconds <= 0: + raise ValueError("lock_timeout_seconds must be positive") + if lock_poll_seconds <= 0: + raise ValueError("lock_poll_seconds must be positive") + + self.root = Path(root).expanduser().resolve() + self.artifact_ttl_seconds = artifact_ttl_seconds + self.stale_temp_seconds = stale_temp_seconds + self.lock_timeout_seconds = lock_timeout_seconds + self.lock_poll_seconds = lock_poll_seconds + self._artifacts = self.root / "artifacts" + self._locks = self.root / "locks" + self._stats_path = self.root / "stats.json" + self._stats_lock_path = self.root / "stats.lock" + self._artifacts.mkdir(parents=True, exist_ok=True) + self._locks.mkdir(parents=True, exist_ok=True) + + @staticmethod + def _validate_request_id(request_id: str) -> None: + if _REQUEST_ID_PATTERN.fullmatch(request_id) is None: + raise ValueError("request_id must be a lowercase SHA-256 digest") + + def artifact_path(self, request_id: str) -> Path: + self._validate_request_id(request_id) + return self._artifacts / request_id[:2] / f"{request_id}.safetensors" + + def _lock_path(self, request_id: str) -> Path: + return self._locks / request_id[:2] / f"{request_id}.lock" + + @contextmanager + def _request_lock( + self, request_id: str, *, timeout_seconds: float | None = None + ) -> Iterator[bool]: + lock_path = self._lock_path(request_id) + lock_path.parent.mkdir(parents=True, exist_ok=True) + descriptor = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o600) + timeout = ( + self.lock_timeout_seconds if timeout_seconds is None else timeout_seconds + ) + deadline = time.monotonic() + timeout + waited = False + acquired = False + try: + while True: + try: + fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB) + acquired = True + break + except BlockingIOError: + waited = True + remaining = deadline - time.monotonic() + if timeout == 0 or remaining <= 0: + raise ArtifactLockTimeoutError( + f"Timed out waiting for hidden-state request {request_id}" + ) from None + time.sleep(min(self.lock_poll_seconds, remaining)) + yield waited + finally: + if acquired: + fcntl.flock(descriptor, fcntl.LOCK_UN) + os.close(descriptor) + + def _read_stats_unlocked(self) -> dict[str, int]: + if not self._stats_path.exists(): + return _empty_stats() + try: + value = json.loads(self._stats_path.read_text()) + except (OSError, json.JSONDecodeError) as error: + raise ArtifactCacheError("Cache accounting file is unreadable") from error + if not isinstance(value, dict) or value.get("schema_version") != _STATS_VERSION: + raise ArtifactCacheError("Cache accounting schema is invalid") + stats = _empty_stats() + for counter in _COUNTERS: + count = value.get(counter) + if not isinstance(count, int) or isinstance(count, bool) or count < 0: + raise ArtifactCacheError( + f"Cache accounting counter {counter!r} is invalid" + ) + stats[counter] = count + return stats + + @contextmanager + def _stats_lock(self, operation: int) -> Iterator[None]: + descriptor = os.open(self._stats_lock_path, os.O_CREAT | os.O_RDWR, 0o600) + try: + fcntl.flock(descriptor, operation) + yield + finally: + fcntl.flock(descriptor, fcntl.LOCK_UN) + os.close(descriptor) + + def _record(self, **deltas: int) -> None: + unknown = set(deltas) - set(_COUNTERS) + if unknown: + raise ValueError(f"Unknown cache accounting counters: {sorted(unknown)}") + if any(delta < 0 for delta in deltas.values()): + raise ValueError("Cache accounting deltas must be non-negative") + with self._stats_lock(fcntl.LOCK_EX): + stats = self._read_stats_unlocked() + for counter, delta in deltas.items(): + stats[counter] += delta + temporary = self.root / f".stats.{os.getpid()}.{uuid.uuid4().hex}.tmp" + try: + with temporary.open("w") as output: + json.dump(stats, output, separators=(",", ":"), sort_keys=True) + output.write("\n") + output.flush() + os.fsync(output.fileno()) + temporary.replace(self._stats_path) + _fsync_directory(self.root) + finally: + temporary.unlink(missing_ok=True) + + def snapshot_stats(self) -> dict[str, int]: + with self._stats_lock(fcntl.LOCK_SH): + return self._read_stats_unlocked() + + def _is_expired(self, path: Path, now: float) -> bool: + return bool( + self.artifact_ttl_seconds is not None + and now - path.stat().st_mtime >= self.artifact_ttl_seconds + ) + + def _remove_stale_temps_for_key(self, request_id: str, now: float) -> int: + artifact_path = self.artifact_path(request_id) + removed = 0 + for path in artifact_path.parent.glob(f".{request_id}.*.tmp"): + try: + if now - path.stat().st_mtime >= self.stale_temp_seconds: + path.unlink(missing_ok=True) + removed += 1 + except FileNotFoundError: + pass + return removed + + @staticmethod + def _validate_tensors(data: Mapping[str, torch.Tensor]) -> None: + if not data or not all( + isinstance(value, torch.Tensor) for value in data.values() + ): + raise ArtifactCacheError("Artifact producer must return a tensor mapping") + + def _publish( + self, request_id: str, data: Mapping[str, torch.Tensor], target: Path + ) -> None: + target.parent.mkdir(parents=True, exist_ok=True) + temporary = target.parent / ( + f".{request_id}.{os.getpid()}.{uuid.uuid4().hex}.tmp" + ) + try: + save_file(dict(data), temporary) + with temporary.open("rb") as published: + os.fsync(published.fileno()) + temporary.replace(target) + _fsync_directory(target.parent) + finally: + temporary.unlink(missing_ok=True) + + def get_or_create( + self, + request_id: str, + create: Callable[[], dict[str, torch.Tensor]], + validate: Callable[[dict[str, torch.Tensor]], None], + ) -> ArtifactResult: + """Load a published artifact or elect one caller to create it.""" + self._validate_request_id(request_id) + try: + lock = self._request_lock(request_id) + with lock as waited: + now = time.time() + stale_temps = self._remove_stale_temps_for_key(request_id, now) + artifact_path = self.artifact_path(request_id) + expired = 0 + invalid = 0 + if artifact_path.exists() and self._is_expired(artifact_path, now): + artifact_path.unlink(missing_ok=True) + expired = 1 + + if artifact_path.exists(): + try: + data = load_file(artifact_path) + validate(data) + except Exception: + artifact_path.unlink(missing_ok=True) + invalid = 1 + else: + self._record( + logical_requests=1, + hits=1, + coalesced_waiters=int(waited), + stale_temps_removed=stale_temps, + expired_artifacts_removed=expired, + ) + return ArtifactResult( + data=data, + request_id=request_id, + path=artifact_path, + cache_hit=True, + coalesced=waited, + ) + + try: + data = create() + self._validate_tensors(data) + validate(data) + except BaseException: + self._record( + logical_requests=1, + misses=1, + generation_failures=1, + retry_generations=int(waited), + invalid_artifacts_removed=invalid, + expired_artifacts_removed=expired, + stale_temps_removed=stale_temps, + ) + raise + try: + self._publish(request_id, data, artifact_path) + except BaseException: + self._record( + logical_requests=1, + misses=1, + publish_failures=1, + retry_generations=int(waited), + invalid_artifacts_removed=invalid, + expired_artifacts_removed=expired, + stale_temps_removed=stale_temps, + ) + raise + + self._record( + logical_requests=1, + misses=1, + publishes=1, + retry_generations=int(waited), + invalid_artifacts_removed=invalid, + expired_artifacts_removed=expired, + stale_temps_removed=stale_temps, + ) + return ArtifactResult( + data=data, + request_id=request_id, + path=artifact_path, + cache_hit=False, + coalesced=False, + ) + except ArtifactLockTimeoutError: + self._record(logical_requests=1, lock_timeouts=1) + raise + + def cleanup_stale(self, *, now: float | None = None) -> dict[str, int]: + """Remove expired artifacts and abandoned temporary publications.""" + current_time = time.time() if now is None else now + expired = 0 + stale_temps = 0 + + for path in self._artifacts.glob("*/*.safetensors"): + request_id = path.stem + if _REQUEST_ID_PATTERN.fullmatch(request_id) is None: + continue + try: + with self._request_lock(request_id, timeout_seconds=0): + if path.exists() and self._is_expired(path, current_time): + path.unlink(missing_ok=True) + expired += 1 + except ArtifactLockTimeoutError: + continue + + for path in self._artifacts.glob("*/.*.tmp"): + name = path.name + request_id = name[1:65] + if _REQUEST_ID_PATTERN.fullmatch(request_id) is None: + continue + try: + with self._request_lock(request_id, timeout_seconds=0): + if ( + path.exists() + and current_time - path.stat().st_mtime + >= self.stale_temp_seconds + ): + path.unlink(missing_ok=True) + stale_temps += 1 + except ArtifactLockTimeoutError: + continue + + if expired or stale_temps: + self._record( + expired_artifacts_removed=expired, + stale_temps_removed=stale_temps, + ) + return { + "expired_artifacts_removed": expired, + "stale_temps_removed": stale_temps, + } diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py index 4e612373b..9549ce971 100644 --- a/src/speculators/train/data.py +++ b/src/speculators/train/data.py @@ -2,6 +2,7 @@ import math import os import random +import uuid import warnings from collections.abc import Callable from os import PathLike @@ -12,9 +13,14 @@ import torch import torch.nn.functional as F # noqa: N812 from datasets import load_from_disk +from safetensors.torch import save_file from torch.utils.data import Dataset from hs_connectors import FileTransfer, HiddenStatesTransfer +from speculators.data_generation.artifact_cache import ( + HiddenStateArtifactCache, + canonical_hidden_state_request_id, +) from speculators.data_generation.offline import check_hidden_states from speculators.data_generation.vllm_client import ( DEFAULT_MAX_RETRIES, @@ -213,6 +219,17 @@ def __getitem__(self, index) -> BatchType | None: return data +def _atomic_save_hs_file(data: dict[str, torch.Tensor], file_path: Path) -> None: + temporary = file_path.parent / ( + f".{file_path.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp" + ) + try: + save_file(data, temporary) + temporary.replace(file_path) + finally: + temporary.unlink(missing_ok=True) + + class ArrowDataset(BaseDataset): def __init__( self, @@ -228,6 +245,10 @@ def __init__( model: str | None = None, request_timeout: float | None = DEFAULT_REQUEST_TIMEOUT, max_retries: int = DEFAULT_MAX_RETRIES, + shared_artifacts_path: str | PathLike | None = None, + shared_artifacts_namespace: str | None = None, + shared_artifacts_ttl_seconds: float | None = 3600.0, + shared_artifacts_lock_timeout_seconds: float = 300.0, ): self.data = load_from_disk(datapath) self.start_file_idx = 0 @@ -252,6 +273,18 @@ def __init__( self.model = model self.request_timeout = request_timeout self.max_retries = max_retries + self.shared_artifacts_namespace = shared_artifacts_namespace + self.artifact_cache = ( + HiddenStateArtifactCache( + shared_artifacts_path, + artifact_ttl_seconds=shared_artifacts_ttl_seconds, + lock_timeout_seconds=shared_artifacts_lock_timeout_seconds, + ) + if shared_artifacts_path is not None + else None + ) + if self.artifact_cache is not None: + self.artifact_cache.cleanup_stale() # Delay super init so that `_compute_approx_lengths` has required data super().__init__(max_len, transform, hidden_states_dtype) @@ -289,6 +322,32 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None: client_item = build_client_item(dataset_item) try: + if self.artifact_cache is not None: + request_id = canonical_hidden_state_request_id( + self.model, # type:ignore[arg-type] + client_item, + namespace=self.shared_artifacts_namespace, + ) + result = self.artifact_cache.get_or_create( + request_id, + lambda: self._generate_shared_hs(dataset_item, client_item), + lambda data: check_hidden_states( + data, dataset_item["input_ids"].tolist() + ), + ) + loaded_hs = result.data + if self.on_generate == "cache" and isinstance( + self.transfer, FileTransfer + ): + file_idx = self._map_to_file_idx(index) + target_path = ( + self.transfer.hidden_states_path + / f"hs_{file_idx}.safetensors" + ) + target_path.parent.mkdir(parents=True, exist_ok=True) + _atomic_save_hs_file(loaded_hs, target_path) + return loaded_hs + handle = generate_hidden_states( self.client, # type:ignore[arg-type] self.model, # type:ignore[arg-type] @@ -320,6 +379,27 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None: return loaded_hs + def _generate_shared_hs( + self, dataset_item: dict, client_item: ClientItem + ) -> dict[str, torch.Tensor]: + handle: str | None = None + try: + handle = generate_hidden_states( + self.client, # type:ignore[arg-type] + self.model, # type:ignore[arg-type] + client_item, + timeout=self.request_timeout, + max_retries=self.max_retries, + ) + loaded_hs = self.transfer.get_generated(handle) + if loaded_hs is None: + raise ValueError(f"Failed to load hidden states for handle {handle}") + check_hidden_states(loaded_hs, dataset_item["input_ids"].tolist()) + return loaded_hs + finally: + if handle is not None: + self.transfer.delete(handle) + def _get_raw_data(self, index): file_idx = self._map_to_file_idx(index) loaded_hs = self.transfer.get_cached(file_idx) diff --git a/src/speculators/train/dataloader.py b/src/speculators/train/dataloader.py index 754e932db..001d6b696 100644 --- a/src/speculators/train/dataloader.py +++ b/src/speculators/train/dataloader.py @@ -81,6 +81,10 @@ def create_train_val_loaders( num_workers: int, prefetch_factor: int, preprocess: Callable[[BatchType], BatchType] | None, + shared_artifacts_path: str | None = None, + shared_artifacts_namespace: str | None = None, + shared_artifacts_ttl_seconds: float | None = 3600.0, + shared_artifacts_lock_timeout_seconds: float = 300.0, train_data_ratio: float = 0.9, ) -> tuple[DataLoader, DataLoader]: """Create training and validation DataLoaders. @@ -127,6 +131,12 @@ def create_train_val_loaders( hidden_states_dtype=hidden_states_dtype, request_timeout=request_timeout, max_retries=max_retries, + shared_artifacts_path=shared_artifacts_path, + shared_artifacts_namespace=shared_artifacts_namespace, + shared_artifacts_ttl_seconds=shared_artifacts_ttl_seconds, + shared_artifacts_lock_timeout_seconds=( + shared_artifacts_lock_timeout_seconds + ), ) val_dataset = ArrowDataset( datapath=data_path, @@ -140,6 +150,12 @@ def create_train_val_loaders( hidden_states_dtype=hidden_states_dtype, request_timeout=request_timeout, max_retries=max_retries, + shared_artifacts_path=shared_artifacts_path, + shared_artifacts_namespace=shared_artifacts_namespace, + shared_artifacts_ttl_seconds=shared_artifacts_ttl_seconds, + shared_artifacts_lock_timeout_seconds=( + shared_artifacts_lock_timeout_seconds + ), ) train_loader = _setup_dataloader( diff --git a/tests/unit/benchmarks/test_independent_consumers.py b/tests/unit/benchmarks/test_independent_consumers.py index 98c72b3f4..23108439e 100644 --- a/tests/unit/benchmarks/test_independent_consumers.py +++ b/tests/unit/benchmarks/test_independent_consumers.py @@ -16,6 +16,7 @@ ProducerSpec, RequestEvent, ScenarioSpec, + analyze_cache_accounting, analyze_consumer_steps, analyze_scenario, canonical_request_key, @@ -237,6 +238,10 @@ def test_analysis_separates_warmup_and_requires_exact_multiplicity(): assert result["valid"] assert result["request_accounting"]["requests"] == 12 assert result["request_accounting"]["qualifying_shared_samples"] == 3 + assert ( + result["request_accounting"]["per_consumer_completions_semantics"] + == "logical_consumer" + ) assert result["steady_state"]["duration_seconds"] == pytest.approx(0.22) assert result["steady_state"]["completions"] == 7 @@ -267,6 +272,101 @@ def test_analysis_fails_closed_on_failed_or_duplicate_completion(): ) +def test_publish_once_analysis_accepts_one_service_call_per_shared_sample(): + events = [ + _event("c0", "warmup", 0.1), + _event("c1", "sample-a", 0.2), + _event("c2", "sample-b", 0.3), + _event("c0", "sample-c", 0.4), + ] + + result = analyze_scenario( + _scenario(multiplicity=1), events, started_at=0.0, finished_at=1.0 + ) + + assert result["valid"] + assert result["request_accounting"]["requests"] == 4 + assert result["request_accounting"]["qualifying_shared_samples"] == 3 + assert ( + result["request_accounting"]["per_consumer_completions_semantics"] + == "service_request_owner" + ) + assert result["steady_state"]["mode"] == "publish_once_service" + assert ( + result["steady_state"]["completions_per_consumer_semantics"] + == "service_request_owner" + ) + assert result["steady_state"]["completions"] == 3 + + +def _cache_stats(**updates: int) -> dict[str, int]: + stats = { + "schema_version": 1, + "logical_requests": 12, + "hits": 8, + "misses": 4, + "coalesced_waiters": 2, + "retry_generations": 0, + "publishes": 4, + "generation_failures": 0, + "publish_failures": 0, + "invalid_artifacts_removed": 0, + "expired_artifacts_removed": 0, + "stale_temps_removed": 0, + "lock_timeouts": 0, + } + stats.update(updates) + return stats + + +def test_cache_accounting_proves_publish_once_fanout(): + result = analyze_cache_accounting( + _scenario(multiplicity=1), _cache_stats(), service_completions=4 + ) + + assert result["valid"] + assert result["stats"]["logical_requests"] == 12 + assert result["stats"]["misses"] == result["stats"]["publishes"] == 4 + assert result["stats"]["hits"] == 8 + + +@pytest.mark.parametrize( + ("updates", "reason"), + [ + ({"schema_version": 2}, "schema_version"), + ({"logical_requests": 11}, "logical_requests"), + ({"hits": 7}, "cache hits"), + ({"misses": 3}, "cache misses"), + ({"publishes": 3}, "cache publishes"), + ({"retry_generations": 1}, "retry_generations"), + ({"generation_failures": 1}, "generation_failures"), + ({"coalesced_waiters": 9}, "coalesced_waiters"), + ], +) +def test_cache_accounting_fails_closed_on_invalid_counters(updates, reason): + result = analyze_cache_accounting( + _scenario(multiplicity=1), + _cache_stats(**updates), + service_completions=4, + ) + + assert not result["valid"] + assert any(reason in message for message in result["invalid_reasons"]) + + +def test_publish_once_requires_cache_accounting(): + publish_once = analyze_cache_accounting( + _scenario(multiplicity=1), None, service_completions=4 + ) + unshared = analyze_cache_accounting( + _scenario(multiplicity=3), None, service_completions=12 + ) + + assert not publish_once["valid"] + assert publish_once["invalid_reasons"] + assert unshared["valid"] + + def test_consumer_step_analysis_excludes_warmup_and_reports_percentiles(tmp_path): log_path = tmp_path / "consumer.log" log_path.write_text( diff --git a/tests/unit/data_generation/test_artifact_cache.py b/tests/unit/data_generation/test_artifact_cache.py new file mode 100644 index 000000000..80f4c7010 --- /dev/null +++ b/tests/unit/data_generation/test_artifact_cache.py @@ -0,0 +1,345 @@ +from __future__ import annotations + +import multiprocessing +import os +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import pytest +import torch +from safetensors.torch import load_file, save_file + +import speculators.data_generation.artifact_cache as artifact_cache_module +from speculators.data_generation.artifact_cache import ( + ArtifactLockTimeoutError, + HiddenStateArtifactCache, + canonical_hidden_state_request_id, +) +from speculators.data_generation.offline import check_hidden_states + + +def _tensors(tokens: tuple[int, ...] = (1, 2, 3)) -> dict[str, torch.Tensor]: + return { + "token_ids": torch.tensor(tokens), + "hidden_states": torch.arange(len(tokens) * 4, dtype=torch.float32).reshape( + len(tokens), 4 + ), + } + + +def _validate(data: dict[str, torch.Tensor]) -> None: + check_hidden_states(data, [1, 2, 3]) + + +def _request_id() -> str: + return canonical_hidden_state_request_id("model", {"input_ids": [1, 2, 3]}) + + +def _spawn_cache_worker(root, ready, start, generation_count, results): + cache = HiddenStateArtifactCache( + root, + artifact_ttl_seconds=None, + lock_timeout_seconds=10, + lock_poll_seconds=0.01, + ) + ready.put(True) + start.wait(10) + + def create(): + with generation_count.get_lock(): + generation_count.value += 1 + time.sleep(0.4) + return _tensors() + + result = cache.get_or_create(_request_id(), create, _validate) + results.put((result.cache_hit, result.data["token_ids"].tolist())) + + +def test_canonical_request_id_is_stable_and_covers_semantics(): + first = canonical_hidden_state_request_id( + "model", + { + "input_ids": [1, 2, 3], + "messages": [{"role": "user", "content": {"text": "hello", "x": 1}}], + }, + namespace="layers:2,18,33", + ) + reordered = canonical_hidden_state_request_id( + "model", + { + "messages": [{"content": {"x": 1, "text": "hello"}, "role": "user"}], + "input_ids": [1, 2, 3], + }, + namespace="layers:2,18,33", + ) + + assert first == reordered + assert first != canonical_hidden_state_request_id( + "model", {"input_ids": [1, 2, 4]}, namespace="layers:2,18,33" + ) + assert first != canonical_hidden_state_request_id( + "model", {"input_ids": [1, 2, 3]}, namespace="layers:2,18,36" + ) + assert first != canonical_hidden_state_request_id( + "other-model", {"input_ids": [1, 2, 3]}, namespace="layers:2,18,33" + ) + + +@pytest.mark.parametrize("tokens", [[], [1, True], [[1, 2]]]) +def test_canonical_request_id_rejects_invalid_tokens(tokens): + with pytest.raises(ValueError, match="input_ids"): + canonical_hidden_state_request_id("model", {"input_ids": tokens}) + + +def test_sequential_request_publishes_once_then_hits(tmp_path): + cache = HiddenStateArtifactCache(tmp_path, artifact_ttl_seconds=None) + generations = 0 + + def create(): + nonlocal generations + generations += 1 + return _tensors() + + first = cache.get_or_create(_request_id(), create, _validate) + second = cache.get_or_create(_request_id(), create, _validate) + + assert generations == 1 + assert not first.cache_hit + assert second.cache_hit + assert torch.equal(first.data["hidden_states"], second.data["hidden_states"]) + assert load_file(first.path)["token_ids"].tolist() == [1, 2, 3] + assert cache.snapshot_stats() == { + "schema_version": 1, + "logical_requests": 2, + "hits": 1, + "misses": 1, + "coalesced_waiters": 0, + "retry_generations": 0, + "publishes": 1, + "generation_failures": 0, + "publish_failures": 0, + "invalid_artifacts_removed": 0, + "expired_artifacts_removed": 0, + "stale_temps_removed": 0, + "lock_timeouts": 0, + } + + +def test_independent_processes_coalesce_one_generation(tmp_path): + context = multiprocessing.get_context("spawn") + ready = context.Queue() + start = context.Event() + generation_count = context.Value("i", 0) + results = context.Queue() + processes = [ + context.Process( + target=_spawn_cache_worker, + args=(str(tmp_path), ready, start, generation_count, results), + ) + for _ in range(3) + ] + try: + for process in processes: + process.start() + for _ in processes: + assert ready.get(timeout=20) + start.set() + for process in processes: + process.join(timeout=20) + assert process.exitcode == 0 + outcomes = [results.get(timeout=5) for _ in processes] + finally: + for process in processes: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + + assert generation_count.value == 1 + assert sum(cache_hit for cache_hit, _tokens in outcomes) == 2 + assert all(tokens == [1, 2, 3] for _cache_hit, tokens in outcomes) + stats = HiddenStateArtifactCache( + tmp_path, artifact_ttl_seconds=None + ).snapshot_stats() + assert stats["logical_requests"] == 3 + assert stats["misses"] == 1 + assert stats["hits"] == 2 + assert stats["publishes"] == 1 + assert stats["coalesced_waiters"] >= 1 + + +def test_failed_owner_allows_waiter_to_retry(tmp_path): + cache = HiddenStateArtifactCache( + tmp_path, + artifact_ttl_seconds=None, + lock_timeout_seconds=5, + lock_poll_seconds=0.01, + ) + start = threading.Barrier(2) + owner_entered = threading.Event() + release_owner = threading.Event() + attempt_lock = threading.Lock() + attempts = 0 + + def create(): + nonlocal attempts + with attempt_lock: + attempts += 1 + attempt = attempts + if attempt == 1: + owner_entered.set() + assert release_owner.wait(5) + raise RuntimeError("producer failed") + return _tensors() + + def request(): + start.wait() + try: + result = cache.get_or_create(_request_id(), create, _validate) + except RuntimeError: + return "failed" + return "hit" if result.cache_hit else "published" + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(request) for _ in range(2)] + assert owner_entered.wait(5) + time.sleep(0.1) + release_owner.set() + outcomes = [future.result(timeout=10) for future in futures] + + assert sorted(outcomes) == ["failed", "published"] + assert attempts == 2 + assert cache.artifact_path(_request_id()).is_file() + stats = cache.snapshot_stats() + assert stats["generation_failures"] == 1 + assert stats["publishes"] == 1 + assert stats["retry_generations"] == 1 + + +def test_partial_publish_is_never_visible_and_is_retried(tmp_path, monkeypatch): + cache = HiddenStateArtifactCache( + tmp_path, artifact_ttl_seconds=None, stale_temp_seconds=0.01 + ) + original_publish = cache._publish + + def fail_publish(request_id, _data, target): + target.parent.mkdir(parents=True, exist_ok=True) + partial = target.parent / f".{request_id}.dead.tmp" + partial.write_bytes(b"partial") + raise OSError("disk write failed") + + monkeypatch.setattr(cache, "_publish", fail_publish) + with pytest.raises(OSError, match="disk write failed"): + cache.get_or_create(_request_id(), _tensors, _validate) + assert not cache.artifact_path(_request_id()).exists() + + partial = next(cache.artifact_path(_request_id()).parent.glob(".*.tmp")) + old = time.time() - 1 + os.utime(partial, (old, old)) + monkeypatch.setattr(cache, "_publish", original_publish) + result = cache.get_or_create(_request_id(), _tensors, _validate) + + assert not result.cache_hit + assert result.path.is_file() + assert not partial.exists() + stats = cache.snapshot_stats() + assert stats["publish_failures"] == 1 + assert stats["stale_temps_removed"] == 1 + + +def test_atomic_name_is_absent_while_temporary_file_is_written(tmp_path, monkeypatch): + cache = HiddenStateArtifactCache(tmp_path, artifact_ttl_seconds=None) + write_started = threading.Event() + finish_write = threading.Event() + original_save = save_file + + def paused_save(data, path): + Path(path).write_bytes(b"partial") + write_started.set() + assert finish_write.wait(5) + original_save(data, path) + + monkeypatch.setattr(artifact_cache_module, "save_file", paused_save) + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit( + cache.get_or_create, _request_id(), _tensors, _validate + ) + assert write_started.wait(5) + assert not cache.artifact_path(_request_id()).exists() + finish_write.set() + result = future.result(timeout=10) + + assert result.path.is_file() + assert load_file(result.path)["token_ids"].tolist() == [1, 2, 3] + + +def test_corrupt_artifact_is_removed_before_regeneration(tmp_path): + cache = HiddenStateArtifactCache(tmp_path, artifact_ttl_seconds=None) + target = cache.artifact_path(_request_id()) + target.parent.mkdir(parents=True) + target.write_bytes(b"not safetensors") + + result = cache.get_or_create(_request_id(), _tensors, _validate) + + assert not result.cache_hit + assert load_file(target)["token_ids"].tolist() == [1, 2, 3] + assert cache.snapshot_stats()["invalid_artifacts_removed"] == 1 + + +def test_cleanup_removes_expired_artifact_and_stale_temp(tmp_path): + cache = HiddenStateArtifactCache( + tmp_path, artifact_ttl_seconds=1, stale_temp_seconds=1 + ) + result = cache.get_or_create(_request_id(), _tensors, _validate) + stale_temp = result.path.parent / f".{_request_id()}.dead.tmp" + stale_temp.write_bytes(b"partial") + old = time.time() - 10 + os.utime(result.path, (old, old)) + os.utime(stale_temp, (old, old)) + + cleanup = cache.cleanup_stale(now=time.time()) + + assert cleanup == { + "expired_artifacts_removed": 1, + "stale_temps_removed": 1, + } + assert not result.path.exists() + assert not stale_temp.exists() + stats = cache.snapshot_stats() + assert stats["expired_artifacts_removed"] == 1 + assert stats["stale_temps_removed"] == 1 + + +def test_lock_timeout_is_counted(tmp_path): + owner = HiddenStateArtifactCache( + tmp_path, artifact_ttl_seconds=None, lock_timeout_seconds=5 + ) + waiter = HiddenStateArtifactCache( + tmp_path, + artifact_ttl_seconds=None, + lock_timeout_seconds=0.05, + lock_poll_seconds=0.01, + ) + locked = threading.Event() + release = threading.Event() + + def hold_lock(): + with owner._request_lock(_request_id()): + locked.set() + assert release.wait(5) + + thread = threading.Thread(target=hold_lock) + thread.start() + assert locked.wait(5) + try: + with pytest.raises(ArtifactLockTimeoutError): + waiter.get_or_create(_request_id(), _tensors, _validate) + finally: + release.set() + thread.join(timeout=5) + + stats = waiter.snapshot_stats() + assert stats["logical_requests"] == 1 + assert stats["lock_timeouts"] == 1 + assert not waiter.artifact_path(_request_id()).exists() diff --git a/tests/unit/train/test_cli_args.py b/tests/unit/train/test_cli_args.py index 38e30ad56..da09e48b9 100644 --- a/tests/unit/train/test_cli_args.py +++ b/tests/unit/train/test_cli_args.py @@ -15,6 +15,36 @@ def _parse(monkeypatch, extra: list[str]): return parse_args() +def test_shared_hidden_state_cache_is_opt_in(monkeypatch): + args = _parse(monkeypatch, []) + + assert args.shared_hidden_states_path is None + assert args.shared_hidden_states_namespace is None + assert args.shared_hidden_states_ttl == 3600.0 + assert args.shared_hidden_states_lock_timeout == 300.0 + + +def test_shared_hidden_state_cache_arguments(monkeypatch): + args = _parse( + monkeypatch, + [ + "--shared-hidden-states-path", + "shared-cache", + "--shared-hidden-states-namespace", + "layers:2,18,33", + "--shared-hidden-states-ttl", + "0", + "--shared-hidden-states-lock-timeout", + "45", + ], + ) + + assert args.shared_hidden_states_path == "shared-cache" + assert args.shared_hidden_states_namespace == "layers:2,18,33" + assert args.shared_hidden_states_ttl == 0 + assert args.shared_hidden_states_lock_timeout == 45 + + # --------------------------------------------------------------------------- # Ensure CLI args flow correctly through vars(args) into get_trainer_kwargs # --------------------------------------------------------------------------- diff --git a/tests/unit/train/test_shared_artifacts.py b/tests/unit/train/test_shared_artifacts.py new file mode 100644 index 000000000..55308078a --- /dev/null +++ b/tests/unit/train/test_shared_artifacts.py @@ -0,0 +1,254 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest +import torch +from datasets import Dataset +from safetensors.torch import load_file, save_file + +import speculators.train.data as data_module +import speculators.train.dataloader as dataloader_module +from speculators.train.data import ArrowDataset + +if TYPE_CHECKING: + from pathlib import Path + + +def _write_dataset(path: Path) -> None: + dataset = Dataset.from_dict( + { + "input_ids": [[1, 2, 3]], + "loss_mask": [[0, 1, 1]], + "seq_len": [3], + } + ).with_format("torch") + dataset.save_to_disk(path) + + +def _hidden_states() -> dict[str, torch.Tensor]: + return { + "token_ids": torch.tensor([1, 2, 3]), + "hidden_states": torch.arange(12, dtype=torch.float32).reshape(3, 4), + } + + +def _arrow_dataset( + data_path: Path, + *, + shared_path: Path | None, + hidden_states_path: Path, + on_generate: str = "delete", +) -> ArrowDataset: + dataset = ArrowDataset( + max_len=128, + datapath=data_path, + hidden_states_path=hidden_states_path, + model="model", + on_missing="generate", + on_generate=on_generate, + shared_artifacts_path=shared_path, + shared_artifacts_ttl_seconds=None, + ) + dataset.client = object() + return dataset + + +def _successful_generator(service_path: Path, calls: list[Path]): + def generate(*_args, **_kwargs): + path = service_path / f"request-{len(calls)}.safetensors" + save_file(_hidden_states(), path) + calls.append(path) + return str(path) + + return generate + + +def test_shared_dataset_requests_publish_once_and_delete_service_temporary( + tmp_path, monkeypatch +): + data_path = tmp_path / "data" + service_path = tmp_path / "service" + shared_path = tmp_path / "shared" + service_path.mkdir() + _write_dataset(data_path) + calls: list[Path] = [] + monkeypatch.setattr( + data_module, + "generate_hidden_states", + _successful_generator(service_path, calls), + ) + first = _arrow_dataset( + data_path, + shared_path=shared_path, + hidden_states_path=tmp_path / "first-index", + ) + second = _arrow_dataset( + data_path, + shared_path=shared_path, + hidden_states_path=tmp_path / "second-index", + ) + + first_data = first._maybe_generate_hs(0) + second_data = second._maybe_generate_hs(0) + + assert first_data is not None + assert second_data is not None + assert torch.equal(first_data["hidden_states"], second_data["hidden_states"]) + assert len(calls) == 1 + assert not calls[0].exists() + assert first.artifact_cache is not None + stats = first.artifact_cache.snapshot_stats() + assert stats["logical_requests"] == 2 + assert stats["misses"] == 1 + assert stats["hits"] == 1 + assert stats["publishes"] == 1 + + +def test_unconfigured_dataset_keeps_existing_per_request_delete_behavior( + tmp_path, monkeypatch +): + data_path = tmp_path / "data" + service_path = tmp_path / "service" + service_path.mkdir() + _write_dataset(data_path) + calls: list[Path] = [] + monkeypatch.setattr( + data_module, + "generate_hidden_states", + _successful_generator(service_path, calls), + ) + first = _arrow_dataset( + data_path, + shared_path=None, + hidden_states_path=tmp_path / "first-index", + ) + second = _arrow_dataset( + data_path, + shared_path=None, + hidden_states_path=tmp_path / "second-index", + ) + + assert first._maybe_generate_hs(0) is not None + assert second._maybe_generate_hs(0) is not None + + assert first.artifact_cache is None + assert len(calls) == 2 + assert all(not path.exists() for path in calls) + + +def test_shared_dataset_does_not_publish_partial_service_artifact( + tmp_path, monkeypatch +): + data_path = tmp_path / "data" + service_path = tmp_path / "service" + shared_path = tmp_path / "shared" + service_path.mkdir() + _write_dataset(data_path) + calls: list[Path] = [] + + def generate(*_args, **_kwargs): + path = service_path / f"request-{len(calls)}.safetensors" + if not calls: + path.write_bytes(b"partial") + else: + save_file(_hidden_states(), path) + calls.append(path) + return str(path) + + monkeypatch.setattr(data_module, "generate_hidden_states", generate) + dataset = _arrow_dataset( + data_path, + shared_path=shared_path, + hidden_states_path=tmp_path / "index", + ) + + with pytest.warns(UserWarning, match="Failed to load/cache"): + assert dataset._maybe_generate_hs(0) is None + assert not calls[0].exists() + assert list((shared_path / "artifacts").glob("*/*.safetensors")) == [] + + loaded = dataset._maybe_generate_hs(0) + + assert loaded is not None + assert len(calls) == 2 + assert not calls[1].exists() + assert dataset.artifact_cache is not None + stats = dataset.artifact_cache.snapshot_stats() + assert stats["generation_failures"] == 1 + assert stats["publishes"] == 1 + + +def test_shared_dataset_preserves_on_generate_index_cache(tmp_path, monkeypatch): + data_path = tmp_path / "data" + service_path = tmp_path / "service" + shared_path = tmp_path / "shared" + index_path = tmp_path / "index" + service_path.mkdir() + _write_dataset(data_path) + calls: list[Path] = [] + monkeypatch.setattr( + data_module, + "generate_hidden_states", + _successful_generator(service_path, calls), + ) + dataset = _arrow_dataset( + data_path, + shared_path=shared_path, + hidden_states_path=index_path, + on_generate="cache", + ) + + assert dataset._maybe_generate_hs(0) is not None + + indexed = index_path / "hs_0.safetensors" + assert indexed.is_file() + assert load_file(indexed)["token_ids"].tolist() == [1, 2, 3] + assert len(calls) == 1 + + +def test_train_and_validation_loaders_share_artifact_configuration(monkeypatch): + dataset_kwargs = [] + + def fake_arrow_dataset(**kwargs): + dataset_kwargs.append(kwargs) + return object() + + monkeypatch.setattr(dataloader_module, "ArrowDataset", fake_arrow_dataset) + monkeypatch.setattr( + dataloader_module, + "_setup_dataloader", + lambda dataset, *_args, **_kwargs: dataset, + ) + + dataloader_module.create_train_val_loaders( + data_path="data", + train_data_ratio=0.9, + total_seq_len=128, + hidden_states_dtype=torch.bfloat16, + noise_std=0.0, + legacy_data=False, + hidden_states_path=None, + vllm_endpoint="http://producer/v1", + on_missing="generate", + on_generate="delete", + verifier_name_or_path="model", + request_timeout=10, + max_retries=2, + hidden_size=4, + num_target_layers=3, + num_workers=0, + prefetch_factor=1, + preprocess=None, + shared_artifacts_path="shared", + shared_artifacts_namespace="layers:2,18,33", + shared_artifacts_ttl_seconds=None, + shared_artifacts_lock_timeout_seconds=45, + ) + + assert len(dataset_kwargs) == 2 + for kwargs in dataset_kwargs: + assert kwargs["shared_artifacts_path"] == "shared" + assert kwargs["shared_artifacts_namespace"] == "layers:2,18,33" + assert kwargs["shared_artifacts_ttl_seconds"] is None + assert kwargs["shared_artifacts_lock_timeout_seconds"] == 45 From 3c2c4d925c27a248c93d0627443689d560752cfa Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Fri, 17 Jul 2026 15:55:22 +0800 Subject: [PATCH 03/20] fix: harden shared artifact cleanup Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- .../independent_consumer_fanout/README.md | 3 ++ .../data_generation/artifact_cache.py | 15 +++++- src/speculators/train/data.py | 5 +- .../data_generation/test_artifact_cache.py | 8 +++- tests/unit/train/test_shared_artifacts.py | 48 +++++++++++++++++++ 5 files changed, 75 insertions(+), 4 deletions(-) diff --git a/benchmarks/independent_consumer_fanout/README.md b/benchmarks/independent_consumer_fanout/README.md index 85c31e5b6..3f13c5e92 100644 --- a/benchmarks/independent_consumer_fanout/README.md +++ b/benchmarks/independent_consumer_fanout/README.md @@ -56,3 +56,6 @@ completion map identify which consumer owned each service miss. They do not repr logical trainer progress: cache logical-request totals and each independent consumer's `consumer_step_times` provide that evidence. The report labels both maps with `service_request_owner` to make this distinction explicit. +Despite their names, `warmup_completions_per_consumer` and +`minimum_steady_completions_per_consumer` are service-wide publication thresholds in +this mode; the corresponding step fields apply separately to every consumer. diff --git a/src/speculators/data_generation/artifact_cache.py b/src/speculators/data_generation/artifact_cache.py index 2d6084cc3..f02a3d747 100644 --- a/src/speculators/data_generation/artifact_cache.py +++ b/src/speculators/data_generation/artifact_cache.py @@ -373,7 +373,7 @@ def get_or_create( raise def cleanup_stale(self, *, now: float | None = None) -> dict[str, int]: - """Remove expired artifacts and abandoned temporary publications.""" + """Remove expired artifacts and abandoned temporary cache writes.""" current_time = time.time() if now is None else now expired = 0 stale_temps = 0 @@ -407,6 +407,19 @@ def cleanup_stale(self, *, now: float | None = None) -> dict[str, int]: except ArtifactLockTimeoutError: continue + with self._stats_lock(fcntl.LOCK_EX): + for path in self.root.glob(".stats.*.tmp"): + try: + if ( + path.exists() + and current_time - path.stat().st_mtime + >= self.stale_temp_seconds + ): + path.unlink(missing_ok=True) + stale_temps += 1 + except FileNotFoundError: + pass + if expired or stale_temps: self._record( expired_artifacts_removed=expired, diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py index 9549ce971..1ed7612f6 100644 --- a/src/speculators/train/data.py +++ b/src/speculators/train/data.py @@ -383,6 +383,7 @@ def _generate_shared_hs( self, dataset_item: dict, client_item: ClientItem ) -> dict[str, torch.Tensor]: handle: str | None = None + retrieved = False try: handle = generate_hidden_states( self.client, # type:ignore[arg-type] @@ -392,12 +393,14 @@ def _generate_shared_hs( max_retries=self.max_retries, ) loaded_hs = self.transfer.get_generated(handle) + retrieved = True if loaded_hs is None: raise ValueError(f"Failed to load hidden states for handle {handle}") check_hidden_states(loaded_hs, dataset_item["input_ids"].tolist()) return loaded_hs finally: - if handle is not None: + # A failed retrieval may still have an in-flight backend writer. + if handle is not None and retrieved: self.transfer.delete(handle) def _get_raw_data(self, index): diff --git a/tests/unit/data_generation/test_artifact_cache.py b/tests/unit/data_generation/test_artifact_cache.py index 80f4c7010..70e8e76c3 100644 --- a/tests/unit/data_generation/test_artifact_cache.py +++ b/tests/unit/data_generation/test_artifact_cache.py @@ -294,21 +294,25 @@ def test_cleanup_removes_expired_artifact_and_stale_temp(tmp_path): result = cache.get_or_create(_request_id(), _tensors, _validate) stale_temp = result.path.parent / f".{_request_id()}.dead.tmp" stale_temp.write_bytes(b"partial") + stale_stats_temp = tmp_path / ".stats.123.dead.tmp" + stale_stats_temp.write_bytes(b"partial") old = time.time() - 10 os.utime(result.path, (old, old)) os.utime(stale_temp, (old, old)) + os.utime(stale_stats_temp, (old, old)) cleanup = cache.cleanup_stale(now=time.time()) assert cleanup == { "expired_artifacts_removed": 1, - "stale_temps_removed": 1, + "stale_temps_removed": 2, } assert not result.path.exists() assert not stale_temp.exists() + assert not stale_stats_temp.exists() stats = cache.snapshot_stats() assert stats["expired_artifacts_removed"] == 1 - assert stats["stale_temps_removed"] == 1 + assert stats["stale_temps_removed"] == 2 def test_lock_timeout_is_counted(tmp_path): diff --git a/tests/unit/train/test_shared_artifacts.py b/tests/unit/train/test_shared_artifacts.py index 55308078a..f43fe49bb 100644 --- a/tests/unit/train/test_shared_artifacts.py +++ b/tests/unit/train/test_shared_artifacts.py @@ -58,6 +58,7 @@ def _successful_generator(service_path: Path, calls: list[Path]): def generate(*_args, **_kwargs): path = service_path / f"request-{len(calls)}.safetensors" save_file(_hidden_states(), path) + (path.parent / f"{path.name}.lock").touch() calls.append(path) return str(path) @@ -97,6 +98,7 @@ def test_shared_dataset_requests_publish_once_and_delete_service_temporary( assert torch.equal(first_data["hidden_states"], second_data["hidden_states"]) assert len(calls) == 1 assert not calls[0].exists() + assert not (calls[0].parent / f"{calls[0].name}.lock").exists() assert first.artifact_cache is not None stats = first.artifact_cache.snapshot_stats() assert stats["logical_requests"] == 2 @@ -179,6 +181,52 @@ def generate(*_args, **_kwargs): assert stats["publishes"] == 1 +def test_shared_dataset_preserves_service_artifact_after_lock_timeout( + tmp_path, monkeypatch +): + data_path = tmp_path / "data" + service_path = tmp_path / "service" + shared_path = tmp_path / "shared" + service_path.mkdir() + _write_dataset(data_path) + artifact_path = service_path / "request.safetensors" + artifact_path.write_bytes(b"partial") + lock_path = artifact_path.parent / f"{artifact_path.name}.lock" + lock_path.touch() + + def time_out_waiting_for_lock(*_args, **_kwargs): + raise TimeoutError("active") + + monkeypatch.setattr( + data_module, + "generate_hidden_states", + lambda *_args, **_kwargs: str(artifact_path), + ) + monkeypatch.setattr( + data_module, + "wait_for_lock", + time_out_waiting_for_lock, + ) + dataset = _arrow_dataset( + data_path, + shared_path=shared_path, + hidden_states_path=tmp_path / "index", + ) + + with pytest.warns(UserWarning, match="Failed to load/cache"): + assert dataset._maybe_generate_hs(0) is None + + assert artifact_path.exists() + assert lock_path.exists() + assert list((shared_path / "artifacts").glob("*/*.safetensors")) == [] + assert dataset.artifact_cache is not None + stats = dataset.artifact_cache.snapshot_stats() + assert stats["logical_requests"] == 1 + assert stats["misses"] == 1 + assert stats["generation_failures"] == 1 + assert stats["publishes"] == 0 + + def test_shared_dataset_preserves_on_generate_index_cache(tmp_path, monkeypatch): data_path = tmp_path / "data" service_path = tmp_path / "service" From 0332621a3859a8591903fb566301f3f2275afd22 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Fri, 17 Jul 2026 22:03:57 +0800 Subject: [PATCH 04/20] fix: harden asynchronous shared artifact fanout Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- .../independent_consumer_fanout/README.md | 29 +- .../config.example.json | 26 +- docs/cli/train.md | 6 +- pyproject.toml | 1 + scripts/benchmark_independent_consumers.py | 7 +- scripts/train.py | 40 +- src/speculators/benchmarks/gpu_monitor.py | 474 ++++++++++++++++++ .../benchmarks/independent_consumers.py | 397 +++++++++------ .../data_generation/artifact_cache.py | 29 ++ src/speculators/train/data.py | 5 + tests/unit/benchmarks/test_gpu_monitor.py | 129 +++++ .../benchmarks/test_independent_consumers.py | 46 ++ .../data_generation/test_artifact_cache.py | 16 + tests/unit/train/test_cli_args.py | 23 + tests/unit/train/test_shared_artifacts.py | 14 + 15 files changed, 1089 insertions(+), 153 deletions(-) create mode 100644 src/speculators/benchmarks/gpu_monitor.py create mode 100644 tests/unit/benchmarks/test_gpu_monitor.py diff --git a/benchmarks/independent_consumer_fanout/README.md b/benchmarks/independent_consumer_fanout/README.md index 3f13c5e92..a25a2b3de 100644 --- a/benchmarks/independent_consumer_fanout/README.md +++ b/benchmarks/independent_consumer_fanout/README.md @@ -20,6 +20,14 @@ only a digest of the request identity, never the returned artifact path. Start from `config.example.json`, replace the model and preprocessed-data placeholders, and keep each consumer command as a direct, single-process `scripts/train.py` launch. +Pin training semantics such as `--optimizer` explicitly so an upstream default change +cannot silently alter a comparison. +Keep multiple DataLoader workers and explicit prefetching for generated data. A single +worker can hold only one unique first-miss request in flight, which serializes producer +prefill and can make otherwise independent consumers wait in lockstep. The example uses +four CPU workers with a prefetch factor of two; these workers do not create additional +GPU compute roles, and the benchmark still rejects more than one compute process on any +assigned GPU. Run the fixture from the repository root: ```bash @@ -29,10 +37,18 @@ python scripts/benchmark_independent_consumers.py \ --report /tmp/speculators-fanout-report.json ``` +Pass `--scenario 1p3c` to run only the configured 1P3C scenario. Omitting it preserves +the default serial 1P1C-then-1P3C comparison. The report records the scenarios that +were actually selected. + The run directory must not already exist. Role logs remain there and the compact report contains the exact command/configuration, package versions, request and valid-completion counts, shared-sample multiplicity, a common post-warmup throughput window, native -per-consumer `profile/step_ms` summaries, makespan, and sampled GPU memory. Environment +per-consumer `profile/step_ms` summaries, makespan, and role-aware NVML utilization and +memory over the common consumer steady-state overlap. Set +`measurement_steps_per_consumer` to use an exact number of post-warmup steps; otherwise +all available post-warmup steps are measured. Startup samples remain in the JSONL stream +but are excluded from steady-state aggregates. Environment values are omitted from the report. The command exits nonzero if a role fails, a GPU is shared or already occupied, a completion is malformed, sample multiplicity is ambiguous, or the common steady-state window is too small. @@ -51,6 +67,17 @@ to every service completion, exactly one miss is published, the other two reques and all failure, retry, cleanup, and timeout counters are zero. Baseline scenarios that do not use the shared-cache placeholder remain valid without cache accounting. +The shared cache is a filesystem data plane, not Mooncake or GPU-direct transport. Its +directory must provide reliable POSIX `flock`, same-filesystem atomic rename, and +directory `fsync` semantics to all consumers. Do not use an arbitrary NFS mount unless +those guarantees have been verified. +It is also not a consumer-centered bounded sliding window. Disabling expiration retains +one artifact for every unique request. With a finite TTL, expired entries are reclaimed +when a dataset opens the cache or when the same key is requested again; a single pass +over previously unseen samples can therefore continue growing on-disk usage. Size the +filesystem and choose the TTL for the maximum expected consumer lag. A throughput run +with a finite dataset is not evidence that long-running cache storage is bounded. + In publish-once mode, `per_consumer_completions` and the steady-state per-consumer completion map identify which consumer owned each service miss. They do not represent logical trainer progress: cache logical-request totals and each independent consumer's diff --git a/benchmarks/independent_consumer_fanout/config.example.json b/benchmarks/independent_consumer_fanout/config.example.json index 163e79459..4bbc12ccd 100644 --- a/benchmarks/independent_consumer_fanout/config.example.json +++ b/benchmarks/independent_consumer_fanout/config.example.json @@ -42,6 +42,8 @@ "1", "--total-seq-len", "3072", + "--optimizer", + "adamw", "--speculator-type", "dflash", "--draft-arch", @@ -63,7 +65,9 @@ "--on-generate", "delete", "--num-workers", - "1" + "4", + "--prefetch-factor", + "2" ] } ], @@ -71,6 +75,7 @@ "minimum_steady_completions_per_consumer": 50, "warmup_steps_per_consumer": 10, "minimum_steady_steps_per_consumer": 50, + "measurement_steps_per_consumer": 50, "minimum_shared_samples": 50, "expected_service_completions_per_shared_sample": 1 }, @@ -95,6 +100,8 @@ "1", "--total-seq-len", "3072", + "--optimizer", + "adamw", "--speculator-type", "dflash", "--draft-arch", @@ -116,7 +123,9 @@ "--on-generate", "delete", "--num-workers", - "1" + "4", + "--prefetch-factor", + "2" ] }, { @@ -137,6 +146,8 @@ "1", "--total-seq-len", "3072", + "--optimizer", + "adamw", "--speculator-type", "dflash", "--draft-arch", @@ -158,7 +169,9 @@ "--on-generate", "delete", "--num-workers", - "1" + "4", + "--prefetch-factor", + "2" ] }, { @@ -179,6 +192,8 @@ "1", "--total-seq-len", "3072", + "--optimizer", + "adamw", "--speculator-type", "dflash", "--draft-arch", @@ -200,7 +215,9 @@ "--on-generate", "delete", "--num-workers", - "1" + "4", + "--prefetch-factor", + "2" ] } ], @@ -208,6 +225,7 @@ "minimum_steady_completions_per_consumer": 50, "warmup_steps_per_consumer": 10, "minimum_steady_steps_per_consumer": 50, + "measurement_steps_per_consumer": 50, "minimum_shared_samples": 50, "expected_service_completions_per_shared_sample": 3 } diff --git a/docs/cli/train.md b/docs/cli/train.md index d9167808c..d00e1e691 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -86,12 +86,16 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--shared-hidden-states-path`** (str, default: `None`) Optional filesystem cache that coalesces identical online hidden-state requests across independent trainers. All participating trainers must use the same path. This does not change `--hidden-states-path`, which remains the per-dataset indexed cache used by `--on-generate cache`. -- **`--shared-hidden-states-namespace`** (str, default: `None`) Optional request-identity namespace for the producer's extraction configuration. Use the same value for trainers sharing artifacts and a different value when target layer IDs or other extraction semantics differ. +- **`--shared-hidden-states-namespace`** (str, default: `None`) Optional additional request-identity namespace for producer extraction settings not represented by the model or target layer IDs. Target layer IDs are fingerprinted automatically. Use the same value for trainers sharing artifacts and a different value for any other extraction semantic that changes the resulting tensors. - **`--shared-hidden-states-ttl`** (float, default: `3600.0`) Seconds to retain shared artifacts before regenerating them. Set to `0` to disable expiration. - **`--shared-hidden-states-lock-timeout`** (float, default: `300.0`) Maximum seconds to wait while another trainer generates and atomically publishes the same artifact. + The shared cache is a filesystem data plane, not Mooncake or GPU-direct transport. Its directory must provide reliable POSIX `flock`, same-filesystem atomic rename, and directory `fsync` semantics to every trainer. Do not assume an arbitrary NFS mount is safe unless those guarantees have been verified. + + This cache is not a consumer-centered bounded sliding window. Setting `--shared-hidden-states-ttl=0` retains one artifact per unique request. With a finite TTL, expired entries are reclaimed when a dataset opens the cache or when the same key is requested again, so a single pass over new samples can still grow on-disk usage. Provision the filesystem and set the TTL according to the maximum expected lag between consumers. + - **`--legacy-data`** (flag) **DEPRECATED.** Use the old data format which stores hidden states alongside token_ids. - **`--total-seq-len`** (int, default: `8192`) Maximum total sequence length for training batches. Note: samples will be packed into batches with total combined sequence length `{total-seq-len}`. diff --git a/pyproject.toml b/pyproject.toml index 1156753c1..8fa781704 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,7 @@ dependencies = [ ] [project.optional-dependencies] +nvml = ["nvidia-ml-py>=12.0.0"] dev = [ # build "build>=1.5.0", diff --git a/scripts/benchmark_independent_consumers.py b/scripts/benchmark_independent_consumers.py index 45c10e759..c76f0063f 100644 --- a/scripts/benchmark_independent_consumers.py +++ b/scripts/benchmark_independent_consumers.py @@ -24,6 +24,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--validate-only", action="store_true", help="Validate config without launching" ) + parser.add_argument( + "--scenario", + choices=("1p1c", "1p3c"), + help="Run only this scenario; by default both run serially", + ) return parser.parse_args() @@ -32,7 +37,7 @@ def main() -> int: config = load_config(args.config) if args.validate_only: return 0 - report = run_benchmark(config, args.run_directory) + report = run_benchmark(config, args.run_directory, scenario_kind=args.scenario) write_report(report, args.report) return 0 if report["valid"] else 1 diff --git a/scripts/train.py b/scripts/train.py index 65f513ba2..3aeed3f7b 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -15,6 +15,9 @@ from transformers.models.qwen3.configuration_qwen3 import Qwen3Config from hs_connectors import HiddenStatesBackend +from speculators.data_generation.artifact_cache import ( + canonical_hidden_state_extraction_namespace, +) from speculators.data_generation.vllm_client import ( DEFAULT_MAX_RETRIES, DEFAULT_REQUEST_TIMEOUT, @@ -591,6 +594,12 @@ def main(args: argparse.Namespace): # noqa: C901 # Get target layer IDs from the model (resolved at model level) num_target_layers = len(draft_model.target_layer_ids) # type: ignore[arg-type] + shared_artifacts_namespace = None + if args.shared_hidden_states_path is not None: + shared_artifacts_namespace = canonical_hidden_state_extraction_namespace( + tuple(int(layer_id) for layer_id in draft_model.target_layer_ids), # type: ignore[union-attr] + user_namespace=args.shared_hidden_states_namespace, + ) if args.speculator_type == "mtp": args.num_speculative_steps = draft_model.config.num_speculative_steps @@ -644,7 +653,7 @@ def main(args: argparse.Namespace): # noqa: C901 request_timeout=args.request_timeout, max_retries=args.max_retries, shared_artifacts_path=args.shared_hidden_states_path, - shared_artifacts_namespace=args.shared_hidden_states_namespace, + shared_artifacts_namespace=shared_artifacts_namespace, shared_artifacts_ttl_seconds=( None if args.shared_hidden_states_ttl == 0 @@ -778,6 +787,31 @@ def validate_draft_init_args( ) +def _validate_shared_hidden_state_args( + parser: argparse.ArgumentParser, args: argparse.Namespace +) -> None: + if args.shared_hidden_states_path is not None: + if not args.shared_hidden_states_path.strip(): + parser.error("--shared-hidden-states-path must be non-empty") + if args.legacy_data: + parser.error( + "--shared-hidden-states-path is incompatible with --legacy-data" + ) + elif args.shared_hidden_states_namespace is not None: + parser.error( + "--shared-hidden-states-namespace requires --shared-hidden-states-path" + ) + if ( + args.shared_hidden_states_namespace is not None + and not args.shared_hidden_states_namespace.strip() + ): + parser.error("--shared-hidden-states-namespace must be non-empty") + if args.shared_hidden_states_ttl < 0: + parser.error("--shared-hidden-states-ttl must be non-negative") + if args.shared_hidden_states_lock_timeout <= 0: + parser.error("--shared-hidden-states-lock-timeout must be positive") + + def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--verifier-name-or-path", type=str, required=True) @@ -916,7 +950,8 @@ def parse_args(): type=str, default=None, help=( - "Optional identity namespace for the producer's extraction configuration. " + "Optional additional identity namespace for the producer's extraction " + "configuration. Target layer IDs are always included automatically. " "Trainers sharing artifacts must use the same value." ), ) @@ -1347,6 +1382,7 @@ def parse_args(): ) args = parser.parse_args() + _validate_shared_hidden_state_args(parser, args) is_eagle3 = args.speculator_type == "eagle3" if args.draft_arch is None: diff --git a/src/speculators/benchmarks/gpu_monitor.py b/src/speculators/benchmarks/gpu_monitor.py new file mode 100644 index 000000000..589385c65 --- /dev/null +++ b/src/speculators/benchmarks/gpu_monitor.py @@ -0,0 +1,474 @@ +"""Role-aware NVML telemetry for independent-consumer benchmarks.""" + +from __future__ import annotations + +import importlib +import json +import os +import statistics +import threading +import time +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Any, Protocol + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator, Mapping, Sequence + from pathlib import Path + from typing import TextIO + + +_MIN_COVERAGE_SAMPLES = 2 + + +class GpuMonitorError(RuntimeError): + """GPU telemetry could not satisfy its benchmark contract.""" + + +@dataclass(frozen=True) +class GpuRoleAssignment: + gpu: int + logical_role: str + process_role: str + + def __post_init__(self) -> None: + if self.gpu < 0: + raise ValueError("gpu index must be non-negative") + if not self.logical_role or not self.process_role: + raise ValueError("GPU role names must be non-empty") + + +@dataclass(frozen=True) +class GpuDevice: + gpu: int + uuid: str + name: str + + +class GpuTelemetryBackend(Protocol): + def open(self, gpu_indices: Sequence[int]) -> Mapping[int, GpuDevice]: ... + + def sample(self, gpu: int) -> Mapping[str, Any]: ... + + def close(self) -> None: ... + + +def _as_text(value: Any) -> str: + if isinstance(value, bytes): + return value.decode("utf-8", errors="replace") + return str(value) + + +def _error_text(error: BaseException) -> str: + message = str(error).replace("\n", " ").strip() + return f"{type(error).__name__}: {message}"[:512] + + +class NvmlBackend: + """Lazy wrapper that never initializes CUDA or Torch.""" + + def __init__(self, module: Any = None) -> None: + self._module = module + self._handles: dict[int, Any] = {} + self._opened = False + + def open(self, gpu_indices: Sequence[int]) -> Mapping[int, GpuDevice]: + if self._opened: + raise GpuMonitorError("NVML backend is already open") + if self._module is None: + try: + self._module = importlib.import_module("pynvml") + except ImportError as error: + raise GpuMonitorError( + "NVML monitoring requires nvidia-ml-py; install the nvml extra" + ) from error + nvml = self._module + try: + nvml.nvmlInit() + self._opened = True + devices: dict[int, GpuDevice] = {} + for gpu in gpu_indices: + handle = nvml.nvmlDeviceGetHandleByIndex(gpu) + self._handles[gpu] = handle + devices[gpu] = GpuDevice( + gpu=gpu, + uuid=_as_text(nvml.nvmlDeviceGetUUID(handle)), + name=_as_text(nvml.nvmlDeviceGetName(handle)), + ) + return devices + except Exception: + self.close() + raise + + def sample(self, gpu: int) -> Mapping[str, Any]: + if gpu not in self._handles: + raise GpuMonitorError(f"NVML GPU {gpu} is not open") + nvml = self._module + handle = self._handles[gpu] + utilization = nvml.nvmlDeviceGetUtilizationRates(handle) + memory = nvml.nvmlDeviceGetMemoryInfo(handle) + processes = [] + for process in nvml.nvmlDeviceGetComputeRunningProcesses(handle): + used_memory = getattr(process, "usedGpuMemory", None) + unavailable = getattr(nvml, "NVML_VALUE_NOT_AVAILABLE", None) + if used_memory == unavailable: + used_memory = None + processes.append( + {"pid": int(process.pid), "used_memory_bytes": used_memory} + ) + return { + "utilization_gpu_pct": int(utilization.gpu), + "utilization_memory_pct": int(utilization.memory), + "memory_used_bytes": int(memory.used), + "memory_free_bytes": int(memory.free), + "memory_total_bytes": int(memory.total), + "compute_processes": processes, + } + + def close(self) -> None: + if not self._opened: + return + self._handles.clear() + self._module.nvmlShutdown() + self._opened = False + + +def _atomic_write_json(path: Path, value: Mapping[str, Any]) -> None: + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + try: + with temporary.open("x", encoding="utf-8") as output: + json.dump(value, output, indent=2, sort_keys=True, allow_nan=False) + output.write("\n") + output.flush() + os.fsync(output.fileno()) + temporary.replace(path) + finally: + temporary.unlink(missing_ok=True) + + +class GpuMonitor: + """Stream NVML samples and retain a bounded health summary.""" + + def __init__( + self, + assignments: Sequence[GpuRoleAssignment], + sample_path: Path, + summary_path: Path, + *, + poll_seconds: float = 0.5, + max_compute_processes: int = 1, + backend: GpuTelemetryBackend | None = None, + ) -> None: + self.assignments = tuple(assignments) + if not self.assignments: + raise ValueError("at least one GPU role assignment is required") + if len({value.gpu for value in self.assignments}) != len(self.assignments): + raise ValueError("GPU role assignments must use distinct GPU indices") + if poll_seconds <= 0: + raise ValueError("poll_seconds must be positive") + if max_compute_processes < 1: + raise ValueError("max_compute_processes must be positive") + self.sample_path = sample_path + self.summary_path = summary_path + self.poll_seconds = poll_seconds + self.max_compute_processes = max_compute_processes + self.backend = backend or NvmlBackend() + self._devices: dict[int, GpuDevice] = {} + self._sample_counts = dict.fromkeys( + (value.gpu for value in self.assignments), 0 + ) + self._max_processes = dict.fromkeys( + (value.gpu for value in self.assignments), 0 + ) + self._errors: list[str] = [] + self._violations: list[str] = [] + self._collection_ms: list[float] = [] + self._poll_overruns = 0 + self._collection_lock = threading.RLock() + self._stop = threading.Event() + self._thread = threading.Thread( + target=self._run, name="speculators-nvml-monitor", daemon=True + ) + self._output: TextIO | None = None + self._started_at_ns: int | None = None + self._ended_at_ns: int | None = None + self._backend_open = False + self._started = False + self._stopped = False + self._summary: dict[str, Any] | None = None + + @property + def errors(self) -> list[str]: + with self._collection_lock: + return list(self._errors) + + @property + def violations(self) -> list[str]: + with self._collection_lock: + return list(self._violations) + + def _close_backend(self) -> None: + if not self._backend_open: + return + try: + self.backend.close() + except Exception as error: # noqa: BLE001 + with self._collection_lock: + self._errors.append(_error_text(error)) + finally: + self._backend_open = False + + def _write(self, value: Mapping[str, Any]) -> None: + output = self._output + if output is None: + raise GpuMonitorError("GPU monitor output is not open") + output.write(json.dumps(value, sort_keys=True, separators=(",", ":"))) + output.write("\n") + output.flush() + + def start(self) -> None: + if self._started: + raise GpuMonitorError("GPU monitor can only be started once") + self._started = True + self.sample_path.parent.mkdir(parents=True, exist_ok=True) + self._output = self.sample_path.open("x", encoding="utf-8", buffering=1) + try: + self._devices = dict( + self.backend.open([value.gpu for value in self.assignments]) + ) + self._backend_open = True + missing = {value.gpu for value in self.assignments} - set(self._devices) + if missing: + raise GpuMonitorError(f"NVML backend omitted GPUs {sorted(missing)}") + except Exception: + self._output.close() + self._output = None + raise + self._started_at_ns = time.monotonic_ns() + self._write( + { + "record_type": "session_start", + "timestamp_monotonic_ns": self._started_at_ns, + "poll_seconds": self.poll_seconds, + "assignments": [asdict(value) for value in self.assignments], + "devices": [asdict(value) for value in self._devices.values()], + } + ) + self._thread.start() + + def _run(self) -> None: # noqa: C901 + try: + while not self._stop.is_set(): + cycle_started = time.monotonic() + timestamp_ns = time.monotonic_ns() + timestamp_unix_ns = time.time_ns() + for assignment in self.assignments: + if self._stop.is_set(): + return + try: + metrics = dict(self.backend.sample(assignment.gpu)) + pids = sorted( + int(value["pid"]) + for value in metrics.get("compute_processes", ()) + ) + with self._collection_lock: + if self._stop.is_set(): + return + self._sample_counts[assignment.gpu] += 1 + self._max_processes[assignment.gpu] = max( + self._max_processes[assignment.gpu], len(pids) + ) + if len(pids) > self.max_compute_processes: + message = ( + f"GPU {assignment.gpu} has {len(pids)} compute " + f"processes: {pids}" + ) + if message not in self._violations: + self._violations.append(message) + self._write( + { + "record_type": "sample", + "timestamp_monotonic_ns": timestamp_ns, + "timestamp_unix_ns": timestamp_unix_ns, + **asdict(assignment), + **metrics, + "compute_pids": pids, + } + ) + except Exception as error: # noqa: BLE001 + with self._collection_lock: + if self._stop.is_set(): + return + message = _error_text(error) + if message not in self._errors: + self._errors.append(message) + elapsed_ms = (time.monotonic() - cycle_started) * 1000.0 + with self._collection_lock: + if self._stop.is_set(): + break + self._collection_ms.append(elapsed_ms) + if elapsed_ms > self.poll_seconds * 1000.0: + self._poll_overruns += 1 + self._stop.wait(max(0.0, self.poll_seconds - elapsed_ms / 1000.0)) + finally: + self._close_backend() + + def stop(self) -> dict[str, Any]: + if not self._started: + raise GpuMonitorError("GPU monitor was not started") + if self._stopped: + if self._summary is None: + raise GpuMonitorError("GPU monitor summary is unavailable") + return self._summary + self._stop.set() + self._thread.join(timeout=max(10.0, self.poll_seconds * 4)) + with self._collection_lock: + if self._thread.is_alive(): + self._errors.append("GPU monitor thread did not stop") + self._ended_at_ns = time.monotonic_ns() + if self._output is not None: + self._write( + { + "record_type": "session_end", + "timestamp_monotonic_ns": self._ended_at_ns, + } + ) + self._output.close() + self._output = None + summary = { + "status": ( + "ok" if not self._errors and not self._violations else "degraded" + ), + "sample_path": str(self.sample_path.resolve()), + "duration_seconds": ( + (self._ended_at_ns - self._started_at_ns) / 1e9 + if self._started_at_ns is not None + else None + ), + "poll_seconds": self.poll_seconds, + "poll_overrun_count": self._poll_overruns, + "collection_latency_ms": { + "mean": statistics.fmean(self._collection_ms) + if self._collection_ms + else None, + "max": max(self._collection_ms) if self._collection_ms else None, + }, + "sample_count_by_gpu": { + str(key): value for key, value in self._sample_counts.items() + }, + "max_compute_processes_by_gpu": { + str(key): value for key, value in self._max_processes.items() + }, + "errors": list(self._errors), + "ownership_violations": list(self._violations), + } + self._summary = summary + self._stopped = True + _atomic_write_json(self.summary_path, summary) + return summary + + +def iter_gpu_samples(path: Path) -> Iterator[dict[str, Any]]: + """Yield complete samples while tolerating one torn final JSONL line.""" + + with path.open(encoding="utf-8") as input_file: + for line_number, line in enumerate(input_file, start=1): + if not line.strip(): + continue + try: + value = json.loads(line) + except json.JSONDecodeError as error: + if not line.endswith("\n"): + return + raise GpuMonitorError( + f"invalid GPU sample JSONL at line {line_number}: {error}" + ) from error + if value.get("record_type") == "sample": + yield value + + +def _percentile(values: Sequence[float], fraction: float) -> float | None: + if not values: + return None + ordered = sorted(values) + index = min(len(ordered) - 1, max(0, int(len(ordered) * fraction + 0.999999) - 1)) + return float(ordered[index]) + + +def summarize_gpu_window( + samples: Iterable[Mapping[str, Any]], + assignments: Sequence[GpuRoleAssignment], + *, + start_monotonic_ns: int, + end_monotonic_ns: int, + low_utilization_pct: int = 10, +) -> dict[str, Any]: + """Aggregate NVML active-time samples in one common consumer interval.""" + + if end_monotonic_ns <= start_monotonic_ns: + raise ValueError("GPU measurement window must have positive duration") + rows: dict[int, list[Mapping[str, Any]]] = {value.gpu: [] for value in assignments} + for sample in samples: + timestamp = sample.get("timestamp_monotonic_ns") + gpu = sample.get("gpu") + if ( + isinstance(timestamp, int) + and not isinstance(timestamp, bool) + and start_monotonic_ns <= timestamp <= end_monotonic_ns + and gpu in rows + ): + rows[gpu].append(sample) + invalid_reasons = [] + per_gpu = {} + by_gpu = {value.gpu: value for value in assignments} + for gpu in sorted(rows): + gpu_rows = rows[gpu] + utilization = [ + float(value["utilization_gpu_pct"]) + for value in gpu_rows + if value.get("utilization_gpu_pct") is not None + ] + timestamps = [int(value["timestamp_monotonic_ns"]) for value in gpu_rows] + process_counts = [len(value.get("compute_pids", ())) for value in gpu_rows] + memory = [ + int(value["memory_used_bytes"]) + for value in gpu_rows + if value.get("memory_used_bytes") is not None + ] + if not gpu_rows: + invalid_reasons.append(f"GPU {gpu} has no samples in the steady window") + if gpu_rows and not any(process_counts): + invalid_reasons.append( + f"GPU {gpu} has no compute process in the steady window" + ) + per_gpu[str(gpu)] = { + **asdict(by_gpu[gpu]), + "sample_count": len(gpu_rows), + "coverage_seconds": ( + (max(timestamps) - min(timestamps)) / 1e9 + if len(timestamps) >= _MIN_COVERAGE_SAMPLES + else 0.0 + ), + "gpu_utilization_pct": { + "mean": statistics.fmean(utilization) if utilization else None, + "p50": _percentile(utilization, 0.50), + "p95": _percentile(utilization, 0.95), + "low_fraction": ( + sum(value < low_utilization_pct for value in utilization) + / len(utilization) + if utilization + else None + ), + }, + "max_memory_used_mib": max(memory) / (1 << 20) if memory else None, + "max_compute_processes": max(process_counts, default=0), + } + return { + "interval": "common_consumer_steady_overlap", + "start_monotonic_ns": start_monotonic_ns, + "end_monotonic_ns": end_monotonic_ns, + "duration_seconds": (end_monotonic_ns - start_monotonic_ns) / 1e9, + "low_utilization_threshold_pct": low_utilization_pct, + "valid": not invalid_reasons, + "invalid_reasons": invalid_reasons, + "per_gpu": per_gpu, + } diff --git a/src/speculators/benchmarks/independent_consumers.py b/src/speculators/benchmarks/independent_consumers.py index 1730ff44b..aeb6477a3 100644 --- a/src/speculators/benchmarks/independent_consumers.py +++ b/src/speculators/benchmarks/independent_consumers.py @@ -17,14 +17,24 @@ from dataclasses import asdict, dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal from urllib.parse import urlsplit -import psutil from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from speculators.benchmarks.gpu_monitor import ( + GpuMonitor as NvmlGpuMonitor, +) +from speculators.benchmarks.gpu_monitor import ( + GpuRoleAssignment, + iter_gpu_samples, + summarize_gpu_window, +) from speculators.data_generation.artifact_cache import HiddenStateArtifactCache +if TYPE_CHECKING: + from collections.abc import Callable + class EvidenceError(RuntimeError): """Raised when a benchmark cannot produce trustworthy evidence.""" @@ -161,6 +171,7 @@ class ScenarioSpec(_StrictModel): minimum_steady_completions_per_consumer: int = Field(default=5, ge=1) warmup_steps_per_consumer: int = Field(default=10, ge=0) minimum_steady_steps_per_consumer: int = Field(default=10, ge=1) + measurement_steps_per_consumer: int | None = Field(default=None, ge=1) minimum_shared_samples: int = Field(default=5, ge=1) expected_service_completions_per_shared_sample: int = Field(ge=1) @@ -194,7 +205,7 @@ class BenchmarkConfig(_StrictModel): scenarios: list[ScenarioSpec] = Field(min_length=2, max_length=2) allowed_gpus: list[int] = Field(min_length=2) proxy_timeout_seconds: float = Field(default=180.0, gt=0) - memory_sample_interval_seconds: float = Field(default=0.25, gt=0) + memory_sample_interval_seconds: float = Field(default=0.5, gt=0) @field_validator("allowed_gpus") @classmethod @@ -731,13 +742,24 @@ def analyze_cache_accounting( ) +@dataclass(frozen=True) +class ConsumerStepEvent: + completed_at_monotonic_ns: int + step_ms: float + + def _percentile(values: list[float], percentile: float) -> float: index = max(0, min(len(values) - 1, int(len(values) * percentile + 0.999999) - 1)) return sorted(values)[index] def analyze_consumer_steps( - log_path: Path, warmup_steps: int, minimum_steady_steps: int + log_path: Path, + warmup_steps: int, + minimum_steady_steps: int, + *, + events: list[ConsumerStepEvent] | None = None, + measurement_steps: int | None = None, ) -> dict[str, Any]: """Extract native trainer step timings and separate warmup from steady state.""" if not log_path.is_file(): @@ -750,20 +772,58 @@ def analyze_consumer_steps( "step_ms_mean": None, "step_ms_p50": None, "step_ms_p95": None, + "steady_started_at_monotonic_ns": None, + "steady_finished_at_monotonic_ns": None, + "steady_duration_seconds": None, + "steady_steps_per_second": None, } - values = [] - with log_path.open(errors="replace") as log_file: - for line in log_file: - values.extend( - float(match.group("value")) - for match in _STEP_TIME_PATTERN.finditer(line) - ) - steady = values[warmup_steps:] + if events is None: + values = [] + with log_path.open(errors="replace") as log_file: + for line in log_file: + values.extend( + float(match.group("value")) + for match in _STEP_TIME_PATTERN.finditer(line) + ) + ordered_events: list[ConsumerStepEvent] = [] + else: + ordered_events = sorted( + events, key=lambda value: value.completed_at_monotonic_ns + ) + values = [value.step_ms for value in ordered_events] + available_steady = values[warmup_steps:] + steady = ( + available_steady[:measurement_steps] + if measurement_steps is not None + else available_steady + ) reasons = [] - if len(steady) < minimum_steady_steps: + if len(available_steady) < minimum_steady_steps: reasons.append( - f"only {len(steady)} steady step timing(s); need {minimum_steady_steps}" + f"only {len(available_steady)} steady step timing(s); " + f"need {minimum_steady_steps}" ) + if measurement_steps is not None and len(steady) < measurement_steps: + reasons.append( + f"only {len(steady)} measured step timing(s); need {measurement_steps}" + ) + started_at = None + finished_at = None + duration = None + if ordered_events and steady: + measured_end_index = warmup_steps + len(steady) - 1 + finished_at = ordered_events[measured_end_index].completed_at_monotonic_ns + if warmup_steps: + started_at = ordered_events[warmup_steps - 1].completed_at_monotonic_ns + else: + started_at = max( + 0, + ordered_events[0].completed_at_monotonic_ns + - int(ordered_events[0].step_ms * 1_000_000), + ) + duration = (finished_at - started_at) / 1e9 + if duration <= 0: + reasons.append("consumer steady-state duration is not positive") return { "valid": not reasons, "invalid_reasons": reasons, @@ -773,6 +833,12 @@ def analyze_consumer_steps( "step_ms_mean": sum(steady) / len(steady) if steady else None, "step_ms_p50": _percentile(steady, 0.50) if steady else None, "step_ms_p95": _percentile(steady, 0.95) if steady else None, + "steady_started_at_monotonic_ns": started_at, + "steady_finished_at_monotonic_ns": finished_at, + "steady_duration_seconds": duration, + "steady_steps_per_second": ( + len(steady) / duration if duration is not None and duration > 0 else None + ), } @@ -835,119 +901,6 @@ def _gpu_snapshot(target_gpus: set[int]) -> _GpuSample: ) -class GpuMonitor: - def __init__(self, target_gpus: set[int], interval_seconds: float) -> None: - self._target_gpus = target_gpus - self._interval = interval_seconds - self._roots: dict[int, int] = {} - self._known_pids: dict[int, set[int]] = defaultdict(set) - self._samples: list[_GpuSample] = [] - self._errors: list[str] = [] - self._lock = threading.Lock() - self._stop = threading.Event() - self._started = False - self._thread = threading.Thread( - target=self._run, name="benchmark-gpu-monitor", daemon=True - ) - - def require_idle(self) -> dict[int, int]: - sample = _gpu_snapshot(self._target_gpus) - occupied = {gpu: pids for gpu, pids in sample.compute_pids.items() if pids} - if occupied: - raise EvidenceError( - f"Target GPUs already have compute processes: {occupied}" - ) - return sample.total_memory_mib - - def set_role_process(self, gpu: int, pid: int) -> None: - with self._lock: - self._roots[gpu] = pid - self._known_pids[gpu].add(pid) - - def start(self) -> None: - self._started = True - self._thread.start() - - def stop(self) -> None: - if not self._started: - return - self._stop.set() - self._thread.join(timeout=max(5.0, self._interval * 4)) - - def _allowed_pids(self, gpu: int) -> set[int]: - with self._lock: - root = self._roots.get(gpu) - known = set(self._known_pids[gpu]) - if root is None: - return known - try: - descendants = {child.pid for child in psutil.Process(root).children(True)} - except (psutil.NoSuchProcess, psutil.AccessDenied): - descendants = set() - allowed = known | {root} | descendants - with self._lock: - self._known_pids[gpu].update(allowed) - return allowed - - def _run(self) -> None: - while not self._stop.wait(self._interval): - try: - sample = _gpu_snapshot(self._target_gpus) - for gpu, pids in sample.compute_pids.items(): - if len(pids) > 1: - self._errors.append( - f"GPU {gpu} has {len(pids)} CUDA compute processes" - ) - foreign = set(pids) - self._allowed_pids(gpu) - if foreign: - self._errors.append( - f"GPU {gpu} has foreign compute PIDs {sorted(foreign)}" - ) - self._samples.append(sample) - except Exception as error: # noqa: BLE001 - self._errors.append(f"{type(error).__name__}: {error}") - - def summarize( - self, started_at: float, finished_at: float, baseline: dict[int, int] - ) -> dict[str, Any]: - samples = [ - sample - for sample in self._samples - if started_at <= sample.captured_at <= finished_at - ] - per_gpu: dict[str, Any] = {} - for gpu in sorted(self._target_gpus): - role_values = [sample.role_memory_mib[gpu] for sample in samples] - total_values = [sample.total_memory_mib[gpu] for sample in samples] - observed = any(sample.compute_pids.get(gpu) for sample in samples) - per_gpu[str(gpu)] = { - "baseline_memory_mib": baseline.get(gpu), - "peak_role_memory_mib": max(role_values) if role_values else None, - "peak_total_memory_mib": max(total_values) if total_values else None, - "max_compute_processes": max( - (len(sample.compute_pids.get(gpu, [])) for sample in samples), - default=0, - ), - "compute_process_observed": observed, - } - errors = list(dict.fromkeys(self._errors)) - if not samples: - errors.append("No memory samples fall inside the steady-state window") - missing = [ - gpu - for gpu, value in per_gpu.items() - if not value["compute_process_observed"] - ] - if missing: - errors.append(f"No compute process observed on GPU(s) {missing}") - return { - "reliable": not errors, - "invalid_reasons": errors, - "sample_count": len(samples), - "per_gpu": per_gpu, - } - - def _render(values: list[str], replacements: dict[str, str]) -> list[str]: rendered: list[str] = [] for original in values: @@ -965,6 +918,7 @@ def __init__( env: dict[str, str], gpu: int, log_path: Path, + line_callback: Callable[[int, str], None] | None = None, ) -> None: process_env = os.environ.copy() for name in _DISTRIBUTED_ENV: @@ -972,19 +926,56 @@ def __init__( process_env.update(env) process_env[_GPU_ENV] = str(gpu) self._log = log_path.open("w", encoding="utf-8") + self._line_callback = line_callback + self._closed = False + self._reader_error: str | None = None self.started_at = time.monotonic() self.finished_at: float | None = None self.process = subprocess.Popen( # noqa: S603 command, env=process_env, - stdout=self._log, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, start_new_session=True, text=True, + bufsize=1, ) + self._reader = threading.Thread( + target=self._pump_output, + name=f"benchmark-log-{self.process.pid}", + daemon=True, + ) + self._reader.start() + + def _pump_output(self) -> None: + stdout = self.process.stdout + if stdout is None: + self._reader_error = "subprocess stdout pipe is unavailable" + return + try: + for line in stdout: + self._log.write(line) + self._log.flush() + if self._line_callback is not None: + self._line_callback(time.monotonic_ns(), line) + except Exception as error: # noqa: BLE001 + self._reader_error = f"{type(error).__name__}: {error}" + finally: + stdout.close() + + @property + def reader_error(self) -> str | None: + return self._reader_error def close_log(self) -> None: + if self._closed: + return + self._reader.join(timeout=5) + if self._reader.is_alive(): + self._reader_error = self._reader_error or "log reader did not stop" + self._log.flush() self._log.close() + self._closed = True def terminate(self, grace_seconds: float = 20.0) -> None: if self.process.poll() is not None: @@ -1031,9 +1022,33 @@ def _run_scenario( # noqa: C901 config.producer.gpu, *(consumer.gpu for consumer in scenario.consumers), } - monitor = GpuMonitor(target_gpus, config.memory_sample_interval_seconds) + assignments = [ + GpuRoleAssignment(config.producer.gpu, "producer", "producer"), + *( + GpuRoleAssignment( + consumer.gpu, + f"consumer:{consumer.consumer_id}", + f"consumer:{consumer.consumer_id}", + ) + for consumer in scenario.consumers + ), + ] + monitor = NvmlGpuMonitor( + assignments, + scenario_dir / "gpu_samples.jsonl", + scenario_dir / "gpu_summary.json", + poll_seconds=config.memory_sample_interval_seconds, + ) + monitor_started = False + monitor_summary: dict[str, Any] = { + "status": "not-started", + "errors": [], + "ownership_violations": [], + } baseline: dict[int, int] = {} ledger = AccountingLedger() + step_events: dict[str, list[ConsumerStepEvent]] = defaultdict(list) + step_events_lock = threading.Lock() producer: _ManagedProcess | None = None consumers: list[_ManagedProcess] = [] proxies: list[AccountingProxy] = [] @@ -1053,7 +1068,13 @@ def _run_scenario( # noqa: C901 } Path(producer_replacements["output_dir"]).mkdir() try: - baseline = monitor.require_idle() + preflight = _gpu_snapshot(target_gpus) + occupied = {gpu: pids for gpu, pids in preflight.compute_pids.items() if pids} + if occupied: + raise EvidenceError( + f"Target GPUs already have compute processes: {occupied}" + ) + baseline = preflight.total_memory_mib producer = _ManagedProcess( _render(config.producer.command, producer_replacements), { @@ -1063,7 +1084,6 @@ def _run_scenario( # noqa: C901 config.producer.gpu, scenario_dir / "producer.log", ) - monitor.set_role_process(config.producer.gpu, producer.process.pid) _wait_for_producer( config.producer.endpoint, producer.process, @@ -1081,6 +1101,7 @@ def _run_scenario( # noqa: C901 proxies.append(proxy) monitor.start() + monitor_started = True started_at = time.monotonic() for consumer, proxy in zip(scenario.consumers, proxies, strict=True): consumer_dir = scenario_dir / consumer.consumer_id @@ -1092,6 +1113,21 @@ def _run_scenario( # noqa: C901 "scenario": scenario.kind, "shared_artifacts_dir": str(shared_artifacts_dir), } + + def capture_step( + timestamp_ns: int, + line: str, + consumer_id: str = consumer.consumer_id, + ) -> None: + matches = list(_STEP_TIME_PATTERN.finditer(line)) + if not matches: + return + with step_events_lock: + step_events[consumer_id].extend( + ConsumerStepEvent(timestamp_ns, float(match.group("value"))) + for match in matches + ) + process = _ManagedProcess( _render(consumer.command, replacements), { @@ -1103,9 +1139,9 @@ def _run_scenario( # noqa: C901 }, consumer.gpu, scenario_dir / f"{consumer.consumer_id}.log", + capture_step, ) consumers.append(process) - monitor.set_role_process(consumer.gpu, process.process.pid) deadline = started_at + scenario.timeout_seconds for consumer, process in zip(scenario.consumers, consumers, strict=True): @@ -1133,13 +1169,27 @@ def _run_scenario( # noqa: C901 runtime_errors.append(f"{type(error).__name__}: {error}") finished_at = time.monotonic() finally: - monitor.stop() + if monitor_started: + try: + monitor_summary = monitor.stop() + except Exception as error: # noqa: BLE001 + runtime_errors.append( + f"GPU monitor shutdown failed: {type(error).__name__}: {error}" + ) for process in consumers: process.terminate() + if process.reader_error is not None: + runtime_errors.append( + f"consumer log reader failed: {process.reader_error}" + ) for proxy in proxies: proxy.close() if producer is not None: producer.terminate() + if producer.reader_error is not None: + runtime_errors.append( + f"producer log reader failed: {producer.reader_error}" + ) analysis = analyze_scenario(scenario, ledger.snapshot(), started_at, finished_at) cache_stats = None @@ -1163,23 +1213,70 @@ def _run_scenario( # noqa: C901 scenario_dir / f"{consumer.consumer_id}.log", scenario.warmup_steps_per_consumer, scenario.minimum_steady_steps_per_consumer, + events=step_events[consumer.consumer_id], + measurement_steps=scenario.measurement_steps_per_consumer, ) consumer_steps[consumer.consumer_id] = step_result runtime_errors.extend( f"{consumer.consumer_id}: {reason}" for reason in step_result["invalid_reasons"] ) - steady = analysis["steady_state"] - memory = monitor.summarize( - steady["started_at_monotonic"], - steady["finished_at_monotonic"], - baseline, - ) + consumer_starts = [ + value["steady_started_at_monotonic_ns"] + for value in consumer_steps.values() + if value["steady_started_at_monotonic_ns"] is not None + ] + consumer_ends = [ + value["steady_finished_at_monotonic_ns"] + for value in consumer_steps.values() + if value["steady_finished_at_monotonic_ns"] is not None + ] + gpu_window: dict[str, Any] = { + "valid": False, + "invalid_reasons": ["common consumer steady-state window is unavailable"], + "per_gpu": {}, + } + if len(consumer_starts) == len(scenario.consumers) and len(consumer_ends) == len( + scenario.consumers + ): + overlap_start = max(consumer_starts) + overlap_end = min(consumer_ends) + try: + gpu_window = summarize_gpu_window( + iter_gpu_samples(monitor.sample_path), + assignments, + start_monotonic_ns=overlap_start, + end_monotonic_ns=overlap_end, + ) + except Exception as error: # noqa: BLE001 + gpu_window = { + "valid": False, + "invalid_reasons": [ + f"GPU steady-window summary failed: {type(error).__name__}: {error}" + ], + "per_gpu": {}, + } + memory = { + "reliable": gpu_window["valid"], + "invalid_reasons": list(gpu_window["invalid_reasons"]), + "per_gpu": { + gpu: { + "baseline_memory_mib": baseline.get(int(gpu)), + "peak_total_memory_mib": value["max_memory_used_mib"], + "peak_role_memory_mib": value["max_memory_used_mib"], + "max_compute_processes": value["max_compute_processes"], + "compute_process_observed": value["max_compute_processes"] > 0, + } + for gpu, value in gpu_window["per_gpu"].items() + }, + } invalid_reasons = [ *runtime_errors, *analysis["invalid_reasons"], *cache_accounting["invalid_reasons"], *memory["invalid_reasons"], + *monitor_summary.get("errors", []), + *monitor_summary.get("ownership_violations", []), ] return { "kind": scenario.kind, @@ -1200,10 +1297,12 @@ def _run_scenario( # noqa: C901 ], "request_accounting": analysis["request_accounting"], "shared_artifact_cache": cache_accounting, - "steady_state": steady, + "steady_state": analysis["steady_state"], "consumer_step_times": consumer_steps, "makespan_seconds": max(finished_at - started_at, 0.0), "memory": memory, + "gpu_monitor": monitor_summary, + "gpu_steady_window": gpu_window, } @@ -1216,12 +1315,21 @@ def _redacted_config(config: BenchmarkConfig) -> dict[str, Any]: return value -def run_benchmark(config: BenchmarkConfig, output_dir: Path) -> dict[str, Any]: - """Run fresh-producer 1P1C and 1P3C scenarios and return one JSON report.""" +def run_benchmark( + config: BenchmarkConfig, + output_dir: Path, + scenario_kind: Literal["1p1c", "1p3c"] | None = None, +) -> dict[str, Any]: + """Run selected fresh-producer scenarios and return one JSON report.""" output_dir.mkdir(parents=True, exist_ok=False) - scenarios = [ - _run_scenario(config, scenario, output_dir) for scenario in config.scenarios + selected = [ + scenario + for scenario in config.scenarios + if scenario_kind is None or scenario.kind == scenario_kind ] + if not selected: + raise ValueError(f"Configured scenarios do not include {scenario_kind!r}") + scenarios = [_run_scenario(config, scenario, output_dir) for scenario in selected] versions = {} for package in ("speculators", "torch", "vllm"): try: @@ -1232,6 +1340,7 @@ def run_benchmark(config: BenchmarkConfig, output_dir: Path) -> dict[str, Any]: "schema_version": 1, "valid": all(scenario["valid"] for scenario in scenarios), "config": _redacted_config(config), + "selected_scenarios": [scenario.kind for scenario in selected], "versions": versions, "scenarios": scenarios, } diff --git a/src/speculators/data_generation/artifact_cache.py b/src/speculators/data_generation/artifact_cache.py index f02a3d747..324cc7ff2 100644 --- a/src/speculators/data_generation/artifact_cache.py +++ b/src/speculators/data_generation/artifact_cache.py @@ -38,6 +38,35 @@ ) +def canonical_hidden_state_extraction_namespace( + target_layer_ids: list[int] | tuple[int, ...], + *, + user_namespace: str | None = None, +) -> str: + """Fingerprint the producer configuration that changes artifact semantics.""" + + if ( + not target_layer_ids + or not all( + isinstance(layer_id, int) and not isinstance(layer_id, bool) + for layer_id in target_layer_ids + ) + or len(set(target_layer_ids)) != len(target_layer_ids) + ): + raise ValueError("target_layer_ids must be non-empty unique integers") + if user_namespace is not None and not user_namespace: + raise ValueError("user_namespace must be non-empty when provided") + identity = { + "schema_version": 1, + "target_layer_ids": list(target_layer_ids), + } + if user_namespace is not None: + identity["user_namespace"] = user_namespace + return json.dumps( + identity, ensure_ascii=True, separators=(",", ":"), sort_keys=True + ) + + class ArtifactCacheError(RuntimeError): """Base error raised by shared hidden-state artifact caching.""" diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py index 1ed7612f6..2c61db537 100644 --- a/src/speculators/train/data.py +++ b/src/speculators/train/data.py @@ -274,6 +274,11 @@ def __init__( self.request_timeout = request_timeout self.max_retries = max_retries self.shared_artifacts_namespace = shared_artifacts_namespace + if shared_artifacts_path is not None and not shared_artifacts_namespace: + raise ValueError( + "shared_artifacts_namespace is required when shared artifacts " + "are enabled" + ) self.artifact_cache = ( HiddenStateArtifactCache( shared_artifacts_path, diff --git a/tests/unit/benchmarks/test_gpu_monitor.py b/tests/unit/benchmarks/test_gpu_monitor.py new file mode 100644 index 000000000..aa1973b87 --- /dev/null +++ b/tests/unit/benchmarks/test_gpu_monitor.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import json +import threading + +from speculators.benchmarks.gpu_monitor import ( + GpuDevice, + GpuMonitor, + GpuRoleAssignment, + iter_gpu_samples, + summarize_gpu_window, +) + + +class _FakeBackend: + def __init__(self) -> None: + self.sampled = threading.Event() + self.closed = False + + def open(self, gpu_indices): + return { + gpu: GpuDevice(gpu=gpu, uuid=f"GPU-{gpu}", name="Fake GPU") + for gpu in gpu_indices + } + + def sample(self, gpu): + self.sampled.set() + return { + "utilization_gpu_pct": 75 + gpu, + "utilization_memory_pct": 40, + "memory_used_bytes": (8 + gpu) << 30, + "memory_free_bytes": 40 << 30, + "memory_total_bytes": 48 << 30, + "compute_processes": [{"pid": 100 + gpu, "used_memory_bytes": 4 << 30}], + } + + def close(self): + self.closed = True + + +class _BlockingBackend(_FakeBackend): + def __init__(self) -> None: + super().__init__() + self.release_sample = threading.Event() + + def sample(self, gpu): + self.sampled.set() + self.release_sample.wait(timeout=5) + return super().sample(gpu) + + +def test_monitor_streams_nvml_samples_and_finalizes_summary(tmp_path): + backend = _FakeBackend() + sample_path = tmp_path / "samples.jsonl" + summary_path = tmp_path / "summary.json" + monitor = GpuMonitor( + [GpuRoleAssignment(0, "producer", "producer")], + sample_path, + summary_path, + poll_seconds=0.01, + backend=backend, + ) + + monitor.start() + assert backend.sampled.wait(1) + summary = monitor.stop() + + assert summary["status"] == "ok" + assert summary["sample_count_by_gpu"]["0"] >= 1 + assert backend.closed + assert list(iter_gpu_samples(sample_path)) + assert json.loads(summary_path.read_text())["status"] == "ok" + + +def test_stop_timeout_does_not_race_a_blocked_sample(tmp_path, monkeypatch): + backend = _BlockingBackend() + monitor = GpuMonitor( + [GpuRoleAssignment(0, "producer", "producer")], + tmp_path / "samples.jsonl", + tmp_path / "summary.json", + poll_seconds=0.01, + backend=backend, + ) + monitor.start() + assert backend.sampled.wait(1) + monkeypatch.setattr(monitor._thread, "join", lambda timeout: None) + + summary = monitor.stop() + + assert summary["status"] == "degraded" + assert summary["sample_count_by_gpu"]["0"] == 0 + assert "GPU monitor thread did not stop" in summary["errors"] + backend.release_sample.set() + threading.Thread.join(monitor._thread, timeout=1) + assert not monitor._thread.is_alive() + assert backend.closed + + +def test_window_summary_excludes_startup_and_reports_active_time(): + assignment = GpuRoleAssignment(2, "consumer:b4", "consumer:b4") + samples = [ + { + "timestamp_monotonic_ns": timestamp, + "gpu": 2, + "utilization_gpu_pct": utilization, + "memory_used_bytes": 10 << 30, + "compute_pids": [123], + } + for timestamp, utilization in ( + (1_000_000_000, 0), + (2_000_000_000, 80), + (3_000_000_000, 100), + ) + ] + + result = summarize_gpu_window( + samples, + [assignment], + start_monotonic_ns=2_000_000_000, + end_monotonic_ns=3_000_000_000, + ) + + gpu = result["per_gpu"]["2"] + assert result["valid"] + assert gpu["sample_count"] == 2 + assert gpu["gpu_utilization_pct"]["mean"] == 90.0 + assert gpu["gpu_utilization_pct"]["p50"] == 80.0 + assert gpu["gpu_utilization_pct"]["p95"] == 100.0 + assert gpu["max_compute_processes"] == 1 diff --git a/tests/unit/benchmarks/test_independent_consumers.py b/tests/unit/benchmarks/test_independent_consumers.py index 23108439e..ed170ef8e 100644 --- a/tests/unit/benchmarks/test_independent_consumers.py +++ b/tests/unit/benchmarks/test_independent_consumers.py @@ -8,11 +8,13 @@ import pytest from pydantic import ValidationError +import speculators.benchmarks.independent_consumers as benchmark_module from speculators.benchmarks.independent_consumers import ( AccountingLedger, AccountingProxy, BenchmarkConfig, ConsumerSpec, + ConsumerStepEvent, ProducerSpec, RequestEvent, ScenarioSpec, @@ -109,6 +111,24 @@ def test_config_requires_serial_one_then_three_consumers(): ) +def test_run_benchmark_can_select_only_1p3c(tmp_path, monkeypatch): + observed = [] + + def fake_run_scenario(_config, scenario, _output_dir): + observed.append(scenario.kind) + return {"valid": True, "kind": scenario.kind} + + monkeypatch.setattr(benchmark_module, "_run_scenario", fake_run_scenario) + + report = benchmark_module.run_benchmark( + _config(), tmp_path / "run", scenario_kind="1p3c" + ) + + assert observed == ["1p3c"] + assert report["selected_scenarios"] == ["1p3c"] + assert [scenario["kind"] for scenario in report["scenarios"]] == ["1p3c"] + + @pytest.mark.parametrize( "command", [ @@ -391,6 +411,32 @@ def test_consumer_step_analysis_excludes_warmup_and_reports_percentiles(tmp_path assert result["step_ms_p95"] == 40.0 +def test_consumer_step_analysis_reports_exact_monotonic_measurement_window(tmp_path): + log_path = tmp_path / "consumer.log" + log_path.write_text("captured by callback\n") + events = [ + ConsumerStepEvent(1_000_000_000, 100.0), + ConsumerStepEvent(2_000_000_000, 20.0), + ConsumerStepEvent(3_000_000_000, 30.0), + ConsumerStepEvent(4_000_000_000, 40.0), + ] + + result = analyze_consumer_steps( + log_path, + warmup_steps=1, + minimum_steady_steps=2, + events=events, + measurement_steps=2, + ) + + assert result["valid"] + assert result["steady_steps"] == 2 + assert result["steady_started_at_monotonic_ns"] == 1_000_000_000 + assert result["steady_finished_at_monotonic_ns"] == 3_000_000_000 + assert result["steady_duration_seconds"] == 2.0 + assert result["steady_steps_per_second"] == 1.0 + + def test_consumer_step_analysis_fails_closed_on_missing_log(tmp_path): result = analyze_consumer_steps( tmp_path / "missing.log", warmup_steps=1, minimum_steady_steps=2 diff --git a/tests/unit/data_generation/test_artifact_cache.py b/tests/unit/data_generation/test_artifact_cache.py index 70e8e76c3..e5f6f75a2 100644 --- a/tests/unit/data_generation/test_artifact_cache.py +++ b/tests/unit/data_generation/test_artifact_cache.py @@ -15,6 +15,7 @@ from speculators.data_generation.artifact_cache import ( ArtifactLockTimeoutError, HiddenStateArtifactCache, + canonical_hidden_state_extraction_namespace, canonical_hidden_state_request_id, ) from speculators.data_generation.offline import check_hidden_states @@ -87,6 +88,21 @@ def test_canonical_request_id_is_stable_and_covers_semantics(): ) +def test_extraction_namespace_fingerprints_layers_and_user_namespace(): + first = canonical_hidden_state_extraction_namespace( + (2, 18, 33), user_namespace="revision-a" + ) + assert first == canonical_hidden_state_extraction_namespace( + (2, 18, 33), user_namespace="revision-a" + ) + assert first != canonical_hidden_state_extraction_namespace( + (2, 18, 36), user_namespace="revision-a" + ) + assert first != canonical_hidden_state_extraction_namespace( + (2, 18, 33), user_namespace="revision-b" + ) + + @pytest.mark.parametrize("tokens", [[], [1, True], [[1, 2]]]) def test_canonical_request_id_rejects_invalid_tokens(tokens): with pytest.raises(ValueError, match="input_ids"): diff --git a/tests/unit/train/test_cli_args.py b/tests/unit/train/test_cli_args.py index da09e48b9..9ba4117db 100644 --- a/tests/unit/train/test_cli_args.py +++ b/tests/unit/train/test_cli_args.py @@ -1,5 +1,7 @@ """Tests for CLI arguments.""" +import pytest + from scripts.train import parse_args from speculators.models.dflash.core import DFlashDraftModel from speculators.models.dspark.core import DSparkDraftModel @@ -45,6 +47,27 @@ def test_shared_hidden_state_cache_arguments(monkeypatch): assert args.shared_hidden_states_lock_timeout == 45 +@pytest.mark.parametrize( + "extra", + [ + ["--shared-hidden-states-path", ""], + ["--shared-hidden-states-namespace", "namespace-only"], + ["--shared-hidden-states-path", "cache", "--legacy-data"], + [ + "--shared-hidden-states-path", + "cache", + "--shared-hidden-states-namespace", + "", + ], + ["--shared-hidden-states-ttl", "-1"], + ["--shared-hidden-states-lock-timeout", "0"], + ], +) +def test_shared_hidden_state_cache_rejects_invalid_combinations(monkeypatch, extra): + with pytest.raises(SystemExit, match="2"): + _parse(monkeypatch, extra) + + # --------------------------------------------------------------------------- # Ensure CLI args flow correctly through vars(args) into get_trainer_kwargs # --------------------------------------------------------------------------- diff --git a/tests/unit/train/test_shared_artifacts.py b/tests/unit/train/test_shared_artifacts.py index f43fe49bb..57db85290 100644 --- a/tests/unit/train/test_shared_artifacts.py +++ b/tests/unit/train/test_shared_artifacts.py @@ -48,12 +48,26 @@ def _arrow_dataset( on_missing="generate", on_generate=on_generate, shared_artifacts_path=shared_path, + shared_artifacts_namespace=("layers:2,18,33" if shared_path else None), shared_artifacts_ttl_seconds=None, ) dataset.client = object() return dataset +def test_shared_dataset_requires_identity_namespace(tmp_path): + data_path = tmp_path / "data" + _write_dataset(data_path) + + with pytest.raises(ValueError, match="shared_artifacts_namespace is required"): + ArrowDataset( + max_len=128, + datapath=data_path, + model="model", + shared_artifacts_path=tmp_path / "shared", + ) + + def _successful_generator(service_path: Path, calls: list[Path]): def generate(*_args, **_kwargs): path = service_path / f"request-{len(calls)}.safetensors" From aa46739074d84c76a5f29b9288edb20cdf3fc426 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 09:14:17 +0800 Subject: [PATCH 05/20] feat: add bounded asynchronous artifact windows Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- docs/cli/train.md | 18 +- scripts/train.py | 97 ++ .../data_generation/artifact_cache.py | 34 + .../data_generation/windowed_artifacts.py | 1464 +++++++++++++++++ src/speculators/train/data.py | 403 ++++- src/speculators/train/dataloader.py | 150 +- src/speculators/train/trainer.py | 200 ++- .../test_windowed_artifacts.py | 499 ++++++ tests/unit/train/test_cli_args.py | 57 + tests/unit/train/test_shared_artifacts.py | 163 +- tests/unit/train/test_windowed_training.py | 231 +++ 11 files changed, 3249 insertions(+), 67 deletions(-) create mode 100644 src/speculators/data_generation/windowed_artifacts.py create mode 100644 tests/unit/data_generation/test_windowed_artifacts.py create mode 100644 tests/unit/train/test_windowed_training.py diff --git a/docs/cli/train.md b/docs/cli/train.md index d00e1e691..5fd91aafb 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -92,9 +92,25 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--shared-hidden-states-lock-timeout`** (float, default: `300.0`) Maximum seconds to wait while another trainer generates and atomically publishes the same artifact. +- **`--shared-hidden-states-consumer-id`** (str, default: `None`) Stable logical trainer identity. Setting this option enables bounded asynchronous windows and requires `--on-missing generate`. Training and validation append separate stream suffixes automatically. + +- **`--shared-hidden-states-lookbehind`** (int, default: `2`) Committed stream positions retained behind this consumer's cursor. + +- **`--shared-hidden-states-lookahead`** (int, default: `16`) Stream positions asynchronously prepared ahead of this consumer's committed cursor. + +- **`--shared-hidden-states-max-inflight`** (int, default: `32`) Maximum waiting or leased positions per consumer. Once the first sample of a packed batch is admitted, the rest of that batch may complete atomically so a batch larger than this value cannot deadlock before trainer ACK. + +- **`--shared-hidden-states-consumer-timeout`** (float, default: `120.0`) Heartbeat timeout before a dead consumer's window and leases are released. + +- **`--shared-hidden-states-claim-timeout`** (float, default: `300.0`) Timeout before an interrupted producer claim can be reassigned. + +- **`--shared-hidden-states-generation-attempts`** (int, default: `3`) Maximum coordinated generation attempts, including expired producer claims. + The shared cache is a filesystem data plane, not Mooncake or GPU-direct transport. Its directory must provide reliable POSIX `flock`, same-filesystem atomic rename, and directory `fsync` semantics to every trainer. Do not assume an arbitrary NFS mount is safe unless those guarantees have been verified. - This cache is not a consumer-centered bounded sliding window. Setting `--shared-hidden-states-ttl=0` retains one artifact per unique request. With a finite TTL, expired entries are reclaimed when a dataset opens the cache or when the same key is requested again, so a single pass over new samples can still grow on-disk usage. Provision the filesystem and set the TTL according to the maximum expected lag between consumers. + Without `--shared-hidden-states-consumer-id`, this remains the legacy TTL cache: setting the TTL to zero retains one artifact per unique request, and a finite TTL does not by itself bound a pass over unseen samples. + + With a consumer ID, SQLite tracks deterministic sampler positions, independent consumer cursors, generation claims, read leases, and the union of live windows. DataLoader workers only acquire and materialize authorized artifacts. The trainer main process advances the cursor after a successful training optimizer boundary or validation forward. Artifacts outside every live window are removed only after all read leases are released. In this mode TTL expiration is disabled; retention is controlled by windows and explicit leases. - **`--legacy-data`** (flag) **DEPRECATED.** Use the old data format which stores hidden states alongside token_ids. diff --git a/scripts/train.py b/scripts/train.py index 3aeed3f7b..1f553cedb 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -660,6 +660,19 @@ def main(args: argparse.Namespace): # noqa: C901 else args.shared_hidden_states_ttl ), shared_artifacts_lock_timeout_seconds=(args.shared_hidden_states_lock_timeout), + shared_artifacts_consumer_id=args.shared_hidden_states_consumer_id, + shared_artifacts_lookbehind=args.shared_hidden_states_lookbehind, + shared_artifacts_lookahead=args.shared_hidden_states_lookahead, + shared_artifacts_max_inflight=args.shared_hidden_states_max_inflight, + shared_artifacts_consumer_timeout_seconds=( + args.shared_hidden_states_consumer_timeout + ), + shared_artifacts_claim_timeout_seconds=( + args.shared_hidden_states_claim_timeout + ), + shared_artifacts_generation_attempts=( + args.shared_hidden_states_generation_attempts + ), hidden_size=hidden_size, num_target_layers=num_target_layers, num_workers=args.num_workers, @@ -787,6 +800,44 @@ def validate_draft_init_args( ) +def _validate_windowed_consumer_args( + parser: argparse.ArgumentParser, args: argparse.Namespace +) -> None: + consumer_id = args.shared_hidden_states_consumer_id + if consumer_id is None: + return + if args.shared_hidden_states_path is None: + parser.error( + "--shared-hidden-states-consumer-id requires --shared-hidden-states-path" + ) + if not consumer_id.strip(): + parser.error("--shared-hidden-states-consumer-id must be non-empty") + if args.on_missing != "generate": + parser.error( + "--shared-hidden-states-consumer-id requires --on-missing generate" + ) + + +def _validate_windowed_shared_hidden_state_args( + parser: argparse.ArgumentParser, args: argparse.Namespace +) -> None: + _validate_windowed_consumer_args(parser, args) + for name in ( + "shared_hidden_states_lookbehind", + "shared_hidden_states_lookahead", + ): + if getattr(args, name) < 0: + parser.error(f"--{name.replace('_', '-')} must be non-negative") + if args.shared_hidden_states_max_inflight < 1: + parser.error("--shared-hidden-states-max-inflight must be at least one") + if args.shared_hidden_states_consumer_timeout <= 0: + parser.error("--shared-hidden-states-consumer-timeout must be positive") + if args.shared_hidden_states_claim_timeout <= 0: + parser.error("--shared-hidden-states-claim-timeout must be positive") + if args.shared_hidden_states_generation_attempts < 1: + parser.error("--shared-hidden-states-generation-attempts must be at least one") + + def _validate_shared_hidden_state_args( parser: argparse.ArgumentParser, args: argparse.Namespace ) -> None: @@ -810,6 +861,7 @@ def _validate_shared_hidden_state_args( parser.error("--shared-hidden-states-ttl must be non-negative") if args.shared_hidden_states_lock_timeout <= 0: parser.error("--shared-hidden-states-lock-timeout must be positive") + _validate_windowed_shared_hidden_state_args(parser, args) def parse_args(): @@ -973,6 +1025,51 @@ def parse_args(): "artifact. Only applies when --shared-hidden-states-path is set." ), ) + parser.add_argument( + "--shared-hidden-states-consumer-id", + type=str, + default=None, + help=( + "Stable logical trainer identity. Setting this enables bounded " + "asynchronous windows and trainer-owned artifact acknowledgements." + ), + ) + parser.add_argument( + "--shared-hidden-states-lookbehind", + type=int, + default=2, + help="Committed positions retained behind each consumer cursor.", + ) + parser.add_argument( + "--shared-hidden-states-lookahead", + type=int, + default=16, + help="Positions asynchronously prepared ahead of each consumer cursor.", + ) + parser.add_argument( + "--shared-hidden-states-max-inflight", + type=int, + default=32, + help="Maximum waiting or leased positions for one logical consumer.", + ) + parser.add_argument( + "--shared-hidden-states-consumer-timeout", + type=float, + default=120.0, + help="Seconds without a heartbeat before releasing a consumer's interests.", + ) + parser.add_argument( + "--shared-hidden-states-claim-timeout", + type=float, + default=300.0, + help="Seconds before an interrupted producer claim can be reassigned.", + ) + parser.add_argument( + "--shared-hidden-states-generation-attempts", + type=int, + default=3, + help="Maximum coordinated generation attempts per artifact.", + ) parser.add_argument( "--legacy-data", action="store_true", diff --git a/src/speculators/data_generation/artifact_cache.py b/src/speculators/data_generation/artifact_cache.py index 324cc7ff2..389a64253 100644 --- a/src/speculators/data_generation/artifact_cache.py +++ b/src/speculators/data_generation/artifact_cache.py @@ -268,6 +268,40 @@ def snapshot_stats(self) -> dict[str, int]: with self._stats_lock(fcntl.LOCK_SH): return self._read_stats_unlocked() + def record_reuse(self) -> None: + """Account for a logical reader served by an existing publication.""" + self._record(logical_requests=1, hits=1) + + def load( + self, + request_id: str, + validate: Callable[[dict[str, torch.Tensor]], None], + ) -> dict[str, torch.Tensor]: + """Load a publication while holding its cross-process file lock.""" + self._validate_request_id(request_id) + with self._request_lock(request_id): + path = self.artifact_path(request_id) + if not path.exists(): + raise ArtifactCacheError(f"Artifact {request_id} is not published") + data = load_file(path) + validate(data) + return data + + def remove(self, request_id: str, *, expected_path: Path | None = None) -> bool: + """Remove one publication under the same lock used by readers/writers.""" + self._validate_request_id(request_id) + with self._request_lock(request_id): + path = self.artifact_path(request_id) + if expected_path is not None and path.resolve() != expected_path.resolve(): + raise ArtifactCacheError( + f"Artifact path mismatch for request {request_id}" + ) + if not path.exists(): + return False + path.unlink() + _fsync_directory(path.parent) + return True + def _is_expired(self, path: Path, now: float) -> bool: return bool( self.artifact_ttl_seconds is not None diff --git a/src/speculators/data_generation/windowed_artifacts.py b/src/speculators/data_generation/windowed_artifacts.py new file mode 100644 index 000000000..555dd7682 --- /dev/null +++ b/src/speculators/data_generation/windowed_artifacts.py @@ -0,0 +1,1464 @@ +"""Bounded asynchronous coordination for shared hidden-state artifacts. + +The coordinator is a single-host control plane. Tensor payloads remain in an +artifact store; SQLite contains only deterministic stream positions, consumer +progress, generation claims, and read leases. +""" + +from __future__ import annotations + +import hashlib +import json +import sqlite3 +import threading +import time +import uuid +from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum, IntEnum +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import os + from collections.abc import Callable, Iterator, Mapping, Sequence + +SCHEMA_VERSION = 1 +DIGEST_LENGTH = 64 +MAX_CONSUMER_ID_LENGTH = 128 + + +class WindowedArtifactError(RuntimeError): + """Base error for bounded artifact coordination.""" + + +class ArtifactGenerationError(WindowedArtifactError): + """A producer exhausted the configured generation attempts.""" + + +class ArtifactState(str, Enum): + ABSENT = "absent" + QUEUED = "queued" + GENERATING = "generating" + READY = "ready" + EVICTING = "evicting" + FAILED = "failed" + + +class ArtifactPriority(IntEnum): + DEMAND = 0 + PREFETCH = 1 + + +def _canonical_json(value: Mapping[str, Any]) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":")) + + +def canonical_stream_id(contract: Mapping[str, Any]) -> str: + """Return a stable identity for a dataset and sampler-order contract.""" + return hashlib.sha256(_canonical_json(contract).encode()).hexdigest() + + +def canonical_position_id( + stream_id: str, + *, + epoch: int, + ordinal: int, + dataset_index: int, + batch_ordinal: int, + batch_start_sequence: int, + batch_end_sequence: int, +) -> str: + value = { + "batch_end_sequence": batch_end_sequence, + "batch_ordinal": batch_ordinal, + "batch_start_sequence": batch_start_sequence, + "dataset_index": dataset_index, + "epoch": epoch, + "ordinal": ordinal, + "stream_id": stream_id, + } + return hashlib.sha256(_canonical_json(value).encode()).hexdigest() + + +@dataclass(frozen=True) +class StreamSampleIndex: + """Dataset index annotated with its deterministic sampler position.""" + + stream_id: str + sequence: int + epoch: int + ordinal: int + dataset_index: int + batch_ordinal: int + batch_start_sequence: int + batch_end_sequence: int + request_id: str + position_id: str + + def __post_init__(self) -> None: + for name in ("stream_id", "request_id", "position_id"): + value = getattr(self, name) + if len(value) != DIGEST_LENGTH or any( + ch not in "0123456789abcdef" for ch in value + ): + raise ValueError(f"{name} must be a lowercase SHA-256 digest") + for name in ( + "sequence", + "epoch", + "ordinal", + "dataset_index", + "batch_ordinal", + "batch_start_sequence", + "batch_end_sequence", + ): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{name} must be a non-negative integer") + if not self.batch_start_sequence <= self.sequence < self.batch_end_sequence: + raise ValueError("sequence must be inside its batch boundaries") + + +@dataclass(frozen=True) +class ArtifactReadLease: + token: str + consumer_id: str + stream_id: str + sequence: int + request_id: str + path: Path + generation: int + cache_hit: bool + wait_seconds: float + + def as_batch_metadata(self) -> dict[str, Any]: + return { + "token": self.token, + "consumer_id": self.consumer_id, + "stream_id": self.stream_id, + "sequence": self.sequence, + "request_id": self.request_id, + "generation": self.generation, + } + + +@dataclass(frozen=True) +class GenerationClaim: + request_id: str + stream_id: str + dataset_index: int + generation: int + priority: ArtifactPriority + + +@dataclass(frozen=True) +class EvictionClaim: + request_id: str + generation: int + path: Path + + +class WindowedArtifactCoordinator: + """Transactional authority for independent consumer windows.""" + + def __init__( + self, + root: str | os.PathLike[str], + *, + poll_seconds: float = 0.02, + consumer_timeout_seconds: float = 120.0, + claim_timeout_seconds: float = 300.0, + max_generation_attempts: int = 3, + clock: Callable[[], float] = time.time, + ) -> None: + if poll_seconds <= 0: + raise ValueError("poll_seconds must be positive") + if consumer_timeout_seconds <= 0: + raise ValueError("consumer_timeout_seconds must be positive") + if claim_timeout_seconds <= 0: + raise ValueError("claim_timeout_seconds must be positive") + if max_generation_attempts < 1: + raise ValueError("max_generation_attempts must be at least one") + self.root = Path(root).expanduser().resolve() + self.root.mkdir(parents=True, exist_ok=True) + self.path = self.root / "windowed-artifacts.sqlite3" + self.poll_seconds = poll_seconds + self.consumer_timeout_seconds = consumer_timeout_seconds + self.claim_timeout_seconds = claim_timeout_seconds + self.max_generation_attempts = max_generation_attempts + self._clock = clock + self._lock = threading.RLock() + self._conn = sqlite3.connect( + self.path, timeout=30.0, isolation_level=None, check_same_thread=False + ) + self._conn.row_factory = sqlite3.Row + self._conn.execute("PRAGMA busy_timeout=30000") + self._conn.execute("PRAGMA synchronous=NORMAL") + self._conn.execute("PRAGMA foreign_keys=ON") + if not self._schema_is_current(): + self._conn.execute("PRAGMA journal_mode=WAL") + self._create_schema() + + def _schema_is_current(self) -> bool: + table = self._conn.execute( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name='coordinator_meta'" + ).fetchone() + if table is None: + return False + row = self._conn.execute( + "SELECT value FROM coordinator_meta WHERE key='schema_version'" + ).fetchone() + if row is None: + return False + observed = int(row["value"]) + if observed != SCHEMA_VERSION: + raise WindowedArtifactError( + f"unsupported coordinator schema version {observed}; " + f"expected {SCHEMA_VERSION}" + ) + return True + + def _create_schema(self) -> None: + schema = """ + CREATE TABLE IF NOT EXISTS coordinator_meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL + ); + CREATE TABLE IF NOT EXISTS streams ( + stream_id TEXT PRIMARY KEY, + contract_json TEXT NOT NULL, + created_at REAL NOT NULL + ); + CREATE TABLE IF NOT EXISTS positions ( + stream_id TEXT NOT NULL, + sequence INTEGER NOT NULL, + epoch INTEGER NOT NULL, + ordinal INTEGER NOT NULL, + dataset_index INTEGER NOT NULL, + batch_ordinal INTEGER NOT NULL, + batch_start_sequence INTEGER NOT NULL, + batch_end_sequence INTEGER NOT NULL, + request_id TEXT NOT NULL, + position_id TEXT NOT NULL UNIQUE, + PRIMARY KEY (stream_id, sequence), + UNIQUE (stream_id, epoch, ordinal), + FOREIGN KEY (stream_id) REFERENCES streams(stream_id) + ); + CREATE INDEX IF NOT EXISTS positions_request + ON positions(request_id, stream_id); + CREATE TABLE IF NOT EXISTS consumers ( + consumer_id TEXT PRIMARY KEY, + stream_id TEXT NOT NULL, + cursor INTEGER NOT NULL, + lookbehind INTEGER NOT NULL, + lookahead INTEGER NOT NULL, + max_inflight INTEGER NOT NULL, + state TEXT NOT NULL, + heartbeat_at REAL NOT NULL, + updated_at REAL NOT NULL, + FOREIGN KEY (stream_id) REFERENCES streams(stream_id) + ); + CREATE TABLE IF NOT EXISTS artifacts ( + request_id TEXT PRIMARY KEY, + state TEXT NOT NULL, + generation INTEGER NOT NULL DEFAULT 0, + path TEXT, + size_bytes INTEGER NOT NULL DEFAULT 0, + priority INTEGER, + queued_at REAL, + claim_owner TEXT, + claim_until REAL, + failures INTEGER NOT NULL DEFAULT 0, + last_error TEXT, + first_reader_accounted INTEGER NOT NULL DEFAULT 0, + updated_at REAL NOT NULL + ); + CREATE INDEX IF NOT EXISTS artifacts_schedule + ON artifacts(state, priority, queued_at); + CREATE TABLE IF NOT EXISTS interests ( + consumer_id TEXT NOT NULL, + stream_id TEXT NOT NULL, + sequence INTEGER NOT NULL, + request_id TEXT NOT NULL, + kind TEXT NOT NULL CHECK(kind IN ('window', 'demand')), + created_at REAL NOT NULL, + PRIMARY KEY (consumer_id, stream_id, sequence, kind), + FOREIGN KEY (consumer_id) REFERENCES consumers(consumer_id) + ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS interests_request ON interests(request_id); + CREATE TABLE IF NOT EXISTS acquisitions ( + token TEXT PRIMARY KEY, + consumer_id TEXT NOT NULL, + stream_id TEXT NOT NULL, + sequence INTEGER NOT NULL, + request_id TEXT NOT NULL, + state TEXT NOT NULL CHECK(state IN ('waiting', 'leased')), + created_at REAL NOT NULL, + updated_at REAL NOT NULL, + FOREIGN KEY (consumer_id) REFERENCES consumers(consumer_id) + ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS acquisitions_request + ON acquisitions(request_id); + CREATE TABLE IF NOT EXISTS completed_positions ( + consumer_id TEXT NOT NULL, + sequence INTEGER NOT NULL, + PRIMARY KEY (consumer_id, sequence), + FOREIGN KEY (consumer_id) REFERENCES consumers(consumer_id) + ON DELETE CASCADE + ); + """ + with self._lock: + self._conn.executescript(schema) + self._conn.execute( + "INSERT OR IGNORE INTO coordinator_meta(key,value) VALUES" + "('schema_version',?),('scheduler_cursor','')," + "('peak_retained_artifacts','0'),('peak_retained_bytes','0')," + "('peak_inflight_acquisitions','0')", + (str(SCHEMA_VERSION),), + ) + + @contextmanager + def _transaction(self) -> Iterator[sqlite3.Connection]: + with self._lock: + self._conn.execute("BEGIN IMMEDIATE") + try: + yield self._conn + except BaseException: + self._conn.rollback() + raise + else: + self._conn.commit() + + def close(self) -> None: + with self._lock: + self._conn.close() + + def __enter__(self) -> WindowedArtifactCoordinator: + return self + + def __exit__(self, *_args: object) -> None: + self.close() + + @staticmethod + def _set_max_meta_locked(conn: sqlite3.Connection, key: str, observed: int) -> None: + row = conn.execute( + "SELECT value FROM coordinator_meta WHERE key=?", (key,) + ).fetchone() + previous = int(row["value"]) if row is not None else 0 + if observed > previous: + conn.execute( + "INSERT OR REPLACE INTO coordinator_meta(key,value) VALUES(?,?)", + (key, str(observed)), + ) + + def _update_high_water_locked(self, conn: sqlite3.Connection) -> None: + retained = conn.execute( + "SELECT COUNT(*) AS count,COALESCE(SUM(size_bytes),0) AS bytes " + "FROM artifacts WHERE state IN (?,?)", + (ArtifactState.READY.value, ArtifactState.EVICTING.value), + ).fetchone() + inflight = int(conn.execute("SELECT COUNT(*) FROM acquisitions").fetchone()[0]) + self._set_max_meta_locked( + conn, "peak_retained_artifacts", int(retained["count"]) + ) + self._set_max_meta_locked(conn, "peak_retained_bytes", int(retained["bytes"])) + self._set_max_meta_locked(conn, "peak_inflight_acquisitions", inflight) + + @staticmethod + def _validate_digest(name: str, value: str) -> None: + if len(value) != DIGEST_LENGTH or any( + ch not in "0123456789abcdef" for ch in value + ): + raise ValueError(f"{name} must be a lowercase SHA-256 digest") + + @staticmethod + def _validate_consumer_id(consumer_id: str) -> None: + if not consumer_id or len(consumer_id) > MAX_CONSUMER_ID_LENGTH: + raise ValueError("consumer_id must contain 1-128 characters") + + def register_stream(self, contract: Mapping[str, Any]) -> str: + contract_json = _canonical_json(contract) + stream_id = canonical_stream_id(contract) + with self._transaction() as conn: + row = conn.execute( + "SELECT contract_json FROM streams WHERE stream_id=?", (stream_id,) + ).fetchone() + if row is not None and row["contract_json"] != contract_json: + raise WindowedArtifactError("stream identity collision") + conn.execute( + "INSERT OR IGNORE INTO streams VALUES(?,?,?)", + (stream_id, contract_json, self._clock()), + ) + return stream_id + + def register_positions(self, samples: Sequence[StreamSampleIndex]) -> None: + if not samples: + return + stream_ids = {sample.stream_id for sample in samples} + if len(stream_ids) != 1: + raise ValueError("all registered positions must belong to one stream") + stream_id = next(iter(stream_ids)) + with self._transaction() as conn: + if ( + conn.execute( + "SELECT 1 FROM streams WHERE stream_id=?", (stream_id,) + ).fetchone() + is None + ): + raise KeyError(f"unknown stream {stream_id!r}") + now = self._clock() + for sample in samples: + expected_position_id = canonical_position_id( + stream_id, + epoch=sample.epoch, + ordinal=sample.ordinal, + dataset_index=sample.dataset_index, + batch_ordinal=sample.batch_ordinal, + batch_start_sequence=sample.batch_start_sequence, + batch_end_sequence=sample.batch_end_sequence, + ) + if sample.position_id != expected_position_id: + raise ValueError( + "sample position identity does not match its fields" + ) + row = conn.execute( + "SELECT * FROM positions WHERE stream_id=? AND sequence=?", + (stream_id, sample.sequence), + ).fetchone() + identity = ( + sample.epoch, + sample.ordinal, + sample.dataset_index, + sample.batch_ordinal, + sample.batch_start_sequence, + sample.batch_end_sequence, + sample.request_id, + sample.position_id, + ) + if row is not None: + observed = tuple( + row[name] + for name in ( + "epoch", + "ordinal", + "dataset_index", + "batch_ordinal", + "batch_start_sequence", + "batch_end_sequence", + "request_id", + "position_id", + ) + ) + if observed != identity: + raise WindowedArtifactError( + "registered stream position changed: " + f"sequence={sample.sequence}" + ) + continue + conn.execute( + "INSERT INTO positions VALUES(?,?,?,?,?,?,?,?,?,?)", + ( + stream_id, + sample.sequence, + *identity, + ), + ) + conn.execute( + "INSERT OR IGNORE INTO artifacts" + "(request_id,state,updated_at) VALUES(?,?,?)", + (sample.request_id, ArtifactState.ABSENT.value, now), + ) + consumers = conn.execute( + "SELECT consumer_id FROM consumers WHERE stream_id=? " + "AND state='active'", + (stream_id,), + ).fetchall() + for row in consumers: + self._refresh_window_locked(conn, row["consumer_id"]) + + def register_consumer( + self, + consumer_id: str, + *, + stream_id: str, + lookbehind: int, + lookahead: int, + max_inflight: int, + cursor: int = 0, + reset: bool = False, + ) -> None: + self._validate_consumer_id(consumer_id) + self._validate_digest("stream_id", stream_id) + for name, value in ( + ("lookbehind", lookbehind), + ("lookahead", lookahead), + ("cursor", cursor), + ): + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{name} must be a non-negative integer") + if isinstance(max_inflight, bool) or not isinstance(max_inflight, int): + raise TypeError("max_inflight must be an integer") + if max_inflight < 1: + raise ValueError("max_inflight must be at least one") + + with self._transaction() as conn: + if ( + conn.execute( + "SELECT 1 FROM streams WHERE stream_id=?", (stream_id,) + ).fetchone() + is None + ): + raise KeyError(f"unknown stream {stream_id!r}") + row = conn.execute( + "SELECT * FROM consumers WHERE consumer_id=?", (consumer_id,) + ).fetchone() + config = (stream_id, lookbehind, lookahead, max_inflight) + now = self._clock() + if row is None: + conn.execute( + "INSERT INTO consumers VALUES(?,?,?,?,?,?,?,?,?)", + ( + consumer_id, + stream_id, + cursor, + lookbehind, + lookahead, + max_inflight, + "active", + now, + now, + ), + ) + else: + observed = tuple( + row[name] + for name in ( + "stream_id", + "lookbehind", + "lookahead", + "max_inflight", + ) + ) + if observed != config: + raise WindowedArtifactError( + f"consumer {consumer_id!r} configuration changed" + ) + if reset: + conn.execute( + "DELETE FROM acquisitions WHERE consumer_id=?", + (consumer_id,), + ) + conn.execute( + "DELETE FROM completed_positions WHERE consumer_id=?", + (consumer_id,), + ) + conn.execute( + "DELETE FROM interests WHERE consumer_id=?", (consumer_id,) + ) + conn.execute( + "UPDATE consumers SET cursor=?,state='active',heartbeat_at=?," + "updated_at=? WHERE consumer_id=?", + (cursor, now, now, consumer_id), + ) + elif row["state"] == "completed" and int(row["cursor"]) == cursor: + conn.execute( + "UPDATE consumers SET state='active',heartbeat_at=?," + "updated_at=? " + "WHERE consumer_id=?", + (now, now, consumer_id), + ) + elif row["state"] != "active": + raise WindowedArtifactError( + f"consumer {consumer_id!r} is {row['state']!r}; resume " + "requires an explicit cursor reset" + ) + elif int(row["cursor"]) != cursor: + raise WindowedArtifactError( + f"consumer {consumer_id!r} cursor is {row['cursor']}, " + f"not requested cursor {cursor}" + ) + else: + conn.execute( + "UPDATE consumers SET heartbeat_at=?,updated_at=? " + "WHERE consumer_id=?", + (now, now, consumer_id), + ) + self._refresh_window_locked(conn, consumer_id) + + def _queue_artifact_locked( + self, + conn: sqlite3.Connection, + request_id: str, + priority: ArtifactPriority, + ) -> None: + row = conn.execute( + "SELECT * FROM artifacts WHERE request_id=?", (request_id,) + ).fetchone() + if row is None: + conn.execute( + "INSERT INTO artifacts(request_id,state,updated_at) VALUES(?,?,?)", + (request_id, ArtifactState.ABSENT.value, self._clock()), + ) + row = conn.execute( + "SELECT * FROM artifacts WHERE request_id=?", (request_id,) + ).fetchone() + state = ArtifactState(row["state"]) + if state == ArtifactState.FAILED and int(row["failures"]) >= ( + self.max_generation_attempts + ): + return + if state in (ArtifactState.ABSENT, ArtifactState.FAILED): + conn.execute( + "UPDATE artifacts SET state=?,priority=?,queued_at=?,claim_owner=NULL," + "claim_until=NULL,updated_at=? WHERE request_id=?", + ( + ArtifactState.QUEUED.value, + int(priority), + self._clock(), + self._clock(), + request_id, + ), + ) + elif state == ArtifactState.QUEUED and ( + row["priority"] is None or int(row["priority"]) > int(priority) + ): + conn.execute( + "UPDATE artifacts SET priority=?,updated_at=? WHERE request_id=?", + (int(priority), self._clock(), request_id), + ) + + @staticmethod + def _retry_priority_locked( + conn: sqlite3.Connection, request_id: str + ) -> ArtifactPriority: + demand = conn.execute( + "SELECT 1 FROM acquisitions WHERE request_id=? LIMIT 1", + (request_id,), + ).fetchone() + return ( + ArtifactPriority.DEMAND if demand is not None else ArtifactPriority.PREFETCH + ) + + def _refresh_window_locked( + self, conn: sqlite3.Connection, consumer_id: str + ) -> None: + consumer = conn.execute( + "SELECT * FROM consumers WHERE consumer_id=?", (consumer_id,) + ).fetchone() + if consumer is None: + raise KeyError(f"unknown consumer {consumer_id!r}") + if consumer["state"] != "active": + return + cursor = int(consumer["cursor"]) + low = max(0, cursor - int(consumer["lookbehind"])) + high = cursor + int(consumer["lookahead"]) + 1 + rows = conn.execute( + "SELECT sequence,request_id FROM positions WHERE stream_id=? " + "AND batch_end_sequence>? AND batch_start_sequence None: + conn.execute( + "UPDATE artifacts SET state=?,priority=NULL,queued_at=NULL,updated_at=? " + "WHERE state IN (?,?) AND NOT EXISTS " + "(SELECT 1 FROM interests i WHERE i.request_id=artifacts.request_id) " + "AND NOT EXISTS (SELECT 1 FROM acquisitions a " + "WHERE a.request_id=artifacts.request_id)", + ( + ArtifactState.ABSENT.value, + self._clock(), + ArtifactState.QUEUED.value, + ArtifactState.FAILED.value, + ), + ) + + def heartbeat(self, consumer_id: str) -> None: + with self._transaction() as conn: + result = conn.execute( + "UPDATE consumers SET heartbeat_at=?,updated_at=? " + "WHERE consumer_id=? AND state='active'", + (self._clock(), self._clock(), consumer_id), + ) + if result.rowcount != 1: + raise KeyError(f"unknown or inactive consumer {consumer_id!r}") + + def recover_expired(self) -> dict[str, int]: + expired_consumers = 0 + expired_claims = 0 + with self._transaction() as conn: + now = self._clock() + consumers = conn.execute( + "SELECT consumer_id FROM consumers WHERE state='active' " + "AND heartbeat_at ArtifactReadLease: + """Wait for an authorized stream position and acquire a read lease.""" + if timeout_seconds is not None and timeout_seconds <= 0: + raise ValueError("timeout_seconds must be positive or None") + started = time.monotonic() + deadline = None if timeout_seconds is None else started + timeout_seconds + token: str | None = None + while token is None: + with self._transaction() as conn: + consumer = conn.execute( + "SELECT * FROM consumers WHERE consumer_id=?", (consumer_id,) + ).fetchone() + if consumer is None or consumer["state"] != "active": + raise WindowedArtifactError( + f"consumer {consumer_id!r} is not active" + ) + position = conn.execute( + "SELECT * FROM positions WHERE stream_id=? AND sequence=?", + (sample.stream_id, sample.sequence), + ).fetchone() + expected = ( + sample.epoch, + sample.ordinal, + sample.dataset_index, + sample.batch_ordinal, + sample.batch_start_sequence, + sample.batch_end_sequence, + sample.request_id, + sample.position_id, + ) + observed = ( + tuple( + position[name] + for name in ( + "epoch", + "ordinal", + "dataset_index", + "batch_ordinal", + "batch_start_sequence", + "batch_end_sequence", + "request_id", + "position_id", + ) + ) + if position is not None + else None + ) + if consumer["stream_id"] != sample.stream_id or observed != expected: + raise WindowedArtifactError( + "requested sample does not match the registered stream position" + ) + now = self._clock() + conn.execute( + "UPDATE consumers SET heartbeat_at=?,updated_at=? " + "WHERE consumer_id=?", + (now, now, consumer_id), + ) + cursor = int(consumer["cursor"]) + low = max(0, cursor - int(consumer["lookbehind"])) + high = cursor + int(consumer["lookahead"]) + inflight = int( + conn.execute( + "SELECT COUNT(*) FROM acquisitions WHERE consumer_id=?", + (consumer_id,), + ).fetchone()[0] + ) + same_batch = ( + conn.execute( + "SELECT 1 FROM acquisitions a JOIN positions p " + "ON p.stream_id=a.stream_id AND p.sequence=a.sequence " + "WHERE a.consumer_id=? AND p.epoch=? " + "AND p.batch_ordinal=? LIMIT 1", + (consumer_id, sample.epoch, sample.batch_ordinal), + ).fetchone() + is not None + ) + inside_window = ( + sample.batch_end_sequence > low + and sample.batch_start_sequence <= high + ) + has_capacity = inflight < int(consumer["max_inflight"]) + if inside_window and (has_capacity or same_batch): + token = uuid.uuid4().hex + conn.execute( + "INSERT INTO acquisitions VALUES(?,?,?,?,?,'waiting',?,?)", + ( + token, + consumer_id, + sample.stream_id, + sample.sequence, + sample.request_id, + now, + now, + ), + ) + conn.execute( + "INSERT OR IGNORE INTO interests VALUES(?,?,?,?,?,?)", + ( + consumer_id, + sample.stream_id, + sample.sequence, + sample.request_id, + "demand", + now, + ), + ) + self._queue_artifact_locked( + conn, sample.request_id, ArtifactPriority.DEMAND + ) + self._update_high_water_locked(conn) + if token is None: + self._sleep_or_timeout(deadline, started, timeout_seconds, sample) + + while True: + with self._transaction() as conn: + acquisition = conn.execute( + "SELECT * FROM acquisitions WHERE token=?", (token,) + ).fetchone() + if acquisition is None: + raise WindowedArtifactError( + f"artifact acquisition {token} was released while waiting" + ) + artifact = conn.execute( + "SELECT * FROM artifacts WHERE request_id=?", + (sample.request_id,), + ).fetchone() + if artifact is None: + raise WindowedArtifactError("artifact metadata disappeared") + state = ArtifactState(artifact["state"]) + if state == ArtifactState.READY: + if not artifact["path"]: + raise WindowedArtifactError("ready artifact has no path") + cache_hit = bool(artifact["first_reader_accounted"]) + conn.execute( + "UPDATE artifacts SET first_reader_accounted=1,updated_at=? " + "WHERE request_id=?", + (self._clock(), sample.request_id), + ) + conn.execute( + "UPDATE acquisitions SET state='leased',updated_at=? " + "WHERE token=?", + (self._clock(), token), + ) + return ArtifactReadLease( + token=token, + consumer_id=consumer_id, + stream_id=sample.stream_id, + sequence=sample.sequence, + request_id=sample.request_id, + path=Path(artifact["path"]), + generation=int(artifact["generation"]), + cache_hit=cache_hit, + wait_seconds=time.monotonic() - started, + ) + if state == ArtifactState.FAILED and int(artifact["failures"]) >= ( + self.max_generation_attempts + ): + conn.execute("DELETE FROM acquisitions WHERE token=?", (token,)) + conn.execute( + "DELETE FROM interests WHERE consumer_id=? AND stream_id=? " + "AND sequence=? AND kind='demand'", + (consumer_id, sample.stream_id, sample.sequence), + ) + raise ArtifactGenerationError( + f"artifact {sample.request_id} failed after " + f"{artifact['failures']} attempts: {artifact['last_error']}" + ) + if state == ArtifactState.FAILED: + self._queue_artifact_locked( + conn, sample.request_id, ArtifactPriority.DEMAND + ) + conn.execute( + "UPDATE consumers SET heartbeat_at=?,updated_at=? " + "WHERE consumer_id=?", + (self._clock(), self._clock(), consumer_id), + ) + self._sleep_or_timeout(deadline, started, timeout_seconds, sample, token) + + def _sleep_or_timeout( + self, + deadline: float | None, + started: float, + timeout_seconds: float | None, + sample: StreamSampleIndex, + token: str | None = None, + ) -> None: + if deadline is not None and time.monotonic() >= deadline: + if token is not None: + self.abandon_tokens(sample.stream_id, [token]) + raise TimeoutError( + f"stream position {sample.sequence} was not ready within " + f"{timeout_seconds:.1f}s" + ) + sleep_seconds = self.poll_seconds + if deadline is not None: + sleep_seconds = min(sleep_seconds, max(0.0, deadline - time.monotonic())) + if time.monotonic() >= started: + time.sleep(sleep_seconds) + + def ack(self, consumer_id: str, leases: Sequence[Mapping[str, Any]]) -> int: + """Commit successful trainer consumption and advance a contiguous cursor.""" + if not leases: + with self._lock: + row = self._conn.execute( + "SELECT cursor FROM consumers WHERE consumer_id=?", (consumer_id,) + ).fetchone() + if row is None: + raise KeyError(f"unknown consumer {consumer_id!r}") + return int(row["cursor"]) + with self._transaction() as conn: + consumer = conn.execute( + "SELECT * FROM consumers WHERE consumer_id=?", (consumer_id,) + ).fetchone() + if consumer is None or consumer["state"] != "active": + raise WindowedArtifactError(f"consumer {consumer_id!r} is not active") + for lease in leases: + token = str(lease["token"]) + row = conn.execute( + "SELECT * FROM acquisitions WHERE token=?", (token,) + ).fetchone() + expected = ( + consumer_id, + str(lease["stream_id"]), + int(lease["sequence"]), + str(lease["request_id"]), + "leased", + ) + observed = ( + ( + row["consumer_id"], + row["stream_id"], + int(row["sequence"]), + row["request_id"], + row["state"], + ) + if row is not None + else None + ) + if observed != expected: + raise WindowedArtifactError( + f"read lease {token!r} is unknown, stale, or mismatched" + ) + conn.execute( + "INSERT OR IGNORE INTO completed_positions VALUES(?,?)", + (consumer_id, int(row["sequence"])), + ) + conn.execute("DELETE FROM acquisitions WHERE token=?", (token,)) + conn.execute( + "DELETE FROM interests WHERE consumer_id=? AND stream_id=? " + "AND sequence=? AND kind='demand'", + (consumer_id, row["stream_id"], int(row["sequence"])), + ) + cursor = int(consumer["cursor"]) + while ( + conn.execute( + "SELECT 1 FROM completed_positions WHERE consumer_id=? " + "AND sequence=?", + (consumer_id, cursor), + ).fetchone() + is not None + ): + conn.execute( + "DELETE FROM completed_positions WHERE consumer_id=? " + "AND sequence=?", + (consumer_id, cursor), + ) + cursor += 1 + now = self._clock() + conn.execute( + "UPDATE consumers SET cursor=?,heartbeat_at=?,updated_at=? " + "WHERE consumer_id=?", + (cursor, now, now, consumer_id), + ) + self._refresh_window_locked(conn, consumer_id) + return cursor + + def abandon(self, consumer_id: str, leases: Sequence[Mapping[str, Any]]) -> None: + self._abandon_tokens(consumer_id, [str(lease["token"]) for lease in leases]) + + def abandon_tokens(self, stream_id: str, tokens: Sequence[str]) -> None: + """Release timed-out waiters when only the stream identity is available.""" + del stream_id + self._abandon_tokens(None, tokens) + + def _abandon_tokens(self, consumer_id: str | None, tokens: Sequence[str]) -> None: + if not tokens: + return + with self._transaction() as conn: + affected: set[str] = set() + for token in tokens: + row = conn.execute( + "SELECT * FROM acquisitions WHERE token=?", (token,) + ).fetchone() + if row is None: + continue + if consumer_id is not None and row["consumer_id"] != consumer_id: + raise WindowedArtifactError( + f"lease {token!r} does not belong to {consumer_id!r}" + ) + affected.add(row["consumer_id"]) + conn.execute("DELETE FROM acquisitions WHERE token=?", (token,)) + conn.execute( + "DELETE FROM interests WHERE consumer_id=? AND stream_id=? " + "AND sequence=? AND kind='demand'", + (row["consumer_id"], row["stream_id"], int(row["sequence"])), + ) + for owner in affected: + self._refresh_window_locked(conn, owner) + + def claim_generation( + self, owner: str, *, stream_id: str, max_claims: int = 1 + ) -> tuple[GenerationClaim, ...]: + if not owner: + raise ValueError("generation owner must be non-empty") + if max_claims < 1: + raise ValueError("max_claims must be at least one") + with self._transaction() as conn: + self._recover_claims_locked(conn) + rows = conn.execute( + "SELECT DISTINCT a.* FROM artifacts a " + "JOIN interests i ON i.request_id=a.request_id " + "WHERE a.state=? AND i.stream_id=? " + "ORDER BY a.priority,a.queued_at,a.request_id LIMIT ?", + ( + ArtifactState.QUEUED.value, + stream_id, + max(64, max_claims * 8), + ), + ).fetchall() + if not rows: + return () + prefetch_rows: list[sqlite3.Row] = [] + by_consumer: dict[str, list[sqlite3.Row]] = {} + for row in rows: + consumers = conn.execute( + "SELECT DISTINCT consumer_id FROM acquisitions " + "WHERE request_id=? ORDER BY consumer_id", + (row["request_id"],), + ).fetchall() + is_demand = int(row["priority"]) == int( + ArtifactPriority.DEMAND + ) and bool(consumers) + if is_demand: + for consumer in consumers: + by_consumer.setdefault(consumer["consumer_id"], []).append(row) + else: + prefetch_rows.append(row) + cursor_row = conn.execute( + "SELECT value FROM coordinator_meta WHERE key='scheduler_cursor'" + ).fetchone() + previous = cursor_row["value"] if cursor_row is not None else "" + consumer_order = sorted(by_consumer) + if previous in consumer_order: + start = (consumer_order.index(previous) + 1) % len(consumer_order) + consumer_order = consumer_order[start:] + consumer_order[:start] + + selected: list[sqlite3.Row] = [] + selected_ids: set[str] = set() + last_consumer = previous + while consumer_order and len(selected) < max_claims: + made_progress = False + for consumer_id in consumer_order: + while by_consumer[consumer_id]: + row = by_consumer[consumer_id].pop(0) + if row["request_id"] not in selected_ids: + selected.append(row) + selected_ids.add(row["request_id"]) + last_consumer = consumer_id + made_progress = True + break + if len(selected) >= max_claims: + break + if not made_progress: + break + for row in prefetch_rows: + if len(selected) >= max_claims: + break + if row["request_id"] not in selected_ids: + selected.append(row) + selected_ids.add(row["request_id"]) + + now = self._clock() + claims: list[GenerationClaim] = [] + for row in selected: + position = conn.execute( + "SELECT dataset_index FROM positions WHERE stream_id=? " + "AND request_id=? ORDER BY sequence LIMIT 1", + (stream_id, row["request_id"]), + ).fetchone() + if position is None: + continue + result = conn.execute( + "UPDATE artifacts SET state=?,claim_owner=?,claim_until=?," + "updated_at=? WHERE request_id=? AND state=?", + ( + ArtifactState.GENERATING.value, + owner, + now + self.claim_timeout_seconds, + now, + row["request_id"], + ArtifactState.QUEUED.value, + ), + ) + if result.rowcount != 1: + continue + claims.append( + GenerationClaim( + request_id=row["request_id"], + stream_id=stream_id, + dataset_index=int(position["dataset_index"]), + generation=int(row["generation"]), + priority=ArtifactPriority(int(row["priority"])), + ) + ) + if last_consumer: + conn.execute( + "UPDATE coordinator_meta SET value=? WHERE key='scheduler_cursor'", + (last_consumer,), + ) + return tuple(claims) + + def _recover_claims_locked(self, conn: sqlite3.Connection) -> None: + now = self._clock() + rows = conn.execute( + "SELECT request_id,failures FROM artifacts WHERE state=? AND claim_until None: + artifact_path = str(Path(path).expanduser().resolve()) + if size_bytes < 0: + raise ValueError("size_bytes must be non-negative") + with self._transaction() as conn: + row = conn.execute( + "SELECT * FROM artifacts WHERE request_id=?", (claim.request_id,) + ).fetchone() + if row is None or ( + row["state"], + row["claim_owner"], + int(row["generation"]), + ) != (ArtifactState.GENERATING.value, owner, claim.generation): + raise WindowedArtifactError( + f"stale generation completion for {claim.request_id}" + ) + conn.execute( + "UPDATE artifacts SET state=?,path=?,size_bytes=?,priority=NULL," + "queued_at=NULL,claim_owner=NULL,claim_until=NULL,failures=0," + "last_error=NULL,first_reader_accounted=0,updated_at=? " + "WHERE request_id=?", + ( + ArtifactState.READY.value, + artifact_path, + size_bytes, + self._clock(), + claim.request_id, + ), + ) + self._update_high_water_locked(conn) + + def fail_generation( + self, owner: str, claim: GenerationClaim, error: BaseException | str + ) -> None: + message = str(error)[:2000] or type(error).__name__ + with self._transaction() as conn: + row = conn.execute( + "SELECT * FROM artifacts WHERE request_id=?", (claim.request_id,) + ).fetchone() + if row is None or ( + row["state"], + row["claim_owner"], + int(row["generation"]), + ) != (ArtifactState.GENERATING.value, owner, claim.generation): + raise WindowedArtifactError( + f"stale generation failure for {claim.request_id}" + ) + failures = int(row["failures"]) + 1 + interested = conn.execute( + "SELECT 1 FROM interests WHERE request_id=? LIMIT 1", + (claim.request_id,), + ).fetchone() + retry = failures < self.max_generation_attempts and interested is not None + priority = ( + self._retry_priority_locked(conn, claim.request_id) if retry else None + ) + conn.execute( + "UPDATE artifacts SET state=?,generation=generation+1,failures=?," + "last_error=?,priority=?,queued_at=?,claim_owner=NULL," + "claim_until=NULL,updated_at=? WHERE request_id=?", + ( + ArtifactState.QUEUED.value if retry else ArtifactState.FAILED.value, + failures, + message, + int(priority) if priority is not None else None, + self._clock() if retry else None, + self._clock(), + claim.request_id, + ), + ) + + def begin_evictions(self, *, limit: int = 64) -> tuple[EvictionClaim, ...]: + if limit < 1: + raise ValueError("eviction limit must be at least one") + with self._transaction() as conn: + rows = conn.execute( + "SELECT * FROM artifacts a WHERE a.state=? " + "AND NOT EXISTS (SELECT 1 FROM interests i " + "WHERE i.request_id=a.request_id) " + "AND NOT EXISTS (SELECT 1 FROM acquisitions q " + "WHERE q.request_id=a.request_id) " + "ORDER BY a.updated_at LIMIT ?", + (ArtifactState.READY.value, limit), + ).fetchall() + claims: list[EvictionClaim] = [] + for row in rows: + if not row["path"]: + continue + result = conn.execute( + "UPDATE artifacts SET state=?,updated_at=? " + "WHERE request_id=? AND state=?", + ( + ArtifactState.EVICTING.value, + self._clock(), + row["request_id"], + ArtifactState.READY.value, + ), + ) + if result.rowcount == 1: + claims.append( + EvictionClaim( + request_id=row["request_id"], + generation=int(row["generation"]), + path=Path(row["path"]), + ) + ) + return tuple(claims) + + def finish_eviction(self, claim: EvictionClaim, *, removed: bool) -> None: + with self._transaction() as conn: + row = conn.execute( + "SELECT * FROM artifacts WHERE request_id=?", (claim.request_id,) + ).fetchone() + if row is None or (row["state"], int(row["generation"])) != ( + ArtifactState.EVICTING.value, + claim.generation, + ): + raise WindowedArtifactError( + f"stale eviction completion for {claim.request_id}" + ) + interested = conn.execute( + "SELECT 1 FROM interests WHERE request_id=? LIMIT 1", + (claim.request_id,), + ).fetchone() + if removed: + state = ( + ArtifactState.QUEUED.value + if interested + else ArtifactState.ABSENT.value + ) + priority = ( + self._retry_priority_locked(conn, claim.request_id) + if interested + else None + ) + conn.execute( + "UPDATE artifacts SET state=?,path=NULL,size_bytes=0,priority=?," + "queued_at=?,updated_at=? WHERE request_id=?", + ( + state, + int(priority) if priority is not None else None, + self._clock() if interested else None, + self._clock(), + claim.request_id, + ), + ) + else: + conn.execute( + "UPDATE artifacts SET state=?,updated_at=? WHERE request_id=?", + (ArtifactState.READY.value, self._clock(), claim.request_id), + ) + + def complete_consumer(self, consumer_id: str) -> None: + with self._transaction() as conn: + row = conn.execute( + "SELECT 1 FROM consumers WHERE consumer_id=?", (consumer_id,) + ).fetchone() + if row is None: + raise KeyError(f"unknown consumer {consumer_id!r}") + conn.execute("DELETE FROM acquisitions WHERE consumer_id=?", (consumer_id,)) + conn.execute("DELETE FROM interests WHERE consumer_id=?", (consumer_id,)) + conn.execute( + "DELETE FROM completed_positions WHERE consumer_id=?", (consumer_id,) + ) + conn.execute( + "UPDATE consumers SET state='completed',updated_at=? " + "WHERE consumer_id=?", + (self._clock(), consumer_id), + ) + self._prune_orphaned_locked(conn) + + def snapshot(self) -> dict[str, Any]: + with self._lock: + consumers = [ + dict(row) + for row in self._conn.execute( + "SELECT * FROM consumers ORDER BY consumer_id" + ).fetchall() + ] + artifact_states = { + row["state"]: int(row["count"]) + for row in self._conn.execute( + "SELECT state,COUNT(*) AS count FROM artifacts GROUP BY state" + ).fetchall() + } + totals = self._conn.execute( + "SELECT COUNT(*) AS count,COALESCE(SUM(size_bytes),0) AS bytes " + "FROM artifacts WHERE state IN (?,?)", + (ArtifactState.READY.value, ArtifactState.EVICTING.value), + ).fetchone() + inflight = int( + self._conn.execute("SELECT COUNT(*) FROM acquisitions").fetchone()[0] + ) + positions = int( + self._conn.execute("SELECT COUNT(*) FROM positions").fetchone()[0] + ) + high_water = { + row["key"].removeprefix("peak_"): int(row["value"]) + for row in self._conn.execute( + "SELECT key,value FROM coordinator_meta WHERE key LIKE 'peak_%'" + ).fetchall() + } + return { + "schema_version": SCHEMA_VERSION, + "positions": positions, + "consumers": consumers, + "artifact_states": artifact_states, + "retained_artifacts": int(totals["count"]), + "retained_bytes": int(totals["bytes"]), + "inflight_acquisitions": inflight, + "high_water": high_water, + } diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py index 2c61db537..82f1dda75 100644 --- a/src/speculators/train/data.py +++ b/src/speculators/train/data.py @@ -1,7 +1,9 @@ +import hashlib import json import math import os import random +import threading import uuid import warnings from collections.abc import Callable @@ -18,6 +20,7 @@ from hs_connectors import FileTransfer, HiddenStatesTransfer from speculators.data_generation.artifact_cache import ( + ArtifactCacheError, HiddenStateArtifactCache, canonical_hidden_state_request_id, ) @@ -28,9 +31,18 @@ ClientItem, generate_hidden_states, ) +from speculators.data_generation.windowed_artifacts import ( + ArtifactReadLease, + GenerationClaim, + StreamSampleIndex, + WindowedArtifactCoordinator, + canonical_stream_id, +) from speculators.train.noise_transforms import TransformTensors BatchType = dict[str, Any] +WINDOWED_LEASE_KEY = "_windowed_artifact_lease" +WINDOWED_BATCH_LEASES_KEY = "_windowed_artifact_leases" def list_files(path): @@ -249,6 +261,13 @@ def __init__( shared_artifacts_namespace: str | None = None, shared_artifacts_ttl_seconds: float | None = 3600.0, shared_artifacts_lock_timeout_seconds: float = 300.0, + shared_artifacts_consumer_id: str | None = None, + shared_artifacts_lookbehind: int = 2, + shared_artifacts_lookahead: int = 16, + shared_artifacts_max_inflight: int = 32, + shared_artifacts_consumer_timeout_seconds: float = 120.0, + shared_artifacts_claim_timeout_seconds: float = 300.0, + shared_artifacts_generation_attempts: int = 3, ): self.data = load_from_disk(datapath) self.start_file_idx = 0 @@ -282,13 +301,49 @@ def __init__( self.artifact_cache = ( HiddenStateArtifactCache( shared_artifacts_path, - artifact_ttl_seconds=shared_artifacts_ttl_seconds, + artifact_ttl_seconds=( + None + if shared_artifacts_consumer_id is not None + else shared_artifacts_ttl_seconds + ), lock_timeout_seconds=shared_artifacts_lock_timeout_seconds, ) if shared_artifacts_path is not None else None ) - if self.artifact_cache is not None: + self.shared_artifacts_path = ( + Path(shared_artifacts_path).expanduser().resolve() + if shared_artifacts_path is not None + else None + ) + self.shared_artifacts_consumer_id = shared_artifacts_consumer_id + self.shared_artifacts_lock_timeout_seconds = ( + shared_artifacts_lock_timeout_seconds + ) + self.shared_artifacts_lookbehind = shared_artifacts_lookbehind + self.shared_artifacts_lookahead = shared_artifacts_lookahead + self.shared_artifacts_max_inflight = shared_artifacts_max_inflight + self.shared_artifacts_consumer_timeout_seconds = ( + shared_artifacts_consumer_timeout_seconds + ) + self.shared_artifacts_claim_timeout_seconds = ( + shared_artifacts_claim_timeout_seconds + ) + self.shared_artifacts_generation_attempts = shared_artifacts_generation_attempts + self.windowed_artifacts_enabled = shared_artifacts_consumer_id is not None + if self.windowed_artifacts_enabled and self.artifact_cache is None: + raise ValueError( + "shared_artifacts_consumer_id requires shared_artifacts_path" + ) + if self.windowed_artifacts_enabled and self.on_missing != "generate": + raise ValueError("windowed artifacts require on_missing='generate'") + self._windowed_stream_id: str | None = None + self._windowed_coordinator: WindowedArtifactCoordinator | None = None + self._windowed_coordinator_pid: int | None = None + self._windowed_producer_thread: threading.Thread | None = None + self._windowed_producer_stop: threading.Event | None = None + self._windowed_producer_error: Exception | None = None + if self.artifact_cache is not None and not self.windowed_artifacts_enabled: self.artifact_cache.cleanup_stale() # Delay super init so that `_compute_approx_lengths` has required data @@ -315,6 +370,281 @@ def _setup_client(self): def __len__(self): return len(self.data) + def _new_windowed_coordinator(self) -> WindowedArtifactCoordinator: + if self.shared_artifacts_path is None: + raise RuntimeError("windowed artifacts are not configured") + return WindowedArtifactCoordinator( + self.shared_artifacts_path, + consumer_timeout_seconds=self.shared_artifacts_consumer_timeout_seconds, + claim_timeout_seconds=self.shared_artifacts_claim_timeout_seconds, + max_generation_attempts=self.shared_artifacts_generation_attempts, + ) + + def _coordinator_for_process(self) -> WindowedArtifactCoordinator: + pid = os.getpid() + if self._windowed_coordinator_pid != pid: + if self._windowed_coordinator is not None: + self._windowed_coordinator.close() + self._windowed_coordinator = self._new_windowed_coordinator() + self._windowed_coordinator_pid = pid + if self._windowed_coordinator is None: + raise RuntimeError("windowed artifact coordinator is unavailable") + return self._windowed_coordinator + + def configure_windowed_stream(self, sampler: Any) -> str: + """Bind this dataset split to a deterministic sampler-order contract.""" + if not self.windowed_artifacts_enabled: + raise RuntimeError("windowed artifacts are not enabled") + lengths_digest = hashlib.sha256() + for length in sampler.lengths: + lengths_digest.update(f"{int(length)}\n".encode()) + contract = { + "batch_max_length": int(sampler.batch_max_length), + "dataset_fingerprint": str( + getattr(self.data, "_fingerprint", "unavailable") + ), + "dataset_length": len(self.data), + "dp_rank": int(sampler.rank), + "dp_size": int(sampler.num_replicas), + "lengths_digest": lengths_digest.hexdigest(), + "namespace": self.shared_artifacts_namespace, + "order": "MultipackDistributedBatchSamplerV2", + "sampler_seed": int(sampler.seed), + "schema_version": 1, + "verifier_model": self.model, + } + stream_id = canonical_stream_id(contract) + with self._new_windowed_coordinator() as coordinator: + observed = coordinator.register_stream(contract) + if observed != stream_id: + raise RuntimeError("coordinator returned an inconsistent stream identity") + self._windowed_stream_id = stream_id + return stream_id + + def windowed_request_id(self, dataset_index: int) -> str: + if not self.model: + raise RuntimeError("windowed artifacts require an explicit verifier model") + dataset_item = self.data[dataset_index] + return canonical_hidden_state_request_id( + self.model, + build_client_item(dataset_item), + namespace=self.shared_artifacts_namespace, + ) + + def prepare_windowed_epoch( + self, + samples: tuple[StreamSampleIndex, ...], + *, + cursor: int, + reset: bool, + ) -> None: + if not self.windowed_artifacts_enabled: + return + if ( + self._windowed_stream_id is None + or self.shared_artifacts_consumer_id is None + ): + raise RuntimeError("windowed stream was not configured by the DataLoader") + if samples and any( + sample.stream_id != self._windowed_stream_id for sample in samples + ): + raise RuntimeError("sampler positions belong to another stream") + with self._new_windowed_coordinator() as coordinator: + coordinator.register_positions(samples) + coordinator.register_consumer( + self.shared_artifacts_consumer_id, + stream_id=self._windowed_stream_id, + lookbehind=self.shared_artifacts_lookbehind, + lookahead=self.shared_artifacts_lookahead, + max_inflight=self.shared_artifacts_max_inflight, + cursor=cursor, + reset=reset, + ) + + def _acquire_windowed_hs( + self, sample: StreamSampleIndex + ) -> tuple[dict[str, torch.Tensor], ArtifactReadLease]: + if self.shared_artifacts_consumer_id is None or self.artifact_cache is None: + raise RuntimeError("windowed artifacts are not configured") + coordinator = self._coordinator_for_process() + lease = coordinator.acquire( + self.shared_artifacts_consumer_id, + sample, + timeout_seconds=self.request_timeout, + ) + dataset_item = self.data[sample.dataset_index] + try: + loaded = self.artifact_cache.load( + sample.request_id, + lambda data: check_hidden_states( + data, dataset_item["input_ids"].tolist() + ), + ) + if lease.cache_hit: + self.artifact_cache.record_reuse() + return loaded, lease + except BaseException: + coordinator.abandon( + self.shared_artifacts_consumer_id, [lease.as_batch_metadata()] + ) + raise + + def ack_windowed_batch(self, leases: list[dict[str, Any]]) -> int | None: + if not self.windowed_artifacts_enabled or not leases: + return None + if self.shared_artifacts_consumer_id is None: + raise RuntimeError("windowed consumer identity is missing") + return self._coordinator_for_process().ack( + self.shared_artifacts_consumer_id, leases + ) + + def abandon_windowed_batch(self, leases: list[dict[str, Any]]) -> None: + if not self.windowed_artifacts_enabled or not leases: + return + if self.shared_artifacts_consumer_id is None: + raise RuntimeError("windowed consumer identity is missing") + self._coordinator_for_process().abandon( + self.shared_artifacts_consumer_id, leases + ) + + def start_windowed_producer(self) -> None: + """Start a trainer-main-process dispatcher after DataLoader workers fork.""" + if not self.windowed_artifacts_enabled: + return + if self._windowed_producer_thread is not None: + if not self._windowed_producer_thread.is_alive(): + error = self._windowed_producer_error + raise RuntimeError("windowed artifact producer stopped") from error + return + if self._windowed_stream_id is None: + raise RuntimeError("windowed stream is not prepared") + self._windowed_producer_stop = threading.Event() + self._windowed_producer_error = None + self._windowed_producer_thread = threading.Thread( + target=self._run_windowed_producer, + name=f"artifact-producer-{self.shared_artifacts_consumer_id}", + daemon=True, + ) + self._windowed_producer_thread.start() + + def _run_windowed_producer(self) -> None: + if ( + self.shared_artifacts_path is None + or self.shared_artifacts_consumer_id is None + or self._windowed_stream_id is None + or self._windowed_producer_stop is None + ): + self._windowed_producer_error = RuntimeError( + "windowed producer started without a complete configuration" + ) + return + owner = f"{self.shared_artifacts_consumer_id}:{os.getpid()}:{uuid.uuid4().hex}" + cache = HiddenStateArtifactCache( + self.shared_artifacts_path, + artifact_ttl_seconds=None, + lock_timeout_seconds=self.shared_artifacts_lock_timeout_seconds, + ) + try: + with self._new_windowed_coordinator() as coordinator: + while not self._windowed_producer_stop.is_set(): + coordinator.heartbeat(self.shared_artifacts_consumer_id) + coordinator.recover_expired() + self._evict_windowed_artifacts(coordinator, cache) + claims = coordinator.claim_generation( + owner, stream_id=self._windowed_stream_id, max_claims=1 + ) + if claims: + for claim in claims: + try: + self._produce_windowed_claim( + coordinator, cache, owner, claim + ) + except Exception as error: # noqa: BLE001 + coordinator.fail_generation(owner, claim, error) + continue + self._windowed_producer_stop.wait(0.02) + except Exception as error: # noqa: BLE001 - background thread boundary + self._windowed_producer_error = error + + def _produce_windowed_claim( + self, + coordinator: WindowedArtifactCoordinator, + cache: HiddenStateArtifactCache, + owner: str, + claim: GenerationClaim, + ) -> None: + expected_request_id = self.windowed_request_id(claim.dataset_index) + if expected_request_id != claim.request_id: + raise RuntimeError("generation claim no longer matches dataset") + dataset_item = self.data[claim.dataset_index] + client_item = build_client_item(dataset_item) + result = cache.get_or_create( + claim.request_id, + lambda: self._materialize_shared_hs( + claim.dataset_index, + dataset_item, + client_item, + ), + lambda data: check_hidden_states(data, dataset_item["input_ids"].tolist()), + ) + coordinator.complete_generation( + owner, + claim, + path=result.path, + size_bytes=result.path.stat().st_size, + ) + + @staticmethod + def _evict_windowed_artifacts( + coordinator: WindowedArtifactCoordinator, + cache: HiddenStateArtifactCache, + ) -> None: + for eviction in coordinator.begin_evictions(limit=16): + try: + removed = cache.remove(eviction.request_id, expected_path=eviction.path) + except (ArtifactCacheError, OSError): + coordinator.finish_eviction(eviction, removed=False) + else: + coordinator.finish_eviction( + eviction, removed=removed or not eviction.path.exists() + ) + + def _materialize_shared_hs( + self, + index: int, + dataset_item: dict, + client_item: ClientItem, + ) -> dict[str, torch.Tensor]: + file_idx = self._map_to_file_idx(index) + cached = self.transfer.get_cached(file_idx) + if cached is not None: + check_hidden_states(cached, dataset_item["input_ids"].tolist()) + return cached + if not self.client: + self._setup_client() + return self._generate_shared_hs(dataset_item, client_item) + + def stop_windowed_producer(self, *, completed: bool = False) -> None: + stop = self._windowed_producer_stop + thread = self._windowed_producer_thread + if stop is not None: + stop.set() + if thread is not None: + thread.join(timeout=30.0) + if thread.is_alive(): + raise RuntimeError("windowed artifact producer did not stop") + self._windowed_producer_thread = None + self._windowed_producer_stop = None + if completed and self.shared_artifacts_consumer_id is not None: + coordinator = self._coordinator_for_process() + coordinator.complete_consumer(self.shared_artifacts_consumer_id) + if self.artifact_cache is not None: + self._evict_windowed_artifacts(coordinator, self.artifact_cache) + if self._windowed_producer_error is not None: + error = self._windowed_producer_error + self._windowed_producer_error = None + raise RuntimeError("windowed artifact producer failed") from error + def _compute_approx_lengths(self) -> list[int]: """Get lengths of the dataset samples.""" return list(self.data.with_format(None)["seq_len"]) @@ -408,44 +738,71 @@ def _generate_shared_hs( if handle is not None and retrieved: self.transfer.delete(handle) - def _get_raw_data(self, index): - file_idx = self._map_to_file_idx(index) - loaded_hs = self.transfer.get_cached(file_idx) + def _load_requested_hidden_states( + self, + dataset_index: int, + windowed_sample: StreamSampleIndex | None, + ) -> tuple[dict[str, torch.Tensor] | None, ArtifactReadLease | None]: + lease: ArtifactReadLease | None = None + file_idx = self._map_to_file_idx(dataset_index) + if windowed_sample is not None: + loaded_hs, lease = self._acquire_windowed_hs(windowed_sample) + else: + loaded_hs = self.transfer.get_cached(file_idx) if loaded_hs is None: match self.on_missing: case "generate": - loaded_hs = self._maybe_generate_hs(index) + loaded_hs = self._maybe_generate_hs(dataset_index) case "skip": - return None + return None, None case "warn": warnings.warn( - f"Failed to load hidden states for sample {index}. Skipping...", + "Failed to load hidden states for sample " + f"{dataset_index}. Skipping...", stacklevel=1, ) - return None + return None, None case "raise": raise RuntimeError( - f"Failed to load hidden states for sample {index}." + f"Failed to load hidden states for sample {dataset_index}." ) + return loaded_hs, lease + + def _get_raw_data(self, index): + windowed_sample = index if isinstance(index, StreamSampleIndex) else None + dataset_index = ( + windowed_sample.dataset_index if windowed_sample is not None else int(index) + ) + loaded_hs, lease = self._load_requested_hidden_states( + dataset_index, windowed_sample + ) if loaded_hs is None: - return loaded_hs + return None # loaded_hs structure: { # "hidden_states": [seq_len, num_layers, hidden_size] # "token_ids": [seq_len] # } - if not torch.equal(loaded_hs["token_ids"], self.data[index]["input_ids"]): + if not torch.equal( + loaded_hs["token_ids"], self.data[dataset_index]["input_ids"] + ): warnings.warn( - f"Loaded token ids {loaded_hs['token_ids']} for index {index} don't" - f"match input ids {self.data[index]['input_ids']}", + f"Loaded token ids {loaded_hs['token_ids']} for index " + f"{dataset_index} don't match input ids " + f"{self.data[dataset_index]['input_ids']}", stacklevel=1, ) + if lease is not None and self.shared_artifacts_consumer_id is not None: + self._coordinator_for_process().abandon( + self.shared_artifacts_consumer_id, + [lease.as_batch_metadata()], + ) return None - return { + result = { "hidden_states": loaded_hs["hidden_states"][:, :-1].flatten( 1 ), # [seq_len, 3 * hidden_size] @@ -453,8 +810,11 @@ def _get_raw_data(self, index): "verifier_last_hidden_states": loaded_hs["hidden_states"][ :, -1 ], # [seq_len, hidden_size] - "loss_mask": self.data[index]["loss_mask"], # [seq_len] + "loss_mask": self.data[dataset_index]["loss_mask"], # [seq_len] } + if lease is not None: + result[WINDOWED_LEASE_KEY] = lease.as_batch_metadata() + return result class SampleFileDataset(BaseDataset): @@ -562,8 +922,15 @@ def create_collate_fn( preprocess: Callable[[BatchType], BatchType] | None = None, ): def collate_fn(batch: list[BatchType | None]) -> BatchType: + # Lease metadata stays on CPU and is never passed through model preprocessing. + valid_batch = [sample for sample in batch if sample is not None] + leases = [ + sample.pop(WINDOWED_LEASE_KEY) + for sample in valid_batch + if WINDOWED_LEASE_KEY in sample + ] # Apply per-sample preprocessing and filter failed samples - batch = [preprocess(b) if preprocess else b for b in batch if b is not None] + batch = [preprocess(b) if preprocess else b for b in valid_batch] if not batch: # Create empty sample which then gets padded to full @@ -615,6 +982,8 @@ def collate_fn(batch: list[BatchType | None]) -> BatchType: ).unsqueeze(0) # shape: [1, max_len] collated_data["document_ids"] = document_ids + if leases: + collated_data[WINDOWED_BATCH_LEASES_KEY] = leases return collated_data diff --git a/src/speculators/train/dataloader.py b/src/speculators/train/dataloader.py index 001d6b696..dcb4eae8e 100644 --- a/src/speculators/train/dataloader.py +++ b/src/speculators/train/dataloader.py @@ -2,7 +2,7 @@ import logging import warnings -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, Protocol if TYPE_CHECKING: from collections.abc import Callable @@ -11,6 +11,10 @@ from torch.utils.data import DataLoader from hs_connectors import HiddenStatesTransfer +from speculators.data_generation.windowed_artifacts import ( + StreamSampleIndex, + canonical_position_id, +) from speculators.train.data import ( ArrowDataset, BaseDataset, @@ -29,6 +33,103 @@ BatchType = dict[str, Any] +class _WindowedDataset(Protocol): + windowed_artifacts_enabled: bool + + def configure_windowed_stream(self, sampler: Any) -> str: ... + + def windowed_request_id(self, dataset_index: int) -> str: ... + + +class WindowedBatchSampler: + """Annotate sampler indices with stable positions in the consumed order.""" + + def __init__( + self, + sampler: MultipackDistributedBatchSamplerV2, + *, + stream_id: str, + request_id_for_index: Callable[[int], str], + ) -> None: + self.sampler = sampler + self.stream_id = stream_id + self.request_id_for_index = request_id_for_index + self.epoch = sampler.epoch + self._epoch_counts: dict[int, int] = {} + self._cached_generated_batches: tuple[int, list[list[StreamSampleIndex]]] = ( + -1, + [], + ) + + def _raw_batches(self, epoch: int) -> list[Any]: + return self.sampler._generate_batches(epoch) # noqa: SLF001 + + def _epoch_count(self, epoch: int) -> int: + if epoch not in self._epoch_counts: + self._epoch_counts[epoch] = sum( + len(batch) for batch in self._raw_batches(epoch) + ) + return self._epoch_counts[epoch] + + def _sequence_offset(self, epoch: int) -> int: + return sum(self._epoch_count(previous) for previous in range(epoch)) + + def _generate_batches(self, epoch: int) -> list[list[StreamSampleIndex]]: + if self._cached_generated_batches[0] == epoch: + return self._cached_generated_batches[1] + offset = self._sequence_offset(epoch) + ordinal = 0 + batches: list[list[StreamSampleIndex]] = [] + for batch_ordinal, raw_batch in enumerate(self._raw_batches(epoch)): + batch: list[StreamSampleIndex] = [] + batch_start_sequence = offset + ordinal + batch_end_sequence = batch_start_sequence + len(raw_batch) + for raw_index in raw_batch: + dataset_index = int(raw_index) + batch.append( + StreamSampleIndex( + stream_id=self.stream_id, + sequence=offset + ordinal, + epoch=epoch, + ordinal=ordinal, + dataset_index=dataset_index, + batch_ordinal=batch_ordinal, + batch_start_sequence=batch_start_sequence, + batch_end_sequence=batch_end_sequence, + request_id=self.request_id_for_index(dataset_index), + position_id=canonical_position_id( + self.stream_id, + epoch=epoch, + ordinal=ordinal, + dataset_index=dataset_index, + batch_ordinal=batch_ordinal, + batch_start_sequence=batch_start_sequence, + batch_end_sequence=batch_end_sequence, + ), + ) + ) + ordinal += 1 + batches.append(batch) + self._epoch_counts[epoch] = ordinal + self._cached_generated_batches = (epoch, batches) + return batches + + def full_epoch_samples(self, epoch: int) -> tuple[StreamSampleIndex, ...]: + return tuple( + sample for batch in self._generate_batches(epoch) for sample in batch + ) + + def __iter__(self): + return iter(self._generate_batches(self.epoch)) + + def __len__(self) -> int: + return len(self._generate_batches(self.epoch)) + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + self.sampler.set_epoch(epoch) + + def _setup_dataloader( dataset: BaseDataset, total_seq_len: int, @@ -38,12 +139,20 @@ def _setup_dataloader( prefetch_factor: int | None = 4, preprocess: Callable[[BatchType], BatchType] | None = None, ) -> DataLoader: - batch_sampler = MultipackDistributedBatchSamplerV2( + batch_sampler: Any = MultipackDistributedBatchSamplerV2( batch_max_length=total_seq_len, lengths=dataset.approx_lengths, num_replicas=get_dp_size(), rank=get_dp_rank(), ) + if getattr(dataset, "windowed_artifacts_enabled", False): + windowed_dataset: _WindowedDataset = dataset # type: ignore[assignment] + stream_id = windowed_dataset.configure_windowed_stream(batch_sampler) + batch_sampler = WindowedBatchSampler( + batch_sampler, + stream_id=stream_id, + request_id_for_index=windowed_dataset.windowed_request_id, + ) use_workers = num_workers > 0 return DataLoader( dataset, @@ -85,6 +194,13 @@ def create_train_val_loaders( shared_artifacts_namespace: str | None = None, shared_artifacts_ttl_seconds: float | None = 3600.0, shared_artifacts_lock_timeout_seconds: float = 300.0, + shared_artifacts_consumer_id: str | None = None, + shared_artifacts_lookbehind: int = 2, + shared_artifacts_lookahead: int = 16, + shared_artifacts_max_inflight: int = 32, + shared_artifacts_consumer_timeout_seconds: float = 120.0, + shared_artifacts_claim_timeout_seconds: float = 300.0, + shared_artifacts_generation_attempts: int = 3, train_data_ratio: float = 0.9, ) -> tuple[DataLoader, DataLoader]: """Create training and validation DataLoaders. @@ -137,6 +253,21 @@ def create_train_val_loaders( shared_artifacts_lock_timeout_seconds=( shared_artifacts_lock_timeout_seconds ), + shared_artifacts_consumer_id=( + f"{shared_artifacts_consumer_id}:train" + if shared_artifacts_consumer_id is not None + else None + ), + shared_artifacts_lookbehind=shared_artifacts_lookbehind, + shared_artifacts_lookahead=shared_artifacts_lookahead, + shared_artifacts_max_inflight=shared_artifacts_max_inflight, + shared_artifacts_consumer_timeout_seconds=( + shared_artifacts_consumer_timeout_seconds + ), + shared_artifacts_claim_timeout_seconds=( + shared_artifacts_claim_timeout_seconds + ), + shared_artifacts_generation_attempts=(shared_artifacts_generation_attempts), ) val_dataset = ArrowDataset( datapath=data_path, @@ -156,6 +287,21 @@ def create_train_val_loaders( shared_artifacts_lock_timeout_seconds=( shared_artifacts_lock_timeout_seconds ), + shared_artifacts_consumer_id=( + f"{shared_artifacts_consumer_id}:val" + if shared_artifacts_consumer_id is not None + else None + ), + shared_artifacts_lookbehind=shared_artifacts_lookbehind, + shared_artifacts_lookahead=shared_artifacts_lookahead, + shared_artifacts_max_inflight=shared_artifacts_max_inflight, + shared_artifacts_consumer_timeout_seconds=( + shared_artifacts_consumer_timeout_seconds + ), + shared_artifacts_claim_timeout_seconds=( + shared_artifacts_claim_timeout_seconds + ), + shared_artifacts_generation_attempts=(shared_artifacts_generation_attempts), ) train_loader = _setup_dataloader( diff --git a/src/speculators/train/trainer.py b/src/speculators/train/trainer.py index d9cfe2b5c..3c0d90988 100644 --- a/src/speculators/train/trainer.py +++ b/src/speculators/train/trainer.py @@ -2,8 +2,9 @@ import logging import time import warnings +from collections.abc import Callable from pathlib import Path -from typing import Literal, NamedTuple +from typing import Literal, NamedTuple, TypeVar import torch import torch.distributed as dist @@ -26,6 +27,7 @@ DistributedCheckpointer, SingleGPUCheckpointer, ) +from speculators.train.data import WINDOWED_BATCH_LEASES_KEY from speculators.train.distributed import ( apply_fully_sharded, get_local_rank, @@ -38,6 +40,7 @@ root_logger = logging.getLogger("speculators") metric_logger = logging.getLogger("speculators.metrics") +_T = TypeVar("_T") class _StepTimer: @@ -186,6 +189,68 @@ def __init__( self.setup_trainer() self.setup_model() self.setup_optimizer() + self._prepared_windowed_datasets: set[int] = set() + + def _prepare_windowed_loader( + self, + loader: DataLoader, + *, + epoch: int, + skip_steps: int, + full_batches: list | None = None, + ): + sampler = loader.batch_sampler + dataset = loader.dataset + if not hasattr(sampler, "full_epoch_samples") or not hasattr( + dataset, "prepare_windowed_epoch" + ): + return iter(loader) + batches = ( + sampler._generate_batches(epoch) # type: ignore[union-attr] # noqa: SLF001 + if full_batches is None + else full_batches + ) + samples = tuple(sample for batch in batches for sample in batch) + if skip_steps < len(batches) and batches[skip_steps]: + cursor = batches[skip_steps][0].sequence + elif samples: + cursor = samples[-1].sequence + 1 + else: + cursor = 0 + dataset_id = id(dataset) + dataset.prepare_windowed_epoch( # type: ignore[union-attr] + samples, + cursor=cursor, + reset=dataset_id not in self._prepared_windowed_datasets, + ) + self._prepared_windowed_datasets.add(dataset_id) + iterator = iter(loader) + dataset.start_windowed_producer() # type: ignore[union-attr] + return iterator + + @staticmethod + def _run_windowed_phase( + loader: DataLoader, operation: Callable[[int], _T], epoch: int + ) -> _T: + completed = False + try: + result = operation(epoch) + completed = True + return result + finally: + dataset = loader.dataset + if hasattr(dataset, "stop_windowed_producer"): + dataset.stop_windowed_producer(completed=completed) + + @staticmethod + def _ack_windowed_batch(dataset, leases: list[dict]) -> None: + if leases and hasattr(dataset, "ack_windowed_batch"): + dataset.ack_windowed_batch(leases) + + @staticmethod + def _abandon_windowed_batch(dataset, leases: list[dict]) -> None: + if leases and hasattr(dataset, "abandon_windowed_batch"): + dataset.abandon_windowed_batch(leases) def _training_state_path(self, epoch: int) -> Path: return self.checkpointer.path / str(epoch) / "training_state.json" @@ -423,11 +488,22 @@ def train_epoch(self, epoch: int): # Capture full-epoch step count before any resume fast-skip mutation. num_steps = len(self.train_loader) + sampler = self.train_loader.batch_sampler + full_batches = ( + sampler._generate_batches(epoch) # type: ignore[union-attr] # noqa: SLF001 + if hasattr(sampler, "full_epoch_samples") + else None + ) # Determine how many batches to skip for mid-epoch resume. skip_steps = self._prepare_resume_skip(epoch) - train_loader = self.train_loader + train_loader = self._prepare_windowed_loader( + self.train_loader, + epoch=epoch, + skip_steps=skip_steps, + full_batches=full_batches, + ) if self.rank == 0: train_loader = tqdm(train_loader, desc=f"Epoch {epoch}") # type: ignore[assignment] @@ -444,34 +520,41 @@ def train_epoch(self, epoch: int): timer.reset(self.global_step % self.config.log_freq == 0) timer.mark_value("start", t_before_fetch) - gpu_batch = { - k: v.to(self.local_rank, non_blocking=True) - if isinstance(v, torch.Tensor) - else v - for k, v in batch.items() - } - - with torch.autocast( - self.device_type, dtype=self.config.hidden_states_dtype - ): - timer.mark("fetch") - _draft_tokens, loss, metrics = self.model( - **gpu_batch, **(self.config.train_call_kwargs or {}) - ) - - timer.mark("fwd") - self._optimizers_zero_grad() - loss.backward() - torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) - - timer.mark("bwd") - self._optimizers_step() + leases = batch.pop(WINDOWED_BATCH_LEASES_KEY, []) + try: + gpu_batch = { + k: v.to(self.local_rank, non_blocking=True) + if isinstance(v, torch.Tensor) + else v + for k, v in batch.items() + } + + with torch.autocast( + self.device_type, dtype=self.config.hidden_states_dtype + ): + timer.mark("fetch") + _draft_tokens, loss, metrics = self.model( + **gpu_batch, **(self.config.train_call_kwargs or {}) + ) - current_lrs = { - type(opt).__name__: opt.param_groups[0]["lr"] for opt in self.optimizers - } - self._schedulers_step() - timer.mark("opt") + timer.mark("fwd") + self._optimizers_zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) + + timer.mark("bwd") + self._optimizers_step() + + current_lrs = { + type(opt).__name__: opt.param_groups[0]["lr"] + for opt in self.optimizers + } + self._schedulers_step() + timer.mark("opt") + except BaseException: + self._abandon_windowed_batch(self.train_loader.dataset, leases) + raise + self._ack_windowed_batch(self.train_loader.dataset, leases) t_before_fetch = timer.now() or time.perf_counter() profile = None @@ -518,26 +601,43 @@ def val_epoch(self, epoch: int) -> dict[str, float] | None: self.model.eval() if hasattr(self.val_loader.batch_sampler, "set_epoch"): self.val_loader.batch_sampler.set_epoch(epoch) # type: ignore[union-attr] - val_loader = self.val_loader + val_sampler = self.val_loader.batch_sampler + full_batches = ( + val_sampler._generate_batches(epoch) # type: ignore[union-attr] # noqa: SLF001 + if hasattr(val_sampler, "full_epoch_samples") + else None + ) + val_loader = self._prepare_windowed_loader( + self.val_loader, + epoch=epoch, + skip_steps=0, + full_batches=full_batches, + ) if self.rank == 0: val_loader = tqdm(val_loader, desc=f"Epoch {epoch}") # type: ignore[assignment] val_metrics: dict[str, float] = {} - num_batches = len(val_loader) + num_batches = len(self.val_loader) for batch in val_loader: - gpu_batch = { - k: v.to(self.local_rank, non_blocking=True) - if isinstance(v, torch.Tensor) - else v - for k, v in batch.items() - } - - with torch.autocast( - self.device_type, dtype=self.config.hidden_states_dtype - ): - _draft_tokens, _loss, metrics = self.model( - **gpu_batch, **(self.config.val_call_kwargs or {}) - ) + leases = batch.pop(WINDOWED_BATCH_LEASES_KEY, []) + try: + gpu_batch = { + k: v.to(self.local_rank, non_blocking=True) + if isinstance(v, torch.Tensor) + else v + for k, v in batch.items() + } + + with torch.autocast( + self.device_type, dtype=self.config.hidden_states_dtype + ): + _draft_tokens, _loss, metrics = self.model( + **gpu_batch, **(self.config.val_call_kwargs or {}) + ) + except BaseException: + self._abandon_windowed_batch(self.val_loader.dataset, leases) + raise + self._ack_windowed_batch(self.val_loader.dataset, leases) if self.is_distributed: for m in metrics.values(): @@ -618,7 +718,7 @@ def run_training(self): n_epochs = self.config.num_epochs for epoch in range(self.current_epoch, n_epochs): root_logger.info(f"Training epoch {epoch + 1}/{n_epochs} started") - self.train_epoch(epoch) + self._run_windowed_phase(self.train_loader, self.train_epoch, epoch) root_logger.info(f"Training epoch {epoch + 1}/{n_epochs} completed") if self.is_distributed: @@ -635,7 +735,9 @@ def run_training(self): root_logger.warning("No val loader, skipping validation epoch") else: root_logger.info(f"Validation epoch {epoch + 1}/{n_epochs} started") - val_metrics = self.val_epoch(epoch) + val_metrics = self._run_windowed_phase( + self.val_loader, self.val_epoch, epoch + ) root_logger.info(f"Validation epoch {epoch + 1}/{n_epochs} completed") if self.is_distributed: @@ -645,3 +747,9 @@ def run_training(self): if self.is_distributed: dist.barrier() + + for loader in (self.train_loader, self.val_loader): + if loader is not None and hasattr(loader.dataset, "stop_windowed_producer"): + loader.dataset.stop_windowed_producer( # type: ignore[union-attr] + completed=True + ) diff --git a/tests/unit/data_generation/test_windowed_artifacts.py b/tests/unit/data_generation/test_windowed_artifacts.py new file mode 100644 index 000000000..87cb0f327 --- /dev/null +++ b/tests/unit/data_generation/test_windowed_artifacts.py @@ -0,0 +1,499 @@ +from __future__ import annotations + +import hashlib +import threading +import time +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from pathlib import Path + +from speculators.data_generation.windowed_artifacts import ( + ArtifactGenerationError, + ArtifactPriority, + StreamSampleIndex, + WindowedArtifactCoordinator, + WindowedArtifactError, + canonical_position_id, + canonical_stream_id, +) + + +def _digest(value: str) -> str: + return hashlib.sha256(value.encode()).hexdigest() + + +def _samples(stream_id: str, count: int) -> tuple[StreamSampleIndex, ...]: + return tuple( + StreamSampleIndex( + stream_id=stream_id, + sequence=index, + epoch=0, + ordinal=index, + dataset_index=index, + batch_ordinal=index, + batch_start_sequence=index, + batch_end_sequence=index + 1, + request_id=_digest(f"request-{index}"), + position_id=canonical_position_id( + stream_id, + epoch=0, + ordinal=index, + dataset_index=index, + batch_ordinal=index, + batch_start_sequence=index, + batch_end_sequence=index + 1, + ), + ) + for index in range(count) + ) + + +def _single_batch_samples(stream_id: str, count: int) -> tuple[StreamSampleIndex, ...]: + return tuple( + StreamSampleIndex( + stream_id=stream_id, + sequence=index, + epoch=0, + ordinal=index, + dataset_index=index, + batch_ordinal=0, + batch_start_sequence=0, + batch_end_sequence=count, + request_id=_digest(f"request-{index}"), + position_id=canonical_position_id( + stream_id, + epoch=0, + ordinal=index, + dataset_index=index, + batch_ordinal=0, + batch_start_sequence=0, + batch_end_sequence=count, + ), + ) + for index in range(count) + ) + + +def _coordinator(tmp_path: Path, **kwargs) -> WindowedArtifactCoordinator: + return WindowedArtifactCoordinator( + tmp_path, + poll_seconds=0.005, + consumer_timeout_seconds=10, + claim_timeout_seconds=10, + **kwargs, + ) + + +def _register( + coordinator: WindowedArtifactCoordinator, + samples: tuple[StreamSampleIndex, ...], + *consumer_ids: str, + contract: dict | None = None, + lookbehind: int = 0, + lookahead: int = 2, + max_inflight: int = 4, +) -> None: + contract = contract or {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + assert coordinator.register_stream(contract) == samples[0].stream_id + coordinator.register_positions(samples) + for consumer_id in consumer_ids: + coordinator.register_consumer( + consumer_id, + stream_id=samples[0].stream_id, + lookbehind=lookbehind, + lookahead=lookahead, + max_inflight=max_inflight, + ) + + +def _publish( + coordinator: WindowedArtifactCoordinator, + stream_id: str, + path: Path, +) -> str: + claim = coordinator.claim_generation("producer", stream_id=stream_id)[0] + artifact = path / f"{claim.request_id}.safetensors" + artifact.write_bytes(b"payload") + coordinator.complete_generation( + "producer", claim, path=artifact, size_bytes=artifact.stat().st_size + ) + return claim.request_id + + +def _wait_until(predicate, timeout: float = 2.0) -> None: + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() >= deadline: + raise TimeoutError("condition was not reached") + time.sleep(0.005) + + +def test_stream_and_position_identity_cover_order_contract(): + contract = { + "dataset_fingerprint": "dataset-a", + "epoch_order": "multipack-v2", + "sampler_seed": 7, + } + stream_id = canonical_stream_id(contract) + + assert stream_id == canonical_stream_id(dict(reversed(tuple(contract.items())))) + assert stream_id != canonical_stream_id({**contract, "sampler_seed": 8}) + assert canonical_position_id( + stream_id, + epoch=2, + ordinal=3, + dataset_index=9, + batch_ordinal=1, + batch_start_sequence=3, + batch_end_sequence=5, + ) != canonical_position_id( + stream_id, + epoch=2, + ordinal=4, + dataset_index=9, + batch_ordinal=1, + batch_start_sequence=3, + batch_end_sequence=5, + ) + + +def test_existing_schema_connection_does_not_wait_for_writer(tmp_path): + owner = _coordinator(tmp_path) + owner._conn.execute("BEGIN IMMEDIATE") + finished = threading.Event() + errors: list[BaseException] = [] + + def connect() -> None: + try: + with _coordinator(tmp_path): + pass + except BaseException as error: # noqa: BLE001 - thread boundary + errors.append(error) + finally: + finished.set() + + thread = threading.Thread(target=connect) + thread.start() + try: + assert finished.wait(0.5), "existing schema initialization attempted a write" + finally: + owner._conn.rollback() + owner.close() + thread.join(timeout=2) + assert not errors + + +def test_completed_consumer_reactivates_at_its_committed_cursor(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 1) + coordinator = _coordinator(tmp_path) + _register(coordinator, samples, "consumer") + + coordinator.complete_consumer("consumer") + coordinator.register_consumer( + "consumer", + stream_id=stream_id, + lookbehind=0, + lookahead=2, + max_inflight=4, + cursor=0, + ) + + assert coordinator.snapshot()["consumers"][0]["state"] == "active" + + +def test_two_consumers_share_publication_but_commit_independently(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 3) + coordinator = _coordinator(tmp_path) + _register(coordinator, samples, "consumer-a", "consumer-b") + _publish(coordinator, stream_id, tmp_path) + + first = coordinator.acquire("consumer-a", samples[0], timeout_seconds=1) + second = coordinator.acquire("consumer-b", samples[0], timeout_seconds=1) + assert not first.cache_hit + assert second.cache_hit + + assert coordinator.ack("consumer-a", [first.as_batch_metadata()]) == 1 + snapshot = coordinator.snapshot() + cursors = {row["consumer_id"]: row["cursor"] for row in snapshot["consumers"]} + assert cursors == {"consumer-a": 1, "consumer-b": 0} + assert snapshot["high_water"] == { + "inflight_acquisitions": 2, + "retained_artifacts": 1, + "retained_bytes": len(b"payload"), + } + assert coordinator.begin_evictions() == () + + assert coordinator.ack("consumer-b", [second.as_batch_metadata()]) == 1 + evictions = coordinator.begin_evictions() + assert [claim.request_id for claim in evictions] == [samples[0].request_id] + + +def test_max_inflight_is_independent_of_dataloader_prefetch(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 3) + coordinator = _coordinator(tmp_path) + _register( + coordinator, + samples, + "consumer", + lookahead=2, + max_inflight=1, + ) + while coordinator.snapshot()["artifact_states"].get("queued", 0): + _publish(coordinator, stream_id, tmp_path) + + first = coordinator.acquire("consumer", samples[0], timeout_seconds=1) + acquired: list = [] + thread = threading.Thread( + target=lambda: acquired.append( + coordinator.acquire("consumer", samples[1], timeout_seconds=2) + ) + ) + thread.start() + time.sleep(0.05) + assert acquired == [] + + coordinator.ack("consumer", [first.as_batch_metadata()]) + thread.join(2) + assert len(acquired) == 1 + coordinator.ack("consumer", [acquired[0].as_batch_metadata()]) + + +def test_authorized_batch_finishes_when_larger_than_window_limits(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _single_batch_samples(stream_id, 3) + coordinator = _coordinator(tmp_path) + _register( + coordinator, + samples, + "consumer", + lookahead=0, + max_inflight=1, + ) + while coordinator.snapshot()["artifact_states"].get("queued", 0): + _publish(coordinator, stream_id, tmp_path) + + leases = [ + coordinator.acquire("consumer", sample, timeout_seconds=1) for sample in samples + ] + assert coordinator.snapshot()["inflight_acquisitions"] == 3 + assert ( + coordinator.ack("consumer", [lease.as_batch_metadata() for lease in leases]) + == 3 + ) + + +def test_demand_claims_are_round_robin_across_consumers(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 6) + coordinator = _coordinator(tmp_path) + _register( + coordinator, + samples, + "a", + "b", + "c", + lookahead=5, + max_inflight=2, + ) + errors: list[Exception] = [] + + def acquire(consumer: str, sample: StreamSampleIndex) -> None: + try: + coordinator.acquire(consumer, sample, timeout_seconds=1) + except (TimeoutError, ArtifactGenerationError): + pass + except Exception as error: # noqa: BLE001 - surface worker failures in test + errors.append(error) + + threads = [ + threading.Thread(target=acquire, args=(consumer, samples[index])) + for index, consumer in enumerate(("a", "b", "c")) + ] + for thread in threads: + thread.start() + _wait_until(lambda: coordinator.snapshot()["inflight_acquisitions"] == 3) + + claims = coordinator.claim_generation("producer", stream_id=stream_id, max_claims=3) + assert {claim.request_id for claim in claims} == { + samples[0].request_id, + samples[1].request_id, + samples[2].request_id, + } + assert all(claim.priority is ArtifactPriority.DEMAND for claim in claims) + for claim in claims: + coordinator.fail_generation("producer", claim, "stop test") + for thread in threads: + thread.join(2) + assert errors == [] + + +def test_generation_failure_retries_then_wakes_waiter(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 1) + coordinator = _coordinator(tmp_path, max_generation_attempts=2) + _register(coordinator, samples, "consumer", lookahead=0) + result: list[Exception] = [] + + def wait() -> None: + try: + coordinator.acquire("consumer", samples[0], timeout_seconds=2) + except Exception as error: # noqa: BLE001 - assert exact async failure below + result.append(error) + + thread = threading.Thread(target=wait) + thread.start() + _wait_until(lambda: coordinator.snapshot()["inflight_acquisitions"] == 1) + first = coordinator.claim_generation("producer", stream_id=stream_id)[0] + coordinator.fail_generation("producer", first, "first failure") + second = coordinator.claim_generation("producer", stream_id=stream_id)[0] + coordinator.fail_generation("producer", second, "terminal failure") + thread.join(2) + + assert len(result) == 1 + assert isinstance(result[0], ArtifactGenerationError) + assert "terminal failure" in str(result[0]) + + +def test_expired_consumer_releases_window_and_read_lease(tmp_path): + now = [100.0] + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 1) + coordinator = WindowedArtifactCoordinator( + tmp_path, + poll_seconds=0.005, + consumer_timeout_seconds=5, + claim_timeout_seconds=10, + clock=lambda: now[0], + ) + _register(coordinator, samples, "consumer", lookahead=0) + _publish(coordinator, stream_id, tmp_path) + coordinator.acquire("consumer", samples[0], timeout_seconds=1) + + now[0] += 6 + assert coordinator.recover_expired() == { + "expired_consumers": 1, + "expired_claims": 0, + } + assert coordinator.snapshot()["inflight_acquisitions"] == 0 + assert len(coordinator.begin_evictions()) == 1 + + +def test_resume_reset_rewinds_cursor_and_clears_uncommitted_leases(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 2) + coordinator = _coordinator(tmp_path) + _register(coordinator, samples, "consumer", lookahead=1, max_inflight=2) + _publish(coordinator, stream_id, tmp_path) + _publish(coordinator, stream_id, tmp_path) + committed = coordinator.acquire("consumer", samples[0], timeout_seconds=1) + assert coordinator.ack("consumer", [committed.as_batch_metadata()]) == 1 + uncommitted = coordinator.acquire("consumer", samples[1], timeout_seconds=1) + assert coordinator.snapshot()["inflight_acquisitions"] == 1 + + coordinator.register_consumer( + "consumer", + stream_id=stream_id, + lookbehind=0, + lookahead=1, + max_inflight=2, + cursor=0, + reset=True, + ) + snapshot = coordinator.snapshot() + assert snapshot["consumers"][0]["cursor"] == 0 + assert snapshot["inflight_acquisitions"] == 0 + with pytest.raises(WindowedArtifactError, match="unknown, stale, or mismatched"): + coordinator.ack("consumer", [uncommitted.as_batch_metadata()]) + replay = coordinator.acquire("consumer", samples[0], timeout_seconds=1) + assert coordinator.ack("consumer", [replay.as_batch_metadata()]) == 1 + + +def test_expired_producer_claim_is_reassigned_with_bounded_attempts(tmp_path): + now = [100.0] + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 1) + coordinator = WindowedArtifactCoordinator( + tmp_path, + poll_seconds=0.005, + consumer_timeout_seconds=30, + claim_timeout_seconds=5, + max_generation_attempts=2, + clock=lambda: now[0], + ) + _register(coordinator, samples, "consumer", lookahead=0) + first = coordinator.claim_generation("producer-a", stream_id=stream_id)[0] + assert first.generation == 0 + + now[0] += 6 + assert coordinator.recover_expired()["expired_claims"] == 1 + second = coordinator.claim_generation("producer-b", stream_id=stream_id)[0] + assert second.generation == 1 + now[0] += 6 + coordinator.recover_expired() + + assert coordinator.claim_generation("producer-c", stream_id=stream_id) == () + assert coordinator.snapshot()["artifact_states"] == {"failed": 1} + + +def _assert_long_stream_retention_bound(tmp_path: Path, count: int) -> None: + contract = { + "dataset_fingerprint": f"dataset-{count}", + "sampler_seed": 0, + } + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, count) + coordinator = _coordinator(tmp_path) + _register( + coordinator, + samples, + "consumer", + contract=contract, + lookbehind=2, + lookahead=16, + max_inflight=32, + ) + retention_bound = 2 + 16 + 1 + + for cursor in (0, count // 2, count - 1): + coordinator.register_consumer( + "consumer", + stream_id=stream_id, + lookbehind=2, + lookahead=16, + max_inflight=32, + cursor=cursor, + reset=True, + ) + for eviction in coordinator.begin_evictions(limit=retention_bound * 2): + coordinator.finish_eviction(eviction, removed=True) + while coordinator.snapshot()["artifact_states"].get("queued", 0): + _publish(coordinator, stream_id, tmp_path) + snapshot = coordinator.snapshot() + assert snapshot["retained_artifacts"] <= retention_bound + assert snapshot["artifact_states"].get("queued", 0) == 0 + for eviction in coordinator.begin_evictions(limit=retention_bound * 2): + coordinator.finish_eviction(eviction, removed=True) + + +def test_10k_position_stream_has_window_bounded_payload_retention(tmp_path): + _assert_long_stream_retention_bound(tmp_path, 10_000) + + +@pytest.mark.slow +def test_100k_position_stream_has_window_bounded_payload_retention(tmp_path): + _assert_long_stream_retention_bound(tmp_path, 100_000) diff --git a/tests/unit/train/test_cli_args.py b/tests/unit/train/test_cli_args.py index 9ba4117db..02da5f790 100644 --- a/tests/unit/train/test_cli_args.py +++ b/tests/unit/train/test_cli_args.py @@ -24,6 +24,10 @@ def test_shared_hidden_state_cache_is_opt_in(monkeypatch): assert args.shared_hidden_states_namespace is None assert args.shared_hidden_states_ttl == 3600.0 assert args.shared_hidden_states_lock_timeout == 300.0 + assert args.shared_hidden_states_consumer_id is None + assert args.shared_hidden_states_lookbehind == 2 + assert args.shared_hidden_states_lookahead == 16 + assert args.shared_hidden_states_max_inflight == 32 def test_shared_hidden_state_cache_arguments(monkeypatch): @@ -47,6 +51,38 @@ def test_shared_hidden_state_cache_arguments(monkeypatch): assert args.shared_hidden_states_lock_timeout == 45 +def test_windowed_shared_hidden_state_arguments(monkeypatch): + args = _parse( + monkeypatch, + [ + "--shared-hidden-states-path", + "shared-cache", + "--shared-hidden-states-consumer-id", + "consumer-a", + "--shared-hidden-states-lookbehind", + "3", + "--shared-hidden-states-lookahead", + "20", + "--shared-hidden-states-max-inflight", + "40", + "--shared-hidden-states-consumer-timeout", + "60", + "--shared-hidden-states-claim-timeout", + "90", + "--shared-hidden-states-generation-attempts", + "4", + ], + ) + + assert args.shared_hidden_states_consumer_id == "consumer-a" + assert args.shared_hidden_states_lookbehind == 3 + assert args.shared_hidden_states_lookahead == 20 + assert args.shared_hidden_states_max_inflight == 40 + assert args.shared_hidden_states_consumer_timeout == 60 + assert args.shared_hidden_states_claim_timeout == 90 + assert args.shared_hidden_states_generation_attempts == 4 + + @pytest.mark.parametrize( "extra", [ @@ -61,6 +97,27 @@ def test_shared_hidden_state_cache_arguments(monkeypatch): ], ["--shared-hidden-states-ttl", "-1"], ["--shared-hidden-states-lock-timeout", "0"], + ["--shared-hidden-states-consumer-id", "consumer"], + [ + "--shared-hidden-states-path", + "cache", + "--shared-hidden-states-consumer-id", + "consumer", + "--on-missing", + "raise", + ], + [ + "--shared-hidden-states-path", + "cache", + "--shared-hidden-states-consumer-id", + "", + ], + ["--shared-hidden-states-lookbehind", "-1"], + ["--shared-hidden-states-lookahead", "-1"], + ["--shared-hidden-states-max-inflight", "0"], + ["--shared-hidden-states-consumer-timeout", "0"], + ["--shared-hidden-states-claim-timeout", "0"], + ["--shared-hidden-states-generation-attempts", "0"], ], ) def test_shared_hidden_state_cache_rejects_invalid_combinations(monkeypatch, extra): diff --git a/tests/unit/train/test_shared_artifacts.py b/tests/unit/train/test_shared_artifacts.py index 57db85290..f02f59e06 100644 --- a/tests/unit/train/test_shared_artifacts.py +++ b/tests/unit/train/test_shared_artifacts.py @@ -9,7 +9,9 @@ import speculators.train.data as data_module import speculators.train.dataloader as dataloader_module -from speculators.train.data import ArrowDataset +from speculators.data_generation.windowed_artifacts import WindowedArtifactCoordinator +from speculators.train.data import WINDOWED_LEASE_KEY, ArrowDataset +from speculators.train.dataloader import WindowedBatchSampler if TYPE_CHECKING: from pathlib import Path @@ -26,6 +28,17 @@ def _write_dataset(path: Path) -> None: dataset.save_to_disk(path) +def _write_multi_dataset(path: Path, count: int = 6) -> None: + dataset = Dataset.from_dict( + { + "input_ids": [[index + 1, index + 2, index + 3] for index in range(count)], + "loss_mask": [[0, 1, 1] for _ in range(count)], + "seq_len": [3 for _ in range(count)], + } + ).with_format("torch") + dataset.save_to_disk(path) + + def _hidden_states() -> dict[str, torch.Tensor]: return { "token_ids": torch.tensor([1, 2, 3]), @@ -55,6 +68,22 @@ def _arrow_dataset( return dataset +class _SingleBatchSampler: + epoch = 0 + seed = 0 + rank = 0 + num_replicas = 1 + batch_max_length = 128 + lengths = [3] + + @staticmethod + def _generate_batches(_epoch: int) -> list[list[int]]: + return [[0]] + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + + def test_shared_dataset_requires_identity_namespace(tmp_path): data_path = tmp_path / "data" _write_dataset(data_path) @@ -68,6 +97,22 @@ def test_shared_dataset_requires_identity_namespace(tmp_path): ) +def test_windowed_dataset_requires_online_generation(tmp_path): + data_path = tmp_path / "data" + _write_dataset(data_path) + + with pytest.raises(ValueError, match="require on_missing='generate'"): + ArrowDataset( + max_len=128, + datapath=data_path, + model="model", + on_missing="raise", + shared_artifacts_path=tmp_path / "shared", + shared_artifacts_namespace="layers:2,18,33", + shared_artifacts_consumer_id="consumer", + ) + + def _successful_generator(service_path: Path, calls: list[Path]): def generate(*_args, **_kwargs): path = service_path / f"request-{len(calls)}.safetensors" @@ -121,6 +166,113 @@ def test_shared_dataset_requests_publish_once_and_delete_service_temporary( assert stats["publishes"] == 1 +def test_windowed_dataset_dispatches_reads_acks_and_cleans_final_window( + tmp_path, monkeypatch +): + data_path = tmp_path / "data" + shared_path = tmp_path / "shared" + _write_dataset(data_path) + dataset = ArrowDataset( + max_len=128, + datapath=data_path, + hidden_states_path=tmp_path / "index", + model="model", + on_missing="generate", + on_generate="delete", + shared_artifacts_path=shared_path, + shared_artifacts_namespace="layers:2,18,33", + shared_artifacts_consumer_id="consumer", + request_timeout=2, + ) + monkeypatch.setattr( + dataset, + "_materialize_shared_hs", + lambda _index, _dataset_item, _client_item: _hidden_states(), + ) + base_sampler = _SingleBatchSampler() + stream_id = dataset.configure_windowed_stream(base_sampler) + sampler = WindowedBatchSampler( + base_sampler, + stream_id=stream_id, + request_id_for_index=dataset.windowed_request_id, + ) + samples = sampler.full_epoch_samples(0) + dataset.prepare_windowed_epoch(samples, cursor=0, reset=True) + dataset.start_windowed_producer() + + item = dataset[samples[0]] + assert item is not None + lease = item.pop(WINDOWED_LEASE_KEY) + assert dataset.ack_windowed_batch([lease]) == 1 + dataset.stop_windowed_producer(completed=True) + + with WindowedArtifactCoordinator(shared_path) as coordinator: + snapshot = coordinator.snapshot() + assert snapshot["retained_artifacts"] == 0 + assert snapshot["consumers"][0]["state"] == "completed" + assert dataset.artifact_cache is not None + stats = dataset.artifact_cache.snapshot_stats() + assert stats["logical_requests"] == 1 + assert stats["publishes"] == 1 + + +@pytest.mark.parametrize("num_workers", [0, 1, 4]) +def test_windowed_scheduling_is_independent_of_dataloader_workers( + tmp_path, monkeypatch, num_workers +): + data_path = tmp_path / "data" + shared_path = tmp_path / "shared" + _write_multi_dataset(data_path) + dataset = ArrowDataset( + max_len=128, + datapath=data_path, + hidden_states_path=tmp_path / "index", + model="model", + on_missing="generate", + on_generate="delete", + shared_artifacts_path=shared_path, + shared_artifacts_namespace="layers:2,18,33", + shared_artifacts_consumer_id="consumer", + request_timeout=5, + ) + + def materialize(_index, dataset_item, _client_item): + tokens = dataset_item["input_ids"] + return { + "token_ids": tokens, + "hidden_states": torch.arange(12, dtype=torch.float32).reshape(3, 4), + } + + monkeypatch.setattr(dataset, "_materialize_shared_hs", materialize) + loader = dataloader_module._setup_dataloader( + dataset, + total_seq_len=6, + hidden_size=1, + num_workers=num_workers, + num_target_layers=3, + prefetch_factor=2, + ) + sampler = loader.batch_sampler + sampler.set_epoch(0) + samples = sampler.full_epoch_samples(0) + dataset.prepare_windowed_epoch(samples, cursor=0, reset=True) + iterator = iter(loader) + dataset.start_windowed_producer() + + seen_sequences = [] + for batch in iterator: + leases = batch.pop(data_module.WINDOWED_BATCH_LEASES_KEY) + seen_sequences.extend(lease["sequence"] for lease in leases) + dataset.ack_windowed_batch(leases) + dataset.stop_windowed_producer(completed=True) + + assert seen_sequences == list(range(len(samples))) + with WindowedArtifactCoordinator(shared_path) as coordinator: + snapshot = coordinator.snapshot() + assert snapshot["retained_artifacts"] == 0 + assert snapshot["consumers"][0]["cursor"] == len(samples) + + def test_unconfigured_dataset_keeps_existing_per_request_delete_behavior( tmp_path, monkeypatch ): @@ -306,6 +458,10 @@ def fake_arrow_dataset(**kwargs): shared_artifacts_namespace="layers:2,18,33", shared_artifacts_ttl_seconds=None, shared_artifacts_lock_timeout_seconds=45, + shared_artifacts_consumer_id="consumer-a", + shared_artifacts_lookbehind=3, + shared_artifacts_lookahead=20, + shared_artifacts_max_inflight=40, ) assert len(dataset_kwargs) == 2 @@ -314,3 +470,8 @@ def fake_arrow_dataset(**kwargs): assert kwargs["shared_artifacts_namespace"] == "layers:2,18,33" assert kwargs["shared_artifacts_ttl_seconds"] is None assert kwargs["shared_artifacts_lock_timeout_seconds"] == 45 + assert kwargs["shared_artifacts_lookbehind"] == 3 + assert kwargs["shared_artifacts_lookahead"] == 20 + assert kwargs["shared_artifacts_max_inflight"] == 40 + assert dataset_kwargs[0]["shared_artifacts_consumer_id"] == "consumer-a:train" + assert dataset_kwargs[1]["shared_artifacts_consumer_id"] == "consumer-a:val" diff --git a/tests/unit/train/test_windowed_training.py b/tests/unit/train/test_windowed_training.py new file mode 100644 index 000000000..9902b0b45 --- /dev/null +++ b/tests/unit/train/test_windowed_training.py @@ -0,0 +1,231 @@ +from __future__ import annotations + +import hashlib +from dataclasses import dataclass +from typing import Any + +import pytest +import torch +from torch import nn + +from speculators.train.data import ( + WINDOWED_BATCH_LEASES_KEY, + WINDOWED_LEASE_KEY, + create_collate_fn, +) +from speculators.train.dataloader import WindowedBatchSampler +from speculators.train.trainer import Trainer, TrainerConfig + + +def _digest(value: str) -> str: + return hashlib.sha256(value.encode()).hexdigest() + + +class _Sampler: + epoch = 0 + seed = 7 + rank = 0 + num_replicas = 1 + batch_max_length = 16 + lengths = [2, 2, 2] + + def __init__(self) -> None: + self.batches = {0: [[2, 0], [1]], 1: [[1], [0, 2]]} + + def _generate_batches(self, epoch: int) -> list[list[int]]: + return self.batches[epoch] + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + + +def test_windowed_batch_sampler_positions_are_stable_across_epochs(): + first = WindowedBatchSampler( + _Sampler(), + stream_id=_digest("stream"), + request_id_for_index=lambda index: _digest(f"request-{index}"), + ) + second = WindowedBatchSampler( + _Sampler(), + stream_id=_digest("stream"), + request_id_for_index=lambda index: _digest(f"request-{index}"), + ) + + epoch_zero = first._generate_batches(0) + epoch_one = first._generate_batches(1) + repeated = second._generate_batches(1) + + assert [sample.dataset_index for batch in epoch_zero for sample in batch] == [ + 2, + 0, + 1, + ] + assert [sample.sequence for batch in epoch_one for sample in batch] == [3, 4, 5] + assert [sample.position_id for batch in epoch_one for sample in batch] == [ + sample.position_id for batch in repeated for sample in batch + ] + assert [ + (sample.batch_start_sequence, sample.batch_end_sequence) + for sample in epoch_one[1] + ] == [(4, 6), (4, 6)] + + +def _sample(lease: dict[str, Any] | None = None) -> dict[str, Any]: + sample: dict[str, Any] = { + "hidden_states": torch.zeros(2, 4), + "input_ids": torch.tensor([1, 2]), + "verifier_last_hidden_states": torch.zeros(2, 2), + "loss_mask": torch.ones(2), + "lengths": torch.tensor([2]), + "position_ids": torch.arange(2), + } + if lease is not None: + sample[WINDOWED_LEASE_KEY] = lease + return sample + + +def test_collate_keeps_artifact_leases_out_of_model_tensors(): + lease = { + "token": "lease", + "consumer_id": "consumer", + "stream_id": _digest("stream"), + "sequence": 0, + "request_id": _digest("request"), + "generation": 0, + } + collate = create_collate_fn(max_len=4, hidden_size=2, num_target_layers=1) + + batch = collate([_sample(lease)]) + + assert batch.pop(WINDOWED_BATCH_LEASES_KEY) == [lease] + assert WINDOWED_LEASE_KEY not in batch + assert all(isinstance(value, torch.Tensor) for value in batch.values()) + + +@dataclass +class _RecordingDataset: + events: list[str] + + def ack_windowed_batch(self, _leases: list[dict[str, Any]]) -> None: + self.events.append("ack") + + def abandon_windowed_batch(self, _leases: list[dict[str, Any]]) -> None: + self.events.append("abandon") + + def stop_windowed_producer(self, *, completed: bool = False) -> None: + self.events.append(f"stop:{completed}") + + +class _Loader: + def __init__(self, dataset: _RecordingDataset, batch: dict[str, Any]) -> None: + self.dataset = dataset + self.batch_sampler = object() + self.batch = batch + + def __iter__(self): + yield self.batch + + def __len__(self) -> int: + return 1 + + +class _Model(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(())) + + def forward(self, **_batch): + loss = self.weight.square() + return None, loss, {"loss": loss.detach()} + + +def _trainer(events: list[str]) -> Trainer: + lease = { + "token": "lease", + "consumer_id": "consumer", + "stream_id": _digest("stream"), + "sequence": 0, + "request_id": _digest("request"), + "generation": 0, + } + batch = { + "input_ids": torch.tensor([[1]]), + "document_ids": torch.tensor([[0]]), + WINDOWED_BATCH_LEASES_KEY: [lease], + } + trainer = Trainer.__new__(Trainer) + trainer.model = _Model() + trainer.config = TrainerConfig( + lr=0.1, + num_epochs=1, + save_path="unused", + hidden_states_dtype=torch.bfloat16, + log_freq=100, + scheduler_type="none", + ) + trainer.local_rank = torch.device("cpu") + trainer.rank = 1 + trainer.is_distributed = False + trainer.device_type = "cpu" + trainer.train_loader = _Loader(_RecordingDataset(events), batch) + trainer.global_step = 1 + trainer.current_epoch = 0 + trainer._resume_local_step = 0 + trainer._prepared_windowed_datasets = set() + trainer.optimizers = [torch.optim.SGD(trainer.model.parameters(), lr=0.1)] + trainer.schedulers = [] + return trainer + + +def test_trainer_acks_only_after_optimizer_step(monkeypatch): + events: list[str] = [] + trainer = _trainer(events) + original_step = trainer._optimizers_step + + def step() -> None: + original_step() + events.append("optimizer") + + monkeypatch.setattr(trainer, "_optimizers_step", step) + trainer.train_epoch(0) + + assert events == ["optimizer", "ack"] + + +def test_optimizer_failure_abandons_without_ack(monkeypatch): + events: list[str] = [] + trainer = _trainer(events) + + def fail() -> None: + raise RuntimeError("optimizer failed") + + monkeypatch.setattr(trainer, "_optimizers_step", fail) + with pytest.raises(RuntimeError, match="optimizer failed"): + trainer.train_epoch(0) + + assert events == ["abandon"] + + +def test_windowed_phase_completes_consumer_after_success(): + events: list[str] = [] + loader = _Loader(_RecordingDataset(events), {}) + + def operation(epoch: int) -> str: + events.append(f"run:{epoch}") + return "result" + + assert Trainer._run_windowed_phase(loader, operation, 3) == "result" + assert events == ["run:3", "stop:True"] + + +def test_windowed_phase_stops_without_completion_after_failure(): + events: list[str] = [] + loader = _Loader(_RecordingDataset(events), {}) + + def operation(_epoch: int) -> None: + events.append("run") + raise RuntimeError("phase failed") + + with pytest.raises(RuntimeError, match="phase failed"): + Trainer._run_windowed_phase(loader, operation, 0) + assert events == ["run", "stop:False"] From 847073c4684bdb541e68412b267876af312b3dd7 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 09:14:36 +0800 Subject: [PATCH 06/20] bench: validate bounded asynchronous fanout Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- .../independent_consumer_fanout/README.md | 23 +++- .../config.example.json | 40 +++++++ .../benchmarks/independent_consumers.py | 110 +++++++++++++++--- .../benchmarks/test_independent_consumers.py | 69 +++++++++++ 4 files changed, 218 insertions(+), 24 deletions(-) diff --git a/benchmarks/independent_consumer_fanout/README.md b/benchmarks/independent_consumer_fanout/README.md index a25a2b3de..38539df37 100644 --- a/benchmarks/independent_consumer_fanout/README.md +++ b/benchmarks/independent_consumer_fanout/README.md @@ -67,16 +67,27 @@ to every service completion, exactly one miss is published, the other two reques and all failure, retry, cleanup, and timeout counters are zero. Baseline scenarios that do not use the shared-cache placeholder remain valid without cache accounting. +For bounded asynchronous fan-out, also pass +`--shared-hidden-states-consumer-id {consumer_id}` and configure lookbehind, +lookahead, and max-inflight limits. The example config enables this mode. DataLoader +workers can prefetch authorized positions, but only the trainer commits cursor progress. +When consumer windows separate, a publication that leaves every live window is evicted; +a lagging consumer may therefore regenerate that request later. The report accepts one +to `consumer_count` service completions per measured key, records the observed +multiplicity histogram and regeneration overhead, and still requires every successful +service completion to match exactly one cache miss and publication. It also adds +`windowed_artifacts`, including each consumer cursor, current artifact states, retained +bytes, in-flight acquisitions, and retained/in-flight high-water marks. + The shared cache is a filesystem data plane, not Mooncake or GPU-direct transport. Its directory must provide reliable POSIX `flock`, same-filesystem atomic rename, and directory `fsync` semantics to all consumers. Do not use an arbitrary NFS mount unless those guarantees have been verified. -It is also not a consumer-centered bounded sliding window. Disabling expiration retains -one artifact for every unique request. With a finite TTL, expired entries are reclaimed -when a dataset opens the cache or when the same key is requested again; a single pass -over previously unseen samples can therefore continue growing on-disk usage. Size the -filesystem and choose the TTL for the maximum expected consumer lag. A throughput run -with a finite dataset is not evidence that long-running cache storage is bounded. +Without a consumer ID, the legacy cache is not bounded by consumer progress: disabling +expiration retains one artifact per unique request, and a finite TTL can still grow over +a pass of unseen samples. With a consumer ID, retention is instead bounded by the union +of live windows, atomic packed batches, and in-flight leases. The focused 10k/100k CPU +state-machine tests validate that this bound is independent of total stream length. In publish-once mode, `per_consumer_completions` and the steady-state per-consumer completion map identify which consumer owned each service miss. They do not represent diff --git a/benchmarks/independent_consumer_fanout/config.example.json b/benchmarks/independent_consumer_fanout/config.example.json index 4bbc12ccd..21caf3f34 100644 --- a/benchmarks/independent_consumer_fanout/config.example.json +++ b/benchmarks/independent_consumer_fanout/config.example.json @@ -64,6 +64,16 @@ "generate", "--on-generate", "delete", + "--shared-hidden-states-path", + "{shared_artifacts_dir}", + "--shared-hidden-states-consumer-id", + "{consumer_id}", + "--shared-hidden-states-lookbehind", + "2", + "--shared-hidden-states-lookahead", + "16", + "--shared-hidden-states-max-inflight", + "32", "--num-workers", "4", "--prefetch-factor", @@ -122,6 +132,16 @@ "generate", "--on-generate", "delete", + "--shared-hidden-states-path", + "{shared_artifacts_dir}", + "--shared-hidden-states-consumer-id", + "{consumer_id}", + "--shared-hidden-states-lookbehind", + "2", + "--shared-hidden-states-lookahead", + "16", + "--shared-hidden-states-max-inflight", + "32", "--num-workers", "4", "--prefetch-factor", @@ -168,6 +188,16 @@ "generate", "--on-generate", "delete", + "--shared-hidden-states-path", + "{shared_artifacts_dir}", + "--shared-hidden-states-consumer-id", + "{consumer_id}", + "--shared-hidden-states-lookbehind", + "2", + "--shared-hidden-states-lookahead", + "16", + "--shared-hidden-states-max-inflight", + "32", "--num-workers", "4", "--prefetch-factor", @@ -214,6 +244,16 @@ "generate", "--on-generate", "delete", + "--shared-hidden-states-path", + "{shared_artifacts_dir}", + "--shared-hidden-states-consumer-id", + "{consumer_id}", + "--shared-hidden-states-lookbehind", + "2", + "--shared-hidden-states-lookahead", + "16", + "--shared-hidden-states-max-inflight", + "32", "--num-workers", "4", "--prefetch-factor", diff --git a/src/speculators/benchmarks/independent_consumers.py b/src/speculators/benchmarks/independent_consumers.py index aeb6477a3..389414fc5 100644 --- a/src/speculators/benchmarks/independent_consumers.py +++ b/src/speculators/benchmarks/independent_consumers.py @@ -31,6 +31,7 @@ summarize_gpu_window, ) from speculators.data_generation.artifact_cache import HiddenStateArtifactCache +from speculators.data_generation.windowed_artifacts import WindowedArtifactCoordinator if TYPE_CHECKING: from collections.abc import Callable @@ -419,7 +420,9 @@ def _forward(self) -> None: 443 if proxy._target.scheme == "https" else 80 ) connection = connection_type( - proxy._target.hostname, port, timeout=proxy._timeout + proxy._target.hostname, + port, + timeout=proxy._timeout, ) headers = { name: value @@ -486,6 +489,8 @@ def analyze_scenario( # noqa: C901 events: list[RequestEvent], started_at: float, finished_at: float, + *, + windowed: bool = False, ) -> dict[str, Any]: """Analyze one scenario and reject incomplete or ambiguous evidence.""" reasons: list[str] = [] @@ -507,7 +512,9 @@ def analyze_scenario( # noqa: C901 reasons.append(f"{invalid_completions} request(s) lacked a valid completion") multiplicity = scenario.expected_service_completions_per_shared_sample - publish_once = multiplicity == 1 and len(expected_ids) > 1 + shared_service = multiplicity == 1 and len(expected_ids) > 1 + bounded_regeneration = shared_service and windowed + publish_once = shared_service and not bounded_regeneration measured: list[RequestEvent] = [] steady_start = started_at steady_end = finished_at @@ -517,7 +524,7 @@ def analyze_scenario( # noqa: C901 ) for consumer_id in expected_ids } - if publish_once: + if shared_service: completed = sorted( (event for event in events if event.valid_completion), key=lambda event: event.completed_at, @@ -569,7 +576,7 @@ def analyze_scenario( # noqa: C901 per_consumer_steady = Counter(event.consumer_id for event in steady_events) if duration <= 0: reasons.append("common steady-state window is empty") - if publish_once: + if shared_service: if len(steady_events) < scenario.minimum_steady_completions_per_consumer: reasons.append( f"service has only {len(steady_events)} completion(s) in the " @@ -603,6 +610,13 @@ def analyze_scenario( # noqa: C901 for key, count in key_counts.items() if count != multiplicity or key_consumers[key] != set(expected_ids) ] + elif bounded_regeneration: + qualifying = [ + key for key, count in key_counts.items() if 1 <= count <= len(expected_ids) + ] + malformed = [ + key for key, count in key_counts.items() if count > len(expected_ids) + ] else: qualifying = [key for key, count in key_counts.items() if count == 1] malformed = [key for key, count in key_counts.items() if count != 1] @@ -611,11 +625,17 @@ def analyze_scenario( # noqa: C901 f"only {len(qualifying)} sample key(s) have expected multiplicity " f"{multiplicity}; need {scenario.minimum_shared_samples}" ) - if malformed: + if malformed and bounded_regeneration: + reasons.append( + f"{len(malformed)} measured sample key(s) exceed the maximum bounded " + f"window multiplicity {len(expected_ids)}" + ) + elif malformed: reasons.append( f"{len(malformed)} measured sample key(s) have ambiguous multiplicity" ) + multiplicity_histogram = Counter(key_counts.values()) safe_key_counts = { key[:16]: count for key, count in sorted(key_counts.items()) if key is not None } @@ -629,14 +649,37 @@ def analyze_scenario( # noqa: C901 "invalid_completions": invalid_completions, "per_consumer_completions": per_consumer_total, "per_consumer_completions_semantics": ( - "service_request_owner" if publish_once else "logical_consumer" + "service_request_owner" if shared_service else "logical_consumer" ), "expected_service_completions_per_shared_sample": multiplicity, + "multiplicity_semantics": ( + "bounded_window_regeneration" if bounded_regeneration else "exact" + ), "qualifying_shared_samples": len(qualifying), + "observed_multiplicity_histogram": { + str(count): samples + for count, samples in sorted(multiplicity_histogram.items()) + }, + "regenerated_sample_keys": sum( + samples + for count, samples in multiplicity_histogram.items() + if count > 1 + ), + "extra_service_completions": sum( + (count - 1) * samples + for count, samples in multiplicity_histogram.items() + if count > 1 + ), "sample_completion_counts": safe_key_counts, }, "steady_state": { - "mode": "publish_once_service" if publish_once else "per_consumer_service", + "mode": ( + "bounded_window_service" + if bounded_regeneration + else "publish_once_service" + if publish_once + else "per_consumer_service" + ), "warmup_completions_per_consumer": ( scenario.warmup_completions_per_consumer ), @@ -646,7 +689,7 @@ def analyze_scenario( # noqa: C901 "completions": len(steady_events), "completions_per_consumer": dict(per_consumer_steady), "completions_per_consumer_semantics": ( - "service_request_owner" if publish_once else "logical_consumer" + "service_request_owner" if shared_service else "logical_consumer" ), "completions_per_second": ( len(steady_events) / duration if duration > 0 else None @@ -655,10 +698,12 @@ def analyze_scenario( # noqa: C901 } -def analyze_cache_accounting( +def analyze_cache_accounting( # noqa: C901 scenario: ScenarioSpec, stats: dict[str, Any] | None, service_completions: int, + *, + windowed: bool = False, ) -> dict[str, Any]: """Validate cache counters against service-level request accounting.""" consumer_count = len(scenario.consumers) @@ -706,17 +751,22 @@ def analyze_cache_accounting( f"cache accounting schema_version={stats['schema_version']}, expected 1" ) - expected_logical_requests = service_completions * consumer_count - expected_hits = service_completions * (consumer_count - 1) - expected = { - "logical_requests": expected_logical_requests, - "hits": expected_hits, - "misses": service_completions, - "publishes": service_completions, - } + expected = {"misses": service_completions, "publishes": service_completions} + if not windowed: + expected.update( + { + "logical_requests": service_completions * consumer_count, + "hits": service_completions * (consumer_count - 1), + } + ) for name, value in expected.items(): if stats[name] != value: reasons.append(f"cache {name}={stats[name]}, expected {value}") + if windowed and stats["logical_requests"] != stats["misses"] + stats["hits"]: + reasons.append( + "windowed cache logical_requests must equal misses + hits, got " + f"{stats['logical_requests']} != {stats['misses']} + {stats['hits']}" + ) for name in ( "retry_generations", "generation_failures", @@ -1059,6 +1109,10 @@ def _run_scenario( # noqa: C901 for consumer in scenario.consumers for value in (*consumer.command, *consumer.env.values()) ) + windowed_artifacts_enabled = any( + "--shared-hidden-states-consumer-id" in consumer.command + for consumer in scenario.consumers + ) started_at = time.monotonic() finished_at = started_at @@ -1191,8 +1245,15 @@ def capture_step( f"producer log reader failed: {producer.reader_error}" ) - analysis = analyze_scenario(scenario, ledger.snapshot(), started_at, finished_at) + analysis = analyze_scenario( + scenario, + ledger.snapshot(), + started_at, + finished_at, + windowed=windowed_artifacts_enabled, + ) cache_stats = None + windowed_snapshot = None if cache_accounting_enabled: try: cache_stats = HiddenStateArtifactCache( @@ -1202,10 +1263,22 @@ def capture_step( runtime_errors.append( f"cache accounting unavailable: {type(error).__name__}: {error}" ) + if windowed_artifacts_enabled: + try: + with WindowedArtifactCoordinator( + shared_artifacts_dir + ) as windowed_coordinator: + windowed_snapshot = windowed_coordinator.snapshot() + except Exception as error: # noqa: BLE001 + runtime_errors.append( + "windowed artifact accounting unavailable: " + f"{type(error).__name__}: {error}" + ) cache_accounting = analyze_cache_accounting( scenario, cache_stats, analysis["request_accounting"]["valid_completions"], + windowed=windowed_artifacts_enabled, ) consumer_steps = {} for consumer in scenario.consumers: @@ -1297,6 +1370,7 @@ def capture_step( ], "request_accounting": analysis["request_accounting"], "shared_artifact_cache": cache_accounting, + "windowed_artifacts": windowed_snapshot, "steady_state": analysis["steady_state"], "consumer_step_times": consumer_steps, "makespan_seconds": max(finished_at - started_at, 0.0), diff --git a/tests/unit/benchmarks/test_independent_consumers.py b/tests/unit/benchmarks/test_independent_consumers.py index ed170ef8e..dd408b902 100644 --- a/tests/unit/benchmarks/test_independent_consumers.py +++ b/tests/unit/benchmarks/test_independent_consumers.py @@ -319,6 +319,58 @@ def test_publish_once_analysis_accepts_one_service_call_per_shared_sample(): assert result["steady_state"]["completions"] == 3 +def test_windowed_analysis_reports_bounded_regeneration(): + events = [ + _event("c0", "warmup", 0.1), + _event("c0", "sample-a", 0.2), + _event("c2", "sample-a", 0.25), + _event("c1", "sample-b", 0.3), + _event("c0", "sample-c", 0.4), + ] + + result = analyze_scenario( + _scenario(multiplicity=1), + events, + started_at=0.0, + finished_at=1.0, + windowed=True, + ) + + assert result["valid"] + accounting = result["request_accounting"] + assert accounting["multiplicity_semantics"] == "bounded_window_regeneration" + assert accounting["observed_multiplicity_histogram"] == {"1": 2, "2": 1} + assert accounting["regenerated_sample_keys"] == 1 + assert accounting["extra_service_completions"] == 1 + assert result["steady_state"]["mode"] == "bounded_window_service" + + +def test_windowed_analysis_rejects_more_generations_than_consumers(): + events = [ + _event("c0", "warmup", 0.1), + _event("c0", "sample-a", 0.2), + _event("c1", "sample-a", 0.25), + _event("c2", "sample-a", 0.3), + _event("c0", "sample-a", 0.35), + _event("c1", "sample-b", 0.4), + _event("c2", "sample-c", 0.45), + ] + + result = analyze_scenario( + _scenario(multiplicity=1), + events, + started_at=0.0, + finished_at=1.0, + windowed=True, + ) + + assert not result["valid"] + assert any( + "exceed the maximum bounded window multiplicity" in reason + for reason in result["invalid_reasons"] + ) + + def _cache_stats(**updates: int) -> dict[str, int]: stats = { "schema_version": 1, @@ -350,6 +402,23 @@ def test_cache_accounting_proves_publish_once_fanout(): assert result["stats"]["hits"] == 8 +def test_windowed_cache_allows_prefetch_ahead_of_consumer_reads(): + stats = _cache_stats(logical_requests=9, hits=5) + + windowed = analyze_cache_accounting( + _scenario(multiplicity=1), + stats, + service_completions=4, + windowed=True, + ) + synchronous = analyze_cache_accounting( + _scenario(multiplicity=1), stats, service_completions=4 + ) + + assert windowed["valid"] + assert not synchronous["valid"] + + @pytest.mark.parametrize( ("updates", "reason"), [ From 9f4cce6de82d56e179f59614560c40567c1467f7 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 09:25:37 +0800 Subject: [PATCH 07/20] fix: align artifact windows with transfer backends Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- src/speculators/train/data.py | 15 +++++++++------ tests/unit/train/test_shared_artifacts.py | 12 +++++++----- 2 files changed, 16 insertions(+), 11 deletions(-) diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py index 82f1dda75..a0a3a8116 100644 --- a/src/speculators/train/data.py +++ b/src/speculators/train/data.py @@ -676,8 +676,7 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None: ): file_idx = self._map_to_file_idx(index) target_path = ( - self.transfer.hidden_states_path - / f"hs_{file_idx}.safetensors" + self.transfer.hidden_states_path / f"hs_{file_idx}.safetensors" ) target_path.parent.mkdir(parents=True, exist_ok=True) _atomic_save_hs_file(loaded_hs, target_path) @@ -718,7 +717,7 @@ def _generate_shared_hs( self, dataset_item: dict, client_item: ClientItem ) -> dict[str, torch.Tensor]: handle: str | None = None - retrieved = False + cleanup_generated = False try: handle = generate_hidden_states( self.client, # type:ignore[arg-type] @@ -727,15 +726,19 @@ def _generate_shared_hs( timeout=self.request_timeout, max_retries=self.max_retries, ) - loaded_hs = self.transfer.get_generated(handle) - retrieved = True + cleanup_generated = True + try: + loaded_hs = self.transfer.get_generated(handle) + except TimeoutError: + cleanup_generated = False + raise if loaded_hs is None: raise ValueError(f"Failed to load hidden states for handle {handle}") check_hidden_states(loaded_hs, dataset_item["input_ids"].tolist()) return loaded_hs finally: # A failed retrieval may still have an in-flight backend writer. - if handle is not None and retrieved: + if handle is not None and cleanup_generated: self.transfer.delete(handle) def _load_requested_hidden_states( diff --git a/tests/unit/train/test_shared_artifacts.py b/tests/unit/train/test_shared_artifacts.py index f02f59e06..4871574dc 100644 --- a/tests/unit/train/test_shared_artifacts.py +++ b/tests/unit/train/test_shared_artifacts.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING +import hs_connectors.transfer as transfer_module import pytest import torch from datasets import Dataset @@ -9,6 +10,7 @@ import speculators.train.data as data_module import speculators.train.dataloader as dataloader_module +from hs_connectors import FileTransfer from speculators.data_generation.windowed_artifacts import WindowedArtifactCoordinator from speculators.train.data import WINDOWED_LEASE_KEY, ArrowDataset from speculators.train.dataloader import WindowedBatchSampler @@ -56,7 +58,7 @@ def _arrow_dataset( dataset = ArrowDataset( max_len=128, datapath=data_path, - hidden_states_path=hidden_states_path, + transfer=FileTransfer(hidden_states_path), model="model", on_missing="generate", on_generate=on_generate, @@ -175,7 +177,7 @@ def test_windowed_dataset_dispatches_reads_acks_and_cleans_final_window( dataset = ArrowDataset( max_len=128, datapath=data_path, - hidden_states_path=tmp_path / "index", + transfer=FileTransfer(tmp_path / "index"), model="model", on_missing="generate", on_generate="delete", @@ -226,7 +228,7 @@ def test_windowed_scheduling_is_independent_of_dataloader_workers( dataset = ArrowDataset( max_len=128, datapath=data_path, - hidden_states_path=tmp_path / "index", + transfer=FileTransfer(tmp_path / "index"), model="model", on_missing="generate", on_generate="delete", @@ -369,7 +371,7 @@ def time_out_waiting_for_lock(*_args, **_kwargs): lambda *_args, **_kwargs: str(artifact_path), ) monkeypatch.setattr( - data_module, + transfer_module, "wait_for_lock", time_out_waiting_for_lock, ) @@ -442,7 +444,7 @@ def fake_arrow_dataset(**kwargs): hidden_states_dtype=torch.bfloat16, noise_std=0.0, legacy_data=False, - hidden_states_path=None, + transfer=None, vllm_endpoint="http://producer/v1", on_missing="generate", on_generate="delete", From d213e1cbf77e4c6cdcd7d13911f67bbfaeb711e9 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 10:39:59 +0800 Subject: [PATCH 08/20] feat: bound asynchronous producer prefetch batches Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- .../config.example.json | 32 ++- docs/cli/train.md | 10 +- scripts/train.py | 56 +++++- .../data_generation/windowed_artifacts.py | 134 +++++++++++-- src/speculators/train/data.py | 94 ++++++++- src/speculators/train/dataloader.py | 19 +- .../test_windowed_artifacts.py | 182 +++++++++++++++++- tests/unit/train/test_cli_args.py | 23 ++- tests/unit/train/test_shared_artifacts.py | 151 +++++++++++++++ 9 files changed, 661 insertions(+), 40 deletions(-) diff --git a/benchmarks/independent_consumer_fanout/config.example.json b/benchmarks/independent_consumer_fanout/config.example.json index 21caf3f34..05b4010c2 100644 --- a/benchmarks/independent_consumer_fanout/config.example.json +++ b/benchmarks/independent_consumer_fanout/config.example.json @@ -71,7 +71,13 @@ "--shared-hidden-states-lookbehind", "2", "--shared-hidden-states-lookahead", - "16", + "40", + "--shared-hidden-states-max-prefetch-per-consumer", + "8", + "--shared-hidden-states-capture-batch-size", + "8", + "--shared-hidden-states-capture-batch-wait", + "0.002", "--shared-hidden-states-max-inflight", "32", "--num-workers", @@ -139,7 +145,13 @@ "--shared-hidden-states-lookbehind", "2", "--shared-hidden-states-lookahead", - "16", + "40", + "--shared-hidden-states-max-prefetch-per-consumer", + "8", + "--shared-hidden-states-capture-batch-size", + "8", + "--shared-hidden-states-capture-batch-wait", + "0.002", "--shared-hidden-states-max-inflight", "32", "--num-workers", @@ -195,7 +207,13 @@ "--shared-hidden-states-lookbehind", "2", "--shared-hidden-states-lookahead", - "16", + "40", + "--shared-hidden-states-max-prefetch-per-consumer", + "8", + "--shared-hidden-states-capture-batch-size", + "8", + "--shared-hidden-states-capture-batch-wait", + "0.002", "--shared-hidden-states-max-inflight", "32", "--num-workers", @@ -251,7 +269,13 @@ "--shared-hidden-states-lookbehind", "2", "--shared-hidden-states-lookahead", - "16", + "40", + "--shared-hidden-states-max-prefetch-per-consumer", + "8", + "--shared-hidden-states-capture-batch-size", + "8", + "--shared-hidden-states-capture-batch-wait", + "0.002", "--shared-hidden-states-max-inflight", "32", "--num-workers", diff --git a/docs/cli/train.md b/docs/cli/train.md index 5fd91aafb..b5082f3d8 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -96,7 +96,13 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--shared-hidden-states-lookbehind`** (int, default: `2`) Committed stream positions retained behind this consumer's cursor. -- **`--shared-hidden-states-lookahead`** (int, default: `16`) Stream positions asynchronously prepared ahead of this consumer's committed cursor. +- **`--shared-hidden-states-lookahead`** (int, default: `40`) Stream positions retained ahead of this consumer's committed cursor. + +- **`--shared-hidden-states-max-prefetch-per-consumer`** (int, default: `8`) Maximum PREFETCH artifacts for one consumer that may be queued or generating at once. Demand requests bypass this bound. This value cannot exceed `lookahead + 1`. + +- **`--shared-hidden-states-capture-batch-size`** (int, default: `8`) Global maximum number of hidden-state captures in flight across all trainer dispatchers sharing the coordinator. + +- **`--shared-hidden-states-capture-batch-wait`** (float, default: `0.002`) Seconds each dispatcher waits before claiming work so newly queued requests can coalesce into a producer batch. - **`--shared-hidden-states-max-inflight`** (int, default: `32`) Maximum waiting or leased positions per consumer. Once the first sample of a packed batch is admitted, the rest of that batch may complete atomically so a batch larger than this value cannot deadlock before trainer ACK. @@ -110,7 +116,7 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ Without `--shared-hidden-states-consumer-id`, this remains the legacy TTL cache: setting the TTL to zero retains one artifact per unique request, and a finite TTL does not by itself bound a pass over unseen samples. - With a consumer ID, SQLite tracks deterministic sampler positions, independent consumer cursors, generation claims, read leases, and the union of live windows. DataLoader workers only acquire and materialize authorized artifacts. The trainer main process advances the cursor after a successful training optimizer boundary or validation forward. Artifacts outside every live window are removed only after all read leases are released. In this mode TTL expiration is disabled; retention is controlled by windows and explicit leases. + With a consumer ID, SQLite tracks deterministic sampler positions, independent consumer cursors, generation claims, read leases, and the union of live windows. Window retention and active prefetch are separate bounds: the full lookahead remains reusable while only the nearest configured prefetches consume producer capacity. DataLoader workers only acquire and materialize authorized artifacts. The trainer main process advances the cursor after a successful training optimizer boundary or validation forward. Artifacts outside every live window are removed only after all read leases are released. In this mode TTL expiration is disabled; retention is controlled by windows and explicit leases. - **`--legacy-data`** (flag) **DEPRECATED.** Use the old data format which stores hidden states alongside token_ids. diff --git a/scripts/train.py b/scripts/train.py index 1f553cedb..f6c5c3c66 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -663,6 +663,15 @@ def main(args: argparse.Namespace): # noqa: C901 shared_artifacts_consumer_id=args.shared_hidden_states_consumer_id, shared_artifacts_lookbehind=args.shared_hidden_states_lookbehind, shared_artifacts_lookahead=args.shared_hidden_states_lookahead, + shared_artifacts_max_prefetch_per_consumer=( + args.shared_hidden_states_max_prefetch_per_consumer + ), + shared_artifacts_capture_batch_size=( + args.shared_hidden_states_capture_batch_size + ), + shared_artifacts_capture_batch_wait_seconds=( + args.shared_hidden_states_capture_batch_wait + ), shared_artifacts_max_inflight=args.shared_hidden_states_max_inflight, shared_artifacts_consumer_timeout_seconds=( args.shared_hidden_states_consumer_timeout @@ -818,6 +827,27 @@ def _validate_windowed_consumer_args( ) +def _validate_windowed_prefetch_args( + parser: argparse.ArgumentParser, args: argparse.Namespace +) -> None: + if args.shared_hidden_states_max_prefetch_per_consumer < 0: + parser.error( + "--shared-hidden-states-max-prefetch-per-consumer must be non-negative" + ) + if ( + args.shared_hidden_states_max_prefetch_per_consumer + > args.shared_hidden_states_lookahead + 1 + ): + parser.error( + "--shared-hidden-states-max-prefetch-per-consumer must not exceed " + "--shared-hidden-states-lookahead + 1" + ) + if args.shared_hidden_states_capture_batch_size < 1: + parser.error("--shared-hidden-states-capture-batch-size must be at least one") + if args.shared_hidden_states_capture_batch_wait < 0: + parser.error("--shared-hidden-states-capture-batch-wait must be non-negative") + + def _validate_windowed_shared_hidden_state_args( parser: argparse.ArgumentParser, args: argparse.Namespace ) -> None: @@ -830,6 +860,7 @@ def _validate_windowed_shared_hidden_state_args( parser.error(f"--{name.replace('_', '-')} must be non-negative") if args.shared_hidden_states_max_inflight < 1: parser.error("--shared-hidden-states-max-inflight must be at least one") + _validate_windowed_prefetch_args(parser, args) if args.shared_hidden_states_consumer_timeout <= 0: parser.error("--shared-hidden-states-consumer-timeout must be positive") if args.shared_hidden_states_claim_timeout <= 0: @@ -1043,9 +1074,32 @@ def parse_args(): parser.add_argument( "--shared-hidden-states-lookahead", type=int, - default=16, + default=40, help="Positions asynchronously prepared ahead of each consumer cursor.", ) + parser.add_argument( + "--shared-hidden-states-max-prefetch-per-consumer", + type=int, + default=8, + help=( + "Maximum queued or generating prefetches for one logical consumer. " + "Demand requests bypass this limit." + ), + ) + parser.add_argument( + "--shared-hidden-states-capture-batch-size", + type=int, + default=8, + help="Maximum concurrent hidden-state captures across all consumers.", + ) + parser.add_argument( + "--shared-hidden-states-capture-batch-wait", + type=float, + default=0.002, + help=( + "Seconds to wait for producer requests to coalesce before claiming a batch." + ), + ) parser.add_argument( "--shared-hidden-states-max-inflight", type=int, diff --git a/src/speculators/data_generation/windowed_artifacts.py b/src/speculators/data_generation/windowed_artifacts.py index 555dd7682..757570d94 100644 --- a/src/speculators/data_generation/windowed_artifacts.py +++ b/src/speculators/data_generation/windowed_artifacts.py @@ -23,7 +23,7 @@ import os from collections.abc import Callable, Iterator, Mapping, Sequence -SCHEMA_VERSION = 1 +SCHEMA_VERSION = 2 DIGEST_LENGTH = 64 MAX_CONSUMER_ID_LENGTH = 128 @@ -252,6 +252,7 @@ def _create_schema(self) -> None: cursor INTEGER NOT NULL, lookbehind INTEGER NOT NULL, lookahead INTEGER NOT NULL, + max_prefetch INTEGER NOT NULL, max_inflight INTEGER NOT NULL, state TEXT NOT NULL, heartbeat_at REAL NOT NULL, @@ -485,6 +486,7 @@ def register_consumer( stream_id: str, lookbehind: int, lookahead: int, + max_prefetch: int, max_inflight: int, cursor: int = 0, reset: bool = False, @@ -494,6 +496,7 @@ def register_consumer( for name, value in ( ("lookbehind", lookbehind), ("lookahead", lookahead), + ("max_prefetch", max_prefetch), ("cursor", cursor), ): if isinstance(value, bool) or not isinstance(value, int) or value < 0: @@ -502,6 +505,8 @@ def register_consumer( raise TypeError("max_inflight must be an integer") if max_inflight < 1: raise ValueError("max_inflight must be at least one") + if max_prefetch > lookahead + 1: + raise ValueError("max_prefetch must not exceed lookahead + 1") with self._transaction() as conn: if ( @@ -514,17 +519,24 @@ def register_consumer( row = conn.execute( "SELECT * FROM consumers WHERE consumer_id=?", (consumer_id,) ).fetchone() - config = (stream_id, lookbehind, lookahead, max_inflight) + config = ( + stream_id, + lookbehind, + lookahead, + max_prefetch, + max_inflight, + ) now = self._clock() if row is None: conn.execute( - "INSERT INTO consumers VALUES(?,?,?,?,?,?,?,?,?)", + "INSERT INTO consumers VALUES(?,?,?,?,?,?,?,?,?,?)", ( consumer_id, stream_id, cursor, lookbehind, lookahead, + max_prefetch, max_inflight, "active", now, @@ -538,6 +550,7 @@ def register_consumer( "stream_id", "lookbehind", "lookahead", + "max_prefetch", "max_inflight", ) ) @@ -674,9 +687,6 @@ def _refresh_window_locked( now, ), ) - self._queue_artifact_locked( - conn, row["request_id"], ArtifactPriority.PREFETCH - ) existing = conn.execute( "SELECT sequence FROM interests WHERE consumer_id=? AND kind='window'", (consumer_id,), @@ -689,6 +699,58 @@ def _refresh_window_locked( (consumer_id, int(row["sequence"])), ) self._prune_orphaned_locked(conn) + self._top_up_prefetch_locked(conn, consumer_id) + + def _top_up_prefetch_locked( + self, conn: sqlite3.Connection, consumer_id: str + ) -> None: + consumer = conn.execute( + "SELECT * FROM consumers WHERE consumer_id=?", (consumer_id,) + ).fetchone() + if consumer is None or consumer["state"] != "active": + return + cap = int(consumer["max_prefetch"]) + if cap == 0: + return + outstanding = int( + conn.execute( + "SELECT COUNT(DISTINCT a.request_id) FROM interests i " + "JOIN artifacts a ON a.request_id=i.request_id " + "WHERE i.consumer_id=? AND i.kind='window' AND a.priority=? " + "AND a.state IN (?,?)", + ( + consumer_id, + int(ArtifactPriority.PREFETCH), + ArtifactState.QUEUED.value, + ArtifactState.GENERATING.value, + ), + ).fetchone()[0] + ) + remaining = max(0, cap - outstanding) + if not remaining: + return + cursor = int(consumer["cursor"]) + high = cursor + int(consumer["lookahead"]) + 1 + candidates = conn.execute( + "SELECT a.request_id,MIN(i.sequence) AS first_sequence " + "FROM interests i JOIN artifacts a ON a.request_id=i.request_id " + "WHERE i.consumer_id=? AND i.kind='window' AND i.sequence>=? " + "AND i.sequence None: conn.execute( @@ -780,6 +842,10 @@ def recover_expired(self) -> dict[str, int]: ) expired_claims += 1 self._prune_orphaned_locked(conn) + for row in conn.execute( + "SELECT consumer_id FROM consumers WHERE state='active'" + ).fetchall(): + self._top_up_prefetch_locked(conn, row["consumer_id"]) return { "expired_consumers": expired_consumers, "expired_claims": expired_claims, @@ -1101,14 +1167,31 @@ def _abandon_tokens(self, consumer_id: str | None, tokens: Sequence[str]) -> Non self._refresh_window_locked(conn, owner) def claim_generation( - self, owner: str, *, stream_id: str, max_claims: int = 1 + self, + owner: str, + *, + stream_id: str, + max_claims: int = 1, + max_active_claims: int | None = None, ) -> tuple[GenerationClaim, ...]: if not owner: raise ValueError("generation owner must be non-empty") if max_claims < 1: raise ValueError("max_claims must be at least one") + if max_active_claims is not None and max_active_claims < 1: + raise ValueError("max_active_claims must be at least one") with self._transaction() as conn: self._recover_claims_locked(conn) + if max_active_claims is not None: + active = int( + conn.execute( + "SELECT COUNT(*) FROM artifacts WHERE state=?", + (ArtifactState.GENERATING.value,), + ).fetchone()[0] + ) + max_claims = min(max_claims, max_active_claims - active) + if max_claims <= 0: + return () rows = conn.execute( "SELECT DISTINCT a.* FROM artifacts a " "JOIN interests i ON i.request_id=a.request_id " @@ -1219,6 +1302,10 @@ def _recover_claims_locked(self, conn: sqlite3.Connection) -> None: (ArtifactState.GENERATING.value, now), ).fetchall() for row in rows: + affected = conn.execute( + "SELECT DISTINCT consumer_id FROM interests WHERE request_id=?", + (row["request_id"],), + ).fetchall() interested = conn.execute( "SELECT 1 FROM interests WHERE request_id=? LIMIT 1", (row["request_id"],), @@ -1242,6 +1329,8 @@ def _recover_claims_locked(self, conn: sqlite3.Connection) -> None: row["request_id"], ), ) + for consumer in affected: + self._top_up_prefetch_locked(conn, consumer["consumer_id"]) def complete_generation( self, @@ -1280,6 +1369,12 @@ def complete_generation( ), ) self._update_high_water_locked(conn) + interested = conn.execute( + "SELECT DISTINCT consumer_id FROM interests WHERE request_id=?", + (claim.request_id,), + ).fetchall() + for consumer in interested: + self._top_up_prefetch_locked(conn, consumer["consumer_id"]) def fail_generation( self, owner: str, claim: GenerationClaim, error: BaseException | str @@ -1298,6 +1393,10 @@ def fail_generation( f"stale generation failure for {claim.request_id}" ) failures = int(row["failures"]) + 1 + affected = conn.execute( + "SELECT DISTINCT consumer_id FROM interests WHERE request_id=?", + (claim.request_id,), + ).fetchall() interested = conn.execute( "SELECT 1 FROM interests WHERE request_id=? LIMIT 1", (claim.request_id,), @@ -1320,6 +1419,8 @@ def fail_generation( claim.request_id, ), ) + for consumer in affected: + self._top_up_prefetch_locked(conn, consumer["consumer_id"]) def begin_evictions(self, *, limit: int = 64) -> tuple[EvictionClaim, ...]: if limit < 1: @@ -1371,27 +1472,22 @@ def finish_eviction(self, claim: EvictionClaim, *, removed: bool) -> None: f"stale eviction completion for {claim.request_id}" ) interested = conn.execute( - "SELECT 1 FROM interests WHERE request_id=? LIMIT 1", + "SELECT DISTINCT consumer_id,kind FROM interests WHERE request_id=?", (claim.request_id,), - ).fetchone() + ).fetchall() if removed: + demand = any(item["kind"] == "demand" for item in interested) state = ( - ArtifactState.QUEUED.value - if interested - else ArtifactState.ABSENT.value - ) - priority = ( - self._retry_priority_locked(conn, claim.request_id) - if interested - else None + ArtifactState.QUEUED.value if demand else ArtifactState.ABSENT.value ) + priority = ArtifactPriority.DEMAND if demand else None conn.execute( "UPDATE artifacts SET state=?,path=NULL,size_bytes=0,priority=?," "queued_at=?,updated_at=? WHERE request_id=?", ( state, int(priority) if priority is not None else None, - self._clock() if interested else None, + self._clock() if demand else None, self._clock(), claim.request_id, ), @@ -1401,6 +1497,8 @@ def finish_eviction(self, claim: EvictionClaim, *, removed: bool) -> None: "UPDATE artifacts SET state=?,updated_at=? WHERE request_id=?", (ArtifactState.READY.value, self._clock(), claim.request_id), ) + for consumer_id in {item["consumer_id"] for item in interested}: + self._top_up_prefetch_locked(conn, consumer_id) def complete_consumer(self, consumer_id: str) -> None: with self._transaction() as conn: diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py index a0a3a8116..19b8bd6fb 100644 --- a/src/speculators/train/data.py +++ b/src/speculators/train/data.py @@ -7,6 +7,7 @@ import uuid import warnings from collections.abc import Callable +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from os import PathLike from pathlib import Path from typing import Any, Literal, cast @@ -45,6 +46,17 @@ WINDOWED_BATCH_LEASES_KEY = "_windowed_artifact_leases" +def _validate_integer_config(name: str, value: object, *, minimum: int) -> None: + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + qualifier = "non-negative" if minimum == 0 else "positive" + raise ValueError(f"{name} must be a {qualifier} integer") + + +def _validate_non_negative_number(name: str, value: object) -> None: + if isinstance(value, bool) or not isinstance(value, (int, float)) or value < 0: + raise ValueError(f"{name} must be non-negative") + + def list_files(path): datapath = [] for root, _directories, files in os.walk(path): @@ -263,7 +275,10 @@ def __init__( shared_artifacts_lock_timeout_seconds: float = 300.0, shared_artifacts_consumer_id: str | None = None, shared_artifacts_lookbehind: int = 2, - shared_artifacts_lookahead: int = 16, + shared_artifacts_lookahead: int = 40, + shared_artifacts_max_prefetch_per_consumer: int = 8, + shared_artifacts_capture_batch_size: int = 8, + shared_artifacts_capture_batch_wait_seconds: float = 0.002, shared_artifacts_max_inflight: int = 32, shared_artifacts_consumer_timeout_seconds: float = 120.0, shared_artifacts_claim_timeout_seconds: float = 300.0, @@ -322,6 +337,32 @@ def __init__( ) self.shared_artifacts_lookbehind = shared_artifacts_lookbehind self.shared_artifacts_lookahead = shared_artifacts_lookahead + _validate_integer_config( + "shared_artifacts_max_prefetch_per_consumer", + shared_artifacts_max_prefetch_per_consumer, + minimum=0, + ) + if shared_artifacts_max_prefetch_per_consumer > shared_artifacts_lookahead + 1: + raise ValueError( + "shared_artifacts_max_prefetch_per_consumer must not exceed " + "shared_artifacts_lookahead + 1" + ) + _validate_integer_config( + "shared_artifacts_capture_batch_size", + shared_artifacts_capture_batch_size, + minimum=1, + ) + _validate_non_negative_number( + "shared_artifacts_capture_batch_wait_seconds", + shared_artifacts_capture_batch_wait_seconds, + ) + self.shared_artifacts_max_prefetch_per_consumer = ( + shared_artifacts_max_prefetch_per_consumer + ) + self.shared_artifacts_capture_batch_size = shared_artifacts_capture_batch_size + self.shared_artifacts_capture_batch_wait_seconds = float( + shared_artifacts_capture_batch_wait_seconds + ) self.shared_artifacts_max_inflight = shared_artifacts_max_inflight self.shared_artifacts_consumer_timeout_seconds = ( shared_artifacts_consumer_timeout_seconds @@ -456,6 +497,7 @@ def prepare_windowed_epoch( stream_id=self._windowed_stream_id, lookbehind=self.shared_artifacts_lookbehind, lookahead=self.shared_artifacts_lookahead, + max_prefetch=self.shared_artifacts_max_prefetch_per_consumer, max_inflight=self.shared_artifacts_max_inflight, cursor=cursor, reset=reset, @@ -545,22 +587,54 @@ def _run_windowed_producer(self) -> None: lock_timeout_seconds=self.shared_artifacts_lock_timeout_seconds, ) try: - with self._new_windowed_coordinator() as coordinator: + if self.client is None: + self._setup_client() + with ( + self._new_windowed_coordinator() as coordinator, + ThreadPoolExecutor( + max_workers=self.shared_artifacts_capture_batch_size, + thread_name_prefix="artifact-capture", + ) as executor, + ): while not self._windowed_producer_stop.is_set(): coordinator.heartbeat(self.shared_artifacts_consumer_id) coordinator.recover_expired() self._evict_windowed_artifacts(coordinator, cache) + if self._windowed_producer_stop.wait( + self.shared_artifacts_capture_batch_wait_seconds + ): + break claims = coordinator.claim_generation( - owner, stream_id=self._windowed_stream_id, max_claims=1 + owner, + stream_id=self._windowed_stream_id, + max_claims=self.shared_artifacts_capture_batch_size, + max_active_claims=self.shared_artifacts_capture_batch_size, ) if claims: - for claim in claims: - try: - self._produce_windowed_claim( - coordinator, cache, owner, claim - ) - except Exception as error: # noqa: BLE001 - coordinator.fail_generation(owner, claim, error) + futures = { + executor.submit( + self._produce_windowed_claim, + coordinator, + cache, + owner, + claim, + ): claim + for claim in claims + } + pending = set(futures) + while pending: + done, pending = wait( + pending, + timeout=1.0, + return_when=FIRST_COMPLETED, + ) + coordinator.heartbeat(self.shared_artifacts_consumer_id) + for future in done: + claim = futures[future] + try: + future.result() + except Exception as error: # noqa: BLE001 + coordinator.fail_generation(owner, claim, error) continue self._windowed_producer_stop.wait(0.02) except Exception as error: # noqa: BLE001 - background thread boundary diff --git a/src/speculators/train/dataloader.py b/src/speculators/train/dataloader.py index dcb4eae8e..cec60bcc7 100644 --- a/src/speculators/train/dataloader.py +++ b/src/speculators/train/dataloader.py @@ -196,7 +196,10 @@ def create_train_val_loaders( shared_artifacts_lock_timeout_seconds: float = 300.0, shared_artifacts_consumer_id: str | None = None, shared_artifacts_lookbehind: int = 2, - shared_artifacts_lookahead: int = 16, + shared_artifacts_lookahead: int = 40, + shared_artifacts_max_prefetch_per_consumer: int = 8, + shared_artifacts_capture_batch_size: int = 8, + shared_artifacts_capture_batch_wait_seconds: float = 0.002, shared_artifacts_max_inflight: int = 32, shared_artifacts_consumer_timeout_seconds: float = 120.0, shared_artifacts_claim_timeout_seconds: float = 300.0, @@ -260,6 +263,13 @@ def create_train_val_loaders( ), shared_artifacts_lookbehind=shared_artifacts_lookbehind, shared_artifacts_lookahead=shared_artifacts_lookahead, + shared_artifacts_max_prefetch_per_consumer=( + shared_artifacts_max_prefetch_per_consumer + ), + shared_artifacts_capture_batch_size=shared_artifacts_capture_batch_size, + shared_artifacts_capture_batch_wait_seconds=( + shared_artifacts_capture_batch_wait_seconds + ), shared_artifacts_max_inflight=shared_artifacts_max_inflight, shared_artifacts_consumer_timeout_seconds=( shared_artifacts_consumer_timeout_seconds @@ -294,6 +304,13 @@ def create_train_val_loaders( ), shared_artifacts_lookbehind=shared_artifacts_lookbehind, shared_artifacts_lookahead=shared_artifacts_lookahead, + shared_artifacts_max_prefetch_per_consumer=( + shared_artifacts_max_prefetch_per_consumer + ), + shared_artifacts_capture_batch_size=shared_artifacts_capture_batch_size, + shared_artifacts_capture_batch_wait_seconds=( + shared_artifacts_capture_batch_wait_seconds + ), shared_artifacts_max_inflight=shared_artifacts_max_inflight, shared_artifacts_consumer_timeout_seconds=( shared_artifacts_consumer_timeout_seconds diff --git a/tests/unit/data_generation/test_windowed_artifacts.py b/tests/unit/data_generation/test_windowed_artifacts.py index 87cb0f327..f0c9ea198 100644 --- a/tests/unit/data_generation/test_windowed_artifacts.py +++ b/tests/unit/data_generation/test_windowed_artifacts.py @@ -94,17 +94,21 @@ def _register( contract: dict | None = None, lookbehind: int = 0, lookahead: int = 2, + max_prefetch: int | None = None, max_inflight: int = 4, ) -> None: contract = contract or {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} assert coordinator.register_stream(contract) == samples[0].stream_id coordinator.register_positions(samples) + if max_prefetch is None: + max_prefetch = min(8, lookahead + 1) for consumer_id in consumer_ids: coordinator.register_consumer( consumer_id, stream_id=samples[0].stream_id, lookbehind=lookbehind, lookahead=lookahead, + max_prefetch=max_prefetch, max_inflight=max_inflight, ) @@ -123,6 +127,19 @@ def _publish( return claim.request_id +def _complete_claim( + coordinator: WindowedArtifactCoordinator, + owner: str, + claim, + path: Path, +) -> None: + artifact = path / f"{claim.request_id}.safetensors" + artifact.write_bytes(b"payload") + coordinator.complete_generation( + owner, claim, path=artifact, size_bytes=artifact.stat().st_size + ) + + def _wait_until(predicate, timeout: float = 2.0) -> None: deadline = time.monotonic() + timeout while not predicate(): @@ -199,6 +216,7 @@ def test_completed_consumer_reactivates_at_its_committed_cursor(tmp_path): stream_id=stream_id, lookbehind=0, lookahead=2, + max_prefetch=3, max_inflight=4, cursor=0, ) @@ -206,6 +224,41 @@ def test_completed_consumer_reactivates_at_its_committed_cursor(tmp_path): assert coordinator.snapshot()["consumers"][0]["state"] == "active" +def test_max_prefetch_is_part_of_the_persisted_resume_contract(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 4) + coordinator = _coordinator(tmp_path) + _register( + coordinator, + samples, + "consumer", + lookahead=3, + max_prefetch=2, + ) + coordinator.close() + + with _coordinator(tmp_path) as resumed: + resumed.register_consumer( + "consumer", + stream_id=stream_id, + lookbehind=0, + lookahead=3, + max_prefetch=2, + max_inflight=4, + ) + assert resumed.snapshot()["consumers"][0]["max_prefetch"] == 2 + with pytest.raises(WindowedArtifactError, match="configuration changed"): + resumed.register_consumer( + "consumer", + stream_id=stream_id, + lookbehind=0, + lookahead=3, + max_prefetch=3, + max_inflight=4, + ) + + def test_two_consumers_share_publication_but_commit_independently(tmp_path): contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} stream_id = canonical_stream_id(contract) @@ -235,6 +288,88 @@ def test_two_consumers_share_publication_but_commit_independently(tmp_path): assert [claim.request_id for claim in evictions] == [samples[0].request_id] +def test_prefetch_cap_limits_active_work_and_tops_up_after_completion(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 50) + coordinator = _coordinator(tmp_path) + _register( + coordinator, + samples, + "consumer", + lookahead=40, + max_prefetch=8, + ) + + snapshot = coordinator.snapshot() + assert snapshot["artifact_states"] == {"absent": 42, "queued": 8} + first = coordinator.claim_generation( + "producer-a", + stream_id=stream_id, + max_claims=8, + max_active_claims=8, + ) + assert len(first) == 8 + assert all(claim.priority is ArtifactPriority.PREFETCH for claim in first) + assert ( + coordinator.claim_generation( + "producer-b", + stream_id=stream_id, + max_claims=8, + max_active_claims=8, + ) + == () + ) + + _complete_claim(coordinator, "producer-a", first[0], tmp_path) + assert coordinator.snapshot()["artifact_states"] == { + "absent": 41, + "generating": 7, + "queued": 1, + "ready": 1, + } + second = coordinator.claim_generation( + "producer-b", + stream_id=stream_id, + max_claims=8, + max_active_claims=8, + ) + assert len(second) == 1 + assert coordinator.snapshot()["artifact_states"]["generating"] == 8 + + +def test_demand_bypasses_full_prefetch_cap(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 50) + coordinator = _coordinator(tmp_path) + _register( + coordinator, + samples, + "consumer", + lookahead=40, + max_prefetch=8, + ) + acquired: list = [] + thread = threading.Thread( + target=lambda: acquired.append( + coordinator.acquire("consumer", samples[20], timeout_seconds=2) + ) + ) + thread.start() + _wait_until(lambda: coordinator.snapshot()["inflight_acquisitions"] == 1) + + claims = coordinator.claim_generation("producer", stream_id=stream_id, max_claims=9) + assert len(claims) == 9 + assert claims[0].request_id == samples[20].request_id + assert claims[0].priority is ArtifactPriority.DEMAND + for claim in claims: + _complete_claim(coordinator, "producer", claim, tmp_path) + thread.join(2) + assert len(acquired) == 1 + coordinator.ack("consumer", [acquired[0].as_batch_metadata()]) + + def test_max_inflight_is_independent_of_dataloader_prefetch(tmp_path): contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} stream_id = canonical_stream_id(contract) @@ -282,9 +417,24 @@ def test_authorized_batch_finishes_when_larger_than_window_limits(tmp_path): while coordinator.snapshot()["artifact_states"].get("queued", 0): _publish(coordinator, stream_id, tmp_path) - leases = [ - coordinator.acquire("consumer", sample, timeout_seconds=1) for sample in samples - ] + def acquire_one(current, results) -> None: + results.append(coordinator.acquire("consumer", current, timeout_seconds=1)) + + leases = [coordinator.acquire("consumer", samples[0], timeout_seconds=1)] + for sample in samples[1:]: + acquired: list = [] + thread = threading.Thread( + target=acquire_one, + args=(sample, acquired), + ) + thread.start() + _wait_until( + lambda: coordinator.snapshot()["inflight_acquisitions"] == len(leases) + 1 + ) + _publish(coordinator, stream_id, tmp_path) + thread.join(2) + assert len(acquired) == 1 + leases.extend(acquired) assert coordinator.snapshot()["inflight_acquisitions"] == 3 assert ( coordinator.ack("consumer", [lease.as_batch_metadata() for lease in leases]) @@ -366,6 +516,30 @@ def wait() -> None: assert "terminal failure" in str(result[0]) +def test_terminal_prefetch_failure_releases_capacity_for_next_position(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 5) + coordinator = _coordinator(tmp_path, max_generation_attempts=1) + _register( + coordinator, + samples, + "consumer", + lookahead=4, + max_prefetch=2, + ) + + claims = coordinator.claim_generation("producer", stream_id=stream_id, max_claims=2) + coordinator.fail_generation("producer", claims[0], "terminal failure") + + assert coordinator.snapshot()["artifact_states"] == { + "absent": 2, + "failed": 1, + "generating": 1, + "queued": 1, + } + + def test_expired_consumer_releases_window_and_read_lease(tmp_path): now = [100.0] contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} @@ -409,6 +583,7 @@ def test_resume_reset_rewinds_cursor_and_clears_uncommitted_leases(tmp_path): stream_id=stream_id, lookbehind=0, lookahead=1, + max_prefetch=2, max_inflight=2, cursor=0, reset=True, @@ -475,6 +650,7 @@ def _assert_long_stream_retention_bound(tmp_path: Path, count: int) -> None: stream_id=stream_id, lookbehind=2, lookahead=16, + max_prefetch=8, max_inflight=32, cursor=cursor, reset=True, diff --git a/tests/unit/train/test_cli_args.py b/tests/unit/train/test_cli_args.py index 02da5f790..1032f87fa 100644 --- a/tests/unit/train/test_cli_args.py +++ b/tests/unit/train/test_cli_args.py @@ -26,7 +26,10 @@ def test_shared_hidden_state_cache_is_opt_in(monkeypatch): assert args.shared_hidden_states_lock_timeout == 300.0 assert args.shared_hidden_states_consumer_id is None assert args.shared_hidden_states_lookbehind == 2 - assert args.shared_hidden_states_lookahead == 16 + assert args.shared_hidden_states_lookahead == 40 + assert args.shared_hidden_states_max_prefetch_per_consumer == 8 + assert args.shared_hidden_states_capture_batch_size == 8 + assert args.shared_hidden_states_capture_batch_wait == 0.002 assert args.shared_hidden_states_max_inflight == 32 @@ -63,6 +66,12 @@ def test_windowed_shared_hidden_state_arguments(monkeypatch): "3", "--shared-hidden-states-lookahead", "20", + "--shared-hidden-states-max-prefetch-per-consumer", + "7", + "--shared-hidden-states-capture-batch-size", + "6", + "--shared-hidden-states-capture-batch-wait", + "0.01", "--shared-hidden-states-max-inflight", "40", "--shared-hidden-states-consumer-timeout", @@ -77,6 +86,9 @@ def test_windowed_shared_hidden_state_arguments(monkeypatch): assert args.shared_hidden_states_consumer_id == "consumer-a" assert args.shared_hidden_states_lookbehind == 3 assert args.shared_hidden_states_lookahead == 20 + assert args.shared_hidden_states_max_prefetch_per_consumer == 7 + assert args.shared_hidden_states_capture_batch_size == 6 + assert args.shared_hidden_states_capture_batch_wait == 0.01 assert args.shared_hidden_states_max_inflight == 40 assert args.shared_hidden_states_consumer_timeout == 60 assert args.shared_hidden_states_claim_timeout == 90 @@ -114,6 +126,15 @@ def test_windowed_shared_hidden_state_arguments(monkeypatch): ], ["--shared-hidden-states-lookbehind", "-1"], ["--shared-hidden-states-lookahead", "-1"], + ["--shared-hidden-states-max-prefetch-per-consumer", "-1"], + [ + "--shared-hidden-states-lookahead", + "0", + "--shared-hidden-states-max-prefetch-per-consumer", + "2", + ], + ["--shared-hidden-states-capture-batch-size", "0"], + ["--shared-hidden-states-capture-batch-wait", "-0.001"], ["--shared-hidden-states-max-inflight", "0"], ["--shared-hidden-states-consumer-timeout", "0"], ["--shared-hidden-states-claim-timeout", "0"], diff --git a/tests/unit/train/test_shared_artifacts.py b/tests/unit/train/test_shared_artifacts.py index 4871574dc..21736675f 100644 --- a/tests/unit/train/test_shared_artifacts.py +++ b/tests/unit/train/test_shared_artifacts.py @@ -1,5 +1,7 @@ from __future__ import annotations +import threading +import time from typing import TYPE_CHECKING import hs_connectors.transfer as transfer_module @@ -86,6 +88,23 @@ def set_epoch(self, epoch: int) -> None: self.epoch = epoch +class _SequentialSampler: + epoch = 0 + seed = 0 + rank = 0 + num_replicas = 1 + batch_max_length = 128 + + def __init__(self, count: int) -> None: + self.lengths = [3] * count + + def _generate_batches(self, _epoch: int) -> list[list[int]]: + return [[index] for index in range(len(self.lengths))] + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + + def test_shared_dataset_requires_identity_namespace(tmp_path): data_path = tmp_path / "data" _write_dataset(data_path) @@ -186,6 +205,7 @@ def test_windowed_dataset_dispatches_reads_acks_and_cleans_final_window( shared_artifacts_consumer_id="consumer", request_timeout=2, ) + dataset.client = object() monkeypatch.setattr( dataset, "_materialize_shared_hs", @@ -237,6 +257,7 @@ def test_windowed_scheduling_is_independent_of_dataloader_workers( shared_artifacts_consumer_id="consumer", request_timeout=5, ) + dataset.client = object() def materialize(_index, dataset_item, _client_item): tokens = dataset_item["input_ids"] @@ -275,6 +296,130 @@ def materialize(_index, dataset_item, _client_item): assert snapshot["consumers"][0]["cursor"] == len(samples) +def test_windowed_producer_runs_bounded_concurrent_capture_batches( + tmp_path, monkeypatch +): + data_path = tmp_path / "data" + shared_path = tmp_path / "shared" + _write_multi_dataset(data_path, count=8) + dataset = ArrowDataset( + max_len=128, + datapath=data_path, + transfer=FileTransfer(tmp_path / "index"), + model="model", + on_missing="generate", + on_generate="delete", + shared_artifacts_path=shared_path, + shared_artifacts_namespace="layers:2,18,33", + shared_artifacts_consumer_id="consumer", + shared_artifacts_lookahead=7, + shared_artifacts_max_prefetch_per_consumer=8, + shared_artifacts_capture_batch_size=4, + shared_artifacts_capture_batch_wait_seconds=0, + request_timeout=5, + ) + dataset.client = object() + lock = threading.Lock() + active = 0 + peak = 0 + + def materialize(_index, dataset_item, _client_item): + nonlocal active, peak + with lock: + active += 1 + peak = max(peak, active) + try: + time.sleep(0.05) + return { + "token_ids": dataset_item["input_ids"], + "hidden_states": torch.arange(12, dtype=torch.float32).reshape(3, 4), + } + finally: + with lock: + active -= 1 + + monkeypatch.setattr(dataset, "_materialize_shared_hs", materialize) + base_sampler = _SequentialSampler(8) + stream_id = dataset.configure_windowed_stream(base_sampler) + sampler = WindowedBatchSampler( + base_sampler, + stream_id=stream_id, + request_id_for_index=dataset.windowed_request_id, + ) + dataset.prepare_windowed_epoch(sampler.full_epoch_samples(0), cursor=0, reset=True) + + dataset.start_windowed_producer() + try: + deadline = time.monotonic() + 5 + with WindowedArtifactCoordinator(shared_path) as coordinator: + while coordinator.snapshot()["artifact_states"].get("ready", 0) != 8: + if time.monotonic() >= deadline: + raise TimeoutError("concurrent capture batch did not complete") + time.sleep(0.01) + finally: + dataset.stop_windowed_producer() + + assert peak == 4 + + +def test_windowed_capture_batch_isolates_one_failed_claim(tmp_path, monkeypatch): + data_path = tmp_path / "data" + shared_path = tmp_path / "shared" + _write_multi_dataset(data_path, count=4) + dataset = ArrowDataset( + max_len=128, + datapath=data_path, + transfer=FileTransfer(tmp_path / "index"), + model="model", + on_missing="generate", + on_generate="delete", + shared_artifacts_path=shared_path, + shared_artifacts_namespace="layers:2,18,33", + shared_artifacts_consumer_id="consumer", + shared_artifacts_lookahead=3, + shared_artifacts_max_prefetch_per_consumer=4, + shared_artifacts_capture_batch_size=4, + shared_artifacts_capture_batch_wait_seconds=0, + shared_artifacts_generation_attempts=1, + request_timeout=5, + ) + dataset.client = object() + + def materialize(index, dataset_item, _client_item): + if index == 1: + raise RuntimeError("isolated capture failure") + return { + "token_ids": dataset_item["input_ids"], + "hidden_states": torch.arange(12, dtype=torch.float32).reshape(3, 4), + } + + monkeypatch.setattr(dataset, "_materialize_shared_hs", materialize) + base_sampler = _SequentialSampler(4) + stream_id = dataset.configure_windowed_stream(base_sampler) + sampler = WindowedBatchSampler( + base_sampler, + stream_id=stream_id, + request_id_for_index=dataset.windowed_request_id, + ) + dataset.prepare_windowed_epoch(sampler.full_epoch_samples(0), cursor=0, reset=True) + + dataset.start_windowed_producer() + try: + deadline = time.monotonic() + 5 + with WindowedArtifactCoordinator(shared_path) as coordinator: + while True: + states = coordinator.snapshot()["artifact_states"] + if states.get("ready", 0) == 3 and states.get("failed", 0) == 1: + break + if time.monotonic() >= deadline: + raise TimeoutError( + "mixed capture batch did not reach terminal state" + ) + time.sleep(0.01) + finally: + dataset.stop_windowed_producer() + + def test_unconfigured_dataset_keeps_existing_per_request_delete_behavior( tmp_path, monkeypatch ): @@ -463,6 +608,9 @@ def fake_arrow_dataset(**kwargs): shared_artifacts_consumer_id="consumer-a", shared_artifacts_lookbehind=3, shared_artifacts_lookahead=20, + shared_artifacts_max_prefetch_per_consumer=7, + shared_artifacts_capture_batch_size=6, + shared_artifacts_capture_batch_wait_seconds=0.01, shared_artifacts_max_inflight=40, ) @@ -474,6 +622,9 @@ def fake_arrow_dataset(**kwargs): assert kwargs["shared_artifacts_lock_timeout_seconds"] == 45 assert kwargs["shared_artifacts_lookbehind"] == 3 assert kwargs["shared_artifacts_lookahead"] == 20 + assert kwargs["shared_artifacts_max_prefetch_per_consumer"] == 7 + assert kwargs["shared_artifacts_capture_batch_size"] == 6 + assert kwargs["shared_artifacts_capture_batch_wait_seconds"] == 0.01 assert kwargs["shared_artifacts_max_inflight"] == 40 assert dataset_kwargs[0]["shared_artifacts_consumer_id"] == "consumer-a:train" assert dataset_kwargs[1]["shared_artifacts_consumer_id"] == "consumer-a:val" From 8f991db3d0229516fecbd5e047839c43a9c59f4e Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 10:40:07 +0800 Subject: [PATCH 09/20] bench: align producer common-window accounting Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- .../independent_consumer_fanout/README.md | 13 +++- .../benchmarks/independent_consumers.py | 71 ++++++++++++++++++- .../benchmarks/test_independent_consumers.py | 37 ++++++++++ 3 files changed, 118 insertions(+), 3 deletions(-) diff --git a/benchmarks/independent_consumer_fanout/README.md b/benchmarks/independent_consumer_fanout/README.md index 38539df37..63a518e5f 100644 --- a/benchmarks/independent_consumer_fanout/README.md +++ b/benchmarks/independent_consumer_fanout/README.md @@ -53,6 +53,12 @@ values are omitted from the report. The command exits nonzero if a role fails, a GPU is shared or already occupied, a completion is malformed, sample multiplicity is ambiguous, or the common steady-state window is too small. +`producer_common_steady` uses that same consumer overlap to report producer requests, +first publications, recaptures, request throughput, and effective unique-sample +throughput. The older `steady_state` field remains the service-wide interval after its +own completion warmup; do not compare its longer train-plus-validation duration with a +producer metric measured only over the common consumer overlap. + For the unshared baseline, `expected_service_completions_per_shared_sample` is one for `1p1c` and three for `1p3c`. A publish-once implementation changes the latter to one; the logical consumer @@ -69,8 +75,11 @@ do not use the shared-cache placeholder remain valid without cache accounting. For bounded asynchronous fan-out, also pass `--shared-hidden-states-consumer-id {consumer_id}` and configure lookbehind, -lookahead, and max-inflight limits. The example config enables this mode. DataLoader -workers can prefetch authorized positions, but only the trainer commits cursor progress. +lookahead, max-prefetch, capture-batch, and max-inflight limits. The example config +uses the aligned `2/40/8/8` window and producer-batch settings. The full lookahead is +retained for reuse while only eight PREFETCH requests per consumer can be queued or +generating; demand bypasses that cap. DataLoader workers can prefetch authorized +positions, but only the trainer commits cursor progress. When consumer windows separate, a publication that leaves every live window is evicted; a lagging consumer may therefore regenerate that request later. The report accepts one to `consumer_count` service completions per measured key, records the observed diff --git a/src/speculators/benchmarks/independent_consumers.py b/src/speculators/benchmarks/independent_consumers.py index 389414fc5..c6cce1685 100644 --- a/src/speculators/benchmarks/independent_consumers.py +++ b/src/speculators/benchmarks/independent_consumers.py @@ -698,6 +698,60 @@ def analyze_scenario( # noqa: C901 } +def analyze_producer_common_window( + events: list[RequestEvent], + *, + start_monotonic_ns: int, + end_monotonic_ns: int, +) -> dict[str, Any]: + """Classify first publications and recaptures in the consumer overlap.""" + duration_ns = end_monotonic_ns - start_monotonic_ns + reasons: list[str] = [] + if duration_ns <= 0: + reasons.append("common consumer steady-state window is empty") + + seen: set[str] = set() + first_publications = 0 + recaptures = 0 + owners: Counter[str] = Counter() + for event in sorted(events, key=lambda item: item.completed_at): + if not event.valid_completion or event.request_key is None: + continue + completed_ns = int(event.completed_at * 1_000_000_000) + recapture = event.request_key in seen + seen.add(event.request_key) + if not start_monotonic_ns <= completed_ns <= end_monotonic_ns: + continue + owners[event.consumer_id] += 1 + if recapture: + recaptures += 1 + else: + first_publications += 1 + + requests = first_publications + recaptures + if not requests: + reasons.append("producer has no completions in the common steady-state window") + duration_seconds = max(duration_ns, 0) / 1_000_000_000 + return { + "valid": not reasons, + "invalid_reasons": reasons, + "interval": "common_consumer_steady_overlap", + "start_monotonic_ns": start_monotonic_ns, + "end_monotonic_ns": end_monotonic_ns, + "duration_seconds": duration_seconds, + "requests": requests, + "first_publications": first_publications, + "recaptures": recaptures, + "requests_per_second": ( + requests / duration_seconds if duration_seconds > 0 else None + ), + "effective_unique_samples_per_second": ( + first_publications / duration_seconds if duration_seconds > 0 else None + ), + "service_request_owners": dict(owners), + } + + def analyze_cache_accounting( # noqa: C901 scenario: ScenarioSpec, stats: dict[str, Any] | None, @@ -1245,9 +1299,10 @@ def capture_step( f"producer log reader failed: {producer.reader_error}" ) + request_events = ledger.snapshot() analysis = analyze_scenario( scenario, - ledger.snapshot(), + request_events, started_at, finished_at, windowed=windowed_artifacts_enabled, @@ -1309,11 +1364,23 @@ def capture_step( "invalid_reasons": ["common consumer steady-state window is unavailable"], "per_gpu": {}, } + producer_window: dict[str, Any] = { + "valid": False, + "invalid_reasons": [ + "common consumer steady-state producer window is unavailable" + ], + "interval": "common_consumer_steady_overlap", + } if len(consumer_starts) == len(scenario.consumers) and len(consumer_ends) == len( scenario.consumers ): overlap_start = max(consumer_starts) overlap_end = min(consumer_ends) + producer_window = analyze_producer_common_window( + request_events, + start_monotonic_ns=overlap_start, + end_monotonic_ns=overlap_end, + ) try: gpu_window = summarize_gpu_window( iter_gpu_samples(monitor.sample_path), @@ -1347,6 +1414,7 @@ def capture_step( *runtime_errors, *analysis["invalid_reasons"], *cache_accounting["invalid_reasons"], + *producer_window["invalid_reasons"], *memory["invalid_reasons"], *monitor_summary.get("errors", []), *monitor_summary.get("ownership_violations", []), @@ -1372,6 +1440,7 @@ def capture_step( "shared_artifact_cache": cache_accounting, "windowed_artifacts": windowed_snapshot, "steady_state": analysis["steady_state"], + "producer_common_steady": producer_window, "consumer_step_times": consumer_steps, "makespan_seconds": max(finished_at - started_at, 0.0), "memory": memory, diff --git a/tests/unit/benchmarks/test_independent_consumers.py b/tests/unit/benchmarks/test_independent_consumers.py index dd408b902..5eb6c0247 100644 --- a/tests/unit/benchmarks/test_independent_consumers.py +++ b/tests/unit/benchmarks/test_independent_consumers.py @@ -20,6 +20,7 @@ ScenarioSpec, analyze_cache_accounting, analyze_consumer_steps, + analyze_producer_common_window, analyze_scenario, canonical_request_key, ) @@ -371,6 +372,42 @@ def test_windowed_analysis_rejects_more_generations_than_consumers(): ) +def test_producer_common_window_separates_first_publications_and_recaptures(): + events = [ + _event("c0", "sample-a", 0.2), + _event("c1", "sample-b", 0.35), + _event("c2", "sample-a", 0.4), + _event("c0", "sample-c", 0.6), + ] + + result = analyze_producer_common_window( + events, + start_monotonic_ns=300_000_000, + end_monotonic_ns=500_000_000, + ) + + assert result["valid"] + assert result["duration_seconds"] == pytest.approx(0.2) + assert result["requests"] == 2 + assert result["first_publications"] == 1 + assert result["recaptures"] == 1 + assert result["requests_per_second"] == pytest.approx(10.0) + assert result["effective_unique_samples_per_second"] == pytest.approx(5.0) + assert result["service_request_owners"] == {"c1": 1, "c2": 1} + + +def test_producer_common_window_rejects_empty_interval(): + result = analyze_producer_common_window( + [], + start_monotonic_ns=500_000_000, + end_monotonic_ns=500_000_000, + ) + + assert not result["valid"] + assert result["requests"] == 0 + assert result["requests_per_second"] is None + + def _cache_stats(**updates: int) -> dict[str, int]: stats = { "schema_version": 1, From 96b3d9b854e15e896b8e393c0e13842db0db8e56 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 11:29:08 +0800 Subject: [PATCH 10/20] bench: align fanout workload with train-only baseline Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- .../independent_consumer_fanout/README.md | 8 +- .../config.example.json | 8 ++ docs/cli/train.md | 2 +- scripts/train.py | 10 +- src/speculators/train/dataloader.py | 130 ++++++++++-------- .../benchmarks/test_independent_consumers.py | 21 +++ tests/unit/train/test_shared_artifacts.py | 46 +++++++ 7 files changed, 162 insertions(+), 63 deletions(-) diff --git a/benchmarks/independent_consumer_fanout/README.md b/benchmarks/independent_consumer_fanout/README.md index 63a518e5f..70c4a41fe 100644 --- a/benchmarks/independent_consumer_fanout/README.md +++ b/benchmarks/independent_consumer_fanout/README.md @@ -53,11 +53,15 @@ values are omitted from the report. The command exits nonzero if a role fails, a GPU is shared or already occupied, a completion is malformed, sample multiplicity is ambiguous, or the common steady-state window is too small. +The example commands set `--train-data-ratio 1.0`, so every input belongs to the +training stream and no validation pass changes window positions or producer request +accounting. Keep this setting fixed when comparing against a train-only baseline. + `producer_common_steady` uses that same consumer overlap to report producer requests, first publications, recaptures, request throughput, and effective unique-sample throughput. The older `steady_state` field remains the service-wide interval after its -own completion warmup; do not compare its longer train-plus-validation duration with a -producer metric measured only over the common consumer overlap. +own completion warmup; do not compare that full-run interval with a producer metric +measured only over the common consumer overlap. For the unshared baseline, `expected_service_completions_per_shared_sample` is one for `1p1c` and three for diff --git a/benchmarks/independent_consumer_fanout/config.example.json b/benchmarks/independent_consumer_fanout/config.example.json index 05b4010c2..36d028564 100644 --- a/benchmarks/independent_consumer_fanout/config.example.json +++ b/benchmarks/independent_consumer_fanout/config.example.json @@ -40,6 +40,8 @@ "{output_dir}/checkpoints", "--epochs", "1", + "--train-data-ratio", + "1.0", "--total-seq-len", "3072", "--optimizer", @@ -114,6 +116,8 @@ "{output_dir}/checkpoints", "--epochs", "1", + "--train-data-ratio", + "1.0", "--total-seq-len", "3072", "--optimizer", @@ -176,6 +180,8 @@ "{output_dir}/checkpoints", "--epochs", "1", + "--train-data-ratio", + "1.0", "--total-seq-len", "3072", "--optimizer", @@ -238,6 +244,8 @@ "{output_dir}/checkpoints", "--epochs", "1", + "--train-data-ratio", + "1.0", "--total-seq-len", "3072", "--optimizer", diff --git a/docs/cli/train.md b/docs/cli/train.md index b5082f3d8..97a8db0a9 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -148,7 +148,7 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--lr`** (float, default: `1e-4`) Learning rate. -- **`--train-data-ratio`** (float, default: `0.9`) Ratio of data to use for training, the rest of the provided data will be used for validation. +- **`--train-data-ratio`** (float, default: `0.9`) Ratio of data to use for training. The rest is used for validation; set this to `1.0` for a train-only run with no validation loader. - **`--no-resume-from-checkpoint`** (flag) Disable automatic checkpoint resumption. Without this flag, this script will automatically load the latest checkpoint in `{save-path}` if one exists. diff --git a/scripts/train.py b/scripts/train.py index f6c5c3c66..bdfa8f782 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -1136,7 +1136,15 @@ def parse_args(): parser.add_argument("--save-path", type=str, default="./output/checkpoints") parser.add_argument("--epochs", type=int, default=20) parser.add_argument("--lr", type=float, default=1e-4) - parser.add_argument("--train-data-ratio", type=float, default=0.9) + parser.add_argument( + "--train-data-ratio", + type=float, + default=0.9, + help=( + "Fraction of data used for training in (0, 1]. " + "Set to 1.0 to disable validation." + ), + ) parser.add_argument("--no-resume-from-checkpoint", action="store_true") parser.add_argument( "--logger", diff --git a/src/speculators/train/dataloader.py b/src/speculators/train/dataloader.py index cec60bcc7..70a08c6e2 100644 --- a/src/speculators/train/dataloader.py +++ b/src/speculators/train/dataloader.py @@ -205,18 +205,22 @@ def create_train_val_loaders( shared_artifacts_claim_timeout_seconds: float = 300.0, shared_artifacts_generation_attempts: int = 3, train_data_ratio: float = 0.9, -) -> tuple[DataLoader, DataLoader]: - """Create training and validation DataLoaders. +) -> tuple[DataLoader, DataLoader | None]: + """Create training and optional validation DataLoaders. Handles dataset construction (legacy vs Arrow) and dataloader wiring. + A ``train_data_ratio`` of 1.0 uses the full dataset for training and returns + no validation loader. Non-data SP ranks get lightweight loaders with no workers (they receive batches via scatter). Reads DP/SP topology from :mod:`speculators.train.distributed`. """ noise_transform = AddUniformNoise(std=noise_std) - if not (0.0 < train_data_ratio < 1.0): - raise ValueError(f"train_data_ratio must be in (0, 1), got {train_data_ratio}") + if not (0.0 < train_data_ratio <= 1.0): + raise ValueError(f"train_data_ratio must be in (0, 1], got {train_data_ratio}") + + val_dataset: BaseDataset | None = None if legacy_data: warnings.warn( @@ -231,11 +235,12 @@ def create_train_val_loaders( transform=noise_transform, hidden_states_dtype=hidden_states_dtype, ) - val_dataset: BaseDataset = SampleFileDataset( - file_list=val_files, - max_len=total_seq_len, - hidden_states_dtype=hidden_states_dtype, - ) + if val_files: + val_dataset = SampleFileDataset( + file_list=val_files, + max_len=total_seq_len, + hidden_states_dtype=hidden_states_dtype, + ) else: train_dataset = ArrowDataset( datapath=data_path, @@ -279,47 +284,52 @@ def create_train_val_loaders( ), shared_artifacts_generation_attempts=(shared_artifacts_generation_attempts), ) - val_dataset = ArrowDataset( - datapath=data_path, - max_len=total_seq_len, - transfer=transfer, - vllm_endpoint=vllm_endpoint, - on_missing=on_missing, - on_generate=on_generate, - split_ratio=train_data_ratio - 1.0, - model=verifier_name_or_path, - hidden_states_dtype=hidden_states_dtype, - request_timeout=request_timeout, - max_retries=max_retries, - shared_artifacts_path=shared_artifacts_path, - shared_artifacts_namespace=shared_artifacts_namespace, - shared_artifacts_ttl_seconds=shared_artifacts_ttl_seconds, - shared_artifacts_lock_timeout_seconds=( - shared_artifacts_lock_timeout_seconds - ), - shared_artifacts_consumer_id=( - f"{shared_artifacts_consumer_id}:val" - if shared_artifacts_consumer_id is not None - else None - ), - shared_artifacts_lookbehind=shared_artifacts_lookbehind, - shared_artifacts_lookahead=shared_artifacts_lookahead, - shared_artifacts_max_prefetch_per_consumer=( - shared_artifacts_max_prefetch_per_consumer - ), - shared_artifacts_capture_batch_size=shared_artifacts_capture_batch_size, - shared_artifacts_capture_batch_wait_seconds=( - shared_artifacts_capture_batch_wait_seconds - ), - shared_artifacts_max_inflight=shared_artifacts_max_inflight, - shared_artifacts_consumer_timeout_seconds=( - shared_artifacts_consumer_timeout_seconds - ), - shared_artifacts_claim_timeout_seconds=( - shared_artifacts_claim_timeout_seconds - ), - shared_artifacts_generation_attempts=(shared_artifacts_generation_attempts), - ) + if train_data_ratio < 1.0: + val_dataset = ArrowDataset( + datapath=data_path, + max_len=total_seq_len, + transfer=transfer, + vllm_endpoint=vllm_endpoint, + on_missing=on_missing, + on_generate=on_generate, + split_ratio=train_data_ratio - 1.0, + model=verifier_name_or_path, + hidden_states_dtype=hidden_states_dtype, + request_timeout=request_timeout, + max_retries=max_retries, + shared_artifacts_path=shared_artifacts_path, + shared_artifacts_namespace=shared_artifacts_namespace, + shared_artifacts_ttl_seconds=shared_artifacts_ttl_seconds, + shared_artifacts_lock_timeout_seconds=( + shared_artifacts_lock_timeout_seconds + ), + shared_artifacts_consumer_id=( + f"{shared_artifacts_consumer_id}:val" + if shared_artifacts_consumer_id is not None + else None + ), + shared_artifacts_lookbehind=shared_artifacts_lookbehind, + shared_artifacts_lookahead=shared_artifacts_lookahead, + shared_artifacts_max_prefetch_per_consumer=( + shared_artifacts_max_prefetch_per_consumer + ), + shared_artifacts_capture_batch_size=( + shared_artifacts_capture_batch_size + ), + shared_artifacts_capture_batch_wait_seconds=( + shared_artifacts_capture_batch_wait_seconds + ), + shared_artifacts_max_inflight=shared_artifacts_max_inflight, + shared_artifacts_consumer_timeout_seconds=( + shared_artifacts_consumer_timeout_seconds + ), + shared_artifacts_claim_timeout_seconds=( + shared_artifacts_claim_timeout_seconds + ), + shared_artifacts_generation_attempts=( + shared_artifacts_generation_attempts + ), + ) train_loader = _setup_dataloader( train_dataset, @@ -330,14 +340,16 @@ def create_train_val_loaders( prefetch_factor=prefetch_factor, preprocess=preprocess, ) - val_loader = _setup_dataloader( - val_dataset, - total_seq_len, - hidden_size, - num_target_layers=num_target_layers, - num_workers=num_workers, - prefetch_factor=prefetch_factor, - preprocess=preprocess, - ) + val_loader = None + if val_dataset is not None: + val_loader = _setup_dataloader( + val_dataset, + total_seq_len, + hidden_size, + num_target_layers=num_target_layers, + num_workers=num_workers, + prefetch_factor=prefetch_factor, + preprocess=preprocess, + ) return train_loader, val_loader diff --git a/tests/unit/benchmarks/test_independent_consumers.py b/tests/unit/benchmarks/test_independent_consumers.py index 5eb6c0247..8d478d4dc 100644 --- a/tests/unit/benchmarks/test_independent_consumers.py +++ b/tests/unit/benchmarks/test_independent_consumers.py @@ -4,6 +4,7 @@ import threading import urllib.request from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path import pytest from pydantic import ValidationError @@ -130,6 +131,26 @@ def fake_run_scenario(_config, scenario, _output_dir): assert [scenario["kind"] for scenario in report["scenarios"]] == ["1p3c"] +def test_example_config_uses_train_only_workloads(): + config_path = ( + Path(__file__).parents[3] + / "benchmarks" + / "independent_consumer_fanout" + / "config.example.json" + ) + config = json.loads(config_path.read_text()) + + commands = [ + consumer["command"] + for scenario in config["scenarios"] + for consumer in scenario["consumers"] + ] + assert len(commands) == 4 + for command in commands: + ratio_index = command.index("--train-data-ratio") + assert command[ratio_index + 1] == "1.0" + + @pytest.mark.parametrize( "command", [ diff --git a/tests/unit/train/test_shared_artifacts.py b/tests/unit/train/test_shared_artifacts.py index 21736675f..fc351fb70 100644 --- a/tests/unit/train/test_shared_artifacts.py +++ b/tests/unit/train/test_shared_artifacts.py @@ -628,3 +628,49 @@ def fake_arrow_dataset(**kwargs): assert kwargs["shared_artifacts_max_inflight"] == 40 assert dataset_kwargs[0]["shared_artifacts_consumer_id"] == "consumer-a:train" assert dataset_kwargs[1]["shared_artifacts_consumer_id"] == "consumer-a:val" + + +def test_train_only_loader_uses_full_dataset_without_validation(monkeypatch): + dataset_kwargs = [] + dataset = object() + + def fake_arrow_dataset(**kwargs): + dataset_kwargs.append(kwargs) + return dataset + + monkeypatch.setattr(dataloader_module, "ArrowDataset", fake_arrow_dataset) + monkeypatch.setattr( + dataloader_module, + "_setup_dataloader", + lambda dataset, *_args, **_kwargs: dataset, + ) + + train_loader, val_loader = dataloader_module.create_train_val_loaders( + data_path="data", + train_data_ratio=1.0, + total_seq_len=128, + hidden_states_dtype=torch.bfloat16, + noise_std=0.0, + legacy_data=False, + transfer=None, + vllm_endpoint="http://producer/v1", + on_missing="generate", + on_generate="delete", + verifier_name_or_path="model", + request_timeout=10, + max_retries=2, + hidden_size=4, + num_target_layers=3, + num_workers=0, + prefetch_factor=1, + preprocess=None, + shared_artifacts_path="shared", + shared_artifacts_namespace="layers:2,18,33", + shared_artifacts_consumer_id="consumer-a", + ) + + assert train_loader is dataset + assert val_loader is None + assert len(dataset_kwargs) == 1 + assert dataset_kwargs[0]["split_ratio"] == 1.0 + assert dataset_kwargs[0]["shared_artifacts_consumer_id"] == "consumer-a:train" From 55e5246274ae229d02d13702a4aa935c805c5f9c Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 13:15:03 +0800 Subject: [PATCH 11/20] perf: optimize DFlash consumer training Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- pyproject.toml | 1 + scripts/train.py | 89 +++++++- src/speculators/models/dflash/core.py | 194 ++++++++++++++++-- src/speculators/models/dflash/metrics.py | 75 +++++++ src/speculators/ops/__init__.py | 1 + .../ops/fused_linear_cross_entropy.py | 148 +++++++++++++ src/speculators/train/optimizers.py | 45 +++- src/speculators/train/trainer.py | 65 +++++- .../integration/models/test_model_forward.py | 34 +++ tests/unit/models/test_dflash_metrics.py | 46 ++++- tests/unit/models/test_dflash_optimized_ce.py | 120 +++++++++++ tests/unit/ops/__init__.py | 1 + .../ops/test_fused_linear_cross_entropy.py | 108 ++++++++++ .../test_fused_linear_cross_entropy_cuda.py | 46 +++++ tests/unit/train/test_cli_args.py | 67 ++++++ tests/unit/train/test_optimizers.py | 175 ++++++++++++++++ 16 files changed, 1194 insertions(+), 21 deletions(-) create mode 100644 src/speculators/ops/__init__.py create mode 100644 src/speculators/ops/fused_linear_cross_entropy.py create mode 100644 tests/unit/models/test_dflash_optimized_ce.py create mode 100644 tests/unit/ops/__init__.py create mode 100644 tests/unit/ops/test_fused_linear_cross_entropy.py create mode 100644 tests/unit/ops/test_fused_linear_cross_entropy_cuda.py create mode 100644 tests/unit/train/test_optimizers.py diff --git a/pyproject.toml b/pyproject.toml index 8fa781704..86bef9534 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,6 +65,7 @@ dependencies = [ [project.optional-dependencies] nvml = ["nvidia-ml-py>=12.0.0"] +liger = ["liger-kernel==0.8.0"] dev = [ # build "build>=1.5.0", diff --git a/scripts/train.py b/scripts/train.py index bdfa8f782..b162e80c3 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -701,6 +701,9 @@ def main(args: argparse.Namespace): # noqa: C901 train_call_kwargs=train_call_kwargs, val_call_kwargs=val_call_kwargs, optimizer=args.optimizer, + adamw_backend=args.adamw_backend, + gradient_clip_backend=args.gradient_clip_backend, + max_grad_norm=args.max_grad_norm, weight_decay=args.weight_decay, muon_lr=args.muon_lr, muon_momentum=args.muon_momentum, @@ -895,7 +898,7 @@ def _validate_shared_hidden_state_args( _validate_windowed_shared_hidden_state_args(parser, args) -def parse_args(): +def parse_args(): # noqa: C901 parser = argparse.ArgumentParser() parser.add_argument("--verifier-name-or-path", type=str, required=True) parser.add_argument( @@ -1346,6 +1349,30 @@ def parse_args(): default=4.0, help="Decay gamma for DFlash/DSpark loss weighting (default: 4.0)", ) + parser.add_argument( + "--dflash-linear-cross-entropy-backend", + choices=["torch", "liger"], + default="torch", + help="DFlash CE backend. Liger is available only for an exactly-CE loss.", + ) + parser.add_argument( + "--dflash-compact-zero-weight-ce-rows", + action=argparse.BooleanOptionalAction, + default=False, + help="Exclude zero-weight DFlash rows before the fused Liger CE kernel.", + ) + parser.add_argument( + "--dflash-label-source", + choices=["verifier_argmax", "input_ids"], + default="verifier_argmax", + help="Hard-label source for the opt-in fused DFlash CE path.", + ) + parser.add_argument( + "--dflash-verifier-argmax-chunk-size", + type=int, + default=0, + help="Verifier LM-head rows per argmax chunk; 0 materializes all rows.", + ) # D-Pace specific arguments (loss weight option + smoothing) parser.add_argument( "--per-position-loss-weight", @@ -1516,6 +1543,24 @@ def parse_args(): "the remaining params (norms, biases, embeddings, lm_head)." ), ) + parser.add_argument( + "--adamw-backend", + choices=["auto", "foreach", "fused"], + default="auto", + help="AdamW execution backend (also applies to AdamW parameters in Muon mode).", + ) + parser.add_argument( + "--gradient-clip-backend", + choices=["torch", "fused_adamw"], + default="torch", + help="Apply clipping explicitly or through the fused AdamW grad scale.", + ) + parser.add_argument( + "--max-grad-norm", + type=float, + default=1.0, + help="Maximum gradient norm (default: 1.0).", + ) parser.add_argument( "--weight-decay", type=float, @@ -1555,7 +1600,47 @@ def parse_args(): provided = explicitly_provided_dests(parser, DECODER_SHAPING_FLAGS) validate_draft_init_args(parser, args, provided) - resolve_loss_config(args.loss_fn) + loss_config = resolve_loss_config(args.loss_fn) + + if args.dflash_linear_cross_entropy_backend == "liger": + if args.speculator_type != "dflash": + parser.error( + "--dflash-linear-cross-entropy-backend=liger requires " + "--speculator-type=dflash" + ) + if set(loss_config) != {"ce"}: + parser.error( + "--dflash-linear-cross-entropy-backend=liger requires an " + "exactly-CE --loss-fn" + ) + if ( + args.dflash_compact_zero_weight_ce_rows + and args.dflash_linear_cross_entropy_backend != "liger" + ): + parser.error( + "--dflash-compact-zero-weight-ce-rows requires the Liger CE backend" + ) + if args.dflash_label_source != "verifier_argmax" and ( + args.dflash_linear_cross_entropy_backend != "liger" + ): + parser.error("--dflash-label-source=input_ids requires the Liger CE backend") + if args.dflash_verifier_argmax_chunk_size < 0: + parser.error("--dflash-verifier-argmax-chunk-size must be non-negative") + if args.dflash_verifier_argmax_chunk_size and ( + args.dflash_linear_cross_entropy_backend != "liger" + ): + parser.error( + "--dflash-verifier-argmax-chunk-size requires the Liger CE backend" + ) + if args.max_grad_norm <= 0: + parser.error("--max-grad-norm must be positive") + if args.gradient_clip_backend == "fused_adamw" and ( + args.optimizer != "adamw" or args.adamw_backend != "fused" + ): + parser.error( + "--gradient-clip-backend=fused_adamw requires " + "--optimizer=adamw --adamw-backend=fused" + ) if args.per_position_loss_weight == "dpace": if args.loss_fn != "ce": diff --git a/src/speculators/models/dflash/core.py b/src/speculators/models/dflash/core.py index 0bad5418e..8b3423ad3 100644 --- a/src/speculators/models/dflash/core.py +++ b/src/speculators/models/dflash/core.py @@ -14,7 +14,7 @@ from speculators.models.attention import create_float_mask from speculators.models.dflash import DFlashSpeculatorConfig from speculators.models.dflash.attention import create_anchor_block_mask_mod -from speculators.models.dflash.metrics import compute_metrics +from speculators.models.dflash.metrics import compute_fused_ce_metrics, compute_metrics from speculators.models.dflash.model_definitions import Qwen3DFlashDecoderLayer from speculators.models.dflash.utils import ( get_base_indices_for_anchored_blocks, @@ -22,6 +22,10 @@ ) from speculators.models.metrics import LossConfig, resolve_loss_config from speculators.models.utils import conditional_torch_compile, resolve_target_layer_ids +from speculators.ops.fused_linear_cross_entropy import ( + frozen_linear_cross_entropy, + validate_liger_installation, +) logger = logging.getLogger(__name__) @@ -241,15 +245,66 @@ def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]: "per_position_loss_weight", "fixed-exp-decay" ) dpace_alpha = kwargs.get("dpace_alpha", 0.5) + linear_ce_backend = kwargs.get("dflash_linear_cross_entropy_backend", "torch") + compact_ce_rows = kwargs.get("dflash_compact_zero_weight_ce_rows", False) + label_source = kwargs.get("dflash_label_source", "verifier_argmax") + verifier_argmax_chunk_size = kwargs.get("dflash_verifier_argmax_chunk_size", 0) + if linear_ce_backend not in {"liger", "torch"}: + raise ValueError( + "dflash_linear_cross_entropy_backend must be 'torch' or 'liger'" + ) + if label_source not in {"input_ids", "verifier_argmax"}: + raise ValueError( + "dflash_label_source must be 'verifier_argmax' or 'input_ids'" + ) + if verifier_argmax_chunk_size < 0: + raise ValueError("dflash_verifier_argmax_chunk_size must be non-negative") + if compact_ce_rows and linear_ce_backend != "liger": + raise ValueError( + "dflash_compact_zero_weight_ce_rows requires " + "dflash_linear_cross_entropy_backend='liger'" + ) + if ( + label_source != "verifier_argmax" or verifier_argmax_chunk_size + ) and linear_ce_backend != "liger": + raise ValueError( + "DFlash label-source and verifier-argmax chunking options require " + "dflash_linear_cross_entropy_backend='liger'" + ) + if linear_ce_backend == "liger": + if set(loss_config) != {"ce"}: + raise ValueError( + "dflash_linear_cross_entropy_backend='liger' requires an " + "exactly-CE --loss-fn configuration" + ) + validate_liger_installation() shared = { "loss_config": loss_config, "gamma": gamma, "max_anchors": max_anchors, "per_position_loss_weight": per_position_loss_weight, "dpace_alpha": dpace_alpha, + "linear_cross_entropy_backend": linear_ce_backend, + "compact_zero_weight_ce_rows": compact_ce_rows, + "label_source": label_source, + "verifier_argmax_chunk_size": verifier_argmax_chunk_size, } return dict(shared), dict(shared) + def prepare_fused_linear_cross_entropy(self, compute_dtype: torch.dtype) -> None: + """Prepare the ignored verifier head as the fused CE compute weight.""" + + if not torch.equal( + self.lm_head.weight.detach(), self.verifier_lm_head.weight.detach() + ): + raise RuntimeError( + "Liger CE requires identical frozen draft and verifier LM heads; " + "the loaded checkpoint and verifier weights differ" + ) + # This head is excluded from checkpoints. Casting it once avoids a persistent + # extra weight copy while matching autocast's compute precision. + self.verifier_lm_head.to(dtype=compute_dtype) + @property def mask_token_id(self) -> int: if self.config.mask_token_id is None: @@ -326,6 +381,10 @@ def _backbone_forward( verifier_last_hidden_states: torch.Tensor, # [1, total_seq_len, hidden_size] document_ids: torch.Tensor, # [1, total_seq_len] position_ids: torch.Tensor | None = None, # [1, total_seq_len] + target_ids_only: bool = False, + materialize_draft_logits: bool = True, + label_source: str = "verifier_argmax", + verifier_argmax_chunk_size: int = 0, **kwargs, ): """Run the anchored-block draft transformer up to the draft logits. @@ -377,16 +436,23 @@ def _backbone_forward( anchor_positions, self.block_size ) # shape: [num_anchors*block_size] - with torch.no_grad(): - verifier_logits = self.verifier_lm_head( - self.verifier_norm(verifier_last_hidden_states) + if target_ids_only: + targets = self._ce_target_ids( + input_ids, + verifier_last_hidden_states, + anchored_block_indices, + label_source=label_source, + verifier_argmax_chunk_size=verifier_argmax_chunk_size, ) - if not self.config.sample_from_anchor: - # False: shift right by 1 so slot j predicts token at position j - verifier_logits = torch.roll(verifier_logits, 1, dims=1) - # else: True, slot k predicts token at position k+1 (next), no shift - targets = verifier_logits[:, anchored_block_indices] - # shape: [1, num_anchors*block_size, draft_vocab_size] + else: + with torch.no_grad(): + verifier_logits = self.verifier_lm_head( + self.verifier_norm(verifier_last_hidden_states) + ) + if not self.config.sample_from_anchor: + # False: shift right so slot j predicts token at position j. + verifier_logits = torch.roll(verifier_logits, 1, dims=1) + targets = verifier_logits[:, anchored_block_indices] for layer_idx, layer in enumerate(self.layers): noise_embedding = layer( @@ -402,8 +468,7 @@ def _backbone_forward( ) hidden = self.norm(noise_embedding) - logits = self.lm_head(hidden) - # shape: [1, num_anchors*block_size, vocab_size] + logits = self.lm_head(hidden) if materialize_draft_logits else None aligned_loss_mask = loss_mask.clone()[:, anchored_block_indices] # shape: [1, num_anchors*block_size] @@ -421,6 +486,50 @@ def _backbone_forward( return hidden, logits, targets, aligned_loss_mask, anchored_block_indices + @torch.no_grad() + def _ce_target_ids( + self, + input_ids: torch.Tensor, + verifier_last_hidden_states: torch.Tensor, + anchored_block_indices: torch.Tensor, + *, + label_source: str, + verifier_argmax_chunk_size: int, + ) -> torch.Tensor: + if label_source == "input_ids": + if self.use_draft_vocab: + raise ValueError( + "dflash_label_source='input_ids' requires the full verifier " + "vocabulary" + ) + label_indices = anchored_block_indices + if self.config.sample_from_anchor: + label_indices = label_indices + 1 + return input_ids[:, label_indices] + if label_source != "verifier_argmax": + raise ValueError(f"unsupported DFlash label source: {label_source!r}") + + normalized = self.verifier_norm(verifier_last_hidden_states) + sequence_length = normalized.shape[1] + if ( + verifier_argmax_chunk_size == 0 + or verifier_argmax_chunk_size >= sequence_length + ): + verifier_ids = self.verifier_lm_head(normalized).argmax(dim=-1) + else: + verifier_ids = torch.cat( + [ + self.verifier_lm_head( + normalized[:, start : start + verifier_argmax_chunk_size] + ).argmax(dim=-1) + for start in range(0, sequence_length, verifier_argmax_chunk_size) + ], + dim=1, + ) + if not self.config.sample_from_anchor: + verifier_ids = torch.roll(verifier_ids, 1, dims=1) + return verifier_ids[:, anchored_block_indices] + @conditional_torch_compile def forward( self, @@ -435,9 +544,14 @@ def forward( max_anchors: int = 3072, per_position_loss_weight: str = "fixed-exp-decay", dpace_alpha: float = 0.5, + linear_cross_entropy_backend: str = "torch", + compact_zero_weight_ce_rows: bool = False, + label_source: str = "verifier_argmax", + verifier_argmax_chunk_size: int = 0, **kwargs, ): - _, logits, targets, aligned_loss_mask, _ = self._backbone_forward( + use_fused_ce = linear_cross_entropy_backend == "liger" + hidden, logits, targets, aligned_loss_mask, _ = self._backbone_forward( hidden_states, input_ids, loss_mask, @@ -445,8 +559,62 @@ def forward( document_ids, position_ids, max_anchors=max_anchors, + target_ids_only=use_fused_ce, + materialize_draft_logits=not use_fused_ce, + label_source=label_source, + verifier_argmax_chunk_size=verifier_argmax_chunk_size, **kwargs, ) + if use_fused_ce: + if loss_config is None or set(loss_config) != {"ce"}: + raise ValueError("Liger DFlash CE requires an exactly-CE loss config") + flat_hidden = hidden.reshape(-1, hidden.shape[-1]) + flat_targets = targets.reshape(-1) + flat_mask = aligned_loss_mask.reshape(-1) + active_indices = None + if compact_zero_weight_ce_rows: + active_indices = torch.nonzero(flat_mask > 0, as_tuple=False).flatten() + if active_indices.numel() > 0: + flat_hidden = flat_hidden.index_select(0, active_indices) + flat_targets = flat_targets.index_select(0, active_indices) + + if active_indices is not None and active_indices.numel() == 0: + loss_per_token = hidden.reshape(-1, hidden.shape[-1]).sum(dim=-1) * 0 + token_accuracy = torch.zeros_like(loss_per_token) + else: + compute_weight = self.verifier_lm_head.weight.detach() + flat_hidden = flat_hidden.to(compute_weight.dtype) + loss_per_token, token_accuracy = frozen_linear_cross_entropy( + flat_hidden, + compute_weight, + flat_targets, + ) + if active_indices is not None: + loss_per_token = loss_per_token.new_zeros( + flat_mask.shape + ).index_copy(0, active_indices, loss_per_token) + token_accuracy = token_accuracy.new_zeros( + flat_mask.shape + ).index_copy(0, active_indices, token_accuracy) + return ( + None, + *compute_fused_ce_metrics( + loss_per_token, + token_accuracy, + aligned_loss_mask, + self.block_size, + loss_weight=loss_config["ce"][1], + gamma=gamma, + per_position_loss_weight=per_position_loss_weight, + dpace_alpha=dpace_alpha, + sample_from_anchor=self.config.sample_from_anchor, + ), + ) + + if linear_cross_entropy_backend != "torch": + raise ValueError( + "linear_cross_entropy_backend must be either 'torch' or 'liger'" + ) loss, metrics = compute_metrics( logits, targets, diff --git a/src/speculators/models/dflash/metrics.py b/src/speculators/models/dflash/metrics.py index 827db77a2..e76c6871d 100644 --- a/src/speculators/models/dflash/metrics.py +++ b/src/speculators/models/dflash/metrics.py @@ -15,6 +15,7 @@ ) _DEFAULT_LOSS_CONFIG: LossConfig = {"kl_div": (kl_div_loss, 1.0)} +_EPS = 1e-5 def compute_metrics( @@ -106,3 +107,77 @@ def compute_metrics( metrics["eal_sum"] = eal metrics["eal_total"] = ones.clone() return loss, metrics + + +def compute_fused_ce_metrics( + loss_per_token: torch.Tensor, + token_accuracy: torch.Tensor, + loss_mask: torch.Tensor, + block_size: int, + *, + loss_weight: float = 1.0, + gamma: float = 4.0, + per_position_loss_weight: str = "fixed-exp-decay", + dpace_alpha: float = 0.5, + sample_from_anchor: bool = False, +) -> tuple[torch.Tensor, dict]: + """Reduce fused hard-label CE outputs with the standard DFlash semantics.""" + + loss_per_token = loss_per_token.reshape_as(loss_mask) + token_accuracy = token_accuracy.reshape_as(loss_mask) + seq_len = loss_mask.shape[1] + pos_idx = torch.arange(seq_len, device=loss_mask.device) + pos_idx = pos_idx.remainder(block_size).unsqueeze(0) + + elementwise_loss = loss_per_token * loss_mask.to(loss_per_token.dtype) + if per_position_loss_weight == "dpace": + decay_mult = dpace_loss_decay( + pos_idx.to(elementwise_loss.dtype), + loss_mask=loss_mask, + block_size=block_size, + dpace_alpha=dpace_alpha, + elementwise_loss=elementwise_loss, + ) + else: + decay_mult = dflash_loss_decay( + pos_idx.to(elementwise_loss.dtype), + gamma=gamma, + sample_from_anchor=sample_from_anchor, + ) + elementwise_loss = elementwise_loss * decay_mult + denominator = loss_mask.to(elementwise_loss.dtype).sum(dim=1) + _EPS + ce_loss = (elementwise_loss.sum(dim=1) / denominator).mean() + loss = ce_loss * loss_weight + + selected = loss_mask.to(torch.bool) + correct = torch.masked_select(token_accuracy.to(torch.bool), selected) + selected_positions = torch.masked_select(pos_idx, selected) + correct_per_pos = torch.zeros( + block_size, dtype=torch.float, device=loss_mask.device + ) + total_per_pos = torch.zeros_like(correct_per_pos) + correct_per_pos.scatter_add_(0, selected_positions, correct.float()) + total_per_pos.scatter_add_( + 0, selected_positions, torch.ones_like(correct, dtype=torch.float) + ) + + ones = torch.tensor(1.0, device=loss_mask.device) + metrics: dict[str, Any] = { + "loss_sum": loss.detach().clone(), + "loss_total": ones, + } + start_pos = 0 if sample_from_anchor else 1 + metrics["full_acc_sum"] = correct_per_pos[start_pos:].sum() + metrics["full_acc_total"] = total_per_pos[start_pos:].sum() + + eal = torch.zeros((), device=loss_mask.device) + cumulative_accuracy = torch.ones((), device=loss_mask.device) + for pos in range(start_pos, block_size): + metrics[f"position_{pos}_acc_sum"] = correct_per_pos[pos] + metrics[f"position_{pos}_acc_total"] = total_per_pos[pos] + position_accuracy = correct_per_pos[pos] / total_per_pos[pos].clamp(min=1.0) + cumulative_accuracy = cumulative_accuracy * position_accuracy + eal = eal + cumulative_accuracy + metrics["eal_sum"] = eal + metrics["eal_total"] = ones.clone() + return loss, metrics diff --git a/src/speculators/ops/__init__.py b/src/speculators/ops/__init__.py new file mode 100644 index 000000000..a9dbc6493 --- /dev/null +++ b/src/speculators/ops/__init__.py @@ -0,0 +1 @@ +"""Optional optimized operators used by speculator training.""" diff --git a/src/speculators/ops/fused_linear_cross_entropy.py b/src/speculators/ops/fused_linear_cross_entropy.py new file mode 100644 index 000000000..6d85c7902 --- /dev/null +++ b/src/speculators/ops/fused_linear_cross_entropy.py @@ -0,0 +1,148 @@ +"""Memory-efficient frozen-linear cross entropy with weighted gradients.""" + +from __future__ import annotations + +import inspect +from functools import lru_cache +from importlib.metadata import PackageNotFoundError, version +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from collections.abc import Callable + +_LIGER_VERSION = "0.8.0" +_LIGER_FORWARD_RESULT_SIZE = 7 +_MATRIX_NDIM = 2 +_TARGET_NDIM = 1 +_REQUIRED_FORWARD_PARAMETERS = { + "_input", + "weight", + "target", + "bias", + "reduction", + "return_token_accuracy", + "return_predicted_tokens", +} + + +@lru_cache(maxsize=1) +def _load_liger_forward() -> Callable: + try: + installed_version = version("liger-kernel") + except PackageNotFoundError as exc: + raise RuntimeError( + "dflash_linear_cross_entropy_backend='liger' requires " + "`pip install 'speculators[liger]'`" + ) from exc + if installed_version != _LIGER_VERSION: + raise RuntimeError( + "dflash_linear_cross_entropy_backend='liger' requires liger-kernel==" + f"{_LIGER_VERSION}, found {installed_version}" + ) + + try: + from liger_kernel.ops.fused_linear_cross_entropy import ( # noqa: PLC0415 + fused_linear_cross_entropy_forward, + ) + except (ImportError, ModuleNotFoundError) as exc: + raise RuntimeError( + "liger-kernel is installed but its fused linear cross entropy " + "operator is unavailable; reinstall `speculators[liger]`" + ) from exc + + parameters = set(inspect.signature(fused_linear_cross_entropy_forward).parameters) + missing = _REQUIRED_FORWARD_PARAMETERS - parameters + if missing: + raise RuntimeError( + "unsupported Liger fused linear cross entropy ABI; missing parameters: " + + ", ".join(sorted(missing)) + ) + return fused_linear_cross_entropy_forward + + +def validate_liger_installation() -> None: + """Fail before training when the pinned Liger operator is unavailable.""" + + _load_liger_forward() + + +class _FrozenLinearCrossEntropy(torch.autograd.Function): + @staticmethod + def forward(ctx, hidden, weight, target, bias): + result = _load_liger_forward()( + _input=hidden, + weight=weight, + target=target, + bias=bias, + reduction="none", + return_token_accuracy=True, + ) + if not isinstance(result, tuple) or len(result) != _LIGER_FORWARD_RESULT_SIZE: + result_count = len(result) if isinstance(result, tuple) else type(result) + raise RuntimeError( + "unsupported Liger fused linear cross entropy return ABI: " + f"expected 7 values, got {result_count}" + ) + + ( + loss, + _z_loss, + token_accuracy, + _predicted_tokens, + grad_input, + _grad_weight, + _grad_bias, + ) = result + ctx.save_for_backward(grad_input.detach()) + ctx.mark_non_differentiable(token_accuracy) + return loss, token_accuracy + + @staticmethod + def backward(ctx, grad_loss, _grad_accuracy): + (grad_input,) = ctx.saved_tensors + if grad_loss is None: + return torch.zeros_like(grad_input), None, None, None + compute_dtype = torch.promote_types(grad_input.dtype, torch.float32) + scaled_grad_input = grad_input.to(compute_dtype) * grad_loss.reshape(-1, 1).to( + compute_dtype + ) + return scaled_grad_input.to(grad_input.dtype), None, None, None + + +@torch.compiler.disable +def frozen_linear_cross_entropy( + hidden: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + bias: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return per-token CE and correctness without materializing full logits.""" + + if ( + hidden.ndim != _MATRIX_NDIM + or weight.ndim != _MATRIX_NDIM + or target.ndim != _TARGET_NDIM + ): + raise ValueError("expected hidden [N,H], weight [V,H], and target [N]") + if hidden.shape[0] != target.shape[0] or hidden.shape[1] != weight.shape[1]: + raise ValueError("hidden, weight, and target shapes are incompatible") + if bias is not None and (bias.ndim != 1 or bias.shape[0] != weight.shape[0]): + raise ValueError("bias shape must match the LM-head vocabulary dimension") + if target.dtype != torch.long: + raise ValueError("target must use torch.long token ids") + if hidden.device != weight.device or hidden.device != target.device: + raise ValueError("hidden, weight, and target must be on the same device") + if bias is not None and bias.device != hidden.device: + raise ValueError("bias must be on the same device as hidden") + if hidden.dtype != weight.dtype: + raise ValueError("hidden and weight must use the same dtype") + if weight.requires_grad or (bias is not None and bias.requires_grad): + raise ValueError("Liger DFlash linear cross entropy requires a frozen LM head") + return _FrozenLinearCrossEntropy.apply( + hidden.contiguous(), weight, target.contiguous(), bias + ) + + +__all__ = ["frozen_linear_cross_entropy", "validate_liger_installation"] diff --git a/src/speculators/train/optimizers.py b/src/speculators/train/optimizers.py index d87fde891..5acb9d78c 100644 --- a/src/speculators/train/optimizers.py +++ b/src/speculators/train/optimizers.py @@ -29,6 +29,46 @@ _MATRIX_NDIM = 2 +def _adamw_backend_kwargs( + named_params: list[tuple[str, Tensor]], backend: str +) -> dict[str, bool]: + if backend == "auto": + return {} + if backend == "foreach": + return {"foreach": True, "fused": False} + if backend == "fused": + if any(param.device.type != "cuda" for _, param in named_params): + raise ValueError("adamw_backend='fused' requires CUDA parameters") + return {"foreach": False, "fused": True} + raise ValueError(f"Unsupported AdamW backend: {backend!r}") + + +def restore_adamw_backend( + optimizers: list[torch.optim.Optimizer], backend: str +) -> None: + """Restore the requested AdamW execution backend after checkpoint loading. + + ``Optimizer.load_state_dict`` restores ``foreach`` and ``fused`` from each saved + parameter group. Those values describe the run that wrote the checkpoint, not the + current run, so they must not silently override the current trainer configuration. + """ + if backend == "auto": + foreach, fused = None, None + elif backend == "foreach": + foreach, fused = True, False + elif backend == "fused": + foreach, fused = False, True + else: + raise ValueError(f"Unsupported AdamW backend: {backend!r}") + + for optimizer in optimizers: + if not isinstance(optimizer, torch.optim.AdamW): + continue + for param_group in optimizer.param_groups: + param_group["foreach"] = foreach + param_group["fused"] = fused + + def split_named_params_for_muon( model: Module, ) -> tuple[list[tuple[str, Tensor]], list[tuple[str, Tensor]]]: @@ -63,11 +103,13 @@ def build_optimizers(model: Module, config) -> list[torch.optim.Optimizer]: "adamw" returns a single optimizer; "muon" returns ``[Muon, AdamW]``. """ if config.optimizer == "adamw": + named_params = list(model.named_parameters()) return [ torch.optim.AdamW( - model.named_parameters(), + named_params, lr=config.lr, weight_decay=config.weight_decay, + **_adamw_backend_kwargs(named_params, config.adamw_backend), ) ] @@ -97,6 +139,7 @@ def build_optimizers(model: Module, config) -> list[torch.optim.Optimizer]: adamw_params, lr=config.lr, weight_decay=config.weight_decay, + **_adamw_backend_kwargs(adamw_params, config.adamw_backend), ) ) if not optimizers: diff --git a/src/speculators/train/trainer.py b/src/speculators/train/trainer.py index 3c0d90988..03bd3aa7c 100644 --- a/src/speculators/train/trainer.py +++ b/src/speculators/train/trainer.py @@ -35,7 +35,7 @@ is_distributed, ) from speculators.train.graceful_shutdown import with_graceful_shutdown -from speculators.train.optimizers import build_optimizers +from speculators.train.optimizers import build_optimizers, restore_adamw_backend from speculators.train.utils import normalize_counted_metrics root_logger = logging.getLogger("speculators") @@ -105,6 +105,9 @@ class TrainerConfig(NamedTuple): train_call_kwargs: dict | None = None val_call_kwargs: dict | None = None optimizer: Literal["adamw", "muon"] = "adamw" + adamw_backend: Literal["auto", "foreach", "fused"] = "auto" + gradient_clip_backend: Literal["torch", "fused_adamw"] = "torch" + max_grad_norm: float = 1.0 weight_decay: float = 0.01 muon_lr: float = 0.02 muon_momentum: float = 0.95 @@ -188,6 +191,7 @@ def __init__( self.setup_trainer() self.setup_model() + self._prepare_model_execution_backends() self.setup_optimizer() self._prepared_windowed_datasets: set[int] = set() @@ -402,6 +406,7 @@ def setup_optimizer(self): last_epoch = -1 if self.resume_from_checkpoint and self.checkpointer.previous_epoch != -1: self.checkpointer.load_optimizer_state_dict(self.model, self.optimizers) + restore_adamw_backend(self.optimizers, self.config.adamw_backend) last_epoch = self.checkpointer.previous_epoch # Setup scheduler(s) — one per optimizer so each optimizer's base LR (e.g. @@ -435,13 +440,65 @@ def make_scheduler(opt: torch.optim.Optimizer): if self.resume_from_checkpoint and self.checkpointer.previous_epoch != -1: self.checkpointer.load_scheduler_state_dict(self.schedulers) + def _prepare_model_execution_backends(self) -> None: + train_kwargs = self.config.train_call_kwargs or {} + if train_kwargs.get("linear_cross_entropy_backend") != "liger": + return + model = ( + self.model.module + if isinstance(self.model, DistributedDataParallel) + else self.model + ) + prepare = getattr(model, "prepare_fused_linear_cross_entropy", None) + if prepare is None: + raise ValueError( + "linear_cross_entropy_backend='liger' is unsupported by this model" + ) + prepare(self.config.hidden_states_dtype) + def _optimizers_zero_grad(self): for opt in self.optimizers: opt.zero_grad() + def _clip_gradients(self): + if self.config.gradient_clip_backend == "torch": + return torch.nn.utils.clip_grad_norm_( + self.model.parameters(), self.config.max_grad_norm + ) + if self.config.gradient_clip_backend != "fused_adamw": + raise ValueError("gradient_clip_backend must be 'torch' or 'fused_adamw'") + if len(self.optimizers) != 1: + raise RuntimeError("fused AdamW gradient clipping requires one optimizer") + optimizer = self.optimizers[0] + group_backends = { + (group.get("foreach"), group.get("fused")) + for group in optimizer.param_groups + } + if group_backends != {(False, True)}: + raise RuntimeError( + "gradient_clip_backend='fused_adamw' requires fused AdamW" + ) + gradients = [ + parameter.grad + for parameter in self.model.parameters() + if parameter.grad is not None + ] + if not gradients: + return torch.tensor(0.0) + grad_norm = torch.nn.utils.get_total_norm(gradients, foreach=True) + optimizer.grad_scale = ((grad_norm + 1e-6) / self.config.max_grad_norm).clamp( + min=1.0 + ) + return grad_norm + def _optimizers_step(self): - for opt in self.optimizers: - opt.step() + try: + for opt in self.optimizers: + opt.step() + finally: + for opt in self.optimizers: + if hasattr(opt, "grad_scale"): + del opt.grad_scale def _schedulers_step(self): for scheduler in self.schedulers: @@ -540,7 +597,7 @@ def train_epoch(self, epoch: int): timer.mark("fwd") self._optimizers_zero_grad() loss.backward() - torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) + self._clip_gradients() timer.mark("bwd") self._optimizers_step() diff --git a/tests/integration/models/test_model_forward.py b/tests/integration/models/test_model_forward.py index f4283fda2..33b0eb6ca 100644 --- a/tests/integration/models/test_model_forward.py +++ b/tests/integration/models/test_model_forward.py @@ -5,6 +5,7 @@ parameter variation tests. """ +import os from collections.abc import Callable from dataclasses import dataclass, field from functools import partial @@ -13,6 +14,7 @@ import pytest import torch +from speculators.models.metrics import resolve_loss_config from speculators.models.mtp import shift_batch_mtp from speculators.models.mtp.core import compute_step_weights from tests.conftest import requires_cuda, requires_transformers_version @@ -248,6 +250,38 @@ def test_boundary_tokens(self, draft_vocab_model): @requires_cuda class TestDFlashParams: + @pytest.mark.skipif( + os.environ.get("SPECULATORS_RUN_LIGER_TESTS") != "1", + reason="set SPECULATORS_RUN_LIGER_TESTS=1 on an exclusive CUDA GPU", + ) + def test_compiled_liger_ce_forward_backward(self): + torch.compiler.reset() + model = make_dflash_model(dtype=torch.float32) + with torch.no_grad(): + model.verifier_lm_head.weight.copy_(model.lm_head.weight) + model.prepare_fused_linear_cross_entropy(torch.bfloat16) + samples = _make_samples([128]) + batch = make_batch(max_len=MAX_LEN, samples=samples, hidden_size=HIDDEN_SIZE) + + with torch.autocast("cuda", dtype=torch.bfloat16): + _, loss, metrics = model( + **batch, + max_anchors=8, + loss_config=resolve_loss_config("ce"), + linear_cross_entropy_backend="liger", + compact_zero_weight_ce_rows=True, + verifier_argmax_chunk_size=32, + ) + loss.backward() + + assert loss.isfinite() + assert metrics["loss_sum"].isfinite() + assert any( + parameter.grad is not None and parameter.grad.isfinite().all() + for parameter in model.parameters() + if parameter.requires_grad + ) + @pytest.mark.parametrize("block_size", [2, 4, 8]) def test_varying_block_size(self, block_size): model = make_dflash_model(block_size=block_size) diff --git a/tests/unit/models/test_dflash_metrics.py b/tests/unit/models/test_dflash_metrics.py index 1922e5b8b..230e3d7c5 100644 --- a/tests/unit/models/test_dflash_metrics.py +++ b/tests/unit/models/test_dflash_metrics.py @@ -4,8 +4,9 @@ import pytest import torch +from torch.nn.functional import cross_entropy -from speculators.models.dflash.metrics import compute_metrics +from speculators.models.dflash.metrics import compute_fused_ce_metrics, compute_metrics from speculators.models.metrics import ( ce_loss, compute_accuracy_multi_step, @@ -337,3 +338,46 @@ def test_counts_match_compute_accuracy(self): for i in range(1, 4): assert torch.isclose(metrics[f"position_{i}_acc_sum"], expected_correct[i]) assert torch.isclose(metrics[f"position_{i}_acc_total"], expected_total[i]) + + +@pytest.mark.parametrize("weighting", ["fixed-exp-decay", "dpace"]) +def test_fused_ce_reduction_matches_standard_metrics(weighting): + torch.manual_seed(17) + logits = torch.randn(1, 12, 19, dtype=torch.double, requires_grad=True) + target_ids = torch.randint(0, logits.shape[-1], (1, logits.shape[1])) + targets = _ids_to_logits(target_ids, logits.shape[-1]).to(logits.dtype) + loss_mask = torch.tensor( + [[0, 1, 1, 1, 0, 1, 0, 1, 0, 1, 1, 0]], dtype=torch.float64 + ) + loss_weight = 0.37 + loss_config = {"ce": (ce_loss, loss_weight)} + + standard_loss, standard_metrics = compute_metrics( + logits, + targets, + loss_mask, + block_size=4, + gamma=3.0, + loss_config=loss_config, + per_position_loss_weight=weighting, + dpace_alpha=0.6, + ) + loss_per_token = cross_entropy( + logits.flatten(0, 1), target_ids.flatten(), reduction="none" + ) + token_accuracy = (logits.argmax(dim=-1) == target_ids).float().flatten() + fused_loss, fused_metrics = compute_fused_ce_metrics( + loss_per_token, + token_accuracy, + loss_mask, + block_size=4, + loss_weight=loss_weight, + gamma=3.0, + per_position_loss_weight=weighting, + dpace_alpha=0.6, + ) + + torch.testing.assert_close(fused_loss, standard_loss) + assert fused_metrics.keys() == standard_metrics.keys() + for name in standard_metrics: + torch.testing.assert_close(fused_metrics[name], standard_metrics[name]) diff --git a/tests/unit/models/test_dflash_optimized_ce.py b/tests/unit/models/test_dflash_optimized_ce.py new file mode 100644 index 000000000..dccaa9253 --- /dev/null +++ b/tests/unit/models/test_dflash_optimized_ce.py @@ -0,0 +1,120 @@ +"""Focused tests for DFlash hard-label CE preparation and target selection.""" + +import pytest +import torch +from transformers import Qwen3Config + +from speculators.config import SpeculatorsConfig, VerifierConfig +from speculators.models.dflash import DFlashSpeculatorConfig +from speculators.models.dflash.core import DFlashDraftModel +from speculators.proposals.greedy import GreedyTokenProposalConfig + + +def _model(*, sample_from_anchor: bool = False) -> DFlashDraftModel: + transformer_config = Qwen3Config( + vocab_size=17, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=4, + max_position_embeddings=32, + ) + config = DFlashSpeculatorConfig( + transformer_layer_config=transformer_config, + draft_vocab_size=17, + block_size=4, + aux_hidden_state_layer_ids=[0], + mask_token_id=1, + sample_from_anchor=sample_from_anchor, + speculators_config=SpeculatorsConfig( + algorithm="dflash", + proposal_methods=[GreedyTokenProposalConfig(speculative_tokens=3)], + default_proposal_method="greedy", + verifier=VerifierConfig(name_or_path=None, architectures=[]), + ), + ) + model = DFlashDraftModel(config) + with torch.no_grad(): + head = torch.randn_like(model.lm_head.weight) + model.lm_head.weight.copy_(head) + model.verifier_lm_head.weight.copy_(head) + model.verifier_norm.weight.fill_(1.0) + return model + + +def test_chunked_verifier_argmax_matches_materialized_argmax(): + torch.manual_seed(11) + model = _model() + hidden = torch.randn(1, 11, model.hidden_size) + input_ids = torch.randint(0, model.draft_vocab_size, (1, 11)) + indices = torch.tensor([1, 2, 5, 7, 8]) + + materialized = model._ce_target_ids( + input_ids, + hidden, + indices, + label_source="verifier_argmax", + verifier_argmax_chunk_size=0, + ) + chunked = model._ce_target_ids( + input_ids, + hidden, + indices, + label_source="verifier_argmax", + verifier_argmax_chunk_size=3, + ) + + assert torch.equal(chunked, materialized) + + +def test_input_id_labels_are_explicit_and_position_aligned(): + model = _model() + input_ids = torch.arange(11).unsqueeze(0) + hidden = torch.randn(1, 11, model.hidden_size) + indices = torch.tensor([1, 2, 5, 7, 8]) + + labels = model._ce_target_ids( + input_ids, + hidden, + indices, + label_source="input_ids", + verifier_argmax_chunk_size=0, + ) + + assert torch.equal(labels, input_ids[:, indices]) + + +def test_sample_from_anchor_input_labels_select_the_next_token(): + model = _model(sample_from_anchor=True) + input_ids = torch.arange(11).unsqueeze(0) + hidden = torch.randn(1, 11, model.hidden_size) + indices = torch.tensor([1, 2, 5, 7, 8]) + + labels = model._ce_target_ids( + input_ids, + hidden, + indices, + label_source="input_ids", + verifier_argmax_chunk_size=0, + ) + + assert torch.equal(labels, input_ids[:, indices + 1]) + + +def test_fused_ce_preparation_reuses_ignored_head_at_compute_dtype(): + model = _model() + model.prepare_fused_linear_cross_entropy(torch.bfloat16) + + assert model.lm_head.weight.dtype == torch.float32 + assert model.verifier_lm_head.weight.dtype == torch.bfloat16 + + +def test_fused_ce_preparation_rejects_mismatched_heads(): + model = _model() + with torch.no_grad(): + model.verifier_lm_head.weight[0, 0].add_(1) + + with pytest.raises(RuntimeError, match="LM heads"): + model.prepare_fused_linear_cross_entropy(torch.bfloat16) diff --git a/tests/unit/ops/__init__.py b/tests/unit/ops/__init__.py new file mode 100644 index 000000000..3b2648031 --- /dev/null +++ b/tests/unit/ops/__init__.py @@ -0,0 +1 @@ +"""Unit tests for optional optimized operators.""" diff --git a/tests/unit/ops/test_fused_linear_cross_entropy.py b/tests/unit/ops/test_fused_linear_cross_entropy.py new file mode 100644 index 000000000..90ee17fb4 --- /dev/null +++ b/tests/unit/ops/test_fused_linear_cross_entropy.py @@ -0,0 +1,108 @@ +"""Correctness guards for the DFlash Liger autograd adapter.""" + +from importlib.metadata import PackageNotFoundError +from unittest import mock + +import pytest +import torch +from torch.nn.functional import cross_entropy + +from speculators.ops import fused_linear_cross_entropy as fused_ce + + +def _reference_forward(*, _input, weight, target, **_kwargs): + logits = _input @ weight.t() + grad_logits = torch.softmax(logits, dim=-1) + grad_logits[torch.arange(target.numel()), target] -= 1 + grad_input = grad_logits @ weight + loss = cross_entropy(logits, target, reduction="none") + accuracy = (logits.argmax(dim=-1) == target).float() + return loss, None, accuracy, None, grad_input, None, None + + +def test_arbitrary_per_token_gradient_matches_torch(): + torch.manual_seed(7) + hidden = torch.randn(9, 6, dtype=torch.double, requires_grad=True) + reference_hidden = hidden.detach().clone().requires_grad_(True) + weight = torch.randn(13, 6, dtype=torch.double) + target = torch.randint(0, 13, (9,)) + token_weights = torch.tensor( + [0.0, 0.3, 1.5, 0.0, 2.1, 0.7, 0.2, 1.0, 0.4], dtype=torch.double + ) + + with mock.patch.object( + fused_ce, "_load_liger_forward", return_value=_reference_forward + ): + loss, accuracy = fused_ce.frozen_linear_cross_entropy(hidden, weight, target) + (loss * token_weights).sum().backward() + + reference_logits = reference_hidden @ weight.t() + reference_loss = cross_entropy(reference_logits, target, reduction="none") + (reference_loss * token_weights).sum().backward() + + torch.testing.assert_close(loss, reference_loss) + torch.testing.assert_close( + accuracy, (reference_logits.argmax(dim=-1) == target).float() + ) + torch.testing.assert_close(hidden.grad, reference_hidden.grad) + + +def test_scales_low_precision_gradient_in_fp32_before_cast(): + saved_gradient = torch.tensor( + [[0.0001001358, -0.00331], [0.00091, 0.02111]], dtype=torch.bfloat16 + ) + + def saved_gradient_forward(**kwargs): + count = kwargs["target"].numel() + return ( + torch.ones(count), + None, + torch.ones(count), + None, + saved_gradient, + None, + None, + ) + + hidden = torch.zeros(2, 2, dtype=torch.bfloat16, requires_grad=True) + weight = torch.zeros(3, 2, dtype=torch.bfloat16) + target = torch.tensor([0, 1]) + token_weights = torch.tensor([0.008334385, 1.33791]) + with mock.patch.object( + fused_ce, "_load_liger_forward", return_value=saved_gradient_forward + ): + loss, _ = fused_ce.frozen_linear_cross_entropy(hidden, weight, target) + (loss * token_weights).sum().backward() + + expected = (saved_gradient.float() * token_weights[:, None]).to(torch.bfloat16) + early_cast = saved_gradient * token_weights.to(torch.bfloat16)[:, None] + torch.testing.assert_close(hidden.grad, expected, rtol=0, atol=0) + assert not torch.equal(expected, early_cast) + + +def test_rejects_trainable_head_and_wrong_target_dtype(): + hidden = torch.randn(2, 3, requires_grad=True) + weight = torch.randn(5, 3, requires_grad=True) + target = torch.tensor([1, 2]) + with pytest.raises(ValueError, match="frozen LM head"): + fused_ce.frozen_linear_cross_entropy(hidden, weight, target) + with pytest.raises(ValueError, match="torch.long"): + fused_ce.frozen_linear_cross_entropy(hidden, weight.detach(), target.float()) + + +def test_missing_dependency_and_wrong_version_fail_early(): + fused_ce._load_liger_forward.cache_clear() + with ( + mock.patch.object( + fused_ce, "version", side_effect=PackageNotFoundError("liger-kernel") + ), + pytest.raises(RuntimeError, match=r"speculators\[liger\]"), + ): + fused_ce._load_liger_forward() + + with ( + mock.patch.object(fused_ce, "version", return_value="0.7.0"), + pytest.raises(RuntimeError, match="requires liger-kernel==0.8.0"), + ): + fused_ce._load_liger_forward() + fused_ce._load_liger_forward.cache_clear() diff --git a/tests/unit/ops/test_fused_linear_cross_entropy_cuda.py b/tests/unit/ops/test_fused_linear_cross_entropy_cuda.py new file mode 100644 index 000000000..98c0ae7d1 --- /dev/null +++ b/tests/unit/ops/test_fused_linear_cross_entropy_cuda.py @@ -0,0 +1,46 @@ +"""Opt-in correctness test for the real Liger Triton CE kernel.""" + +import os + +import pytest +import torch +from torch.nn.functional import cosine_similarity, cross_entropy + +from speculators.ops.fused_linear_cross_entropy import frozen_linear_cross_entropy + +pytestmark = pytest.mark.skipif( + os.environ.get("SPECULATORS_RUN_LIGER_TESTS") != "1" + or not torch.cuda.is_available(), + reason="set SPECULATORS_RUN_LIGER_TESTS=1 on an exclusive CUDA GPU", +) + + +def test_bf16_loss_accuracy_and_weighted_hidden_gradient(): + torch.manual_seed(123) + token_count, hidden_size, vocab_size = 67, 128, 1024 + hidden = torch.randn( + token_count, + hidden_size, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + weight = torch.randn(vocab_size, hidden_size, device="cuda", dtype=torch.bfloat16) + target = torch.randint(vocab_size, (token_count,), device="cuda") + token_weights = torch.linspace(0.0, 1.75, token_count, device="cuda") + + loss, accuracy = frozen_linear_cross_entropy(hidden, weight, target) + (loss * token_weights).sum().backward() + fused_grad = hidden.grad.float().clone() + + reference_hidden = hidden.detach().clone().requires_grad_(True) + reference_logits = reference_hidden @ weight.t() + reference_loss = cross_entropy(reference_logits, target, reduction="none") + (reference_loss * token_weights).sum().backward() + reference_grad = reference_hidden.grad.float() + + assert (loss - reference_loss.float()).abs().mean().item() < 0.1 + assert torch.equal(accuracy, (reference_logits.argmax(dim=-1) == target).float()) + cosine = cosine_similarity(fused_grad.flatten(), reference_grad.flatten(), dim=0) + assert cosine.item() > 0.999 + assert torch.isfinite(fused_grad).all() diff --git a/tests/unit/train/test_cli_args.py b/tests/unit/train/test_cli_args.py index 1032f87fa..d0f9bc070 100644 --- a/tests/unit/train/test_cli_args.py +++ b/tests/unit/train/test_cli_args.py @@ -300,3 +300,70 @@ def test_no_norm_before_fc_flag(monkeypatch): def test_no_norm_output_flag(monkeypatch): args = _parse(monkeypatch, ["--no-norm-output"]) assert args.norm_output is False + + +def test_consumer_optimization_defaults_preserve_existing_backends(monkeypatch): + args = _parse(monkeypatch, ["--speculator-type", "dflash"]) + + assert args.dflash_linear_cross_entropy_backend == "torch" + assert not args.dflash_compact_zero_weight_ce_rows + assert args.dflash_label_source == "verifier_argmax" + assert args.dflash_verifier_argmax_chunk_size == 0 + assert args.adamw_backend == "auto" + assert args.gradient_clip_backend == "torch" + assert args.max_grad_norm == 1.0 + + +def test_consumer_optimization_arguments_flow_to_dflash(monkeypatch): + args = _parse( + monkeypatch, + [ + "--speculator-type", + "dflash", + "--loss-fn", + "ce", + "--dflash-linear-cross-entropy-backend", + "liger", + "--dflash-compact-zero-weight-ce-rows", + "--dflash-label-source", + "input_ids", + "--dflash-verifier-argmax-chunk-size", + "512", + "--optimizer", + "adamw", + "--adamw-backend", + "fused", + "--gradient-clip-backend", + "fused_adamw", + "--max-grad-norm", + "0.75", + ], + ) + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + "speculators.models.dflash.core.validate_liger_installation", lambda: None + ) + train_kw, val_kw = DFlashDraftModel.get_trainer_kwargs(**vars(args)) + + assert train_kw["linear_cross_entropy_backend"] == "liger" + assert train_kw["compact_zero_weight_ce_rows"] is True + assert train_kw["label_source"] == "input_ids" + assert train_kw["verifier_argmax_chunk_size"] == 512 + assert val_kw == train_kw + assert args.adamw_backend == "fused" + assert args.gradient_clip_backend == "fused_adamw" + assert args.max_grad_norm == 0.75 + + +@pytest.mark.parametrize( + "extra", + [ + ["--speculator-type", "dflash", "--dflash-compact-zero-weight-ce-rows"], + ["--speculator-type", "dflash", "--dflash-label-source", "input_ids"], + ["--speculator-type", "dflash", "--dflash-verifier-argmax-chunk-size", "-1"], + ["--gradient-clip-backend", "fused_adamw"], + ], +) +def test_consumer_optimization_rejects_invalid_combinations(monkeypatch, extra): + with pytest.raises(SystemExit, match="2"): + _parse(monkeypatch, extra) diff --git a/tests/unit/train/test_optimizers.py b/tests/unit/train/test_optimizers.py new file mode 100644 index 000000000..f40cb5796 --- /dev/null +++ b/tests/unit/train/test_optimizers.py @@ -0,0 +1,175 @@ +"""Focused optimizer execution-backend and fused clipping tests.""" + +from types import SimpleNamespace + +import pytest +import torch + +from speculators.train.optimizers import build_optimizers, restore_adamw_backend +from speculators.train.trainer import Trainer, TrainerConfig + + +def _optimizer_config(**overrides): + values = { + "optimizer": "adamw", + "adamw_backend": "auto", + "lr": 1e-3, + "weight_decay": 0.01, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_foreach_adamw_backend_is_explicit(): + model = torch.nn.Linear(4, 3) + (optimizer,) = build_optimizers(model, _optimizer_config(adamw_backend="foreach")) + + assert optimizer.param_groups[0]["foreach"] is True + assert optimizer.param_groups[0]["fused"] is False + + +def test_fused_adamw_rejects_cpu_parameters(): + with pytest.raises(ValueError, match="CUDA parameters"): + build_optimizers( + torch.nn.Linear(4, 3), _optimizer_config(adamw_backend="fused") + ) + + +@pytest.mark.parametrize( + ("backend", "expected_foreach", "expected_fused"), + [ + ("auto", None, None), + ("foreach", True, False), + ("fused", False, True), + ], +) +def test_restore_adamw_backend_overrides_checkpoint_param_groups( + backend, expected_foreach, expected_fused +): + parameter = torch.nn.Parameter(torch.ones(2)) + checkpoint_optimizer = torch.optim.AdamW([parameter], foreach=True, fused=False) + resumed_optimizer = torch.optim.AdamW([parameter]) + resumed_optimizer.load_state_dict(checkpoint_optimizer.state_dict()) + + assert resumed_optimizer.param_groups[0]["foreach"] is True + assert resumed_optimizer.param_groups[0]["fused"] is False + + restore_adamw_backend([resumed_optimizer], backend) + + assert resumed_optimizer.param_groups[0]["foreach"] is expected_foreach + assert resumed_optimizer.param_groups[0]["fused"] is expected_fused + + +def test_trainer_restores_requested_backend_after_checkpoint_load(): + class _Checkpoint: + previous_epoch = 3 + + @staticmethod + def load_optimizer_state_dict(model, optimizers): + del model + optimizers[0].param_groups[0]["foreach"] = None + optimizers[0].param_groups[0]["fused"] = None + + trainer = Trainer.__new__(Trainer) + trainer.model = torch.nn.Linear(4, 3) + trainer.config = TrainerConfig( + lr=1e-3, + num_epochs=1, + save_path="unused", + resume_from_checkpoint=True, + adamw_backend="foreach", + scheduler_type="none", + ) + trainer.resume_from_checkpoint = True + trainer.checkpointer = _Checkpoint() + + trainer.setup_optimizer() + + assert trainer.optimizers[0].param_groups[0]["foreach"] is True + assert trainer.optimizers[0].param_groups[0]["fused"] is False + + +class _RecordingFusedOptimizer: + def __init__(self, fail: bool = False): + self.param_groups = [{"foreach": False, "fused": True}] + self.fail = fail + self.saw_grad_scale = False + + def step(self): + self.saw_grad_scale = hasattr(self, "grad_scale") + if self.fail: + raise RuntimeError("step failed") + + +def _fused_clip_trainer(optimizer) -> Trainer: + trainer = Trainer.__new__(Trainer) + trainer.model = torch.nn.Linear(3, 2, bias=False) + for parameter in trainer.model.parameters(): + parameter.grad = torch.full_like(parameter, 2.0) + trainer.config = TrainerConfig( + lr=1e-3, + num_epochs=1, + save_path="unused", + optimizer="adamw", + adamw_backend="fused", + gradient_clip_backend="fused_adamw", + max_grad_norm=0.5, + ) + trainer.optimizers = [optimizer] + trainer.local_rank = torch.device("cpu") + return trainer + + +def test_fused_clip_sets_scale_and_cleans_it_after_step(): + optimizer = _RecordingFusedOptimizer() + trainer = _fused_clip_trainer(optimizer) + + norm = trainer._clip_gradients() + expected_scale = ((norm + 1e-6) / trainer.config.max_grad_norm).clamp(min=1) + torch.testing.assert_close(optimizer.grad_scale, expected_scale) + trainer._optimizers_step() + + assert optimizer.saw_grad_scale + assert not hasattr(optimizer, "grad_scale") + + +def test_fused_clip_cleans_scale_when_optimizer_step_fails(): + optimizer = _RecordingFusedOptimizer(fail=True) + trainer = _fused_clip_trainer(optimizer) + trainer._clip_gradients() + + with pytest.raises(RuntimeError, match="step failed"): + trainer._optimizers_step() + + assert not hasattr(optimizer, "grad_scale") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_fused_adamw_clip_matches_explicit_clipping_on_cuda(): + initial = torch.randn(19, device="cuda") + reference = torch.nn.Parameter(initial.clone()) + fused = torch.nn.Parameter(initial.clone()) + gradient = torch.linspace(-3, 4, initial.numel(), device="cuda") + reference.grad = gradient.clone() + fused.grad = gradient.clone() + reference_optimizer = torch.optim.AdamW([reference], lr=3e-4, fused=True) + fused_optimizer = torch.optim.AdamW([fused], lr=3e-4, fused=True) + + torch.nn.utils.clip_grad_norm_([reference], 0.7) + reference_optimizer.step() + grad_norm = torch.nn.utils.get_total_norm([fused.grad], foreach=True) + fused_optimizer.grad_scale = ((grad_norm + 1e-6) / 0.7).clamp(min=1) + try: + fused_optimizer.step() + finally: + del fused_optimizer.grad_scale + + torch.testing.assert_close(fused, reference, rtol=1e-6, atol=1e-7) + torch.testing.assert_close( + fused_optimizer.state[fused]["exp_avg"], + reference_optimizer.state[reference]["exp_avg"], + ) + torch.testing.assert_close( + fused_optimizer.state[fused]["exp_avg_sq"], + reference_optimizer.state[reference]["exp_avg_sq"], + ) From 91b074f29d28a980b3705643575b2b05be070677 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 13:37:59 +0800 Subject: [PATCH 12/20] chore: prepare fanout port for upstream review Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- .../independent_consumer_fanout/README.md | 101 +++--------------- .../config.example.json | 16 +++ docs/cli/train.md | 16 ++- pyproject.toml | 2 +- .../benchmarks/independent_consumers.py | 17 +-- .../data_generation/artifact_cache.py | 4 +- src/speculators/models/attention.py | 3 +- src/speculators/train/data.py | 4 +- src/speculators/train/dataloader.py | 10 +- src/speculators/train/optimizers.py | 55 +++++----- src/speculators/train/trainer.py | 31 ++++-- .../benchmarks/test_independent_consumers.py | 3 +- tests/unit/models/test_dflash_metrics.py | 3 +- .../test_fused_linear_cross_entropy_cuda.py | 2 + tests/unit/train/test_optimizers.py | 12 ++- tests/unit/train/test_shared_artifacts.py | 17 +-- tests/unit/train/test_windowed_training.py | 8 +- 17 files changed, 147 insertions(+), 157 deletions(-) diff --git a/benchmarks/independent_consumer_fanout/README.md b/benchmarks/independent_consumer_fanout/README.md index 70c4a41fe..e29ae0471 100644 --- a/benchmarks/independent_consumer_fanout/README.md +++ b/benchmarks/independent_consumer_fanout/README.md @@ -3,32 +3,13 @@ This fixture runs two fresh-server scenarios in order: 1. one vLLM hidden-state producer and one single-process trainer (`1p1c`); -2. one fresh producer and three independently launched single-process trainers - (`1p3c`). +2. one fresh producer and three independently launched single-process trainers (`1p3c`). -The launcher rejects distributed launchers and any DP, TP, PP, SP, or process-count -option other than one. It also requires distinct physical GPU indices for every role -that overlaps in time. `CUDA_VISIBLE_DEVICES` and distributed rank variables are owned -by the launcher and cannot be supplied by a role config. +The launcher rejects distributed launchers and any DP, TP, PP, SP, or process-count option other than one. It also requires distinct physical GPU indices for every role that overlaps in time. `CUDA_VISIBLE_DEVICES` and distributed rank variables are owned by the launcher and cannot be supplied by a role config. -Each trainer receives its own accounting endpoint through the `{endpoint}` command -placeholder. Other placeholders are `{consumer_id}`, `{output_dir}`, `{scenario}`, and -the scenario-local `{shared_artifacts_dir}`. -The proxy forwards non-streaming OpenAI requests to vLLM and records a completion only -when the response is successful and contains a hidden-state artifact path. It stores -only a digest of the request identity, never the returned artifact path. +Each trainer receives its own accounting endpoint through the `{endpoint}` command placeholder. Other placeholders are `{consumer_id}`, `{output_dir}`, `{scenario}`, and the scenario-local `{shared_artifacts_dir}`. The proxy forwards non-streaming OpenAI requests to vLLM and records a completion only when the response is successful and contains a hidden-state artifact path. It stores only a digest of the request identity, never the returned artifact path. -Start from `config.example.json`, replace the model and preprocessed-data placeholders, -and keep each consumer command as a direct, single-process `scripts/train.py` launch. -Pin training semantics such as `--optimizer` explicitly so an upstream default change -cannot silently alter a comparison. -Keep multiple DataLoader workers and explicit prefetching for generated data. A single -worker can hold only one unique first-miss request in flight, which serializes producer -prefill and can make otherwise independent consumers wait in lockstep. The example uses -four CPU workers with a prefetch factor of two; these workers do not create additional -GPU compute roles, and the benchmark still rejects more than one compute process on any -assigned GPU. -Run the fixture from the repository root: +Start from `config.example.json`, replace the model and preprocessed-data placeholders, and keep each consumer command as a direct, single-process `scripts/train.py` launch. Pin training semantics such as `--optimizer` explicitly so an upstream default change cannot silently alter a comparison. The example also pins fused AdamW and its integrated gradient clipping, which preserve the AdamW update while avoiding a separate gradient-rescaling pass on CUDA. Keep multiple DataLoader workers and explicit prefetching for generated data. A single worker can hold only one unique first-miss request in flight, which serializes producer prefill and can make otherwise independent consumers wait in lockstep. The example uses four CPU workers with a prefetch factor of two; these workers do not create additional GPU compute roles, and the benchmark still rejects more than one compute process on any assigned GPU. Run the fixture from the repository root: ```bash python scripts/benchmark_independent_consumers.py \ @@ -37,76 +18,20 @@ python scripts/benchmark_independent_consumers.py \ --report /tmp/speculators-fanout-report.json ``` -Pass `--scenario 1p3c` to run only the configured 1P3C scenario. Omitting it preserves -the default serial 1P1C-then-1P3C comparison. The report records the scenarios that -were actually selected. +Pass `--scenario 1p3c` to run only the configured 1P3C scenario. Omitting it preserves the default serial 1P1C-then-1P3C comparison. The report records the scenarios that were actually selected. -The run directory must not already exist. Role logs remain there and the compact report -contains the exact command/configuration, package versions, request and valid-completion -counts, shared-sample multiplicity, a common post-warmup throughput window, native -per-consumer `profile/step_ms` summaries, makespan, and role-aware NVML utilization and -memory over the common consumer steady-state overlap. Set -`measurement_steps_per_consumer` to use an exact number of post-warmup steps; otherwise -all available post-warmup steps are measured. Startup samples remain in the JSONL stream -but are excluded from steady-state aggregates. Environment -values are omitted from the report. The command exits nonzero -if a role fails, a GPU is shared or already occupied, a completion is malformed, sample -multiplicity is ambiguous, or the common steady-state window is too small. +The run directory must not already exist. Role logs remain there and the compact report contains the exact command/configuration, package versions, request and valid-completion counts, shared-sample multiplicity, a common post-warmup throughput window, native per-consumer `profile/step_ms` summaries, makespan, and role-aware NVML utilization and memory over the common consumer steady-state overlap. Set `measurement_steps_per_consumer` to use an exact number of post-warmup steps; otherwise all available post-warmup steps are measured. Startup samples remain in the JSONL stream but are excluded from steady-state aggregates. Environment values are omitted from the report. The command exits nonzero if a role fails, a GPU is shared or already occupied, a completion is malformed, sample multiplicity is ambiguous, or the common steady-state window is too small. -The example commands set `--train-data-ratio 1.0`, so every input belongs to the -training stream and no validation pass changes window positions or producer request -accounting. Keep this setting fixed when comparing against a train-only baseline. +The example commands set `--train-data-ratio 1.0`, so every input belongs to the training stream and no validation pass changes window positions or producer request accounting. Keep this setting fixed when comparing against a train-only baseline. -`producer_common_steady` uses that same consumer overlap to report producer requests, -first publications, recaptures, request throughput, and effective unique-sample -throughput. The older `steady_state` field remains the service-wide interval after its -own completion warmup; do not compare that full-run interval with a producer metric -measured only over the common consumer overlap. +`producer_common_steady` uses that same consumer overlap to report producer requests, first publications, recaptures, request throughput, and effective unique-sample throughput. The older `steady_state` field remains the service-wide interval after its own completion warmup; do not compare that full-run interval with a producer metric measured only over the common consumer overlap. -For the unshared baseline, -`expected_service_completions_per_shared_sample` is one for `1p1c` and three for -`1p3c`. A publish-once implementation changes the latter to one; the logical consumer -commands and all other workload settings must remain equivalent. +For the unshared baseline, `expected_service_completions_per_shared_sample` is one for `1p1c` and three for `1p3c`. A publish-once implementation changes the latter to one; the logical consumer commands and all other workload settings must remain equivalent. -To measure publish-once fan-out, pass the same cache to every consumer with -`--shared-hidden-states-path {shared_artifacts_dir}` and set the `1p3c` expected service -multiplicity to one. The report then includes aggregate logical request, hit, miss, -coalesced-waiter, retry, publish, failure, cleanup, and timeout counters under -`shared_artifact_cache`. The run fails closed unless three logical requests correspond -to every service completion, exactly one miss is published, the other two requests hit, -and all failure, retry, cleanup, and timeout counters are zero. Baseline scenarios that -do not use the shared-cache placeholder remain valid without cache accounting. +To measure publish-once fan-out, pass the same cache to every consumer with `--shared-hidden-states-path {shared_artifacts_dir}` and set the `1p3c` expected service multiplicity to one. The report then includes aggregate logical request, hit, miss, coalesced-waiter, retry, publish, failure, cleanup, and timeout counters under `shared_artifact_cache`. The run fails closed unless three logical requests correspond to every service completion, exactly one miss is published, the other two requests hit, and all failure, retry, cleanup, and timeout counters are zero. Baseline scenarios that do not use the shared-cache placeholder remain valid without cache accounting. -For bounded asynchronous fan-out, also pass -`--shared-hidden-states-consumer-id {consumer_id}` and configure lookbehind, -lookahead, max-prefetch, capture-batch, and max-inflight limits. The example config -uses the aligned `2/40/8/8` window and producer-batch settings. The full lookahead is -retained for reuse while only eight PREFETCH requests per consumer can be queued or -generating; demand bypasses that cap. DataLoader workers can prefetch authorized -positions, but only the trainer commits cursor progress. -When consumer windows separate, a publication that leaves every live window is evicted; -a lagging consumer may therefore regenerate that request later. The report accepts one -to `consumer_count` service completions per measured key, records the observed -multiplicity histogram and regeneration overhead, and still requires every successful -service completion to match exactly one cache miss and publication. It also adds -`windowed_artifacts`, including each consumer cursor, current artifact states, retained -bytes, in-flight acquisitions, and retained/in-flight high-water marks. +For bounded asynchronous fan-out, also pass `--shared-hidden-states-consumer-id {consumer_id}` and configure lookbehind, lookahead, max-prefetch, capture-batch, and max-inflight limits. The example config uses the aligned `2/40/8/8` window and producer-batch settings. The full lookahead is retained for reuse while only eight PREFETCH requests per consumer can be queued or generating; demand bypasses that cap. DataLoader workers can prefetch authorized positions, but only the trainer commits cursor progress. When consumer windows separate, a publication that leaves every live window is evicted; a lagging consumer may therefore regenerate that request later. The report accepts one to `consumer_count` service completions per measured key, records the observed multiplicity histogram and regeneration overhead, and still requires every successful service completion to match exactly one cache miss and publication. It also adds `windowed_artifacts`, including each consumer cursor, current artifact states, retained bytes, in-flight acquisitions, and retained/in-flight high-water marks. -The shared cache is a filesystem data plane, not Mooncake or GPU-direct transport. Its -directory must provide reliable POSIX `flock`, same-filesystem atomic rename, and -directory `fsync` semantics to all consumers. Do not use an arbitrary NFS mount unless -those guarantees have been verified. -Without a consumer ID, the legacy cache is not bounded by consumer progress: disabling -expiration retains one artifact per unique request, and a finite TTL can still grow over -a pass of unseen samples. With a consumer ID, retention is instead bounded by the union -of live windows, atomic packed batches, and in-flight leases. The focused 10k/100k CPU -state-machine tests validate that this bound is independent of total stream length. +The shared cache is a filesystem data plane, not Mooncake or GPU-direct transport. Its directory must provide reliable POSIX `flock`, same-filesystem atomic rename, and directory `fsync` semantics to all consumers. Do not use an arbitrary NFS mount unless those guarantees have been verified. Without a consumer ID, the legacy cache is not bounded by consumer progress: disabling expiration retains one artifact per unique request, and a finite TTL can still grow over a pass of unseen samples. With a consumer ID, retention is instead bounded by the union of live windows, atomic packed batches, and in-flight leases. The focused 10k/100k CPU state-machine tests validate that this bound is independent of total stream length. -In publish-once mode, `per_consumer_completions` and the steady-state per-consumer -completion map identify which consumer owned each service miss. They do not represent -logical trainer progress: cache logical-request totals and each independent consumer's -`consumer_step_times` provide that evidence. The report labels both maps with -`service_request_owner` to make this distinction explicit. -Despite their names, `warmup_completions_per_consumer` and -`minimum_steady_completions_per_consumer` are service-wide publication thresholds in -this mode; the corresponding step fields apply separately to every consumer. +In publish-once mode, `per_consumer_completions` and the steady-state per-consumer completion map identify which consumer owned each service miss. They do not represent logical trainer progress: cache logical-request totals and each independent consumer's `consumer_step_times` provide that evidence. The report labels both maps with `service_request_owner` to make this distinction explicit. Despite their names, `warmup_completions_per_consumer` and `minimum_steady_completions_per_consumer` are service-wide publication thresholds in this mode; the corresponding step fields apply separately to every consumer. diff --git a/benchmarks/independent_consumer_fanout/config.example.json b/benchmarks/independent_consumer_fanout/config.example.json index 36d028564..7eca0a32c 100644 --- a/benchmarks/independent_consumer_fanout/config.example.json +++ b/benchmarks/independent_consumer_fanout/config.example.json @@ -46,6 +46,10 @@ "3072", "--optimizer", "adamw", + "--adamw-backend", + "fused", + "--gradient-clip-backend", + "fused_adamw", "--speculator-type", "dflash", "--draft-arch", @@ -122,6 +126,10 @@ "3072", "--optimizer", "adamw", + "--adamw-backend", + "fused", + "--gradient-clip-backend", + "fused_adamw", "--speculator-type", "dflash", "--draft-arch", @@ -186,6 +194,10 @@ "3072", "--optimizer", "adamw", + "--adamw-backend", + "fused", + "--gradient-clip-backend", + "fused_adamw", "--speculator-type", "dflash", "--draft-arch", @@ -250,6 +262,10 @@ "3072", "--optimizer", "adamw", + "--adamw-backend", + "fused", + "--gradient-clip-backend", + "fused_adamw", "--speculator-type", "dflash", "--draft-arch", diff --git a/docs/cli/train.md b/docs/cli/train.md index 97a8db0a9..17edb8064 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -168,6 +168,12 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--optimizer`** (str, default: `"muon"`) Optimizer to use. Options: `adamw`, `muon`. The `muon` option applies the Muon optimizer to 2D weight matrices and AdamW to the remaining parameters (norms, biases, embeddings, lm_head). +- **`--adamw-backend`** (str, default: `"auto"`) AdamW execution backend. Options: `auto`, `foreach`, `fused`. This also applies to the AdamW parameter group in Muon mode. The fused backend requires CUDA parameters. + +- **`--gradient-clip-backend`** (str, default: `"torch"`) Gradient clipping implementation. Options: `torch`, `fused_adamw`. The fused option avoids a separate gradient-scaling kernel and requires both `--optimizer adamw` and `--adamw-backend fused`. + +- **`--max-grad-norm`** (float, default: `1.0`) Maximum gradient norm used by either clipping backend. + - **`--weight-decay`** (float, default: `0.01`) Weight decay for the AdamW optimizer (and the AdamW group in muon mode). - **`--muon-lr`** (float, default: `10*lr`) Learning rate for the Muon (2D weights) group. Only used with `--optimizer muon`. Defaults to 10× the `--lr` value. @@ -208,10 +214,18 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--sample-from-anchor`** / **`--no-sample-from-anchor`** (bool, default: algorithm-specific) Whether to sample from the anchor position. `True`: sample from anchor and all mask positions (default for dspark, produces block_size tokens). `False`: anchor is bonus token (default for dflash, produces block_size-1 tokens). -- **`--max-anchors`** (int, default: `256`) Maximum anchor positions for DFlash training. +- **`--max-anchors`** (int, default: `3072`) Maximum anchor positions for DFlash training. - **`--dflash-decay-gamma`** (float, default: `4.0`) Decay gamma for DFlash loss weighting. +- **`--dflash-linear-cross-entropy-backend`** (str, default: `"torch"`) DFlash cross-entropy backend. Options: `torch`, `liger`. Liger avoids materializing draft logits and requires an exactly-CE loss configuration plus the optional `speculators[liger]` dependency. + +- **`--dflash-compact-zero-weight-ce-rows`** / **`--no-dflash-compact-zero-weight-ce-rows`** (bool, default: `False`) Exclude masked, zero-weight rows before the fused Liger CE kernel. Requires `--dflash-linear-cross-entropy-backend liger`. + +- **`--dflash-label-source`** (str, default: `"verifier_argmax"`) Hard-label source for the opt-in Liger CE path. Options: `verifier_argmax`, `input_ids`. The default preserves DFlash verifier-target semantics; `input_ids` is an explicitly different training target and requires the full verifier vocabulary. + +- **`--dflash-verifier-argmax-chunk-size`** (int, default: `0`) Number of verifier LM-head rows processed per argmax chunk in the Liger CE path. `0` materializes the complete verifier logits; a positive value reduces their peak memory. Requires the Liger CE backend. + ### Sliding Window Attention Arguments All speculator types (except `mtp`) use sliding window attention on all draft layers by default. diff --git a/pyproject.toml b/pyproject.toml index 86bef9534..431452833 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -146,7 +146,7 @@ exclude = ["venv", "build", "dist"] follow_imports = 'silent' [[tool.mypy.overrides]] -module = ["datasets.*", "transformers.*", "setuptools.*", "setuptools_git_versioning.*", "vllm.*"] +module = ["datasets.*", "liger_kernel.*", "transformers.*", "setuptools.*", "setuptools_git_versioning.*", "vllm.*"] ignore_missing_imports=true [tool.ruff] diff --git a/src/speculators/benchmarks/independent_consumers.py b/src/speculators/benchmarks/independent_consumers.py index c6cce1685..1a4c70e22 100644 --- a/src/speculators/benchmarks/independent_consumers.py +++ b/src/speculators/benchmarks/independent_consumers.py @@ -419,8 +419,11 @@ def _forward(self) -> None: port = proxy._target.port or ( 443 if proxy._target.scheme == "https" else 80 ) + hostname = proxy._target.hostname + if hostname is None: + raise ValueError("accounting proxy target must include a host") connection = connection_type( - proxy._target.hostname, + hostname, port, timeout=proxy._timeout, ) @@ -882,7 +885,7 @@ def analyze_consumer_steps( "steady_steps_per_second": None, } if events is None: - values = [] + values: list[float] = [] with log_path.open(errors="replace") as log_file: for line in log_file: values.extend( @@ -987,16 +990,16 @@ def _gpu_snapshot(target_gpus: set[int]) -> _GpuSample: compute_pids: dict[int, list[int]] = defaultdict(list) for row in _run_nvidia_smi("compute-apps=gpu_uuid,pid,used_gpu_memory"): uuid, raw_pid, raw_memory = (part.strip() for part in row.split(",", 2)) - index = uuid_to_index.get(uuid) - if index is None: + gpu_index = uuid_to_index.get(uuid) + if gpu_index is None: continue try: pid = int(raw_pid) memory = int(raw_memory) except ValueError as error: raise EvidenceError(f"Unparseable nvidia-smi compute row: {row}") from error - compute_pids[index].append(pid) - role_memory[index] += memory + compute_pids[gpu_index].append(pid) + role_memory[gpu_index] += memory return _GpuSample( captured_at=time.monotonic(), total_memory_mib=total_memory, @@ -1473,7 +1476,7 @@ def run_benchmark( if not selected: raise ValueError(f"Configured scenarios do not include {scenario_kind!r}") scenarios = [_run_scenario(config, scenario, output_dir) for scenario in selected] - versions = {} + versions: dict[str, str | None] = {} for package in ("speculators", "torch", "vllm"): try: versions[package] = importlib.metadata.version(package) diff --git a/src/speculators/data_generation/artifact_cache.py b/src/speculators/data_generation/artifact_cache.py index 389a64253..b3fadce5f 100644 --- a/src/speculators/data_generation/artifact_cache.py +++ b/src/speculators/data_generation/artifact_cache.py @@ -15,8 +15,6 @@ import torch from safetensors.torch import load_file, save_file -from speculators.data_generation.vllm_client import ClientItem - if TYPE_CHECKING: from collections.abc import Callable, Iterator, Mapping @@ -86,7 +84,7 @@ class ArtifactResult: def canonical_hidden_state_request_id( model: str, - client_item: ClientItem, + client_item: Mapping[str, Any], *, namespace: str | None = None, ) -> str: diff --git a/src/speculators/models/attention.py b/src/speculators/models/attention.py index 4467700a4..98ca6f3ec 100644 --- a/src/speculators/models/attention.py +++ b/src/speculators/models/attention.py @@ -5,6 +5,7 @@ """ from collections.abc import Callable +from typing import cast import torch from torch.nn.attention.flex_attention import ( @@ -61,7 +62,7 @@ def flex_attention_forward( enable_gqa=enable_gqa, scale=scaling, ) - attention_output: torch.Tensor = flex_attention_output + attention_output = cast("torch.Tensor", flex_attention_output) attention_output = attention_output.transpose(1, 2).contiguous() return attention_output, None diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py index 19b8bd6fb..f9a26e3d7 100644 --- a/src/speculators/train/data.py +++ b/src/speculators/train/data.py @@ -730,6 +730,7 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None: dataset_item = self.data[index] client_item = build_client_item(dataset_item) + loaded_hs: dict[str, torch.Tensor] | None try: if self.artifact_cache is not None: request_id = canonical_hidden_state_request_id( @@ -821,6 +822,7 @@ def _load_requested_hidden_states( windowed_sample: StreamSampleIndex | None, ) -> tuple[dict[str, torch.Tensor] | None, ArtifactReadLease | None]: lease: ArtifactReadLease | None = None + loaded_hs: dict[str, torch.Tensor] | None file_idx = self._map_to_file_idx(dataset_index) if windowed_sample is not None: loaded_hs, lease = self._acquire_windowed_hs(windowed_sample) @@ -1020,7 +1022,7 @@ def collate_fn(batch: list[BatchType | None]) -> BatchType: empty = preprocess(empty) batch = [empty] - collated_data = {} + collated_data: BatchType = {} for key in batch[0]: # type: ignore[union-attr] # Concatenate the tensors along the seq (0th) dimension collated_data[key] = torch.cat([b[key] for b in batch], dim=0) # type: ignore[index] diff --git a/src/speculators/train/dataloader.py b/src/speculators/train/dataloader.py index 70a08c6e2..d7e3e352a 100644 --- a/src/speculators/train/dataloader.py +++ b/src/speculators/train/dataloader.py @@ -41,12 +41,20 @@ def configure_windowed_stream(self, sampler: Any) -> str: ... def windowed_request_id(self, dataset_index: int) -> str: ... +class _BatchSampler(Protocol): + epoch: int + + def _generate_batches(self, epoch: int) -> list[Any]: ... + + def set_epoch(self, epoch: int) -> None: ... + + class WindowedBatchSampler: """Annotate sampler indices with stable positions in the consumed order.""" def __init__( self, - sampler: MultipackDistributedBatchSamplerV2, + sampler: _BatchSampler, *, stream_id: str, request_id_for_index: Callable[[int], str], diff --git a/src/speculators/train/optimizers.py b/src/speculators/train/optimizers.py index 5acb9d78c..7036348b6 100644 --- a/src/speculators/train/optimizers.py +++ b/src/speculators/train/optimizers.py @@ -13,6 +13,7 @@ """ import logging +from collections.abc import Iterable import torch from torch import Tensor @@ -29,18 +30,31 @@ _MATRIX_NDIM = 2 -def _adamw_backend_kwargs( - named_params: list[tuple[str, Tensor]], backend: str -) -> dict[str, bool]: - if backend == "auto": - return {} - if backend == "foreach": - return {"foreach": True, "fused": False} - if backend == "fused": - if any(param.device.type != "cuda" for _, param in named_params): +def _build_adamw( + named_params: Iterable[tuple[str, Tensor]], config +) -> torch.optim.AdamW: + params = list(named_params) + if config.adamw_backend == "auto": + return torch.optim.AdamW(params, lr=config.lr, weight_decay=config.weight_decay) + if config.adamw_backend == "foreach": + return torch.optim.AdamW( + params, + lr=config.lr, + weight_decay=config.weight_decay, + foreach=True, + fused=False, + ) + if config.adamw_backend == "fused": + if any(param.device.type != "cuda" for _, param in params): raise ValueError("adamw_backend='fused' requires CUDA parameters") - return {"foreach": False, "fused": True} - raise ValueError(f"Unsupported AdamW backend: {backend!r}") + return torch.optim.AdamW( + params, + lr=config.lr, + weight_decay=config.weight_decay, + foreach=False, + fused=True, + ) + raise ValueError(f"Unsupported AdamW backend: {config.adamw_backend!r}") def restore_adamw_backend( @@ -103,15 +117,7 @@ def build_optimizers(model: Module, config) -> list[torch.optim.Optimizer]: "adamw" returns a single optimizer; "muon" returns ``[Muon, AdamW]``. """ if config.optimizer == "adamw": - named_params = list(model.named_parameters()) - return [ - torch.optim.AdamW( - named_params, - lr=config.lr, - weight_decay=config.weight_decay, - **_adamw_backend_kwargs(named_params, config.adamw_backend), - ) - ] + return [_build_adamw(model.named_parameters(), config)] if config.optimizer == "muon": muon_params, adamw_params = split_named_params_for_muon(model) @@ -134,14 +140,7 @@ def build_optimizers(model: Module, config) -> list[torch.optim.Optimizer]: ) ) if adamw_params: - optimizers.append( - torch.optim.AdamW( - adamw_params, - lr=config.lr, - weight_decay=config.weight_decay, - **_adamw_backend_kwargs(adamw_params, config.adamw_backend), - ) - ) + optimizers.append(_build_adamw(adamw_params, config)) if not optimizers: raise ValueError("No trainable parameters found to optimize.") return optimizers diff --git a/src/speculators/train/trainer.py b/src/speculators/train/trainer.py index 03bd3aa7c..cf682d8c3 100644 --- a/src/speculators/train/trainer.py +++ b/src/speculators/train/trainer.py @@ -4,7 +4,7 @@ import warnings from collections.abc import Callable from pathlib import Path -from typing import Literal, NamedTuple, TypeVar +from typing import Any, Literal, NamedTuple, Protocol, TypeVar, cast import torch import torch.distributed as dist @@ -43,6 +43,18 @@ _T = TypeVar("_T") +class _WindowedDataset(Protocol): + def prepare_windowed_epoch( + self, samples: tuple[Any, ...], *, cursor: int, reset: bool + ) -> None: ... + + def start_windowed_producer(self) -> None: ... + + +class _FusedOptimizer(Protocol): + grad_scale: torch.Tensor + + class _StepTimer: # Each mark()/now() forces a cuda.synchronize to capture true GPU time. # This serialises the CUDA pipeline, so profiled steps are slower; keep @@ -205,10 +217,12 @@ def _prepare_windowed_loader( ): sampler = loader.batch_sampler dataset = loader.dataset - if not hasattr(sampler, "full_epoch_samples") or not hasattr( - dataset, "prepare_windowed_epoch" + if not hasattr(sampler, "full_epoch_samples") or not all( + hasattr(dataset, method) + for method in ("prepare_windowed_epoch", "start_windowed_producer") ): return iter(loader) + windowed_dataset = cast("_WindowedDataset", dataset) batches = ( sampler._generate_batches(epoch) # type: ignore[union-attr] # noqa: SLF001 if full_batches is None @@ -222,14 +236,14 @@ def _prepare_windowed_loader( else: cursor = 0 dataset_id = id(dataset) - dataset.prepare_windowed_epoch( # type: ignore[union-attr] + windowed_dataset.prepare_windowed_epoch( samples, cursor=cursor, reset=dataset_id not in self._prepared_windowed_datasets, ) self._prepared_windowed_datasets.add(dataset_id) iterator = iter(loader) - dataset.start_windowed_producer() # type: ignore[union-attr] + windowed_dataset.start_windowed_producer() return iterator @staticmethod @@ -486,9 +500,10 @@ def _clip_gradients(self): if not gradients: return torch.tensor(0.0) grad_norm = torch.nn.utils.get_total_norm(gradients, foreach=True) - optimizer.grad_scale = ((grad_norm + 1e-6) / self.config.max_grad_norm).clamp( - min=1.0 - ) + fused_optimizer = cast("_FusedOptimizer", optimizer) + fused_optimizer.grad_scale = ( + (grad_norm + 1e-6) / self.config.max_grad_norm + ).clamp(min=1.0) return grad_norm def _optimizers_step(self): diff --git a/tests/unit/benchmarks/test_independent_consumers.py b/tests/unit/benchmarks/test_independent_consumers.py index 8d478d4dc..9d3baaf00 100644 --- a/tests/unit/benchmarks/test_independent_consumers.py +++ b/tests/unit/benchmarks/test_independent_consumers.py @@ -5,6 +5,7 @@ import urllib.request from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path +from typing import Literal import pytest from pydantic import ValidationError @@ -36,7 +37,7 @@ def _consumer(consumer_id: str, gpu: int, command: list[str] | None = None): def _scenario( - kind: str = "1p3c", + kind: Literal["1p1c", "1p3c"] = "1p3c", *, multiplicity: int | None = None, warmup: int = 1, diff --git a/tests/unit/models/test_dflash_metrics.py b/tests/unit/models/test_dflash_metrics.py index 230e3d7c5..90755ce8e 100644 --- a/tests/unit/models/test_dflash_metrics.py +++ b/tests/unit/models/test_dflash_metrics.py @@ -8,6 +8,7 @@ from speculators.models.dflash.metrics import compute_fused_ce_metrics, compute_metrics from speculators.models.metrics import ( + LossConfig, ce_loss, compute_accuracy_multi_step, dflash_loss_decay, @@ -350,7 +351,7 @@ def test_fused_ce_reduction_matches_standard_metrics(weighting): [[0, 1, 1, 1, 0, 1, 0, 1, 0, 1, 1, 0]], dtype=torch.float64 ) loss_weight = 0.37 - loss_config = {"ce": (ce_loss, loss_weight)} + loss_config: LossConfig = {"ce": (ce_loss, loss_weight)} standard_loss, standard_metrics = compute_metrics( logits, diff --git a/tests/unit/ops/test_fused_linear_cross_entropy_cuda.py b/tests/unit/ops/test_fused_linear_cross_entropy_cuda.py index 98c0ae7d1..64caeea6e 100644 --- a/tests/unit/ops/test_fused_linear_cross_entropy_cuda.py +++ b/tests/unit/ops/test_fused_linear_cross_entropy_cuda.py @@ -31,12 +31,14 @@ def test_bf16_loss_accuracy_and_weighted_hidden_gradient(): loss, accuracy = frozen_linear_cross_entropy(hidden, weight, target) (loss * token_weights).sum().backward() + assert hidden.grad is not None fused_grad = hidden.grad.float().clone() reference_hidden = hidden.detach().clone().requires_grad_(True) reference_logits = reference_hidden @ weight.t() reference_loss = cross_entropy(reference_logits, target, reduction="none") (reference_loss * token_weights).sum().backward() + assert reference_hidden.grad is not None reference_grad = reference_hidden.grad.float() assert (loss - reference_loss.float()).abs().mean().item() < 0.1 diff --git a/tests/unit/train/test_optimizers.py b/tests/unit/train/test_optimizers.py index f40cb5796..ddff9ebb9 100644 --- a/tests/unit/train/test_optimizers.py +++ b/tests/unit/train/test_optimizers.py @@ -1,6 +1,7 @@ """Focused optimizer execution-backend and fused clipping tests.""" from types import SimpleNamespace +from typing import Any, cast import pytest import torch @@ -70,7 +71,7 @@ def load_optimizer_state_dict(model, optimizers): optimizers[0].param_groups[0]["foreach"] = None optimizers[0].param_groups[0]["fused"] = None - trainer = Trainer.__new__(Trainer) + trainer = cast("Any", Trainer.__new__(Trainer)) trainer.model = torch.nn.Linear(4, 3) trainer.config = TrainerConfig( lr=1e-3, @@ -90,6 +91,8 @@ def load_optimizer_state_dict(model, optimizers): class _RecordingFusedOptimizer: + grad_scale: torch.Tensor + def __init__(self, fail: bool = False): self.param_groups = [{"foreach": False, "fused": True}] self.fail = fail @@ -102,7 +105,7 @@ def step(self): def _fused_clip_trainer(optimizer) -> Trainer: - trainer = Trainer.__new__(Trainer) + trainer = cast("Any", Trainer.__new__(Trainer)) trainer.model = torch.nn.Linear(3, 2, bias=False) for parameter in trainer.model.parameters(): parameter.grad = torch.full_like(parameter, 2.0) @@ -158,11 +161,12 @@ def test_fused_adamw_clip_matches_explicit_clipping_on_cuda(): torch.nn.utils.clip_grad_norm_([reference], 0.7) reference_optimizer.step() grad_norm = torch.nn.utils.get_total_norm([fused.grad], foreach=True) - fused_optimizer.grad_scale = ((grad_norm + 1e-6) / 0.7).clamp(min=1) + fused_optimizer_with_scale = cast("Any", fused_optimizer) + fused_optimizer_with_scale.grad_scale = ((grad_norm + 1e-6) / 0.7).clamp(min=1) try: fused_optimizer.step() finally: - del fused_optimizer.grad_scale + del fused_optimizer_with_scale.grad_scale torch.testing.assert_close(fused, reference, rtol=1e-6, atol=1e-7) torch.testing.assert_close( diff --git a/tests/unit/train/test_shared_artifacts.py b/tests/unit/train/test_shared_artifacts.py index fc351fb70..421fabe7a 100644 --- a/tests/unit/train/test_shared_artifacts.py +++ b/tests/unit/train/test_shared_artifacts.py @@ -2,7 +2,7 @@ import threading import time -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Literal, cast import hs_connectors.transfer as transfer_module import pytest @@ -55,7 +55,7 @@ def _arrow_dataset( *, shared_path: Path | None, hidden_states_path: Path, - on_generate: str = "delete", + on_generate: Literal["cache", "delete"] = "delete", ) -> ArrowDataset: dataset = ArrowDataset( max_len=128, @@ -68,7 +68,7 @@ def _arrow_dataset( shared_artifacts_namespace=("layers:2,18,33" if shared_path else None), shared_artifacts_ttl_seconds=None, ) - dataset.client = object() + dataset.client = cast("Any", object()) return dataset @@ -205,7 +205,7 @@ def test_windowed_dataset_dispatches_reads_acks_and_cleans_final_window( shared_artifacts_consumer_id="consumer", request_timeout=2, ) - dataset.client = object() + dataset.client = cast("Any", object()) monkeypatch.setattr( dataset, "_materialize_shared_hs", @@ -257,7 +257,7 @@ def test_windowed_scheduling_is_independent_of_dataloader_workers( shared_artifacts_consumer_id="consumer", request_timeout=5, ) - dataset.client = object() + dataset.client = cast("Any", object()) def materialize(_index, dataset_item, _client_item): tokens = dataset_item["input_ids"] @@ -276,13 +276,14 @@ def materialize(_index, dataset_item, _client_item): prefetch_factor=2, ) sampler = loader.batch_sampler + assert isinstance(sampler, WindowedBatchSampler) sampler.set_epoch(0) samples = sampler.full_epoch_samples(0) dataset.prepare_windowed_epoch(samples, cursor=0, reset=True) iterator = iter(loader) dataset.start_windowed_producer() - seen_sequences = [] + seen_sequences: list[int] = [] for batch in iterator: leases = batch.pop(data_module.WINDOWED_BATCH_LEASES_KEY) seen_sequences.extend(lease["sequence"] for lease in leases) @@ -318,7 +319,7 @@ def test_windowed_producer_runs_bounded_concurrent_capture_batches( shared_artifacts_capture_batch_wait_seconds=0, request_timeout=5, ) - dataset.client = object() + dataset.client = cast("Any", object()) lock = threading.Lock() active = 0 peak = 0 @@ -383,7 +384,7 @@ def test_windowed_capture_batch_isolates_one_failed_claim(tmp_path, monkeypatch) shared_artifacts_generation_attempts=1, request_timeout=5, ) - dataset.client = object() + dataset.client = cast("Any", object()) def materialize(index, dataset_item, _client_item): if index == 1: diff --git a/tests/unit/train/test_windowed_training.py b/tests/unit/train/test_windowed_training.py index 9902b0b45..7a4e21dc7 100644 --- a/tests/unit/train/test_windowed_training.py +++ b/tests/unit/train/test_windowed_training.py @@ -2,7 +2,7 @@ import hashlib from dataclasses import dataclass -from typing import Any +from typing import Any, cast import pytest import torch @@ -153,7 +153,7 @@ def _trainer(events: list[str]) -> Trainer: "document_ids": torch.tensor([[0]]), WINDOWED_BATCH_LEASES_KEY: [lease], } - trainer = Trainer.__new__(Trainer) + trainer = cast("Any", Trainer.__new__(Trainer)) trainer.model = _Model() trainer.config = TrainerConfig( lr=0.1, @@ -214,7 +214,7 @@ def operation(epoch: int) -> str: events.append(f"run:{epoch}") return "result" - assert Trainer._run_windowed_phase(loader, operation, 3) == "result" + assert Trainer._run_windowed_phase(cast("Any", loader), operation, 3) == "result" assert events == ["run:3", "stop:True"] @@ -227,5 +227,5 @@ def operation(_epoch: int) -> None: raise RuntimeError("phase failed") with pytest.raises(RuntimeError, match="phase failed"): - Trainer._run_windowed_phase(loader, operation, 0) + Trainer._run_windowed_phase(cast("Any", loader), operation, 0) assert events == ["run", "stop:False"] From 0223dcc43a8fe2f2fbfa4c65fbe317a39ab75717 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 14:14:54 +0800 Subject: [PATCH 13/20] fix: validate Flex Attention return values Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- src/speculators/models/attention.py | 5 +-- tests/unit/models/test_attention.py | 52 +++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) create mode 100644 tests/unit/models/test_attention.py diff --git a/src/speculators/models/attention.py b/src/speculators/models/attention.py index 98ca6f3ec..6dfa9ff9c 100644 --- a/src/speculators/models/attention.py +++ b/src/speculators/models/attention.py @@ -5,7 +5,6 @@ """ from collections.abc import Callable -from typing import cast import torch from torch.nn.attention.flex_attention import ( @@ -62,7 +61,9 @@ def flex_attention_forward( enable_gqa=enable_gqa, scale=scaling, ) - attention_output = cast("torch.Tensor", flex_attention_output) + if not isinstance(flex_attention_output, torch.Tensor): + raise TypeError("Flex Attention unexpectedly returned auxiliary output") + attention_output = flex_attention_output attention_output = attention_output.transpose(1, 2).contiguous() return attention_output, None diff --git a/tests/unit/models/test_attention.py b/tests/unit/models/test_attention.py new file mode 100644 index 000000000..29508b91f --- /dev/null +++ b/tests/unit/models/test_attention.py @@ -0,0 +1,52 @@ +import pytest +import torch + +from speculators.models import attention + + +def test_flex_attention_forward_returns_transposed_tensor( + monkeypatch: pytest.MonkeyPatch, +) -> None: + query = torch.randn(1, 2, 3, 4) + key = torch.randn(1, 1, 3, 4) + value = torch.randn(1, 1, 3, 4) + flex_output = torch.randn(1, 2, 3, 4) + mask = object() + + def fake_flex_attention(*args, **kwargs): + assert args == (query, key, value) + assert kwargs == { + "score_mod": None, + "block_mask": mask, + "enable_gqa": True, + "scale": 0.5, + } + return flex_output + + monkeypatch.setattr(attention, "flex_attention", fake_flex_attention) + + output, weights = attention.flex_attention_forward( + torch.nn.Identity(), query, key, value, mask, scaling=0.5 + ) + + torch.testing.assert_close(output, flex_output.transpose(1, 2)) + assert output.is_contiguous() + assert weights is None + + +def test_flex_attention_forward_rejects_auxiliary_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + tensor = torch.randn(1, 1, 2, 4) + monkeypatch.setattr( + attention, + "flex_attention", + lambda *args, **kwargs: (tensor, torch.zeros(1)), + ) + + with pytest.raises( + TypeError, match="Flex Attention unexpectedly returned auxiliary output" + ): + attention.flex_attention_forward( + torch.nn.Identity(), tensor, tensor, tensor, attention_mask=None + ) From 2af645e2165f9736ffbbd9f5cb16b45a679b0465 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 14:14:54 +0800 Subject: [PATCH 14/20] docs: document fanout telemetry dependency Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- benchmarks/independent_consumer_fanout/README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/benchmarks/independent_consumer_fanout/README.md b/benchmarks/independent_consumer_fanout/README.md index e29ae0471..8a2ae7cca 100644 --- a/benchmarks/independent_consumer_fanout/README.md +++ b/benchmarks/independent_consumer_fanout/README.md @@ -11,6 +11,8 @@ Each trainer receives its own accounting endpoint through the `{endpoint}` comma Start from `config.example.json`, replace the model and preprocessed-data placeholders, and keep each consumer command as a direct, single-process `scripts/train.py` launch. Pin training semantics such as `--optimizer` explicitly so an upstream default change cannot silently alter a comparison. The example also pins fused AdamW and its integrated gradient clipping, which preserve the AdamW update while avoiding a separate gradient-rescaling pass on CUDA. Keep multiple DataLoader workers and explicit prefetching for generated data. A single worker can hold only one unique first-miss request in flight, which serializes producer prefill and can make otherwise independent consumers wait in lockstep. The example uses four CPU workers with a prefetch factor of two; these workers do not create additional GPU compute roles, and the benchmark still rejects more than one compute process on any assigned GPU. Run the fixture from the repository root: +Install the benchmark telemetry dependency with `pip install -e '.[nvml]'`. + ```bash python scripts/benchmark_independent_consumers.py \ benchmarks/independent_consumer_fanout/config.example.json \ From 967ddf589b3a4f1151f9d413718d2a1cbe62b95b Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 17:13:32 +0800 Subject: [PATCH 15/20] fix: address asynchronous fanout review findings Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- .../independent_consumer_fanout/README.md | 2 +- .../config.example.json | 2 +- scripts/benchmark_independent_consumers.py | 15 +- src/speculators/benchmarks/gpu_monitor.py | 93 +++++++++--- .../benchmarks/independent_consumers.py | 54 ++++++- .../data_generation/windowed_artifacts.py | 135 +++++++++++++++++- src/speculators/train/data.py | 129 ++++++++++------- src/speculators/train/dataloader.py | 5 +- tests/unit/benchmarks/test_gpu_monitor.py | 90 ++++++++++++ .../benchmarks/test_independent_consumers.py | 81 +++++++++++ .../test_windowed_artifacts.py | 111 ++++++++++++++ tests/unit/models/test_dflash_optimized_ce.py | 43 +++++- .../ops/test_fused_linear_cross_entropy.py | 4 +- tests/unit/train/test_shared_artifacts.py | 67 ++++++++- 14 files changed, 734 insertions(+), 97 deletions(-) diff --git a/benchmarks/independent_consumer_fanout/README.md b/benchmarks/independent_consumer_fanout/README.md index 8a2ae7cca..36952109b 100644 --- a/benchmarks/independent_consumer_fanout/README.md +++ b/benchmarks/independent_consumer_fanout/README.md @@ -22,7 +22,7 @@ python scripts/benchmark_independent_consumers.py \ Pass `--scenario 1p3c` to run only the configured 1P3C scenario. Omitting it preserves the default serial 1P1C-then-1P3C comparison. The report records the scenarios that were actually selected. -The run directory must not already exist. Role logs remain there and the compact report contains the exact command/configuration, package versions, request and valid-completion counts, shared-sample multiplicity, a common post-warmup throughput window, native per-consumer `profile/step_ms` summaries, makespan, and role-aware NVML utilization and memory over the common consumer steady-state overlap. Set `measurement_steps_per_consumer` to use an exact number of post-warmup steps; otherwise all available post-warmup steps are measured. Startup samples remain in the JSONL stream but are excluded from steady-state aggregates. Environment values are omitted from the report. The command exits nonzero if a role fails, a GPU is shared or already occupied, a completion is malformed, sample multiplicity is ambiguous, or the common steady-state window is too small. +The run directory must not already exist. Role logs remain there and the compact report contains the exact command/configuration, package versions, request and valid-completion counts, shared-sample multiplicity, a common post-warmup throughput window, native per-consumer `profile/step_ms` summaries, makespan, and role-aware NVML utilization and memory over the common consumer steady-state overlap. `peak_total_memory_mib` is device-wide memory, while `peak_role_memory_mib` sums the assigned GPU's compute-process memory; the window is invalid if either process-memory coverage or the minimum two-sample positive-duration coverage is incomplete. Set `measurement_steps_per_consumer` to use an exact number of post-warmup steps; otherwise all available post-warmup steps are measured. Startup samples remain in the JSONL stream but are excluded from steady-state aggregates. Environment values are omitted from the report. The command exits nonzero if a role fails, a GPU is shared or already occupied, a completion is malformed, sample multiplicity is ambiguous, or the common steady-state window is too small. The example commands set `--train-data-ratio 1.0`, so every input belongs to the training stream and no validation pass changes window positions or producer request accounting. Keep this setting fixed when comparing against a train-only baseline. diff --git a/benchmarks/independent_consumer_fanout/config.example.json b/benchmarks/independent_consumer_fanout/config.example.json index 7eca0a32c..af0635e10 100644 --- a/benchmarks/independent_consumer_fanout/config.example.json +++ b/benchmarks/independent_consumer_fanout/config.example.json @@ -315,7 +315,7 @@ "minimum_steady_steps_per_consumer": 50, "measurement_steps_per_consumer": 50, "minimum_shared_samples": 50, - "expected_service_completions_per_shared_sample": 3 + "expected_service_completions_per_shared_sample": 1 } ], "allowed_gpus": [0, 1, 2, 3], diff --git a/scripts/benchmark_independent_consumers.py b/scripts/benchmark_independent_consumers.py index c76f0063f..e6c90ef73 100644 --- a/scripts/benchmark_independent_consumers.py +++ b/scripts/benchmark_independent_consumers.py @@ -16,11 +16,9 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("config", type=Path, help="Strict benchmark JSON config") parser.add_argument( - "--run-directory", type=Path, required=True, help="New directory for role logs" - ) - parser.add_argument( - "--report", type=Path, required=True, help="Path for the compact JSON report" + "--run-directory", type=Path, help="New directory for role logs" ) + parser.add_argument("--report", type=Path, help="Path for the compact JSON report") parser.add_argument( "--validate-only", action="store_true", help="Validate config without launching" ) @@ -29,7 +27,12 @@ def parse_args() -> argparse.Namespace: choices=("1p1c", "1p3c"), help="Run only this scenario; by default both run serially", ) - return parser.parse_args() + args = parser.parse_args() + if not args.validate_only and (args.run_directory is None or args.report is None): + parser.error( + "--run-directory and --report are required unless --validate-only is set" + ) + return args def main() -> int: @@ -37,6 +40,8 @@ def main() -> int: config = load_config(args.config) if args.validate_only: return 0 + if args.run_directory is None or args.report is None: + raise RuntimeError("benchmark output arguments were not validated") report = run_benchmark(config, args.run_directory, scenario_kind=args.scenario) write_report(report, args.report) return 0 if report["valid"] else 1 diff --git a/src/speculators/benchmarks/gpu_monitor.py b/src/speculators/benchmarks/gpu_monitor.py index 589385c65..2f1ece24d 100644 --- a/src/speculators/benchmarks/gpu_monitor.py +++ b/src/speculators/benchmarks/gpu_monitor.py @@ -63,6 +63,19 @@ def _error_text(error: BaseException) -> str: return f"{type(error).__name__}: {message}"[:512] +def _compute_process_memory_used_bytes(row: Mapping[str, Any]) -> int | None: + processes = row.get("compute_processes") + if processes is None: + return None + total = 0 + for process in processes: + value = process.get("used_memory_bytes") + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + return None + total += value + return total + + class NvmlBackend: """Lazy wrapper that never initializes CUDA or Torch.""" @@ -217,6 +230,17 @@ def _close_backend(self) -> None: finally: self._backend_open = False + def _close_output(self) -> None: + if self._output is None: + return + try: + self._output.close() + except OSError as error: + with self._collection_lock: + self._errors.append(_error_text(error)) + finally: + self._output = None + def _write(self, value: Mapping[str, Any]) -> None: output = self._output if output is None: @@ -232,28 +256,34 @@ def start(self) -> None: self.sample_path.parent.mkdir(parents=True, exist_ok=True) self._output = self.sample_path.open("x", encoding="utf-8", buffering=1) try: + # Treat open as attempted before invoking the backend so partial NVML + # initialization is also closed when a custom backend raises. + self._backend_open = True self._devices = dict( self.backend.open([value.gpu for value in self.assignments]) ) - self._backend_open = True missing = {value.gpu for value in self.assignments} - set(self._devices) if missing: raise GpuMonitorError(f"NVML backend omitted GPUs {sorted(missing)}") + self._started_at_ns = time.monotonic_ns() + self._write( + { + "record_type": "session_start", + "timestamp_monotonic_ns": self._started_at_ns, + "poll_seconds": self.poll_seconds, + "assignments": [asdict(value) for value in self.assignments], + "devices": [asdict(value) for value in self._devices.values()], + } + ) + self._thread.start() except Exception: - self._output.close() - self._output = None + self._stop.set() + if self._thread.is_alive(): + self._thread.join(timeout=max(10.0, self.poll_seconds * 4)) + if not self._thread.is_alive(): + self._close_backend() + self._close_output() raise - self._started_at_ns = time.monotonic_ns() - self._write( - { - "record_type": "session_start", - "timestamp_monotonic_ns": self._started_at_ns, - "poll_seconds": self.poll_seconds, - "assignments": [asdict(value) for value in self.assignments], - "devices": [asdict(value) for value in self._devices.values()], - } - ) - self._thread.start() def _run(self) -> None: # noqa: C901 try: @@ -332,8 +362,7 @@ def stop(self) -> dict[str, Any]: "timestamp_monotonic_ns": self._ended_at_ns, } ) - self._output.close() - self._output = None + self._close_output() summary = { "status": ( "ok" if not self._errors and not self._violations else "degraded" @@ -429,17 +458,38 @@ def summarize_gpu_window( ] timestamps = [int(value["timestamp_monotonic_ns"]) for value in gpu_rows] process_counts = [len(value.get("compute_pids", ())) for value in gpu_rows] - memory = [ + device_memory = [ int(value["memory_used_bytes"]) for value in gpu_rows if value.get("memory_used_bytes") is not None ] + process_memory = [ + value + for row in gpu_rows + if (value := _compute_process_memory_used_bytes(row)) is not None + ] + unavailable_process_memory = len(gpu_rows) - len(process_memory) if not gpu_rows: invalid_reasons.append(f"GPU {gpu} has no samples in the steady window") + elif len(gpu_rows) < _MIN_COVERAGE_SAMPLES: + invalid_reasons.append( + f"GPU {gpu} has {len(gpu_rows)} sample(s); at least " + f"{_MIN_COVERAGE_SAMPLES} are required" + ) + elif max(timestamps) <= min(timestamps): + invalid_reasons.append(f"GPU {gpu} has zero sample coverage") if gpu_rows and not any(process_counts): invalid_reasons.append( f"GPU {gpu} has no compute process in the steady window" ) + if unavailable_process_memory: + invalid_reasons.append( + f"GPU {gpu} lacks compute-process memory for " + f"{unavailable_process_memory}/{len(gpu_rows)} sample(s)" + ) + max_device_memory_mib = ( + max(device_memory) / (1 << 20) if device_memory else None + ) per_gpu[str(gpu)] = { **asdict(by_gpu[gpu]), "sample_count": len(gpu_rows), @@ -459,7 +509,14 @@ def summarize_gpu_window( else None ), }, - "max_memory_used_mib": max(memory) / (1 << 20) if memory else None, + # Preserve this legacy field as device-wide memory. + "max_memory_used_mib": max_device_memory_mib, + "max_device_memory_used_mib": max_device_memory_mib, + "max_compute_process_memory_used_mib": ( + max(process_memory) / (1 << 20) if process_memory else None + ), + "compute_process_memory_sample_count": len(process_memory), + "compute_process_memory_unavailable_samples": (unavailable_process_memory), "max_compute_processes": max(process_counts, default=0), } return { diff --git a/src/speculators/benchmarks/independent_consumers.py b/src/speculators/benchmarks/independent_consumers.py index 1a4c70e22..50b0d8a3b 100644 --- a/src/speculators/benchmarks/independent_consumers.py +++ b/src/speculators/benchmarks/independent_consumers.py @@ -60,6 +60,11 @@ class EvidenceError(RuntimeError): "srun", "torchrun", } +_DISTRIBUTED_LAUNCHER_MODULES = { + "accelerate.commands.launch", + "torch.distributed.launch", + "torch.distributed.run", +} _PARALLEL_SIZE_OPTIONS = { "-dp", "-pp", @@ -88,16 +93,30 @@ def _option_value(command: list[str], index: int) -> tuple[str, int]: return command[index + 1], index + 1 +def _command_has_option(command: list[str], option: str) -> bool: + return any(token.split("=", 1)[0] == option for token in command) + + +def _python_module(command: list[str]) -> str | None: + for index, argument in enumerate(command[1:-1], start=1): + if argument == "-m": + return command[index + 1] + if argument == "--" or not argument.startswith("-"): + return None + return None + + def _validate_single_process_command(command: list[str]) -> list[str]: executable = Path(command[0]).name if executable in _DISTRIBUTED_LAUNCHERS: raise ValueError( f"Distributed launcher {executable!r} is not an independent consumer" ) - if command[1:3] == ["-m", "torch.distributed.run"]: - raise ValueError("torch.distributed.run is not an independent consumer") if not executable.startswith("python"): raise ValueError("Benchmark roles must be direct Python commands") + module = _python_module(command) + if module in _DISTRIBUTED_LAUNCHER_MODULES: + raise ValueError(f"{module} is not an independent consumer") if "--fsdp-shard" in command: raise ValueError("--fsdp-shard requires one distributed trainer") @@ -198,6 +217,19 @@ def validate_consumers(self) -> ScenarioSpec: "Expected service multiplicity must be either publish-once (1) or " f"one completion per consumer ({expected_consumers})" ) + windowed = any( + _command_has_option(consumer.command, "--shared-hidden-states-consumer-id") + for consumer in self.consumers + ) + if ( + windowed + and expected_consumers > 1 + and self.expected_service_completions_per_shared_sample != 1 + ): + raise ValueError( + "Windowed multi-consumer scenarios require publish-once service " + "multiplicity (1)" + ) return self @@ -1089,11 +1121,13 @@ def terminate(self, grace_seconds: float = 20.0) -> None: self.finished_at = self.finished_at or time.monotonic() self.close_log() return - os.killpg(self.process.pid, signal.SIGTERM) + with suppress(ProcessLookupError): + os.killpg(self.process.pid, signal.SIGTERM) try: self.process.wait(timeout=grace_seconds) except subprocess.TimeoutExpired: - os.killpg(self.process.pid, signal.SIGKILL) + with suppress(ProcessLookupError): + os.killpg(self.process.pid, signal.SIGKILL) self.process.wait(timeout=10) self.finished_at = time.monotonic() self.close_log() @@ -1167,7 +1201,7 @@ def _run_scenario( # noqa: C901 for value in (*consumer.command, *consumer.env.values()) ) windowed_artifacts_enabled = any( - "--shared-hidden-states-consumer-id" in consumer.command + _command_has_option(consumer.command, "--shared-hidden-states-consumer-id") for consumer in scenario.consumers ) started_at = time.monotonic() @@ -1405,8 +1439,14 @@ def capture_step( "per_gpu": { gpu: { "baseline_memory_mib": baseline.get(int(gpu)), - "peak_total_memory_mib": value["max_memory_used_mib"], - "peak_role_memory_mib": value["max_memory_used_mib"], + "peak_total_memory_mib": value["max_device_memory_used_mib"], + "peak_role_memory_mib": value["max_compute_process_memory_used_mib"], + "process_memory_sample_count": value[ + "compute_process_memory_sample_count" + ], + "process_memory_unavailable_samples": value[ + "compute_process_memory_unavailable_samples" + ], "max_compute_processes": value["max_compute_processes"], "compute_process_observed": value["max_compute_processes"] > 0, } diff --git a/src/speculators/data_generation/windowed_artifacts.py b/src/speculators/data_generation/windowed_artifacts.py index 757570d94..9e8275735 100644 --- a/src/speculators/data_generation/windowed_artifacts.py +++ b/src/speculators/data_generation/windowed_artifacts.py @@ -478,6 +478,11 @@ def register_positions(self, samples: Sequence[StreamSampleIndex]) -> None: ).fetchall() for row in consumers: self._refresh_window_locked(conn, row["consumer_id"]) + self._prune_positions_locked( + conn, + stream_id, + incoming_start=min(sample.sequence for sample in samples), + ) def register_consumer( self, @@ -767,6 +772,53 @@ def _prune_orphaned_locked(self, conn: sqlite3.Connection) -> None: ), ) + def _prune_positions_locked( + self, + conn: sqlite3.Connection, + stream_id: str, + *, + incoming_start: int | None = None, + ) -> int: + active = conn.execute( + "SELECT cursor,lookbehind FROM consumers WHERE stream_id=? " + "AND state='active'", + (stream_id,), + ).fetchall() + if active: + boundary = min( + max(0, int(row["cursor"]) - int(row["lookbehind"])) for row in active + ) + if incoming_start is not None: + boundary = min(boundary, incoming_start) + elif incoming_start is not None: + # A new epoch is registered before its first consumer is reactivated. + boundary = incoming_start + else: + inactive = conn.execute( + "SELECT cursor FROM consumers WHERE stream_id=?", + (stream_id,), + ).fetchall() + if not inactive: + return 0 + boundary = min(int(row["cursor"]) for row in inactive) + if boundary <= 0: + return 0 + result = conn.execute( + "DELETE FROM positions WHERE stream_id=? AND batch_end_sequence<=? " + "AND NOT EXISTS (SELECT 1 FROM interests i " + "WHERE i.stream_id=positions.stream_id " + "AND i.sequence=positions.sequence) " + "AND NOT EXISTS (SELECT 1 FROM acquisitions a " + "WHERE a.stream_id=positions.stream_id " + "AND a.sequence=positions.sequence) " + "AND NOT EXISTS (SELECT 1 FROM completed_positions cp " + "JOIN consumers c ON c.consumer_id=cp.consumer_id " + "WHERE c.stream_id=positions.stream_id " + "AND cp.sequence=positions.sequence)", + (stream_id, boundary), + ) + return result.rowcount + def heartbeat(self, consumer_id: str) -> None: with self._transaction() as conn: result = conn.execute( @@ -1295,6 +1347,86 @@ def claim_generation( ) return tuple(claims) + def renew_generation_claims( + self, owner: str, claims: Sequence[GenerationClaim] + ) -> int: + """Extend live generation leases without changing their generation.""" + if not owner: + raise ValueError("generation owner must be non-empty") + if not claims: + return 0 + with self._transaction() as conn: + now = self._clock() + for claim in claims: + result = conn.execute( + "UPDATE artifacts SET claim_until=?,updated_at=? " + "WHERE request_id=? AND state=? AND claim_owner=? " + "AND generation=?", + ( + now + self.claim_timeout_seconds, + now, + claim.request_id, + ArtifactState.GENERATING.value, + owner, + claim.generation, + ), + ) + if result.rowcount != 1: + raise WindowedArtifactError( + f"stale generation renewal for {claim.request_id}" + ) + return len(claims) + + def release_generation_claims( + self, owner: str, claims: Sequence[GenerationClaim] + ) -> int: + """Release unfinished work during cooperative producer shutdown.""" + if not owner: + raise ValueError("generation owner must be non-empty") + if not claims: + return 0 + released = 0 + with self._transaction() as conn: + now = self._clock() + for claim in claims: + row = conn.execute( + "SELECT state,claim_owner,generation FROM artifacts " + "WHERE request_id=?", + (claim.request_id,), + ).fetchone() + if row is None or ( + row["state"], + row["claim_owner"], + int(row["generation"]), + ) != (ArtifactState.GENERATING.value, owner, claim.generation): + continue + interested = conn.execute( + "SELECT 1 FROM interests WHERE request_id=? LIMIT 1", + (claim.request_id,), + ).fetchone() + priority = ( + self._retry_priority_locked(conn, claim.request_id) + if interested is not None + else None + ) + conn.execute( + "UPDATE artifacts SET state=?,generation=generation+1," + "priority=?,queued_at=?,claim_owner=NULL,claim_until=NULL," + "last_error=NULL,updated_at=? WHERE request_id=?", + ( + ArtifactState.QUEUED.value + if interested is not None + else ArtifactState.ABSENT.value, + int(priority) if priority is not None else None, + now if interested is not None else None, + now, + claim.request_id, + ), + ) + released += 1 + self._prune_orphaned_locked(conn) + return released + def _recover_claims_locked(self, conn: sqlite3.Connection) -> None: now = self._clock() rows = conn.execute( @@ -1503,7 +1635,7 @@ def finish_eviction(self, claim: EvictionClaim, *, removed: bool) -> None: def complete_consumer(self, consumer_id: str) -> None: with self._transaction() as conn: row = conn.execute( - "SELECT 1 FROM consumers WHERE consumer_id=?", (consumer_id,) + "SELECT stream_id FROM consumers WHERE consumer_id=?", (consumer_id,) ).fetchone() if row is None: raise KeyError(f"unknown consumer {consumer_id!r}") @@ -1518,6 +1650,7 @@ def complete_consumer(self, consumer_id: str) -> None: (self._clock(), consumer_id), ) self._prune_orphaned_locked(conn) + self._prune_positions_locked(conn, row["stream_id"]) def snapshot(self) -> dict[str, Any]: with self._lock: diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py index f9a26e3d7..098d10bbc 100644 --- a/src/speculators/train/data.py +++ b/src/speculators/train/data.py @@ -589,64 +589,88 @@ def _run_windowed_producer(self) -> None: try: if self.client is None: self._setup_client() - with ( - self._new_windowed_coordinator() as coordinator, - ThreadPoolExecutor( - max_workers=self.shared_artifacts_capture_batch_size, - thread_name_prefix="artifact-capture", - ) as executor, - ): - while not self._windowed_producer_stop.is_set(): - coordinator.heartbeat(self.shared_artifacts_consumer_id) - coordinator.recover_expired() - self._evict_windowed_artifacts(coordinator, cache) - if self._windowed_producer_stop.wait( - self.shared_artifacts_capture_batch_wait_seconds - ): - break - claims = coordinator.claim_generation( - owner, - stream_id=self._windowed_stream_id, - max_claims=self.shared_artifacts_capture_batch_size, - max_active_claims=self.shared_artifacts_capture_batch_size, - ) - if claims: - futures = { - executor.submit( - self._produce_windowed_claim, - coordinator, - cache, - owner, - claim, - ): claim - for claim in claims - } - pending = set(futures) - while pending: - done, pending = wait( - pending, - timeout=1.0, - return_when=FIRST_COMPLETED, + executor = ThreadPoolExecutor( + max_workers=self.shared_artifacts_capture_batch_size, + thread_name_prefix="artifact-capture", + ) + try: + with self._new_windowed_coordinator() as coordinator: + while not self._windowed_producer_stop.is_set(): + coordinator.heartbeat(self.shared_artifacts_consumer_id) + coordinator.recover_expired() + self._evict_windowed_artifacts(coordinator, cache) + if self._windowed_producer_stop.wait( + self.shared_artifacts_capture_batch_wait_seconds + ): + break + claims = coordinator.claim_generation( + owner, + stream_id=self._windowed_stream_id, + max_claims=self.shared_artifacts_capture_batch_size, + max_active_claims=self.shared_artifacts_capture_batch_size, + ) + if claims: + self._run_windowed_claim_batch( + coordinator, cache, executor, owner, claims ) - coordinator.heartbeat(self.shared_artifacts_consumer_id) - for future in done: - claim = futures[future] - try: - future.result() - except Exception as error: # noqa: BLE001 - coordinator.fail_generation(owner, claim, error) - continue - self._windowed_producer_stop.wait(0.02) + continue + self._windowed_producer_stop.wait(0.02) + finally: + executor.shutdown(wait=False, cancel_futures=True) except Exception as error: # noqa: BLE001 - background thread boundary self._windowed_producer_error = error - def _produce_windowed_claim( + def _run_windowed_claim_batch( self, coordinator: WindowedArtifactCoordinator, cache: HiddenStateArtifactCache, + executor: ThreadPoolExecutor, owner: str, - claim: GenerationClaim, + claims: tuple[GenerationClaim, ...], ) -> None: + stop = self._windowed_producer_stop + consumer_id = self.shared_artifacts_consumer_id + if stop is None or consumer_id is None: + raise RuntimeError("windowed producer lifecycle is unavailable") + pending = { + executor.submit(self._produce_windowed_claim, cache, claim): claim + for claim in claims + } + renewal_wait = min(1.0, self.shared_artifacts_claim_timeout_seconds / 3.0) + try: + while pending and not stop.is_set(): + done, _ = wait( + set(pending), + timeout=renewal_wait, + return_when=FIRST_COMPLETED, + ) + coordinator.heartbeat(consumer_id) + for future in done: + claim = pending.pop(future) + try: + path, size_bytes = future.result() + except Exception as error: # noqa: BLE001 + coordinator.fail_generation(owner, claim, error) + else: + coordinator.complete_generation( + owner, + claim, + path=path, + size_bytes=size_bytes, + ) + if pending: + coordinator.renew_generation_claims(owner, tuple(pending.values())) + finally: + if pending: + for future in pending: + future.cancel() + coordinator.release_generation_claims(owner, tuple(pending.values())) + + def _produce_windowed_claim( + self, + cache: HiddenStateArtifactCache, + claim: GenerationClaim, + ) -> tuple[Path, int]: expected_request_id = self.windowed_request_id(claim.dataset_index) if expected_request_id != claim.request_id: raise RuntimeError("generation claim no longer matches dataset") @@ -661,12 +685,7 @@ def _produce_windowed_claim( ), lambda data: check_hidden_states(data, dataset_item["input_ids"].tolist()), ) - coordinator.complete_generation( - owner, - claim, - path=result.path, - size_bytes=result.path.stat().st_size, - ) + return result.path, result.path.stat().st_size @staticmethod def _evict_windowed_artifacts( diff --git a/src/speculators/train/dataloader.py b/src/speculators/train/dataloader.py index d7e3e352a..59afcd126 100644 --- a/src/speculators/train/dataloader.py +++ b/src/speculators/train/dataloader.py @@ -250,6 +250,7 @@ def create_train_val_loaders( hidden_states_dtype=hidden_states_dtype, ) else: + dp_rank = get_dp_rank() train_dataset = ArrowDataset( datapath=data_path, max_len=total_seq_len, @@ -270,7 +271,7 @@ def create_train_val_loaders( shared_artifacts_lock_timeout_seconds ), shared_artifacts_consumer_id=( - f"{shared_artifacts_consumer_id}:train" + f"{shared_artifacts_consumer_id}:dp{dp_rank}:train" if shared_artifacts_consumer_id is not None else None ), @@ -312,7 +313,7 @@ def create_train_val_loaders( shared_artifacts_lock_timeout_seconds ), shared_artifacts_consumer_id=( - f"{shared_artifacts_consumer_id}:val" + f"{shared_artifacts_consumer_id}:dp{dp_rank}:val" if shared_artifacts_consumer_id is not None else None ), diff --git a/tests/unit/benchmarks/test_gpu_monitor.py b/tests/unit/benchmarks/test_gpu_monitor.py index aa1973b87..5c418f1ca 100644 --- a/tests/unit/benchmarks/test_gpu_monitor.py +++ b/tests/unit/benchmarks/test_gpu_monitor.py @@ -2,6 +2,9 @@ import json import threading +from unittest import mock + +import pytest from speculators.benchmarks.gpu_monitor import ( GpuDevice, @@ -72,6 +75,44 @@ def test_monitor_streams_nvml_samples_and_finalizes_summary(tmp_path): assert json.loads(summary_path.read_text())["status"] == "ok" +def test_start_write_failure_closes_backend_and_output(tmp_path, monkeypatch): + backend = _FakeBackend() + monitor = GpuMonitor( + [GpuRoleAssignment(0, "producer", "producer")], + tmp_path / "samples.jsonl", + tmp_path / "summary.json", + backend=backend, + ) + monkeypatch.setattr(monitor, "_write", mock.Mock(side_effect=OSError("disk full"))) + + with pytest.raises(OSError, match="disk full"): + monitor.start() + + assert backend.closed + assert monitor._output is None + assert not monitor._thread.is_alive() + + +def test_start_thread_failure_closes_backend_and_output(tmp_path, monkeypatch): + backend = _FakeBackend() + monitor = GpuMonitor( + [GpuRoleAssignment(0, "producer", "producer")], + tmp_path / "samples.jsonl", + tmp_path / "summary.json", + backend=backend, + ) + monkeypatch.setattr( + monitor._thread, "start", mock.Mock(side_effect=RuntimeError("no thread")) + ) + + with pytest.raises(RuntimeError, match="no thread"): + monitor.start() + + assert backend.closed + assert monitor._output is None + assert not monitor._thread.is_alive() + + def test_stop_timeout_does_not_race_a_blocked_sample(tmp_path, monkeypatch): backend = _BlockingBackend() monitor = GpuMonitor( @@ -105,6 +146,7 @@ def test_window_summary_excludes_startup_and_reports_active_time(): "utilization_gpu_pct": utilization, "memory_used_bytes": 10 << 30, "compute_pids": [123], + "compute_processes": [{"pid": 123, "used_memory_bytes": 8 << 30}], } for timestamp, utilization in ( (1_000_000_000, 0), @@ -127,3 +169,51 @@ def test_window_summary_excludes_startup_and_reports_active_time(): assert gpu["gpu_utilization_pct"]["p50"] == 80.0 assert gpu["gpu_utilization_pct"]["p95"] == 100.0 assert gpu["max_compute_processes"] == 1 + assert gpu["max_memory_used_mib"] == 10 << 10 + assert gpu["max_device_memory_used_mib"] == 10 << 10 + assert gpu["max_compute_process_memory_used_mib"] == 8 << 10 + assert gpu["compute_process_memory_sample_count"] == 2 + assert gpu["compute_process_memory_unavailable_samples"] == 0 + + +def test_window_summary_rejects_one_sample_zero_coverage_and_missing_memory(): + assignment = GpuRoleAssignment(2, "consumer:b4", "consumer:b4") + sample = { + "timestamp_monotonic_ns": 2_000_000_000, + "gpu": 2, + "utilization_gpu_pct": 80, + "memory_used_bytes": 10 << 30, + "compute_pids": [123], + "compute_processes": [{"pid": 123, "used_memory_bytes": 8 << 30}], + } + one = summarize_gpu_window( + [sample], + [assignment], + start_monotonic_ns=1_000_000_000, + end_monotonic_ns=3_000_000_000, + ) + assert not one["valid"] + assert "at least 2" in one["invalid_reasons"][0] + + zero = summarize_gpu_window( + [sample, dict(sample)], + [assignment], + start_monotonic_ns=1_000_000_000, + end_monotonic_ns=3_000_000_000, + ) + assert not zero["valid"] + assert any("zero sample coverage" in reason for reason in zero["invalid_reasons"]) + + missing = dict(sample, timestamp_monotonic_ns=3_000_000_000) + missing["compute_processes"] = [{"pid": 123, "used_memory_bytes": None}] + incomplete = summarize_gpu_window( + [sample, missing], + [assignment], + start_monotonic_ns=1_000_000_000, + end_monotonic_ns=3_000_000_000, + ) + assert not incomplete["valid"] + assert any( + "lacks compute-process memory" in reason + for reason in incomplete["invalid_reasons"] + ) diff --git a/tests/unit/benchmarks/test_independent_consumers.py b/tests/unit/benchmarks/test_independent_consumers.py index 9d3baaf00..55d6eecc9 100644 --- a/tests/unit/benchmarks/test_independent_consumers.py +++ b/tests/unit/benchmarks/test_independent_consumers.py @@ -11,6 +11,7 @@ from pydantic import ValidationError import speculators.benchmarks.independent_consumers as benchmark_module +from scripts import benchmark_independent_consumers as benchmark_cli from speculators.benchmarks.independent_consumers import ( AccountingLedger, AccountingProxy, @@ -140,6 +141,7 @@ def test_example_config_uses_train_only_workloads(): / "config.example.json" ) config = json.loads(config_path.read_text()) + validated = benchmark_module.load_config(config_path) commands = [ consumer["command"] @@ -150,6 +152,22 @@ def test_example_config_uses_train_only_workloads(): for command in commands: ratio_index = command.index("--train-data-ratio") assert command[ratio_index + 1] == "1.0" + assert validated.scenarios[1].expected_service_completions_per_shared_sample == 1 + + +def test_benchmark_cli_output_paths_are_conditional(monkeypatch, tmp_path): + config = tmp_path / "config.json" + monkeypatch.setattr( + "sys.argv", + ["benchmark_independent_consumers.py", str(config), "--validate-only"], + ) + args = benchmark_cli.parse_args() + assert args.run_directory is None + assert args.report is None + + monkeypatch.setattr("sys.argv", ["benchmark_independent_consumers.py", str(config)]) + with pytest.raises(SystemExit, match="2"): + benchmark_cli.parse_args() @pytest.mark.parametrize( @@ -157,6 +175,8 @@ def test_example_config_uses_train_only_workloads(): [ ["torchrun", "--nproc-per-node", "3", "trainer.py"], ["python", "-m", "torch.distributed.run", "trainer.py"], + ["python", "-m", "torch.distributed.launch", "trainer.py"], + ["python", "-u", "-m", "accelerate.commands.launch", "trainer.py"], ["python", "trainer.py", "--tensor-parallel-size=3"], ["python", "trainer.py", "-tp", "3"], ["python", "trainer.py", "--data-parallel-size", "3"], @@ -172,6 +192,67 @@ def test_config_rejects_distributed_consumer_commands(command): _consumer("c0", 1, command) +def test_windowed_1p3c_requires_publish_once_multiplicity(): + consumers = [ + _consumer( + f"c{index}", + index + 1, + [ + "python", + "trainer.py", + f"--shared-hidden-states-consumer-id=c{index}", + ], + ) + for index in range(3) + ] + with pytest.raises(ValidationError, match="publish-once"): + ScenarioSpec( + kind="1p3c", + consumers=consumers, + expected_service_completions_per_shared_sample=3, + ) + scenario = ScenarioSpec( + kind="1p3c", + consumers=consumers, + expected_service_completions_per_shared_sample=1, + ) + assert scenario.expected_service_completions_per_shared_sample == 1 + + +@pytest.mark.parametrize("times_out", [False, True]) +def test_managed_process_termination_tolerates_exit_races(monkeypatch, times_out): + class _Process: + pid = 123 + + def __init__(self): + self.wait_calls = 0 + + def poll(self): + return None + + def wait(self, timeout): + self.wait_calls += 1 + if times_out and self.wait_calls == 1: + raise benchmark_module.subprocess.TimeoutExpired("consumer", timeout) + return 0 + + managed = object.__new__(benchmark_module._ManagedProcess) + managed.process = _Process() + managed.finished_at = None + closed = [] + managed.close_log = lambda: closed.append(True) + + def process_exited(*_args): + raise ProcessLookupError + + monkeypatch.setattr(benchmark_module.os, "killpg", process_exited) + managed.terminate(grace_seconds=0.01) + + assert managed.process.wait_calls == (2 if times_out else 1) + assert managed.finished_at is not None + assert closed == [True] + + def test_config_rejects_gpu_sharing_and_owned_environment(): with pytest.raises(ValidationError, match="more than one consumer"): ScenarioSpec( diff --git a/tests/unit/data_generation/test_windowed_artifacts.py b/tests/unit/data_generation/test_windowed_artifacts.py index f0c9ea198..b35da893f 100644 --- a/tests/unit/data_generation/test_windowed_artifacts.py +++ b/tests/unit/data_generation/test_windowed_artifacts.py @@ -625,6 +625,117 @@ def test_expired_producer_claim_is_reassigned_with_bounded_attempts(tmp_path): assert coordinator.snapshot()["artifact_states"] == {"failed": 1} +def test_generation_claim_renewal_prevents_live_work_reassignment(tmp_path): + now = [100.0] + stream_id = canonical_stream_id( + {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + ) + samples = _samples(stream_id, 1) + coordinator = WindowedArtifactCoordinator( + tmp_path, + claim_timeout_seconds=5, + clock=lambda: now[0], + ) + _register(coordinator, samples, "consumer", lookahead=0) + claim = coordinator.claim_generation("producer", stream_id=stream_id)[0] + + now[0] += 4 + assert coordinator.renew_generation_claims("producer", [claim]) == 1 + now[0] += 2 + assert coordinator.recover_expired()["expired_claims"] == 0 + assert coordinator.snapshot()["artifact_states"] == {"generating": 1} + + assert coordinator.release_generation_claims("producer", [claim]) == 1 + replacement = coordinator.claim_generation("replacement", stream_id=stream_id)[0] + assert replacement.generation == claim.generation + 1 + + +def test_position_metadata_is_bounded_across_epochs(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + coordinator = _coordinator(tmp_path) + assert coordinator.register_stream(contract) == stream_id + cursor = 0 + + for epoch in range(3): + samples = tuple( + StreamSampleIndex( + stream_id=stream_id, + sequence=cursor + ordinal, + epoch=epoch, + ordinal=ordinal, + dataset_index=ordinal, + batch_ordinal=ordinal, + batch_start_sequence=cursor + ordinal, + batch_end_sequence=cursor + ordinal + 1, + request_id=_digest(f"request-{ordinal}"), + position_id=canonical_position_id( + stream_id, + epoch=epoch, + ordinal=ordinal, + dataset_index=ordinal, + batch_ordinal=ordinal, + batch_start_sequence=cursor + ordinal, + batch_end_sequence=cursor + ordinal + 1, + ), + ) + for ordinal in range(4) + ) + coordinator.register_positions(samples) + coordinator.register_consumer( + "consumer", + stream_id=stream_id, + lookbehind=1, + lookahead=3, + max_prefetch=4, + max_inflight=4, + cursor=cursor, + reset=epoch == 0, + ) + while coordinator.snapshot()["artifact_states"].get("queued", 0): + _publish(coordinator, stream_id, tmp_path) + for sample in samples: + lease = coordinator.acquire("consumer", sample, timeout_seconds=1) + cursor = coordinator.ack("consumer", [lease.as_batch_metadata()]) + coordinator.complete_consumer("consumer") + assert coordinator.snapshot()["positions"] == 0 + + +def test_late_consumer_registration_preserves_incoming_epoch_positions(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 4) + coordinator = _coordinator(tmp_path) + _register( + coordinator, + samples, + "fast", + contract=contract, + lookahead=3, + max_prefetch=4, + ) + while coordinator.snapshot()["artifact_states"].get("queued", 0): + _publish(coordinator, stream_id, tmp_path) + for sample in samples[:2]: + lease = coordinator.acquire("fast", sample, timeout_seconds=1) + coordinator.ack("fast", [lease.as_batch_metadata()]) + + # Another process registers the same epoch before registering its cursor. + coordinator.register_positions(samples) + assert coordinator.snapshot()["positions"] == len(samples) + coordinator.register_consumer( + "late", + stream_id=stream_id, + lookbehind=0, + lookahead=3, + max_prefetch=4, + max_inflight=4, + cursor=0, + ) + lease = coordinator.acquire("late", samples[0], timeout_seconds=1) + assert lease.sequence == 0 + + def _assert_long_stream_retention_bound(tmp_path: Path, count: int) -> None: contract = { "dataset_fingerprint": f"dataset-{count}", diff --git a/tests/unit/models/test_dflash_optimized_ce.py b/tests/unit/models/test_dflash_optimized_ce.py index dccaa9253..5adbec674 100644 --- a/tests/unit/models/test_dflash_optimized_ce.py +++ b/tests/unit/models/test_dflash_optimized_ce.py @@ -10,7 +10,9 @@ from speculators.proposals.greedy import GreedyTokenProposalConfig -def _model(*, sample_from_anchor: bool = False) -> DFlashDraftModel: +def _model( + *, sample_from_anchor: bool = False, draft_vocab_size: int = 17 +) -> DFlashDraftModel: transformer_config = Qwen3Config( vocab_size=17, hidden_size=8, @@ -23,7 +25,7 @@ def _model(*, sample_from_anchor: bool = False) -> DFlashDraftModel: ) config = DFlashSpeculatorConfig( transformer_layer_config=transformer_config, - draft_vocab_size=17, + draft_vocab_size=draft_vocab_size, block_size=4, aux_hidden_state_layer_ids=[0], mask_token_id=1, @@ -103,6 +105,43 @@ def test_sample_from_anchor_input_labels_select_the_next_token(): assert torch.equal(labels, input_ids[:, indices + 1]) +def test_reduced_draft_vocab_rejects_input_id_labels(): + model = _model(draft_vocab_size=5) + input_ids = torch.tensor([[0, 16]]) + hidden = torch.randn(1, 2, model.hidden_size) + + with pytest.raises(ValueError, match="requires the full verifier vocabulary"): + model._ce_target_ids( + input_ids, + hidden, + torch.tensor([0, 1]), + label_source="input_ids", + verifier_argmax_chunk_size=0, + ) + + +def test_reduced_vocab_verifier_argmax_returns_direct_draft_boundary_ids(): + model = _model(sample_from_anchor=True, draft_vocab_size=5) + model.verifier_norm = torch.nn.Identity() + with torch.no_grad(): + model.verifier_lm_head.weight.zero_() + model.verifier_lm_head.weight[0, 0] = 1 + model.verifier_lm_head.weight[-1, 1] = 1 + hidden = torch.zeros(1, 2, model.hidden_size) + hidden[0, 0, 0] = 2 + hidden[0, 1, 1] = 2 + + labels = model._ce_target_ids( + torch.tensor([[0, 16]]), + hidden, + torch.tensor([0, 1]), + label_source="verifier_argmax", + verifier_argmax_chunk_size=1, + ) + + assert torch.equal(labels, torch.tensor([[0, model.draft_vocab_size - 1]])) + + def test_fused_ce_preparation_reuses_ignored_head_at_compute_dtype(): model = _model() model.prepare_fused_linear_cross_entropy(torch.bfloat16) diff --git a/tests/unit/ops/test_fused_linear_cross_entropy.py b/tests/unit/ops/test_fused_linear_cross_entropy.py index 90ee17fb4..e1858c051 100644 --- a/tests/unit/ops/test_fused_linear_cross_entropy.py +++ b/tests/unit/ops/test_fused_linear_cross_entropy.py @@ -86,7 +86,7 @@ def test_rejects_trainable_head_and_wrong_target_dtype(): target = torch.tensor([1, 2]) with pytest.raises(ValueError, match="frozen LM head"): fused_ce.frozen_linear_cross_entropy(hidden, weight, target) - with pytest.raises(ValueError, match="torch.long"): + with pytest.raises(ValueError, match=r"torch\.long"): fused_ce.frozen_linear_cross_entropy(hidden, weight.detach(), target.float()) @@ -102,7 +102,7 @@ def test_missing_dependency_and_wrong_version_fail_early(): with ( mock.patch.object(fused_ce, "version", return_value="0.7.0"), - pytest.raises(RuntimeError, match="requires liger-kernel==0.8.0"), + pytest.raises(RuntimeError, match=r"requires liger-kernel==0\.8\.0"), ): fused_ce._load_liger_forward() fused_ce._load_liger_forward.cache_clear() diff --git a/tests/unit/train/test_shared_artifacts.py b/tests/unit/train/test_shared_artifacts.py index 421fabe7a..b48a4b6da 100644 --- a/tests/unit/train/test_shared_artifacts.py +++ b/tests/unit/train/test_shared_artifacts.py @@ -363,6 +363,66 @@ def materialize(_index, dataset_item, _client_item): assert peak == 4 +def test_windowed_producer_stop_releases_blocked_claim(tmp_path, monkeypatch): + data_path = tmp_path / "data" + shared_path = tmp_path / "shared" + _write_multi_dataset(data_path, count=1) + dataset = ArrowDataset( + max_len=128, + datapath=data_path, + transfer=FileTransfer(tmp_path / "index"), + model="model", + on_missing="generate", + on_generate="delete", + shared_artifacts_path=shared_path, + shared_artifacts_namespace="layers:2,18,33", + shared_artifacts_consumer_id="consumer", + shared_artifacts_lookahead=0, + shared_artifacts_max_prefetch_per_consumer=1, + shared_artifacts_capture_batch_size=1, + shared_artifacts_capture_batch_wait_seconds=0, + shared_artifacts_claim_timeout_seconds=0.3, + request_timeout=5, + ) + dataset.client = cast("Any", object()) + entered = threading.Event() + release = threading.Event() + finished = threading.Event() + + def materialize(_index, dataset_item, _client_item): + entered.set() + try: + release.wait(timeout=5) + return { + "token_ids": dataset_item["input_ids"], + "hidden_states": torch.arange(12, dtype=torch.float32).reshape(3, 4), + } + finally: + finished.set() + + monkeypatch.setattr(dataset, "_materialize_shared_hs", materialize) + base_sampler = _SequentialSampler(1) + stream_id = dataset.configure_windowed_stream(base_sampler) + sampler = WindowedBatchSampler( + base_sampler, + stream_id=stream_id, + request_id_for_index=dataset.windowed_request_id, + ) + dataset.prepare_windowed_epoch(sampler.full_epoch_samples(0), cursor=0, reset=True) + + dataset.start_windowed_producer() + try: + assert entered.wait(timeout=2) + started = time.monotonic() + dataset.stop_windowed_producer() + assert time.monotonic() - started < 2 + with WindowedArtifactCoordinator(shared_path) as coordinator: + assert coordinator.snapshot()["artifact_states"] == {"queued": 1} + finally: + release.set() + assert finished.wait(timeout=2) + + def test_windowed_capture_batch_isolates_one_failed_claim(tmp_path, monkeypatch): data_path = tmp_path / "data" shared_path = tmp_path / "shared" @@ -577,6 +637,7 @@ def fake_arrow_dataset(**kwargs): return object() monkeypatch.setattr(dataloader_module, "ArrowDataset", fake_arrow_dataset) + monkeypatch.setattr(dataloader_module, "get_dp_rank", lambda: 2) monkeypatch.setattr( dataloader_module, "_setup_dataloader", @@ -627,8 +688,8 @@ def fake_arrow_dataset(**kwargs): assert kwargs["shared_artifacts_capture_batch_size"] == 6 assert kwargs["shared_artifacts_capture_batch_wait_seconds"] == 0.01 assert kwargs["shared_artifacts_max_inflight"] == 40 - assert dataset_kwargs[0]["shared_artifacts_consumer_id"] == "consumer-a:train" - assert dataset_kwargs[1]["shared_artifacts_consumer_id"] == "consumer-a:val" + assert dataset_kwargs[0]["shared_artifacts_consumer_id"] == "consumer-a:dp2:train" + assert dataset_kwargs[1]["shared_artifacts_consumer_id"] == "consumer-a:dp2:val" def test_train_only_loader_uses_full_dataset_without_validation(monkeypatch): @@ -674,4 +735,4 @@ def fake_arrow_dataset(**kwargs): assert val_loader is None assert len(dataset_kwargs) == 1 assert dataset_kwargs[0]["split_ratio"] == 1.0 - assert dataset_kwargs[0]["shared_artifacts_consumer_id"] == "consumer-a:train" + assert dataset_kwargs[0]["shared_artifacts_consumer_id"] == "consumer-a:dp0:train" From 2d7156d9cf9f2e5a051329e8c39e99f309bace24 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 22:49:26 +0800 Subject: [PATCH 16/20] fix: harden asynchronous artifact lifecycle Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- THIRD_PARTY_NOTICES | 31 ++ docs/cli/train.md | 6 +- pyproject.toml | 2 +- scripts/train.py | 28 ++ .../data_generation/artifact_cache.py | 53 +++- .../data_generation/windowed_artifacts.py | 137 +++++++++ src/speculators/train/data.py | 176 +++++++++-- src/speculators/train/dataloader.py | 14 + src/speculators/train/trainer.py | 14 +- .../data_generation/test_artifact_cache.py | 84 ++++++ .../test_windowed_artifacts.py | 87 ++++++ tests/unit/train/test_cli_args.py | 10 + tests/unit/train/test_optimizers.py | 17 +- tests/unit/train/test_shared_artifacts.py | 274 +++++++++++++++++- tests/unit/train/test_windowed_training.py | 27 ++ 15 files changed, 917 insertions(+), 43 deletions(-) create mode 100644 THIRD_PARTY_NOTICES diff --git a/THIRD_PARTY_NOTICES b/THIRD_PARTY_NOTICES new file mode 100644 index 000000000..52ebbc82d --- /dev/null +++ b/THIRD_PARTY_NOTICES @@ -0,0 +1,31 @@ +Third-Party Notices +=================== + +Bounded asynchronous producer-consumer fan-out +----------------------------------------------- + +Portions of the bounded asynchronous hidden-state fan-out implementation are +adapted from work published in the sgl-project/SpecForge pull request #707: + +https://github.com/sgl-project/SpecForge/pull/707 + +SpecForge is Copyright (c) 2025 sgl-project and is distributed under the MIT +License: + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/docs/cli/train.md b/docs/cli/train.md index 17edb8064..420c6372d 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -80,7 +80,7 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--vllm-endpoint`** (str, default: `"http://localhost:8000/v1"`) vLLM endpoint address for generating hidden states on-demand (online training). Ignored if `--on-missing` is not set to `generate`. -- **`--request-timeout`** (float, default: `180.0`) Timeout in seconds for each individual vLLM request. +- **`--request-timeout`** (float, default: `120.0`) Timeout in seconds for each individual vLLM request. - **`--max-retries`** (int, default: `3`) Maximum number of retry attempts per vLLM request on failure. @@ -110,6 +110,10 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--shared-hidden-states-claim-timeout`** (float, default: `300.0`) Timeout before an interrupted producer claim can be reassigned. +- **`--shared-hidden-states-acquire-timeout`** (float, default: derived) Maximum total seconds a consumer waits for one windowed artifact. By default this covers every configured request attempt, all exponential retry backoff, and a five-second scheduling margin. Set an explicit value to override that derived deadline. + +- **`--shared-hidden-states-lease-timeout`** (float, default: `3600.0`) Timeout before an unacknowledged artifact read lease is released. Increase this only when one legitimate training step can exceed the default. + - **`--shared-hidden-states-generation-attempts`** (int, default: `3`) Maximum coordinated generation attempts, including expired producer claims. The shared cache is a filesystem data plane, not Mooncake or GPU-direct transport. Its directory must provide reliable POSIX `flock`, same-filesystem atomic rename, and directory `fsync` semantics to every trainer. Do not assume an arbitrary NFS mount is safe unless those guarantees have been verified. diff --git a/pyproject.toml b/pyproject.toml index 431452833..af8d2a482 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ description = "A unified library for creating, representing, and storing specula readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.10" license = "Apache-2.0" -license-files = ["LICENSE"] +license-files = ["LICENSE", "THIRD_PARTY_NOTICES"] authors = [ { name = "Red Hat"} ] keywords = [ "speculative decoding", diff --git a/scripts/train.py b/scripts/train.py index b162e80c3..cd974e8f0 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -679,6 +679,12 @@ def main(args: argparse.Namespace): # noqa: C901 shared_artifacts_claim_timeout_seconds=( args.shared_hidden_states_claim_timeout ), + shared_artifacts_acquire_timeout_seconds=( + args.shared_hidden_states_acquire_timeout + ), + shared_artifacts_lease_timeout_seconds=( + args.shared_hidden_states_lease_timeout + ), shared_artifacts_generation_attempts=( args.shared_hidden_states_generation_attempts ), @@ -868,6 +874,13 @@ def _validate_windowed_shared_hidden_state_args( parser.error("--shared-hidden-states-consumer-timeout must be positive") if args.shared_hidden_states_claim_timeout <= 0: parser.error("--shared-hidden-states-claim-timeout must be positive") + if ( + args.shared_hidden_states_acquire_timeout is not None + and args.shared_hidden_states_acquire_timeout <= 0 + ): + parser.error("--shared-hidden-states-acquire-timeout must be positive") + if args.shared_hidden_states_lease_timeout <= 0: + parser.error("--shared-hidden-states-lease-timeout must be positive") if args.shared_hidden_states_generation_attempts < 1: parser.error("--shared-hidden-states-generation-attempts must be at least one") @@ -1121,6 +1134,21 @@ def parse_args(): # noqa: C901 default=300.0, help="Seconds before an interrupted producer claim can be reassigned.", ) + parser.add_argument( + "--shared-hidden-states-acquire-timeout", + type=float, + default=None, + help=( + "Maximum seconds a consumer waits for one artifact. By default this " + "covers all request attempts, retry backoff, and a small margin." + ), + ) + parser.add_argument( + "--shared-hidden-states-lease-timeout", + type=float, + default=3600.0, + help="Seconds before an unacknowledged read lease is released.", + ) parser.add_argument( "--shared-hidden-states-generation-attempts", type=int, diff --git a/src/speculators/data_generation/artifact_cache.py b/src/speculators/data_generation/artifact_cache.py index b3fadce5f..c58a63260 100644 --- a/src/speculators/data_generation/artifact_cache.py +++ b/src/speculators/data_generation/artifact_cache.py @@ -3,6 +3,7 @@ import fcntl import hashlib import json +import logging import os import re import time @@ -19,6 +20,8 @@ from collections.abc import Callable, Iterator, Mapping _REQUEST_ID_PATTERN = re.compile(r"^[0-9a-f]{64}$") +_LOCK_STRIPE_HEX_LENGTH = 3 +logger = logging.getLogger(__name__) _STATS_VERSION = 1 _COUNTERS = ( "logical_requests", @@ -176,7 +179,9 @@ def artifact_path(self, request_id: str) -> Path: return self._artifacts / request_id[:2] / f"{request_id}.safetensors" def _lock_path(self, request_id: str) -> Path: - return self._locks / request_id[:2] / f"{request_id}.lock" + # A fixed set of stripes keeps lock metadata bounded for long streams. + stripe = request_id[:_LOCK_STRIPE_HEX_LENGTH] + return self._locks / stripe[:2] / f"{stripe}.lock" @contextmanager def _request_lock( @@ -281,8 +286,20 @@ def load( path = self.artifact_path(request_id) if not path.exists(): raise ArtifactCacheError(f"Artifact {request_id} is not published") - data = load_file(path) - validate(data) + try: + data = load_file(path) + validate(data) + except Exception: + removed = False + try: + path.unlink(missing_ok=True) + removed = True + except OSError: + logger.exception("Failed to remove invalid artifact %s", request_id) + if removed: + self._fsync_after_commit(path.parent, operation="artifact removal") + self._record(invalid_artifacts_removed=1) + raise return data def remove(self, request_id: str, *, expected_path: Path | None = None) -> bool: @@ -297,9 +314,22 @@ def remove(self, request_id: str, *, expected_path: Path | None = None) -> bool: if not path.exists(): return False path.unlink() - _fsync_directory(path.parent) + self._fsync_after_commit(path.parent, operation="artifact removal") return True + @staticmethod + def _fsync_after_commit(directory: Path, *, operation: str) -> None: + try: + _fsync_directory(directory) + except OSError: + logger.warning( + "Directory fsync failed after committed %s in %s; the current " + "namespace is consistent but crash durability is not guaranteed", + operation, + directory, + exc_info=True, + ) + def _is_expired(self, path: Path, now: float) -> bool: return bool( self.artifact_ttl_seconds is not None @@ -337,7 +367,7 @@ def _publish( with temporary.open("rb") as published: os.fsync(published.fileno()) temporary.replace(target) - _fsync_directory(target.parent) + self._fsync_after_commit(target.parent, operation="artifact publication") finally: temporary.unlink(missing_ok=True) @@ -357,9 +387,11 @@ def get_or_create( artifact_path = self.artifact_path(request_id) expired = 0 invalid = 0 + removed_publication = False if artifact_path.exists() and self._is_expired(artifact_path, now): artifact_path.unlink(missing_ok=True) expired = 1 + removed_publication = True if artifact_path.exists(): try: @@ -368,6 +400,7 @@ def get_or_create( except Exception: artifact_path.unlink(missing_ok=True) invalid = 1 + removed_publication = True else: self._record( logical_requests=1, @@ -384,6 +417,11 @@ def get_or_create( coalesced=waited, ) + if removed_publication: + self._fsync_after_commit( + artifact_path.parent, operation="artifact removal" + ) + try: data = create() self._validate_tensors(data) @@ -438,6 +476,7 @@ def cleanup_stale(self, *, now: float | None = None) -> dict[str, int]: current_time = time.time() if now is None else now expired = 0 stale_temps = 0 + changed_artifact_dirs: set[Path] = set() for path in self._artifacts.glob("*/*.safetensors"): request_id = path.stem @@ -447,10 +486,14 @@ def cleanup_stale(self, *, now: float | None = None) -> dict[str, int]: with self._request_lock(request_id, timeout_seconds=0): if path.exists() and self._is_expired(path, current_time): path.unlink(missing_ok=True) + changed_artifact_dirs.add(path.parent) expired += 1 except ArtifactLockTimeoutError: continue + for directory in changed_artifact_dirs: + self._fsync_after_commit(directory, operation="expired artifact cleanup") + for path in self._artifacts.glob("*/.*.tmp"): name = path.name request_id = name[1:65] diff --git a/src/speculators/data_generation/windowed_artifacts.py b/src/speculators/data_generation/windowed_artifacts.py index 9e8275735..44312148d 100644 --- a/src/speculators/data_generation/windowed_artifacts.py +++ b/src/speculators/data_generation/windowed_artifacts.py @@ -3,6 +3,9 @@ The coordinator is a single-host control plane. Tensor payloads remain in an artifact store; SQLite contains only deterministic stream positions, consumer progress, generation claims, and read leases. + +Portions are adapted from the SpecForge asynchronous fan-out work referenced +in the repository's third-party notices. """ from __future__ import annotations @@ -168,6 +171,7 @@ def __init__( poll_seconds: float = 0.02, consumer_timeout_seconds: float = 120.0, claim_timeout_seconds: float = 300.0, + lease_timeout_seconds: float = 3600.0, max_generation_attempts: int = 3, clock: Callable[[], float] = time.time, ) -> None: @@ -177,6 +181,8 @@ def __init__( raise ValueError("consumer_timeout_seconds must be positive") if claim_timeout_seconds <= 0: raise ValueError("claim_timeout_seconds must be positive") + if lease_timeout_seconds <= 0: + raise ValueError("lease_timeout_seconds must be positive") if max_generation_attempts < 1: raise ValueError("max_generation_attempts must be at least one") self.root = Path(root).expanduser().resolve() @@ -185,6 +191,7 @@ def __init__( self.poll_seconds = poll_seconds self.consumer_timeout_seconds = consumer_timeout_seconds self.claim_timeout_seconds = claim_timeout_seconds + self.lease_timeout_seconds = lease_timeout_seconds self.max_generation_attempts = max_generation_attempts self._clock = clock self._lock = threading.RLock() @@ -198,6 +205,10 @@ def __init__( if not self._schema_is_current(): self._conn.execute("PRAGMA journal_mode=WAL") self._create_schema() + self._conn.execute( + "CREATE INDEX IF NOT EXISTS acquisitions_position " + "ON acquisitions(stream_id,sequence)" + ) def _schema_is_current(self) -> bool: table = self._conn.execute( @@ -659,6 +670,38 @@ def _retry_priority_locked( ArtifactPriority.DEMAND if demand is not None else ArtifactPriority.PREFETCH ) + def _transition_missing_artifact_locked( + self, + conn: sqlite3.Connection, + request_id: str, + *, + error: str, + now: float, + ) -> None: + interested = conn.execute( + "SELECT DISTINCT consumer_id FROM interests WHERE request_id=?", + (request_id,), + ).fetchall() + priority = self._retry_priority_locked(conn, request_id) if interested else None + conn.execute( + "UPDATE artifacts SET state=?,generation=generation+1,path=NULL," + "size_bytes=0,priority=?,queued_at=?,claim_owner=NULL,claim_until=NULL," + "failures=0,last_error=?,first_reader_accounted=0,updated_at=? " + "WHERE request_id=?", + ( + ArtifactState.QUEUED.value + if interested + else ArtifactState.ABSENT.value, + int(priority) if priority is not None else None, + now if interested else None, + error[:2000], + now, + request_id, + ), + ) + for row in interested: + self._top_up_prefetch_locked(conn, row["consumer_id"]) + def _refresh_window_locked( self, conn: sqlite3.Connection, consumer_id: str ) -> None: @@ -832,6 +875,8 @@ def heartbeat(self, consumer_id: str) -> None: def recover_expired(self) -> dict[str, int]: expired_consumers = 0 expired_claims = 0 + expired_leases = 0 + expired_evictions = 0 with self._transaction() as conn: now = self._clock() consumers = conn.execute( @@ -857,6 +902,22 @@ def recover_expired(self) -> dict[str, int]: (consumer_id,), ) expired_consumers += 1 + affected_consumers: set[str] = set() + leases = conn.execute( + "SELECT * FROM acquisitions WHERE state='leased' AND updated_at dict[str, int]: ), ) expired_claims += 1 + evictions = conn.execute( + "SELECT request_id,path FROM artifacts WHERE state=? AND updated_at dict[str, int]: return { "expired_consumers": expired_consumers, "expired_claims": expired_claims, + "expired_leases": expired_leases, + "expired_evictions": expired_evictions, } def acquire( @@ -1108,6 +1190,61 @@ def _sleep_or_timeout( if time.monotonic() >= started: time.sleep(sleep_seconds) + def invalidate_artifact( + self, lease: ArtifactReadLease, error: BaseException | str + ) -> bool: + """Invalidate a corrupt READY publication and wake all readers to retry.""" + + message = str(error)[:2000] or type(error).__name__ + with self._transaction() as conn: + acquisition = conn.execute( + "SELECT * FROM acquisitions WHERE token=?", (lease.token,) + ).fetchone() + expected_acquisition = ( + lease.consumer_id, + lease.stream_id, + lease.sequence, + lease.request_id, + "leased", + ) + observed_acquisition = ( + ( + acquisition["consumer_id"], + acquisition["stream_id"], + int(acquisition["sequence"]), + acquisition["request_id"], + acquisition["state"], + ) + if acquisition is not None + else None + ) + if observed_acquisition != expected_acquisition: + return False + + artifact = conn.execute( + "SELECT * FROM artifacts WHERE request_id=?", (lease.request_id,) + ).fetchone() + if artifact is None or ( + artifact["state"], + int(artifact["generation"]), + Path(artifact["path"]) if artifact["path"] else None, + ) != (ArtifactState.READY.value, lease.generation, lease.path): + return False + + now = self._clock() + conn.execute( + "UPDATE acquisitions SET state='waiting',updated_at=? " + "WHERE request_id=?", + (now, lease.request_id), + ) + self._transition_missing_artifact_locked( + conn, + lease.request_id, + error=f"invalid artifact: {message}", + now=now, + ) + return True + def ack(self, consumer_id: str, leases: Sequence[Mapping[str, Any]]) -> int: """Commit successful trainer consumption and advance a contiguous cursor.""" if not leases: diff --git a/src/speculators/train/data.py b/src/speculators/train/data.py index 098d10bbc..74d535178 100644 --- a/src/speculators/train/data.py +++ b/src/speculators/train/data.py @@ -1,13 +1,15 @@ import hashlib import json +import logging import math import os import random import threading +import time import uuid import warnings from collections.abc import Callable -from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait +from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait from os import PathLike from pathlib import Path from typing import Any, Literal, cast @@ -29,6 +31,7 @@ from speculators.data_generation.vllm_client import ( DEFAULT_MAX_RETRIES, DEFAULT_REQUEST_TIMEOUT, + RETRY_BACKOFF_BASE, ClientItem, generate_hidden_states, ) @@ -37,6 +40,7 @@ GenerationClaim, StreamSampleIndex, WindowedArtifactCoordinator, + WindowedArtifactError, canonical_stream_id, ) from speculators.train.noise_transforms import TransformTensors @@ -44,6 +48,7 @@ BatchType = dict[str, Any] WINDOWED_LEASE_KEY = "_windowed_artifact_lease" WINDOWED_BATCH_LEASES_KEY = "_windowed_artifact_leases" +logger = logging.getLogger(__name__) def _validate_integer_config(name: str, value: object, *, minimum: int) -> None: @@ -57,6 +62,25 @@ def _validate_non_negative_number(name: str, value: object) -> None: raise ValueError(f"{name} must be non-negative") +def _resolve_artifact_acquire_timeout( + configured_timeout: float | None, + request_timeout: float | None, + max_retries: int, +) -> float | None: + if configured_timeout is not None: + if configured_timeout <= 0: + raise ValueError( + "shared_artifacts_acquire_timeout_seconds must be positive or None" + ) + return float(configured_timeout) + if request_timeout is None: + return None + retry_backoff = sum( + RETRY_BACKOFF_BASE**attempt for attempt in range(1, max_retries + 1) + ) + return float(request_timeout) * (max_retries + 1) + retry_backoff + 5.0 + + def list_files(path): datapath = [] for root, _directories, files in os.walk(path): @@ -282,6 +306,8 @@ def __init__( shared_artifacts_max_inflight: int = 32, shared_artifacts_consumer_timeout_seconds: float = 120.0, shared_artifacts_claim_timeout_seconds: float = 300.0, + shared_artifacts_acquire_timeout_seconds: float | None = None, + shared_artifacts_lease_timeout_seconds: float = 3600.0, shared_artifacts_generation_attempts: int = 3, ): self.data = load_from_disk(datapath) @@ -307,6 +333,7 @@ def __init__( self.model = model self.request_timeout = request_timeout self.max_retries = max_retries + _validate_integer_config("max_retries", max_retries, minimum=0) self.shared_artifacts_namespace = shared_artifacts_namespace if shared_artifacts_path is not None and not shared_artifacts_namespace: raise ValueError( @@ -370,6 +397,18 @@ def __init__( self.shared_artifacts_claim_timeout_seconds = ( shared_artifacts_claim_timeout_seconds ) + self.shared_artifacts_acquire_timeout_seconds = ( + _resolve_artifact_acquire_timeout( + shared_artifacts_acquire_timeout_seconds, + request_timeout, + max_retries, + ) + ) + if shared_artifacts_lease_timeout_seconds <= 0: + raise ValueError("shared_artifacts_lease_timeout_seconds must be positive") + self.shared_artifacts_lease_timeout_seconds = float( + shared_artifacts_lease_timeout_seconds + ) self.shared_artifacts_generation_attempts = shared_artifacts_generation_attempts self.windowed_artifacts_enabled = shared_artifacts_consumer_id is not None if self.windowed_artifacts_enabled and self.artifact_cache is None: @@ -418,6 +457,7 @@ def _new_windowed_coordinator(self) -> WindowedArtifactCoordinator: self.shared_artifacts_path, consumer_timeout_seconds=self.shared_artifacts_consumer_timeout_seconds, claim_timeout_seconds=self.shared_artifacts_claim_timeout_seconds, + lease_timeout_seconds=self.shared_artifacts_lease_timeout_seconds, max_generation_attempts=self.shared_artifacts_generation_attempts, ) @@ -509,27 +549,58 @@ def _acquire_windowed_hs( if self.shared_artifacts_consumer_id is None or self.artifact_cache is None: raise RuntimeError("windowed artifacts are not configured") coordinator = self._coordinator_for_process() - lease = coordinator.acquire( - self.shared_artifacts_consumer_id, - sample, - timeout_seconds=self.request_timeout, - ) dataset_item = self.data[sample.dataset_index] - try: - loaded = self.artifact_cache.load( - sample.request_id, - lambda data: check_hidden_states( - data, dataset_item["input_ids"].tolist() - ), + timeout = self.shared_artifacts_acquire_timeout_seconds + deadline = time.monotonic() + timeout if timeout is not None else None + while True: + remaining = ( + max(0.0, deadline - time.monotonic()) if deadline is not None else None ) - if lease.cache_hit: - self.artifact_cache.record_reuse() - return loaded, lease - except BaseException: - coordinator.abandon( - self.shared_artifacts_consumer_id, [lease.as_batch_metadata()] + if remaining is not None and remaining <= 0: + raise TimeoutError( + f"stream position {sample.sequence} was not ready within " + f"{timeout:.1f}s" + ) + lease = coordinator.acquire( + self.shared_artifacts_consumer_id, + sample, + timeout_seconds=remaining, ) - raise + try: + loaded = self.artifact_cache.load( + sample.request_id, + lambda data: check_hidden_states( + data, dataset_item["input_ids"].tolist() + ), + ) + except Exception as error: # noqa: BLE001 - invalid publication boundary + try: + coordinator.invalidate_artifact(lease, error) + finally: + coordinator.abandon( + self.shared_artifacts_consumer_id, + [lease.as_batch_metadata()], + ) + logger.warning( + "Invalidated shared artifact %s after load failure; retrying", + sample.request_id, + exc_info=True, + ) + continue + except BaseException: + coordinator.abandon( + self.shared_artifacts_consumer_id, [lease.as_batch_metadata()] + ) + raise + try: + if lease.cache_hit: + self.artifact_cache.record_reuse() + return loaded, lease + except BaseException: + coordinator.abandon( + self.shared_artifacts_consumer_id, [lease.as_batch_metadata()] + ) + raise def ack_windowed_batch(self, leases: list[dict[str, Any]]) -> int | None: if not self.windowed_artifacts_enabled or not leases: @@ -586,6 +657,7 @@ def _run_windowed_producer(self) -> None: artifact_ttl_seconds=None, lock_timeout_seconds=self.shared_artifacts_lock_timeout_seconds, ) + next_cache_maintenance = 0.0 try: if self.client is None: self._setup_client() @@ -596,6 +668,16 @@ def _run_windowed_producer(self) -> None: try: with self._new_windowed_coordinator() as coordinator: while not self._windowed_producer_stop.is_set(): + now = time.monotonic() + if now >= next_cache_maintenance: + try: + cache.cleanup_stale() + except (ArtifactCacheError, OSError): + logger.warning( + "Windowed artifact cache maintenance failed", + exc_info=True, + ) + next_cache_maintenance = now + cache.stale_temp_seconds coordinator.heartbeat(self.shared_artifacts_consumer_id) coordinator.recover_expired() self._evict_windowed_artifacts(coordinator, cache) @@ -648,24 +730,50 @@ def _run_windowed_claim_batch( for future in done: claim = pending.pop(future) try: - path, size_bytes = future.result() - except Exception as error: # noqa: BLE001 - coordinator.fail_generation(owner, claim, error) - else: - coordinator.complete_generation( - owner, - claim, - path=path, - size_bytes=size_bytes, + self._finish_windowed_claim(coordinator, owner, claim, future) + except WindowedArtifactError: + logger.warning( + "Discarding stale generation result for %s", + claim.request_id, + exc_info=True, ) if pending: - coordinator.renew_generation_claims(owner, tuple(pending.values())) + for future, claim in tuple(pending.items()): + try: + coordinator.renew_generation_claims(owner, (claim,)) + except WindowedArtifactError: + pending.pop(future) + future.cancel() + logger.warning( + "Stopped tracking stale generation claim for %s", + claim.request_id, + exc_info=True, + ) finally: if pending: for future in pending: future.cancel() coordinator.release_generation_claims(owner, tuple(pending.values())) + @staticmethod + def _finish_windowed_claim( + coordinator: WindowedArtifactCoordinator, + owner: str, + claim: GenerationClaim, + future: Future[tuple[Path, int]], + ) -> None: + try: + path, size_bytes = future.result() + except Exception as error: # noqa: BLE001 - isolate one capture + coordinator.fail_generation(owner, claim, error) + return + coordinator.complete_generation( + owner, + claim, + path=path, + size_bytes=size_bytes, + ) + def _produce_windowed_claim( self, cache: HiddenStateArtifactCache, @@ -696,11 +804,17 @@ def _evict_windowed_artifacts( try: removed = cache.remove(eviction.request_id, expected_path=eviction.path) except (ArtifactCacheError, OSError): - coordinator.finish_eviction(eviction, removed=False) - else: + removed = False + try: coordinator.finish_eviction( eviction, removed=removed or not eviction.path.exists() ) + except WindowedArtifactError: + logger.warning( + "Discarding stale eviction result for %s", + eviction.request_id, + exc_info=True, + ) def _materialize_shared_hs( self, diff --git a/src/speculators/train/dataloader.py b/src/speculators/train/dataloader.py index 59afcd126..de1359f6b 100644 --- a/src/speculators/train/dataloader.py +++ b/src/speculators/train/dataloader.py @@ -211,6 +211,8 @@ def create_train_val_loaders( shared_artifacts_max_inflight: int = 32, shared_artifacts_consumer_timeout_seconds: float = 120.0, shared_artifacts_claim_timeout_seconds: float = 300.0, + shared_artifacts_acquire_timeout_seconds: float | None = None, + shared_artifacts_lease_timeout_seconds: float = 3600.0, shared_artifacts_generation_attempts: int = 3, train_data_ratio: float = 0.9, ) -> tuple[DataLoader, DataLoader | None]: @@ -291,6 +293,12 @@ def create_train_val_loaders( shared_artifacts_claim_timeout_seconds=( shared_artifacts_claim_timeout_seconds ), + shared_artifacts_acquire_timeout_seconds=( + shared_artifacts_acquire_timeout_seconds + ), + shared_artifacts_lease_timeout_seconds=( + shared_artifacts_lease_timeout_seconds + ), shared_artifacts_generation_attempts=(shared_artifacts_generation_attempts), ) if train_data_ratio < 1.0: @@ -335,6 +343,12 @@ def create_train_val_loaders( shared_artifacts_claim_timeout_seconds=( shared_artifacts_claim_timeout_seconds ), + shared_artifacts_acquire_timeout_seconds=( + shared_artifacts_acquire_timeout_seconds + ), + shared_artifacts_lease_timeout_seconds=( + shared_artifacts_lease_timeout_seconds + ), shared_artifacts_generation_attempts=( shared_artifacts_generation_attempts ), diff --git a/src/speculators/train/trainer.py b/src/speculators/train/trainer.py index cf682d8c3..2e648d1c2 100644 --- a/src/speculators/train/trainer.py +++ b/src/speculators/train/trainer.py @@ -258,7 +258,15 @@ def _run_windowed_phase( finally: dataset = loader.dataset if hasattr(dataset, "stop_windowed_producer"): - dataset.stop_windowed_producer(completed=completed) + try: + dataset.stop_windowed_producer(completed=completed) + except Exception: + if completed: + raise + root_logger.exception( + "Windowed producer cleanup failed while the training phase " + "was already unwinding" + ) @staticmethod def _ack_windowed_batch(dataset, leases: list[dict]) -> None: @@ -502,8 +510,8 @@ def _clip_gradients(self): grad_norm = torch.nn.utils.get_total_norm(gradients, foreach=True) fused_optimizer = cast("_FusedOptimizer", optimizer) fused_optimizer.grad_scale = ( - (grad_norm + 1e-6) / self.config.max_grad_norm - ).clamp(min=1.0) + ((grad_norm + 1e-6) / self.config.max_grad_norm).clamp(min=1.0).float() + ) return grad_norm def _optimizers_step(self): diff --git a/tests/unit/data_generation/test_artifact_cache.py b/tests/unit/data_generation/test_artifact_cache.py index e5f6f75a2..da97c34da 100644 --- a/tests/unit/data_generation/test_artifact_cache.py +++ b/tests/unit/data_generation/test_artifact_cache.py @@ -143,6 +143,16 @@ def create(): } +def test_request_locks_use_a_bounded_stripe_set(tmp_path): + cache = HiddenStateArtifactCache(tmp_path, artifact_ttl_seconds=None) + request_ids = [f"{index % 4096:03x}{index:061x}" for index in range(5000)] + + lock_paths = {cache._lock_path(request_id) for request_id in request_ids} + + assert len(lock_paths) == 4096 + assert all(path.parent.parent == tmp_path / "locks" for path in lock_paths) + + def test_independent_processes_coalesce_one_generation(tmp_path): context = multiprocessing.get_context("spawn") ready = context.Queue() @@ -303,6 +313,80 @@ def test_corrupt_artifact_is_removed_before_regeneration(tmp_path): assert cache.snapshot_stats()["invalid_artifacts_removed"] == 1 +def test_load_removes_invalid_publication(tmp_path): + cache = HiddenStateArtifactCache(tmp_path, artifact_ttl_seconds=None) + target = cache.artifact_path(_request_id()) + target.parent.mkdir(parents=True) + save_file(_tensors((4, 5, 6)), target) + + with pytest.raises(ValueError, match="token"): + cache.load(_request_id(), _validate) + + assert not target.exists() + assert cache.snapshot_stats()["invalid_artifacts_removed"] == 1 + + +def test_invalid_load_fsyncs_removal_even_when_accounting_fails(tmp_path, monkeypatch): + cache = HiddenStateArtifactCache(tmp_path, artifact_ttl_seconds=None) + target = cache.artifact_path(_request_id()) + target.parent.mkdir(parents=True) + save_file(_tensors((4, 5, 6)), target) + fsyncs = [] + monkeypatch.setattr( + cache, + "_fsync_after_commit", + lambda directory, **_kwargs: fsyncs.append(directory), + ) + monkeypatch.setattr( + cache, "_record", lambda **_kwargs: (_ for _ in ()).throw(OSError("stats")) + ) + + with pytest.raises(OSError, match="stats"): + cache.load(_request_id(), _validate) + + assert not target.exists() + assert fsyncs == [target.parent] + + +def test_post_publish_directory_fsync_failure_keeps_visible_success( + tmp_path, monkeypatch, caplog +): + cache = HiddenStateArtifactCache(tmp_path, artifact_ttl_seconds=None) + original_fsync = artifact_cache_module._fsync_directory + + def fail_artifact_fsync(path): + if Path(path) != tmp_path: + raise OSError("directory fsync failed") + original_fsync(path) + + monkeypatch.setattr(artifact_cache_module, "_fsync_directory", fail_artifact_fsync) + result = cache.get_or_create(_request_id(), _tensors, _validate) + + assert result.path.is_file() + assert cache.snapshot_stats()["publishes"] == 1 + assert cache.snapshot_stats()["publish_failures"] == 0 + assert "crash durability is not guaranteed" in caplog.text + + +def test_post_remove_directory_fsync_failure_keeps_removed_result( + tmp_path, monkeypatch, caplog +): + cache = HiddenStateArtifactCache(tmp_path, artifact_ttl_seconds=None) + result = cache.get_or_create(_request_id(), _tensors, _validate) + original_fsync = artifact_cache_module._fsync_directory + + def fail_artifact_fsync(path): + if Path(path) != tmp_path: + raise OSError("directory fsync failed") + original_fsync(path) + + monkeypatch.setattr(artifact_cache_module, "_fsync_directory", fail_artifact_fsync) + + assert cache.remove(_request_id(), expected_path=result.path) + assert not result.path.exists() + assert "crash durability is not guaranteed" in caplog.text + + def test_cleanup_removes_expired_artifact_and_stale_temp(tmp_path): cache = HiddenStateArtifactCache( tmp_path, artifact_ttl_seconds=1, stale_temp_seconds=1 diff --git a/tests/unit/data_generation/test_windowed_artifacts.py b/tests/unit/data_generation/test_windowed_artifacts.py index b35da893f..f614dc33a 100644 --- a/tests/unit/data_generation/test_windowed_artifacts.py +++ b/tests/unit/data_generation/test_windowed_artifacts.py @@ -560,11 +560,98 @@ def test_expired_consumer_releases_window_and_read_lease(tmp_path): assert coordinator.recover_expired() == { "expired_consumers": 1, "expired_claims": 0, + "expired_leases": 0, + "expired_evictions": 0, } assert coordinator.snapshot()["inflight_acquisitions"] == 0 assert len(coordinator.begin_evictions()) == 1 +def test_invalid_ready_artifact_requeues_all_readers(tmp_path): + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 1) + coordinator = _coordinator(tmp_path) + _register(coordinator, samples, "consumer-a", "consumer-b", lookahead=0) + _publish(coordinator, stream_id, tmp_path) + first = coordinator.acquire("consumer-a", samples[0], timeout_seconds=1) + second = coordinator.acquire("consumer-b", samples[0], timeout_seconds=1) + + assert coordinator.invalidate_artifact(first, "corrupt payload") + assert not coordinator.invalidate_artifact(second, "same stale payload") + assert coordinator.snapshot()["artifact_states"] == {"queued": 1} + assert coordinator.snapshot()["inflight_acquisitions"] == 2 + with pytest.raises(WindowedArtifactError, match="unknown, stale, or mismatched"): + coordinator.ack("consumer-a", [first.as_batch_metadata()]) + + coordinator.abandon("consumer-a", [first.as_batch_metadata()]) + coordinator.abandon("consumer-b", [second.as_batch_metadata()]) + claim = coordinator.claim_generation("producer", stream_id=stream_id)[0] + assert claim.generation == first.generation + 1 + _complete_claim(coordinator, "producer", claim, tmp_path) + replay = coordinator.acquire("consumer-a", samples[0], timeout_seconds=1) + assert coordinator.ack("consumer-a", [replay.as_batch_metadata()]) == 1 + + +def test_expired_read_lease_is_reaped_while_consumer_remains_active(tmp_path): + now = [100.0] + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 1) + coordinator = WindowedArtifactCoordinator( + tmp_path, + poll_seconds=0.005, + consumer_timeout_seconds=30, + claim_timeout_seconds=10, + lease_timeout_seconds=5, + clock=lambda: now[0], + ) + _register(coordinator, samples, "consumer", lookahead=0) + _publish(coordinator, stream_id, tmp_path) + stale = coordinator.acquire("consumer", samples[0], timeout_seconds=1) + + now[0] += 6 + coordinator.heartbeat("consumer") + assert coordinator.recover_expired() == { + "expired_consumers": 0, + "expired_claims": 0, + "expired_leases": 1, + "expired_evictions": 0, + } + with pytest.raises(WindowedArtifactError, match="unknown, stale, or mismatched"): + coordinator.ack("consumer", [stale.as_batch_metadata()]) + replay = coordinator.acquire("consumer", samples[0], timeout_seconds=1) + assert coordinator.ack("consumer", [replay.as_batch_metadata()]) == 1 + + +def test_expired_eviction_reconciles_visible_and_removed_paths(tmp_path): + now = [100.0] + contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} + stream_id = canonical_stream_id(contract) + samples = _samples(stream_id, 1) + coordinator = WindowedArtifactCoordinator( + tmp_path, + poll_seconds=0.005, + consumer_timeout_seconds=30, + claim_timeout_seconds=5, + clock=lambda: now[0], + ) + _register(coordinator, samples, "consumer", lookahead=0) + _publish(coordinator, stream_id, tmp_path) + coordinator.complete_consumer("consumer") + + coordinator.begin_evictions()[0] + now[0] += 6 + assert coordinator.recover_expired()["expired_evictions"] == 1 + assert coordinator.snapshot()["artifact_states"] == {"ready": 1} + + second = coordinator.begin_evictions()[0] + second.path.unlink() + now[0] += 6 + assert coordinator.recover_expired()["expired_evictions"] == 1 + assert coordinator.snapshot()["artifact_states"] == {"absent": 1} + + def test_resume_reset_rewinds_cursor_and_clears_uncommitted_leases(tmp_path): contract = {"dataset_fingerprint": "dataset-a", "sampler_seed": 0} stream_id = canonical_stream_id(contract) diff --git a/tests/unit/train/test_cli_args.py b/tests/unit/train/test_cli_args.py index d0f9bc070..4526b53e9 100644 --- a/tests/unit/train/test_cli_args.py +++ b/tests/unit/train/test_cli_args.py @@ -31,6 +31,8 @@ def test_shared_hidden_state_cache_is_opt_in(monkeypatch): assert args.shared_hidden_states_capture_batch_size == 8 assert args.shared_hidden_states_capture_batch_wait == 0.002 assert args.shared_hidden_states_max_inflight == 32 + assert args.shared_hidden_states_acquire_timeout is None + assert args.shared_hidden_states_lease_timeout == 3600.0 def test_shared_hidden_state_cache_arguments(monkeypatch): @@ -78,6 +80,10 @@ def test_windowed_shared_hidden_state_arguments(monkeypatch): "60", "--shared-hidden-states-claim-timeout", "90", + "--shared-hidden-states-acquire-timeout", + "180", + "--shared-hidden-states-lease-timeout", + "600", "--shared-hidden-states-generation-attempts", "4", ], @@ -92,6 +98,8 @@ def test_windowed_shared_hidden_state_arguments(monkeypatch): assert args.shared_hidden_states_max_inflight == 40 assert args.shared_hidden_states_consumer_timeout == 60 assert args.shared_hidden_states_claim_timeout == 90 + assert args.shared_hidden_states_acquire_timeout == 180 + assert args.shared_hidden_states_lease_timeout == 600 assert args.shared_hidden_states_generation_attempts == 4 @@ -138,6 +146,8 @@ def test_windowed_shared_hidden_state_arguments(monkeypatch): ["--shared-hidden-states-max-inflight", "0"], ["--shared-hidden-states-consumer-timeout", "0"], ["--shared-hidden-states-claim-timeout", "0"], + ["--shared-hidden-states-acquire-timeout", "0"], + ["--shared-hidden-states-lease-timeout", "0"], ["--shared-hidden-states-generation-attempts", "0"], ], ) diff --git a/tests/unit/train/test_optimizers.py b/tests/unit/train/test_optimizers.py index ddff9ebb9..9a0f573ae 100644 --- a/tests/unit/train/test_optimizers.py +++ b/tests/unit/train/test_optimizers.py @@ -130,12 +130,25 @@ def test_fused_clip_sets_scale_and_cleans_it_after_step(): norm = trainer._clip_gradients() expected_scale = ((norm + 1e-6) / trainer.config.max_grad_norm).clamp(min=1) torch.testing.assert_close(optimizer.grad_scale, expected_scale) + assert optimizer.grad_scale.dtype == torch.float32 trainer._optimizers_step() assert optimizer.saw_grad_scale assert not hasattr(optimizer, "grad_scale") +def test_fused_clip_uses_fp32_scale_for_bfloat16_gradients(): + optimizer = _RecordingFusedOptimizer() + trainer = _fused_clip_trainer(optimizer) + trainer.model.to(dtype=torch.bfloat16) + for parameter in trainer.model.parameters(): + parameter.grad = torch.full_like(parameter, 2.0) + + trainer._clip_gradients() + + assert optimizer.grad_scale.dtype == torch.float32 + + def test_fused_clip_cleans_scale_when_optimizer_step_fails(): optimizer = _RecordingFusedOptimizer(fail=True) trainer = _fused_clip_trainer(optimizer) @@ -162,7 +175,9 @@ def test_fused_adamw_clip_matches_explicit_clipping_on_cuda(): reference_optimizer.step() grad_norm = torch.nn.utils.get_total_norm([fused.grad], foreach=True) fused_optimizer_with_scale = cast("Any", fused_optimizer) - fused_optimizer_with_scale.grad_scale = ((grad_norm + 1e-6) / 0.7).clamp(min=1) + fused_optimizer_with_scale.grad_scale = ( + ((grad_norm + 1e-6) / 0.7).clamp(min=1).float() + ) try: fused_optimizer.step() finally: diff --git a/tests/unit/train/test_shared_artifacts.py b/tests/unit/train/test_shared_artifacts.py index b48a4b6da..61e8e8776 100644 --- a/tests/unit/train/test_shared_artifacts.py +++ b/tests/unit/train/test_shared_artifacts.py @@ -2,6 +2,7 @@ import threading import time +from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any, Literal, cast import hs_connectors.transfer as transfer_module @@ -13,7 +14,12 @@ import speculators.train.data as data_module import speculators.train.dataloader as dataloader_module from hs_connectors import FileTransfer -from speculators.data_generation.windowed_artifacts import WindowedArtifactCoordinator +from speculators.data_generation.windowed_artifacts import ( + ArtifactPriority, + GenerationClaim, + WindowedArtifactCoordinator, + WindowedArtifactError, +) from speculators.train.data import WINDOWED_LEASE_KEY, ArrowDataset from speculators.train.dataloader import WindowedBatchSampler @@ -134,6 +140,35 @@ def test_windowed_dataset_requires_online_generation(tmp_path): ) +def test_windowed_dataset_derives_and_overrides_artifact_timeouts(tmp_path): + data_path = tmp_path / "data" + _write_dataset(data_path) + + defaults = ArrowDataset(max_len=128, datapath=data_path, model="model") + derived = ArrowDataset( + max_len=128, + datapath=data_path, + model="model", + request_timeout=10, + max_retries=2, + ) + explicit = ArrowDataset( + max_len=128, + datapath=data_path, + model="model", + request_timeout=10, + max_retries=2, + shared_artifacts_acquire_timeout_seconds=73, + shared_artifacts_lease_timeout_seconds=91, + ) + + assert defaults.shared_artifacts_acquire_timeout_seconds == 499 + assert derived.shared_artifacts_acquire_timeout_seconds == 41 + assert derived.shared_artifacts_lease_timeout_seconds == 3600 + assert explicit.shared_artifacts_acquire_timeout_seconds == 73 + assert explicit.shared_artifacts_lease_timeout_seconds == 91 + + def _successful_generator(service_path: Path, calls: list[Path]): def generate(*_args, **_kwargs): path = service_path / f"request-{len(calls)}.safetensors" @@ -238,6 +273,127 @@ def test_windowed_dataset_dispatches_reads_acks_and_cleans_final_window( assert stats["publishes"] == 1 +def test_windowed_dataset_invalidates_and_regenerates_corrupt_ready_artifact( + tmp_path, monkeypatch +): + data_path = tmp_path / "data" + shared_path = tmp_path / "shared" + _write_dataset(data_path) + dataset = ArrowDataset( + max_len=128, + datapath=data_path, + transfer=FileTransfer(tmp_path / "index"), + model="model", + on_missing="generate", + on_generate="delete", + shared_artifacts_path=shared_path, + shared_artifacts_namespace="layers:2,18,33", + shared_artifacts_consumer_id="consumer", + request_timeout=2, + ) + dataset.client = cast("Any", object()) + monkeypatch.setattr( + dataset, + "_materialize_shared_hs", + lambda _index, _dataset_item, _client_item: _hidden_states(), + ) + base_sampler = _SingleBatchSampler() + stream_id = dataset.configure_windowed_stream(base_sampler) + sampler = WindowedBatchSampler( + base_sampler, + stream_id=stream_id, + request_id_for_index=dataset.windowed_request_id, + ) + samples = sampler.full_epoch_samples(0) + dataset.prepare_windowed_epoch(samples, cursor=0, reset=True) + assert dataset.artifact_cache is not None + corrupt_path = dataset.artifact_cache.artifact_path(samples[0].request_id) + corrupt_path.parent.mkdir(parents=True) + save_file( + { + "token_ids": torch.tensor([4, 5, 6]), + "hidden_states": torch.zeros(3, 4), + }, + corrupt_path, + ) + with WindowedArtifactCoordinator(shared_path) as coordinator: + claim = coordinator.claim_generation("initial", stream_id=stream_id)[0] + coordinator.complete_generation( + "initial", + claim, + path=corrupt_path, + size_bytes=corrupt_path.stat().st_size, + ) + + dataset.start_windowed_producer() + try: + item = dataset[samples[0]] + assert item is not None + assert item["input_ids"].tolist() == [1, 2, 3] + lease = item.pop(WINDOWED_LEASE_KEY) + assert dataset.ack_windowed_batch([lease]) == 1 + finally: + dataset.stop_windowed_producer(completed=True) + + stats = dataset.artifact_cache.snapshot_stats() + assert stats["invalid_artifacts_removed"] == 1 + assert stats["publishes"] == 1 + + +def test_windowed_producer_cleans_abandoned_cache_temporary(tmp_path, monkeypatch): + data_path = tmp_path / "data" + shared_path = tmp_path / "shared" + _write_dataset(data_path) + dataset = ArrowDataset( + max_len=128, + datapath=data_path, + transfer=FileTransfer(tmp_path / "index"), + model="model", + on_missing="generate", + on_generate="delete", + shared_artifacts_path=shared_path, + shared_artifacts_namespace="layers:2,18,33", + shared_artifacts_consumer_id="consumer", + request_timeout=2, + ) + dataset.client = cast("Any", object()) + monkeypatch.setattr( + dataset, + "_materialize_shared_hs", + lambda _index, _dataset_item, _client_item: _hidden_states(), + ) + base_sampler = _SingleBatchSampler() + stream_id = dataset.configure_windowed_stream(base_sampler) + sampler = WindowedBatchSampler( + base_sampler, + stream_id=stream_id, + request_id_for_index=dataset.windowed_request_id, + ) + dataset.prepare_windowed_epoch(sampler.full_epoch_samples(0), cursor=0, reset=True) + assert dataset.artifact_cache is not None + temporary = ( + dataset.artifact_cache.artifact_path("a" * 64).parent + / f".{('a' * 64)}.dead.tmp" + ) + temporary.parent.mkdir(parents=True, exist_ok=True) + temporary.write_bytes(b"partial") + old = time.time() - dataset.artifact_cache.stale_temp_seconds - 1 + temporary.touch() + data_module.os.utime(temporary, (old, old)) + + dataset.start_windowed_producer() + try: + deadline = time.monotonic() + 2 + while temporary.exists(): + if time.monotonic() >= deadline: + raise TimeoutError("stale cache temporary was not cleaned") + time.sleep(0.01) + finally: + dataset.stop_windowed_producer() + + assert dataset.artifact_cache.snapshot_stats()["stale_temps_removed"] == 1 + + @pytest.mark.parametrize("num_workers", [0, 1, 4]) def test_windowed_scheduling_is_independent_of_dataloader_workers( tmp_path, monkeypatch, num_workers @@ -481,6 +637,118 @@ def materialize(index, dataset_item, _client_item): dataset.stop_windowed_producer() +@pytest.mark.parametrize("stale_operation", ["complete", "fail"]) +def test_windowed_capture_batch_isolates_stale_terminal_result( + tmp_path, monkeypatch, stale_operation +): + dataset = cast("Any", ArrowDataset.__new__(ArrowDataset)) + dataset._windowed_producer_stop = threading.Event() + dataset.shared_artifacts_consumer_id = "consumer" + dataset.shared_artifacts_claim_timeout_seconds = 0.03 + claims = ( + GenerationClaim( + request_id="a" * 64, + stream_id="1" * 64, + dataset_index=0, + generation=1, + priority=ArtifactPriority.DEMAND, + ), + GenerationClaim( + request_id="b" * 64, + stream_id="1" * 64, + dataset_index=1, + generation=1, + priority=ArtifactPriority.DEMAND, + ), + ) + completed = [] + + def produce(_cache, claim): + if stale_operation == "fail" and claim is claims[0]: + raise RuntimeError("capture failed") + return tmp_path / f"{claim.request_id}.safetensors", 1 + + class _Coordinator: + @staticmethod + def heartbeat(_consumer_id): + return None + + @staticmethod + def complete_generation(_owner, claim, **_kwargs): + if stale_operation == "complete" and claim is claims[0]: + raise WindowedArtifactError("stale completion") + completed.append(claim) + + @staticmethod + def fail_generation(_owner, claim, _error): + if stale_operation == "fail" and claim is claims[0]: + raise WindowedArtifactError("stale failure") + + @staticmethod + def renew_generation_claims(_owner, _claims): + return None + + @staticmethod + def release_generation_claims(_owner, _claims): + return None + + monkeypatch.setattr(dataset, "_produce_windowed_claim", produce) + with ThreadPoolExecutor(max_workers=2) as executor: + dataset._run_windowed_claim_batch( + _Coordinator(), cast("Any", object()), executor, "owner", claims + ) + + assert claims[1] in completed + + +def test_windowed_capture_batch_drops_only_stale_renewal(monkeypatch): + dataset = cast("Any", ArrowDataset.__new__(ArrowDataset)) + dataset._windowed_producer_stop = threading.Event() + dataset.shared_artifacts_consumer_id = "consumer" + dataset.shared_artifacts_claim_timeout_seconds = 0.03 + release = threading.Event() + claims = ( + GenerationClaim("a" * 64, "1" * 64, 0, 1, ArtifactPriority.DEMAND), + GenerationClaim("b" * 64, "1" * 64, 1, 1, ArtifactPriority.DEMAND), + ) + completed = [] + + def produce(_cache, claim): + assert release.wait(2) + return cast("Any", object()), claim.dataset_index + + class _Coordinator: + @staticmethod + def heartbeat(_consumer_id): + return None + + @staticmethod + def complete_generation(_owner, claim, **_kwargs): + completed.append(claim) + + @staticmethod + def fail_generation(_owner, _claim, _error): + raise AssertionError("generation should not fail") + + @staticmethod + def renew_generation_claims(_owner, renewed): + if renewed[0] is claims[0]: + raise WindowedArtifactError("stale renewal") + release.set() + + @staticmethod + def release_generation_claims(_owner, _claims): + return None + + monkeypatch.setattr(dataset, "_produce_windowed_claim", produce) + with ThreadPoolExecutor(max_workers=2) as executor: + dataset._run_windowed_claim_batch( + _Coordinator(), cast("Any", object()), executor, "owner", claims + ) + + assert completed == [claims[1]] + + def test_unconfigured_dataset_keeps_existing_per_request_delete_behavior( tmp_path, monkeypatch ): @@ -674,6 +942,8 @@ def fake_arrow_dataset(**kwargs): shared_artifacts_capture_batch_size=6, shared_artifacts_capture_batch_wait_seconds=0.01, shared_artifacts_max_inflight=40, + shared_artifacts_acquire_timeout_seconds=75, + shared_artifacts_lease_timeout_seconds=600, ) assert len(dataset_kwargs) == 2 @@ -688,6 +958,8 @@ def fake_arrow_dataset(**kwargs): assert kwargs["shared_artifacts_capture_batch_size"] == 6 assert kwargs["shared_artifacts_capture_batch_wait_seconds"] == 0.01 assert kwargs["shared_artifacts_max_inflight"] == 40 + assert kwargs["shared_artifacts_acquire_timeout_seconds"] == 75 + assert kwargs["shared_artifacts_lease_timeout_seconds"] == 600 assert dataset_kwargs[0]["shared_artifacts_consumer_id"] == "consumer-a:dp2:train" assert dataset_kwargs[1]["shared_artifacts_consumer_id"] == "consumer-a:dp2:val" diff --git a/tests/unit/train/test_windowed_training.py b/tests/unit/train/test_windowed_training.py index 7a4e21dc7..6085b32a8 100644 --- a/tests/unit/train/test_windowed_training.py +++ b/tests/unit/train/test_windowed_training.py @@ -105,6 +105,7 @@ def test_collate_keeps_artifact_leases_out_of_model_tensors(): @dataclass class _RecordingDataset: events: list[str] + stop_error: Exception | None = None def ack_windowed_batch(self, _leases: list[dict[str, Any]]) -> None: self.events.append("ack") @@ -114,6 +115,8 @@ def abandon_windowed_batch(self, _leases: list[dict[str, Any]]) -> None: def stop_windowed_producer(self, *, completed: bool = False) -> None: self.events.append(f"stop:{completed}") + if self.stop_error is not None: + raise self.stop_error class _Loader: @@ -229,3 +232,27 @@ def operation(_epoch: int) -> None: with pytest.raises(RuntimeError, match="phase failed"): Trainer._run_windowed_phase(cast("Any", loader), operation, 0) assert events == ["run", "stop:False"] + + +def test_windowed_cleanup_error_does_not_replace_phase_failure(caplog): + events: list[str] = [] + dataset = _RecordingDataset(events, stop_error=RuntimeError("producer failed")) + loader = _Loader(dataset, {}) + + def operation(_epoch: int) -> None: + raise ValueError("training failed") + + with pytest.raises(ValueError, match="training failed"): + Trainer._run_windowed_phase(cast("Any", loader), operation, 0) + assert events == ["stop:False"] + assert "cleanup failed" in caplog.text + + +def test_windowed_cleanup_error_propagates_after_success(): + events: list[str] = [] + dataset = _RecordingDataset(events, stop_error=RuntimeError("producer failed")) + loader = _Loader(dataset, {}) + + with pytest.raises(RuntimeError, match="producer failed"): + Trainer._run_windowed_phase(cast("Any", loader), lambda _epoch: None, 0) + assert events == ["stop:True"] From 47c280534dd997eff75564142ee3e1743e271c91 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 22:49:52 +0800 Subject: [PATCH 17/20] fix: preserve fanout benchmark failure evidence Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- .../independent_consumer_fanout/README.md | 2 +- src/speculators/benchmarks/gpu_monitor.py | 25 ++- .../benchmarks/independent_consumers.py | 86 ++++++-- tests/unit/benchmarks/test_gpu_monitor.py | 31 +++ .../benchmarks/test_independent_consumers.py | 183 +++++++++++++++++- 5 files changed, 296 insertions(+), 31 deletions(-) diff --git a/benchmarks/independent_consumer_fanout/README.md b/benchmarks/independent_consumer_fanout/README.md index 36952109b..f5c3ec950 100644 --- a/benchmarks/independent_consumer_fanout/README.md +++ b/benchmarks/independent_consumer_fanout/README.md @@ -26,7 +26,7 @@ The run directory must not already exist. Role logs remain there and the compact The example commands set `--train-data-ratio 1.0`, so every input belongs to the training stream and no validation pass changes window positions or producer request accounting. Keep this setting fixed when comparing against a train-only baseline. -`producer_common_steady` uses that same consumer overlap to report producer requests, first publications, recaptures, request throughput, and effective unique-sample throughput. The older `steady_state` field remains the service-wide interval after its own completion warmup; do not compare that full-run interval with a producer metric measured only over the common consumer overlap. +`producer_common_steady` uses that same consumer overlap to report producer requests, first publications, recaptures, request throughput, and effective unique-sample throughput. Request multiplicity is also computed only from this steady interval. A complete cross-consumer sample cohort split by the start or end boundary is excluded rather than misclassified, and `boundary_sample_keys_excluded` records that count; a cohort already complete inside the interval is not changed by later requests outside it. The older `steady_state` field remains the service-wide interval after its own completion warmup; do not compare that full-run interval with a producer metric measured only over the common consumer overlap. For the unshared baseline, `expected_service_completions_per_shared_sample` is one for `1p1c` and three for `1p3c`. A publish-once implementation changes the latter to one; the logical consumer commands and all other workload settings must remain equivalent. diff --git a/src/speculators/benchmarks/gpu_monitor.py b/src/speculators/benchmarks/gpu_monitor.py index 2f1ece24d..d2c0edeba 100644 --- a/src/speculators/benchmarks/gpu_monitor.py +++ b/src/speculators/benchmarks/gpu_monitor.py @@ -356,13 +356,24 @@ def stop(self) -> dict[str, Any]: self._errors.append("GPU monitor thread did not stop") self._ended_at_ns = time.monotonic_ns() if self._output is not None: - self._write( - { - "record_type": "session_end", - "timestamp_monotonic_ns": self._ended_at_ns, - } - ) - self._close_output() + try: + self._write( + { + "record_type": "session_end", + "timestamp_monotonic_ns": self._ended_at_ns, + } + ) + except Exception as error: # noqa: BLE001 - preserve shutdown + self._errors.append( + f"session_end write failed: {_error_text(error)}" + ) + finally: + try: + self._close_output() + except Exception as error: # noqa: BLE001 - preserve summary + self._errors.append( + f"GPU sample output close failed: {_error_text(error)}" + ) summary = { "status": ( "ok" if not self._errors and not self._violations else "degraded" diff --git a/src/speculators/benchmarks/independent_consumers.py b/src/speculators/benchmarks/independent_consumers.py index 50b0d8a3b..a874f5eb3 100644 --- a/src/speculators/benchmarks/independent_consumers.py +++ b/src/speculators/benchmarks/independent_consumers.py @@ -62,6 +62,7 @@ class EvidenceError(RuntimeError): } _DISTRIBUTED_LAUNCHER_MODULES = { "accelerate.commands.launch", + "deepspeed", "torch.distributed.launch", "torch.distributed.run", } @@ -442,6 +443,7 @@ def _forward(self) -> None: ] response_body = b'{"error":"accounting proxy upstream failure"}' valid_completion = False + connection: http.client.HTTPConnection | None = None try: connection_type = ( http.client.HTTPSConnection @@ -474,8 +476,6 @@ def _forward(self) -> None: response_reason = response.reason response_body = response.read() response_headers = response.getheaders() - connection.close() - valid_completion = bool( is_completion and request_key @@ -484,6 +484,12 @@ def _forward(self) -> None: ) except Exception as error: # noqa: BLE001 request_error = request_error or type(error).__name__ + finally: + if connection is not None: + try: + connection.close() + except Exception as error: # noqa: BLE001 - preserve response + request_error = request_error or type(error).__name__ if is_completion: proxy._ledger.add( @@ -627,13 +633,34 @@ def analyze_scenario( # noqa: C901 ) key_counts = Counter( - event.request_key for event in measured if event.request_key is not None + event.request_key for event in steady_events if event.request_key is not None ) key_consumers: dict[str, set[str]] = defaultdict(set) - for event in measured: + for event in steady_events: if event.request_key is not None: key_consumers[event.request_key].add(event.consumer_id) + boundary_keys: set[str] = set() + if multiplicity == len(expected_ids): + measured_key_counts = Counter( + event.request_key for event in measured if event.request_key is not None + ) + measured_key_consumers: dict[str, set[str]] = defaultdict(set) + for event in measured: + if event.request_key is not None: + measured_key_consumers[event.request_key].add(event.consumer_id) + expected_consumers = set(expected_ids) + boundary_keys = { + key + for key, count in key_counts.items() + if count < multiplicity + and measured_key_counts[key] == multiplicity + and measured_key_consumers[key] == expected_consumers + } + for key in boundary_keys: + del key_counts[key] + key_consumers.pop(key, None) + if multiplicity == len(expected_ids): qualifying = [ key @@ -691,6 +718,7 @@ def analyze_scenario( # noqa: C901 "bounded_window_regeneration" if bounded_regeneration else "exact" ), "qualifying_shared_samples": len(qualifying), + "boundary_sample_keys_excluded": len(boundary_keys), "observed_multiplicity_histogram": { str(count): samples for count, samples in sorted(multiplicity_histogram.items()) @@ -1128,7 +1156,12 @@ def terminate(self, grace_seconds: float = 20.0) -> None: except subprocess.TimeoutExpired: with suppress(ProcessLookupError): os.killpg(self.process.pid, signal.SIGKILL) - self.process.wait(timeout=10) + try: + self.process.wait(timeout=10) + except subprocess.TimeoutExpired: + self._reader_error = self._reader_error or ( + "subprocess did not exit within 10s after SIGKILL" + ) self.finished_at = time.monotonic() self.close_log() @@ -1322,15 +1355,30 @@ def capture_step( f"GPU monitor shutdown failed: {type(error).__name__}: {error}" ) for process in consumers: - process.terminate() + try: + process.terminate() + except Exception as error: # noqa: BLE001 - continue role cleanup + runtime_errors.append( + f"consumer cleanup failed: {type(error).__name__}: {error}" + ) if process.reader_error is not None: runtime_errors.append( f"consumer log reader failed: {process.reader_error}" ) for proxy in proxies: - proxy.close() + try: + proxy.close() + except Exception as error: # noqa: BLE001 - continue role cleanup + runtime_errors.append( + f"proxy cleanup failed: {type(error).__name__}: {error}" + ) if producer is not None: - producer.terminate() + try: + producer.terminate() + except Exception as error: # noqa: BLE001 - report incomplete cleanup + runtime_errors.append( + f"producer cleanup failed: {type(error).__name__}: {error}" + ) if producer.reader_error is not None: runtime_errors.append( f"producer log reader failed: {producer.reader_error}" @@ -1462,23 +1510,27 @@ def capture_step( *monitor_summary.get("errors", []), *monitor_summary.get("ownership_violations", []), ] - return { - "kind": scenario.kind, - "valid": not invalid_reasons, - "invalid_reasons": invalid_reasons, - "consumer_processes": [ + consumer_processes = [] + for index, consumer in enumerate(scenario.consumers): + process = consumers[index] if index < len(consumers) else None + consumer_processes.append( { "consumer_id": consumer.consumer_id, "gpu": consumer.gpu, - "return_code": process.process.returncode, + "started": process is not None, + "return_code": process.process.returncode if process else None, "runtime_seconds": ( process.finished_at - process.started_at - if process.finished_at is not None + if process is not None and process.finished_at is not None else None ), } - for consumer, process in zip(scenario.consumers, consumers, strict=False) - ], + ) + return { + "kind": scenario.kind, + "valid": not invalid_reasons, + "invalid_reasons": invalid_reasons, + "consumer_processes": consumer_processes, "request_accounting": analysis["request_accounting"], "shared_artifact_cache": cache_accounting, "windowed_artifacts": windowed_snapshot, diff --git a/tests/unit/benchmarks/test_gpu_monitor.py b/tests/unit/benchmarks/test_gpu_monitor.py index 5c418f1ca..d6376a7c9 100644 --- a/tests/unit/benchmarks/test_gpu_monitor.py +++ b/tests/unit/benchmarks/test_gpu_monitor.py @@ -137,6 +137,37 @@ def test_stop_timeout_does_not_race_a_blocked_sample(tmp_path, monkeypatch): assert backend.closed +def test_stop_write_failure_still_closes_and_writes_degraded_summary( + tmp_path, monkeypatch +): + backend = _FakeBackend() + summary_path = tmp_path / "summary.json" + monitor = GpuMonitor( + [GpuRoleAssignment(0, "producer", "producer")], + tmp_path / "samples.jsonl", + summary_path, + poll_seconds=0.01, + backend=backend, + ) + monitor.start() + assert backend.sampled.wait(1) + original_write = monitor._write + + def fail_session_end(value): + if value.get("record_type") == "session_end": + raise OSError("disk full") + original_write(value) + + monkeypatch.setattr(monitor, "_write", fail_session_end) + summary = monitor.stop() + + assert summary["status"] == "degraded" + assert any("session_end write failed" in error for error in summary["errors"]) + assert monitor._output is None + assert backend.closed + assert json.loads(summary_path.read_text())["status"] == "degraded" + + def test_window_summary_excludes_startup_and_reports_active_time(): assignment = GpuRoleAssignment(2, "consumer:b4", "consumer:b4") samples = [ diff --git a/tests/unit/benchmarks/test_independent_consumers.py b/tests/unit/benchmarks/test_independent_consumers.py index 55d6eecc9..040b8f2d2 100644 --- a/tests/unit/benchmarks/test_independent_consumers.py +++ b/tests/unit/benchmarks/test_independent_consumers.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import socket import threading import urllib.request from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -177,6 +178,7 @@ def test_benchmark_cli_output_paths_are_conditional(monkeypatch, tmp_path): ["python", "-m", "torch.distributed.run", "trainer.py"], ["python", "-m", "torch.distributed.launch", "trainer.py"], ["python", "-u", "-m", "accelerate.commands.launch", "trainer.py"], + ["python", "-m", "deepspeed", "trainer.py"], ["python", "trainer.py", "--tensor-parallel-size=3"], ["python", "trainer.py", "-tp", "3"], ["python", "trainer.py", "--data-parallel-size", "3"], @@ -219,8 +221,8 @@ def test_windowed_1p3c_requires_publish_once_multiplicity(): assert scenario.expected_service_completions_per_shared_sample == 1 -@pytest.mark.parametrize("times_out", [False, True]) -def test_managed_process_termination_tolerates_exit_races(monkeypatch, times_out): +@pytest.mark.parametrize("timeout_count", [0, 1, 2]) +def test_managed_process_termination_tolerates_exit_races(monkeypatch, timeout_count): class _Process: pid = 123 @@ -232,13 +234,14 @@ def poll(self): def wait(self, timeout): self.wait_calls += 1 - if times_out and self.wait_calls == 1: + if self.wait_calls <= timeout_count: raise benchmark_module.subprocess.TimeoutExpired("consumer", timeout) return 0 managed = object.__new__(benchmark_module._ManagedProcess) managed.process = _Process() managed.finished_at = None + managed._reader_error = None closed = [] managed.close_log = lambda: closed.append(True) @@ -248,9 +251,13 @@ def process_exited(*_args): monkeypatch.setattr(benchmark_module.os, "killpg", process_exited) managed.terminate(grace_seconds=0.01) - assert managed.process.wait_calls == (2 if times_out else 1) + assert managed.process.wait_calls == min(timeout_count + 1, 2) assert managed.finished_at is not None assert closed == [True] + if timeout_count == 2: + assert "after SIGKILL" in managed.reader_error + else: + assert managed.reader_error is None def test_config_rejects_gpu_sharing_and_owned_environment(): @@ -356,12 +363,60 @@ def test_accounting_proxy_counts_validated_hidden_state_completion(): upstream_thread.join(timeout=2) +def test_accounting_proxy_closes_upstream_connection_after_failure( + monkeypatch, +): + connections = [] + connection_base = benchmark_module.http.client.HTTPConnection + + class _FailingConnection(connection_base): + def __init__(self, *_args, **_kwargs): + self.closed = False + connections.append(self) + + def request(self, *_args, **_kwargs): + raise OSError("upstream disconnected") + + def close(self): + self.closed = True + + monkeypatch.setattr( + benchmark_module.http.client, "HTTPConnection", _FailingConnection + ) + proxy = AccountingProxy( + "http://127.0.0.1:1/v1", "c0", AccountingLedger(), timeout_seconds=1 + ) + proxy.start() + try: + body = json.dumps( + {"model": "model", "prompt": [1, 2, 3], "max_tokens": 1} + ).encode() + host, port = proxy._server.server_address + with socket.create_connection((host, port), timeout=2) as client: + client.sendall( + b"POST /v1/completions HTTP/1.1\r\n" + b"Host: localhost\r\n" + b"Content-Type: application/json\r\n" + + f"Content-Length: {len(body)}\r\n".encode() + + b"Connection: close\r\n\r\n" + + body + ) + response = client.recv(4096) + assert b" 502 " in response.split(b"\r\n", 1)[0] + finally: + proxy.close() + + assert len(connections) == 1 + assert connections[0].closed + + def test_analysis_separates_warmup_and_requires_exact_multiplicity(): result = analyze_scenario(_scenario(), _valid_events(), 0.0, 1.0) assert result["valid"] assert result["request_accounting"]["requests"] == 12 - assert result["request_accounting"]["qualifying_shared_samples"] == 3 + assert result["request_accounting"]["qualifying_shared_samples"] == 2 + assert result["request_accounting"]["boundary_sample_keys_excluded"] == 1 assert ( result["request_accounting"]["per_consumer_completions_semantics"] == "logical_consumer" @@ -370,6 +425,31 @@ def test_analysis_separates_warmup_and_requires_exact_multiplicity(): assert result["steady_state"]["completions"] == 7 +def test_analysis_ignores_multiplicity_outside_common_steady_window(): + events = [ + _event("c0", "warmup-0", 0.10), + _event("c1", "warmup-1", 0.15), + _event("c2", "warmup-2", 0.18), + _event("c0", "sample-a", 0.20), + _event("c1", "sample-a", 0.25), + _event("c2", "sample-a", 0.30), + _event("c0", "sample-b", 0.50), + _event("c1", "sample-b", 0.50), + _event("c2", "sample-b", 0.50), + _event("c0", "sample-a", 0.90), + ] + + result = analyze_scenario( + _scenario(warmup=1, minimum_steady=1, minimum_shared=2), + events, + started_at=0.0, + finished_at=1.0, + ) + + assert result["valid"] + assert result["request_accounting"]["observed_multiplicity_histogram"] == {"3": 2} + + def test_analysis_fails_closed_on_missing_consumer_evidence(): events = [event for event in _valid_events() if event.consumer_id != "c2"] @@ -391,7 +471,7 @@ def test_analysis_fails_closed_on_failed_or_duplicate_completion(): assert not result["valid"] assert result["request_accounting"]["invalid_completions"] == 1 - assert any( + assert not any( "ambiguous multiplicity" in reason for reason in result["invalid_reasons"] ) @@ -653,3 +733,94 @@ def test_consumer_step_analysis_fails_closed_on_missing_log(tmp_path): assert not result["valid"] assert result["invalid_reasons"] == ["consumer log is missing"] + + +def test_partial_startup_reports_every_role_and_isolates_cleanup_failures( # noqa: C901 + tmp_path, monkeypatch +): + terminated = [] + closed_proxies = [] + + class _Process: + returncode = None + + @staticmethod + def poll(): + return None + + class _FakeManagedProcess: + def __init__(self, _command, _env, _gpu, log_path, _line_callback=None): + self.role = log_path.stem + if self.role == "c1": + raise RuntimeError("c1 failed to start") + self.process = _Process() + self.started_at = 1.0 + self.finished_at = None + self.reader_error = None + + def terminate(self): + terminated.append(self.role) + if self.role == "c0": + raise RuntimeError("c0 cleanup failed") + self.finished_at = 2.0 + + class _FakeProxy: + def __init__(self, _target, consumer_id, _ledger, _timeout): + self.consumer_id = consumer_id + self.endpoint = f"http://proxy/{consumer_id}" + + def start(self): + return None + + def close(self): + closed_proxies.append(self.consumer_id) + if self.consumer_id == "c0": + raise RuntimeError("c0 proxy cleanup failed") + + class _FakeMonitor: + def __init__(self, _assignments, sample_path, _summary_path, **_kwargs): + self.sample_path = sample_path + + def start(self): + return None + + def stop(self): + return {"status": "ok", "errors": [], "ownership_violations": []} + + monkeypatch.setattr(benchmark_module, "_ManagedProcess", _FakeManagedProcess) + monkeypatch.setattr(benchmark_module, "AccountingProxy", _FakeProxy) + monkeypatch.setattr(benchmark_module, "NvmlGpuMonitor", _FakeMonitor) + monkeypatch.setattr(benchmark_module, "_wait_for_producer", lambda *_args: None) + monkeypatch.setattr( + benchmark_module, + "_gpu_snapshot", + lambda gpus: benchmark_module._GpuSample( + captured_at=0.0, + total_memory_mib=dict.fromkeys(gpus, 0), + role_memory_mib=dict.fromkeys(gpus, 0), + compute_pids={gpu: [] for gpu in gpus}, + ), + ) + + config = _config() + report = benchmark_module._run_scenario( + config, config.scenarios[1], tmp_path / "run" + ) + + assert [role["started"] for role in report["consumer_processes"]] == [ + True, + False, + False, + ] + assert [role["consumer_id"] for role in report["consumer_processes"]] == [ + "c0", + "c1", + "c2", + ] + assert terminated == ["c0", "producer"] + assert closed_proxies == ["c0", "c1", "c2"] + assert any("c1 failed to start" in reason for reason in report["invalid_reasons"]) + assert any( + "consumer cleanup failed" in reason for reason in report["invalid_reasons"] + ) + assert any("proxy cleanup failed" in reason for reason in report["invalid_reasons"]) From 5dfac24a4e44ec53e5fa4497f74f31c81236d682 Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sat, 18 Jul 2026 23:27:08 +0800 Subject: [PATCH 18/20] fix: satisfy quality checks and clarify third-party notices Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- THIRD_PARTY_NOTICES | 23 ++++++++++++++----- .../benchmarks/independent_consumers.py | 13 +++++++---- .../benchmarks/test_independent_consumers.py | 12 ++++++---- tests/unit/models/test_dflash_optimized_ce.py | 4 +++- tests/unit/train/test_optimizers.py | 2 +- 5 files changed, 36 insertions(+), 18 deletions(-) diff --git a/THIRD_PARTY_NOTICES b/THIRD_PARTY_NOTICES index 52ebbc82d..802e10952 100644 --- a/THIRD_PARTY_NOTICES +++ b/THIRD_PARTY_NOTICES @@ -1,16 +1,27 @@ Third-Party Notices =================== -Bounded asynchronous producer-consumer fan-out ------------------------------------------------ +SpecForge bounded asynchronous window coordinator +-------------------------------------------------- -Portions of the bounded asynchronous hidden-state fan-out implementation are -adapted from work published in the sgl-project/SpecForge pull request #707: +Portions of the bounded asynchronous window coordinator are adapted from +``specforge/runtime/data_plane/windowed_capture.py`` and +``specforge/runtime/data_plane/windowed_capture_runtime.py`` published in the +sgl-project/SpecForge pull request #707: https://github.com/sgl-project/SpecForge/pull/707 -SpecForge is Copyright (c) 2025 sgl-project and is distributed under the MIT -License: +Copyright 2024 The SpecForge team. These portions are licensed under the Apache +License, Version 2.0. The complete Apache-2.0 terms are included in the +repository ``LICENSE`` file. + +SpecForge fused linear cross entropy +------------------------------------ + +Portions of the fused linear cross-entropy adapter are adapted from +``specforge/ops/fused_linear_cross_entropy.py`` in the same SpecForge pull +request. SpecForge is Copyright (c) 2025 sgl-project and is distributed under +the MIT License: Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/src/speculators/benchmarks/independent_consumers.py b/src/speculators/benchmarks/independent_consumers.py index a874f5eb3..876c46a8c 100644 --- a/src/speculators/benchmarks/independent_consumers.py +++ b/src/speculators/benchmarks/independent_consumers.py @@ -1512,16 +1512,19 @@ def capture_step( ] consumer_processes = [] for index, consumer in enumerate(scenario.consumers): - process = consumers[index] if index < len(consumers) else None + managed_process = consumers[index] if index < len(consumers) else None consumer_processes.append( { "consumer_id": consumer.consumer_id, "gpu": consumer.gpu, - "started": process is not None, - "return_code": process.process.returncode if process else None, + "started": managed_process is not None, + "return_code": ( + managed_process.process.returncode if managed_process else None + ), "runtime_seconds": ( - process.finished_at - process.started_at - if process is not None and process.finished_at is not None + managed_process.finished_at - managed_process.started_at + if managed_process is not None + and managed_process.finished_at is not None else None ), } diff --git a/tests/unit/benchmarks/test_independent_consumers.py b/tests/unit/benchmarks/test_independent_consumers.py index 040b8f2d2..4cba1ffa6 100644 --- a/tests/unit/benchmarks/test_independent_consumers.py +++ b/tests/unit/benchmarks/test_independent_consumers.py @@ -1,12 +1,13 @@ from __future__ import annotations +import http.client import json import socket import threading import urllib.request from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path -from typing import Literal +from typing import Any, Literal import pytest from pydantic import ValidationError @@ -238,7 +239,7 @@ def wait(self, timeout): raise benchmark_module.subprocess.TimeoutExpired("consumer", timeout) return 0 - managed = object.__new__(benchmark_module._ManagedProcess) + managed: Any = object.__new__(benchmark_module._ManagedProcess) managed.process = _Process() managed.finished_at = None managed._reader_error = None @@ -367,9 +368,8 @@ def test_accounting_proxy_closes_upstream_connection_after_failure( monkeypatch, ): connections = [] - connection_base = benchmark_module.http.client.HTTPConnection - class _FailingConnection(connection_base): + class _FailingConnection(http.client.HTTPConnection): def __init__(self, *_args, **_kwargs): self.closed = False connections.append(self) @@ -391,7 +391,9 @@ def close(self): body = json.dumps( {"model": "model", "prompt": [1, 2, 3], "max_tokens": 1} ).encode() - host, port = proxy._server.server_address + address = proxy._server.server_address + host = str(address[0]) + port = int(address[1]) with socket.create_connection((host, port), timeout=2) as client: client.sendall( b"POST /v1/completions HTTP/1.1\r\n" diff --git a/tests/unit/models/test_dflash_optimized_ce.py b/tests/unit/models/test_dflash_optimized_ce.py index 5adbec674..81581357c 100644 --- a/tests/unit/models/test_dflash_optimized_ce.py +++ b/tests/unit/models/test_dflash_optimized_ce.py @@ -1,5 +1,7 @@ """Focused tests for DFlash hard-label CE preparation and target selection.""" +from typing import Any, cast + import pytest import torch from transformers import Qwen3Config @@ -122,7 +124,7 @@ def test_reduced_draft_vocab_rejects_input_id_labels(): def test_reduced_vocab_verifier_argmax_returns_direct_draft_boundary_ids(): model = _model(sample_from_anchor=True, draft_vocab_size=5) - model.verifier_norm = torch.nn.Identity() + model.verifier_norm = cast("Any", torch.nn.Identity()) with torch.no_grad(): model.verifier_lm_head.weight.zero_() model.verifier_lm_head.weight[0, 0] = 1 diff --git a/tests/unit/train/test_optimizers.py b/tests/unit/train/test_optimizers.py index 9a0f573ae..30869e462 100644 --- a/tests/unit/train/test_optimizers.py +++ b/tests/unit/train/test_optimizers.py @@ -140,7 +140,7 @@ def test_fused_clip_sets_scale_and_cleans_it_after_step(): def test_fused_clip_uses_fp32_scale_for_bfloat16_gradients(): optimizer = _RecordingFusedOptimizer() trainer = _fused_clip_trainer(optimizer) - trainer.model.to(dtype=torch.bfloat16) + cast("Any", trainer.model).to(dtype=torch.bfloat16) for parameter in trainer.model.parameters(): parameter.grad = torch.full_like(parameter, 2.0) From 933a6005c4b6d585e27bf75bd4c087b4addb612f Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Sun, 19 Jul 2026 00:22:56 +0800 Subject: [PATCH 19/20] fix: support non-windowed training loaders Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- src/speculators/train/trainer.py | 11 +++++------ tests/unit/train/test_windowed_training.py | 4 ++++ 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/speculators/train/trainer.py b/src/speculators/train/trainer.py index 2e648d1c2..090c337a1 100644 --- a/src/speculators/train/trainer.py +++ b/src/speculators/train/trainer.py @@ -256,8 +256,8 @@ def _run_windowed_phase( completed = True return result finally: - dataset = loader.dataset - if hasattr(dataset, "stop_windowed_producer"): + dataset = getattr(loader, "dataset", None) + if dataset is not None and hasattr(dataset, "stop_windowed_producer"): try: dataset.stop_windowed_producer(completed=completed) except Exception: @@ -829,7 +829,6 @@ def run_training(self): dist.barrier() for loader in (self.train_loader, self.val_loader): - if loader is not None and hasattr(loader.dataset, "stop_windowed_producer"): - loader.dataset.stop_windowed_producer( # type: ignore[union-attr] - completed=True - ) + dataset = getattr(loader, "dataset", None) + if dataset is not None and hasattr(dataset, "stop_windowed_producer"): + dataset.stop_windowed_producer(completed=True) diff --git a/tests/unit/train/test_windowed_training.py b/tests/unit/train/test_windowed_training.py index 6085b32a8..6a6586a90 100644 --- a/tests/unit/train/test_windowed_training.py +++ b/tests/unit/train/test_windowed_training.py @@ -221,6 +221,10 @@ def operation(epoch: int) -> str: assert events == ["run:3", "stop:True"] +def test_windowed_phase_supports_loader_without_dataset(): + assert Trainer._run_windowed_phase(cast("Any", []), lambda epoch: epoch, 3) == 3 + + def test_windowed_phase_stops_without_completion_after_failure(): events: list[str] = [] loader = _Loader(_RecordingDataset(events), {}) From cd23186957d28d78e46c38d652873b40466d2b4e Mon Sep 17 00:00:00 2001 From: heiheiha798 <2300012738@stu.pku.edu.cn> Date: Mon, 20 Jul 2026 20:47:04 +0800 Subject: [PATCH 20/20] fix: harden benchmark report utilities Signed-off-by: heiheiha798 <2300012738@stu.pku.edu.cn> --- src/speculators/benchmarks/_statistics.py | 12 +++++++++++ src/speculators/benchmarks/gpu_monitor.py | 14 ++++--------- .../benchmarks/independent_consumers.py | 21 +++++++++++-------- .../benchmarks/test_independent_consumers.py | 17 +++++++++++++++ 4 files changed, 45 insertions(+), 19 deletions(-) create mode 100644 src/speculators/benchmarks/_statistics.py diff --git a/src/speculators/benchmarks/_statistics.py b/src/speculators/benchmarks/_statistics.py new file mode 100644 index 000000000..5239ba593 --- /dev/null +++ b/src/speculators/benchmarks/_statistics.py @@ -0,0 +1,12 @@ +from collections.abc import Sequence + + +def percentile(values: Sequence[float], fraction: float) -> float | None: + if not values: + return None + ordered = sorted(values) + index = min( + len(ordered) - 1, + max(0, int(len(ordered) * fraction + 0.999999) - 1), + ) + return float(ordered[index]) diff --git a/src/speculators/benchmarks/gpu_monitor.py b/src/speculators/benchmarks/gpu_monitor.py index d2c0edeba..0c03b7d15 100644 --- a/src/speculators/benchmarks/gpu_monitor.py +++ b/src/speculators/benchmarks/gpu_monitor.py @@ -11,6 +11,8 @@ from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Any, Protocol +from speculators.benchmarks._statistics import percentile + if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Mapping, Sequence from pathlib import Path @@ -426,14 +428,6 @@ def iter_gpu_samples(path: Path) -> Iterator[dict[str, Any]]: yield value -def _percentile(values: Sequence[float], fraction: float) -> float | None: - if not values: - return None - ordered = sorted(values) - index = min(len(ordered) - 1, max(0, int(len(ordered) * fraction + 0.999999) - 1)) - return float(ordered[index]) - - def summarize_gpu_window( samples: Iterable[Mapping[str, Any]], assignments: Sequence[GpuRoleAssignment], @@ -511,8 +505,8 @@ def summarize_gpu_window( ), "gpu_utilization_pct": { "mean": statistics.fmean(utilization) if utilization else None, - "p50": _percentile(utilization, 0.50), - "p95": _percentile(utilization, 0.95), + "p50": percentile(utilization, 0.50), + "p95": percentile(utilization, 0.95), "low_fraction": ( sum(value < low_utilization_pct for value in utilization) / len(utilization) diff --git a/src/speculators/benchmarks/independent_consumers.py b/src/speculators/benchmarks/independent_consumers.py index 876c46a8c..af7bb3672 100644 --- a/src/speculators/benchmarks/independent_consumers.py +++ b/src/speculators/benchmarks/independent_consumers.py @@ -22,6 +22,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from speculators.benchmarks._statistics import percentile from speculators.benchmarks.gpu_monitor import ( GpuMonitor as NvmlGpuMonitor, ) @@ -915,11 +916,6 @@ class ConsumerStepEvent: step_ms: float -def _percentile(values: list[float], percentile: float) -> float: - index = max(0, min(len(values) - 1, int(len(values) * percentile + 0.999999) - 1)) - return sorted(values)[index] - - def analyze_consumer_steps( log_path: Path, warmup_steps: int, @@ -998,8 +994,8 @@ def analyze_consumer_steps( "warmup_steps": min(warmup_steps, len(values)), "steady_steps": len(steady), "step_ms_mean": sum(steady) / len(steady) if steady else None, - "step_ms_p50": _percentile(steady, 0.50) if steady else None, - "step_ms_p95": _percentile(steady, 0.95) if steady else None, + "step_ms_p50": percentile(steady, 0.50), + "step_ms_p95": percentile(steady, 0.95), "steady_started_at_monotonic_ns": started_at, "steady_finished_at_monotonic_ns": finished_at, "steady_duration_seconds": duration, @@ -1594,8 +1590,15 @@ def load_config(path: Path) -> BenchmarkConfig: def write_report(report: dict[str, Any], path: Path) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") - temporary.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") - temporary.replace(path) + try: + with temporary.open("w", encoding="utf-8") as output: + json.dump(report, output, indent=2, sort_keys=True) + output.write("\n") + output.flush() + os.fsync(output.fileno()) + temporary.replace(path) + finally: + temporary.unlink(missing_ok=True) def event_as_dict(event: RequestEvent) -> dict[str, Any]: diff --git a/tests/unit/benchmarks/test_independent_consumers.py b/tests/unit/benchmarks/test_independent_consumers.py index 4cba1ffa6..8277ac90b 100644 --- a/tests/unit/benchmarks/test_independent_consumers.py +++ b/tests/unit/benchmarks/test_independent_consumers.py @@ -157,6 +157,23 @@ def test_example_config_uses_train_only_workloads(): assert validated.scenarios[1].expected_service_completions_per_shared_sample == 1 +def test_write_report_syncs_complete_json_before_replace(tmp_path, monkeypatch): + report_path = tmp_path / "report.json" + temporary = report_path.with_suffix(report_path.suffix + ".tmp") + synced_payloads = [] + + def record_fsync(_file_descriptor): + synced_payloads.append(temporary.read_text(encoding="utf-8")) + + monkeypatch.setattr(benchmark_module.os, "fsync", record_fsync) + + benchmark_module.write_report({"valid": True}, report_path) + + assert synced_payloads == ['{\n "valid": true\n}\n'] + assert json.loads(report_path.read_text(encoding="utf-8")) == {"valid": True} + assert not temporary.exists() + + def test_benchmark_cli_output_paths_are_conditional(monkeypatch, tmp_path): config = tmp_path / "config.json" monkeypatch.setattr(