From 72cd6fc831d2ab82ddc453a677213656ae8e523a Mon Sep 17 00:00:00 2001 From: GentleCold Date: Sat, 5 Sep 2026 06:30:52 +0800 Subject: [PATCH] feat(benchmarks): make daser prefetch opt-in - Add an explicit DaseR prefetch flag with a default worker limit of two\n- Preserve numeric worker overrides and record effective settings in manifests\n- Cover CLI precedence, legacy manifests, and connector payload propagation\n- Document the opt-in behavior in benchmark and development guides --- benchmarks/README.md | 8 ++ benchmarks/bench_start_servers.py | 27 ++++- benchmarks/run_bench.py | 39 ++++++- benchmarks/utils/servers.py | 56 +++++++++- docs/development.md | 5 + tests/unit/test_benchmark_unified_utils.py | 118 +++++++++++++++++++++ 6 files changed, 244 insertions(+), 9 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 3caaecf..6cd12a0 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -110,6 +110,14 @@ and `/infer` endpoints instead of the OpenAI completions endpoint. If fails fast and asks you to select only `baseline,lmcache,daser-prefix` or use `--load-generator internal`. +Scheduler-side DaseR prefetch is disabled by default. Add +`--daser-prefetch` to opt in; it uses two worker threads. The compatibility +override `--daser-prefetch-max-requests N` accepts any non-negative value and +takes precedence, so `--daser-prefetch-max-requests 0` explicitly disables +prefetch. Both options are available on `run_bench.py` and +`bench_start_servers.py`; the effective value and enabled state are written to +each DaseR manifest. + The vLLM-bench-specific knobs are: | Option | Meaning | diff --git a/benchmarks/bench_start_servers.py b/benchmarks/bench_start_servers.py index 7fbdff0..99ae88a 100644 --- a/benchmarks/bench_start_servers.py +++ b/benchmarks/bench_start_servers.py @@ -12,7 +12,10 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) from benchmarks.utils.constants import BLOCK_TOKENS -from benchmarks.utils.servers import ServerManager +from benchmarks.utils.servers import ( + ServerManager, + resolve_daser_prefetch_max_requests, +) from benchmarks.utils.sizing import parse_size_bytes from benchmarks.utils.system import apply_gpu_selection @@ -43,7 +46,21 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--block-size", type=int, default=BLOCK_TOKENS) 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("--daser-prefetch-max-requests", type=int, default=0) + parser.add_argument( + "--daser-prefetch", + action="store_true", + help="Enable DaseR scheduler-side prefetch (default worker limit: 2).", + ) + parser.add_argument( + "--daser-prefetch-max-requests", + type=int, + default=None, + metavar="N", + help=( + "Expert override for DaseR prefetch workers; zero explicitly " + "disables prefetch and takes precedence over --daser-prefetch." + ), + ) parser.add_argument( "--cache-reuse-mode", choices=("chunk", "prefix"), default="chunk" ) @@ -63,6 +80,10 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: async def main_async(args: argparse.Namespace) -> None: """Start services and write manifest.""" + prefetch_max_requests = resolve_daser_prefetch_max_requests( + args.daser_prefetch, + args.daser_prefetch_max_requests, + ) selected_gpu = apply_gpu_selection(args.gpu_id) or args.gpu_id run_id = args.run_id or time.strftime("%Y%m%d_%H%M%S") manager = ServerManager( @@ -88,7 +109,7 @@ async def main_async(args: argparse.Namespace) -> None: skip_l2=args.skip_l2, tensor_parallel_size=args.tensor_parallel_size, trust_remote_code=args.trust_remote_code, - daser_prefetch_max_requests=args.daser_prefetch_max_requests, + daser_prefetch_max_requests=prefetch_max_requests, ) manifest = await manager.start() print(f"manifest={args.store_dir}/manifest.json") diff --git a/benchmarks/run_bench.py b/benchmarks/run_bench.py index 98c84a8..7071349 100644 --- a/benchmarks/run_bench.py +++ b/benchmarks/run_bench.py @@ -20,7 +20,11 @@ from benchmarks.utils import vllm_bench from benchmarks.utils.constants import BLOCK_TOKENS from benchmarks.utils.datasets import add_dataset_cli_args -from benchmarks.utils.servers import BenchmarkManifest, stop_from_pid_file +from benchmarks.utils.servers import ( + BenchmarkManifest, + resolve_daser_prefetch_max_requests, + stop_from_pid_file, +) _DASER_METRICS_SETTLE_SECONDS = 2.0 _BACKEND_CLEANUP_SETTLE_SECONDS = 2.0 @@ -81,6 +85,7 @@ class RunBenchArgs: bench_seed: vLLM bench random seed. bench_burstiness: vLLM bench burstiness factor. evict: Whether to enable L2 and eviction sizing. + daser_prefetch: Explicitly enable scheduler-side DaseR prefetch. daser_prefetch_max_requests: Maximum concurrent DaseR prefetches. prometheus_url: Optional Prometheus base URL for scrape diagnostics. @@ -118,6 +123,7 @@ class RunBenchArgs: bench_seed: int = 42 bench_burstiness: float = 1.0 evict: bool = False + daser_prefetch: bool = False daser_prefetch_max_requests: int = 0 prometheus_url: str = "http://127.0.0.1:9090" @@ -190,7 +196,21 @@ def parse_args(argv: list[str] | None = None) -> RunBenchArgs: 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("--daser-prefetch-max-requests", type=int, default=0) + parser.add_argument( + "--daser-prefetch", + action="store_true", + help="Enable DaseR scheduler-side prefetch (default worker limit: 2).", + ) + parser.add_argument( + "--daser-prefetch-max-requests", + type=int, + default=None, + metavar="N", + help=( + "Expert override for DaseR prefetch workers; accepts zero to " + "explicitly disable and takes precedence over --daser-prefetch." + ), + ) parser.add_argument( "--prometheus-url", default="http://127.0.0.1:9090", @@ -200,6 +220,13 @@ def parse_args(argv: list[str] | None = None) -> RunBenchArgs: ), ) args = parser.parse_args(argv) + try: + effective_prefetch_max_requests = resolve_daser_prefetch_max_requests( + args.daser_prefetch, + args.daser_prefetch_max_requests, + ) + except ValueError as exc: + parser.error(str(exc)) parsed = RunBenchArgs( backend=args.backend, load_generator=args.load_generator, @@ -231,7 +258,8 @@ def parse_args(argv: list[str] | None = None) -> RunBenchArgs: bench_seed=args.bench_seed, bench_burstiness=args.bench_burstiness, evict=args.evict, - daser_prefetch_max_requests=args.daser_prefetch_max_requests, + daser_prefetch=args.daser_prefetch, + daser_prefetch_max_requests=effective_prefetch_max_requests, prometheus_url=args.prometheus_url, ) try: @@ -275,6 +303,8 @@ def run_benchmark(args: RunBenchArgs) -> Path: _print_kv("dataset", args.dataset) _print_kv("max_samples", args.max_samples) _print_kv("block_size", args.block_size) + _print_kv("daser_prefetch_enabled", args.daser_prefetch_max_requests > 0) + _print_kv("daser_prefetch_max_requests", args.daser_prefetch_max_requests) _print_kv("output", prepare_path) if args.load_generator in ("vllm-bench", "vllm-bench-prefix"): prepare = {"config": vllm_bench.prepare_config(args, run_root)} @@ -394,6 +424,7 @@ def _validate_run_args(args: RunBenchArgs) -> None: non_negative_ints = { "max_num_batched_tokens": args.max_num_batched_tokens, "max_context_tokens": args.max_context_tokens, + "daser_prefetch_max_requests": args.daser_prefetch_max_requests, } for name, value in non_negative_ints.items(): if value < 0: @@ -542,7 +573,7 @@ def _start_command( "--l2-size", str(derived_l2), ] - if backend_run.backend == "daser" and args.daser_prefetch_max_requests: + if backend_run.backend == "daser": command.extend( [ "--daser-prefetch-max-requests", diff --git a/benchmarks/utils/servers.py b/benchmarks/utils/servers.py index f6cf784..fd6ad9a 100644 --- a/benchmarks/utils/servers.py +++ b/benchmarks/utils/servers.py @@ -27,10 +27,38 @@ LMCACHE_HTTP_PORT = 8080 LMCACHE_MP_CONNECTOR_NAME = "DaseRBenchLMCacheMPConnector" LMCACHE_MP_CONNECTOR_MODULE = "benchmarks.utils.lmcache_connector_shim" +DEFAULT_DASER_PREFETCH_MAX_REQUESTS = 2 REPO_ROOT = Path(__file__).resolve().parents[2] LMCACHE_REPO_ROOT = REPO_ROOT.parent / "LMCache" +def resolve_daser_prefetch_max_requests( + prefetch_enabled: bool, + override: int | None, +) -> int: + """Resolve the effective scheduler prefetch worker limit. + + Args: + prefetch_enabled: Whether the explicit boolean opt-in was supplied. + override: Optional numeric compatibility override. Zero explicitly + disables prefetch even when the boolean flag is present. + + Returns: + Effective non-negative worker limit. The opt-in default is two. + + Raises: + ValueError: If ``override`` is negative. + + Thread-safety: + Pure function; safe to call from CLI parsing or service setup. + """ + if override is not None: + if override < 0: + raise ValueError("daser prefetch worker limit must be non-negative") + return override + return DEFAULT_DASER_PREFETCH_MAX_REQUESTS if prefetch_enabled else 0 + + @dataclass(frozen=True) class ServiceEndpoint: """HTTP endpoint for a benchmark service. @@ -62,6 +90,8 @@ class BenchmarkManifest: log_dir: Log directory. pid_file: JSON file containing subprocess PIDs. block_size: vLLM KV block size in tokens. + prefetch_enabled: Whether scheduler prefetch is effectively enabled. + prefetch_max_requests: Effective scheduler prefetch worker limit. Thread-safety: Immutable value object. @@ -79,6 +109,8 @@ class BenchmarkManifest: log_dir: str pid_file: str block_size: int = BLOCK_TOKENS + prefetch_enabled: bool = False + prefetch_max_requests: int = 0 def write(self, path: str | Path) -> None: """Write manifest JSON atomically enough for local benchmark use.""" @@ -104,6 +136,10 @@ def read(cls, path: str | Path) -> "BenchmarkManifest": } payload["endpoints"] = endpoints payload.setdefault("block_size", 16) # Legacy manifests predate this field. + payload.setdefault( + "prefetch_enabled", bool(payload.get("prefetch_max_requests", 0)) + ) + payload.setdefault("prefetch_max_requests", 0) return cls(**payload) @@ -161,6 +197,8 @@ def __init__( """ if tensor_parallel_size <= 0: raise ValueError("tensor_parallel_size must be positive") + if daser_prefetch_max_requests < 0: + raise ValueError("daser_prefetch_max_requests must be non-negative") self.run_id = run_id self.backend = backend self.model = model @@ -234,6 +272,8 @@ def manifest(self) -> BenchmarkManifest: log_dir=str(self.log_dir), pid_file=str(self.pid_file), block_size=self.block_size, + prefetch_enabled=self.daser_prefetch_max_requests > 0, + prefetch_max_requests=self.daser_prefetch_max_requests, ) async def start_lmcache_mp_server(self) -> None: @@ -347,7 +387,19 @@ def lmcache_process_env(self) -> dict[str, str]: async def start_vllm_daser(self) -> None: """Start vLLM with DaseR connector.""" - kv_config = { + await self._start_vllm("vllm_daser.log", self.daser_kv_transfer_config()) + + def daser_kv_transfer_config(self) -> dict[str, Any]: + """Return the vLLM KV transfer payload for a DaseR service. + + Returns: + JSON-serialisable vLLM connector configuration, including the + effective prefetch state used by the scheduler role. + + Thread-safety: + Pure calculation over immutable service configuration. + """ + return { "kv_connector": "DaserConnector", "kv_connector_module_path": "daser.connector.daser_connector", "kv_role": "kv_both", @@ -355,9 +407,9 @@ async def start_vllm_daser(self) -> None: "socket_path": str(self.socket_path), "cache_reuse_mode": self.reuse_mode, "prefetch_max_requests": self.daser_prefetch_max_requests, + "prefetch_enabled": self.daser_prefetch_max_requests > 0, }, } - await self._start_vllm("vllm_daser.log", kv_config) async def start_daser_server(self) -> None: """Start DaseR HTTP + IPC server.""" diff --git a/docs/development.md b/docs/development.md index 7f6ba75..a4f3c55 100644 --- a/docs/development.md +++ b/docs/development.md @@ -187,6 +187,11 @@ under one run directory. Use a specific backend name to run only one row of the matrix; `--backend daser` still honors `--cache-reuse-mode` for compatibility. `--cache-reuse-mode` is only passed to DaseR rows; baseline and LMCache use the same prepared prompts and record `reuse_mode: none`. +Scheduler-side DaseR prefetch is off by default. Pass `--daser-prefetch` to +enable it with two workers, or use the compatibility override +`--daser-prefetch-max-requests N` (including `0` to force it off). The numeric +override takes precedence over the boolean flag, and the effective state is +recorded in the DaseR manifest. The Python entry point prints stage separators for prepare, backend start, cold/warm load, a final `== COMPARISON SUMMARY ==` with cold, warm, correctness, elapsed-time, and throughput fields, and the final `run_root`. diff --git a/tests/unit/test_benchmark_unified_utils.py b/tests/unit/test_benchmark_unified_utils.py index 3fb7612..50024ed 100644 --- a/tests/unit/test_benchmark_unified_utils.py +++ b/tests/unit/test_benchmark_unified_utils.py @@ -22,6 +22,7 @@ _serialise_phase, _should_dedup_context, ) +from benchmarks.bench_start_servers import parse_args as parse_start_server_args from benchmarks.run_bench import ( BackendRun, RunBenchArgs, @@ -77,6 +78,7 @@ BenchmarkManifest, ServerManager, ServiceEndpoint, + resolve_daser_prefetch_max_requests, ) from benchmarks.utils.sizing import ( BenchmarkCapacityLimits, @@ -686,6 +688,48 @@ def test_manifest_read_defaults_legacy_block_size(tmp_path: Path) -> None: loaded = BenchmarkManifest.read(path) assert loaded.block_size == 16 + assert loaded.prefetch_enabled is False + assert loaded.prefetch_max_requests == 0 + + +def test_bench_start_prefetch_parser_uses_same_precedence() -> None: + """Direct service startup accepts the opt-in and numeric override.""" + common = [ + "--backend", + "daser", + "--model", + "/models/qwen", + "--store-dir", + "/data/zwt/daser_test/bench", + ] + + disabled = parse_start_server_args(common) + enabled = parse_start_server_args([*common, "--daser-prefetch"]) + explicit_zero = parse_start_server_args( + [*common, "--daser-prefetch", "--daser-prefetch-max-requests", "0"] + ) + + assert ( + resolve_daser_prefetch_max_requests( + disabled.daser_prefetch, + disabled.daser_prefetch_max_requests, + ) + == 0 + ) + assert ( + resolve_daser_prefetch_max_requests( + enabled.daser_prefetch, + enabled.daser_prefetch_max_requests, + ) + == 2 + ) + assert ( + resolve_daser_prefetch_max_requests( + explicit_zero.daser_prefetch, + explicit_zero.daser_prefetch_max_requests, + ) + == 0 + ) def test_daser_noevict_start_uses_l1_only_mode(tmp_path: Path) -> None: @@ -1651,6 +1695,80 @@ def test_run_bench_backend_parser_rejects_legacy_aliases() -> None: ) +def test_daser_prefetch_cli_defaults_off_and_opt_in_uses_two_workers() -> None: + """The boolean prefetch flag is opt-in and resolves to two workers.""" + common = [ + "--backend", + "daser-prefix", + "--model", + "/models/qwen", + "--store-dir", + "/data/zwt/daser_test/bench", + ] + + disabled = parse_args(common) + enabled = parse_args([*common, "--daser-prefetch"]) + explicit_zero = parse_args( + [*common, "--daser-prefetch", "--daser-prefetch-max-requests", "0"] + ) + + assert disabled.daser_prefetch is False + assert disabled.daser_prefetch_max_requests == 0 + assert enabled.daser_prefetch is True + assert enabled.daser_prefetch_max_requests == 2 + assert explicit_zero.daser_prefetch is True + assert explicit_zero.daser_prefetch_max_requests == 0 + + +def test_daser_prefetch_rejects_negative_override() -> None: + """Negative prefetch worker overrides fail during CLI parsing.""" + with pytest.raises(SystemExit): + parse_args( + [ + "--backend", + "daser-prefix", + "--model", + "/models/qwen", + "--store-dir", + "/data/zwt/daser_test/bench", + "--daser-prefetch-max-requests", + "-1", + ] + ) + + +def test_daser_prefetch_resolution_and_manifest_payload() -> None: + """Effective prefetch state reaches the manifest and vLLM payload.""" + assert resolve_daser_prefetch_max_requests(False, None) == 0 + assert resolve_daser_prefetch_max_requests(True, None) == 2 + assert resolve_daser_prefetch_max_requests(True, 4) == 4 + assert resolve_daser_prefetch_max_requests(True, 0) == 0 + + manager = ServerManager( + run_id="run1", + backend="daser", + model="/models/qwen", + store_dir="/data/zwt/daser_test/bench", + gpu_id="0", + gpu_util=0.85, + max_num_seqs=8, + l1_size_bytes=1024, + l2_size_bytes=2048, + daser_prefetch_max_requests=2, + ) + manifest = manager.manifest() + payload = manager.daser_kv_transfer_config() + + assert manifest.prefetch_enabled is True + assert manifest.prefetch_max_requests == 2 + assert payload["kv_connector_extra_config"] == { + "socket_path": "/data/zwt/daser_test/bench/daser.sock", + "cache_reuse_mode": "chunk", + "prefetch_max_requests": 2, + "prefetch_enabled": True, + } + + def test_run_bench_parser_rejects_invalid_numeric_args() -> None: """Invalid benchmark numeric knobs fail before orchestration starts.""" with pytest.raises(SystemExit):