diff --git a/benchmarks/README.md b/benchmarks/README.md index ff331e8..eb6ae5e 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -148,6 +148,37 @@ interleaves samples by dataset, reducing queue-order bias. `--max-samples` limits each selected JSONL file independently. LongBench is intended for long-text performance validation rather than fast correctness smoke tests. +### Adding a dataset + +Datasets live behind a registry in `utils/datasets.py`, so the runner and load +generator never branch on a dataset name. To add one, subclass +`BenchmarkDataset`, set `name`, implement `load()` (return `BenchmarkSample` +objects), declare your own CLI options in `add_cli_args`, build yourself in +`from_args`, and decorate the class with `@register_dataset`: + +```python +@register_dataset +class MyDataset(BenchmarkDataset): + name = "mydata" + + @classmethod + def add_cli_args(cls, parser): + parser.add_argument("--mydata-path") + + @classmethod + def from_args(cls, args): + return cls(args.mydata_path, max_samples=args.max_samples) + + def load(self): + return [BenchmarkSample(...), ...] +``` + +`--dataset mydata` then works in both `run_bench.py` and `bench_load.py` with no +edits to either: `add_dataset_cli_args` registers the options and +`build_dataset` dispatches by name. Prompt construction, sizing, load +generation, and metrics already consume only the normalized `BenchmarkSample` +shape. + ## Prompt Construction All service benchmarks use the same RAG chat prompt shape: diff --git a/benchmarks/bench_load.py b/benchmarks/bench_load.py index 1cfaeae..066d386 100644 --- a/benchmarks/bench_load.py +++ b/benchmarks/bench_load.py @@ -20,9 +20,8 @@ slot_size_for_block_tokens, ) from benchmarks.utils.datasets import ( - BenchmarkSample, - ImdbDataset, - LongBenchDataset, + add_dataset_cli_args, + build_dataset, dedup_by_context, interleave_samples, ) @@ -65,10 +64,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: ) parser.add_argument("--manifest", default=None) parser.add_argument("--prepared-config", default=None) - parser.add_argument("--dataset", choices=("imdb", "longbench"), required=True) - parser.add_argument("--imdb") - parser.add_argument("--longbench-dir") - parser.add_argument("--datasets", default=None) + add_dataset_cli_args(parser) parser.add_argument("--max-samples", type=int, default=20) parser.add_argument("--block-size", type=int, default=BLOCK_TOKENS) parser.add_argument("--max-context-tokens", type=int, default=0) @@ -104,7 +100,7 @@ async def main_async(args: argparse.Namespace) -> None: raise ValueError("--model is required with --prepare-only") if store_dir is None: raise ValueError("--store-dir is required with --prepare-only") - samples = _load_samples(args) + samples = build_dataset(args).load() from transformers import AutoConfig, AutoTokenizer @@ -260,21 +256,6 @@ async def main_async(args: argparse.Namespace) -> None: print(f"results={args.out}") -def _load_samples(args: argparse.Namespace) -> list[BenchmarkSample]: - if args.dataset == "imdb": - if not args.imdb: - raise ValueError("--imdb is required for --dataset imdb") - return ImdbDataset(args.imdb, max_samples=args.max_samples).load() - if not args.longbench_dir: - raise ValueError("--longbench-dir is required for --dataset longbench") - datasets = None - if args.datasets: - datasets = [item.strip() for item in args.datasets.split(",") if item.strip()] - return LongBenchDataset( - args.longbench_dir, datasets=datasets, max_samples=args.max_samples - ).load() - - def _should_dedup_context(args: argparse.Namespace) -> bool: """Return whether duplicate contexts should be removed from the workload. diff --git a/benchmarks/run_bench.py b/benchmarks/run_bench.py index 25777bc..2ef4249 100644 --- a/benchmarks/run_bench.py +++ b/benchmarks/run_bench.py @@ -4,10 +4,8 @@ from __future__ import annotations import argparse -import asyncio from dataclasses import dataclass import json -import math from pathlib import Path import shlex import subprocess @@ -19,22 +17,9 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1])) -from benchmarks.utils.constants import ( - COMPARISON_IOURING_MEM, - slot_size_for_block_tokens, -) -from benchmarks.utils.loadgen import ( - _wait_lmcache_quiescent, - backend_server_hit_rate, - collect_phase_metrics, -) +from benchmarks.utils import vllm_bench +from benchmarks.utils.datasets import add_dataset_cli_args from benchmarks.utils.servers import BenchmarkManifest, stop_from_pid_file -from benchmarks.utils.sizing import ( - BenchmarkCapacityLimits, - derive_benchmark_sizing, - derive_capacity_limits, - format_capacity, -) _DASER_METRICS_SETTLE_SECONDS = 2.0 _BACKEND_ROWS = ("baseline", "lmcache", "daser-prefix", "daser-chunk") @@ -170,12 +155,9 @@ def parse_args(argv: list[str] | None = None) -> RunBenchArgs: choices=("internal", "vllm-bench"), default="internal", ) - parser.add_argument("--dataset", choices=("imdb", "longbench"), default="longbench") + add_dataset_cli_args(parser, default="longbench") parser.add_argument("--model", required=True) parser.add_argument("--store-dir", required=True) - parser.add_argument("--imdb") - parser.add_argument("--longbench-dir") - parser.add_argument("--datasets", default=None) parser.add_argument("--max-samples", type=int, default=20) parser.add_argument("--gpu-id", default="auto") parser.add_argument("--gpu-util", type=float, default=0.85) @@ -269,14 +251,14 @@ def run_benchmark(args: RunBenchArgs) -> Path: _print_kv("dataset", "vllm-bench-random") _print_kv("bench_num_prompts", args.bench_num_prompts) _print_kv("bench_input_len", args.bench_input_len) - _print_kv("bench_output_len", _bench_output_len(args)) + _print_kv("bench_output_len", vllm_bench.bench_output_len(args)) else: _print_kv("dataset", args.dataset) _print_kv("max_samples", args.max_samples) _print_kv("block_size", args.block_size) _print_kv("output", prepare_path) if args.load_generator == "vllm-bench": - prepare = {"config": _bench_prepare_config(args, run_root)} + prepare = {"config": vllm_bench.prepare_config(args, run_root)} prepare_path.write_text(json.dumps(prepare, indent=2), encoding="utf-8") print(json.dumps(prepare["config"], indent=2), flush=True) print(f"prepare={prepare_path}", flush=True) @@ -394,35 +376,7 @@ def _validate_run_args(args: RunBenchArgs) -> None: if args.gpu_util <= 0.0 or args.gpu_util > 1.0: raise ValueError("gpu_util must be in (0, 1]") if args.load_generator == "vllm-bench": - if args.bench_num_prompts <= 0: - raise ValueError("bench_num_prompts must be positive") - if args.bench_input_len <= 0: - raise ValueError("bench_input_len must be positive") - if _bench_output_len(args) <= 0: - raise ValueError("bench_output_len must be positive") - if _bench_max_concurrency(args) <= 0: - raise ValueError("bench_max_concurrency must be positive") - if args.bench_random_prefix_len < 0: - raise ValueError("bench_random_prefix_len must be non-negative") - if args.bench_random_range_ratio < 0.0: - raise ValueError("bench_random_range_ratio must be non-negative") - if args.bench_burstiness <= 0.0: - raise ValueError("bench_burstiness must be positive") - _validate_bench_request_rate(args.bench_request_rate) - - -def _validate_bench_request_rate(value: str) -> None: - """Validate the vLLM bench request-rate argument.""" - if value == "inf": - return - try: - rate = float(value) - except ValueError as exc: - raise ValueError( - "bench_request_rate must be 'inf' or a positive number" - ) from exc - if rate <= 0.0 or math.isinf(rate) or math.isnan(rate): - raise ValueError("bench_request_rate must be 'inf' or a positive number") + vllm_bench.validate_args(args) def _run_backend( @@ -465,7 +419,15 @@ def _run_backend( _print_stage("cold/warm load", backend_run.label) _print_kv("output", result_path) if args.load_generator == "vllm-bench": - _run_vllm_bench_load(args, manifest, backend_run, backend_dir, result_path) + vllm_bench.run_load( + args, + manifest, + backend_run, + backend_dir, + result_path, + run_command=_run_command, + print_kv=_print_kv, + ) else: _run_command(_load_command(args, backend_dir, prepare_path, result_path)) if manifest is not None: @@ -479,373 +441,6 @@ def _run_backend( return result_path -def _bench_prepare_config(args: RunBenchArgs, run_root: Path) -> dict[str, Any]: - """Build prepare config for synthetic vLLM bench random workloads. - - Args: - args: Benchmark runner arguments. - run_root: Run root used for capacity probing. - - Returns: - JSON-serializable prepare config. - - Thread-safety: - Reads current disk and host memory state through sizing helpers. - """ - prompt_tokens = _bench_max_prompt_tokens(args) - max_prompt_blocks = max(1, math.ceil(prompt_tokens / args.block_size)) - total_blocks = args.bench_num_prompts * max_prompt_blocks - slot_size = slot_size_for_block_tokens(args.block_size) - sizing = derive_benchmark_sizing( - total_blocks=total_blocks, - max_prompt_blocks=max_prompt_blocks, - slot_size=slot_size, - mode=COMPARISON_IOURING_MEM, - evict=args.evict, - capacity_limits=_bench_capacity_limits(args, run_root), - ) - return { - "dataset": "vllm-bench-random", - "num_samples": args.bench_num_prompts, - "max_inflight": _bench_max_concurrency(args), - "gen_params": { - "max_tokens": _bench_output_len(args), - "temperature": 0.0, - "top_p": 1.0, - "seed": args.bench_seed, - }, - "total_prompt_tokens": args.bench_num_prompts * prompt_tokens, - "total_blocks": total_blocks, - "max_prompt_blocks": max_prompt_blocks, - "max_prompt_tokens": prompt_tokens, - "block_size": args.block_size, - "bench_num_prompts": args.bench_num_prompts, - "bench_input_len": args.bench_input_len, - "bench_output_len": _bench_output_len(args), - "bench_request_rate": args.bench_request_rate, - "bench_max_concurrency": _bench_max_concurrency(args), - "bench_random_prefix_len": args.bench_random_prefix_len, - "bench_random_range_ratio": args.bench_random_range_ratio, - "bench_seed": args.bench_seed, - "derived_l1_size_bytes": sizing.daser_l1_bytes, - "derived_l1_size": format_capacity(sizing.daser_l1_bytes), - "derived_l2_size_bytes": sizing.daser_l2_bytes, - "derived_l2_size": format_capacity(sizing.daser_l2_bytes), - "lmcache_l1_gb": sizing.lmcache_cpu_gb, - "lmcache_l2_gb": sizing.lmcache_disk_gb, - "capacity_capped": sizing.capacity_capped, - "evict": args.evict, - "planned_skip_l2": not args.evict, - } - - -def _bench_capacity_limits( - args: RunBenchArgs, - run_root: Path, -) -> BenchmarkCapacityLimits: - return derive_capacity_limits(run_root) - - -def _bench_max_prompt_tokens(args: RunBenchArgs) -> int: - variable_tokens = math.ceil( - args.bench_input_len * (1.0 + args.bench_random_range_ratio) - ) - return max(1, args.bench_random_prefix_len + variable_tokens) - - -def _bench_output_len(args: RunBenchArgs) -> int: - if args.bench_output_len is not None: - return args.bench_output_len - return args.gen_max_tokens - - -def _bench_max_concurrency(args: RunBenchArgs) -> int: - if args.bench_max_concurrency is not None: - return args.bench_max_concurrency - return args.max_inflight - - -def _run_vllm_bench_load( - args: RunBenchArgs, - manifest: BenchmarkManifest | None, - backend_run: BackendRun, - backend_dir: Path, - result_path: Path, -) -> None: - if manifest is None: - manifest = BenchmarkManifest.read(backend_dir / "manifest.json") - if backend_run.backend == "vllm": - raw = backend_dir / "vllm_bench_baseline.json" - baseline_metrics, baseline_hit_rate = _run_vllm_bench_phase( - args, - manifest, - raw, - ) - baseline_summary = _normalise_vllm_bench_result(raw) - _apply_vllm_bench_phase_metrics(baseline_summary, baseline_hit_rate) - result = { - "manifest": _manifest_payload(manifest), - "result": { - "baseline": { - "summary": baseline_summary, - "metrics": baseline_metrics, - } - }, - } - else: - cold_raw = backend_dir / "vllm_bench_cold.json" - warm_raw = backend_dir / "vllm_bench_warm.json" - cold_metrics, cold_hit_rate = _run_vllm_bench_phase( - args, - manifest, - cold_raw, - ) - if backend_run.backend == "lmcache": - _print_kv("lmcache_warm_wait", "quiescent") - asyncio.run(_wait_lmcache_quiescent(manifest, settle_seconds=0.0)) - elif backend_run.backend == "daser": - _drain_daser(manifest) - warm_metrics, warm_hit_rate = _run_vllm_bench_phase( - args, - manifest, - warm_raw, - ) - cold_summary = _normalise_vllm_bench_result(cold_raw) - warm_summary = _normalise_vllm_bench_result(warm_raw) - _apply_vllm_bench_phase_metrics(cold_summary, cold_hit_rate) - _apply_vllm_bench_phase_metrics(warm_summary, warm_hit_rate) - result = { - "manifest": _manifest_payload(manifest), - "result": { - "cold": {"summary": cold_summary, "metrics": cold_metrics}, - "warm": {"summary": warm_summary, "metrics": warm_metrics}, - }, - "correctness": _compare_vllm_bench_outputs(cold_raw, warm_raw), - } - result_path.write_text(json.dumps(result, indent=2), encoding="utf-8") - - -def _run_vllm_bench_phase( - args: RunBenchArgs, - manifest: BenchmarkManifest, - raw_path: Path, -) -> tuple[dict[str, Any], float | None]: - before_metrics = asyncio.run(collect_phase_metrics(manifest)) - _run_command(_vllm_bench_command(args, manifest.endpoints["vllm"], raw_path)) - return _collect_vllm_bench_phase_metrics(manifest, before_metrics) - - -def _collect_vllm_bench_phase_metrics( - manifest: BenchmarkManifest, - before_metrics: dict[str, Any] | None, -) -> tuple[dict[str, Any], float | None]: - """Collect vLLM bench phase metrics and backend token hit rate. - - Args: - manifest: Started benchmark service manifest. - before_metrics: Optional pre-phase metric snapshot. If absent, the - returned metrics are an empty delta. - - Returns: - Phase metric deltas and the backend token-level cache hit ratio. - - Thread-safety: - Runs asynchronous metric collection in this orchestration process. - """ - if before_metrics is None: - before_metrics = { - "vllm_prometheus": {}, - "backend_prometheus": {}, - "backend_status": {}, - } - metrics = asyncio.run(collect_phase_metrics(manifest, before_metrics)) - hit_ratios = metrics.get("hit_ratios", {}) if isinstance(metrics, dict) else {} - return metrics, backend_server_hit_rate(hit_ratios) - - -def _apply_vllm_bench_phase_metrics( - summary: dict[str, Any], - backend_hit_rate: float | None, -) -> None: - if backend_hit_rate is not None: - summary["backend_server_cache_hit_rate"] = backend_hit_rate - - -def _vllm_bench_command( - args: RunBenchArgs, - endpoint: Any, - raw_path: Path, -) -> list[str]: - """Build a vLLM bench serve command for one benchmark phase.""" - return [ - "vllm", - "bench", - "serve", - "--backend", - "openai", - "--base-url", - endpoint.url, - "--endpoint", - "/v1/completions", - "--model", - args.model, - "--dataset-name", - "random", - "--num-prompts", - str(args.bench_num_prompts), - "--input-len", - str(args.bench_input_len), - "--output-len", - str(_bench_output_len(args)), - "--request-rate", - str(args.bench_request_rate), - "--max-concurrency", - str(_bench_max_concurrency(args)), - "--random-prefix-len", - str(args.bench_random_prefix_len), - "--random-range-ratio", - str(args.bench_random_range_ratio), - "--seed", - str(args.bench_seed), - "--burstiness", - str(args.bench_burstiness), - "--temperature", - "0.0", - "--top-p", - "1.0", - "--percentile-metrics", - "ttft,tpot,itl,e2el", - "--save-result", - "--save-detailed", - "--result-dir", - str(raw_path.parent), - "--result-filename", - raw_path.name, - ] - - -def _normalise_vllm_bench_result(path: Path) -> dict[str, Any]: - """Convert a vLLM bench JSON result to the benchmark summary shape.""" - payload = json.loads(path.read_text(encoding="utf-8")) - duration_s = float(_first_number(payload, ("duration", "benchmark_duration"), 0.0)) - prompt_tokens = int(_first_number(payload, ("total_input_tokens",), 0)) - completion_tokens = int(_first_number(payload, ("total_output_tokens",), 0)) - return { - "num_requests": int(_first_number(payload, ("completed",), 0)), - "num_errors": int(_first_number(payload, ("failed",), 0)), - "ttft_ms_mean": float(_first_number(payload, ("mean_ttft_ms",), 0.0)), - "latency_ms_mean": float( - _first_number( - payload, - ("mean_e2el_ms", "mean_latency_ms", "mean_ttft_ms"), - 0.0, - ) - ), - "phase_elapsed_ms": duration_s * 1000.0, - "phase_prompt_tok_per_s": prompt_tokens / duration_s if duration_s > 0 else 0.0, - "prompt_tokens_total": prompt_tokens, - "completion_tokens_total": completion_tokens, - } - - -def _compare_vllm_bench_outputs(cold_path: Path, warm_path: Path) -> dict[str, Any]: - """Compare detailed vLLM bench generated text across cold and warm phases.""" - cold = json.loads(cold_path.read_text(encoding="utf-8")) - warm = json.loads(warm_path.read_text(encoding="utf-8")) - cold_texts = _generated_texts_from_vllm_bench(cold) - warm_texts = _generated_texts_from_vllm_bench(warm) - if cold_texts is None or warm_texts is None: - return { - "cold_warm_exact_match": { - "available": False, - "matches": 0, - "total": 0, - "accuracy": None, - "reason": "vLLM bench result did not include generated text details", - } - } - paired = min(len(cold_texts), len(warm_texts)) - total = max(len(cold_texts), len(warm_texts)) - matches = sum(1 for idx in range(paired) if cold_texts[idx] == warm_texts[idx]) - return { - "cold_warm_exact_match": { - "available": True, - "matches": matches, - "total": total, - "accuracy": matches / total if total else None, - "length_mismatch": len(cold_texts) != len(warm_texts), - } - } - - -def _generated_texts_from_vllm_bench(payload: dict[str, Any]) -> list[str] | None: - """Extract generated texts from known vLLM bench result shapes.""" - generated_texts = payload.get("generated_texts") - if isinstance(generated_texts, list) and all( - isinstance(text, str) for text in generated_texts - ): - return generated_texts - outputs = payload.get("outputs") - if not isinstance(outputs, list): - return None - texts: list[str] = [] - for output in outputs: - if not isinstance(output, dict): - return None - text = output.get("generated_text") - if not isinstance(text, str): - return None - texts.append(text) - return texts - - -def _first_number( - payload: dict[str, Any], - keys: tuple[str, ...], - default: float, -) -> float: - for key in keys: - value = payload.get(key) - if isinstance(value, int | float): - return float(value) - return default - - -def _manifest_payload(manifest: BenchmarkManifest) -> dict[str, Any]: - return { - "run_id": manifest.run_id, - "backend": manifest.backend, - "reuse_mode": manifest.reuse_mode, - "model": manifest.model, - "store_dir": manifest.store_dir, - "l1_size_bytes": manifest.l1_size_bytes, - "l2_size_bytes": manifest.l2_size_bytes, - "skip_l2": manifest.skip_l2, - "endpoints": { - name: {"url": endpoint.url} for name, endpoint in manifest.endpoints.items() - }, - "log_dir": manifest.log_dir, - "pid_file": manifest.pid_file, - "block_size": manifest.block_size, - } - - -def _wait_with_message(label: str, seconds: float) -> None: - _print_kv(label, seconds) - time.sleep(seconds) - - -def _drain_daser(manifest: BenchmarkManifest) -> None: - endpoint = manifest.endpoints.get("daser") - if endpoint is None: - return - try: - response = httpx.post(f"{endpoint.url}/drain", timeout=30.0) - response.raise_for_status() - except Exception as exc: # noqa: BLE001 - _print_kv("daser_drain_status", f"unavailable ({exc})") - - def _prepare_command( args: RunBenchArgs, run_root: Path, prepare_path: Path ) -> list[str]: diff --git a/benchmarks/utils/constants.py b/benchmarks/utils/constants.py index 16a6a36..7b117c3 100644 --- a/benchmarks/utils/constants.py +++ b/benchmarks/utils/constants.py @@ -9,9 +9,7 @@ HEAD_DIM: int = 128 NUM_LAYERS: int = 36 DTYPE_BYTES: int = 2 -BENCHMARK_SEED: int = 42 -COMPARISON_GDS = "gds-vs-lmcache-local-ssd" COMPARISON_IOURING_MEM = "iouring-mem-vs-lmcache-local-ssd-mem" DEFAULT_SYSTEM_PROMPT: str = ( diff --git a/benchmarks/utils/datasets.py b/benchmarks/utils/datasets.py index bd67074..67e3f16 100644 --- a/benchmarks/utils/datasets.py +++ b/benchmarks/utils/datasets.py @@ -4,6 +4,7 @@ from __future__ import annotations from abc import ABC, abstractmethod +import argparse import csv from dataclasses import dataclass import json @@ -36,7 +37,16 @@ class BenchmarkSample: class BenchmarkDataset(ABC): - """Abstract benchmark dataset loader.""" + """Abstract benchmark dataset loader. + + Subclasses own their dataset name, their CLI arguments, and how to build + themselves from parsed args, so adding a dataset needs no edits to the + benchmark runner or load generator: implement the three hooks and register + the class with ``register_dataset``. + """ + + #: Registry key; also the value accepted by ``--dataset``. + name: str = "" @abstractmethod def load(self) -> list[BenchmarkSample]: @@ -49,10 +59,102 @@ def load(self) -> list[BenchmarkSample]: Implementations perform local file reads and keep no global state. """ + @classmethod # noqa: B027 + def add_cli_args(cls, parser: argparse.ArgumentParser) -> None: + """Register this dataset's CLI arguments on ``parser``. + + Args: + parser: argument parser shared by all datasets. Implementations add + only their own options and tolerate being called alongside + sibling datasets, so option names must not collide. + """ + + @classmethod + @abstractmethod + def from_args(cls, args: argparse.Namespace) -> "BenchmarkDataset": + """Build a dataset instance from parsed CLI args. + + Args: + args: parsed argparse namespace containing this dataset's options + plus the shared ``max_samples``. + + Returns: + A ready-to-load dataset instance. + """ + + +_DATASET_REGISTRY: dict[str, type[BenchmarkDataset]] = {} + + +def register_dataset(cls: type[BenchmarkDataset]) -> type[BenchmarkDataset]: + """Register a dataset class under its ``name`` for CLI dispatch. + + Args: + cls: dataset class with a non-empty ``name``. + + Returns: + The same class, so this can be used as a decorator. + + Raises: + ValueError: if ``name`` is empty or already registered. + """ + if not cls.name: + raise ValueError(f"{cls.__name__} must set a non-empty name") + if cls.name in _DATASET_REGISTRY: + raise ValueError(f"dataset name already registered: {cls.name}") + _DATASET_REGISTRY[cls.name] = cls + return cls + +def dataset_names() -> tuple[str, ...]: + """Return registered dataset names in registration order.""" + return tuple(_DATASET_REGISTRY) + + +def add_dataset_cli_args( + parser: argparse.ArgumentParser, default: str | None = None +) -> None: + """Add ``--dataset`` plus every registered dataset's own CLI args. + + Args: + parser: parser to extend with the dataset selector and per-dataset + options. + default: default dataset name; falls back to the first registered. + """ + names = dataset_names() + parser.add_argument( + "--dataset", + choices=names, + default=default if default is not None else (names[0] if names else None), + ) + for cls in _DATASET_REGISTRY.values(): + cls.add_cli_args(parser) + + +def build_dataset(args: argparse.Namespace) -> BenchmarkDataset: + """Build the selected dataset from parsed args via the registry. + + Args: + args: parsed argparse namespace with a ``dataset`` selector. + + Returns: + The dataset instance for ``args.dataset``. + + Raises: + ValueError: if ``args.dataset`` is not registered. + """ + cls = _DATASET_REGISTRY.get(args.dataset) + if cls is None: + raise ValueError(f"unknown dataset: {args.dataset}") + return cls.from_args(args) + + +@register_dataset class ImdbDataset(BenchmarkDataset): """Load IMDB CSV reviews as benchmark samples.""" + name = "imdb" + def __init__( self, path: str | Path, @@ -70,6 +172,18 @@ def __init__( self._max_samples = max_samples self._question = question + @classmethod + def add_cli_args(cls, parser: argparse.ArgumentParser) -> None: + """Register the IMDB CSV path option.""" + parser.add_argument("--imdb", help="IMDB CSV path with a review column") + + @classmethod + def from_args(cls, args: argparse.Namespace) -> "ImdbDataset": + """Build an IMDB dataset from parsed args.""" + if not getattr(args, "imdb", None): + raise ValueError("--imdb is required for --dataset imdb") + return cls(args.imdb, max_samples=args.max_samples) + def load(self) -> list[BenchmarkSample]: """Load IMDB reviews as benchmark samples.""" if not self._path.is_file(): @@ -95,9 +209,12 @@ def load(self) -> list[BenchmarkSample]: return samples +@register_dataset class LongBenchDataset(BenchmarkDataset): """Load one or more LongBench JSONL files.""" + name = "longbench" + def __init__( self, data_dir: str | Path, @@ -116,6 +233,30 @@ def __init__( self._datasets = datasets self._max_samples = max_samples + @classmethod + def add_cli_args(cls, parser: argparse.ArgumentParser) -> None: + """Register the LongBench directory and dataset-subset options.""" + parser.add_argument( + "--longbench-dir", help="Directory containing LongBench *.jsonl files" + ) + parser.add_argument( + "--datasets", + default=None, + help="Comma-separated LongBench dataset names; default discovers all", + ) + + @classmethod + def from_args(cls, args: argparse.Namespace) -> "LongBenchDataset": + """Build a LongBench dataset from parsed args.""" + if not getattr(args, "longbench_dir", None): + raise ValueError("--longbench-dir is required for --dataset longbench") + datasets = None + if getattr(args, "datasets", None): + datasets = [ + item.strip() for item in args.datasets.split(",") if item.strip() + ] + return cls(args.longbench_dir, datasets=datasets, max_samples=args.max_samples) + def load(self) -> list[BenchmarkSample]: """Load LongBench JSONL records as benchmark samples.""" if not self._data_dir.is_dir(): diff --git a/benchmarks/utils/loadgen.py b/benchmarks/utils/loadgen.py index 16e5b08..5bafde9 100644 --- a/benchmarks/utils/loadgen.py +++ b/benchmarks/utils/loadgen.py @@ -176,7 +176,6 @@ async def run_daser_prefix( max_inflight: int, gen_params: dict[str, Any], timeout: float, - settle_seconds: float = 0.0, ) -> dict[str, Any]: """Run DaseR prefix cold and warm full-prompt phases.""" before_cold = await collect_phase_metrics(manifest) @@ -195,7 +194,7 @@ async def run_daser_prefix( metrics=await collect_phase_metrics(manifest, before_cold), elapsed_ms=cold_elapsed_ms, ) - await _wait_daser_drained(manifest, settle_seconds) + await _wait_daser_drained(manifest) before_warm = await collect_phase_metrics(manifest) warm, warm_elapsed_ms = await _run_vllm_phase_requests( manifest, @@ -512,13 +511,10 @@ def _lmcache_is_quiescent(status: dict[str, Any]) -> bool: return all(int(mapping.get(key, 0)) == 0 for mapping, key in zero_fields) -async def _wait_daser_drained( - manifest: BenchmarkManifest, settle_seconds: float -) -> None: +async def _wait_daser_drained(manifest: BenchmarkManifest) -> None: daser = manifest.endpoints.get("daser") if daser is None: return - del settle_seconds async with httpx.AsyncClient(timeout=httpx.Timeout(30.0)) as client: response = await client.post(f"{daser.url}/drain") response.raise_for_status() diff --git a/benchmarks/utils/prompts.py b/benchmarks/utils/prompts.py index 0c3e603..554f1e1 100644 --- a/benchmarks/utils/prompts.py +++ b/benchmarks/utils/prompts.py @@ -189,26 +189,6 @@ def build_prompt_payloads( return build_prompts(tokenizer, samples) -def count_prompt_tokens(tokenizer: Any, prompts: list[str]) -> list[int]: - """Count prompt tokens without adding special tokens. - - Args: - tokenizer: Hugging Face tokenizer or compatible object. - prompts: Prompt strings. - - Returns: - Token counts aligned with prompts. - - Thread-safety: - Depends on tokenizer implementation; this function keeps no state. - """ - counts: list[int] = [] - for prompt in prompts: - encoded = tokenizer(prompt, add_special_tokens=False) - counts.append(len(encoded["input_ids"])) - return counts - - def count_prompt_payload_tokens( tokenizer: Any, prompts: list[str | list[int]], diff --git a/benchmarks/utils/vllm_bench.py b/benchmarks/utils/vllm_bench.py new file mode 100644 index 0000000..52ab536 --- /dev/null +++ b/benchmarks/utils/vllm_bench.py @@ -0,0 +1,442 @@ +# SPDX-License-Identifier: Apache-2.0 +"""vLLM ``bench serve`` load generator for the benchmark runner. + +This module isolates the synthetic ``--load-generator vllm-bench`` path so the +orchestrator in ``run_bench.py`` only forks on which generator to use, not on +how each one works. It drives vLLM's built-in random workload through +``vllm bench serve`` and normalizes the result into the shared summary shape. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from dataclasses import asdict +import json +import math +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import httpx + +from benchmarks.utils.constants import ( + COMPARISON_IOURING_MEM, + slot_size_for_block_tokens, +) +from benchmarks.utils.loadgen import ( + _wait_lmcache_quiescent, + backend_server_hit_rate, + collect_phase_metrics, +) +from benchmarks.utils.servers import BenchmarkManifest +from benchmarks.utils.sizing import ( + derive_benchmark_sizing, + derive_capacity_limits, + format_capacity, +) + +if TYPE_CHECKING: + from benchmarks.run_bench import BackendRun, RunBenchArgs + + +def bench_output_len(args: RunBenchArgs) -> int: + """Return the resolved vLLM bench output length. + + Args: + args: Benchmark runner arguments. + + Returns: + ``bench_output_len`` when set, else the shared ``gen_max_tokens``. + """ + if args.bench_output_len is not None: + return args.bench_output_len + return args.gen_max_tokens + + +def bench_max_concurrency(args: RunBenchArgs) -> int: + """Return the resolved vLLM bench max concurrency. + + Args: + args: Benchmark runner arguments. + + Returns: + ``bench_max_concurrency`` when set, else the shared ``max_inflight``. + """ + if args.bench_max_concurrency is not None: + return args.bench_max_concurrency + return args.max_inflight + + +def validate_args(args: RunBenchArgs) -> None: + """Validate vLLM bench arguments with clear preflight errors. + + Args: + args: Benchmark runner arguments. + + Raises: + ValueError: If any vLLM bench argument is out of range. + """ + if args.bench_num_prompts <= 0: + raise ValueError("bench_num_prompts must be positive") + if args.bench_input_len <= 0: + raise ValueError("bench_input_len must be positive") + if bench_output_len(args) <= 0: + raise ValueError("bench_output_len must be positive") + if bench_max_concurrency(args) <= 0: + raise ValueError("bench_max_concurrency must be positive") + if args.bench_random_prefix_len < 0: + raise ValueError("bench_random_prefix_len must be non-negative") + if args.bench_random_range_ratio < 0.0: + raise ValueError("bench_random_range_ratio must be non-negative") + if args.bench_burstiness <= 0.0: + raise ValueError("bench_burstiness must be positive") + _validate_request_rate(args.bench_request_rate) + + +def _validate_request_rate(value: str) -> None: + """Validate the vLLM bench request-rate argument.""" + if value == "inf": + return + try: + rate = float(value) + except ValueError as exc: + raise ValueError( + "bench_request_rate must be 'inf' or a positive number" + ) from exc + if rate <= 0.0 or math.isinf(rate) or math.isnan(rate): + raise ValueError("bench_request_rate must be 'inf' or a positive number") + + +def _max_prompt_tokens(args: RunBenchArgs) -> int: + variable_tokens = math.ceil( + args.bench_input_len * (1.0 + args.bench_random_range_ratio) + ) + return max(1, args.bench_random_prefix_len + variable_tokens) + + +def prepare_config(args: RunBenchArgs, run_root: Path) -> dict[str, Any]: + """Build the prepare config for synthetic vLLM bench random workloads. + + Args: + args: Benchmark runner arguments. + run_root: Run root used for capacity probing. + + Returns: + JSON-serializable prepare config. + + Thread-safety: + Reads current disk and host memory state through sizing helpers. + """ + prompt_tokens = _max_prompt_tokens(args) + max_prompt_blocks = max(1, math.ceil(prompt_tokens / args.block_size)) + total_blocks = args.bench_num_prompts * max_prompt_blocks + slot_size = slot_size_for_block_tokens(args.block_size) + sizing = derive_benchmark_sizing( + total_blocks=total_blocks, + max_prompt_blocks=max_prompt_blocks, + slot_size=slot_size, + mode=COMPARISON_IOURING_MEM, + evict=args.evict, + capacity_limits=derive_capacity_limits(run_root), + ) + return { + "dataset": "vllm-bench-random", + "num_samples": args.bench_num_prompts, + "max_inflight": bench_max_concurrency(args), + "gen_params": { + "max_tokens": bench_output_len(args), + "temperature": 0.0, + "top_p": 1.0, + "seed": args.bench_seed, + }, + "total_prompt_tokens": args.bench_num_prompts * prompt_tokens, + "total_blocks": total_blocks, + "max_prompt_blocks": max_prompt_blocks, + "max_prompt_tokens": prompt_tokens, + "block_size": args.block_size, + "bench_num_prompts": args.bench_num_prompts, + "bench_input_len": args.bench_input_len, + "bench_output_len": bench_output_len(args), + "bench_request_rate": args.bench_request_rate, + "bench_max_concurrency": bench_max_concurrency(args), + "bench_random_prefix_len": args.bench_random_prefix_len, + "bench_random_range_ratio": args.bench_random_range_ratio, + "bench_seed": args.bench_seed, + "derived_l1_size_bytes": sizing.daser_l1_bytes, + "derived_l1_size": format_capacity(sizing.daser_l1_bytes), + "derived_l2_size_bytes": sizing.daser_l2_bytes, + "derived_l2_size": format_capacity(sizing.daser_l2_bytes), + "lmcache_l1_gb": sizing.lmcache_cpu_gb, + "lmcache_l2_gb": sizing.lmcache_disk_gb, + "capacity_capped": sizing.capacity_capped, + "evict": args.evict, + "planned_skip_l2": not args.evict, + } + + +def run_load( + args: RunBenchArgs, + manifest: BenchmarkManifest | None, + backend_run: BackendRun, + backend_dir: Path, + result_path: Path, + *, + run_command: Callable[[list[str]], Any], + print_kv: Callable[[str, Any], None], +) -> None: + """Run cold/warm (or baseline) vLLM bench phases and write results.json. + + Args: + args: Benchmark runner arguments. + manifest: Started service manifest; read from disk when None. + backend_run: Resolved backend row. + backend_dir: Per-backend output directory. + result_path: results.json destination. + run_command: Callable that runs a subprocess command (injected). + print_kv: Callable that prints a key/value progress line (injected). + """ + if manifest is None: + manifest = BenchmarkManifest.read(backend_dir / "manifest.json") + if backend_run.backend == "vllm": + raw = backend_dir / "vllm_bench_baseline.json" + baseline_metrics, baseline_hit_rate = _run_phase( + args, manifest, raw, run_command=run_command + ) + baseline_summary = _normalise_result(raw) + _apply_phase_metrics(baseline_summary, baseline_hit_rate) + result = { + "manifest": asdict(manifest), + "result": { + "baseline": { + "summary": baseline_summary, + "metrics": baseline_metrics, + } + }, + } + else: + cold_raw = backend_dir / "vllm_bench_cold.json" + warm_raw = backend_dir / "vllm_bench_warm.json" + cold_metrics, cold_hit_rate = _run_phase( + args, manifest, cold_raw, run_command=run_command + ) + if backend_run.backend == "lmcache": + print_kv("lmcache_warm_wait", "quiescent") + asyncio.run(_wait_lmcache_quiescent(manifest, settle_seconds=0.0)) + elif backend_run.backend == "daser": + _drain_daser(manifest, print_kv=print_kv) + warm_metrics, warm_hit_rate = _run_phase( + args, manifest, warm_raw, run_command=run_command + ) + cold_summary = _normalise_result(cold_raw) + warm_summary = _normalise_result(warm_raw) + _apply_phase_metrics(cold_summary, cold_hit_rate) + _apply_phase_metrics(warm_summary, warm_hit_rate) + result = { + "manifest": asdict(manifest), + "result": { + "cold": {"summary": cold_summary, "metrics": cold_metrics}, + "warm": {"summary": warm_summary, "metrics": warm_metrics}, + }, + "correctness": _compare_outputs(cold_raw, warm_raw), + } + result_path.write_text(json.dumps(result, indent=2), encoding="utf-8") + + +def _run_phase( + args: RunBenchArgs, + manifest: BenchmarkManifest, + raw_path: Path, + *, + run_command: Callable[[list[str]], Any], +) -> tuple[dict[str, Any], float | None]: + before_metrics = asyncio.run(collect_phase_metrics(manifest)) + run_command(_bench_command(args, manifest.endpoints["vllm"], raw_path)) + return _collect_phase_metrics(manifest, before_metrics) + + +def _collect_phase_metrics( + manifest: BenchmarkManifest, + before_metrics: dict[str, Any] | None, +) -> tuple[dict[str, Any], float | None]: + """Collect vLLM bench phase metric deltas and backend token hit rate. + + Args: + manifest: Started benchmark service manifest. + before_metrics: Optional pre-phase metric snapshot; absent yields an + empty delta. + + Returns: + Phase metric deltas and the backend token-level cache hit ratio. + """ + if before_metrics is None: + before_metrics = { + "vllm_prometheus": {}, + "backend_prometheus": {}, + "backend_status": {}, + } + metrics = asyncio.run(collect_phase_metrics(manifest, before_metrics)) + hit_ratios = metrics.get("hit_ratios", {}) if isinstance(metrics, dict) else {} + return metrics, backend_server_hit_rate(hit_ratios) + + +def _apply_phase_metrics( + summary: dict[str, Any], + backend_hit_rate: float | None, +) -> None: + if backend_hit_rate is not None: + summary["backend_server_cache_hit_rate"] = backend_hit_rate + + +def _drain_daser( + manifest: BenchmarkManifest, + *, + print_kv: Callable[[str, Any], None], +) -> None: + endpoint = manifest.endpoints.get("daser") + if endpoint is None: + return + try: + response = httpx.post(f"{endpoint.url}/drain", timeout=30.0) + response.raise_for_status() + except Exception as exc: # noqa: BLE001 + print_kv("daser_drain_status", f"unavailable ({exc})") + + +def _bench_command( + args: RunBenchArgs, + endpoint: Any, + raw_path: Path, +) -> list[str]: + """Build a ``vllm bench serve`` command for one benchmark phase.""" + return [ + "vllm", + "bench", + "serve", + "--backend", + "openai", + "--base-url", + endpoint.url, + "--endpoint", + "/v1/completions", + "--model", + args.model, + "--dataset-name", + "random", + "--num-prompts", + str(args.bench_num_prompts), + "--input-len", + str(args.bench_input_len), + "--output-len", + str(bench_output_len(args)), + "--request-rate", + str(args.bench_request_rate), + "--max-concurrency", + str(bench_max_concurrency(args)), + "--random-prefix-len", + str(args.bench_random_prefix_len), + "--random-range-ratio", + str(args.bench_random_range_ratio), + "--seed", + str(args.bench_seed), + "--burstiness", + str(args.bench_burstiness), + "--temperature", + "0.0", + "--top-p", + "1.0", + "--percentile-metrics", + "ttft,tpot,itl,e2el", + "--save-result", + "--save-detailed", + "--result-dir", + str(raw_path.parent), + "--result-filename", + raw_path.name, + ] + + +def _normalise_result(path: Path) -> dict[str, Any]: + """Convert a vLLM bench JSON result to the benchmark summary shape.""" + payload = json.loads(path.read_text(encoding="utf-8")) + duration_s = float(_first_number(payload, ("duration", "benchmark_duration"), 0.0)) + prompt_tokens = int(_first_number(payload, ("total_input_tokens",), 0)) + completion_tokens = int(_first_number(payload, ("total_output_tokens",), 0)) + return { + "num_requests": int(_first_number(payload, ("completed",), 0)), + "num_errors": int(_first_number(payload, ("failed",), 0)), + "ttft_ms_mean": float(_first_number(payload, ("mean_ttft_ms",), 0.0)), + "latency_ms_mean": float( + _first_number( + payload, + ("mean_e2el_ms", "mean_latency_ms", "mean_ttft_ms"), + 0.0, + ) + ), + "phase_elapsed_ms": duration_s * 1000.0, + "phase_prompt_tok_per_s": prompt_tokens / duration_s if duration_s > 0 else 0.0, + "prompt_tokens_total": prompt_tokens, + "completion_tokens_total": completion_tokens, + } + + +def _compare_outputs(cold_path: Path, warm_path: Path) -> dict[str, Any]: + """Compare detailed vLLM bench generated text across cold and warm phases.""" + cold = json.loads(cold_path.read_text(encoding="utf-8")) + warm = json.loads(warm_path.read_text(encoding="utf-8")) + cold_texts = _generated_texts(cold) + warm_texts = _generated_texts(warm) + if cold_texts is None or warm_texts is None: + return { + "cold_warm_exact_match": { + "available": False, + "matches": 0, + "total": 0, + "accuracy": None, + "reason": "vLLM bench result did not include generated text details", + } + } + paired = min(len(cold_texts), len(warm_texts)) + total = max(len(cold_texts), len(warm_texts)) + matches = sum(1 for idx in range(paired) if cold_texts[idx] == warm_texts[idx]) + return { + "cold_warm_exact_match": { + "available": True, + "matches": matches, + "total": total, + "accuracy": matches / total if total else None, + "length_mismatch": len(cold_texts) != len(warm_texts), + } + } + + +def _generated_texts(payload: dict[str, Any]) -> list[str] | None: + """Extract generated texts from known vLLM bench result shapes.""" + generated_texts = payload.get("generated_texts") + if isinstance(generated_texts, list) and all( + isinstance(text, str) for text in generated_texts + ): + return generated_texts + outputs = payload.get("outputs") + if not isinstance(outputs, list): + return None + texts: list[str] = [] + for output in outputs: + if not isinstance(output, dict): + return None + text = output.get("generated_text") + if not isinstance(text, str): + return None + texts.append(text) + return texts + + +def _first_number( + payload: dict[str, Any], + keys: tuple[str, ...], + default: float, +) -> float: + for key in keys: + value = payload.get(key) + if isinstance(value, int | float): + return float(value) + return default diff --git a/daser/config.py b/daser/config.py index 0c2c2b3..5ff2331 100644 --- a/daser/config.py +++ b/daser/config.py @@ -8,6 +8,12 @@ BLOCK_TOKENS = 16 DEFAULT_IOURING_L1_BYTES = 1024 * 1024 * 1024 +#: Cache reuse mode selecting the retrieval index + position encoder pair. +CACHE_REUSE_PREFIX = "prefix" +CACHE_REUSE_CHUNK = "chunk" +CACHE_REUSE_MODES = (CACHE_REUSE_PREFIX, CACHE_REUSE_CHUNK) +DEFAULT_CACHE_REUSE_MODE = CACHE_REUSE_CHUNK + @dataclass(frozen=True) class ModelGeometry: @@ -43,11 +49,6 @@ def slot_size_for_block_tokens(self, block_tokens: int) -> int: * self.dtype_bytes ) - @property - def slot_size(self) -> int: - """Return bytes required for one default vLLM KV block.""" - return self.slot_size_for_block_tokens(BLOCK_TOKENS) - def _dtype_bytes(dtype: object) -> int: """Return storage bytes for a HuggingFace dtype string. @@ -148,7 +149,7 @@ class DaserConfig: log_level: str = "INFO" block_tokens: int = BLOCK_TOKENS - cache_reuse_mode: str = "chunk" + cache_reuse_mode: str = DEFAULT_CACHE_REUSE_MODE transfer_mode: str = "iouring" l1_size_bytes: int = DEFAULT_IOURING_L1_BYTES skip_l2: bool = False diff --git a/daser/connector/daser_connector.py b/daser/connector/daser_connector.py index 6c6910e..9568941 100644 --- a/daser/connector/daser_connector.py +++ b/daser/connector/daser_connector.py @@ -100,10 +100,6 @@ def __init__( self._skip_l2: bool = bool(extra.get("skip_l2", False)) self._cache_reuse_strategy: CacheReuseStrategy self._set_cache_reuse_strategy(str(extra.get("cache_reuse_mode", "chunk"))) - cache_config = getattr(vllm_config, "cache_config", None) - self._vllm_prefix_caching_enabled = bool( - getattr(cache_config, "enable_prefix_caching", False) - ) self._runtime_config_ready = False self._rope_base: float = 10000.0 self._rope_rotary_dim: int = 0 @@ -128,7 +124,6 @@ def __init__( self._transfer_mode = str(extra.get("transfer_mode", "iouring")) self._ipc_load_async = IPCClientAsync(self._socket_path) self._ipc_store_async = IPCClientAsync(self._socket_path) - self._ipc_async = self._ipc_store_async self._kv_caches: dict[str, torch.Tensor] = {} self._layer_names: list[str] = [] self._layer_idx_map: dict[str, int] = {} @@ -142,7 +137,6 @@ def __init__( self._pending_finished_saves: dict[str, Any] = {} self._load_loop = asyncio.new_event_loop() self._store_loop = asyncio.new_event_loop() - self._bg_loop = self._store_loop self._load_thread = threading.Thread( target=self._run_load_loop, daemon=True, diff --git a/daser/connector/ipc_client.py b/daser/connector/ipc_client.py index ec8834b..0e514b9 100644 --- a/daser/connector/ipc_client.py +++ b/daser/connector/ipc_client.py @@ -14,7 +14,43 @@ logger = init_logger(__name__) -class IPCClientSync: +def _raise_on_error(result: dict[str, Any]) -> dict[str, Any]: + """Raise when the server returned an error frame, else return the result. + + Args: + result: decoded response frame from the server. + + Returns: + The same ``result`` dict when it carries no ``error`` field. + + Raises: + RuntimeError: if the response contains an ``error`` field. + """ + if "error" in result: + raise RuntimeError(f"[IPC] server error: {result['error']}") + return result + + +class _IPCClientBase: + """Shared connection policy for the sync and async IPC clients. + + Holds the socket path and the common retry/error contract. Concrete + subclasses provide the transport-specific ``call`` (blocking socket vs + asyncio streams); both reset their connection and retry once on transport + failure so a restarted server does not wedge the client. + + Args: + socket_path: Unix socket path of the DaseR server. + """ + + #: Number of attempts (one retry) before surfacing a transport failure. + _MAX_ATTEMPTS = 2 + + def __init__(self, socket_path: str) -> None: + self._path = socket_path + + +class IPCClientSync(_IPCClientBase): """Synchronous blocking IPC client for scheduler-side calls. Uses a persistent blocking Unix socket that is connected lazily on @@ -31,7 +67,7 @@ class IPCClientSync: """ def __init__(self, socket_path: str) -> None: - self._path = socket_path + super().__init__(socket_path) self._sock: socket.socket | None = None self._lock = threading.Lock() @@ -64,7 +100,7 @@ def call(self, payload: dict[str, Any]) -> dict[str, Any]: """ raw = pack_frame(payload) with self._lock: - for attempt in range(2): + for attempt in range(self._MAX_ATTEMPTS): if self._sock is None: self._sock = self._connect() try: @@ -73,11 +109,9 @@ def call(self, payload: dict[str, Any]) -> dict[str, Any]: break except (ConnectionError, OSError, BrokenPipeError) as exc: self._reset() - if attempt == 1: + if attempt == self._MAX_ATTEMPTS - 1: raise RuntimeError(f"[IPC] transport failure: {exc}") from exc - if "error" in result: - raise RuntimeError(f"[IPC] server error: {result['error']}") - return result + return _raise_on_error(result) def close(self) -> None: """Close the persistent socket if open.""" @@ -154,41 +188,13 @@ def record_external_prefix_cache(self, queries: int, hits: int) -> None: } ) - def match_and_alloc( - self, tokens: list[int], chunk_key: str, model_id: str - ) -> dict[str, Any]: - """Combined lookup + alloc in one RPC. - - On a cache hit the server returns the matching chunks and no - allocation; on a miss it allocates a slot for the block-aligned - prefix and returns the allocation info. Either way the scheduler - gets both possible futures in a single round trip. - - Args: - tokens: full prompt token IDs. - chunk_key: client-computed hash of the block-aligned prefix; - empty string disables miss-path allocation. - model_id: model identifier. - - Returns: - Dict with "chunks" (list[dict]) and "alloc" (dict|None). - """ - return self.call( - { - "op": "match_and_alloc", - "tokens": tokens, - "chunk_key": chunk_key, - "model_id": model_id, - } - ) - def alloc_chunk( self, chunk_key: str, token_count: int, model_id: str ) -> dict[str, Any]: """Allocate a slot for a new chunk. Args: - chunk_key: SHA256 hex of the token IDs. + chunk_key: xxh3_128 hex of the token IDs. token_count: number of tokens in the chunk. model_id: model identifier. @@ -235,7 +241,7 @@ def commit_chunk(self, chunk_key: str) -> None: """Mark a chunk as committed (GDS write complete). Args: - chunk_key: SHA256 hex of the chunk's token IDs. + chunk_key: xxh3_128 hex of the chunk's token IDs. """ self.call({"op": "commit_chunk", "chunk_key": chunk_key}) @@ -248,15 +254,6 @@ def transfer_drain(self) -> None: """ self.call({"op": "transfer_drain"}) - def init_transfer(self) -> None: - """Initialize the server-owned transfer layer. - - Thread-safety: - Uses the same lock-protected blocking RPC path as other scheduler - calls. - """ - self.call({"op": "init_transfer"}) - def commit_stats(self) -> dict[str, int]: """Return server-side connector commit counters. @@ -296,7 +293,7 @@ def evict_chunk(self, chunk_key: str) -> None: """Evict a chunk from the DaseR index. Args: - chunk_key: SHA256 hex of the chunk's token IDs. + chunk_key: xxh3_128 hex of the chunk's token IDs. """ self.call({"op": "evict_chunk", "chunk_key": chunk_key}) @@ -323,7 +320,7 @@ def release_chunk_writer( ) -class IPCClientAsync: +class IPCClientAsync(_IPCClientBase): """Asyncio IPC client for worker-side calls. Args: @@ -331,7 +328,7 @@ class IPCClientAsync: """ def __init__(self, socket_path: str) -> None: - self._path = socket_path + super().__init__(socket_path) self._reader: asyncio.StreamReader | None = None self._writer: asyncio.StreamWriter | None = None self._lock = asyncio.Lock() @@ -365,7 +362,7 @@ async def call(self, payload: dict[str, Any]) -> dict[str, Any]: """ raw = pack_frame(payload) async with self._lock: - for attempt in range(2): + for attempt in range(self._MAX_ATTEMPTS): try: reader, writer = await self._connect() writer.write(raw) @@ -374,12 +371,10 @@ async def call(self, payload: dict[str, Any]) -> dict[str, Any]: break except (ConnectionError, OSError, asyncio.IncompleteReadError) as exc: await self._reset() - if attempt == 1: + if attempt == self._MAX_ATTEMPTS - 1: raise RuntimeError(f"[IPC] transport failure: {exc}") from exc - if "error" in result: - raise RuntimeError(f"[IPC] server error: {result['error']}") - return result + return _raise_on_error(result) async def close(self) -> None: """Close the persistent async connection.""" diff --git a/daser/connector/metadata.py b/daser/connector/metadata.py index d592dbd..33eacf3 100644 --- a/daser/connector/metadata.py +++ b/daser/connector/metadata.py @@ -10,7 +10,7 @@ class ReqLoadSpec: """Load specification for one request. Attributes: - chunk_key: SHA256 of the cached token sequence. + chunk_key: xxh3_128 of the cached token sequence. start_slot: first DaseR slot for this chunk. num_slots: number of slots in the chunk. block_ids: vLLM block IDs allocated to hold the loaded KV. @@ -36,7 +36,7 @@ class ReqStoreSpec: """Store specification for one request. Attributes: - chunk_key: SHA256 of this request's token sequence. + chunk_key: xxh3_128 of this request's token sequence. start_slot: first DaseR slot allocated for this chunk. num_slots: number of slots allocated. block_ids: vLLM block IDs whose KV to save. diff --git a/daser/connector/reuse.py b/daser/connector/reuse.py index 476aa9e..2920199 100644 --- a/daser/connector/reuse.py +++ b/daser/connector/reuse.py @@ -7,6 +7,7 @@ from typing import Any # First Party +from daser.config import CACHE_REUSE_CHUNK, CACHE_REUSE_PREFIX from daser.connector.helpers import ( ROLLING_PREFIX_SEED, PendingStore, @@ -331,9 +332,9 @@ def build_cache_reuse_strategy( Raises: ValueError: if cache_reuse_mode is unknown. """ - if cache_reuse_mode == "chunk": + if cache_reuse_mode == CACHE_REUSE_CHUNK: return ChunkReuseStrategy(block_tokens) - if cache_reuse_mode == "prefix": + if cache_reuse_mode == CACHE_REUSE_PREFIX: return PrefixReuseStrategy(block_tokens) raise ValueError(f"unknown cache reuse mode: {cache_reuse_mode}") diff --git a/daser/connector/staging.py b/daser/connector/staging.py index 9690d28..cf664bd 100644 --- a/daser/connector/staging.py +++ b/daser/connector/staging.py @@ -695,33 +695,28 @@ def copy_staging_to_kv_cache( kv_tensor = kv_caches.get(layer_name) if kv_tensor is None: continue - if kv_tensor.dim() >= 2: - sample = kv_tensor[:, block_ids[0]] - src = ( - staging_by_layer[:, layer_idx, :] - .view(kv_tensor.dtype) - .view(num_slots, *sample.shape) - ) - if block_range is None: - if block_index is None: - raise RuntimeError("block_index is required for non-contiguous IDs") - kv_tensor.index_copy_(1, block_index, src.movedim(0, 1)) - else: - start, stop = block_range - kv_tensor[:, start:stop].copy_(src.movedim(0, 1)) + # KV tensors are either block-major ([blocks, ...]) or kv-major + # ([2, blocks, ...]); block_dim points at the block axis. + block_dim = 1 if kv_tensor.dim() >= 2 else 0 + sample = ( + kv_tensor[:, block_ids[0]] if block_dim == 1 else kv_tensor[block_ids[0]] + ) + src = ( + staging_by_layer[:, layer_idx, :] + .view(kv_tensor.dtype) + .view(num_slots, *sample.shape) + ) + # staging is slot-major (slots first); align it to the block axis. + src = src.movedim(0, block_dim) + if block_range is None: + if block_index is None: + raise RuntimeError("block_index is required for non-contiguous IDs") + kv_tensor.index_copy_(block_dim, block_index, src) else: - sample = kv_tensor[block_ids[0]] - src = ( - staging_by_layer[:, layer_idx, :] - .view(kv_tensor.dtype) - .view(num_slots, *sample.shape) - ) - if block_range is None: - if block_index is None: - raise RuntimeError("block_index is required for non-contiguous IDs") - kv_tensor.index_copy_(0, block_index, src) + start, stop = block_range + if block_dim == 1: + kv_tensor[:, start:stop].copy_(src) else: - start, stop = block_range kv_tensor[start:stop].copy_(src) copies += 1 return copies diff --git a/daser/connector/worker.py b/daser/connector/worker.py index a26f614..7d86270 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -91,6 +91,27 @@ class _DeferredFinishedSave: future: Any | None = None +@dataclass +class _SaveFuture: + """One background save future and the staging it keeps alive. + + Attributes: + future: future returned by ``asyncio.run_coroutine_threadsafe``. + staging_bytes: GPU staging bytes held alive until completion. + lease: optional reusable staging lease released after completion. + """ + + future: Any + staging_bytes: int + lease: CudaStagingLease | None + + def release(self) -> None: + """Release the reusable staging lease, if any.""" + if self.lease is not None: + self.lease.release() + self.lease = None + + def _cuda_allocation_base_and_offset(device_ptr: int) -> tuple[int, int]: """Return CUDA allocation base pointer and byte offset for ``device_ptr``. @@ -307,19 +328,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: is_neox_style=bool(getattr(self, "_rope_is_neox_style", True)), ) - if ( - self._ensure_transfer_ready() - and getattr(self, "_ipc_load_async", None) is not None - and getattr(self, "_ipc_store_async", None) is not None - and getattr(self, "_load_loop", None) is not None - and getattr(self, "_store_loop", None) is not None - ): - self._submit_load_coroutine(self._ipc_load_async.init_transfer()).result( - timeout=120.0 - ) - self._submit_store_coroutine(self._ipc_store_async.init_transfer()).result( - timeout=120.0 - ) + self._init_server_transfer() def register_cross_layers_kv_cache( self, @@ -345,7 +354,6 @@ def register_cross_layers_kv_cache( if not layer_names: layer_count = int(kv_cache.shape[1]) if kv_cache.dim() >= 2 else 0 layer_names = [f"layer.{idx}" for idx in range(layer_count)] - self._cross_layers_attn_backend = attn_backend self._kv_caches = {CROSS_LAYER_KV_CACHE_KEY: kv_cache} self._layer_names = layer_names self._layer_idx_map = {name: idx for idx, name in enumerate(self._layer_names)} @@ -403,19 +411,7 @@ def register_cross_layers_kv_cache( rope_base=float(getattr(self, "_rope_base", 10000.0)), is_neox_style=bool(getattr(self, "_rope_is_neox_style", True)), ) - if ( - self._ensure_transfer_ready() - and getattr(self, "_ipc_load_async", None) is not None - and getattr(self, "_ipc_store_async", None) is not None - and getattr(self, "_load_loop", None) is not None - and getattr(self, "_store_loop", None) is not None - ): - self._submit_load_coroutine(self._ipc_load_async.init_transfer()).result( - timeout=120.0 - ) - self._submit_store_coroutine(self._ipc_store_async.init_transfer()).result( - timeout=120.0 - ) + self._init_server_transfer() def bind_connector_metadata(self, connector_metadata: DaserConnectorMeta) -> None: """Receive scheduler metadata before each forward pass. @@ -427,7 +423,6 @@ def bind_connector_metadata(self, connector_metadata: DaserConnectorMeta) -> Non self._meta = connector_metadata self._reap_save_futures(block=False) self._pending_commits = set() - self._clear_save_state() for spec in connector_metadata.reqs_to_store.values(): if spec.block_ids: self._pending_commits.add(spec.chunk_key) @@ -619,7 +614,6 @@ def wait_for_save(self) -> None: save.reqs_to_store[req_id] = spec if spec.chunk_key in commit_keys: save.commit_keys.add(spec.chunk_key) - self._clear_save_state() self._pending_commits.clear() def get_finished( @@ -688,10 +682,6 @@ def _run_store_loop(self) -> None: asyncio.set_event_loop(self._store_loop) self._store_loop.run_forever() - def _run_bg_loop(self) -> None: - """Run the backward-compatible background store asyncio IO loop.""" - self._run_store_loop() - def _submit_load_coroutine(self, coro: Any) -> Any: """Submit foreground load work to the dedicated load event loop. @@ -705,7 +695,7 @@ def _submit_load_coroutine(self, coro: Any) -> Any: Called from vLLM worker threads. A load-only loop prevents cache-hit reads from queueing behind background store coroutines. """ - loop = getattr(self, "_load_loop", getattr(self, "_bg_loop", None)) + loop = self._load_loop return asyncio.run_coroutine_threadsafe(coro, loop) def _submit_store_coroutine(self, coro: Any) -> Any: @@ -721,7 +711,7 @@ def _submit_store_coroutine(self, coro: Any) -> Any: Called from vLLM worker threads. Store work is serialized on the store loop and does not occupy the foreground load loop. """ - loop = getattr(self, "_store_loop", getattr(self, "_bg_loop", None)) + loop = self._store_loop return asyncio.run_coroutine_threadsafe(coro, loop) def _ensure_transfer_ready(self) -> bool: @@ -743,9 +733,28 @@ def _ensure_transfer_ready(self) -> bool: logger.info("[CONNECTOR] server transfer mode=%s", self._transfer_mode) return True - def _clear_save_state(self) -> None: - """Clear worker-side per-forward save state.""" - return + def _init_server_transfer(self) -> None: + """Initialize the server-owned transfer layer on both IO loops. + + Async/thread-safety: + Called from a vLLM worker thread after KV-cache registration. Does + nothing until the server transfer config and both async IO loops + are ready. + """ + if not ( + self._ensure_transfer_ready() + and getattr(self, "_ipc_load_async", None) is not None + and getattr(self, "_ipc_store_async", None) is not None + and getattr(self, "_load_loop", None) is not None + and getattr(self, "_store_loop", None) is not None + ): + return + self._submit_load_coroutine(self._ipc_load_async.init_transfer()).result( + timeout=120.0 + ) + self._submit_store_coroutine(self._ipc_store_async.init_transfer()).result( + timeout=120.0 + ) def _reap_save_futures(self, block: bool) -> None: """Collect completed background save tasks. @@ -754,26 +763,17 @@ def _reap_save_futures(self, block: bool) -> None: block: If True, wait for every pending save. If False, collect only tasks that are already complete. """ - remaining = [] - pending_bytes = getattr(self, "_pending_save_staging_bytes", 0) + remaining: list[_SaveFuture] = [] + pending_bytes = self._pending_save_staging_bytes for record in self._save_futures: - if isinstance(record, tuple): - future = record[0] - staging_bytes = int(record[1]) if len(record) > 1 else 0 - staging_lease = record[2] if len(record) > 2 else None - else: - future = record - staging_bytes = 0 - staging_lease = None - if block or future.done(): + if block or record.future.done(): try: - future.result(timeout=120.0) + record.future.result(timeout=120.0) finally: - pending_bytes = max(0, pending_bytes - staging_bytes) - if staging_lease is not None: - staging_lease.release() + pending_bytes = max(0, pending_bytes - record.staging_bytes) + record.release() else: - remaining.append((future, staging_bytes, staging_lease)) + remaining.append(record) self._save_futures = remaining self._pending_save_staging_bytes = pending_bytes @@ -795,10 +795,10 @@ def _track_save_future( Called on the worker thread. Completion is collected by ``_reap_save_futures``. """ - self._pending_save_staging_bytes = ( - getattr(self, "_pending_save_staging_bytes", 0) + staging_bytes + self._pending_save_staging_bytes += staging_bytes + self._save_futures.append( + _SaveFuture(future=future, staging_bytes=staging_bytes, lease=staging_lease) ) - self._save_futures.append((future, staging_bytes, staging_lease)) def _wait_for_save_staging_capacity(self, nbytes: int) -> None: """Apply backpressure before allocating another store staging buffer. @@ -816,28 +816,16 @@ def _wait_for_save_staging_capacity(self, nbytes: int) -> None: or DEFAULT_PENDING_STORE_STAGING_BYTES, nbytes, ) - while ( - getattr(self, "_pending_save_staging_bytes", 0) + nbytes > limit - and self._save_futures - ): + while self._pending_save_staging_bytes + nbytes > limit and self._save_futures: record = self._save_futures.pop(0) - if isinstance(record, tuple): - future = record[0] - staging_bytes = int(record[1]) if len(record) > 1 else 0 - staging_lease = record[2] if len(record) > 2 else None - else: - future = record - staging_bytes = 0 - staging_lease = None try: - future.result(timeout=120.0) + record.future.result(timeout=120.0) finally: self._pending_save_staging_bytes = max( 0, - getattr(self, "_pending_save_staging_bytes", 0) - staging_bytes, + self._pending_save_staging_bytes - record.staging_bytes, ) - if staging_lease is not None: - staging_lease.release() + record.release() self._reap_save_futures(block=False) def _acquire_staging( diff --git a/daser/metrics.py b/daser/metrics.py index e1fe758..f354f34 100644 --- a/daser/metrics.py +++ b/daser/metrics.py @@ -65,6 +65,32 @@ def _format_labels(labels: LabelKey) -> str: return f"{{{body}}}" +def _render_scalar( + name: str, + description: str, + kind: str, + samples: list[tuple[LabelKey, float]], +) -> list[str]: + """Render a counter or gauge in Prometheus text format. + + Args: + name: Metric name. + description: HELP text. + kind: Prometheus metric type (``counter`` or ``gauge``). + samples: label-keyed sample values, already sorted. + + Returns: + Text lines without trailing newlines. + """ + lines = [ + f"# HELP {name} {description}", + f"# TYPE {name} {kind}", + ] + for labels, value in samples: + lines.append(f"{name}{_format_labels(labels)} {_format_float(value)}") + return lines + + @dataclass class Counter: """Monotonic Prometheus counter. @@ -103,13 +129,7 @@ def render(self) -> list[str]: """ with self._lock: samples = sorted(self._values.items()) - lines = [ - f"# HELP {self.name} {self.description}", - f"# TYPE {self.name} counter", - ] - for labels, value in samples: - lines.append(f"{self.name}{_format_labels(labels)} {_format_float(value)}") - return lines + return _render_scalar(self.name, self.description, "counter", samples) @dataclass @@ -151,15 +171,6 @@ def inc(self, amount: float = 1.0, labels: Labels | None = None) -> None: with self._lock: self._values[key] = self._values.get(key, 0.0) + amount - def dec(self, amount: float = 1.0, labels: Labels | None = None) -> None: - """Decrease a labeled gauge series. - - Args: - amount: Decrement amount. - labels: Optional low-cardinality labels. - """ - self.inc(-amount, labels=labels) - def render(self) -> list[str]: """Render this gauge in Prometheus text format. @@ -168,13 +179,7 @@ def render(self) -> list[str]: """ with self._lock: samples = sorted(self._values.items()) - lines = [ - f"# HELP {self.name} {self.description}", - f"# TYPE {self.name} gauge", - ] - for labels, value in samples: - lines.append(f"{self.name}{_format_labels(labels)} {_format_float(value)}") - return lines + return _render_scalar(self.name, self.description, "gauge", samples) @dataclass diff --git a/daser/ops/rope_apply.py b/daser/ops/rope_apply.py index 2d33fe8..731cd79 100644 --- a/daser/ops/rope_apply.py +++ b/daser/ops/rope_apply.py @@ -19,12 +19,11 @@ None, ] TileLangCacheKey = tuple[Any, int, int, bool] -TileLangKVCacheKey = tuple[Any, tuple[int, ...], int, bool] -TileLangRestoreCacheKey = tuple[Any, tuple[int, ...], int, bool] +TileLangShapeCacheKey = tuple[Any, tuple[int, ...], int, bool] _kernel_cache: dict[TileLangCacheKey, TileLangFn] = {} -_kv_table_kernel_cache: dict[TileLangKVCacheKey, TileLangTableFn] = {} -_restore_table_kernel_cache: dict[TileLangRestoreCacheKey, TileLangRestoreTableFn] = {} +_kv_table_kernel_cache: dict[TileLangShapeCacheKey, TileLangTableFn] = {} +_restore_table_kernel_cache: dict[TileLangShapeCacheKey, TileLangRestoreTableFn] = {} def clear_rope_apply_cache() -> None: @@ -62,6 +61,10 @@ def apply_rope_delta_to_key_block( Returns: None. ``key_block`` is modified in place. + Raises: + ValueError: if ``rotary_dim`` exceeds the head dimension, or the block + is not a contiguous CUDA tensor with the expected shape. + Async/thread-safety: Launches a CUDA kernel on the current stream. The kernel cache is process-wide and safe for normal single worker-thread use. @@ -69,7 +72,7 @@ def apply_rope_delta_to_key_block( if delta == 0 or rotary_dim <= 0: return if key_block.shape[-1] < rotary_dim: - return + raise ValueError("rotary_dim must not exceed head_dim") if key_block.device.type != "cuda": raise ValueError("TileLang RoPE apply requires a CUDA tensor") if key_block.dim() < 3: @@ -111,6 +114,9 @@ def apply_rope_delta_to_kv_key_block( Returns: None. Only the key slice ``kv_block[:, :, 0]`` is modified in place. + Raises: + ValueError: if ``rotary_dim`` exceeds the head dimension. + Async/thread-safety: Launches CUDA work on the current PyTorch stream. The cosine/sine tables are built once per call and passed to the TileLang table kernel. @@ -118,7 +124,7 @@ def apply_rope_delta_to_kv_key_block( if delta == 0 or rotary_dim <= 0: return if kv_block.shape[-1] < rotary_dim: - return + raise ValueError("rotary_dim must not exceed head_dim") cos_table, sin_table = build_rope_delta_tables( kv_block.device, @@ -159,30 +165,13 @@ def apply_rope_delta_to_kv_key_block_table( Launches one CUDA kernel on the current stream and uses a process-wide kernel cache for normal single worker-thread use. """ - if kv_block.device.type != "cuda": - raise ValueError("TileLang table RoPE backend requires a CUDA tensor") - if not kv_block.is_contiguous(): - raise ValueError("TileLang table RoPE backend requires contiguous staging") - if kv_block.dim() != 6 or kv_block.shape[2] != 2: - raise ValueError("TileLang table RoPE backend requires [blocks,layers,2,...]") - if rotary_dim <= 0 or rotary_dim % 2 != 0: - raise ValueError( - "TileLang table RoPE backend requires a positive even rotary_dim" - ) - if kv_block.shape[-1] < rotary_dim: - raise ValueError("rotary_dim must not exceed head_dim") - expected = rotary_dim // 2 - if cos_table.shape != (expected,) or sin_table.shape != (expected,): - raise ValueError("RoPE tables must have shape [rotary_dim // 2]") - if ( - cos_table.device != kv_block.device - or sin_table.device != kv_block.device - or cos_table.dtype != torch.float32 - or sin_table.dtype != torch.float32 - or not cos_table.is_contiguous() - or not sin_table.is_contiguous() - ): - raise ValueError("RoPE tables must be contiguous fp32 tensors on the KV device") + _validate_kv_and_tables( + kv_block, + cos_table, + sin_table, + rotary_dim, + "TileLang table RoPE backend", + ) kernel = _get_tilelang_kv_table_kernel( kv_block.dtype, @@ -219,34 +208,21 @@ def restore_cross_layer_kv_cache_table( Launches one CUDA kernel on the current stream and uses a process-wide kernel cache for normal single worker-thread use. """ - if staging_kv.device.type != "cuda" or dst_kv.device.type != "cuda": + if dst_kv.device.type != "cuda": raise ValueError("TileLang table fused restore requires CUDA tensors") - if not staging_kv.is_contiguous() or not dst_kv.is_contiguous(): + if not dst_kv.is_contiguous(): raise ValueError("TileLang table fused restore requires contiguous tensors") if staging_kv.shape != dst_kv.shape: raise ValueError("source and destination KV shapes must match") if staging_kv.dtype != dst_kv.dtype: raise ValueError("source and destination KV dtypes must match") - if staging_kv.dim() != 6 or staging_kv.shape[2] != 2: - raise ValueError("TileLang table fused restore requires [blocks,layers,2,...]") - if rotary_dim <= 0 or rotary_dim % 2 != 0: - raise ValueError( - "TileLang table fused restore requires a positive even rotary_dim" - ) - if staging_kv.shape[-1] < rotary_dim: - raise ValueError("rotary_dim must not exceed head_dim") - expected = rotary_dim // 2 - if cos_table.shape != (expected,) or sin_table.shape != (expected,): - raise ValueError("RoPE tables must have shape [rotary_dim // 2]") - if ( - cos_table.device != staging_kv.device - or sin_table.device != staging_kv.device - or cos_table.dtype != torch.float32 - or sin_table.dtype != torch.float32 - or not cos_table.is_contiguous() - or not sin_table.is_contiguous() - ): - raise ValueError("RoPE tables must be contiguous fp32 tensors on the KV device") + _validate_kv_and_tables( + staging_kv, + cos_table, + sin_table, + rotary_dim, + "TileLang table fused restore", + ) kernel = _get_tilelang_restore_table_kernel( staging_kv.dtype, @@ -290,6 +266,75 @@ def build_rope_delta_tables( return freqs.cos().contiguous(), freqs.sin().contiguous() +def _compile_cached(cache: dict[Any, Any], key: Any, build: Callable[[], Any]) -> Any: + """Return a cached compiled TileLang kernel, compiling on first use. + + Args: + cache: per-kernel cache dict keyed by layout. + key: cache key identifying the kernel layout. + build: zero-arg callable returning the TileLang prim_func to compile. + + Returns: + The compiled, cached TileLang kernel. + + Async/thread-safety: + The cache is process-wide and safe for normal single worker-thread use. + """ + cached = cache.get(key) + if cached is not None: + return cached + + import tilelang + + kernel = tilelang.compile(build(), target="cuda", execution_backend="cython") + cache[key] = kernel + return kernel + + +def _validate_kv_and_tables( + kv: torch.Tensor, + cos_table: torch.Tensor, + sin_table: torch.Tensor, + rotary_dim: int, + label: str, +) -> None: + """Validate a 6-D KV tensor and its RoPE cosine/sine tables. + + Args: + kv: KV tensor expected to be a contiguous CUDA tensor with shape + ``[blocks, layers, 2, block_tokens, heads, head_dim]``. + cos_table: fp32 cosine table with ``rotary_dim // 2`` values. + sin_table: fp32 sine table with ``rotary_dim // 2`` values. + rotary_dim: number of head dimensions covered by RoPE. + label: backend label used in error messages. + + Raises: + ValueError: if any tensor fails the device/shape/dtype contract. + """ + if kv.device.type != "cuda": + raise ValueError(f"{label} requires a CUDA tensor") + if not kv.is_contiguous(): + raise ValueError(f"{label} requires contiguous staging") + if kv.dim() != 6 or kv.shape[2] != 2: + raise ValueError(f"{label} requires [blocks,layers,2,...]") + if rotary_dim <= 0 or rotary_dim % 2 != 0: + raise ValueError(f"{label} requires a positive even rotary_dim") + if kv.shape[-1] < rotary_dim: + raise ValueError("rotary_dim must not exceed head_dim") + expected = rotary_dim // 2 + if cos_table.shape != (expected,) or sin_table.shape != (expected,): + raise ValueError("RoPE tables must have shape [rotary_dim // 2]") + if ( + cos_table.device != kv.device + or sin_table.device != kv.device + or cos_table.dtype != torch.float32 + or sin_table.dtype != torch.float32 + or not cos_table.is_contiguous() + or not sin_table.is_contiguous() + ): + raise ValueError("RoPE tables must be contiguous fp32 tensors on the KV device") + + def _get_tilelang_kernel( dtype: Any, head_dim: int, @@ -297,25 +342,16 @@ def _get_tilelang_kernel( is_neox_style: bool, ) -> TileLangFn: """Return a cached dynamic-shape TileLang kernel for the requested layout.""" - key = (dtype, head_dim, rotary_dim, is_neox_style) - cached = _kernel_cache.get(key) - if cached is not None: - return cached - - import tilelang - - kernel = tilelang.compile( - _build_tilelang_kernel( + return _compile_cached( + _kernel_cache, + (dtype, head_dim, rotary_dim, is_neox_style), + lambda: _build_tilelang_kernel( head_dim=head_dim, rotary_dim=rotary_dim, is_neox_style=is_neox_style, dtype=_tilelang_dtype(dtype), ), - target="cuda", - execution_backend="cython", ) - _kernel_cache[key] = kernel - return kernel def _get_tilelang_restore_table_kernel( @@ -325,25 +361,16 @@ def _get_tilelang_restore_table_kernel( is_neox_style: bool, ) -> TileLangRestoreTableFn: """Return a cached fused restore kernel using precomputed RoPE tables.""" - key = (dtype, shape[1:], rotary_dim, is_neox_style) - cached = _restore_table_kernel_cache.get(key) - if cached is not None: - return cached - - import tilelang - - kernel = tilelang.compile( - _build_tilelang_restore_table_kernel( + return _compile_cached( + _restore_table_kernel_cache, + (dtype, shape[1:], rotary_dim, is_neox_style), + lambda: _build_tilelang_restore_table_kernel( static_shape=shape[1:], rotary_dim=rotary_dim, is_neox_style=is_neox_style, dtype=_tilelang_dtype(dtype), ), - target="cuda", - execution_backend="cython", ) - _restore_table_kernel_cache[key] = kernel - return kernel def _get_tilelang_kv_table_kernel( @@ -353,25 +380,16 @@ def _get_tilelang_kv_table_kernel( is_neox_style: bool, ) -> TileLangTableFn: """Return a cached TileLang kernel using precomputed RoPE tables.""" - key = (dtype, shape[1:], rotary_dim, is_neox_style) - cached = _kv_table_kernel_cache.get(key) - if cached is not None: - return cached - - import tilelang - - kernel = tilelang.compile( - _build_tilelang_kv_table_kernel( + return _compile_cached( + _kv_table_kernel_cache, + (dtype, shape[1:], rotary_dim, is_neox_style), + lambda: _build_tilelang_kv_table_kernel( static_shape=shape[1:], rotary_dim=rotary_dim, is_neox_style=is_neox_style, dtype=_tilelang_dtype(dtype), ), - target="cuda", - execution_backend="cython", ) - _kv_table_kernel_cache[key] = kernel - return kernel def _tilelang_dtype(dtype: Any) -> str: diff --git a/daser/position/base.py b/daser/position/base.py index 08e48f9..6143f28 100644 --- a/daser/position/base.py +++ b/daser/position/base.py @@ -13,9 +13,11 @@ class PositionEncoder(ABC): Manages position offsets so that KV computed at one position can be reused when the chunk is loaded at a (potentially different) position. - The first implementation is FixedOffsetEncoder, which stores the - position offset at insert time and returns it unchanged at load time. - Future implementations may support dynamic position remapping. + FixedOffsetEncoder stores the position offset at insert time and returns + it unchanged at load time (rolling-prefix reuse). ChunkPositionEncoder + returns a target-aware offset so block-aligned chunks can be reused at the + prompt offset where they are matched. Future implementations may support + dynamic position remapping. """ @abstractmethod diff --git a/daser/retrieval/base.py b/daser/retrieval/base.py index 3ada5b2..9da6a61 100644 --- a/daser/retrieval/base.py +++ b/daser/retrieval/base.py @@ -27,10 +27,19 @@ class RetrievalIndex(ABC): Implementations map token sequences to stored ChunkMeta objects. PrefixHashIndex uses chained rolling-prefix keys for slot-granular exact - prefix reuse. - Future implementations may use vector similarity or hybrid strategies. + prefix reuse; ChunkReuseIndex matches block-aligned document chunks at + arbitrary prompt offsets. Future implementations may use vector similarity + or hybrid strategies. + + The base provides concrete ``insert``/``remove`` operating on a + subclass-owned ``_index`` dict keyed by ``chunk_key``. Subclasses keep + ``lookup`` abstract, initialize ``self._index`` in ``__init__``, and may + maintain secondary structures through the ``_on_insert``/``_on_remove`` + hooks. """ + _index: dict[str, ChunkMeta] + @abstractmethod async def lookup(self, tokens: list[int], model_id: str) -> list[RetrievalMatch]: """Find cached chunks matching the given token sequence. @@ -45,20 +54,41 @@ async def lookup(self, tokens: list[int], model_id: str) -> list[RetrievalMatch] """ ... - @abstractmethod async def insert(self, meta: ChunkMeta) -> None: """Add a committed chunk to the retrieval index. Args: meta: ChunkMeta describing the stored chunk. + + Async/thread-safety: + Records ``meta`` in the primary index and notifies ``_on_insert``. """ - ... + self._index[meta.chunk_key] = meta + self._on_insert(meta) - @abstractmethod async def remove(self, chunk_key: str) -> None: """Remove an evicted chunk from the retrieval index. Args: - chunk_key: key of the chunk to remove. + chunk_key: key of the chunk to remove; ignored when absent. + + Async/thread-safety: + Drops the chunk from the primary index and notifies ``_on_remove``. + """ + meta = self._index.pop(chunk_key, None) + self._on_remove(chunk_key, meta) + + def _on_insert(self, meta: ChunkMeta) -> None: # noqa: B027 + """Update subclass secondary structures after a primary insert. + + Args: + meta: chunk metadata just inserted into the primary index. + """ + + def _on_remove(self, chunk_key: str, meta: ChunkMeta | None) -> None: # noqa: B027 + """Update subclass secondary structures after a primary removal. + + Args: + chunk_key: key removed from the primary index. + meta: removed chunk metadata, or None when the key was absent. """ - ... diff --git a/daser/retrieval/chunk_reuse.py b/daser/retrieval/chunk_reuse.py index 82f4bd8..ca5d6c6 100644 --- a/daser/retrieval/chunk_reuse.py +++ b/daser/retrieval/chunk_reuse.py @@ -61,25 +61,24 @@ async def lookup(self, tokens: list[int], model_id: str) -> list[RetrievalMatch] start += matched_tokens if matched_tokens else self._block_tokens return matches - async def insert(self, meta: ChunkMeta) -> None: - """Insert a committed chunk into the chunk reuse index. + def _on_insert(self, meta: ChunkMeta) -> None: + """Index a committed chunk by its token count after a primary insert. Args: meta: committed chunk metadata. """ - self._index[meta.chunk_key] = meta bucket = self._by_token_count.setdefault(meta.token_count, {}) bucket[meta.chunk_key] = meta self._refresh_token_counts() logger.debug("[INDEX] chunk insert key=%s", meta.chunk_key[:8]) - async def remove(self, chunk_key: str) -> None: - """Remove an evicted chunk from the chunk reuse index. + def _on_remove(self, chunk_key: str, meta: ChunkMeta | None) -> None: + """Drop an evicted chunk from the token-count index. Args: - chunk_key: key to remove; ignored when absent. + chunk_key: key removed from the primary index. + meta: removed chunk metadata, or None when the key was absent. """ - meta = self._index.pop(chunk_key, None) if meta is not None: bucket = self._by_token_count.get(meta.token_count) if bucket is not None: diff --git a/daser/retrieval/prefix.py b/daser/retrieval/prefix.py index 38aea15..2b6d350 100644 --- a/daser/retrieval/prefix.py +++ b/daser/retrieval/prefix.py @@ -94,21 +94,3 @@ def flush_run() -> None: run_target_start = target_token_start flush_run() return matches - - async def insert(self, meta: ChunkMeta) -> None: - """Insert a chunk into the prefix index. - - Args: - meta: ChunkMeta keyed by a rolling-prefix slot key. - """ - self._index[meta.chunk_key] = meta - logger.debug("[INDEX] insert chunk_key=%s", meta.chunk_key[:8]) - - async def remove(self, chunk_key: str) -> None: - """Remove a chunk from the prefix index. - - Args: - chunk_key: key to remove; silently ignored if not present. - """ - self._index.pop(chunk_key, None) - logger.debug("[INDEX] remove chunk_key=%s", chunk_key[:8]) diff --git a/daser/server/__main__.py b/daser/server/__main__.py index 1da28f1..82e3fb2 100644 --- a/daser/server/__main__.py +++ b/daser/server/__main__.py @@ -12,7 +12,15 @@ import uvicorn # First Party -from daser.config import BLOCK_TOKENS, DEFAULT_IOURING_L1_BYTES, DaserConfig +from daser.config import ( + BLOCK_TOKENS, + CACHE_REUSE_CHUNK, + CACHE_REUSE_MODES, + CACHE_REUSE_PREFIX, + DEFAULT_CACHE_REUSE_MODE, + DEFAULT_IOURING_L1_BYTES, + DaserConfig, +) from daser.logging import init_logger from daser.position.base import PositionEncoder from daser.position.chunk_position import ChunkPositionEncoder @@ -180,8 +188,8 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--log-level", default="INFO") parser.add_argument( "--cache-reuse-mode", - choices=("prefix", "chunk"), - default="chunk", + choices=CACHE_REUSE_MODES, + default=DEFAULT_CACHE_REUSE_MODE, help="Cache reuse strategy: chunk enables block-aligned chunk reuse " "inside RAG prompts; prefix enables rolling-prefix slot reuse.", ) @@ -346,7 +354,7 @@ def _build_http_config(args: argparse.Namespace) -> HTTPServerConfig: tokenizer=args.model_path, block_tokens=int(args.block_tokens), cache_reuse_mode=args.cache_reuse_mode, - align_document_chunks=args.cache_reuse_mode == "chunk", + align_document_chunks=args.cache_reuse_mode == CACHE_REUSE_CHUNK, transfer_mode=args.transfer_mode, ) @@ -366,11 +374,11 @@ def _build_index_components( Raises: ValueError: if cache_reuse_mode is unknown. """ - if cache_reuse_mode == "prefix": + if cache_reuse_mode == CACHE_REUSE_PREFIX: return PrefixHashIndex(block_tokens=block_tokens), FixedOffsetEncoder( fixed_offset=0 ) - if cache_reuse_mode == "chunk": + if cache_reuse_mode == CACHE_REUSE_CHUNK: return ChunkReuseIndex(block_tokens=block_tokens), ChunkPositionEncoder( initial_offset=0 ) diff --git a/daser/server/chunk_lifecycle.py b/daser/server/chunk_lifecycle.py new file mode 100644 index 0000000..9d905fc --- /dev/null +++ b/daser/server/chunk_lifecycle.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Chunk commit/eviction/ownership state for the DaseR control plane.""" + +# Standard +import asyncio + + +class ChunkLifecycle: + """Track commit, write-ownership, eviction, and commit-waiter state. + + ServerCore keeps several parallel sets in lockstep: which chunk keys are + committed and visible to lookup, which have an active store writer, and + which were evicted. This class owns those sets and the commit-waiter + futures so the transitions stay consistent in one place. + + Async/thread-safety: + All methods run on the server asyncio event loop. Commit waiters are + completed from the same loop by ``mark_committed``. + """ + + def __init__(self) -> None: + self._committed: set[str] = set() + self._write_owners: set[str] = set() + self._evicted: set[str] = set() + self._commit_waiters: dict[str, set[asyncio.Future[None]]] = {} + + def is_committed(self, chunk_key: str) -> bool: + """Return whether ``chunk_key`` is committed and visible to lookup.""" + return chunk_key in self._committed + + def is_evicted(self, chunk_key: str) -> bool: + """Return whether ``chunk_key`` was evicted.""" + return chunk_key in self._evicted + + @property + def write_owners(self) -> set[str]: + """Return the live write-owner key set (used for reuse predicates).""" + return self._write_owners + + @property + def committed(self) -> set[str]: + """Return the committed key set (used for reuse predicates).""" + return self._committed + + def mark_write_owner(self, chunk_key: str) -> None: + """Record that a store writer claimed ``chunk_key``.""" + self._write_owners.add(chunk_key) + + def mark_committed(self, chunk_key: str) -> None: + """Mark ``chunk_key`` committed, claim ownership, and wake waiters.""" + self._committed.add(chunk_key) + self._write_owners.add(chunk_key) + self._notify_commit_waiters(chunk_key) + + def mark_evicted(self, chunk_key: str) -> None: + """Drop committed/owner state for ``chunk_key`` and record eviction.""" + self._committed.discard(chunk_key) + self._write_owners.discard(chunk_key) + self._evicted.add(chunk_key) + + def discard_owner(self, chunk_key: str) -> None: + """Release a store-writer claim without committing or evicting.""" + self._write_owners.discard(chunk_key) + + def discard(self, chunk_key: str) -> None: + """Drop committed and write-owner state without recording eviction.""" + self._committed.discard(chunk_key) + self._write_owners.discard(chunk_key) + + async def wait_for_committed( + self, + chunk_keys: list[str], + timeout_s: float, + ) -> None: + """Wait until every key in ``chunk_keys`` is committed. + + Args: + chunk_keys: keys to await; duplicates and already-committed keys + are ignored. + timeout_s: maximum seconds to wait. + + Raises: + TimeoutError: if any key is still uncommitted at timeout. + """ + pending = [ + key for key in dict.fromkeys(chunk_keys) if key not in self._committed + ] + if not pending: + return + + loop = asyncio.get_running_loop() + futures: dict[str, asyncio.Future[None]] = {} + for key in pending: + future: asyncio.Future[None] = loop.create_future() + self._commit_waiters.setdefault(key, set()).add(future) + futures[key] = future + try: + await asyncio.wait_for( + asyncio.gather(*futures.values()), + timeout=timeout_s, + ) + except asyncio.TimeoutError as exc: + raise TimeoutError("timed out waiting for committed chunks") from exc + finally: + for key, future in futures.items(): + waiters = self._commit_waiters.get(key) + if waiters is None: + continue + waiters.discard(future) + if not waiters: + self._commit_waiters.pop(key, None) + + def _notify_commit_waiters(self, chunk_key: str) -> None: + """Complete and clear any commit waiters for ``chunk_key``.""" + waiters = self._commit_waiters.pop(chunk_key, set()) + for future in waiters: + if not future.done(): + future.set_result(None) diff --git a/daser/server/core.py b/daser/server/core.py index c699847..c3c3096 100644 --- a/daser/server/core.py +++ b/daser/server/core.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Standard -import asyncio from dataclasses import dataclass import math from typing import Any, Optional @@ -11,6 +10,7 @@ from daser.metrics import REGISTRY, MetricsRegistry from daser.position.base import PositionEncoder from daser.retrieval.base import RetrievalIndex, RetrievalMatch +from daser.server.chunk_lifecycle import ChunkLifecycle from daser.server.chunk_manager import ChunkManager from daser.server.doc_registry import DocEntry, DocRegistry from daser.server.metadata_store import ChunkMeta @@ -226,14 +226,11 @@ def __init__( self._slot_size = slot_size self._block_tokens = block_tokens self._metrics = metrics_registry or REGISTRY - self._evicted_chunk_keys: set[str] = set() + self._lifecycle = ChunkLifecycle() self._commit_requests = 0 self._late_evicted_commits = 0 self._lookup_requests = 0 self._lookup_hits = 0 - self._committed_chunk_keys: set[str] = set() - self._write_owner_chunk_keys: set[str] = set() - self._commit_waiters: dict[str, set[asyncio.Future[None]]] = {} self._record_capacity_metrics() @property @@ -249,8 +246,7 @@ async def rebuild_retrieval_index(self) -> None: """ for meta in list(self._cm.store.iter_chunks()): await self._ri.insert(meta) - self._committed_chunk_keys.add(meta.chunk_key) - self._write_owner_chunk_keys.add(meta.chunk_key) + self._lifecycle.mark_committed(meta.chunk_key) async def lookup(self, tokens: list[int], model_id: str) -> list[ChunkInfo]: """Look up cached chunks for token IDs. @@ -335,22 +331,11 @@ async def alloc_chunk( Async/thread-safety: Performs in-memory mutation on the server event loop. """ - num_slots = math.ceil(token_count / self._block_tokens) - should_write = not self._has_store_owner(chunk_key, token_count, model_id) - meta = await self._alloc_or_get_chunk( - chunk_key=chunk_key, - token_count=token_count, - num_slots=num_slots, - model_id=model_id, - ) - if should_write: - self._write_owner_chunk_keys.add(chunk_key) - return self._allocation( - meta, - token_count=token_count, - num_slots=num_slots, - skipped=not should_write, + allocations = await self.alloc_chunks( + [{"chunk_key": chunk_key, "token_count": token_count}], + model_id, ) + return allocations[0] async def alloc_chunks( self, @@ -373,7 +358,7 @@ async def alloc_chunks( for chunk in chunks: chunk_key = str(chunk["chunk_key"]) token_count = int(chunk["token_count"]) - num_slots = math.ceil(token_count / self._block_tokens) + num_slots = self._slots_for(token_count) should_write = not self._has_store_owner(chunk_key, token_count, model_id) meta = await self._alloc_or_get_chunk( chunk_key=chunk_key, @@ -382,7 +367,7 @@ async def alloc_chunks( model_id=model_id, ) if should_write: - self._write_owner_chunk_keys.add(chunk_key) + self._lifecycle.mark_write_owner(chunk_key) allocations.append( self._allocation( meta, @@ -417,7 +402,7 @@ async def match_and_alloc( aligned = (len(tokens) // self._block_tokens) * self._block_tokens if aligned == 0: return MatchAndAllocResult(chunks=[], alloc=None) - num_slots = math.ceil(aligned / self._block_tokens) + num_slots = self._slots_for(aligned) should_write = not self._has_store_owner(chunk_key, aligned, model_id) meta = await self._alloc_or_get_chunk( chunk_key=chunk_key, @@ -426,7 +411,7 @@ async def match_and_alloc( model_id=model_id, ) if should_write: - self._write_owner_chunk_keys.add(chunk_key) + self._lifecycle.mark_write_owner(chunk_key) return MatchAndAllocResult( chunks=[], alloc=self._allocation( @@ -451,7 +436,7 @@ async def commit_chunk(self, chunk_key: str) -> None: """ meta = self._cm.store.get(chunk_key) if meta is None: - if chunk_key in self._evicted_chunk_keys: + if self._lifecycle.is_evicted(chunk_key): self._commit_requests += 1 self._late_evicted_commits += 1 self._metrics.counter( @@ -465,9 +450,7 @@ async def commit_chunk(self, chunk_key: str) -> None: return raise ValueError(f"chunk_key not found: {chunk_key}") await self._ri.insert(meta) - self._committed_chunk_keys.add(chunk_key) - self._write_owner_chunk_keys.add(chunk_key) - self._notify_commit_waiters(chunk_key) + self._lifecycle.mark_committed(chunk_key) self._commit_requests += 1 self._metrics.counter( "daser_cache_committed_chunks_total", @@ -490,7 +473,7 @@ def is_chunk_committed(self, chunk_key: str) -> bool: Reads in-memory state on the server event loop. It performs no blocking I/O. """ - return chunk_key in self._committed_chunk_keys + return self._lifecycle.is_committed(chunk_key) def is_chunk_reusable( self, @@ -514,13 +497,8 @@ def is_chunk_reusable( Reads in-memory state on the server event loop. It performs no blocking I/O. """ - meta = self._cm.store.get(chunk_key) - if meta is None or chunk_key not in self._committed_chunk_keys: - return False - return ( - meta.token_count == token_count - and meta.num_slots == math.ceil(token_count / self._block_tokens) - and meta.model_id == model_id + return self._meta_matches( + chunk_key, token_count, model_id, self._lifecycle.committed ) async def wait_for_committed_chunks( @@ -541,35 +519,7 @@ async def wait_for_committed_chunks( Must run on the server event loop. Waiters are completed by ``commit_chunk`` on the same event loop. """ - pending = [ - key - for key in dict.fromkeys(chunk_keys) - if key not in self._committed_chunk_keys - ] - if not pending: - return - - loop = asyncio.get_running_loop() - futures: dict[str, asyncio.Future[None]] = {} - for key in pending: - future: asyncio.Future[None] = loop.create_future() - self._commit_waiters.setdefault(key, set()).add(future) - futures[key] = future - try: - await asyncio.wait_for( - asyncio.gather(*futures.values()), - timeout=timeout_s, - ) - except asyncio.TimeoutError as exc: - raise TimeoutError("timed out waiting for committed chunks") from exc - finally: - for key, future in futures.items(): - waiters = self._commit_waiters.get(key) - if waiters is None: - continue - waiters.discard(future) - if not waiters: - self._commit_waiters.pop(key, None) + await self._lifecycle.wait_for_committed(chunk_keys, timeout_s) async def commit_stats(self) -> dict[str, int]: """Return connector commit counters for benchmark synchronization. @@ -661,11 +611,11 @@ async def release_chunk_writer( Async/thread-safety: Performs in-memory mutation on the server event loop. """ - if chunk_key in self._committed_chunk_keys: + if self._lifecycle.is_committed(chunk_key): return False if not self.is_current_allocation(chunk_key, start_slot, num_slots): return False - self._write_owner_chunk_keys.discard(chunk_key) + self._lifecycle.discard_owner(chunk_key) return True async def evict_chunk(self, chunk_key: str) -> None: @@ -682,9 +632,7 @@ async def evict_chunk(self, chunk_key: str) -> None: if meta is not None: self._mark_chunk_evicted_in_docs(meta) self._cm.store.remove(chunk_key) - self._committed_chunk_keys.discard(chunk_key) - self._write_owner_chunk_keys.discard(chunk_key) - self._evicted_chunk_keys.add(chunk_key) + self._lifecycle.mark_evicted(chunk_key) self._metrics.counter( "daser_cache_evicted_chunks_total", "Chunks evicted from cache metadata.", @@ -840,9 +788,7 @@ async def _drain_ring_evictions(self) -> None: """ for chunk_key in self._cm.drain_evicted_chunk_keys(): await self._ri.remove(chunk_key) - self._committed_chunk_keys.discard(chunk_key) - self._write_owner_chunk_keys.discard(chunk_key) - self._evicted_chunk_keys.add(chunk_key) + self._lifecycle.mark_evicted(chunk_key) self._metrics.counter( "daser_cache_evicted_chunks_total", "Chunks evicted from cache metadata.", @@ -863,13 +809,6 @@ def _record_capacity_metrics(self) -> None: "Currently used L2 store bytes.", ).set(used_slots * self._slot_size) - def _notify_commit_waiters(self, chunk_key: str) -> None: - """Wake coroutines waiting for a chunk commit.""" - waiters = self._commit_waiters.pop(chunk_key, set()) - for future in waiters: - if not future.done(): - future.set_result(None) - def _require_doc_registry(self) -> DocRegistry: """Return the attached DocRegistry or raise a public operation error.""" registry = self._cm.doc_registry @@ -901,16 +840,14 @@ def _detach_doc_from_chunk(self, doc_id: str, chunk_key: str) -> bool: """ meta = self._cm.store.get(chunk_key) if meta is None: - self._committed_chunk_keys.discard(chunk_key) - self._write_owner_chunk_keys.discard(chunk_key) + self._lifecycle.discard(chunk_key) return False if doc_id in meta.doc_ids: meta.doc_ids.remove(doc_id) if meta.doc_ids: return False self._cm.store.remove(chunk_key) - self._committed_chunk_keys.discard(chunk_key) - self._write_owner_chunk_keys.discard(chunk_key) + self._lifecycle.discard(chunk_key) return True def _has_store_owner( @@ -930,12 +867,39 @@ def _has_store_owner( True for committed chunks and for uncommitted allocations whose first writer has already claimed the store target. """ + return self._meta_matches( + chunk_key, token_count, model_id, self._lifecycle.write_owners + ) + + def _slots_for(self, token_count: int) -> int: + """Return the number of KV slots needed for ``token_count`` tokens.""" + return math.ceil(token_count / self._block_tokens) + + def _meta_matches( + self, + chunk_key: str, + token_count: int, + model_id: str, + membership: set[str], + ) -> bool: + """Return whether stored metadata matches and is in ``membership``. + + Args: + chunk_key: cache key to inspect. + token_count: expected token count. + model_id: expected model identifier. + membership: set the key must belong to (committed or write-owner). + + Returns: + True when a metadata entry exists for the key, is in ``membership``, + and matches the token count, slot count, and model. + """ meta = self._cm.store.get(chunk_key) - if meta is None or chunk_key not in self._write_owner_chunk_keys: + if meta is None or chunk_key not in membership: return False return ( meta.token_count == token_count - and meta.num_slots == math.ceil(token_count / self._block_tokens) + and meta.num_slots == self._slots_for(token_count) and meta.model_id == model_id ) diff --git a/daser/server/http/app.py b/daser/server/http/app.py index 7db2c4f..45a41cd 100644 --- a/daser/server/http/app.py +++ b/daser/server/http/app.py @@ -15,6 +15,10 @@ from pydantic import BaseModel, Field # First Party +from daser.config import ( + CACHE_REUSE_PREFIX, + DEFAULT_CACHE_REUSE_MODE, +) from daser.logging import init_logger from daser.metrics import REGISTRY, MetricsRegistry from daser.server.core import ServerCore @@ -38,8 +42,6 @@ class HTTPServerConfig: block_tokens: vLLM block size. system_prompt: fixed prefix before document prompts. doc_separator: separator inserted between documents. - task_separator: separator inserted before task text. - answer_separator: separator inserted after task text before generation. cache_reuse_mode: cache reuse strategy selected by the DaseR server. align_document_chunks: when True, insert padding tokens before each document so document chunks begin on vLLM block boundaries. @@ -54,9 +56,7 @@ class HTTPServerConfig: "the following documents.\n\n" ) doc_separator: str = "\n\n---\n\n" - task_separator: str = "\n\n---\nTask: " - answer_separator: str = "\nAnswer: " - cache_reuse_mode: str = "chunk" + cache_reuse_mode: str = DEFAULT_CACHE_REUSE_MODE align_document_chunks: bool = False transfer_mode: str = "iouring" @@ -630,7 +630,7 @@ async def drain_endpoint() -> dict[str, bool]: @app.post("/documents", status_code=201) async def upload_document(req: UploadRequest) -> dict[str, Any]: """Upload a document, prefill chunk KV, and register it.""" - if cfg.cache_reuse_mode == "prefix": + if cfg.cache_reuse_mode == CACHE_REUSE_PREFIX: raise HTTPException( status_code=400, detail="document upload requires chunk cache reuse mode", diff --git a/daser/server/http/chunker.py b/daser/server/http/chunker.py index 13fc114..adc817c 100644 --- a/daser/server/http/chunker.py +++ b/daser/server/http/chunker.py @@ -63,18 +63,6 @@ def block_tokens(self) -> int: """vLLM block size in tokens.""" return self._block_tokens - def pad_to_chunk_boundary(self, tokens: list[int], pad_token: int) -> list[int]: - """Return a copy of ``tokens`` padded to a vLLM block boundary. - - Args: - tokens: tokenized input. - pad_token: token ID used to fill the final partial block. - - Returns: - Padded token list. Empty input remains empty. - """ - return self.pad_to_block_boundary(tokens, pad_token) - def pad_to_block_boundary(self, tokens: list[int], pad_token: int) -> list[int]: """Return a copy of ``tokens`` padded to a vLLM block boundary. @@ -126,7 +114,7 @@ def chunk( """ chunks: list[TokenChunk] = [] chunked_tokens = ( - self.pad_to_chunk_boundary(tokens, pad_token) + self.pad_to_block_boundary(tokens, pad_token) if pad_token is not None else list(tokens) ) diff --git a/daser/server/http/vllm_client.py b/daser/server/http/vllm_client.py index e296f45..ec778a2 100644 --- a/daser/server/http/vllm_client.py +++ b/daser/server/http/vllm_client.py @@ -58,6 +58,18 @@ async def close(self) -> None: await self._client.aclose() self._client = None + def _ensure_client(self) -> httpx.AsyncClient: + """Return the HTTP client, opening it lazily on first use. + + Returns: + The live ``httpx.AsyncClient`` bound to the vLLM base URL. + """ + client = self._client + if client is None: + client = httpx.AsyncClient(base_url=self._base_url, timeout=self._timeout) + self._client = client + return client + async def _post_completions(self, payload: dict[str, Any]) -> dict[str, Any]: """POST to ``/v1/completions`` and return the JSON body. @@ -67,11 +79,7 @@ async def _post_completions(self, payload: dict[str, Any]) -> dict[str, Any]: Returns: Parsed JSON response. """ - client = self._client - if client is None: - client = httpx.AsyncClient(base_url=self._base_url, timeout=self._timeout) - self._client = client - resp = await client.post("/v1/completions", json=payload) + resp = await self._ensure_client().post("/v1/completions", json=payload) resp.raise_for_status() return resp.json() @@ -95,42 +103,6 @@ async def prefill(self, tokens: list[int]) -> None: } await self._post_completions(payload) - async def completion( - self, - tokens: list[int], - gen_params: Optional[dict[str, Any]] = None, - kv_transfer_params: Optional[dict[str, Any]] = None, - ) -> dict[str, Any]: - """Run a normal completion for the supplied tokens. - - Args: - tokens: token IDs forming the prompt. - gen_params: optional OpenAI-style generation parameters - (max_tokens, temperature, top_p, ...). Unknown keys - are forwarded untouched; vLLM decides what to accept. - kv_transfer_params: optional per-request KV-transfer hints - forwarded to vLLM as the top-level ``kv_transfer_params`` - field. vLLM exposes the dict on ``request.kv_transfer_params`` - so the connector can adjust per-request KV behavior (e.g. - skipping persistence of task-prompt KV). When ``None`` - the field is omitted to preserve the prior request shape. - - Returns: - Parsed OpenAI-format completion response. - """ - payload: dict[str, Any] = { - "model": self._model, - "prompt": tokens, - "max_tokens": 256, - "temperature": 0.7, - "stream": False, - } - if gen_params: - payload.update(gen_params) - if kv_transfer_params is not None: - payload["kv_transfer_params"] = kv_transfer_params - return await self._post_completions(payload) - async def completion_with_ttft( self, tokens: list[int], @@ -168,10 +140,7 @@ async def completion_with_ttft( if kv_transfer_params is not None: payload["kv_transfer_params"] = kv_transfer_params - client = self._client - if client is None: - client = httpx.AsyncClient(base_url=self._base_url, timeout=self._timeout) - self._client = client + client = self._ensure_client() text_parts: list[str] = [] usage: dict[str, Any] = {} @@ -214,12 +183,8 @@ async def health(self) -> bool: Returns: True on HTTP 200, False otherwise. """ - client = self._client - if client is None: - client = httpx.AsyncClient(base_url=self._base_url, timeout=self._timeout) - self._client = client try: - resp = await client.get("/health") + resp = await self._ensure_client().get("/health") return resp.status_code == 200 except Exception as exc: # noqa: BLE001 logger.warning("[SERVICE] vLLM health check failed: %s", exc) @@ -235,11 +200,7 @@ async def list_models(self) -> list[str]: Uses the client's asyncio HTTP session and must be awaited from an event loop. """ - client = self._client - if client is None: - client = httpx.AsyncClient(base_url=self._base_url, timeout=self._timeout) - self._client = client - resp = await client.get("/v1/models") + resp = await self._ensure_client().get("/v1/models") resp.raise_for_status() body = resp.json() return [str(model["id"]) for model in body.get("data", [])] diff --git a/daser/server/ipc/server.py b/daser/server/ipc/server.py index d702433..e73a6ae 100644 --- a/daser/server/ipc/server.py +++ b/daser/server/ipc/server.py @@ -3,6 +3,7 @@ # Standard import asyncio from collections import OrderedDict +from collections.abc import Awaitable, Callable import contextlib from dataclasses import asdict import os @@ -125,6 +126,26 @@ def __init__( self._cuda_ipc_cache: OrderedDict[ tuple[int, int, int, int | None], "_CachedCudaArray" ] = OrderedDict() + self._op_handlers: dict[ + str, Callable[[dict[str, Any]], Awaitable[dict[str, Any]]] + ] = { + "lookup": self._op_lookup, + "record_external_prefix_cache": self._op_record_external_prefix_cache, + "get_runtime_config": self._op_get_runtime_config, + "alloc_chunk": self._op_alloc_chunk, + "alloc_chunks": self._op_alloc_chunks, + "match_and_alloc": self._op_match_and_alloc, + "commit_chunk": self._op_commit_chunk, + "commit_chunks": self._op_commit_chunks, + "commit_stats": self._op_commit_stats, + "live_allocations": self._op_live_allocations, + "transfer_drain": self._op_transfer_drain, + "init_transfer": self._op_init_transfer, + "transfer_store": self._transfer_store, + "transfer_load": self._transfer_load, + "evict_chunk": self._op_evict_chunk, + "release_chunk_writer": self._op_release_chunk_writer, + } async def start(self) -> None: """Start listening on the Unix socket. @@ -161,9 +182,7 @@ async def drain_transfer(self) -> None: when present. """ transfer = self._ensure_transfer() - drain = getattr(transfer, "drain", None) - if drain is not None: - await drain() + await transfer.drain() async def stop(self) -> None: """Stop the server and remove the socket file. @@ -198,9 +217,7 @@ async def close(self) -> None: rejected. """ if self._transfer is not None: - drain = getattr(self._transfer, "drain", None) - if drain is not None: - await drain() + await self._transfer.drain() self._transfer.close() self._transfer = None self._close_cuda_ipc_cache() @@ -256,103 +273,13 @@ async def _dispatch(self, msg: dict[str, Any]) -> dict[str, Any]: started = time.perf_counter() status = "error" try: - if op == "lookup": - chunks = await self._core.lookup(msg["tokens"], msg["model_id"]) - if "external_prefix_queries" in msg: - queries = int(msg.get("external_prefix_queries", 0)) - await self._core.record_external_prefix_cache( - queries=queries, - hits=_external_prefix_hits( - chunks, - num_computed_tokens=int(msg.get("num_computed_tokens", 0)), - queries=queries, - ), - ) - status = "ok" - return {"chunks": [chunk.to_dict() for chunk in chunks]} - if op == "record_external_prefix_cache": - await self._core.record_external_prefix_cache( - queries=int(msg.get("queries", 0)), - hits=int(msg.get("hits", 0)), - ) - status = "ok" - return {"ok": True} - if op == "get_runtime_config": - status = "ok" - return {"runtime_config": dict(self._runtime_config)} - if op == "alloc_chunk": - alloc = await self._core.alloc_chunk( - msg["chunk_key"], int(msg["token_count"]), msg["model_id"] - ) - status = "ok" - return alloc.to_dict(include_chunk_key=False) - if op == "alloc_chunks": - allocs = await self._core.alloc_chunks( - list(msg.get("chunks", [])), msg["model_id"] - ) - status = "ok" - return { - "allocations": [ - alloc.to_dict(include_chunk_key=True) for alloc in allocs - ] - } - if op == "match_and_alloc": - result = await self._core.match_and_alloc( - msg["tokens"], msg.get("chunk_key", ""), msg["model_id"] - ) - status = "ok" - return result.to_dict() - if op == "commit_chunk": - await self._core.commit_chunk(msg["chunk_key"]) - status = "ok" - return {"ok": True} - if op == "commit_chunks": - for chunk_key in msg.get("chunk_keys", []): - await self._core.commit_chunk(chunk_key) - status = "ok" - return {"ok": True} - if op == "commit_stats": - status = "ok" - return {"commit_stats": await self._core.commit_stats()} - if op == "live_allocations": - live = await self._core.live_allocations( - list(msg.get("allocations", [])) - ) - status = "ok" - return {"chunk_keys": live} - if op == "transfer_drain": - transfer = self._transfer - if transfer is not None: - drain = getattr(transfer, "drain", None) - if drain is not None: - await drain() - status = "ok" - return {"ok": True} - if op == "init_transfer": - self._ensure_transfer() - status = "ok" - return {"ok": True} - if op == "transfer_store": - response = await self._transfer_store(msg) - status = "ok" if response.get("ok") is True else "error" - return response - if op == "transfer_load": - response = await self._transfer_load(msg) - status = "ok" if response.get("ok") is True else "error" - return response - if op == "evict_chunk": - await self._core.evict_chunk(msg["chunk_key"]) - status = "ok" - return {"ok": True} - if op == "release_chunk_writer": - released = await self._core.release_chunk_writer( - chunk_key=str(msg["chunk_key"]), - start_slot=int(msg["start_slot"]), - num_slots=int(msg["num_slots"]), - ) - status = "ok" - return {"released": released} - return {"error": f"unknown op: {op}"} + handler = self._op_handlers.get(op) + if handler is None: + return {"error": f"unknown op: {op}"} + response = await handler(msg) + ok = response.get("ok", True) is not False and "error" not in response + status = "ok" if ok else "error" + return response except Exception as exc: # noqa: BLE001 logger.exception("[IPC] request failed: %s", exc) return {"error": str(exc)} @@ -368,6 +295,104 @@ async def _dispatch(self, msg: dict[str, Any]) -> dict[str, Any]: buckets=(0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0), ).observe(elapsed, labels=ipc_labels) + async def _op_lookup(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle a ``lookup`` request, recording external prefix counters.""" + chunks = await self._core.lookup(msg["tokens"], msg["model_id"]) + if "external_prefix_queries" in msg: + queries = int(msg.get("external_prefix_queries", 0)) + await self._core.record_external_prefix_cache( + queries=queries, + hits=_external_prefix_hits( + chunks, + num_computed_tokens=int(msg.get("num_computed_tokens", 0)), + queries=queries, + ), + ) + return {"chunks": [chunk.to_dict() for chunk in chunks]} + + async def _op_record_external_prefix_cache( + self, msg: dict[str, Any] + ) -> dict[str, Any]: + """Handle a standalone ``record_external_prefix_cache`` request.""" + await self._core.record_external_prefix_cache( + queries=int(msg.get("queries", 0)), + hits=int(msg.get("hits", 0)), + ) + return {"ok": True} + + async def _op_get_runtime_config(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle a ``get_runtime_config`` request.""" + return {"runtime_config": dict(self._runtime_config)} + + async def _op_alloc_chunk(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle a single ``alloc_chunk`` request.""" + alloc = await self._core.alloc_chunk( + msg["chunk_key"], int(msg["token_count"]), msg["model_id"] + ) + return alloc.to_dict(include_chunk_key=False) + + async def _op_alloc_chunks(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle a batched ``alloc_chunks`` request.""" + allocs = await self._core.alloc_chunks( + list(msg.get("chunks", [])), msg["model_id"] + ) + return { + "allocations": [alloc.to_dict(include_chunk_key=True) for alloc in allocs] + } + + async def _op_match_and_alloc(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle a ``match_and_alloc`` request.""" + result = await self._core.match_and_alloc( + msg["tokens"], msg.get("chunk_key", ""), msg["model_id"] + ) + return result.to_dict() + + async def _op_commit_chunk(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle a single ``commit_chunk`` request.""" + await self._core.commit_chunk(msg["chunk_key"]) + return {"ok": True} + + async def _op_commit_chunks(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle a batched ``commit_chunks`` request.""" + for chunk_key in msg.get("chunk_keys", []): + await self._core.commit_chunk(chunk_key) + return {"ok": True} + + async def _op_commit_stats(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle a ``commit_stats`` request.""" + return {"commit_stats": await self._core.commit_stats()} + + async def _op_live_allocations(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle a ``live_allocations`` request.""" + live = await self._core.live_allocations(list(msg.get("allocations", []))) + return {"chunk_keys": live} + + async def _op_transfer_drain(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle a ``transfer_drain`` request.""" + transfer = self._transfer + if transfer is not None: + await transfer.drain() + return {"ok": True} + + async def _op_init_transfer(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle an ``init_transfer`` request.""" + self._ensure_transfer() + return {"ok": True} + + async def _op_evict_chunk(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle an ``evict_chunk`` request.""" + await self._core.evict_chunk(msg["chunk_key"]) + return {"ok": True} + + async def _op_release_chunk_writer(self, msg: dict[str, Any]) -> dict[str, Any]: + """Handle a ``release_chunk_writer`` request.""" + released = await self._core.release_chunk_writer( + chunk_key=str(msg["chunk_key"]), + start_slot=int(msg["start_slot"]), + num_slots=int(msg["num_slots"]), + ) + return {"released": released} + async def _transfer_store(self, msg: dict[str, Any]) -> dict[str, Any]: """Store one or more spans through the server-owned transfer layer. @@ -413,19 +438,10 @@ async def _transfer_store(self, msg: dict[str, Any]) -> dict[str, Any]: store_spans = ( _coalesce_transfer_spans(live_spans) - if bool(getattr(transfer, "coalesce_store_spans", False)) + if transfer.coalesce_store_spans else live_spans ) - grouped_store = getattr(transfer, "store_bytes_grouped", None) - if grouped_store is not None: - total = await grouped_store(buffer, store_spans) - else: - for span in store_spans: - source_offset = int(span.get("source_offset", 0)) - nbytes = int(span["nbytes"]) - file_offset = int(span["file_offset"]) - src = buffer[source_offset : source_offset + nbytes] - total += await transfer.store_bytes(src, file_offset, nbytes) + total = await transfer.store_bytes_grouped(buffer, store_spans) finally: if isinstance(buffer, _UncachedCudaArray): buffer.close() @@ -469,23 +485,10 @@ async def _transfer_load(self, msg: dict[str, Any]) -> dict[str, Any]: sync_ms = 0.0 close_ms = 0.0 try: - stats_before = getattr(transfer, "stats", None) - before = asdict(stats_before) if stats_before is not None else {} + before = asdict(transfer.stats) started = time.perf_counter() load_start = time.perf_counter() - grouped_load = getattr(transfer, "load_bytes_grouped", None) - if grouped_load is not None: - total = await grouped_load(buffer, spans) - else: - for span in spans: - target_offset = int(span.get("target_offset", 0)) - nbytes = int(span["nbytes"]) - file_offset = int(span["file_offset"]) - if isinstance(buffer, bytearray): - dst = memoryview(buffer)[target_offset : target_offset + nbytes] - else: - dst = buffer[target_offset : target_offset + nbytes] - total += await transfer.load_bytes(dst, file_offset, nbytes) + total = await transfer.load_bytes_grouped(buffer, spans) load_ms = (time.perf_counter() - load_start) * 1000 synchronize = getattr(buffer, "synchronize", None) if synchronize is not None: @@ -493,8 +496,7 @@ async def _transfer_load(self, msg: dict[str, Any]) -> dict[str, Any]: synchronize() sync_ms = (time.perf_counter() - sync_start) * 1000 elapsed_ms = (time.perf_counter() - started) * 1000 - stats_after = getattr(transfer, "stats", None) - after = asdict(stats_after) if stats_after is not None else {} + after = asdict(transfer.stats) stats_delta = { key: int(after.get(key, 0)) - int(before.get(key, 0)) for key in set(before) | set(after) @@ -595,11 +597,9 @@ def _record_l1_metrics(self) -> None: transfer = self._transfer if transfer is None: return - stats = getattr(transfer, "stats", None) - if stats is None: - return - current_hits = getattr(stats, "l1_hits", 0) - current_misses = getattr(stats, "l1_misses", 0) + stats = transfer.stats + current_hits = stats.l1_hits + current_misses = stats.l1_misses prev_hits = getattr(self, "_prev_l1_hits", 0) prev_misses = getattr(self, "_prev_l1_misses", 0) delta_hits = current_hits - prev_hits @@ -614,7 +614,7 @@ def _record_l1_metrics(self) -> None: ).inc(delta_misses) self._prev_l1_hits = current_hits self._prev_l1_misses = current_misses - l1_used = getattr(transfer, "_l1_used", 0) + l1_used = transfer.l1_bytes_used l1_capacity = int(self._runtime_config.get("l1_size_bytes", 0)) self._metrics.gauge("daser_l1_bytes_used", "L1 memory cache bytes in use.").set( l1_used diff --git a/daser/server/metadata_store.py b/daser/server/metadata_store.py index dc50e1a..03d8a85 100644 --- a/daser/server/metadata_store.py +++ b/daser/server/metadata_store.py @@ -197,9 +197,8 @@ def iter_chunks(self): def save(self, path: str) -> None: """Serialize the index to a msgpack file. - Writes ring buffer state (head, tail) are not stored here — they - are passed in by ChunkManager. Call ChunkManager.save() instead - to persist the full state. + Ring buffer state (head, tail) is not stored here — it is owned by + ChunkManager. Call ChunkManager.save() instead to persist full state. Args: path: absolute path to write the msgpack file. @@ -207,10 +206,7 @@ def save(self, path: str) -> None: payload = { "total_slots": self._total_slots, "chunk_index": {k: asdict(v) for k, v in self._chunk_index.items()}, - "slot_map": [ - {"kind": e.kind, "chunk_key": e.chunk_key, "num_slots": e.num_slots} - for e in self._slot_map - ], + "slot_map": [asdict(e) for e in self._slot_map], } with open(path, "wb") as f: f.write(msgpack.packb(payload, use_bin_type=True)) diff --git a/daser/transfer/__init__.py b/daser/transfer/__init__.py index 975fb84..749975d 100644 --- a/daser/transfer/__init__.py +++ b/daser/transfer/__init__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -from daser.transfer.base import TransferLayer, TransferMode, TransferStats +from daser.transfer.base import TransferLayer, TransferStats from daser.transfer.iouring import TieredIOUringTransferLayer __all__ = [ @@ -8,7 +8,6 @@ "TieredIOUringTransferLayer", "TransferBackend", "TransferLayer", - "TransferMode", "TransferStats", ] diff --git a/daser/transfer/base.py b/daser/transfer/base.py index 21a338e..c72bb55 100644 --- a/daser/transfer/base.py +++ b/daser/transfer/base.py @@ -3,17 +3,9 @@ # Standard from abc import ABC, abstractmethod from dataclasses import dataclass -import enum from typing import Any -class TransferMode(enum.Enum): - """Server-selected transfer implementation.""" - - GDS = "gds" - IOURING = "iouring" - - @dataclass class TransferStats: """Common transfer counters. @@ -37,12 +29,44 @@ class TransferLayer(ABC): Concrete implementations hide whether storage is GDS direct-to-GPU or io_uring through pinned host memory. + Capability surface: + Backends advertise optional behavior through attributes and overridable + methods rather than ad-hoc duck typing. ``coalesce_store_spans`` lets a + backend opt into adjacent-span coalescing; ``stats`` and + ``l1_bytes_used`` expose tiering counters; ``drain`` waits for + background work; ``store_bytes_grouped``/``load_bytes_grouped`` execute + multi-span batches and default to looping over the single-span methods. + Async/thread-safety: Methods are coroutine-compatible and should be called from the server asyncio event loop. Implementations may offload blocking syscalls to an executor. """ + #: When True the server coalesces adjacent store spans before dispatch. + coalesce_store_spans: bool = False + + @property + def stats(self) -> TransferStats: + """Return tiering counters for this backend. + + Returns: + The backend's ``_stats`` snapshot when it maintains one, otherwise + zeroed counters for backends without a memory tier. + """ + stats = getattr(self, "_stats", None) + return stats if stats is not None else TransferStats() + + @property + def l1_bytes_used(self) -> int: + """Return L1 memory-tier bytes currently resident. + + Returns: + The backend's ``_l1_used`` counter, or 0 for backends without a + memory tier. + """ + return int(getattr(self, "_l1_used", 0)) + @abstractmethod async def load_bytes(self, dst: Any, file_offset: int, nbytes: int) -> int: """Load bytes into ``dst``. @@ -69,6 +93,63 @@ async def store_bytes(self, src: Any, file_offset: int, nbytes: int) -> int: Number of bytes stored. """ + async def store_bytes_grouped(self, src: Any, spans: list[dict[str, int]]) -> int: + """Store multiple spans from ``src``. + + Args: + src: readable buffer or GPU array spanning all source offsets. + spans: span dicts with ``source_offset``, ``nbytes``, ``file_offset``. + + Returns: + Total number of bytes stored. + + Async/thread-safety: + Default implementation loops over ``store_bytes``. Backends may + override to batch the spans into a single backend call. + """ + total = 0 + for span in spans: + source_offset = int(span.get("source_offset", 0)) + nbytes = int(span["nbytes"]) + file_offset = int(span["file_offset"]) + total += await self.store_bytes( + src[source_offset : source_offset + nbytes], file_offset, nbytes + ) + return total + + async def load_bytes_grouped(self, dst: Any, spans: list[dict[str, int]]) -> int: + """Load multiple spans into ``dst``. + + Args: + dst: writable buffer or GPU array spanning all target offsets. + spans: span dicts with ``target_offset``, ``nbytes``, ``file_offset``. + + Returns: + Total number of bytes loaded. + + Async/thread-safety: + Default implementation loops over ``load_bytes``. Backends may + override to batch the spans into a single backend call. + """ + total = 0 + view = memoryview(dst) if isinstance(dst, (bytearray, bytes)) else dst + for span in spans: + target_offset = int(span.get("target_offset", 0)) + nbytes = int(span["nbytes"]) + file_offset = int(span["file_offset"]) + total += await self.load_bytes( + view[target_offset : target_offset + nbytes], file_offset, nbytes + ) + return total + + async def drain(self) -> None: # noqa: B027 + """Wait for any background transfer work to complete. + + Async/thread-safety: + Default implementation is an intentional no-op for backends with no + background work. Backends with async write-back override this. + """ + @abstractmethod def close(self) -> None: """Release file handles and backend resources.""" diff --git a/daser/transfer/iouring/copy_ops.py b/daser/transfer/iouring/copy_ops.py new file mode 100644 index 0000000..7c86187 --- /dev/null +++ b/daser/transfer/iouring/copy_ops.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Stateless copy/marshalling helpers for the io_uring transfer layer. + +These functions move bytes between CPU buffers, CuPy device arrays, and the +pinned host slices used by the L1 tier. They hold no tiering state and take no +locks, so they are safe to call from either the event loop or executor threads +as long as the caller owns the buffers passed in. +""" + +# Standard +from typing import Any + +# First Party +from daser.transfer.iouring.pinned_pool import PinnedMemorySlice + +#: One grouped copy chunk: (target_offset, source_slice, source_offset, nbytes). +CopyChunk = tuple[int, PinnedMemorySlice, int, int] + + +def cuda_array_ptr(dst: Any) -> int | None: + """Return a CUDA device pointer for a CuPy-like array destination. + + Args: + dst: candidate destination object. + + Returns: + Integer device pointer when ``dst`` exposes the CuPy ``data.ptr`` + interface, otherwise None. + """ + data = getattr(dst, "data", None) + ptr = getattr(data, "ptr", None) + if ptr is None: + return None + return int(ptr) + + +def slice_dst(dst: Any, offset: int, nbytes: int) -> Any: + """Return a writable destination slice for ``[offset, offset + nbytes)``. + + Args: + dst: writable buffer or CuPy ndarray. + offset: byte offset into ``dst``. + nbytes: number of bytes the slice must cover. + + Returns: + A sliced view appropriate for the destination type. + """ + if hasattr(dst, "set"): + try: + return dst[offset : offset + nbytes] + except (TypeError, KeyError, IndexError): + if offset == 0: + return dst + raise + if isinstance(dst, bytearray | memoryview): + return memoryview(dst).cast("B")[offset : offset + nbytes] + try: + return dst[offset : offset + nbytes] + except (TypeError, KeyError, IndexError): + pass + return memoryview(dst).cast("B")[offset : offset + nbytes] + + +def slice_src(src: Any, offset: int, nbytes: int) -> Any: + """Return a readable source slice for ``[offset, offset + nbytes)``. + + Args: + src: readable buffer or CuPy ndarray. + offset: byte offset into ``src``. + nbytes: number of bytes the slice must cover. + + Returns: + A sliced view appropriate for the source type. + """ + if hasattr(src, "get"): + return src[offset : offset + nbytes] + try: + return src[offset : offset + nbytes] + except (TypeError, KeyError, IndexError): + pass + return memoryview(src).cast("B")[offset : offset + nbytes] + + +def copy_src_to_pinned( + src: Any, + pinned: PinnedMemorySlice, + target_offset: int, + nbytes: int, +) -> None: + """Copy bytes from a CPU or CuPy source into pinned host memory. + + Args: + src: readable byte buffer or CuPy ndarray. + pinned: destination slice leased from the L1 pool. + target_offset: byte offset into ``pinned`` to write at. + nbytes: number of bytes to copy. + """ + if hasattr(src, "data") and getattr(src.data, "ptr", None) is not None: + from cupy.cuda import runtime + + runtime.memcpy( + pinned.ptr_at(target_offset), + int(src.data.ptr), + nbytes, + runtime.memcpyDeviceToHost, + ) + return + pinned.view()[target_offset : target_offset + nbytes] = memoryview(src).cast("B")[ + :nbytes + ] + + +def coalesce_copy_chunks(chunks: list[CopyChunk]) -> list[CopyChunk]: + """Merge adjacent copies sharing one source slice and contiguous offsets. + + Args: + chunks: grouped copy chunks to merge. + + Returns: + Chunks sorted by target offset with adjacent runs merged. + """ + ordered = sorted(chunks, key=lambda item: item[0]) + merged: list[CopyChunk] = [] + for target_offset, data, source_offset, nbytes in ordered: + if not merged: + merged.append((target_offset, data, source_offset, nbytes)) + continue + prev_target, prev_data, prev_source, prev_nbytes = merged[-1] + if ( + prev_data is data + and target_offset == prev_target + prev_nbytes + and source_offset == prev_source + prev_nbytes + ): + merged[-1] = (prev_target, prev_data, prev_source, prev_nbytes + nbytes) + continue + merged.append((target_offset, data, source_offset, nbytes)) + return merged + + +def copy_grouped_to_cuda_dst(dst: Any, chunks: list[CopyChunk]) -> None: + """Copy grouped pinned ranges into a CUDA destination. + + Args: + dst: CuPy ndarray destination. + chunks: grouped copy chunks from the L1 tier. + """ + from cupy.cuda import runtime + + ordered = sorted(chunks, key=lambda item: item[0]) + merged: list[tuple[int, int, int]] = [] + for target_offset, data, source_offset, nbytes in ordered: + source_ptr = data.ptr_at(source_offset) + if not merged: + merged.append((target_offset, source_ptr, nbytes)) + continue + prev_target, prev_source, prev_nbytes = merged[-1] + if ( + target_offset == prev_target + prev_nbytes + and source_ptr == prev_source + prev_nbytes + ): + merged[-1] = (prev_target, prev_source, prev_nbytes + nbytes) + continue + merged.append((target_offset, source_ptr, nbytes)) + + for target_offset, source_ptr, nbytes in merged: + target = slice_dst(dst, target_offset, nbytes) + dst_ptr = cuda_array_ptr(target) + if dst_ptr is None: + raise TypeError("grouped CUDA copy target lost CUDA array interface") + runtime.memcpyAsync( + dst_ptr, + source_ptr, + nbytes, + runtime.memcpyHostToDevice, + 0, + ) + + +def copy_grouped_to_dst(dst: Any, chunks: list[CopyChunk]) -> None: + """Copy source chunks into the destination without staging repacks. + + Args: + dst: writable buffer or CuPy ndarray destination. + chunks: grouped copy chunks from the L1 tier. + """ + if not chunks: + return + first_target = slice_dst(dst, chunks[0][0], chunks[0][3]) + if cuda_array_ptr(first_target) is not None: + copy_grouped_to_cuda_dst(dst, chunks) + return + + for target_offset, data, source_offset, nbytes in coalesce_copy_chunks(chunks): + target = slice_dst(dst, target_offset, nbytes) + if hasattr(target, "set"): + import numpy + + host = numpy.frombuffer( + data.view()[source_offset : source_offset + nbytes], + dtype=numpy.uint8, + count=nbytes, + ) + target.set(host) + continue + dst_view = memoryview(dst).cast("B") + dst_view[target_offset : target_offset + nbytes] = data.view()[ + source_offset : source_offset + nbytes + ] diff --git a/daser/transfer/iouring/l1_cache.py b/daser/transfer/iouring/l1_cache.py new file mode 100644 index 0000000..e7254b0 --- /dev/null +++ b/daser/transfer/iouring/l1_cache.py @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Range-keyed pinned-host LRU cache for the io_uring L1 tier. + +This cache maps L2 byte ranges to pinned-memory slices and enforces an LRU +capacity bound. It is backend-agnostic and holds no io_uring or asyncio state: +the transfer-layer orchestrator owns the metadata lock and calls these methods +with it held. Because an in-flight L2 write can pin a pool slice, the cache +asks the orchestrator whether a slice is still pinned through an injected +predicate before closing it on eviction. +""" + +# Standard +import bisect +from collections import OrderedDict +from collections.abc import Callable +from dataclasses import dataclass + +# First Party +from daser.replacement import LRUReplacementPolicy +from daser.transfer.iouring.pinned_pool import PinnedMemoryPool, PinnedMemorySlice + + +@dataclass(frozen=True) +class L1RangeHit: + """One L1-backed subrange inside a requested load span.""" + + target_offset: int + key: tuple[int, int] + data: PinnedMemorySlice + source_offset: int + nbytes: int + + +class L1Cache: + """Pinned-host LRU cache keyed by ``(file_offset, nbytes)`` byte ranges. + + Args: + l1_bytes: maximum resident bytes in the memory tier. + alignment: pinned-pool allocation alignment for O_DIRECT compatibility. + pinned_predicate: returns True when a ``(key, slice)`` pair is still + owned by an in-flight L2 write and must not be closed on eviction. + + Async/thread-safety: + Not internally synchronized. All methods assume the orchestrator's + metadata lock is held; ``register_pool_waiter`` futures are resolved by + ``notify_pool_waiters`` after space frees up. + """ + + def __init__( + self, + l1_bytes: int, + alignment: int, + pinned_predicate: Callable[[tuple[int, int], PinnedMemorySlice], bool], + ) -> None: + self._l1_bytes = l1_bytes + self._pool = PinnedMemoryPool(l1_bytes, alignment=alignment) + self._entries: OrderedDict[tuple[int, int], PinnedMemorySlice] = OrderedDict() + self._starts: list[int] = [] + self._by_start: dict[int, tuple[int, int]] = {} + self._used = 0 + self._policy = LRUReplacementPolicy[tuple[int, int]]() + self._pool_waiters: list[object] = [] + self._is_pinned = pinned_predicate + + @property + def bytes_used(self) -> int: + """Return resident L1 bytes.""" + return self._used + + def get(self, key: tuple[int, int]) -> PinnedMemorySlice | None: + """Return the resident slice for ``key`` or None.""" + return self._entries.get(key) + + def contains_slice(self, data: PinnedMemorySlice) -> bool: + """Return whether ``data`` is currently a resident L1 slice.""" + return data in self._entries.values() + + def resident_slice_ids(self) -> set[int]: + """Return ``id()`` of every resident slice for liveness checks.""" + return {id(buffer) for buffer in self._entries.values()} + + def close(self) -> None: + """Release the pinned pool backing this cache.""" + self._pool.close() + + def find( + self, file_offset: int + ) -> tuple[tuple[int, int], PinnedMemorySlice, int] | None: + """Return the cached range containing ``file_offset``. + + Args: + file_offset: L2 byte offset to locate. + + Returns: + ``(key, slice, source_offset)`` when a resident range covers the + offset, otherwise None. + """ + idx = bisect.bisect_right(self._starts, file_offset) - 1 + if idx < 0: + return None + start = self._starts[idx] + key = self._by_start.get(start) + if key is None: + return None + data = self._entries.get(key) + if data is None: + return None + if file_offset < key[0] + key[1]: + return key, data, file_offset - key[0] + return None + + def resolve_subranges( + self, + target_offset: int, + file_offset: int, + nbytes: int, + ) -> tuple[list[L1RangeHit], list[dict[str, int]]]: + """Split a load span into cached slices and uncached gaps. + + Args: + target_offset: destination byte offset matching ``file_offset``. + file_offset: L2 byte offset where the requested span starts. + nbytes: requested byte count. + + Returns: + A pair of L1 hit slices and L2 miss gaps in ascending file-offset + order. + """ + hits: list[L1RangeHit] = [] + misses: list[dict[str, int]] = [] + request_end = file_offset + nbytes + cursor = file_offset + while cursor < request_end: + hit = self.find(cursor) + if hit is not None: + key, data, source_offset = hit + covered = min(key[0] + key[1], request_end) - cursor + hits.append( + L1RangeHit( + target_offset=target_offset + (cursor - file_offset), + key=key, + data=data, + source_offset=source_offset, + nbytes=covered, + ) + ) + cursor += covered + continue + + next_idx = bisect.bisect_left(self._starts, cursor) + next_start = ( + self._starts[next_idx] if next_idx < len(self._starts) else request_end + ) + gap_end = min(next_start, request_end) + if gap_end <= cursor: + gap_end = request_end + misses.append( + { + "target_offset": target_offset + (cursor - file_offset), + "file_offset": cursor, + "nbytes": gap_end - cursor, + } + ) + cursor = gap_end + return hits, misses + + def record_hits(self, hits: list[L1RangeHit]) -> None: + """Refresh LRU recency for hit slices. + + Args: + hits: slices returned by ``resolve_subranges``. + """ + for hit in hits: + self._policy.access(hit.key) + self._entries.move_to_end(hit.key) + + def touch(self, key: tuple[int, int]) -> None: + """Refresh LRU recency for one resident key (used on in-place stores).""" + self._policy.access(key) + self._entries.move_to_end(key) + + def put(self, key: tuple[int, int], data: PinnedMemorySlice) -> None: + """Insert bytes into L1 after dropping overlapping ranges. + + Args: + key: ``(file_offset, nbytes)`` range key. + data: pinned slice holding the range's bytes. + """ + self.drop_overlapping(key[0], key[1]) + self._insert_entry(key, data) + + def reserve( + self, + key: tuple[int, int], + nbytes: int, + *, + drop_overlaps: bool = True, + preserve_overlaps: bool = False, + ) -> PinnedMemorySlice | None: + """Try to reserve pinned space for a store or promoted load. + + Args: + key: range key being inserted. + nbytes: logical bytes needed. + drop_overlaps: drop resident ranges overlapping ``key`` first. + preserve_overlaps: keep the non-overlapping remainder of dropped + ranges when ``drop_overlaps`` is set. + + Returns: + A pinned slice, or None when the pool is exhausted and no further + victim can be evicted (the caller must wait for an in-flight L2 + write to free its slice, then retry). + + Raises: + ValueError: if ``nbytes`` exceeds the L1 capacity. + """ + if nbytes > self._l1_bytes: + raise ValueError( + f"range {nbytes} bytes exceeds L1 capacity {self._l1_bytes}" + ) + if drop_overlaps: + self.drop_overlapping(key[0], key[1], preserve_remainder=preserve_overlaps) + data = self._pool.allocate(nbytes) + while data is None: + victim = self._policy.evict() + if victim is None: + return None + removed = self._entries.pop(victim, None) + self._remove_index(victim) + if removed is not None: + self._used -= len(removed) + self.release(victim, removed) + data = self._pool.allocate(nbytes) + return data + + def reserve_or_raise( + self, + key: tuple[int, int], + nbytes: int, + *, + preserve_overlaps: bool = False, + ) -> PinnedMemorySlice: + """Reserve pinned space when no in-flight L2 writer can block reuse. + + Args: + key: range key being inserted. + nbytes: logical bytes needed. + preserve_overlaps: keep non-overlapping remainders of dropped ranges. + + Returns: + A pinned slice. + + Raises: + MemoryError: if the pool cannot satisfy the request. + """ + data = self.reserve(key, nbytes, preserve_overlaps=preserve_overlaps) + if data is None: + raise MemoryError( + f"could not reserve {nbytes} pinned L1 bytes from " + f"{self._l1_bytes} byte pool" + ) + return data + + def release(self, key: tuple[int, int], data: PinnedMemorySlice) -> None: + """Close an evicted L1 slice unless an L2 write still owns it. + + Args: + key: range key being released. + data: slice removed from the cache. + """ + if self._is_pinned(key, data): + return + data.close() + self.notify_pool_waiters() + + def register_pool_waiter(self, waiter: object) -> None: + """Register a future to wake when pool space or metadata changes.""" + self._pool_waiters.append(waiter) + + def notify_pool_waiters(self) -> None: + """Wake futures waiting for L1 pool metadata or free-space changes.""" + waiters = self._pool_waiters + self._pool_waiters = [] + for waiter in waiters: + if not waiter.done(): # type: ignore[attr-defined] + waiter.set_result(None) # type: ignore[attr-defined] + + def drop_overlapping( + self, + file_offset: int, + nbytes: int, + *, + preserve_remainder: bool = False, + ) -> None: + """Remove L1 entries overlapping a newly written byte range. + + Args: + file_offset: start of the overwritten range. + nbytes: length of the overwritten range. + preserve_remainder: re-insert non-overlapping fragments of dropped + ranges as fresh L1 entries. + """ + end = file_offset + nbytes + victims = [ + key + for key in self._entries + if key[0] < end and file_offset < key[0] + key[1] + ] + for victim in victims: + removed = self._entries.pop(victim, None) + self._remove_index(victim) + self._policy.remove(victim) + preserved = ( + self._preserve_non_overlapping(victim, removed, file_offset, end) + if preserve_remainder and removed is not None + else [] + ) + if removed is not None: + self._used -= len(removed) + self.release(victim, removed) + for preserved_key, payload in preserved: + self._put_preserved_fragment(preserved_key, payload) + + def _insert_entry(self, key: tuple[int, int], data: PinnedMemorySlice) -> None: + """Insert one non-overlapping entry and enforce capacity.""" + if len(data) > self._l1_bytes: + return + self._entries[key] = data + self._insert_index(key) + self._entries.move_to_end(key) + self._policy.insert(key) + self._used += len(data) + self.notify_pool_waiters() + while self._used > self._l1_bytes: + victim = self._policy.evict() + if victim is None: + break + removed = self._entries.pop(victim, None) + self._remove_index(victim) + if removed is not None: + self._used -= len(removed) + self.release(victim, removed) + + def _preserve_non_overlapping( + self, + key: tuple[int, int], + data: PinnedMemorySlice, + overlap_start: int, + overlap_end: int, + ) -> list[tuple[tuple[int, int], bytes]]: + """Return old L1 fragments that fall outside an overwrite range.""" + key_start, key_size = key + key_end = key_start + key_size + fragments: list[tuple[tuple[int, int], bytes]] = [] + view = data.view() + if key_start < overlap_start: + keep = overlap_start - key_start + fragments.append(((key_start, keep), bytes(view[:keep]))) + if overlap_end < key_end: + source_offset = overlap_end - key_start + keep = key_end - overlap_end + fragments.append( + ( + (overlap_end, keep), + bytes(view[source_offset : source_offset + keep]), + ) + ) + return fragments + + def _put_preserved_fragment(self, key: tuple[int, int], payload: bytes) -> None: + """Insert one fragment copied out of an overwritten L1 range.""" + if not payload: + return + data = self.reserve(key, len(payload), drop_overlaps=False) + if data is None: + raise MemoryError( + f"could not preserve {len(payload)} L1 bytes from overwritten range" + ) + try: + data.view()[: len(payload)] = payload + except BaseException: + data.close() + raise + self._insert_entry(key, data) + + def _insert_index(self, key: tuple[int, int]) -> None: + """Add one range to the start-offset lookup index.""" + start = key[0] + existing = self._by_start.get(start) + if existing == key: + return + if existing is not None: + self._remove_index(existing) + bisect.insort(self._starts, start) + self._by_start[start] = key + + def _remove_index(self, key: tuple[int, int]) -> None: + """Remove one range from the start-offset lookup index.""" + start = key[0] + if self._by_start.get(start) != key: + return + del self._by_start[start] + idx = bisect.bisect_left(self._starts, start) + if idx < len(self._starts) and self._starts[idx] == start: + self._starts.pop(idx) diff --git a/daser/transfer/iouring/l2_engine.py b/daser/transfer/iouring/l2_engine.py new file mode 100644 index 0000000..9ec6d9e --- /dev/null +++ b/daser/transfer/iouring/l2_engine.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""io_uring positioned I/O engine for the L2 SSD tier. + +This engine owns the O_DIRECT file descriptor, the pool of native io_uring +rings, and the executor used to offload blocking completion waits. It holds no +tiering or L1 state: callers pass absolute byte offsets and memoryviews. Write +scheduling and the pending-write bookkeeping that ties a write to an L1 pool +slice stay in the transfer-layer orchestrator. +""" + +# Standard +import concurrent.futures +import os +import threading + +# First Party +from daser.transfer.iouring.native import NativeIOUring + + +class L2IoEngine: + """Round-robin io_uring positioned reads/writes against the L2 store file. + + Args: + path: pre-allocated or creatable L2 store file. + l2_bytes: SSD-tier capacity; the file is truncated to this size. + io_workers: number of native io_uring rings and executor threads. + + Async/thread-safety: + Ring selection is serialized by an internal lock. Blocking syscalls are + meant to be offloaded onto ``executor`` by the caller, which awaits the + completion on its event loop. + """ + + def __init__(self, path: str, l2_bytes: int, io_workers: int) -> None: + if l2_bytes <= 0: + raise ValueError("l2_bytes must be positive") + if io_workers <= 0: + raise ValueError("io_workers must be positive") + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + with open(path, "a+b") as f: + f.truncate(l2_bytes) + self._fd = os.open(path, os.O_RDWR | os.O_DIRECT) + self._urings = [NativeIOUring(entries=64) for _ in range(io_workers)] + self._executor = concurrent.futures.ThreadPoolExecutor( + max_workers=io_workers, + thread_name_prefix="daser-iouring", + ) + self._uring_lock = threading.Lock() + self._next_uring_index = 0 + + @property + def ring_count(self) -> int: + """Return the number of io_uring rings (used to bound read batches).""" + return len(self._urings) + + @property + def executor(self) -> concurrent.futures.ThreadPoolExecutor: + """Return the executor used to offload blocking io_uring waits.""" + return self._executor + + def next_uring(self) -> NativeIOUring: + """Return the next native io_uring instance for one L2 operation.""" + with self._uring_lock: + uring = self._urings[self._next_uring_index] + self._next_uring_index = (self._next_uring_index + 1) % len(self._urings) + return uring + + def read_into(self, file_offset: int, dst: memoryview, uring: NativeIOUring) -> int: + """Blocking io_uring L2 read into ``dst``. + + Args: + file_offset: byte offset in the L2 store file. + dst: writable memoryview to fill. + uring: ring to submit on (from ``next_uring``). + + Returns: + Number of bytes read. + """ + return uring.read_into(self._fd, file_offset, dst) + + def write(self, file_offset: int, data: memoryview, uring: NativeIOUring) -> int: + """Blocking io_uring L2 write of ``data``. + + Args: + file_offset: byte offset in the L2 store file. + data: readable memoryview to persist. + uring: ring to submit on (from ``next_uring``). + + Returns: + Number of bytes written. + """ + return uring.write(self._fd, file_offset, data) + + def close(self) -> None: + """Shut down the executor, close rings, and close the file descriptor.""" + self._executor.shutdown(wait=True) + for uring in self._urings: + uring.close() + os.close(self._fd) diff --git a/daser/transfer/iouring/layer.py b/daser/transfer/iouring/layer.py index 119f8a4..95a82a7 100644 --- a/daser/transfer/iouring/layer.py +++ b/daser/transfer/iouring/layer.py @@ -2,37 +2,22 @@ # Standard import asyncio -import bisect -from collections import OrderedDict -import concurrent.futures -from dataclasses import dataclass -import os -import threading from typing import Any # First Party from daser.logging import init_logger -from daser.replacement import LRUReplacementPolicy from daser.transfer.base import TransferLayer, TransferStats +from daser.transfer.iouring import copy_ops +from daser.transfer.iouring.l1_cache import L1Cache, L1RangeHit +from daser.transfer.iouring.l2_engine import L2IoEngine from daser.transfer.iouring.native import NativeIOUring -from daser.transfer.iouring.pinned_pool import PinnedMemoryPool, PinnedMemorySlice +from daser.transfer.iouring.pinned_pool import PinnedMemorySlice logger = init_logger(__name__) _DIRECT_IO_ALIGNMENT = 4096 -@dataclass(frozen=True) -class _L1RangeHit: - """One L1-backed subrange inside a requested load span.""" - - target_offset: int - key: tuple[int, int] - data: PinnedMemorySlice - source_offset: int - nbytes: int - - class TieredIOUringTransferLayer(TransferLayer): """Async L1 pinned-memory + L2 SSD transfer layer. @@ -70,42 +55,21 @@ def __init__( raise ValueError("l1_bytes must not exceed l2_bytes") if io_workers <= 0: raise ValueError("io_workers must be positive") - self._skip_l2 = skip_l2 - self._path = path - self._fd: int | None = None - self._urings: list[NativeIOUring] = [] - self._io_executor: concurrent.futures.ThreadPoolExecutor | None = None - self._uring_lock = threading.Lock() - self._next_uring_index = 0 + self._l2: L2IoEngine | None = None if not skip_l2: - parent = os.path.dirname(path) - if parent: - os.makedirs(parent, exist_ok=True) - with open(path, "a+b") as f: - f.truncate(l2_bytes) - self._fd = os.open(path, os.O_RDWR | os.O_DIRECT) - self._urings = [NativeIOUring(entries=64) for _ in range(io_workers)] - self._io_executor = concurrent.futures.ThreadPoolExecutor( - max_workers=io_workers, - thread_name_prefix="daser-iouring", - ) + self._l2 = L2IoEngine(path, l2_bytes, io_workers) self._l1_bytes = l1_bytes self._l2_bytes = l2_bytes - self._pool = PinnedMemoryPool( + self._pending_l2: dict[tuple[int, int], asyncio.Task[None]] = {} + self._pending_l2_buffers: dict[tuple[int, int], PinnedMemorySlice] = {} + self._l1 = L1Cache( l1_bytes, alignment=_DIRECT_IO_ALIGNMENT, + pinned_predicate=self._is_pinned_by_l2, ) - self._l1: OrderedDict[tuple[int, int], PinnedMemorySlice] = OrderedDict() - self._l1_starts: list[int] = [] - self._l1_by_start: dict[int, tuple[int, int]] = {} - self._l1_used = 0 - self._policy = LRUReplacementPolicy[tuple[int, int]]() - self._pending_l2: dict[tuple[int, int], asyncio.Task[None]] = {} - self._pending_l2_buffers: dict[tuple[int, int], PinnedMemorySlice] = {} - self._pool_waiters: list[asyncio.Future[None]] = [] self._l2_errors: list[BaseException] = [] self._lock = asyncio.Lock() - self.stats = TransferStats() + self._stats = TransferStats() logger.info( "[TRANSFER:iouring] path=%s l1=%d l2=%d direct_io=%s " "io_workers=%d skip_l2=%s", @@ -130,26 +94,24 @@ async def load_bytes(self, dst: Any, file_offset: int, nbytes: int) -> int: """ self._check_range(file_offset, nbytes) pending: list[asyncio.Task[None]] = [] - l1_hits: list[_L1RangeHit] = [] + l1_hits: list[L1RangeHit] = [] misses: list[dict[str, int]] = [] async with self._lock: self._raise_l2_error_locked() - l1_hits, misses = self._resolve_l1_subranges_locked( + l1_hits, misses = self._l1.resolve_subranges( target_offset=0, file_offset=file_offset, nbytes=nbytes, ) - if self._skip_l2 and misses: - self.stats.l1_misses += len(misses) + if self._l2 is None and misses: + self._stats.l1_misses += len(misses) raise KeyError( "skip_l2 cache miss for range " f"[{file_offset}, {file_offset + nbytes})" ) if l1_hits: - self._record_l1_hits_locked( - l1_hits, - hit_count=1 if self._skip_l2 else None, - ) + self._l1.record_hits(l1_hits) + self._stats.l1_hits += 1 if self._l2 is None else len(l1_hits) self._copy_grouped_to_dst( dst, [ @@ -163,7 +125,7 @@ async def load_bytes(self, dst: Any, file_offset: int, nbytes: int) -> int: ], ) for miss in misses: - self.stats.l1_misses += 1 + self._stats.l1_misses += 1 pending.extend( self._find_pending_l2_locked( int(miss["file_offset"]), @@ -207,22 +169,20 @@ async def load_bytes_grouped( nbytes = int(span["nbytes"]) self._check_range(file_offset, nbytes) total += nbytes - l1_hits, span_misses = self._resolve_l1_subranges_locked( + l1_hits, span_misses = self._l1.resolve_subranges( target_offset=target_offset, file_offset=file_offset, nbytes=nbytes, ) - if self._skip_l2 and span_misses: - self.stats.l1_misses += 1 + if self._l2 is None and span_misses: + self._stats.l1_misses += 1 raise KeyError( "skip_l2 cache miss for range " f"[{file_offset}, {file_offset + nbytes})" ) if l1_hits: - self._record_l1_hits_locked( - l1_hits, - hit_count=1 if self._skip_l2 else None, - ) + self._l1.record_hits(l1_hits) + self._stats.l1_hits += 1 if self._l2 is None else len(l1_hits) merged_l1.extend( ( hit.target_offset, @@ -233,7 +193,7 @@ async def load_bytes_grouped( for hit in l1_hits ) for miss in span_misses: - self.stats.l1_misses += 1 + self._stats.l1_misses += 1 pending.extend( self._find_pending_l2_locked( int(miss["file_offset"]), @@ -269,17 +229,16 @@ async def store_bytes(self, src: Any, file_offset: int, nbytes: int) -> int: """ self._check_range(file_offset, nbytes) key = (file_offset, nbytes) - if self._skip_l2: + if self._l2 is None: async with self._lock: - hit = self._find_l1_locked(file_offset, nbytes) + hit = self._l1.find(file_offset) if hit is not None: hit_key, cached, target_offset = hit if target_offset + nbytes <= len(cached): self._copy_src_to_pinned_at(src, cached, target_offset, nbytes) - self._policy.access(hit_key) - self._l1.move_to_end(hit_key) + self._l1.touch(hit_key) return nbytes - data = self._reserve_l1_buffer_locked_or_raise( + data = self._l1.reserve_or_raise( key, nbytes, preserve_overlaps=True, @@ -289,7 +248,7 @@ async def store_bytes(self, src: Any, file_offset: int, nbytes: int) -> int: except BaseException: data.close() raise - self._put_l1_locked(key, data) + self._l1.put(key, data) return nbytes data = await self._reserve_l1_buffer(key, nbytes) @@ -301,7 +260,7 @@ async def store_bytes(self, src: Any, file_offset: int, nbytes: int) -> int: async with self._lock: self._raise_l2_error_locked() previous = self._find_pending_l2_locked(file_offset, nbytes) - self._put_l1_locked(key, data) + self._l1.put(key, data) task = self._schedule_l2_write_locked( key, file_offset, @@ -333,7 +292,7 @@ async def store_bytes_grouped( L1 metadata once for the whole group because no pending L2 writer can retain evicted pool slices. """ - if self._skip_l2: + if self._l2 is None: return await self._store_bytes_grouped_l1_only(src, spans) total = 0 @@ -352,7 +311,7 @@ async def drain(self) -> None: Must be called from the owning asyncio event loop before shutdown when durable L2 contents are required. """ - if self._skip_l2: + if self._l2 is None: return None while True: async with self._lock: @@ -368,16 +327,21 @@ def close(self) -> None: task.cancel() self._pending_l2.clear() for pending_buffer in self._pending_l2_buffers.values(): - if pending_buffer not in self._l1.values(): + if not self._l1.contains_slice(pending_buffer): pending_buffer.close() self._pending_l2_buffers.clear() - if self._io_executor is not None: - self._io_executor.shutdown(wait=True) - for uring in self._urings: - uring.close() - if self._fd is not None: - os.close(self._fd) - self._pool.close() + if self._l2 is not None: + self._l2.close() + self._l1.close() + + def _is_pinned_by_l2(self, key: tuple[int, int], data: PinnedMemorySlice) -> bool: + """Return whether ``data`` is still owned by an in-flight L2 write.""" + return self._pending_l2_buffers.get(key) is data + + @property + def l1_bytes_used(self) -> int: + """Return resident L1 memory-tier bytes.""" + return self._l1.bytes_used def _check_range(self, file_offset: int, nbytes: int) -> None: """Validate an L2 byte range. @@ -391,7 +355,7 @@ def _check_range(self, file_offset: int, nbytes: int) -> None: """ if file_offset < 0 or nbytes < 0: raise ValueError("file_offset and nbytes must be non-negative") - if self._skip_l2: + if self._l2 is None: if nbytes > self._l1_bytes: raise ValueError( f"range {nbytes} bytes exceeds L1 capacity {self._l1_bytes}" @@ -411,221 +375,6 @@ def _check_range(self, file_offset: int, nbytes: int) -> None: f"{_DIRECT_IO_ALIGNMENT} bytes: offset={file_offset} nbytes={nbytes}" ) - def _put_l1_locked(self, key: tuple[int, int], data: PinnedMemorySlice) -> None: - """Insert bytes into L1 after dropping overlapping ranges.""" - self._drop_overlapping_l1_locked(key[0], key[1]) - self._insert_l1_entry_locked(key, data) - - def _insert_l1_entry_locked( - self, - key: tuple[int, int], - data: PinnedMemorySlice, - ) -> None: - """Insert one non-overlapping L1 entry and enforce capacity.""" - if len(data) > self._l1_bytes: - return - self._l1[key] = data - self._insert_l1_index_locked(key) - self._l1.move_to_end(key) - self._policy.insert(key) - self._l1_used += len(data) - self._notify_pool_waiters_locked() - while self._l1_used > self._l1_bytes: - victim = self._policy.evict() - if victim is None: - break - removed = self._l1.pop(victim, None) - self._remove_l1_index_locked(victim) - if removed is not None: - self._l1_used -= len(removed) - self._release_l1_buffer_locked(victim, removed) - - def _drop_overlapping_l1_locked( - self, - file_offset: int, - nbytes: int, - *, - preserve_remainder: bool = False, - ) -> None: - """Remove L1 entries that overlap a newly written byte range.""" - end = file_offset + nbytes - victims = [ - key for key in self._l1 if key[0] < end and file_offset < key[0] + key[1] - ] - for victim in victims: - removed = self._l1.pop(victim, None) - self._remove_l1_index_locked(victim) - self._policy.remove(victim) - preserved = ( - self._preserve_non_overlapping_l1_bytes( - victim, - removed, - file_offset, - end, - ) - if preserve_remainder and removed is not None - else [] - ) - if removed is not None: - self._l1_used -= len(removed) - self._release_l1_buffer_locked(victim, removed) - for preserved_key, payload in preserved: - self._put_preserved_l1_fragment_locked(preserved_key, payload) - - def _preserve_non_overlapping_l1_bytes( - self, - key: tuple[int, int], - data: PinnedMemorySlice, - overlap_start: int, - overlap_end: int, - ) -> list[tuple[tuple[int, int], bytes]]: - """Return old L1 fragments that fall outside an overwrite range.""" - key_start, key_size = key - key_end = key_start + key_size - fragments: list[tuple[tuple[int, int], bytes]] = [] - view = data.view() - if key_start < overlap_start: - keep = overlap_start - key_start - fragments.append(((key_start, keep), bytes(view[:keep]))) - if overlap_end < key_end: - source_offset = overlap_end - key_start - keep = key_end - overlap_end - fragments.append( - ( - (overlap_end, keep), - bytes(view[source_offset : source_offset + keep]), - ) - ) - return fragments - - def _put_preserved_l1_fragment_locked( - self, - key: tuple[int, int], - payload: bytes, - ) -> None: - """Insert one preserved fragment copied out of an overwritten L1 range.""" - if not payload: - return - data, wait_for = self._reserve_l1_buffer_locked( - key, - len(payload), - drop_overlaps=False, - ) - if data is None: - if wait_for is not None: - raise RuntimeError("unexpected pending L2 write preserving L1 range") - raise MemoryError( - f"could not preserve {len(payload)} L1 bytes from overwritten range" - ) - try: - data.view()[: len(payload)] = payload - except BaseException: - data.close() - raise - self._insert_l1_entry_locked(key, data) - - def _find_l1_locked( - self, - file_offset: int, - nbytes: int, - ) -> tuple[tuple[int, int], PinnedMemorySlice, int] | None: - """Return the cached L1 range containing the requested start offset.""" - idx = bisect.bisect_right(self._l1_starts, file_offset) - 1 - if idx < 0: - return None - start = self._l1_starts[idx] - key = self._l1_by_start.get(start) - if key is None: - return None - data = self._l1.get(key) - if data is None: - return None - if file_offset < key[0] + key[1]: - return key, data, file_offset - key[0] - return None - - def _resolve_l1_subranges_locked( - self, - target_offset: int, - file_offset: int, - nbytes: int, - ) -> tuple[list[_L1RangeHit], list[dict[str, int]]]: - """Split a load span into cached L1 slices and uncached gaps. - - Args: - target_offset: destination byte offset corresponding to - ``file_offset``. - file_offset: L2 byte offset where the requested span starts. - nbytes: requested byte count. - - Returns: - A pair of L1 hit slices and L2 miss gaps in ascending file-offset - order. Metadata is read under ``_lock`` and no buffers are mutated. - - Async/thread-safety: - Must be called with the transfer metadata lock held. - """ - hits: list[_L1RangeHit] = [] - misses: list[dict[str, int]] = [] - request_end = file_offset + nbytes - cursor = file_offset - while cursor < request_end: - hit = self._find_l1_locked(cursor, request_end - cursor) - if hit is not None: - key, data, source_offset = hit - covered = min(key[0] + key[1], request_end) - cursor - hits.append( - _L1RangeHit( - target_offset=target_offset + (cursor - file_offset), - key=key, - data=data, - source_offset=source_offset, - nbytes=covered, - ) - ) - cursor += covered - continue - - next_idx = bisect.bisect_left(self._l1_starts, cursor) - next_start = ( - self._l1_starts[next_idx] - if next_idx < len(self._l1_starts) - else request_end - ) - gap_end = min(next_start, request_end) - if gap_end <= cursor: - gap_end = request_end - misses.append( - { - "target_offset": target_offset + (cursor - file_offset), - "file_offset": cursor, - "nbytes": gap_end - cursor, - } - ) - cursor = gap_end - return hits, misses - - def _record_l1_hits_locked( - self, - hits: list[_L1RangeHit], - hit_count: int | None = None, - ) -> None: - """Update replacement state and stats for L1 hit slices. - - Args: - hits: L1 slices returned by ``_resolve_l1_subranges_locked``. - - Returns: - None. - - Async/thread-safety: - Must be called with the transfer metadata lock held. - """ - for hit in hits: - self._policy.access(hit.key) - self._l1.move_to_end(hit.key) - self.stats.l1_hits += len(hits) if hit_count is None else hit_count - def _find_pending_l2_locked( self, file_offset: int, @@ -646,9 +395,9 @@ def _read_l2_into( uring: NativeIOUring, ) -> int: """Blocking io_uring L2 read into pinned memory.""" - if self._fd is None: + if self._l2 is None: raise RuntimeError("L2 reads are disabled when skip_l2 is true") - return uring.read_into(self._fd, file_offset, dst.view()) + return self._l2.read_into(file_offset, dst.view(), uring) def _write_l2( self, @@ -657,9 +406,9 @@ def _write_l2( uring: NativeIOUring, ) -> None: """Blocking io_uring L2 write.""" - if self._fd is None: + if self._l2 is None: raise RuntimeError("L2 writes are disabled when skip_l2 is true") - written = uring.write(self._fd, file_offset, data.view()) + written = self._l2.write(file_offset, data.view(), uring) if written != len(data): raise IOError(f"short io_uring write: {written} != {len(data)}") @@ -671,7 +420,7 @@ def _schedule_l2_write_locked( previous: list[asyncio.Task[None]], ) -> asyncio.Task[None]: """Schedule one L2 write and start independent IO immediately.""" - if self._io_executor is None: + if self._l2 is None: raise RuntimeError("L2 writes are disabled when skip_l2 is true") if previous: return asyncio.create_task( @@ -681,7 +430,7 @@ def _schedule_l2_write_locked( loop = asyncio.get_event_loop() uring = self._next_uring() future = loop.run_in_executor( - self._io_executor, + self._l2.executor, self._write_l2, file_offset, data, @@ -699,7 +448,7 @@ async def _track_l2_write( try: await future async with self._lock: - self.stats.l2_writes += 1 + self._stats.l2_writes += 1 except BaseException as exc: async with self._lock: self._l2_errors.append(exc) @@ -728,17 +477,17 @@ async def _write_l2_async( if previous: await asyncio.gather(*previous) loop = asyncio.get_event_loop() - if self._io_executor is None: + if self._l2 is None: raise RuntimeError("L2 writes are disabled when skip_l2 is true") await loop.run_in_executor( - self._io_executor, + self._l2.executor, self._write_l2, file_offset, data, self._next_uring(), ) async with self._lock: - self.stats.l2_writes += 1 + self._stats.l2_writes += 1 except BaseException as exc: async with self._lock: self._l2_errors.append(exc) @@ -780,7 +529,7 @@ async def _load_l2_miss_batch( ) -> None: """Read one bounded L2 miss batch and promote it to L1.""" loop = asyncio.get_event_loop() - if self._io_executor is None: + if self._l2 is None: raise RuntimeError("L2 reads are disabled when skip_l2 is true") reads: list[tuple[dict[str, int], PinnedMemorySlice]] = [] try: @@ -793,7 +542,7 @@ async def _load_l2_miss_batch( await asyncio.gather( *( loop.run_in_executor( - self._io_executor, + self._l2.executor, self._read_l2_into, int(span["file_offset"]), pinned, @@ -820,12 +569,12 @@ async def _load_l2_miss_batch( for span, pinned in reads: nbytes = int(span["nbytes"]) key = (int(span["file_offset"]), nbytes) - self.stats.l2_reads += 1 - self._put_l1_locked(key, pinned) + self._stats.l2_reads += 1 + self._l1.put(key, pinned) except BaseException: live_buffers = set() async with self._lock: - live_buffers = {id(buffer) for buffer in self._l1.values()} + live_buffers = self._l1.resident_slice_ids() for _span, pinned in reads: if id(pinned) not in live_buffers: pinned.close() @@ -837,12 +586,13 @@ def _next_l2_miss_batch( start: int, ) -> list[dict[str, int]]: """Return a miss batch bounded by L1 capacity and io_uring count.""" + ring_count = self._l2.ring_count if self._l2 is not None else 1 batch: list[dict[str, int]] = [] batch_bytes = 0 for span in misses[start:]: nbytes = int(span["nbytes"]) if batch and ( - batch_bytes + nbytes > self._l1_bytes or len(batch) >= len(self._urings) + batch_bytes + nbytes > self._l1_bytes or len(batch) >= ring_count ): break batch.append(span) @@ -851,47 +601,9 @@ def _next_l2_miss_batch( def _next_uring(self) -> NativeIOUring: """Return the next native io_uring instance for one L2 operation.""" - if not self._urings: + if self._l2 is None: raise RuntimeError("io_uring rings are disabled when skip_l2 is true") - with self._uring_lock: - uring = self._urings[self._next_uring_index] - self._next_uring_index = (self._next_uring_index + 1) % len(self._urings) - return uring - - def _copy_pinned_to_dst( - self, - dst: Any, - data: PinnedMemorySlice, - source_offset: int, - nbytes: int, - ) -> None: - """Copy pinned host bytes into a CPU or CUDA destination.""" - target = self._slice_dst(dst, 0, nbytes) - dst_ptr = self._cuda_array_ptr(target) - if dst_ptr is not None: - from cupy.cuda import runtime - - runtime.memcpyAsync( - dst_ptr, - data.ptr_at(source_offset), - nbytes, - runtime.memcpyHostToDevice, - 0, - ) - return - if hasattr(target, "set"): - import numpy - - host = numpy.frombuffer( - data.view()[source_offset : source_offset + nbytes], - dtype=numpy.uint8, - count=nbytes, - ) - target.set(host) - return - memoryview(dst).cast("B")[:nbytes] = data.view()[ - source_offset : source_offset + nbytes - ] + return self._l2.next_uring() def _copy_grouped_to_dst( self, @@ -899,130 +611,11 @@ def _copy_grouped_to_dst( chunks: list[tuple[int, PinnedMemorySlice, int, int]], ) -> None: """Copy source chunks into the destination without staging repacks.""" - if not chunks: - return - first_target = self._slice_dst(dst, chunks[0][0], chunks[0][3]) - if self._cuda_array_ptr(first_target) is not None: - self._copy_grouped_to_cuda_dst(dst, chunks) - return - - for target_offset, data, source_offset, nbytes in self._coalesce_copy_chunks( - chunks - ): - target = self._slice_dst(dst, target_offset, nbytes) - if hasattr(target, "set"): - import numpy - - host = numpy.frombuffer( - data.view()[source_offset : source_offset + nbytes], - dtype=numpy.uint8, - count=nbytes, - ) - target.set(host) - continue - dst_view = memoryview(dst).cast("B") - dst_view[target_offset : target_offset + nbytes] = data.view()[ - source_offset : source_offset + nbytes - ] - - def _copy_grouped_to_cuda_dst( - self, - dst: Any, - chunks: list[tuple[int, PinnedMemorySlice, int, int]], - ) -> None: - """Copy grouped pinned ranges into a CUDA destination.""" - from cupy.cuda import runtime - - ordered = sorted(chunks, key=lambda item: item[0]) - merged: list[tuple[int, int, int]] = [] - for target_offset, data, source_offset, nbytes in ordered: - source_ptr = data.ptr_at(source_offset) - if not merged: - merged.append((target_offset, source_ptr, nbytes)) - continue - prev_target, prev_source, prev_nbytes = merged[-1] - if ( - target_offset == prev_target + prev_nbytes - and source_ptr == prev_source + prev_nbytes - ): - merged[-1] = (prev_target, prev_source, prev_nbytes + nbytes) - continue - merged.append((target_offset, source_ptr, nbytes)) - - for target_offset, source_ptr, nbytes in merged: - target = self._slice_dst(dst, target_offset, nbytes) - dst_ptr = self._cuda_array_ptr(target) - if dst_ptr is None: - raise TypeError("grouped CUDA copy target lost CUDA array interface") - runtime.memcpyAsync( - dst_ptr, - source_ptr, - nbytes, - runtime.memcpyHostToDevice, - 0, - ) - - def _coalesce_copy_chunks( - self, - chunks: list[tuple[int, PinnedMemorySlice, int, int]], - ) -> list[tuple[int, PinnedMemorySlice, int, int]]: - """Merge adjacent L1-hit copies with contiguous source and target.""" - ordered = sorted(chunks, key=lambda item: item[0]) - merged: list[tuple[int, PinnedMemorySlice, int, int]] = [] - for target_offset, data, source_offset, nbytes in ordered: - if not merged: - merged.append((target_offset, data, source_offset, nbytes)) - continue - prev_target, prev_data, prev_source, prev_nbytes = merged[-1] - if ( - prev_data is data - and target_offset == prev_target + prev_nbytes - and source_offset == prev_source + prev_nbytes - ): - merged[-1] = ( - prev_target, - prev_data, - prev_source, - prev_nbytes + nbytes, - ) - continue - merged.append((target_offset, data, source_offset, nbytes)) - return merged - - def _slice_dst(self, dst: Any, offset: int, nbytes: int) -> Any: - """Return a writable destination slice.""" - if hasattr(dst, "set"): - try: - return dst[offset : offset + nbytes] - except (TypeError, KeyError, IndexError): - if offset == 0: - return dst - raise - if isinstance(dst, bytearray | memoryview): - return memoryview(dst).cast("B")[offset : offset + nbytes] - try: - return dst[offset : offset + nbytes] - except (TypeError, KeyError, IndexError): - pass - return memoryview(dst).cast("B")[offset : offset + nbytes] + copy_ops.copy_grouped_to_dst(dst, chunks) def _slice_src(self, src: Any, offset: int, nbytes: int) -> Any: """Return a readable source slice.""" - if hasattr(src, "get"): - return src[offset : offset + nbytes] - try: - return src[offset : offset + nbytes] - except (TypeError, KeyError, IndexError): - pass - return memoryview(src).cast("B")[offset : offset + nbytes] - - def _cuda_array_ptr(self, dst: Any) -> int | None: - """Return a CUDA device pointer for a CuPy-like array destination.""" - data = getattr(dst, "data", None) - ptr = getattr(data, "ptr", None) - if ptr is None: - return None - return int(ptr) + return copy_ops.slice_src(src, offset, nbytes) def _copy_src_to_pinned( self, @@ -1030,24 +623,8 @@ def _copy_src_to_pinned( pinned: PinnedMemorySlice, nbytes: int, ) -> None: - """Copy bytes from a CPU or CuPy source into pinned host memory. - - Args: - src: readable byte buffer or CuPy ndarray. - pinned: destination slice leased from the L1 pool. - nbytes: number of bytes to copy. - """ - if hasattr(src, "data") and getattr(src.data, "ptr", None) is not None: - from cupy.cuda import runtime - - runtime.memcpy( - pinned.ptr_at(0), - int(src.data.ptr), - nbytes, - runtime.memcpyDeviceToHost, - ) - return - pinned.view()[:nbytes] = memoryview(src).cast("B")[:nbytes] + """Copy bytes from a CPU or CuPy source into pinned host memory.""" + copy_ops.copy_src_to_pinned(src, pinned, 0, nbytes) def _copy_src_to_pinned_at( self, @@ -1057,19 +634,7 @@ def _copy_src_to_pinned_at( nbytes: int, ) -> None: """Copy bytes from a CPU or CUDA source into pinned host memory.""" - if hasattr(src, "data") and getattr(src.data, "ptr", None) is not None: - from cupy.cuda import runtime - - runtime.memcpy( - pinned.ptr_at(target_offset), - int(src.data.ptr), - nbytes, - runtime.memcpyDeviceToHost, - ) - return - pinned.view()[target_offset : target_offset + nbytes] = memoryview(src).cast( - "B" - )[:nbytes] + copy_ops.copy_src_to_pinned(src, pinned, target_offset, nbytes) async def _store_bytes_grouped_l1_only( self, @@ -1086,7 +651,7 @@ async def _store_bytes_grouped_l1_only( file_offset = int(span["file_offset"]) self._check_range(file_offset, nbytes) key = (file_offset, nbytes) - hit = self._find_l1_locked(file_offset, nbytes) + hit = self._l1.find(file_offset) source = self._slice_src(src, source_offset, nbytes) if hit is not None: hit_key, cached, target_offset = hit @@ -1094,11 +659,10 @@ async def _store_bytes_grouped_l1_only( self._copy_src_to_pinned_at( source, cached, target_offset, nbytes ) - self._policy.access(hit_key) - self._l1.move_to_end(hit_key) + self._l1.touch(hit_key) total += nbytes continue - data = self._reserve_l1_buffer_locked_or_raise( + data = self._l1.reserve_or_raise( key, nbytes, preserve_overlaps=True, @@ -1108,7 +672,7 @@ async def _store_bytes_grouped_l1_only( except BaseException: data.close() raise - self._put_l1_locked(key, data) + self._l1.put(key, data) total += nbytes return total @@ -1134,110 +698,15 @@ async def _reserve_l1_buffer( while True: async with self._lock: self._raise_l2_error_locked() - data, wait_for = self._reserve_l1_buffer_locked(key, nbytes) + data = self._l1.reserve(key, nbytes) if data is not None: return data + # Pool exhausted with no evictable victim: an in-flight L2 + # write still pins a slice. Wait for one to finish, then retry. + wait_for: asyncio.Future[None] | None = next( + iter(self._pending_l2.values()), None + ) if wait_for is None: wait_for = asyncio.get_event_loop().create_future() - self._pool_waiters.append(wait_for) - if wait_for is None: - raise MemoryError( - f"could not reserve {nbytes} pinned L1 bytes from " - f"{self._l1_bytes} byte pool" - ) + self._l1.register_pool_waiter(wait_for) await wait_for - - def _reserve_l1_buffer_locked( - self, - key: tuple[int, int], - nbytes: int, - *, - drop_overlaps: bool = True, - preserve_overlaps: bool = False, - ) -> tuple[PinnedMemorySlice | None, asyncio.Task[None] | None]: - """Try to reserve pinned L1 space for a new store or promoted load.""" - if nbytes > self._l1_bytes: - raise ValueError( - f"range {nbytes} bytes exceeds L1 capacity {self._l1_bytes}" - ) - if drop_overlaps: - self._drop_overlapping_l1_locked( - key[0], - key[1], - preserve_remainder=preserve_overlaps, - ) - data = self._pool.allocate(nbytes) - while data is None: - victim = self._policy.evict() - if victim is None: - pending = next(iter(self._pending_l2.values()), None) - return None, pending - removed = self._l1.pop(victim, None) - self._remove_l1_index_locked(victim) - if removed is not None: - self._l1_used -= len(removed) - self._release_l1_buffer_locked(victim, removed) - data = self._pool.allocate(nbytes) - return data, None - - def _reserve_l1_buffer_locked_or_raise( - self, - key: tuple[int, int], - nbytes: int, - *, - preserve_overlaps: bool = False, - ) -> PinnedMemorySlice: - """Reserve pinned L1 space when no pending L2 writer can block reuse.""" - data, wait_for = self._reserve_l1_buffer_locked( - key, - nbytes, - preserve_overlaps=preserve_overlaps, - ) - if data is not None: - return data - if wait_for is not None: - raise RuntimeError("unexpected pending L2 write in skip_l2 mode") - raise MemoryError( - f"could not reserve {nbytes} pinned L1 bytes from " - f"{self._l1_bytes} byte pool" - ) - - def _release_l1_buffer_locked( - self, - key: tuple[int, int], - data: PinnedMemorySlice, - ) -> None: - """Close an L1 buffer unless an L2 write still owns it.""" - if self._pending_l2_buffers.get(key) is data: - return - data.close() - self._notify_pool_waiters_locked() - - def _insert_l1_index_locked(self, key: tuple[int, int]) -> None: - """Add one L1 range to the start-offset lookup index.""" - start = key[0] - existing = self._l1_by_start.get(start) - if existing == key: - return - if existing is not None: - self._remove_l1_index_locked(existing) - bisect.insort(self._l1_starts, start) - self._l1_by_start[start] = key - - def _remove_l1_index_locked(self, key: tuple[int, int]) -> None: - """Remove one L1 range from the start-offset lookup index.""" - start = key[0] - if self._l1_by_start.get(start) != key: - return - del self._l1_by_start[start] - idx = bisect.bisect_left(self._l1_starts, start) - if idx < len(self._l1_starts) and self._l1_starts[idx] == start: - self._l1_starts.pop(idx) - - def _notify_pool_waiters_locked(self) -> None: - """Wake tasks waiting for L1 pool metadata or free-space changes.""" - waiters = self._pool_waiters - self._pool_waiters = [] - for waiter in waiters: - if not waiter.done(): - waiter.set_result(None) diff --git a/daser/transfer/iouring/native.py b/daser/transfer/iouring/native.py index bc02ea9..e2d4539 100644 --- a/daser/transfer/iouring/native.py +++ b/daser/transfer/iouring/native.py @@ -175,38 +175,6 @@ def __init__(self, entries: int = 8) -> None: offset=_IORING_OFF_SQES, ) - def read(self, fd: int, file_offset: int, nbytes: int) -> bytes: - """Read bytes at a file offset through io_uring. - - Args: - fd: Open file descriptor. - file_offset: Byte offset in the file. - nbytes: Number of bytes to read. - - Returns: - Bytes read from the file. - - Thread-safety: - Serialized by the wrapper lock. The call blocks the caller until - the io_uring completion arrives. - """ - if nbytes == 0: - return b"" - buf = bytearray(nbytes) - view = memoryview(buf) - cursor = 0 - while cursor < nbytes: - chunk = min(_MAX_RW_COUNT, nbytes - cursor) - self._submit_and_wait( - _IORING_OP_READ, - fd, - file_offset + cursor, - view[cursor : cursor + chunk], - chunk, - ) - cursor += chunk - return bytes(buf) - def read_into(self, fd: int, file_offset: int, dst: memoryview) -> int: """Read bytes at a file offset into a writable buffer. diff --git a/daser/transfer/iouring/pinned_pool.py b/daser/transfer/iouring/pinned_pool.py index 94809d5..8987cd7 100644 --- a/daser/transfer/iouring/pinned_pool.py +++ b/daser/transfer/iouring/pinned_pool.py @@ -31,24 +31,6 @@ def __init__(self, size: int) -> None: self._memory = cupy.cuda.alloc_pinned_memory(size if size > 0 else 1) self._closed = False - @classmethod - def from_bytes(cls, data: bytes | bytearray | memoryview) -> "PinnedMemoryBuffer": - """Allocate a pinned buffer and initialize it from bytes. - - Args: - data: Source byte data. - - Returns: - Pinned memory buffer containing a copy of ``data``. - - Thread-safety: - Allocates and initializes a new independent buffer. - """ - view = memoryview(data).cast("B") - buf = cls(len(view)) - buf.view()[: len(view)] = view - return buf - def view(self) -> memoryview: """Return a byte-addressable memoryview for the pinned buffer. @@ -78,17 +60,6 @@ def ptr_at(self, offset: int = 0) -> int: raise ValueError("offset is outside pinned buffer") return int(self._memory.ptr) + offset - def to_bytes(self) -> bytes: - """Return a bytes copy of the pinned buffer contents. - - Returns: - Immutable bytes containing the buffer contents. - - Thread-safety: - Reads the current contents without internal locking. - """ - return bytes(self.view()) - def close(self) -> None: """Release the CUDA pinned pages. diff --git a/docs/design/components.md b/docs/design/components.md index 4163d72..00a9824 100644 --- a/docs/design/components.md +++ b/docs/design/components.md @@ -8,16 +8,18 @@ | `SchedulerConnectorMixin` | vLLM scheduler | `daser/connector/scheduler.py`;负责 lookup、pending load/store 跟踪、slot 分配和 connector metadata 构造 | | `WorkerConnectorMixin` | vLLM worker | `daser/connector/worker.py`;负责 KV cache 注册、CUDA IPC handle 导出、后台 IPC loop | | `CudaStagingPool` | vLLM worker | `daser/connector/staging.py`;负责 GDS 和 iouring 共享的 bounded slot-major GPU staging 复用 | -| `IPCClientSync` | vLLM scheduler | 阻塞式 Unix socket 客户端,用于 `get_runtime_config`、`match_and_alloc`、`alloc_chunk` | -| `IPCClientAsync` | vLLM worker | asyncio Unix socket 客户端,用于 `transfer_store`、`transfer_load`、`commit_chunk` | +| `IPCClientSync` | vLLM scheduler | 阻塞式 Unix socket 客户端,用于 `get_runtime_config`、`lookup`、`alloc_chunk` | +| `IPCClientAsync` | vLLM worker | asyncio Unix socket 客户端,用于 `transfer_store`、`transfer_load`、`commit_chunks` | | `TransferLayer` | DaseR | `daser/transfer/base.py`;server-owned KV 数据传输抽象,由 `IPCServer` 按 runtime config 初始化 | | `GDSTransferLayer` | DaseR | `daser/transfer/gds/`;封装 kvikio cuFile / compat IO;backend 在初始化时选定,运行期不可切换 | -| `TieredIOUringTransferLayer` | DaseR | `daser/transfer/iouring/`;L1 pinned-memory + L2 SSD transfer,L1 使用 LRU replacement;`--skip-l2` 时禁用 L2 store/io_uring 路径 | +| `TieredIOUringTransferLayer` | DaseR | `daser/transfer/iouring/layer.py`;编排 L1 pinned-memory + L2 SSD transfer。组合 `L1Cache` + 可选 `L2IoEngine`,拷贝逻辑在 `copy_ops`;`--skip-l2` 即不构造 `L2IoEngine` | +| `L1Cache` / `L2IoEngine` | DaseR | `daser/transfer/iouring/l1_cache.py`(range-keyed pinned-host LRU 缓存)/ `l2_engine.py`(io_uring positioned I/O);编排层组合二者,pinned-predicate 单向解耦 | | `ReplacementPolicy` | DaseR | `daser/replacement/`;通用替换策略抽象,当前实现为 `LRUReplacementPolicy` | | `python -m daser.server` | DaseR | CLI 入口;解析配置,构造 `ServerCore`,启动 HTTP server 和 IPC server,关机保存 index | | `HTTP server` | DaseR | `daser/server/http/`;FastAPI routes、tokenize/chunk、vLLM HTTP 调用、文档 API 和 `/infer` | -| `IPCServer` | DaseR | `daser/server/ipc/server.py`;Unix socket lifecycle、msgpack framing、connector op dispatch | +| `IPCServer` | DaseR | `daser/server/ipc/server.py`;Unix socket lifecycle、msgpack framing、op→handler dispatch table | | `ServerCore` | DaseR | 共享控制面核心;管理 chunk、doc、retrieval、position 状态 | +| `ChunkLifecycle` | DaseR | `daser/server/chunk_lifecycle.py`;chunk commit/写者归属/淘汰状态 + commit waiter,由 `ServerCore` 持有 | | `ChunkManager` | DaseR | ring buffer slot 分配、淘汰、引用计数和持久化 | | `MetadataStore` | DaseR | `chunk_index` 和 `slot_map` 内存状态 | | `DocRegistry` | DaseR | `doc_id -> DocEntry` 文档状态 | @@ -32,7 +34,8 @@ KV bytes 的 load/store,而是把请求分发给更窄的组件:slot 分配交给 `ChunkManager`,metadata 查询和持久化状态交给 `MetadataStore`,文档生命周期 交给 `DocRegistry`,cache lookup 交给 `RetrievalIndex`,position offset 交给 -`PositionEncoder`。 +`PositionEncoder`。chunk 的 commit/写者归属/淘汰状态和 commit waiter 由 +`ChunkLifecycle` 集中维护,保证这几个并行集合的状态转换一致。 `ChunkManager` 和 `MetadataStore` 刻意保持分离。`MetadataStore` 是纯状态容器, 保存 `chunk_key -> ChunkMeta` 的 `chunk_index` 和 `slot_id -> SlotEntry` 的 @@ -118,17 +121,23 @@ layer_offset(slot_i, layer_idx) = slot_i * slot_size + layer_idx * layer_size ```python class RetrievalIndex(ABC): async def lookup(self, tokens: list[int], model_id: str) -> list[RetrievalMatch]: ... - async def insert(self, meta: ChunkMeta) -> None: ... - async def remove(self, chunk_key: str) -> None: ... + async def insert(self, meta: ChunkMeta) -> None: ... # base: primary _index + async def remove(self, chunk_key: str) -> None: ... # base: primary _index + def _on_insert(self, meta: ChunkMeta) -> None: ... # hook for secondary index + def _on_remove(self, chunk_key, meta) -> None: ... # hook for secondary index ``` +base 持有主 `_index`(按 `chunk_key`)并实现 `insert`/`remove`,子类只覆盖 +`lookup` 这个真正分叉的方法,需要维护二级结构时覆盖 `_on_insert`/`_on_remove` +钩子。 + 实现: - `PrefixHashIndex`:用 `h_i = H(h_{i-1}, block_tokens_i)` 计算 rolling prefix key,每个命中对应一个 KV slot,并返回从 prompt 开头连续命中的 - slot 列表。 + slot 列表。只用主 `_index`。 - `ChunkReuseIndex`:返回多个 block-aligned chunk 命中,用于文档 chunk - 复用。 + 复用;用 `_on_insert`/`_on_remove` 维护按 token-count 分桶的二级索引。 ### PositionEncoder @@ -150,23 +159,40 @@ position 组件。`ServerCore` 只依赖 ABC。 ```python class TransferLayer(ABC): - async def load_bytes(self, dst: Any, file_offset: int, nbytes: int) -> int: ... - async def store_bytes(self, src: Any, file_offset: int, nbytes: int) -> int: ... - def close(self) -> None: ... + coalesce_store_spans: bool = False # backend opts into span coalescing + + @property + def stats(self) -> TransferStats: ... # tiering counters (zeroed by default) + @property + def l1_bytes_used(self) -> int: ... # resident L1 bytes (0 by default) + + async def load_bytes(self, dst, file_offset, nbytes) -> int: ... # abstract + async def store_bytes(self, src, file_offset, nbytes) -> int: ... # abstract + async def load_bytes_grouped(self, dst, spans) -> int: ... # default: loop + async def store_bytes_grouped(self, src, spans) -> int: ... # default: loop + async def drain(self) -> None: ... # default no-op; override for write-back + def close(self) -> None: ... # abstract ``` +可选能力(`coalesce_store_spans`、`stats`、`l1_bytes_used`、`drain`、grouped +方法)都在 ABC 上声明并带默认实现,`IPCServer` 直接按接口调用,不做 `getattr` +探测,也不读后端私有属性。 + 实现: - `GDSTransferLayer`:server 通过 CUDA IPC 打开 worker staging buffer, - 再用 kvikio/cuFile 在 GPU buffer 和 SSD file 之间直接传输。 -- `TieredIOUringTransferLayer`:server 通过 CUDA IPC 打开 worker staging - buffer,把 bytes 放入预分配的 pinned host L1 pool,随后异步写入 L2 SSD; - load 时先查 L1,miss 再从 L2 读入并 promote 到 L1。L2 文件使用 - `O_DIRECT` 打开,所有 L2 offset 和 byte count 都要求 4096-byte 对齐。 - Pinned L1 pool 在 transfer 初始化时一次性分配,后续 store/load 热路径只 - lease pool slice。`--skip-l2` 时使用同一个 transfer 实现和同一套 - pinned-memory copy path,但不打开 SSD 文件,load miss 直接报错,因为没有 L2 - 文件可读回。该模式和 GDS 冲突。 + 再用 kvikio/cuFile 在 GPU buffer 和 SSD file 之间直接传输。无 L1 tier, + `stats`/`l1_bytes_used` 取 ABC 默认零值,`drain` 为 no-op。 +- `TieredIOUringTransferLayer`:编排层,组合两个内聚组件—— + `L1Cache`(`l1_cache.py`,range-keyed pinned-host LRU 缓存)和可选的 + `L2IoEngine`(`l2_engine.py`,io_uring positioned I/O);纯拷贝/marshalling + 逻辑在 `copy_ops.py`。store 把 bytes 放入 L1 pool 立即对 load 可见,再异步 + 写入 L2 SSD;load 先查 L1,miss 再经 L2 读入并 promote 回 L1。L2 文件用 + `O_DIRECT` 打开,offset/byte count 要求 4096-byte 对齐。L1 与 L2 通过注入的 + pinned-predicate 单向解耦(在途 L2 写 pin 住的 L1 slice 不被淘汰 close), + 写回调度胶水留在编排层。`--skip-l2` 即「不构造 `L2IoEngine`」 + (`self._l2 is None`):只写 L1、load miss 直接报错、不打开 SSD 文件。该模式和 + GDS 冲突。 connector 不感知具体 transfer 实现,只发送 `transfer_store` / `transfer_load` IPC 请求和 CUDA IPC handle。 diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index 5a7ab18..8a6487a 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -140,7 +140,6 @@ def __init__(self, ipc_client) -> None: self._pending_loads = {} self._pending_alloc = {} self._ipc_sync = ipc_client - self._vllm_prefix_caching_enabled = False self.refresh_count = 0 def _refresh_runtime_config(self) -> None: @@ -1009,7 +1008,6 @@ class DummyRequest: ipc_client = DummyIPCClient() connector = _SchedulerProbe(ipc_client) connector.mark_runtime_ready("served-model") - connector._vllm_prefix_caching_enabled = True # noqa: SLF001 assert connector.get_num_new_matched_tokens( DummyRequest(), diff --git a/tests/server/test_chunker.py b/tests/server/test_chunker.py index 412f146..6ce9e1d 100644 --- a/tests/server/test_chunker.py +++ b/tests/server/test_chunker.py @@ -4,11 +4,11 @@ from daser.server.http.chunker import Chunker, hash_tokens -def test_pad_to_chunk_boundary_extends_tail_without_mutating_input() -> None: +def test_pad_to_block_boundary_extends_tail_without_mutating_input() -> None: chunker = Chunker(block_tokens=4) tokens = [1, 2, 3, 4, 5] - padded = chunker.pad_to_chunk_boundary(tokens, pad_token=0) + padded = chunker.pad_to_block_boundary(tokens, pad_token=0) assert tokens == [1, 2, 3, 4, 5] assert padded == [1, 2, 3, 4, 5, 0, 0, 0] diff --git a/tests/server/test_http_server.py b/tests/server/test_http_server.py index 13b436f..f96e917 100644 --- a/tests/server/test_http_server.py +++ b/tests/server/test_http_server.py @@ -169,8 +169,6 @@ def _make_client( block_tokens=4, system_prompt="S:", doc_separator="|", - task_separator="? ", - answer_separator="! ", ), core, tokenizer=FakeTokenizer(), @@ -377,8 +375,6 @@ def test_chunk_reuse_lifespan_prewarms_fixed_segments() -> None: block_tokens=BLOCK_TOKENS, system_prompt="S:", doc_separator="|", - task_separator="? ", - answer_separator="! ", align_document_chunks=True, ), core, @@ -435,8 +431,6 @@ def test_chunk_reuse_lifespan_skips_restored_fixed_segments() -> None: block_tokens=BLOCK_TOKENS, system_prompt="S:", doc_separator="|", - task_separator="? ", - answer_separator="! ", align_document_chunks=True, ), core, @@ -495,8 +489,6 @@ def test_upload_document_rejects_prefix_mode() -> None: block_tokens=BLOCK_TOKENS, system_prompt="S:", doc_separator="|", - task_separator="? ", - answer_separator="! ", cache_reuse_mode="prefix", align_document_chunks=False, ), @@ -528,8 +520,6 @@ def test_upload_document_waits_for_committed_store() -> None: block_tokens=BLOCK_TOKENS, system_prompt="S:", doc_separator="|", - task_separator="? ", - answer_separator="! ", ), core, tokenizer=FakeTokenizer(), @@ -632,8 +622,6 @@ def test_infer_chat_template_disables_thinking_before_tokenization() -> None: block_tokens=4, system_prompt="S:", doc_separator="|", - task_separator="? ", - answer_separator="! ", ), core, tokenizer=tokenizer, @@ -791,8 +779,6 @@ def test_chunk_reuse_infer_uses_contiguous_prewarmed_padded_segments() -> None: block_tokens=BLOCK_TOKENS, system_prompt="S:", doc_separator="|", - task_separator="? ", - answer_separator="! ", align_document_chunks=True, ), core, @@ -872,8 +858,6 @@ class PadTokenizer(FakeTokenizer): block_tokens=BLOCK_TOKENS, system_prompt="S:", doc_separator="|", - task_separator="? ", - answer_separator="! ", align_document_chunks=True, ), core, @@ -899,8 +883,6 @@ def test_chunk_reuse_uses_one_block_aligned_chunk_per_prompt_segment() -> None: block_tokens=BLOCK_TOKENS, system_prompt="S:", doc_separator="|", - task_separator="? ", - answer_separator="! ", align_document_chunks=True, ), core, @@ -964,8 +946,6 @@ def test_chunk_reuse_repeated_separator_keeps_hits_contiguous_before_suffix() -> block_tokens=BLOCK_TOKENS, system_prompt="S:", doc_separator="|", - task_separator="? ", - answer_separator="! ", align_document_chunks=True, ), core, diff --git a/tests/server/test_ipc_server.py b/tests/server/test_ipc_server.py index 7c8380b..907528b 100644 --- a/tests/server/test_ipc_server.py +++ b/tests/server/test_ipc_server.py @@ -19,6 +19,7 @@ from daser.server.doc_registry import DocRegistry from daser.server.ipc import IPCServer from daser.server.metadata_store import MetadataStore +from daser.transfer.base import TransferLayer SLOT_SIZE = 4096 BLOCK_TOKENS = 4 @@ -495,7 +496,13 @@ def fake_open_cuda_ipc_buffer(**kwargs: Any) -> FakeOpened: opened_buffers.append(opened) return opened - class FakeTransfer: + class FakeTransfer(TransferLayer): + async def load_bytes(self, dst: Any, file_offset: int, nbytes: int) -> int: + return 0 + + async def store_bytes(self, src: Any, file_offset: int, nbytes: int) -> int: + return 0 + async def load_bytes_grouped( self, _dst: Any, @@ -503,6 +510,9 @@ async def load_bytes_grouped( ) -> int: return 0 + def close(self) -> None: + pass + def fake_ensure_transfer(_server: IPCServer) -> FakeTransfer: return FakeTransfer() @@ -562,10 +572,16 @@ async def test_stop_accepting_closes_listener_before_transfer( ) -> None: events: list[str] = [] - class FakeTransfer: + class FakeTransfer(TransferLayer): def __init__(self, **_kwargs: Any) -> None: events.append("ensure_transfer") + async def load_bytes(self, dst: Any, file_offset: int, nbytes: int) -> int: + return nbytes + + async def store_bytes(self, src: Any, file_offset: int, nbytes: int) -> int: + return nbytes + async def drain(self) -> None: events.append("drain") @@ -604,7 +620,7 @@ async def test_eager_transfer_initialization_avoids_lazy_init_on_first_request( """Eagerly initializing the transfer layer creates it before any request.""" init_events: list[str] = [] - class FakeTransfer: + class FakeTransfer(TransferLayer): def __init__(self, **_kwargs: Any) -> None: init_events.append("transfer_created") @@ -657,12 +673,24 @@ async def test_skip_l2_selects_iouring_transfer_without_store_path( """skip_l2 should wire IPC transfer ops to iouring with its L2 tier disabled.""" init_kwargs: list[dict[str, Any]] = [] - class FakeTieredIOUringTransfer: + class FakeTieredIOUringTransfer(TransferLayer): coalesce_store_spans = True def __init__(self, **kwargs: Any) -> None: init_kwargs.append(kwargs) + async def load_bytes(self, dst: Any, file_offset: int, nbytes: int) -> int: + return await self.load_bytes_grouped( + dst, + [{"target_offset": 0, "file_offset": file_offset, "nbytes": nbytes}], + ) + + async def store_bytes(self, src: Any, file_offset: int, nbytes: int) -> int: + return await self.store_bytes_grouped( + src, + [{"source_offset": 0, "file_offset": file_offset, "nbytes": nbytes}], + ) + async def store_bytes_grouped( self, src: Any, diff --git a/tests/server/test_vllm_client.py b/tests/server/test_vllm_client.py index 8529f98..6c16fb2 100644 --- a/tests/server/test_vllm_client.py +++ b/tests/server/test_vllm_client.py @@ -70,56 +70,6 @@ async def aiter_lines(self): return _Resp() -def _make_client() -> tuple[VLLMClient, _CapturingClient]: - vllm = VLLMClient(base_url="http://localhost:8000", model="dummy") - fake = _CapturingClient() - vllm._client = fake # noqa: SLF001 — test-only injection - return vllm, fake - - -@pytest.mark.asyncio -async def test_completion_omits_kv_transfer_params_by_default() -> None: - vllm, fake = _make_client() - - await vllm.completion([1, 2, 3]) - - assert len(fake.posts) == 1 - _, body = fake.posts[0] - assert "kv_transfer_params" not in body, ( - "Default completion call must not include kv_transfer_params " - "so existing vLLM endpoints see the same request shape." - ) - - -@pytest.mark.asyncio -async def test_completion_forwards_kv_transfer_params() -> None: - vllm, fake = _make_client() - - await vllm.completion( - [1, 2, 3], - kv_transfer_params={"daser_skip_save": True}, - ) - - _, body = fake.posts[0] - assert body.get("kv_transfer_params") == {"daser_skip_save": True} - - -@pytest.mark.asyncio -async def test_completion_merges_gen_params_and_kv_transfer_params() -> None: - vllm, fake = _make_client() - - await vllm.completion( - [1, 2, 3], - gen_params={"max_tokens": 8, "temperature": 0.1}, - kv_transfer_params={"daser_skip_save": True}, - ) - - _, body = fake.posts[0] - assert body["max_tokens"] == 8 - assert body["temperature"] == pytest.approx(0.1) - assert body["kv_transfer_params"] == {"daser_skip_save": True} - - @pytest.mark.asyncio async def test_completion_with_ttft_records_first_token(monkeypatch) -> None: vllm = VLLMClient(base_url="http://localhost:8000", model="dummy") diff --git a/tests/test_config.py b/tests/test_config.py index af8bc57..b792e88 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -37,7 +37,10 @@ def test_model_geometry_from_path_reads_hf_config(tmp_path: Path) -> None: assert geometry.head_dim == 128 assert geometry.num_layers == 28 assert geometry.dtype_bytes == 2 - assert geometry.slot_size == 4 * 128 * 2 * 28 * BLOCK_TOKENS * 2 + assert ( + geometry.slot_size_for_block_tokens(BLOCK_TOKENS) + == 4 * 128 * 2 * 28 * BLOCK_TOKENS * 2 + ) def test_model_geometry_uses_explicit_head_dim_and_float32_dtype( @@ -60,7 +63,10 @@ def test_model_geometry_uses_explicit_head_dim_and_float32_dtype( assert geometry.head_dim == 64 assert geometry.dtype_bytes == 4 - assert geometry.slot_size == 8 * 64 * 2 * 2 * BLOCK_TOKENS * 4 + assert ( + geometry.slot_size_for_block_tokens(BLOCK_TOKENS) + == 8 * 64 * 2 * 2 * BLOCK_TOKENS * 4 + ) def test_daser_config_derives_paths_and_slot_size(tmp_path: Path) -> None: diff --git a/tests/transfer/test_tiered_iouring_transfer.py b/tests/transfer/test_tiered_iouring_transfer.py index 6dd0345..812d0fe 100644 --- a/tests/transfer/test_tiered_iouring_transfer.py +++ b/tests/transfer/test_tiered_iouring_transfer.py @@ -192,7 +192,9 @@ def test_native_iouring_splits_large_positioned_io(tmp_path, monkeypatch) -> Non uring = NativeIOUring(entries=8) try: assert uring.write(fd, 1, b"abcdefghijkl") == 12 - assert uring.read(fd, 1, 12) == b"abcdefghijkl" + dst = bytearray(12) + assert uring.read_into(fd, 1, memoryview(dst)) == 12 + assert bytes(dst) == b"abcdefghijkl" assert path.read_bytes()[1:13] == b"abcdefghijkl" finally: uring.close() diff --git a/tests/unit/test_benchmark_unified_utils.py b/tests/unit/test_benchmark_unified_utils.py index a0965f3..6a50096 100644 --- a/tests/unit/test_benchmark_unified_utils.py +++ b/tests/unit/test_benchmark_unified_utils.py @@ -24,20 +24,16 @@ from benchmarks.run_bench import ( BackendRun, RunBenchArgs, - _bench_prepare_config, - _collect_vllm_bench_phase_metrics, - _compare_vllm_bench_outputs, _expand_backend_runs, - _normalise_vllm_bench_result, _probe_daser_metrics, _run_command, _should_probe_daser_metrics, _stage_title, _validate_backend_runs, - _vllm_bench_command, parse_args, run_benchmark, ) +from benchmarks.utils import vllm_bench # First Party from benchmarks.utils.constants import BYTES_PER_GIB, COMPARISON_IOURING_MEM @@ -82,6 +78,12 @@ format_capacity, parse_size_bytes, ) +from benchmarks.utils.vllm_bench import ( + _bench_command, + _collect_phase_metrics, + _compare_outputs, + _normalise_result, +) class _Tokenizer: @@ -140,6 +142,50 @@ def test_imdb_dataset_uses_review_as_context(tmp_path: Path) -> None: assert samples[0].question == "Summarize the sentiment of this review." +def test_dataset_registry_dispatches_by_name(tmp_path: Path) -> None: + """The registry builds the selected dataset from CLI args with zero + runner-side branching.""" + # First Party + import argparse + + from benchmarks.utils.datasets import ( + add_dataset_cli_args, + build_dataset, + dataset_names, + ) + + assert "imdb" in dataset_names() + assert "longbench" in dataset_names() + + imdb = tmp_path / "imdb.csv" + imdb.write_text('review\n"a good movie"\n') + + parser = argparse.ArgumentParser() + add_dataset_cli_args(parser) + parser.add_argument("--max-samples", type=int, default=20) + + args = parser.parse_args(["--dataset", "imdb", "--imdb", str(imdb)]) + dataset = build_dataset(args) + assert dataset.name == "imdb" + assert dataset.load()[0].context == "a good movie" + + +def test_dataset_registry_reports_missing_required_arg(tmp_path: Path) -> None: + """A dataset selected without its required path raises a clear error.""" + # First Party + import argparse + + from benchmarks.utils.datasets import add_dataset_cli_args, build_dataset + + parser = argparse.ArgumentParser() + add_dataset_cli_args(parser) + parser.add_argument("--max-samples", type=int, default=20) + + args = parser.parse_args(["--dataset", "longbench"]) + with pytest.raises(ValueError, match="--longbench-dir is required"): + build_dataset(args) + + def test_build_full_prompt_replaces_single_documents_marker() -> None: """Prompt builder injects context into the chat-template document slot.""" prompt = build_full_prompt(_Tokenizer(), "doc body", "answer?") @@ -2025,7 +2071,7 @@ def test_vllm_bench_prepare_config_uses_synthetic_lengths(tmp_path: Path) -> Non bench_random_range_ratio=0.5, ) - config = _bench_prepare_config(args, tmp_path) + config = vllm_bench.prepare_config(args, tmp_path) assert config["dataset"] == "vllm-bench-random" assert config["num_samples"] == 10 @@ -2056,7 +2102,7 @@ def test_vllm_bench_command_uses_random_dataset(tmp_path: Path) -> None: ) raw_path = tmp_path / "raw.json" - command = _vllm_bench_command( + command = _bench_command( args, ServiceEndpoint("http://127.0.0.1:8001"), raw_path, @@ -2100,7 +2146,7 @@ def test_vllm_bench_normalises_result_summary(tmp_path: Path) -> None: ) ) - summary = _normalise_vllm_bench_result(raw) + summary = _normalise_result(raw) assert summary["num_requests"] == 10 assert summary["ttft_ms_mean"] == 100.0 @@ -2134,7 +2180,7 @@ def test_vllm_bench_compares_detailed_outputs(tmp_path: Path) -> None: ) ) - correctness = _compare_vllm_bench_outputs(cold, warm) + correctness = _compare_outputs(cold, warm) assert correctness == { "cold_warm_exact_match": { @@ -2158,7 +2204,7 @@ def test_vllm_bench_compares_outputs_detail_shape(tmp_path: Path) -> None: json.dumps({"outputs": [{"generated_text": "A"}, {"generated_text": "X"}]}) ) - correctness = _compare_vllm_bench_outputs(cold, warm) + correctness = _compare_outputs(cold, warm) assert correctness["cold_warm_exact_match"]["available"] is True assert correctness["cold_warm_exact_match"]["matches"] == 1 @@ -2174,7 +2220,7 @@ def test_vllm_bench_correctness_counts_length_mismatch(tmp_path: Path) -> None: cold.write_text(json.dumps({"generated_texts": ["A", "B", "C"]})) warm.write_text(json.dumps({"generated_texts": ["A", "B"]})) - correctness = _compare_vllm_bench_outputs(cold, warm) + correctness = _compare_outputs(cold, warm) assert correctness["cold_warm_exact_match"]["available"] is True assert correctness["cold_warm_exact_match"]["matches"] == 2 @@ -2192,7 +2238,7 @@ def test_vllm_bench_correctness_marks_missing_details_unavailable( cold.write_text(json.dumps({"completed": 2})) warm.write_text(json.dumps({"completed": 2})) - correctness = _compare_vllm_bench_outputs(cold, warm) + correctness = _compare_outputs(cold, warm) assert correctness["cold_warm_exact_match"]["available"] is False assert correctness["cold_warm_exact_match"]["total"] == 0 @@ -2260,11 +2306,11 @@ async def fake_collect( pid_file="/bench/pids.json", ) monkeypatch.setattr( - "benchmarks.run_bench.collect_phase_metrics", + "benchmarks.utils.vllm_bench.collect_phase_metrics", fake_collect, ) - metrics, hit_rate = _collect_vllm_bench_phase_metrics(manifest, before_metrics) + metrics, hit_rate = _collect_phase_metrics(manifest, before_metrics) assert metrics["backend_prometheus"]["daser_cache_requested_tokens_total"] == 2000 assert metrics["backend_prometheus"]["daser_cache_matched_tokens_total"] == 1500 @@ -2344,17 +2390,14 @@ def fake_run_command(command: list[str]) -> None: monkeypatch.setattr("benchmarks.run_bench._run_command", fake_run_command) monkeypatch.setattr("benchmarks.run_bench.stop_from_pid_file", lambda _path: None) - monkeypatch.setattr( - "benchmarks.run_bench._wait_with_message", - lambda *_args, **_kwargs: (_ for _ in ()).throw( - AssertionError("vLLM bench LMCache warm wait must not use fixed sleep") - ), - ) monkeypatch.setattr( "benchmarks.run_bench.time.strftime", lambda _fmt: "20260102_030405", ) - monkeypatch.setattr("benchmarks.run_bench._drain_daser", lambda _manifest: None) + monkeypatch.setattr( + "benchmarks.utils.vllm_bench._drain_daser", + lambda _manifest, **_kwargs: None, + ) monkeypatch.setattr( "benchmarks.run_bench._probe_daser_metrics", lambda *_args, **_kwargs: None, @@ -2367,7 +2410,7 @@ async def fake_wait_lmcache_quiescent( lmcache_waits.append(f"{manifest.backend}:{settle_seconds}") monkeypatch.setattr( - "benchmarks.run_bench._wait_lmcache_quiescent", + "benchmarks.utils.vllm_bench._wait_lmcache_quiescent", fake_wait_lmcache_quiescent, ) @@ -2404,7 +2447,7 @@ async def fake_collect_phase_metrics( } monkeypatch.setattr( - "benchmarks.run_bench.collect_phase_metrics", + "benchmarks.utils.vllm_bench.collect_phase_metrics", fake_collect_phase_metrics, ) diff --git a/tests/unit/test_rope_apply_contract.py b/tests/unit/test_rope_apply_contract.py new file mode 100644 index 0000000..0c74c16 --- /dev/null +++ b/tests/unit/test_rope_apply_contract.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 + +# Standard +from dataclasses import dataclass + +# Third Party +import pytest + +# First Party +from daser.ops import rope_apply + + +@dataclass +class _FakeTensor: + """Minimal stand-in exposing only what the rope shape guard reads. + + The ``rotary_dim`` boundary check runs before any real tensor op or CUDA + dispatch, so the contract is verified without importing a real torch + runtime (the CPU CI installs a torch stub with no tensor ops). + """ + + shape: tuple[int, ...] + + +def test_apply_rope_delta_to_key_block_rejects_rotary_dim_over_head_dim() -> None: + """rotary_dim larger than head_dim is a misconfiguration and must raise. + + Previously this silently returned, masking a configuration error that + would leave KV unrotated. + """ + key_block = _FakeTensor(shape=(4, 2, 8)) + + with pytest.raises(ValueError, match="rotary_dim must not exceed head_dim"): + rope_apply.apply_rope_delta_to_key_block( + key_block, + delta=4, + rope_base=10000.0, + rotary_dim=16, + is_neox_style=True, + ) + + +def test_apply_rope_delta_to_kv_key_block_rejects_rotary_dim_over_head_dim() -> None: + """The KV-block wrapper raises on rotary_dim larger than head_dim.""" + kv_block = _FakeTensor(shape=(1, 2, 2, 4, 2, 8)) + + with pytest.raises(ValueError, match="rotary_dim must not exceed head_dim"): + rope_apply.apply_rope_delta_to_kv_key_block( + kv_block, + delta=4, + rope_base=10000.0, + rotary_dim=16, + is_neox_style=True, + ) + + +def test_apply_rope_delta_noop_returns_without_shape_check() -> None: + """delta==0 and rotary_dim<=0 remain no-ops regardless of head_dim.""" + key_block = _FakeTensor(shape=(4, 2, 8)) + + # delta == 0 short-circuits before the shape guard. + rope_apply.apply_rope_delta_to_key_block( + key_block, + delta=0, + rope_base=10000.0, + rotary_dim=16, + is_neox_style=True, + ) + # rotary_dim <= 0 short-circuits before the shape guard. + rope_apply.apply_rope_delta_to_key_block( + key_block, + delta=4, + rope_base=10000.0, + rotary_dim=0, + is_neox_style=True, + )