diff --git a/benchmarks/README.md b/benchmarks/README.md index c43ff4a..2e9dc2b 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -74,10 +74,60 @@ python benchmarks/run_bench.py \ `--backend all` runs the full comparison matrix sequentially: `baseline`, `lmcache`, `daser-chunk`, and `daser-prefix`. Use `--backend baseline`, `--backend lmcache`, `--backend daser-chunk`, or `--backend daser-prefix` for a -single backend. `--backend daser` still starts DaseR with the mode selected by -`--cache-reuse-mode`, and `--backend vllm` is accepted as a compatibility alias -for `baseline`. `--cache-reuse-mode` is DaseR-only; baseline and LMCache use the -same prompts but record `reuse_mode` as `none` in their manifests. +single backend. Use comma-separated rows such as +`--backend baseline,lmcache,daser-prefix` to run a subset. Baseline and LMCache +use the same prompts but record `reuse_mode` as `none` in their manifests. + +Use synthetic vLLM benchmark traffic when you need to control prompt length +directly instead of loading IMDB or LongBench records. This mode starts the same +services but sends load with `vllm bench serve` against the OpenAI-compatible +`/v1/completions` endpoint: + +```bash +python benchmarks/run_bench.py \ + --backend baseline,lmcache,daser-prefix \ + --load-generator vllm-bench \ + --model /data//model/models/Qwen/Qwen3-8B \ + --store-dir /data//daser_bench/vllmbench_len8192 \ + --gpu-id 2 \ + --gpu-util 0.85 \ + --max-num-seqs 32 \ + --block-size 128 \ + --bench-num-prompts 1000 \ + --bench-input-len 8192 \ + --bench-output-len 1 \ + --bench-request-rate inf \ + --bench-max-concurrency 16 \ + --bench-seed 42 +``` + +`--backend baseline,lmcache,daser-prefix` runs the OpenAI-compatible rows. It +omits `daser-chunk` because DaseR chunk mode uses DaseR-specific `/documents` +and `/infer` endpoints instead of the OpenAI completions endpoint. If +`--load-generator vllm-bench` is combined with a selection that includes +`daser-chunk`, such as `--backend all` or `--backend daser-chunk`, the runner +fails fast and asks you to select only `baseline,lmcache,daser-prefix` or use +`--load-generator internal`. + +The vLLM-bench-specific knobs are: + +| Option | Meaning | +|--------|---------| +| `--bench-num-prompts` | Number of random prompts to send | +| `--bench-input-len` | Random dataset input length passed to `vllm bench serve` | +| `--bench-output-len` | Random dataset output length; defaults to `--gen-max-tokens` | +| `--bench-request-rate` | Requests per second; `inf` sends all requests immediately | +| `--bench-max-concurrency` | Maximum concurrent requests; defaults to `--max-inflight` | +| `--bench-random-prefix-len` | Fixed prefix tokens before random context tokens | +| `--bench-random-range-ratio` | Symmetric random length range ratio | +| `--bench-seed` | Seed reused for cold and warm phases | +| `--bench-burstiness` | Request arrival burstiness passed to `vllm bench serve` | + +For LMCache and DaseR prefix, cold and warm phases run `vllm bench serve` twice +with the same seed and length parameters so the second pass can reuse the first +pass. Raw vLLM bench JSON is saved beside each backend's `results.json` as +`vllm_bench_baseline.json`, `vllm_bench_cold.json`, and +`vllm_bench_warm.json`. Generation defaults are deterministic across backends: `--gen-temperature` defaults to `0.0`, `--gen-top-p` defaults to `1.0`, and `--gen-seed` defaults @@ -255,16 +305,20 @@ cache hit rates from multiple sources: - `vllm_external_prefix_cache_hit_rate`: vLLM Prometheus external prefix cache hit ratio from `vllm:external_prefix_cache_*` counter deltas. - `backend_server_cache_hit_rate`: summary hit ratio used for backend - comparison. DaseR reports its internal - `daser_external_prefix_cache_*` counters, which the connector records with - the same queried-token / accepted-token semantics as vLLM's - `vllm:external_prefix_cache_*` counters. LMCache reports MP server token hit - counters when available and falls back to request counters only when token - counters are absent. DaseR control-plane lookup counters are still kept in raw - metrics as `daser_prometheus_tokens` and `daser_prometheus_requests`. + comparison. DaseR reports token-level + `daser_cache_matched_tokens_total / daser_cache_requested_tokens_total` + counter deltas. LMCache reports MP server token hit counters when available + and falls back to status prefetch token counters. Request-level counters and + external-prefix counters are kept in raw metrics for diagnostics but are not + used as the summary backend hit rate. - `metrics`: raw vLLM Prometheus, backend Prometheus, backend status counter deltas, and all named hit-ratio candidates. +`--load-generator vllm-bench` writes raw vLLM bench JSON files and wraps each +phase with the same backend metric snapshots as the internal load generator, so +its cold/warm summaries include token-level `backend_server_cache_hit_rate` +when the backend exposes the required counters. + For datasets with answer labels, each summary also includes `answer_contains_accuracy`; datasets without labels report `null`. diff --git a/benchmarks/bench_load.py b/benchmarks/bench_load.py index 4211301..1cfaeae 100644 --- a/benchmarks/bench_load.py +++ b/benchmarks/bench_load.py @@ -17,7 +17,7 @@ from benchmarks.utils.constants import ( BLOCK_TOKENS, COMPARISON_IOURING_MEM, - SLOT_SIZE, + slot_size_for_block_tokens, ) from benchmarks.utils.datasets import ( BenchmarkSample, @@ -29,6 +29,7 @@ from benchmarks.utils.loadgen import ( PhaseResult, RequestResult, + backend_server_hit_rate, run_daser_chunk, run_daser_prefix, run_lmcache, @@ -69,6 +70,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--longbench-dir") parser.add_argument("--datasets", default=None) 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) parser.add_argument("--no-dedup-context", action="store_true") parser.add_argument("--max-inflight", type=int, default=32) @@ -91,6 +93,11 @@ async def main_async(args: argparse.Namespace) -> None: if args.prepare_only else BenchmarkManifest.read(_required(args.manifest, "--manifest")) ) + if manifest is not None and args.block_size != manifest.block_size: + raise ValueError( + f"--block-size={args.block_size} does not match manifest " + f"block_size={manifest.block_size}" + ) model = args.model if args.prepare_only else manifest.model store_dir = args.store_dir if args.prepare_only else manifest.store_dir if model is None: @@ -105,7 +112,10 @@ async def main_async(args: argparse.Namespace) -> None: model_config = AutoConfig.from_pretrained(model, trust_remote_code=True) chunk_aligned_prompts = True prompts = build_prompt_payloads( - tokenizer, samples, chunk_aligned=chunk_aligned_prompts + tokenizer, + samples, + chunk_aligned=chunk_aligned_prompts, + block_tokens=args.block_size, ) token_counts = count_prompt_payload_tokens(tokenizer, prompts) effective_max_context_tokens = _effective_max_context_tokens( @@ -120,22 +130,29 @@ async def main_async(args: argparse.Namespace) -> None: if _should_dedup_context(args): samples = dedup_by_context(samples) prompts = build_prompt_payloads( - tokenizer, samples, chunk_aligned=chunk_aligned_prompts + tokenizer, + samples, + chunk_aligned=chunk_aligned_prompts, + block_tokens=args.block_size, ) token_counts = count_prompt_payload_tokens(tokenizer, prompts) samples = interleave_samples(samples) prompts = build_prompt_payloads( - tokenizer, samples, chunk_aligned=chunk_aligned_prompts + tokenizer, + samples, + chunk_aligned=chunk_aligned_prompts, + block_tokens=args.block_size, ) token_counts = count_prompt_payload_tokens(tokenizer, prompts) - total_blocks, max_prompt_blocks = workload_blocks(token_counts, BLOCK_TOKENS) + total_blocks, max_prompt_blocks = workload_blocks(token_counts, args.block_size) + slot_size = slot_size_for_block_tokens(args.block_size) sizing = None if args.prepare_only: capacity_limits = _capacity_limits(args, store_dir) sizing = derive_benchmark_sizing( total_blocks=total_blocks, max_prompt_blocks=max_prompt_blocks, - slot_size=SLOT_SIZE, + slot_size=slot_size, mode=COMPARISON_IOURING_MEM, evict=args.evict, capacity_limits=capacity_limits, @@ -161,6 +178,7 @@ async def main_async(args: argparse.Namespace) -> None: total_prompt_tokens=sum(token_counts), total_blocks=total_blocks, max_prompt_blocks=max_prompt_blocks, + block_size=args.block_size, evict=args.evict, sizing=sizing, ) @@ -391,6 +409,7 @@ def _common_config_for_run( total_blocks: int, max_prompt_blocks: int, evict: bool, + block_size: int, sizing: BenchmarkSizing | None, ) -> dict[str, Any]: """Build benchmark config without re-inferring capacities during load. @@ -406,6 +425,7 @@ def _common_config_for_run( total_prompt_tokens: Total selected prompt tokens. total_blocks: Total selected KV blocks. max_prompt_blocks: Largest selected prompt in KV blocks. + block_size: vLLM KV block size in tokens. evict: Whether eviction sizing was requested. sizing: Prepare-time sizing, required for prepare-only invocations. @@ -429,6 +449,7 @@ def _common_config_for_run( "total_prompt_tokens": total_prompt_tokens, "total_blocks": total_blocks, "max_prompt_blocks": max_prompt_blocks, + "block_size": block_size, } if prepare_only: @@ -442,6 +463,7 @@ def _common_config_for_run( "total_prompt_tokens": total_prompt_tokens, "total_blocks": total_blocks, "max_prompt_blocks": max_prompt_blocks, + "block_size": block_size, "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, @@ -463,6 +485,7 @@ def _common_config_for_run( "total_prompt_tokens": total_prompt_tokens, "total_blocks": total_blocks, "max_prompt_blocks": max_prompt_blocks, + "block_size": block_size, "derived_l1_size_bytes": manifest.l1_size_bytes, "derived_l1_size": format_capacity(manifest.l1_size_bytes), "derived_l2_size_bytes": manifest.l2_size_bytes, @@ -549,22 +572,7 @@ def _serialise_phase( def _backend_server_hit_rate(hit_ratios: dict[str, Any]) -> float | None: - if ( - hit_ratios.get("daser_external_prefix") is not None - or hit_ratios.get("daser_prometheus_tokens") is not None - or hit_ratios.get("daser_prometheus_requests") is not None - ): - ratio = hit_ratios.get("daser_external_prefix") - return float(ratio) if ratio is not None else None - for key in ( - "lmcache_prometheus_lookup", - "lmcache_prometheus_retrieve", - "lmcache_status_prefetch", - ): - ratio = hit_ratios.get(key) - if ratio is not None: - return float(ratio) - return None + return backend_server_hit_rate(hit_ratios) def main(argv: list[str] | None = None) -> None: diff --git a/benchmarks/bench_start_servers.py b/benchmarks/bench_start_servers.py index f9a6eea..7033425 100644 --- a/benchmarks/bench_start_servers.py +++ b/benchmarks/bench_start_servers.py @@ -37,6 +37,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--max-num-seqs", type=int, default=32) parser.add_argument("--max-num-batched-tokens", type=int, default=0) parser.add_argument("--max-model-len", type=int, default=0) + parser.add_argument("--block-size", type=int, default=16) parser.add_argument("--l1-size", type=parse_size_bytes, default="256gib") parser.add_argument("--l2-size", type=parse_size_bytes, default="300gib") parser.add_argument( @@ -73,6 +74,7 @@ async def main_async(args: argparse.Namespace) -> None: max_num_batched_tokens=( args.max_num_batched_tokens if args.max_num_batched_tokens > 0 else None ), + block_size=args.block_size, reuse_mode=args.cache_reuse_mode, transfer_mode=args.transfer_mode, vllm_port=args.vllm_port, diff --git a/benchmarks/run_bench.py b/benchmarks/run_bench.py index dcdf92d..629859b 100644 --- a/benchmarks/run_bench.py +++ b/benchmarks/run_bench.py @@ -4,8 +4,10 @@ from __future__ import annotations import argparse +import asyncio from dataclasses import dataclass import json +import math from pathlib import Path import shlex import subprocess @@ -17,9 +19,22 @@ 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 backend_server_hit_rate, collect_phase_metrics 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") +_BACKEND_SELECTIONS = ("all", *_BACKEND_ROWS) @dataclass(frozen=True) @@ -45,8 +60,8 @@ class RunBenchArgs: """Arguments for the benchmark orchestration entrypoint. Args: - backend: Requested backend row or ``all``. - cache_reuse_mode: Compatibility reuse mode for ``--backend daser``. + backend: Requested backend row, comma-separated row list, or ``all``. + load_generator: Load generator implementation. dataset: Dataset family. model: Model path. store_dir: Parent directory for the generated run root. @@ -58,9 +73,19 @@ class RunBenchArgs: gpu_util: vLLM GPU memory utilization. max_num_seqs: vLLM maximum sequence concurrency. max_num_batched_tokens: Optional vLLM scheduler token budget. + block_size: vLLM KV block size in tokens. max_inflight: HTTP load generator concurrency. gen_max_tokens: Maximum generated tokens. max_context_tokens: Prompt token ceiling; 0 infers from model metadata. + bench_num_prompts: vLLM bench random prompt count. + bench_input_len: vLLM bench random input length. + bench_output_len: vLLM bench random output length. + bench_request_rate: vLLM bench request rate. + bench_max_concurrency: vLLM bench max in-flight requests. + bench_random_prefix_len: Fixed prefix length for vLLM random dataset. + bench_random_range_ratio: vLLM random input/output length range ratio. + bench_seed: vLLM bench random seed. + bench_burstiness: vLLM bench burstiness factor. evict: Whether to enable L2 and eviction sizing. prometheus_url: Optional Prometheus base URL for scrape diagnostics. @@ -69,7 +94,7 @@ class RunBenchArgs: """ backend: str = "all" - cache_reuse_mode: str = "chunk" + load_generator: str = "internal" dataset: str = "longbench" model: str = "" store_dir: str = "" @@ -81,13 +106,38 @@ class RunBenchArgs: gpu_util: float = 0.85 max_num_seqs: int = 32 max_num_batched_tokens: int = 0 + block_size: int = 16 max_inflight: int = 32 gen_max_tokens: int = 128 max_context_tokens: int = 0 + bench_num_prompts: int = 1000 + bench_input_len: int = 1024 + bench_output_len: int | None = None + bench_request_rate: str = "inf" + bench_max_concurrency: int | None = None + bench_random_prefix_len: int = 0 + bench_random_range_ratio: float = 0.0 + bench_seed: int = 42 + bench_burstiness: float = 1.0 evict: bool = False prometheus_url: str = "http://127.0.0.1:9090" +def _backend_selection(value: str) -> str: + """Validate and normalize a benchmark backend selection.""" + value = value.strip() + if value == "all": + return value + names = [name.strip() for name in value.split(",")] + invalid = [name for name in names if name not in _BACKEND_ROWS] + if not names or any(not name for name in names) or invalid: + valid = ", ".join(_BACKEND_SELECTIONS) + raise argparse.ArgumentTypeError( + f"unknown backend selection: {value}; choose from {valid}" + ) + return ",".join(names) + + def parse_args(argv: list[str] | None = None) -> RunBenchArgs: """Parse CLI arguments. @@ -103,19 +153,18 @@ def parse_args(argv: list[str] | None = None) -> RunBenchArgs: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--backend", - choices=( - "all", - "baseline", - "vllm", - "lmcache", - "daser", - "daser-chunk", - "daser-prefix", - ), default="all", + metavar="{all,baseline,lmcache,daser-prefix,daser-chunk}[,...]", + type=_backend_selection, + help=( + "Benchmark rows to run. Use all, one row, or a comma-separated " + "subset such as baseline,lmcache,daser-prefix." + ), ) parser.add_argument( - "--cache-reuse-mode", choices=("chunk", "prefix"), default="chunk" + "--load-generator", + choices=("internal", "vllm-bench"), + default="internal", ) parser.add_argument("--dataset", choices=("imdb", "longbench"), default="longbench") parser.add_argument("--model", required=True) @@ -128,9 +177,19 @@ def parse_args(argv: list[str] | None = None) -> RunBenchArgs: parser.add_argument("--gpu-util", type=float, default=0.85) parser.add_argument("--max-num-seqs", type=int, default=32) parser.add_argument("--max-num-batched-tokens", type=int, default=0) + parser.add_argument("--block-size", type=int, default=16) parser.add_argument("--max-inflight", type=int, default=32) parser.add_argument("--gen-max-tokens", type=int, default=128) parser.add_argument("--max-context-tokens", type=int, default=0) + parser.add_argument("--bench-num-prompts", type=int, default=1000) + parser.add_argument("--bench-input-len", type=int, default=1024) + parser.add_argument("--bench-output-len", type=int, default=None) + parser.add_argument("--bench-request-rate", default="inf") + parser.add_argument("--bench-max-concurrency", type=int, default=None) + parser.add_argument("--bench-random-prefix-len", type=int, default=0) + parser.add_argument("--bench-random-range-ratio", type=float, default=0.0) + parser.add_argument("--bench-seed", type=int, default=42) + parser.add_argument("--bench-burstiness", type=float, default=1.0) parser.add_argument("--evict", action="store_true") parser.add_argument( "--prometheus-url", @@ -141,9 +200,9 @@ def parse_args(argv: list[str] | None = None) -> RunBenchArgs: ), ) args = parser.parse_args(argv) - return RunBenchArgs( + parsed = RunBenchArgs( backend=args.backend, - cache_reuse_mode=args.cache_reuse_mode, + load_generator=args.load_generator, dataset=args.dataset, model=args.model, store_dir=args.store_dir, @@ -155,12 +214,27 @@ def parse_args(argv: list[str] | None = None) -> RunBenchArgs: gpu_util=args.gpu_util, max_num_seqs=args.max_num_seqs, max_num_batched_tokens=args.max_num_batched_tokens, + block_size=args.block_size, max_inflight=args.max_inflight, gen_max_tokens=args.gen_max_tokens, max_context_tokens=args.max_context_tokens, + bench_num_prompts=args.bench_num_prompts, + bench_input_len=args.bench_input_len, + bench_output_len=args.bench_output_len, + bench_request_rate=args.bench_request_rate, + bench_max_concurrency=args.bench_max_concurrency, + bench_random_prefix_len=args.bench_random_prefix_len, + bench_random_range_ratio=args.bench_random_range_ratio, + bench_seed=args.bench_seed, + bench_burstiness=args.bench_burstiness, evict=args.evict, prometheus_url=args.prometheus_url, ) + try: + _validate_run_args(parsed) + except ValueError as exc: + parser.error(str(exc)) + return parsed def run_benchmark(args: RunBenchArgs) -> Path: @@ -177,17 +251,34 @@ def run_benchmark(args: RunBenchArgs) -> Path: under ``args.store_dir``, and owns cleanup for pid files below the run root. """ + _validate_run_args(args) run_id = time.strftime("%Y%m%d_%H%M%S") run_root = Path(args.store_dir).expanduser() / f"run_{run_id}" run_root.mkdir(parents=True, exist_ok=True) prepare_path = run_root / "prepare.json" + backend_runs = _expand_backend_runs(args.backend) + _validate_backend_runs(backend_runs, load_generator=args.load_generator) _print_stage("prepare") - _print_kv("dataset", args.dataset) - _print_kv("max_samples", args.max_samples) + _print_kv("load_generator", args.load_generator) + if args.load_generator == "vllm-bench": + _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)) + 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) - _run_command(_prepare_command(args, run_root, prepare_path)) - prepare = json.loads(prepare_path.read_text(encoding="utf-8")) + if args.load_generator == "vllm-bench": + prepare = {"config": _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) + else: + _run_command(_prepare_command(args, run_root, prepare_path)) + prepare = json.loads(prepare_path.read_text(encoding="utf-8")) config = prepare["config"] derived_l1 = int(config["derived_l1_size_bytes"]) derived_l2 = int(config["derived_l2_size_bytes"]) @@ -196,9 +287,7 @@ def run_benchmark(args: RunBenchArgs) -> Path: result_paths: list[tuple[BackendRun, Path]] = [] try: - for backend_run in _expand_backend_runs( - args.backend, default_reuse_mode=args.cache_reuse_mode - ): + for backend_run in backend_runs: result_paths.append( ( backend_run, @@ -222,12 +311,11 @@ def run_benchmark(args: RunBenchArgs) -> Path: return run_root -def _expand_backend_runs(backend: str, *, default_reuse_mode: str) -> list[BackendRun]: +def _expand_backend_runs(backend: str) -> list[BackendRun]: """Resolve a requested backend into concrete benchmark rows. Args: backend: User-facing backend name. - default_reuse_mode: Reuse mode for compatibility ``daser`` requests. Returns: Concrete backend rows in execution order. @@ -235,24 +323,102 @@ def _expand_backend_runs(backend: str, *, default_reuse_mode: str) -> list[Backe Thread-safety: Pure helper. """ - if backend == "all": + row_map = { + "baseline": BackendRun("baseline", "vllm", "none"), + "lmcache": BackendRun("lmcache", "lmcache", "none"), + "daser-chunk": BackendRun("daser-chunk", "daser", "chunk"), + "daser-prefix": BackendRun("daser-prefix", "daser", "prefix"), + } + if backend.strip() == "all": return [ - BackendRun("baseline", "vllm", "none"), - BackendRun("lmcache", "lmcache", "none"), - BackendRun("daser-chunk", "daser", "chunk"), - BackendRun("daser-prefix", "daser", "prefix"), + row_map["baseline"], + row_map["lmcache"], + row_map["daser-chunk"], + row_map["daser-prefix"], ] - if backend in ("baseline", "vllm"): - return [BackendRun("baseline", "vllm", "none")] - if backend == "lmcache": - return [BackendRun("lmcache", "lmcache", "none")] - if backend == "daser": - return [BackendRun("daser", "daser", default_reuse_mode)] - if backend == "daser-chunk": - return [BackendRun("daser-chunk", "daser", "chunk")] - if backend == "daser-prefix": - return [BackendRun("daser-prefix", "daser", "prefix")] - raise ValueError(f"unknown backend: {backend}") + names = _backend_selection(backend).split(",") + return [row_map[name] for name in names] + + +def _validate_backend_runs( + backend_runs: list[BackendRun], + *, + load_generator: str, +) -> None: + """Validate backend rows against the selected load generator. + + Args: + backend_runs: Resolved benchmark rows. + load_generator: Selected load generator name. + + Raises: + ValueError: If a backend row is incompatible. + + Thread-safety: + Pure helper. + """ + if load_generator != "vllm-bench": + return + unsupported = [run.label for run in backend_runs if run.label == "daser-chunk"] + if unsupported: + raise ValueError( + "vllm-bench load generator does not support daser-chunk; " + "select baseline,lmcache,daser-prefix or use --load-generator internal" + ) + + +def _validate_run_args(args: RunBenchArgs) -> None: + """Validate benchmark runner arguments with clear preflight errors.""" + positive_ints = { + "block_size": args.block_size, + "max_num_seqs": args.max_num_seqs, + "max_inflight": args.max_inflight, + "gen_max_tokens": args.gen_max_tokens, + } + for name, value in positive_ints.items(): + if value <= 0: + raise ValueError(f"{name} must be positive") + non_negative_ints = { + "max_num_batched_tokens": args.max_num_batched_tokens, + "max_context_tokens": args.max_context_tokens, + } + for name, value in non_negative_ints.items(): + if value < 0: + raise ValueError(f"{name} must be non-negative") + if args.max_samples <= 0 and args.load_generator == "internal": + raise ValueError("max_samples must be positive") + 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") def _run_backend( @@ -294,7 +460,10 @@ def _run_backend( result_path = backend_dir / "results.json" _print_stage("cold/warm load", backend_run.label) _print_kv("output", result_path) - _run_command(_load_command(args, backend_dir, prepare_path, result_path)) + if args.load_generator == "vllm-bench": + _run_vllm_bench_load(args, manifest, backend_run, backend_dir, result_path) + else: + _run_command(_load_command(args, backend_dir, prepare_path, result_path)) if manifest is not None: _probe_daser_metrics( manifest, @@ -306,6 +475,372 @@ 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": + _wait_with_message("lmcache_warm_settle_s", 10.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]: @@ -323,6 +858,8 @@ def _prepare_command( str(args.max_samples), "--max-inflight", str(args.max_inflight), + "--block-size", + str(args.block_size), "--gen-max-tokens", str(args.gen_max_tokens), "--max-context-tokens", @@ -361,6 +898,8 @@ def _start_command( str(args.max_num_seqs), "--max-num-batched-tokens", str(args.max_num_batched_tokens), + "--block-size", + str(args.block_size), "--l1-size", str(derived_l1), "--l2-size", @@ -392,6 +931,8 @@ def _load_command( str(args.max_samples), "--max-inflight", str(args.max_inflight), + "--block-size", + str(args.block_size), "--gen-max-tokens", str(args.gen_max_tokens), "--max-context-tokens", @@ -498,9 +1039,15 @@ def _add_correctness_fields( exact_match = correctness.get("cold_warm_exact_match") if not isinstance(exact_match, dict): return + available = exact_match.get("available") + if available is not None: + fields["cold_warm_exact_match_available"] = available accuracy = exact_match.get("accuracy") if accuracy is not None: fields["cold_warm_exact_match_accuracy"] = accuracy + reason = exact_match.get("reason") + if reason is not None: + fields["cold_warm_exact_match_reason"] = reason def _run_command(command: list[str]) -> None: diff --git a/benchmarks/utils/constants.py b/benchmarks/utils/constants.py index 0a3c008..16a6a36 100644 --- a/benchmarks/utils/constants.py +++ b/benchmarks/utils/constants.py @@ -9,7 +9,6 @@ HEAD_DIM: int = 128 NUM_LAYERS: int = 36 DTYPE_BYTES: int = 2 -SLOT_SIZE: int = NUM_KV_HEADS * HEAD_DIM * 2 * NUM_LAYERS * BLOCK_TOKENS * DTYPE_BYTES BENCHMARK_SEED: int = 42 COMPARISON_GDS = "gds-vs-lmcache-local-ssd" @@ -19,3 +18,21 @@ "You are a helpful assistant answering questions using the following documents.\n\n" ) DEFAULT_IMDB_QUESTION: str = "Summarize the sentiment of this review." + + +def slot_size_for_block_tokens(block_tokens: int) -> int: + """Return bytes required for one model KV block. + + Args: + block_tokens: Number of tokens in a vLLM KV block. + + Returns: + Slot size in bytes for the benchmark model geometry. + + Thread-safety: + Pure calculation over constants. + """ + return NUM_KV_HEADS * HEAD_DIM * 2 * NUM_LAYERS * block_tokens * DTYPE_BYTES + + +SLOT_SIZE: int = slot_size_for_block_tokens(BLOCK_TOKENS) diff --git a/benchmarks/utils/loadgen.py b/benchmarks/utils/loadgen.py index 54e7af2..cad1f38 100644 --- a/benchmarks/utils/loadgen.py +++ b/benchmarks/utils/loadgen.py @@ -12,6 +12,7 @@ import httpx +from benchmarks.utils.constants import BLOCK_TOKENS from benchmarks.utils.datasets import BenchmarkSample from benchmarks.utils.metrics import ( compute_metric_delta, @@ -92,6 +93,7 @@ async def run_vllm_phase( gen_params, timeout, chunk_aligned_prompts=chunk_aligned_prompts, + block_tokens=manifest.block_size, ) return PhaseResult( requests=requests, @@ -184,6 +186,7 @@ async def run_daser_prefix( gen_params, timeout, chunk_aligned_prompts=True, + block_tokens=manifest.block_size, ) cold_phase = PhaseResult( requests=cold, @@ -200,6 +203,7 @@ async def run_daser_prefix( gen_params, timeout, chunk_aligned_prompts=True, + block_tokens=manifest.block_size, ) warm_phase = PhaseResult( requests=warm, @@ -224,6 +228,7 @@ async def run_lmcache( tokenizer, samples, chunk_aligned=chunk_aligned_prompts, + block_tokens=manifest.block_size, ) before_cold = await collect_phase_metrics(manifest) cold, cold_elapsed_ms = await _run_vllm_phase_requests( @@ -270,12 +275,14 @@ async def _run_vllm_phase_requests( timeout: float, chunk_aligned_prompts: bool = False, prompts: list[str | list[int]] | None = None, + block_tokens: int = BLOCK_TOKENS, ) -> tuple[list[RequestResult], float]: if prompts is None: prompts = build_prompt_payloads( tokenizer, samples, chunk_aligned=chunk_aligned_prompts, + block_tokens=block_tokens, ) sem = asyncio.Semaphore(max_inflight) async with httpx.AsyncClient(timeout=httpx.Timeout(timeout)) as client: @@ -434,6 +441,30 @@ def _metric_hit_ratios(metrics: dict[str, dict[str, float]]) -> dict[str, Any]: } +def backend_server_hit_rate(hit_ratios: dict[str, Any]) -> float | None: + """Return the backend token-level cache hit ratio used for comparison. + + Args: + hit_ratios: Named hit-ratio candidates from phase metric deltas. + + Returns: + Token-level backend cache hit ratio when available. + + Thread-safety: + Pure helper. + """ + for key in ( + "daser_prometheus_tokens", + "lmcache_prometheus_lookup", + "lmcache_prometheus_retrieve", + "lmcache_status_prefetch", + ): + ratio = hit_ratios.get(key) + if ratio is not None: + return float(ratio) + return None + + async def _wait_lmcache_quiescent( manifest: BenchmarkManifest, settle_seconds: float ) -> None: diff --git a/benchmarks/utils/prompts.py b/benchmarks/utils/prompts.py index c139650..0c3e603 100644 --- a/benchmarks/utils/prompts.py +++ b/benchmarks/utils/prompts.py @@ -159,6 +159,7 @@ def build_prompt_payloads( tokenizer: Any, samples: list[BenchmarkSample], chunk_aligned: bool = False, + block_tokens: int = BLOCK_TOKENS, ) -> list[str | list[int]]: """Build prompt payloads for completions requests. @@ -167,6 +168,7 @@ def build_prompt_payloads( samples: Benchmark samples. chunk_aligned: when True, return token-ID prompts using DaseR chunk padding semantics; otherwise return plain prompt strings. + block_tokens: vLLM block size used for chunk alignment. Returns: Prompt strings or token ID lists aligned with samples. @@ -176,7 +178,12 @@ def build_prompt_payloads( """ if chunk_aligned: return [ - build_chunk_aligned_prompt_ids(tokenizer, sample.context, sample.question) + build_chunk_aligned_prompt_ids( + tokenizer, + sample.context, + sample.question, + block_tokens=block_tokens, + ) for sample in samples ] return build_prompts(tokenizer, samples) diff --git a/benchmarks/utils/servers.py b/benchmarks/utils/servers.py index 40d3b5c..2d53cc9 100644 --- a/benchmarks/utils/servers.py +++ b/benchmarks/utils/servers.py @@ -59,6 +59,7 @@ class BenchmarkManifest: endpoints: Named service endpoints. log_dir: Log directory. pid_file: JSON file containing subprocess PIDs. + block_size: vLLM KV block size in tokens. Thread-safety: Immutable value object. @@ -75,6 +76,7 @@ class BenchmarkManifest: endpoints: dict[str, ServiceEndpoint] log_dir: str pid_file: str + block_size: int = BLOCK_TOKENS def write(self, path: str | Path) -> None: """Write manifest JSON atomically enough for local benchmark use.""" @@ -99,6 +101,7 @@ def read(cls, path: str | Path) -> "BenchmarkManifest": for name, endpoint in payload["endpoints"].items() } payload["endpoints"] = endpoints + payload.setdefault("block_size", BLOCK_TOKENS) return cls(**payload) @@ -117,6 +120,7 @@ def __init__( l1_size_bytes: int, l2_size_bytes: int, max_num_batched_tokens: int | None = None, + block_size: int = BLOCK_TOKENS, reuse_mode: str = "chunk", transfer_mode: str = "iouring", vllm_port: int = 8001, @@ -136,6 +140,7 @@ def __init__( gpu_util: vLLM GPU memory utilization. max_num_seqs: vLLM max_num_seqs. max_num_batched_tokens: Optional vLLM scheduler token budget. + block_size: vLLM KV block size in tokens. l1_size_bytes: L1 size. l2_size_bytes: L2 size. reuse_mode: DaseR cache reuse mode. @@ -154,6 +159,7 @@ def __init__( self.gpu_util = gpu_util self.max_num_seqs = max_num_seqs self.max_num_batched_tokens = max_num_batched_tokens + self.block_size = block_size self.l1_size_bytes = l1_size_bytes self.l2_size_bytes = l2_size_bytes self.reuse_mode = reuse_mode @@ -214,6 +220,7 @@ def manifest(self) -> BenchmarkManifest: endpoints=endpoints, log_dir=str(self.log_dir), pid_file=str(self.pid_file), + block_size=self.block_size, ) async def start_lmcache_mp_server(self) -> None: @@ -251,7 +258,7 @@ def _lmcache_mp_server_command(self) -> list[str]: "--port", str(LMCACHE_MP_PORT), "--chunk-size", - str(BLOCK_TOKENS), + str(self.block_size), "--max-workers", "4", "--l1-size-gb", @@ -350,6 +357,8 @@ def _daser_server_command(self) -> list[str]: str(self.daser_port), "--socket-path", str(self.socket_path), + "--block-tokens", + str(self.block_size), ] if self.skip_l2: cmd.append("--skip-l2") @@ -417,6 +426,8 @@ def vllm_command( "--no-enable-prefix-caching", "--generation-config", "vllm", + "--block-size", + str(self.block_size), ] if self.max_model_len is not None and self.max_model_len > 0: cmd.extend(["--max-model-len", str(self.max_model_len)]) diff --git a/daser/config.py b/daser/config.py index 52865d8..0c2c2b3 100644 --- a/daser/config.py +++ b/daser/config.py @@ -25,18 +25,29 @@ class ModelGeometry: num_layers: int dtype_bytes: int - @property - def slot_size(self) -> int: - """Return bytes required for one vLLM KV block across all layers.""" + def slot_size_for_block_tokens(self, block_tokens: int) -> int: + """Return bytes required for one vLLM KV block across all layers. + + Args: + block_tokens: Number of tokens in each vLLM KV block. + + Returns: + Bytes required for one KV slot with ``block_tokens`` tokens. + """ return ( self.num_kv_heads * self.head_dim * 2 # K and V * self.num_layers - * BLOCK_TOKENS + * block_tokens * 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. @@ -178,7 +189,9 @@ def resolved_slot_size(self) -> int: Returns: Slot size in bytes. """ - return model_geometry_from_path(self.model_path).slot_size + return model_geometry_from_path(self.model_path).slot_size_for_block_tokens( + self.block_tokens + ) def runtime_config(self) -> dict[str, object]: """Return connector runtime config owned by DaseR server. diff --git a/daser/connector/daser_connector.py b/daser/connector/daser_connector.py index 75d9f78..6c6910e 100644 --- a/daser/connector/daser_connector.py +++ b/daser/connector/daser_connector.py @@ -121,6 +121,7 @@ def __init__( self._pending_loads: dict[str, dict[str, Any]] = {} self._pending_stores: dict[str, dict[str, Any]] = {} self._pending_alloc: dict[str, PendingStore] = {} + self._pending_async_saves: set[str] = set() self._req_tokens: dict[str, list[int]] = {} else: self._transfer_ready = False @@ -138,6 +139,7 @@ def __init__( self._pending_store_staging_limit_bytes = 0 self._staging_pool = None self._pending_commits: set[str] = set() + 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 @@ -235,20 +237,6 @@ def _set_cache_reuse_strategy(self, cache_reuse_mode: str) -> None: self._block_tokens, ) - def _discard_pending_request(self, req_id: str) -> None: - """Clear scheduler-side pending state for a request. - - Args: - req_id: vLLM request ID. - """ - self._pending_loads.pop(req_id, None) - if req_id in self._pending_stores: - self._drop_pending_store(req_id) - for pending_req_id in list(self._pending_stores): - if pending_req_id.startswith(f"{req_id}:store:"): - self._drop_pending_store(pending_req_id) - self._pending_alloc.pop(req_id, None) - def _init_rope_config(self, vllm_config: "VllmConfig") -> None: """Extract default RoPE settings from vLLM model config. diff --git a/daser/connector/helpers.py b/daser/connector/helpers.py index f50a068..44f8d1a 100644 --- a/daser/connector/helpers.py +++ b/daser/connector/helpers.py @@ -50,6 +50,61 @@ def rolling_prefix_key(prev_key: str, block_tokens: list[int]) -> str: return h.hexdigest() +def rolling_prefix_keys( + tokens: list[int], + block_tokens: int, + seed: str = ROLLING_PREFIX_SEED, + start_slot: int = 0, + initial_key: str | None = None, +) -> list[str]: + """Return chained rolling-prefix keys for a token sequence. + + Args: + tokens: full prompt token IDs. + block_tokens: number of token IDs in one KV slot. + seed: initial rolling-prefix seed. + start_slot: slot index where key generation starts. + initial_key: optional key immediately before ``start_slot``. + + Returns: + Rolling keys for every full block from ``start_slot`` onward. + + Async/thread-safety: + Pure CPU helper with no shared mutable state; safe to call from any + thread or asyncio task. + """ + if block_tokens <= 0: + raise ValueError("block_tokens must be positive") + aligned = (len(tokens) // block_tokens) * block_tokens + start = start_slot * block_tokens + if start < 0 or start > aligned: + return [] + key_bytes = bytes.fromhex(initial_key or seed) + keys: list[str] = [] + for offset in range(start, aligned, block_tokens): + h = xxhash.xxh3_128() + h.update(key_bytes) + h.update(bytes(array.array("i", tokens[offset : offset + block_tokens]))) + key_bytes = h.digest() + keys.append(key_bytes.hex()) + return keys + + +def base_req_id(req_id: str) -> str: + """Return the original request ID for synthetic scheduler sub-work IDs. + + Args: + req_id: Request ID or scheduler-generated sub-work ID. + + Returns: + Base vLLM request ID. + + Thread-safety: + Pure string helper with no shared state. + """ + return req_id.split(":store:", 1)[0] + + @dataclass class PendingStore: """Scheduler-side state for a prompt KV store that may span steps. diff --git a/daser/connector/reuse.py b/daser/connector/reuse.py index f644afe..476aa9e 100644 --- a/daser/connector/reuse.py +++ b/daser/connector/reuse.py @@ -11,7 +11,7 @@ ROLLING_PREFIX_SEED, PendingStore, hash_tokens, - rolling_prefix_key, + rolling_prefix_keys, ) from daser.logging import init_logger @@ -243,13 +243,20 @@ def allocate_store( """ requested_tokens = pending_store.token_count num_slots = math.ceil(requested_tokens / self._block_tokens) - key = pending_store.rolling_key or ROLLING_PREFIX_SEED slot_i = pending_store.rolling_slot_index + key = pending_store.rolling_key or ROLLING_PREFIX_SEED + keys = rolling_prefix_keys( + tokens, + self._block_tokens, + start_slot=slot_i, + initial_key=key, + ) run: list[tuple[int, str]] = [] - while slot_i < num_slots and slot_i < len(pending_store.block_ids): - start = slot_i * self._block_tokens - key = rolling_prefix_key(key, tokens[start : start + self._block_tokens]) + for next_key in keys: + if slot_i >= num_slots or slot_i >= len(pending_store.block_ids): + break + key = next_key if slot_i >= pending_store.start_slot_index: store_id = f"{req_id}:store:{slot_i}" if not owner.has_pending_store(store_id): diff --git a/daser/connector/scheduler.py b/daser/connector/scheduler.py index c48068d..b0f8826 100644 --- a/daser/connector/scheduler.py +++ b/daser/connector/scheduler.py @@ -12,7 +12,7 @@ from vllm.v1.request import Request # First Party -from daser.connector.helpers import PendingStore +from daser.connector.helpers import PendingStore, base_req_id from daser.connector.metadata import DaserConnectorMeta, ReqLoadSpec, ReqStoreSpec from daser.connector.reuse import build_cache_reuse_strategy from daser.logging import init_logger @@ -21,15 +21,8 @@ def _base_req_id(req_id: str) -> str: - """Return the original request ID for synthetic scheduler sub-work IDs. - - Args: - req_id: Request ID or scheduler-generated sub-work ID. - - Returns: - Base vLLM request ID. - """ - return req_id.split(":store:", 1)[0] + """Compatibility wrapper for tests importing the scheduler-private helper.""" + return base_req_id(req_id) def _store_slot_index(req_id: str) -> int | None: @@ -556,6 +549,9 @@ def build_connector_meta( if meta.reqs_to_store: meta.reqs_to_store = self._filter_live_store_specs(meta.reqs_to_store) + pending_async_saves = self._pending_async_save_ids() + for req_id in meta.reqs_to_store: + pending_async_saves.add(_base_req_id(req_id)) if logger.isEnabledFor(logging.DEBUG): for req_id, spec in meta.reqs_to_load.items(): @@ -589,8 +585,10 @@ def _drop_preempted_pending_state( Runs on the scheduler thread before metadata is handed to workers. """ preempted_req_ids = getattr(scheduler_output, "preempted_req_ids", set()) + pending_async_saves = self._pending_async_save_ids() for req_id in preempted_req_ids: base_req_id = str(req_id) + pending_async_saves.discard(base_req_id) for pending_req_id in list(self._pending_loads): if _matches_request_or_store_id(pending_req_id, base_req_id): self._pending_loads.pop(pending_req_id, None) @@ -856,6 +854,36 @@ def _drop_pending_store(self, req_id: str) -> None: except Exception as exc: # noqa: BLE001 logger.warning("[CONNECTOR] release_chunk_writer failed: %s", exc) + def _discard_pending_request(self, req_id: str) -> None: + """Clear scheduler-side pending state for a request. + + Args: + req_id: vLLM request ID. + """ + self._pending_loads.pop(req_id, None) + if req_id in self._pending_stores: + self._drop_pending_store(req_id) + for pending_req_id in list(self._pending_stores): + if pending_req_id.startswith(f"{req_id}:store:"): + self._drop_pending_store(pending_req_id) + self._pending_alloc.pop(req_id, None) + + def _pending_async_save_ids(self) -> set[str]: + """Return request IDs whose worker-side saves are still pending. + + Returns: + Mutable set of base vLLM request IDs. + + Thread-safety: + Runs on the scheduler thread. The lazy initialization supports + tests and mixin probes that do not call ``DaserConnector.__init__``. + """ + pending = getattr(self, "_pending_async_saves", None) + if pending is None: + pending = set() + self._pending_async_saves = pending + return pending + def _maybe_allocate_pending_store( self, req_id: str, pending_store: PendingStore ) -> None: @@ -886,8 +914,28 @@ def request_finished( block_ids: block IDs being freed. Returns: - (False, None) - no async cleanup needed. + (True, None) when DaseR is still storing this request's KV blocks, + otherwise (False, None). """ + del block_ids + if request.request_id in self._pending_async_save_ids(): + return True, None self._req_tokens.pop(request.request_id, None) self._discard_pending_request(request.request_id) return False, None + + def update_connector_output(self, connector_output: Any) -> None: + """Update scheduler state from worker-side transfer completions. + + Args: + connector_output: vLLM KVConnectorOutput carrying finished request + IDs from workers. + + Async/thread-safety: + Runs on vLLM's scheduler thread after worker connector polling. + """ + pending_async_saves = self._pending_async_save_ids() + for req_id in getattr(connector_output, "finished_sending", None) or (): + pending_async_saves.discard(req_id) + self._req_tokens.pop(req_id, None) + self._discard_pending_request(req_id) diff --git a/daser/connector/worker.py b/daser/connector/worker.py index ea4d423..a26f614 100644 --- a/daser/connector/worker.py +++ b/daser/connector/worker.py @@ -4,6 +4,7 @@ # Standard import asyncio +from dataclasses import dataclass import os import time from typing import TYPE_CHECKING, Any @@ -19,8 +20,10 @@ from vllm.forward_context import ForwardContext # First Party +from daser.connector.helpers import base_req_id from daser.connector.metadata import ( DaserConnectorMeta, + ReqStoreSpec, StoreWriteSpan, ) from daser.connector.staging import ( @@ -78,6 +81,16 @@ _ROPE_WARMUP_BLOCKS = 1 +@dataclass +class _DeferredFinishedSave: + """Store work held until vLLM reports a request as finished.""" + + commit_keys: set[str] + reqs_to_store: dict[str, ReqStoreSpec] + submitted: bool = False + future: Any | None = None + + def _cuda_allocation_base_and_offset(device_ptr: int) -> tuple[int, int]: """Return CUDA allocation base pointer and byte offset for ``device_ptr``. @@ -586,39 +599,26 @@ def save_kv_layer( ) def wait_for_save(self) -> None: - """Wait for submitted layer stores and commit visible chunks.""" + """Queue stores until vLLM reports request completion.""" if self._meta is None: return commit_keys = list(self._pending_commits) reqs_to_store = dict(self._meta.reqs_to_store) if commit_keys and reqs_to_store: - batch_futures = [] - batches = _build_staging_store_batches( - reqs_to_store, - self._slot_size, - max_batch_bytes=( - self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES - ), - ) - for block_ids, spans in batches: - staged = self._stage_store_batch(block_ids, spans) - if staged is None: - continue - future = self._submit_store_coroutine( - self._write_cuda_buffer( - buffer=staged.buffer, - ready_event=staged.ready_event, - spans=staged.spans, - ) - ) - self._track_save_future(future, staged.buffer.nbytes, staged.lease) - batch_futures.append(future) - if batch_futures: - commit_future = self._submit_store_coroutine( - self._commit_after_store_futures(batch_futures, commit_keys), - ) - self._track_save_future(commit_future, 0, None) + pending_finished = getattr(self, "_pending_finished_saves", None) + if pending_finished is None: + pending_finished = {} + self._pending_finished_saves = pending_finished + for req_id, spec in reqs_to_store.items(): + base_id = base_req_id(req_id) + save = pending_finished.get(base_id) + if save is None: + save = _DeferredFinishedSave(commit_keys=set(), reqs_to_store={}) + pending_finished[base_id] = save + 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() @@ -631,16 +631,41 @@ def get_finished( finished_req_ids: Request IDs that vLLM finished in this step. Returns: - ``(None, None)`` because DaseR does not take ownership of request - blocks beyond the current vLLM lifecycle. + Finished-saving request IDs and no async receiving IDs. """ self._reap_save_futures(block=False) - return None, None + pending_finished = getattr(self, "_pending_finished_saves", {}) + if not pending_finished: + return None, None + + finished_sending: set[str] = set() + candidates = set(finished_req_ids) + candidates.update( + req_id for req_id, save in pending_finished.items() if save.submitted + ) + for req_id in list(candidates): + save = pending_finished.get(req_id) + if save is None: + continue + if not save.submitted: + save.future = self._submit_finished_save(save) + save.submitted = True + future = save.future + if future is None: + finished_sending.add(req_id) + del pending_finished[req_id] + elif future.done(): + future.result(timeout=120.0) + finished_sending.add(req_id) + del pending_finished[req_id] + return finished_sending or None, None def shutdown(self) -> None: """Stop the background IO loop.""" if self._role != KVConnectorRole.WORKER: return + for req_id in list(getattr(self, "_pending_finished_saves", {})): + self.get_finished({req_id}) self._reap_save_futures(block=True) load_client = getattr(self, "_ipc_load_async", None) store_client = getattr(self, "_ipc_store_async", None) @@ -905,7 +930,6 @@ def _stage_store_batch( slot_size=self._slot_size, block_index=block_index, ) - _synchronize_cuda_tensor(staging) return StagedStoreBatch( buffer=staging, ready_event=_record_cuda_event(staging), @@ -913,6 +937,47 @@ def _stage_store_batch( lease=staging_lease, ) + def _submit_finished_save(self, save: _DeferredFinishedSave) -> Any | None: + """Submit one request's deferred KV store after request completion. + + Args: + save: Deferred store plan built during ``wait_for_save``. + + Returns: + Future that completes after all store batches are committed, or + ``None`` when no store batch could be staged. + + Async/thread-safety: + Called by vLLM's worker thread from ``get_finished`` while vLLM is + still holding the finished request's KV blocks. + """ + batch_futures = [] + batches = _build_staging_store_batches( + save.reqs_to_store, + self._slot_size, + max_batch_bytes=(self._store_staging_bytes or DEFAULT_STORE_STAGING_BYTES), + ) + for block_ids, spans in batches: + staged = self._stage_store_batch(block_ids, spans) + if staged is None: + continue + future = self._submit_store_coroutine( + self._write_cuda_buffer( + buffer=staged.buffer, + ready_event=staged.ready_event, + spans=staged.spans, + ) + ) + self._track_save_future(future, staged.buffer.nbytes, staged.lease) + batch_futures.append(future) + if not batch_futures: + return None + commit_future = self._submit_store_coroutine( + self._commit_after_store_futures(batch_futures, sorted(save.commit_keys)), + ) + self._track_save_future(commit_future, 0, None) + return commit_future + async def _commit_after_store_futures( self, batch_futures: list[Any], diff --git a/daser/retrieval/prefix.py b/daser/retrieval/prefix.py index 542f831..38aea15 100644 --- a/daser/retrieval/prefix.py +++ b/daser/retrieval/prefix.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # First Party -from daser.connector.helpers import ROLLING_PREFIX_SEED, rolling_prefix_key +from daser.connector.helpers import rolling_prefix_keys from daser.logging import init_logger from daser.retrieval.base import RetrievalIndex, RetrievalMatch from daser.server.metadata_store import ChunkMeta @@ -44,7 +44,6 @@ async def lookup(self, tokens: list[int], model_id: str) -> list[RetrievalMatch] matches: list[RetrievalMatch] = [] run_meta: ChunkMeta | None = None run_target_start = 0 - key = ROLLING_PREFIX_SEED aligned = (len(tokens) // self._block_tokens) * self._block_tokens def flush_run() -> None: @@ -59,8 +58,9 @@ def flush_run() -> None: run_meta = None run_target_start = 0 - for slot_i, start in enumerate(range(0, aligned, self._block_tokens)): - key = rolling_prefix_key(key, tokens[start : start + self._block_tokens]) + for slot_i, key in enumerate( + rolling_prefix_keys(tokens[:aligned], self._block_tokens) + ): meta = self._index.get(key) if meta is None or meta.model_id != model_id: break diff --git a/daser/server/__main__.py b/daser/server/__main__.py index e74b7e8..1da28f1 100644 --- a/daser/server/__main__.py +++ b/daser/server/__main__.py @@ -205,6 +205,12 @@ def _parse_args() -> argparse.Namespace: help="Use volatile L1 memory only: do not allocate daser.store and do " "not persist daser.index. Incompatible with --transfer-mode=gds.", ) + parser.add_argument( + "--block-tokens", + type=int, + default=BLOCK_TOKENS, + help="vLLM KV block size in tokens. Must match vLLM --block-size.", + ) return parser.parse_args() @@ -303,6 +309,7 @@ def _build_daser_config(args: argparse.Namespace) -> DaserConfig: total_store_bytes=total_store_bytes, ipc_socket_path=args.socket_path, log_level=args.log_level, + block_tokens=int(args.block_tokens), cache_reuse_mode=args.cache_reuse_mode, transfer_mode=transfer_mode, l1_size_bytes=l1_size, @@ -337,7 +344,7 @@ def _build_http_config(args: argparse.Namespace) -> HTTPServerConfig: vllm_base_url=args.vllm_base_url, model=getattr(args, "vllm_model_id", None) or args.model_path, tokenizer=args.model_path, - block_tokens=BLOCK_TOKENS, + block_tokens=int(args.block_tokens), cache_reuse_mode=args.cache_reuse_mode, align_document_chunks=args.cache_reuse_mode == "chunk", transfer_mode=args.transfer_mode, diff --git a/daser/transfer/memory/layer.py b/daser/transfer/memory/layer.py index 5c82d99..a30ff00 100644 --- a/daser/transfer/memory/layer.py +++ b/daser/transfer/memory/layer.py @@ -222,16 +222,34 @@ async def store_bytes_grouped( Total number of bytes stored. Async/thread-safety: - Calls ``store_bytes`` for each span so replacement ordering matches - single-span stores. + Serializes L1 metadata once for the whole group while preserving + per-span replacement ordering. """ total = 0 - for span in spans: - source_offset = int(span.get("source_offset", 0)) - nbytes = int(span["nbytes"]) - file_offset = int(span["file_offset"]) - source = self._slice_src(src, source_offset, nbytes) - total += await self.store_bytes(source, file_offset, nbytes) + async with self._lock: + for span in spans: + source_offset = int(span.get("source_offset", 0)) + nbytes = int(span["nbytes"]) + file_offset = int(span["file_offset"]) + self._check_range(file_offset, nbytes) + key = (file_offset, nbytes) + hit = self._find_l1_locked(file_offset, nbytes) + source = self._slice_src(src, source_offset, nbytes) + if hit is not None: + hit_key, cached, target_offset = hit + self._copy_src_to_pinned_at(source, cached, target_offset, nbytes) + self._policy.access(hit_key) + self._l1.move_to_end(hit_key) + total += nbytes + continue + data = self._reserve_l1_buffer_locked(key, nbytes) + try: + self._copy_src_to_pinned_at(source, data, 0, nbytes) + except BaseException: + data.close() + raise + self._put_l1_locked(key, data) + total += nbytes return total async def drain(self) -> None: diff --git a/tests/connector/test_daser_connector.py b/tests/connector/test_daser_connector.py index 1abb1bc..5a7ab18 100644 --- a/tests/connector/test_daser_connector.py +++ b/tests/connector/test_daser_connector.py @@ -16,6 +16,7 @@ PendingStore, hash_tokens, rolling_prefix_key, + rolling_prefix_keys, ) from daser.connector.metadata import ( DaserConnectorMeta, @@ -64,7 +65,7 @@ from daser.connector.staging import ( synchronize_cuda_tensor as _synchronize_cuda_tensor, ) -from daser.connector.worker import WorkerConnectorMixin +from daser.connector.worker import WorkerConnectorMixin, _DeferredFinishedSave BLOCK_TOKENS = 4 NUM_LAYERS = 2 @@ -83,6 +84,16 @@ def rolling_keys(tokens: list[int], block_tokens: int) -> list[str]: return keys +def test_rolling_prefix_keys_match_single_step_helper() -> None: + """Batched rolling-prefix keys preserve the existing key sequence.""" + tokens = list(range(32)) + + assert rolling_prefix_keys(tokens, block_tokens=8) == rolling_keys( + tokens, + block_tokens=8, + ) + + class _RuntimeConfigProbe(DaserConnector): """Test connector exposing runtime config state through public properties.""" @@ -153,6 +164,7 @@ def __init__(self) -> None: self._pending_loads = {} self._pending_stores = {} self._pending_alloc = {} + self._pending_async_saves = set() self._req_tokens = {} self._model_id = "m" self.alloc_calls: list[tuple[str, int, str]] = [] @@ -236,6 +248,27 @@ def seed_pending_store_spec(self, req_id: str, spec: ReqStoreSpec) -> None: "token_count": spec.token_count, } + def request_finished_for_test(self, req_id: str): + """Expose scheduler request-finished behavior through a test request.""" + + class _Request: + request_id = req_id + + return self.request_finished(_Request(), []) + + def update_connector_output_for_test(self, finished_sending: set[str]): + """Expose connector output handling for async save completions.""" + + class _Output: + def __init__(self, ids: set[str]) -> None: + self.finished_sending = ids + + self.update_connector_output(_Output(finished_sending)) + + def has_req_tokens(self, req_id: str) -> bool: + """Return whether scheduler token state is still held for a request.""" + return req_id in self._req_tokens + @property def pending_state(self) -> tuple[dict, dict]: """Return pending allocation and store state for assertions.""" @@ -306,6 +339,125 @@ async def store_via_public_helper(self) -> None: await self._transfer_store_cuda() +class _FinishedSaveProbe(WorkerConnectorMixin): + """Worker probe for finished-request save scheduling.""" + + def __init__(self) -> None: + self._meta = None + self._pending_commits = set() + self._pending_finished_saves = {} + self._save_futures = [] + self._pending_save_staging_bytes = 0 + self._slot_size = 32 + self._store_staging_bytes = 128 + self._pending_store_staging_limit_bytes = 128 + self._layer_names = ["layer.0"] + self._kv_caches = {"layer.0": torch.empty(1)} + self.staged_batches: list[tuple[list[int], list[StoreWriteSpan]]] = [] + self.submitted = 0 + self.tracked: list[tuple[int, object | None]] = [] + self.committed_after: list[tuple[int, list[str]]] = [] + + def _clear_save_state(self) -> None: + return + + def _reap_save_futures(self, block: bool) -> None: + if block: + for future, _bytes, lease in self._save_futures: + future.result(timeout=120.0) + if lease is not None: + lease.release() + self._save_futures = [] + + def _stage_store_batch( + self, + block_ids: list[int], + spans: list[StoreWriteSpan], + ): + self.staged_batches.append((list(block_ids), list(spans))) + + class _Staged: + buffer = torch.empty(len(block_ids) * 32, dtype=torch.uint8) + ready_event = None + lease = object() + + def __init__(self, spans): + self.spans = spans + + return _Staged(spans) + + def _submit_store_coroutine(self, coro): + self.submitted += 1 + if getattr(getattr(coro, "cr_code", None), "co_name", "") == ( + "_commit_after_store_futures" + ): + asyncio.get_event_loop().run_until_complete(coro) + else: + coro.close() + + class _Future: + def __init__(self, value): + self._value = value + + def done(self) -> bool: + return True + + def result(self, timeout: float): + del timeout + return self._value + + return _Future(["stored"]) + + def _track_save_future( + self, + future, + staging_bytes: int, + staging_lease, + ) -> None: + self.tracked.append((staging_bytes, staging_lease)) + self._save_futures.append((future, staging_bytes, staging_lease)) + + async def _commit_after_store_futures(self, batch_futures, commit_keys): + self.committed_after.append((len(batch_futures), list(commit_keys))) + + def set_pending_meta( + self, + reqs_to_store: dict[str, ReqStoreSpec], + commit_keys: set[str], + ) -> None: + """Seed worker metadata and commit keys for save tests.""" + self._meta = DaserConnectorMeta(reqs_to_store=reqs_to_store) + self._pending_commits = commit_keys + + def seed_finished_save( + self, + req_id: str, + reqs_to_store: dict[str, ReqStoreSpec], + commit_keys: list[str], + ) -> None: + """Seed a deferred finished save for worker completion tests.""" + self._pending_finished_saves[req_id] = _DeferredFinishedSave( + commit_keys=set(commit_keys), + reqs_to_store=reqs_to_store, + ) + + def pending_finished_save_ids(self) -> set[str]: + """Return request IDs with deferred save work.""" + return set(self._pending_finished_saves) + + def pending_commit_keys(self) -> set[str]: + """Return pending worker commit keys.""" + return set(self._pending_commits) + + def set_submit_store_coroutine(self, submitter) -> None: + """Replace the store coroutine submitter for worker tests.""" + self._submit_store_coroutine = submitter + + def disable_store_staging(self) -> None: + """Make staging fail for store lifecycle regression tests.""" + self._stage_store_batch = lambda _block_ids, _spans: None + + class _LoopProbe(WorkerConnectorMixin): """Minimal worker probe exposing background loop selection.""" @@ -536,6 +688,181 @@ def fake_submit_load_coroutine(coro): assert "start_load_kv timing" not in caplog.text +def test_wait_for_save_defers_store_until_request_finished() -> None: + """Cold stores should be snapshotted after vLLM reports request finish.""" + connector = _FinishedSaveProbe() + connector.set_pending_meta( + { + "req": ReqStoreSpec( + chunk_key="stored", + start_slot=0, + num_slots=2, + block_ids=[4, 5], + file_offset=0, + token_count=8, + ) + }, + {"stored"}, + ) + + connector.wait_for_save() + + assert connector.staged_batches == [] + assert connector.submitted == 0 + assert connector.pending_finished_save_ids() == {"req"} + assert connector.pending_commit_keys() == set() + + finished_sending, finished_recving = connector.get_finished({"req"}) + + assert finished_recving is None + assert finished_sending == {"req"} + assert [block_ids for block_ids, _spans in connector.staged_batches] == [[4, 5]] + assert connector.submitted == 2 + assert connector.committed_after == [(1, ["stored"])] + + +def test_wait_for_save_groups_prefix_slot_stores_by_base_request() -> None: + """Synthetic prefix slot stores should finish with their base request.""" + connector = _FinishedSaveProbe() + connector.set_pending_meta( + { + "req:store:0": ReqStoreSpec( + chunk_key="stored-0", + start_slot=0, + num_slots=1, + block_ids=[4], + file_offset=0, + token_count=4, + ), + "req:store:1": ReqStoreSpec( + chunk_key="stored-1", + start_slot=1, + num_slots=1, + block_ids=[5], + file_offset=32, + token_count=4, + ), + }, + {"stored-0", "stored-1"}, + ) + + connector.wait_for_save() + + assert connector.pending_finished_save_ids() == {"req"} + + finished_sending, finished_recving = connector.get_finished({"req"}) + + assert finished_recving is None + assert finished_sending == {"req"} + assert [block_ids for block_ids, _spans in connector.staged_batches] == [[4, 5]] + assert connector.committed_after == [(1, ["stored-0", "stored-1"])] + + +def test_get_finished_holds_blocks_until_deferred_store_completes() -> None: + """Worker should not release finished request blocks before store is done.""" + connector = _FinishedSaveProbe() + connector.seed_finished_save( + "req", + { + "req": ReqStoreSpec( + chunk_key="stored", + start_slot=0, + num_slots=1, + block_ids=[4], + file_offset=0, + token_count=4, + ) + }, + ["stored"], + ) + + class _PendingFuture: + def done(self) -> bool: + return False + + def submit_pending(coro): + coro.close() + return _PendingFuture() + + connector.set_submit_store_coroutine(submit_pending) + + finished_sending, finished_recving = connector.get_finished({"req"}) + + assert finished_recving is None + assert finished_sending is None + assert connector.staged_batches + assert "req" in connector.pending_finished_save_ids() + + +def test_get_finished_reports_completed_deferred_store_on_later_step() -> None: + """Completed saves should be reported even after the original finish step.""" + connector = _FinishedSaveProbe() + pending_future = None + + class _PendingFuture: + def __init__(self) -> None: + self.complete = False + + def done(self) -> bool: + return self.complete + + def result(self, timeout: float): + del timeout + return None + + def submit_pending(coro): + nonlocal pending_future + coro.close() + pending_future = _PendingFuture() + return pending_future + + connector.seed_finished_save( + "req", + { + "req": ReqStoreSpec( + chunk_key="stored", + start_slot=0, + num_slots=1, + block_ids=[4], + file_offset=0, + token_count=4, + ) + }, + ["stored"], + ) + connector.set_submit_store_coroutine(submit_pending) + + assert connector.get_finished({"req"}) == (None, None) + assert pending_future is not None + pending_future.complete = True + + assert connector.get_finished(set()) == ({"req"}, None) + + +def test_get_finished_releases_request_when_no_store_batch_can_be_staged() -> None: + """A skipped staging batch should not keep scheduler blocks forever.""" + connector = _FinishedSaveProbe() + connector.seed_finished_save( + "req", + { + "req": ReqStoreSpec( + chunk_key="stored", + start_slot=0, + num_slots=1, + block_ids=[4], + file_offset=0, + token_count=4, + ) + }, + ["stored"], + ) + connector.disable_store_staging() + + assert connector.get_finished({"req"}) == ({"req"}, None) + assert connector.pending_finished_save_ids() == set() + assert connector.submitted == 0 + + def test_worker_transfer_ready_allows_skip_l2_without_store_path() -> None: """L1-only mode has no store path but still has a valid transfer config.""" connector = _WorkerProbe("") @@ -1587,7 +1914,7 @@ def stage_store_batch(self, block_ids: list[int], spans: list[StoreWriteSpan]): assert calls == [] -def test_stage_store_batch_synchronizes_staging_before_return(monkeypatch): +def test_stage_store_batch_records_ready_event_without_synchronizing(monkeypatch): from daser.connector import worker class Probe(WorkerConnectorMixin): @@ -1612,11 +1939,13 @@ def stage_store_batch(self, block_ids: list[int], spans: list[StoreWriteSpan]): return self._stage_store_batch(block_ids, spans) synced = [] + recorded = object() monkeypatch.setattr( worker, "_synchronize_cuda_tensor", lambda tensor: synced.append(tensor), ) + monkeypatch.setattr(worker, "_record_cuda_event", lambda tensor: recorded) probe = Probe() staged = probe.stage_store_batch( @@ -1625,7 +1954,8 @@ def stage_store_batch(self, block_ids: list[int], spans: list[StoreWriteSpan]): ) assert staged is not None - assert synced == [staged.buffer] + assert staged.ready_event is recorded + assert synced == [] def test_stage_store_batch_keeps_dynamic_rope_warmup_out_of_store_path(monkeypatch): @@ -1905,12 +2235,19 @@ def __init__(self) -> None: super().__init__() self.use_prefix_reuse_strategy() - calls: list[tuple[str, list[int]]] = [] + calls: list[tuple[list[int], int, str | None, int]] = [] monkeypatch.setattr( - "daser.connector.reuse.rolling_prefix_key", - lambda prev_key, block: ( - calls.append((prev_key, list(block))) or rolling_prefix_key(prev_key, block) + "daser.connector.reuse.rolling_prefix_keys", + lambda tokens, block_tokens, initial_key=None, start_slot=0, **kwargs: ( + calls.append((list(tokens), block_tokens, initial_key, start_slot)) + or rolling_prefix_keys( + tokens, + block_tokens, + initial_key=initial_key, + start_slot=start_slot, + **kwargs, + ) ), ) @@ -1925,15 +2262,83 @@ def __init__(self) -> None: connector.maybe_allocate_store_for_test("req") _, pending_stores = connector.pending_state - assert calls == [ - (ROLLING_PREFIX_SEED, tokens[:BLOCK_TOKENS]), - (key0, tokens[BLOCK_TOKENS : BLOCK_TOKENS * 2]), - (key1, tokens[BLOCK_TOKENS * 2 : BLOCK_TOKENS * 3]), - ] + assert calls == [(tokens, BLOCK_TOKENS, ROLLING_PREFIX_SEED, 0)] assert connector.alloc_calls == [("batch", 3, "m")] assert pending_stores["req:store:2"]["chunk_key"] == key2 +def test_request_finished_keeps_request_until_store_finishes() -> None: + """Scheduler should hold finished requests while worker stores KV.""" + connector = _AllocatingSchedulerProbe() + connector.seed_tokens("req", [1] * 8) + connector.seed_pending_store_spec( + "req", + ReqStoreSpec( + chunk_key="live-store", + start_slot=0, + num_slots=2, + block_ids=[10, 11], + file_offset=0, + token_count=8, + ), + ) + + class Output: + num_scheduled_tokens = {"req": 8} + + meta = connector.build_connector_meta(Output()) + + assert list(meta.reqs_to_store) == ["req"] + assert connector.request_finished_for_test("req") == (True, None) + assert connector.has_req_tokens("req") + + connector.update_connector_output_for_test({"req"}) + + assert not connector.has_req_tokens("req") + assert connector.pending_state == ({}, {}) + + +def test_request_finished_keeps_prefix_store_request_until_store_finishes() -> None: + """Synthetic prefix store IDs should hold their base request lifecycle.""" + connector = _AllocatingSchedulerProbe() + connector.seed_tokens("req", [1] * 8) + connector.seed_pending_store_spec( + "req:store:0", + ReqStoreSpec( + chunk_key="live-store-0", + start_slot=0, + num_slots=1, + block_ids=[10], + file_offset=0, + token_count=4, + ), + ) + connector.seed_pending_store_spec( + "req:store:1", + ReqStoreSpec( + chunk_key="live-store-1", + start_slot=1, + num_slots=1, + block_ids=[11], + file_offset=32, + token_count=4, + ), + ) + + class Output: + num_scheduled_tokens = {"req": 8} + + meta = connector.build_connector_meta(Output()) + + assert sorted(meta.reqs_to_store) == ["req:store:0", "req:store:1"] + assert connector.request_finished_for_test("req") == (True, None) + + connector.update_connector_output_for_test({"req"}) + + assert not connector.has_req_tokens("req") + assert connector.pending_state == ({}, {}) + + def test_prefix_store_allocation_skips_committed_duplicate_slot() -> None: """Rolling-prefix stores should not enqueue writes for committed slots.""" @@ -2349,6 +2754,25 @@ class Output: assert meta.reqs_to_store == {"a": specs["a"]} +def test_stale_filtered_store_does_not_hold_finished_request() -> None: + """Scheduler should only wait for stores actually sent to workers.""" + connector = _AllocatingSchedulerProbe() + connector.seed_tokens("req", [1] * 4) + connector.seed_pending_store_spec( + "req", + ReqStoreSpec("stale-key", 0, 1, [10], 0, 4), + ) + + class Output: + num_scheduled_tokens = {"req": 4} + scheduled_cached_reqs = None + + meta = connector.build_connector_meta(Output()) + + assert meta.reqs_to_store == {} + assert connector.request_finished_for_test("req") == (False, None) + + def test_hash_tokens_deterministic(): tokens = [1, 2, 3, 4] assert hash_tokens(tokens) == hash_tokens(tokens) diff --git a/tests/test_config.py b/tests/test_config.py index e72fbbc..af8bc57 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -90,6 +90,32 @@ def test_daser_config_derives_paths_and_slot_size(tmp_path: Path) -> None: assert cfg.aligned_store_bytes == cfg.resolved_slot_size() * 4 +def test_daser_config_uses_configured_block_tokens_for_slot_size( + tmp_path: Path, +) -> None: + """Resolved slot size follows the configured vLLM block size.""" + model_path = tmp_path / "model" + store_dir = tmp_path / "store" + _write_model_config( + model_path, + { + "hidden_size": 1024, + "num_attention_heads": 8, + "num_key_value_heads": 4, + "num_hidden_layers": 28, + "torch_dtype": "float16", + }, + ) + cfg = DaserConfig( + model_path=str(model_path), + store_dir=str(store_dir), + block_tokens=128, + ) + + assert cfg.resolved_slot_size() == 4 * 128 * 2 * 28 * 128 * 2 + assert cfg.runtime_config()["block_tokens"] == 128 + + def test_runtime_config_reuses_server_parameters(tmp_path: Path) -> None: model_path = tmp_path / "model" store_dir = tmp_path / "store" diff --git a/tests/transfer/test_l1_only_transfer.py b/tests/transfer/test_l1_only_transfer.py index 859dec6..be984af 100644 --- a/tests/transfer/test_l1_only_transfer.py +++ b/tests/transfer/test_l1_only_transfer.py @@ -109,6 +109,64 @@ def test_l1_only_transfer_grouped_loads_l1_ranges() -> None: assert layer.stats.l2_writes == 0 +def test_l1_only_grouped_store_uses_single_lock_pass(monkeypatch) -> None: + """Grouped L1 stores should avoid per-span store_bytes overhead.""" + layer = L1OnlyTransferLayer(l1_bytes=ALIGNMENT * 3) + store_bytes_calls = 0 + lock_entries = 0 + original_store_bytes = layer.store_bytes + original_lock = layer._lock # noqa: SLF001 + + async def counted_store_bytes(src, file_offset: int, nbytes: int) -> int: + nonlocal store_bytes_calls + store_bytes_calls += 1 + return await original_store_bytes(src, file_offset, nbytes) + + class CountedLock: + def __init__(self, lock): + self._lock = lock + + async def __aenter__(self): + nonlocal lock_entries + lock_entries += 1 + return await self._lock.__aenter__() + + async def __aexit__(self, exc_type, exc, tb): + return await self._lock.__aexit__(exc_type, exc, tb) + + monkeypatch.setattr(layer, "store_bytes", counted_store_bytes) + monkeypatch.setattr(layer, "_lock", CountedLock(original_lock)) + try: + written = _run( + layer.store_bytes_grouped( + _block(b"a") + _block(b"b") + _block(b"c"), + [ + {"source_offset": 0, "file_offset": 0, "nbytes": ALIGNMENT}, + { + "source_offset": ALIGNMENT, + "file_offset": ALIGNMENT, + "nbytes": ALIGNMENT, + }, + { + "source_offset": ALIGNMENT * 2, + "file_offset": ALIGNMENT * 2, + "nbytes": ALIGNMENT, + }, + ], + ) + ) + dst = bytearray(ALIGNMENT * 3) + loaded = _run(layer.load_bytes(dst, file_offset=0, nbytes=ALIGNMENT * 3)) + finally: + layer.close() + + assert written == ALIGNMENT * 3 + assert loaded == ALIGNMENT * 3 + assert bytes(dst) == bytes(_block(b"a") + _block(b"b") + _block(b"c")) + assert store_bytes_calls == 0 + assert lock_entries == 2 + + def test_l1_only_transfer_overwrite_preserves_adjacent_coalesced_ranges() -> None: """Overwriting part of a coalesced L1 range should keep neighbors loadable.""" layer = L1OnlyTransferLayer(l1_bytes=ALIGNMENT * 4) diff --git a/tests/unit/test_benchmark_unified_utils.py b/tests/unit/test_benchmark_unified_utils.py index 4a92341..29e5232 100644 --- a/tests/unit/test_benchmark_unified_utils.py +++ b/tests/unit/test_benchmark_unified_utils.py @@ -7,6 +7,9 @@ from pathlib import Path import subprocess import sys +from typing import Any + +import pytest from benchmarks.bench_load import ( _add_phase_comparison, @@ -20,11 +23,18 @@ 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, ) @@ -284,26 +294,25 @@ def test_daser_prometheus_token_hit_ratio() -> None: ) -def test_daser_summary_hit_rate_uses_vllm_external_prefix() -> None: - """DaseR summary hit rate uses DaseR's vLLM-equivalent internal counters.""" +def test_daser_summary_hit_rate_uses_token_counters() -> None: + """DaseR backend comparison should use token-level cache hit ratio.""" assert ( _backend_server_hit_rate( { "daser_external_prefix": 0.93, - "daser_prometheus_tokens": 1.0, + "daser_prometheus_tokens": 0.75, "daser_prometheus_requests": 1.0, } ) - == 0.93 + == 0.75 ) -def test_daser_summary_hit_rate_ignores_control_plane_lookup_ratio() -> None: - """DaseR lookup counters are diagnostics, not external-prefix hit ratio.""" +def test_daser_summary_hit_rate_ignores_request_level_ratio() -> None: + """DaseR request counters are diagnostics, not token-level hit ratio.""" assert ( _backend_server_hit_rate( { - "daser_prometheus_tokens": 1.0, "daser_prometheus_requests": 1.0, } ) @@ -408,6 +417,7 @@ def test_load_config_uses_prepared_sizing_over_runtime_limits(tmp_path: Path) -> total_prompt_tokens=956512, total_blocks=59747, max_prompt_blocks=2314, + block_size=16, evict=False, sizing=None, ) @@ -444,6 +454,7 @@ def test_load_config_falls_back_to_manifest_sizing_without_prepare() -> None: total_prompt_tokens=956512, total_blocks=59747, max_prompt_blocks=2314, + block_size=16, evict=False, sizing=None, ) @@ -563,6 +574,7 @@ def test_manifest_round_trip(tmp_path: Path) -> None: }, log_dir=str(tmp_path / "logs"), pid_file=str(tmp_path / "pids.json"), + block_size=128, ) path = tmp_path / "manifest.json" @@ -572,6 +584,35 @@ def test_manifest_round_trip(tmp_path: Path) -> None: assert loaded == manifest +def test_manifest_read_defaults_legacy_block_size(tmp_path: Path) -> None: + """Older benchmark manifests remain readable after block-size tracking.""" + path = tmp_path / "manifest.json" + path.write_text( + json.dumps( + { + "run_id": "run1", + "backend": "daser", + "reuse_mode": "chunk", + "model": "/models/qwen", + "store_dir": str(tmp_path), + "l1_size_bytes": 1024, + "l2_size_bytes": 2048, + "skip_l2": True, + "endpoints": { + "vllm": {"url": "http://127.0.0.1:8001"}, + "daser": {"url": "http://127.0.0.1:2026"}, + }, + "log_dir": str(tmp_path / "logs"), + "pid_file": str(tmp_path / "pids.json"), + } + ) + ) + + loaded = BenchmarkManifest.read(path) + + assert loaded.block_size == 16 + + def test_daser_noevict_start_uses_l1_only_mode(tmp_path: Path) -> None: """DaseR no-evict starts without L2 sizing arguments.""" manager = ServerManager( @@ -1000,9 +1041,11 @@ async def fake_run_vllm_phase_requests( _timeout, chunk_aligned_prompts=False, prompts=None, + block_tokens=16, ): del prompts chunk_aligned_values.append(bool(chunk_aligned_prompts)) + assert block_tokens == 128 return ( [ RequestResult( @@ -1041,6 +1084,7 @@ async def fake_run_vllm_phase_requests( }, log_dir="/logs", pid_file="/pids.json", + block_size=128, ) await run_daser_prefix( @@ -1311,7 +1355,7 @@ def test_run_bench_entrypoint_hides_manual_cache_size_flags() -> None: def test_run_bench_entrypoint_names_backend_matrix() -> None: """The e2e benchmark entrypoint exposes the full comparison matrix.""" - runs = _expand_backend_runs("all", default_reuse_mode="chunk") + runs = _expand_backend_runs("all") assert [run.label for run in runs] == [ "baseline", @@ -1327,6 +1371,75 @@ def test_run_bench_entrypoint_names_backend_matrix() -> None: ] +def test_run_bench_parses_comma_separated_backend_rows() -> None: + """The runner accepts a comma-separated subset of benchmark rows.""" + args = parse_args( + [ + "--backend", + "baseline,lmcache,daser-prefix", + "--model", + "/models/qwen", + "--store-dir", + "/data/zwt/daser_test/bench", + ] + ) + + assert args.backend == "baseline,lmcache,daser-prefix" + assert _expand_backend_runs(args.backend) == [ + BackendRun("baseline", "vllm", "none"), + BackendRun("lmcache", "lmcache", "none"), + BackendRun("daser-prefix", "daser", "prefix"), + ] + + +def test_run_bench_backend_parser_rejects_legacy_aliases() -> None: + """Backend choices stay limited to canonical benchmark row names.""" + for backend in ("vllm", "daser", "all-openai"): + with pytest.raises(SystemExit): + parse_args( + [ + "--backend", + backend, + "--model", + "/models/qwen", + "--store-dir", + "/data/zwt/daser_test/bench", + ] + ) + + +def test_run_bench_parser_rejects_invalid_numeric_args() -> None: + """Invalid benchmark numeric knobs fail before orchestration starts.""" + with pytest.raises(SystemExit): + parse_args( + [ + "--backend", + "baseline", + "--model", + "/models/qwen", + "--store-dir", + "/data/zwt/daser_test/bench", + "--block-size", + "0", + ] + ) + with pytest.raises(SystemExit): + parse_args( + [ + "--backend", + "baseline", + "--load-generator", + "vllm-bench", + "--model", + "/models/qwen", + "--store-dir", + "/data/zwt/daser_test/bench", + "--bench-request-rate", + "nan", + ] + ) + + def test_run_bench_shell_entrypoint_is_removed() -> None: """The benchmark entrypoint should live in Python, not a shell wrapper.""" assert not (REPO_ROOT / "benchmarks" / "run_bench.sh").exists() @@ -1820,6 +1933,7 @@ def fake_run_command(command: list[str]) -> None: }, "correctness": { "cold_warm_exact_match": { + "available": True, "accuracy": 0.9, } }, @@ -1864,6 +1978,7 @@ def fake_run_command(command: list[str]) -> None: }, "correctness": { "cold_warm_exact_match": { + "available": True, "accuracy": 1.0, } }, @@ -1893,6 +2008,7 @@ def fake_run_command(command: list[str]) -> None: longbench_dir="/data/longbench", datasets="triviaqa", max_samples=1, + block_size=128, ) ) @@ -1909,6 +2025,7 @@ def fake_run_command(command: list[str]) -> None: assert "lmcache:" in captured assert "cold_ttft_ms_mean: 40.0" in captured assert "warm_ttft_ms_mean: 12.5" in captured + assert "cold_warm_exact_match_available: True" in captured assert "cold_warm_exact_match_accuracy: 0.9" in captured assert "daser-chunk:" in captured assert "cold_uploaded_documents: 1" in captured @@ -1918,6 +2035,22 @@ def fake_run_command(command: list[str]) -> None: assert "warm_answer_contains_accuracy: 1.0" in captured assert f"run_root: {run_root}" in captured assert any(command[1] == "benchmarks/bench_load.py" for command in commands) + block_size_commands = [ + command + for command in commands + if any( + item.endswith( + ( + "bench_load.py", + "bench_start_servers.py", + ) + ) + for item in command + ) + ] + assert block_size_commands + for command in block_size_commands: + assert command[command.index("--block-size") + 1] == "128" start_commands = [ command for command in commands @@ -1945,6 +2078,470 @@ def test_run_bench_stage_title_formats_backend_names() -> None: ) +def test_run_bench_explicit_openai_subset_excludes_daser_chunk() -> None: + """Comma-separated OpenAI-compatible rows can omit DaseR chunk.""" + runs = _expand_backend_runs("baseline,lmcache,daser-prefix") + + assert runs == [ + BackendRun("baseline", "vllm", "none"), + BackendRun("lmcache", "lmcache", "none"), + BackendRun("daser-prefix", "daser", "prefix"), + ] + + +def test_vllm_bench_rejects_chunk_backends() -> None: + """vLLM bench load generation cannot exercise DaseR chunk endpoints.""" + with pytest.raises(ValueError, match="daser-chunk"): + _validate_backend_runs( + _expand_backend_runs("all"), + load_generator="vllm-bench", + ) + with pytest.raises(ValueError, match="daser-chunk"): + _validate_backend_runs( + _expand_backend_runs("daser-chunk"), + load_generator="vllm-bench", + ) + + +def test_run_benchmark_validates_direct_args_before_creating_run_dir( + tmp_path: Path, +) -> None: + """Direct RunBenchArgs callers get the same preflight validation as CLI.""" + with pytest.raises(ValueError, match="block_size"): + run_benchmark( + RunBenchArgs( + backend="baseline", + model="/models/qwen", + store_dir=str(tmp_path), + block_size=0, + ) + ) + + assert list(tmp_path.iterdir()) == [] + + +def test_vllm_bench_prepare_config_uses_synthetic_lengths(tmp_path: Path) -> None: + """Synthetic vLLM bench sizing uses configured input length and block size.""" + args = RunBenchArgs( + backend="baseline,lmcache,daser-prefix", + model="/models/qwen", + store_dir=str(tmp_path), + load_generator="vllm-bench", + block_size=128, + bench_num_prompts=10, + bench_input_len=1024, + bench_output_len=1, + bench_random_prefix_len=128, + bench_random_range_ratio=0.5, + ) + + config = _bench_prepare_config(args, tmp_path) + + assert config["dataset"] == "vllm-bench-random" + assert config["num_samples"] == 10 + assert config["block_size"] == 128 + assert config["bench_input_len"] == 1024 + assert config["bench_output_len"] == 1 + assert config["max_prompt_tokens"] == 1664 + assert config["max_prompt_blocks"] == 13 + assert config["total_blocks"] == 130 + + +def test_vllm_bench_command_uses_random_dataset(tmp_path: Path) -> None: + """vLLM bench commands target completions with deterministic random load.""" + args = RunBenchArgs( + backend="daser-prefix", + model="/models/qwen", + store_dir=str(tmp_path), + load_generator="vllm-bench", + bench_num_prompts=12, + bench_input_len=4096, + bench_output_len=2, + bench_max_concurrency=4, + bench_request_rate="inf", + bench_seed=123, + bench_burstiness=2.0, + bench_random_prefix_len=256, + bench_random_range_ratio=0.25, + ) + raw_path = tmp_path / "raw.json" + + command = _vllm_bench_command( + args, + ServiceEndpoint("http://127.0.0.1:8001"), + raw_path, + ) + + assert command[:3] == ["vllm", "bench", "serve"] + assert command[command.index("--backend") + 1] == "openai" + assert command[command.index("--base-url") + 1] == "http://127.0.0.1:8001" + assert command[command.index("--endpoint") + 1] == "/v1/completions" + assert command[command.index("--dataset-name") + 1] == "random" + assert command[command.index("--num-prompts") + 1] == "12" + assert command[command.index("--input-len") + 1] == "4096" + assert command[command.index("--output-len") + 1] == "2" + assert command[command.index("--max-concurrency") + 1] == "4" + assert command[command.index("--request-rate") + 1] == "inf" + assert command[command.index("--seed") + 1] == "123" + assert command[command.index("--burstiness") + 1] == "2.0" + assert command[command.index("--temperature") + 1] == "0.0" + assert command[command.index("--top-p") + 1] == "1.0" + assert command[command.index("--percentile-metrics") + 1] == "ttft,tpot,itl,e2el" + assert "--save-detailed" in command + assert command[command.index("--random-prefix-len") + 1] == "256" + assert command[command.index("--random-range-ratio") + 1] == "0.25" + assert command[command.index("--result-filename") + 1] == raw_path.name + + +def test_vllm_bench_normalises_result_summary(tmp_path: Path) -> None: + """vLLM bench JSON metrics are converted into benchmark summaries.""" + raw = tmp_path / "raw.json" + raw.write_text( + json.dumps( + { + "completed": 10, + "total_input_tokens": 10240, + "total_output_tokens": 10, + "benchmark_duration": 2.5, + "mean_ttft_ms": 100.0, + "mean_tpot_ms": 3.0, + "mean_e2el_ms": 120.0, + } + ) + ) + + summary = _normalise_vllm_bench_result(raw) + + assert summary["num_requests"] == 10 + assert summary["ttft_ms_mean"] == 100.0 + assert summary["latency_ms_mean"] == 120.0 + assert summary["phase_elapsed_ms"] == 2500.0 + assert summary["phase_prompt_tok_per_s"] == 4096.0 + assert summary["prompt_tokens_total"] == 10240 + assert summary["completion_tokens_total"] == 10 + + +def test_vllm_bench_compares_detailed_outputs(tmp_path: Path) -> None: + """Detailed vLLM bench JSON enables cold/warm exact-match correctness.""" + cold = tmp_path / "cold.json" + warm = tmp_path / "warm.json" + cold.write_text( + json.dumps( + { + "completed": 3, + "generated_texts": ["A", "B", "C"], + "errors": [None, None, None], + } + ) + ) + warm.write_text( + json.dumps( + { + "completed": 3, + "generated_texts": ["A", "X", "C"], + "errors": [None, None, None], + } + ) + ) + + correctness = _compare_vllm_bench_outputs(cold, warm) + + assert correctness == { + "cold_warm_exact_match": { + "available": True, + "matches": 2, + "total": 3, + "accuracy": 2 / 3, + "length_mismatch": False, + } + } + + +def test_vllm_bench_compares_outputs_detail_shape(tmp_path: Path) -> None: + """vLLM bench output-list details also support correctness comparison.""" + cold = tmp_path / "cold.json" + warm = tmp_path / "warm.json" + cold.write_text( + json.dumps({"outputs": [{"generated_text": "A"}, {"generated_text": "B"}]}) + ) + warm.write_text( + json.dumps({"outputs": [{"generated_text": "A"}, {"generated_text": "X"}]}) + ) + + correctness = _compare_vllm_bench_outputs(cold, warm) + + assert correctness["cold_warm_exact_match"]["available"] is True + assert correctness["cold_warm_exact_match"]["matches"] == 1 + assert correctness["cold_warm_exact_match"]["total"] == 2 + assert correctness["cold_warm_exact_match"]["accuracy"] == 0.5 + assert correctness["cold_warm_exact_match"]["length_mismatch"] is False + + +def test_vllm_bench_correctness_counts_length_mismatch(tmp_path: Path) -> None: + """Missing cold/warm outputs should lower exact-match accuracy.""" + cold = tmp_path / "cold.json" + warm = tmp_path / "warm.json" + 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) + + assert correctness["cold_warm_exact_match"]["available"] is True + assert correctness["cold_warm_exact_match"]["matches"] == 2 + assert correctness["cold_warm_exact_match"]["total"] == 3 + assert correctness["cold_warm_exact_match"]["accuracy"] == 2 / 3 + assert correctness["cold_warm_exact_match"]["length_mismatch"] is True + + +def test_vllm_bench_correctness_marks_missing_details_unavailable( + tmp_path: Path, +) -> None: + """Missing detailed outputs should not look like a zero-sample pass.""" + cold = tmp_path / "cold.json" + warm = tmp_path / "warm.json" + cold.write_text(json.dumps({"completed": 2})) + warm.write_text(json.dumps({"completed": 2})) + + correctness = _compare_vllm_bench_outputs(cold, warm) + + assert correctness["cold_warm_exact_match"]["available"] is False + assert correctness["cold_warm_exact_match"]["total"] == 0 + assert correctness["cold_warm_exact_match"]["accuracy"] is None + assert "reason" in correctness["cold_warm_exact_match"] + + +def test_vllm_bench_phase_metrics_report_backend_token_hit_rate(monkeypatch) -> None: + """vLLM bench phases record backend token hit rates from metric deltas.""" + before_metrics = { + "vllm_prometheus": {}, + "backend_prometheus": { + "daser_cache_requested_tokens_total": 1000, + "daser_cache_matched_tokens_total": 100, + }, + "backend_status": {}, + } + after_metrics = { + "vllm_prometheus": {}, + "backend_prometheus": { + "daser_cache_requested_tokens_total": 3000, + "daser_cache_matched_tokens_total": 1600, + "daser_external_prefix_cache_queries_total": 10, + "daser_external_prefix_cache_hits_total": 10, + 'daser_cache_lookup_total{result="hit"}': 10, + "daser_cache_lookup_total": 10, + }, + "backend_status": {}, + } + + async def fake_collect( + manifest: BenchmarkManifest, + before: dict[str, Any] | None = None, + ) -> dict[str, Any]: + del manifest + if before is None: + return before_metrics + return { + "vllm_prometheus": {}, + "backend_prometheus": { + key: value - before["backend_prometheus"].get(key, 0.0) + for key, value in after_metrics["backend_prometheus"].items() + }, + "backend_status": {}, + "hit_ratios": { + "daser_prometheus_tokens": 0.75, + "daser_external_prefix_cache_queries_total": 10, + "daser_external_prefix_cache_hits_total": 10, + "daser_external_prefix": 1.0, + "daser_prometheus_requests": 1.0, + }, + } + + manifest = BenchmarkManifest( + run_id="run1", + backend="daser", + reuse_mode="prefix", + model="/models/qwen", + store_dir="/bench", + l1_size_bytes=1024, + l2_size_bytes=2048, + skip_l2=True, + endpoints={"vllm": ServiceEndpoint("http://127.0.0.1:8001")}, + log_dir="/bench/logs", + pid_file="/bench/pids.json", + ) + monkeypatch.setattr( + "benchmarks.run_bench.collect_phase_metrics", + fake_collect, + ) + + metrics, hit_rate = _collect_vllm_bench_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 + assert metrics["hit_ratios"]["daser_prometheus_tokens"] == 0.75 + assert metrics["hit_ratios"]["daser_external_prefix"] == 1.0 + assert metrics["hit_ratios"]["daser_prometheus_requests"] == 1.0 + assert hit_rate == 0.75 + + +def test_run_bench_vllm_bench_entrypoint_runs_openai_rows( + tmp_path: Path, + monkeypatch, + capsys, +) -> None: + """vLLM bench mode starts OpenAI-compatible rows and writes a summary.""" + run_root = tmp_path / "run_20260102_030405" + commands: list[list[str]] = [] + + def fake_run_command(command: list[str]) -> None: + commands.append(command) + if any(item.endswith("bench_start_servers.py") for item in command): + store_dir = Path(command[command.index("--store-dir") + 1]) + store_dir.mkdir(parents=True, exist_ok=True) + backend = command[command.index("--backend") + 1] + reuse_mode = "none" + if backend == "daser": + reuse_mode = command[command.index("--cache-reuse-mode") + 1] + endpoints = {"vllm": {"url": "http://127.0.0.1:8001"}} + if backend == "daser": + endpoints["daser"] = {"url": "http://127.0.0.1:2026"} + (store_dir / "manifest.json").write_text( + json.dumps( + { + "run_id": "run1", + "backend": backend, + "reuse_mode": reuse_mode, + "model": "/models/qwen", + "store_dir": str(store_dir), + "l1_size_bytes": 1024, + "l2_size_bytes": 2048, + "skip_l2": True, + "endpoints": endpoints, + "log_dir": str(store_dir / "logs"), + "pid_file": str(store_dir / "pids.json"), + "block_size": 128, + } + ) + ) + return + if command[:3] == ["vllm", "bench", "serve"]: + result_dir = Path(command[command.index("--result-dir") + 1]) + result_name = command[command.index("--result-filename") + 1] + backend_label = result_dir.name + generated_texts = ["A", "B", "C"] + if ( + backend_label == "daser-prefix" + and result_name == "vllm_bench_warm.json" + ): + generated_texts = ["A", "B", "X"] + result_dir.mkdir(parents=True, exist_ok=True) + (result_dir / result_name).write_text( + json.dumps( + { + "duration": 2.0, + "completed": 3, + "failed": 0, + "total_input_tokens": 3072, + "total_output_tokens": 3, + "mean_ttft_ms": 11.0, + "mean_e2el_ms": 13.0, + "generated_texts": generated_texts, + "errors": [None, None, 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.time.strftime", + lambda _fmt: "20260102_030405", + ) + monkeypatch.setattr("benchmarks.run_bench._drain_daser", lambda _manifest: None) + monkeypatch.setattr( + "benchmarks.run_bench._probe_daser_metrics", + lambda *_args, **_kwargs: None, + ) + + async def fake_collect_phase_metrics( + manifest: BenchmarkManifest, + before_metrics: dict[str, Any] | None = None, + ) -> dict[str, Any]: + del before_metrics + if manifest.backend == "lmcache": + return { + "vllm_prometheus": {}, + "backend_prometheus": { + "lmcache_mp_lookup_requested_tokens_total": 3000.0, + "lmcache_mp_lookup_hit_tokens_total": 2400.0, + }, + "backend_status": {}, + "hit_ratios": {"lmcache_prometheus_lookup": 0.8}, + } + if manifest.backend == "daser": + return { + "vllm_prometheus": {}, + "backend_prometheus": { + "daser_cache_requested_tokens_total": 3000.0, + "daser_cache_matched_tokens_total": 2550.0, + }, + "backend_status": {}, + "hit_ratios": {"daser_prometheus_tokens": 0.85}, + } + return { + "vllm_prometheus": {}, + "backend_prometheus": {}, + "backend_status": {}, + "hit_ratios": {}, + } + + monkeypatch.setattr( + "benchmarks.run_bench.collect_phase_metrics", + fake_collect_phase_metrics, + ) + + result = run_benchmark( + RunBenchArgs( + backend="baseline,lmcache,daser-prefix", + load_generator="vllm-bench", + model="/models/qwen", + store_dir=str(tmp_path), + block_size=128, + bench_num_prompts=3, + bench_input_len=1024, + bench_output_len=1, + bench_max_concurrency=2, + ) + ) + + captured = capsys.readouterr().out + assert result == run_root + assert "load_generator: vllm-bench" in captured + assert "bench_input_len: 1024" in captured + assert "baseline_ttft_ms_mean: 11.0" in captured + assert "warm_ttft_ms_mean: 11.0" in captured + assert "warm_backend_cache_hit_rate: 0.8" in captured + assert "warm_backend_cache_hit_rate: 0.85" in captured + assert "cold_warm_exact_match_accuracy: 1.0" in captured + assert "cold_warm_exact_match_accuracy: 0.6666666666666666" in captured + assert len([cmd for cmd in commands if cmd[:3] == ["vllm", "bench", "serve"]]) == 5 + for command in commands: + if command[:3] == ["vllm", "bench", "serve"]: + assert "--save-detailed" in command + assert not any( + "daser-chunk" in str(item) for command in commands for item in command + ) + lmcache = json.loads((run_root / "lmcache" / "results.json").read_text()) + daser_prefix = json.loads((run_root / "daser-prefix" / "results.json").read_text()) + assert lmcache["correctness"]["cold_warm_exact_match"]["accuracy"] == 1.0 + assert daser_prefix["correctness"]["cold_warm_exact_match"]["accuracy"] == 2 / 3 + assert lmcache["result"]["warm"]["summary"]["backend_server_cache_hit_rate"] == 0.8 + assert ( + daser_prefix["result"]["warm"]["summary"]["backend_server_cache_hit_rate"] + == 0.85 + ) + + def test_run_bench_probes_daser_metrics_only_for_daser_backends() -> None: """Only DaseR benchmark rows need a DaseR metrics readiness probe.""" assert _should_probe_daser_metrics(BackendRun("daser", "daser", "chunk")) @@ -2023,6 +2620,31 @@ def test_vllm_start_can_override_max_num_batched_tokens(tmp_path: Path) -> None: assert command[command.index("--max-num-batched-tokens") + 1] == "32768" +def test_server_commands_propagate_custom_block_size(tmp_path: Path) -> None: + """Custom benchmark block size reaches vLLM, LMCache, and DaseR.""" + manager = ServerManager( + run_id="run1", + backend="daser", + model="/models/qwen", + store_dir=tmp_path, + gpu_id="2", + gpu_util=0.85, + max_num_seqs=32, + l1_size_bytes=1024**3, + l2_size_bytes=2 * 1024**3, + block_size=128, + ) + + vllm_command = manager.vllm_command(None) + daser_command = manager._daser_server_command() # noqa: SLF001 + lmcache_command = manager._lmcache_mp_server_command() # noqa: SLF001 + + assert vllm_command[vllm_command.index("--block-size") + 1] == "128" + assert daser_command[daser_command.index("--block-tokens") + 1] == "128" + assert lmcache_command[lmcache_command.index("--chunk-size") + 1] == "128" + assert manager.manifest().block_size == 128 + + def test_lmcache_metrics_use_http_server_endpoint() -> None: """LMCache MP metrics are exposed by the HTTP server, not port 9090.""" manifest = BenchmarkManifest(