vllm-benchmark 0.1.0__py3-none-any.whl

vllm_bench/__init__.py

@@ -0,0 +1,3 @@
+"""Script-driven vLLM benchmarking."""
+
+__version__ = "0.1.0"

vllm_bench/artifacts.py

@@ -0,0 +1,126 @@
+"""Number result directories and save the exact experiment inputs."""
+
+import hashlib
+import os
+import re
+import subprocess
+import time
+import zipfile
+from pathlib import Path
+
+import yaml
+
+from . import __version__
+from .contracts import Trial
+from .json_io import write_json
+from .progress import log
+
+EXCLUDED = {
+    ".git",
+    ".venv",
+    "venv",
+    "env",
+    "node_modules",
+    "results",
+    ".vllm-bench",
+    "__pycache__",
+    ".cache",
+    ".pytest_cache",
+    ".ruff_cache",
+    "build",
+    "dist",
+    ".wheels",
+}
+
+
+def project_files(root):
+    try:
+        names = (
+            subprocess.check_output(
+                ["git", "ls-files", "-z", "--cached", "--others", "--exclude-standard"],
+                cwd=root,
+                stderr=subprocess.DEVNULL,
+            )
+            .decode()
+            .split("\0")
+        )
+    except subprocess.CalledProcessError:
+        names = []
+        for folder, directories, files in os.walk(root):
+            directories[:] = [name for name in directories if name not in EXCLUDED]
+            names.extend(str((Path(folder) / name).relative_to(root)) for name in files)
+    for name in sorted(set(names)):
+        path = root / name
+        if name == ".vllm_bench/state.json":
+            continue
+        if not name or not path.is_file() or path.is_symlink():
+            continue
+        if any(
+            part in EXCLUDED or part.startswith(".env") or part.endswith(".egg-info")
+            for part in Path(name).parts
+        ):
+            continue
+        if any((root / parent).is_symlink() for parent in Path(name).parents):
+            continue
+        if path.suffix not in {".pyc", ".so", ".dylib", ".whl"}:
+            yield path, name
+
+
+def trials(config):
+    return [
+        Trial(benchmark=name, scenario=scenario, repeat=repeat)
+        for name, benchmark in config.benchmarks.items()
+        for scenario in benchmark.scenarios
+        for repeat in range(1, config.runs + 1)
+    ]
+
+
+def prepare_run(config, root):
+    parent = root / "results" / f"{config.provider}-{config.name}"
+    parent.mkdir(parents=True, exist_ok=True)
+    last = max(
+        (
+            int(match[1])
+            for path in parent.iterdir()
+            if path.is_dir() and (match := re.fullmatch(r"run(\d+)", path.name))
+        ),
+        default=0,
+    )
+    output = parent / f"run{last + 1:03d}"
+    output.mkdir()
+    log(f"Result directory: {output}")
+    (output / "config.yaml").write_text(
+        yaml.safe_dump(config.model_dump(exclude_none=True), sort_keys=False)
+    )
+    if not config.skip_volume_mount:
+        archive = output / "project.zip"
+        log("Saving project snapshot (excluding results, environments, and caches)...")
+        count = 0
+        last_update = time.monotonic()
+        with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as bundle:
+            for path, name in project_files(root):
+                bundle.write(path, name)
+                count += 1
+                if time.monotonic() - last_update >= 5:
+                    log(f"Project snapshot: {count} files archived so far...")
+                    last_update = time.monotonic()
+        with archive.open("rb") as stream:
+            digest = hashlib.file_digest(stream, "sha256").hexdigest()
+        (output / "project.sha256").write_text(f"{digest}  project.zip\n")
+        log(
+            f"Snapshot saved: {count} files, {archive.stat().st_size / 1024**2:.1f} MiB"
+        )
+    else:
+        log("Project snapshot skipped: skip_volume_mount is enabled.")
+    write_json(
+        output / "run.json",
+        {
+            "schema_version": 1,
+            "cli_version": __version__,
+            "provider": config.provider,
+            "image": config.base_image,
+            "trials": [trial.model_dump() for trial in trials(config)],
+        },
+    )
+    log("Saved effective config and run manifest.")
+    return output

vllm_bench/bench_util.py

@@ -0,0 +1,124 @@
+"""The small public interface used by startup and benchmark scripts."""
+
+import json
+import os
+import subprocess
+from functools import lru_cache
+from pathlib import Path
+
+
+@lru_cache(maxsize=1)
+def _context():
+    return json.loads(Path(os.environ["VLLM_BENCH_CONTEXT"]).read_text())
+
+
+def get_global_args():
+    return dict(_context()["global_args"])
+
+
+def get_service_args(name=None):
+    """Global arguments with overrides for the current or named service group."""
+    if name is None and (service := get_service()):
+        name = service["group"]
+    overrides = _context()["service_args"][name] if name is not None else {}
+    return {**get_global_args(), **overrides}
+
+
+def get_global_arg(name, default=None):
+    return _context()["global_args"].get(name, default)
+
+
+def get_scenario_arg(name, default=None):
+    return _context()["scenario_args"].get(name, default)
+
+
+def get_benchmark_arg(name, default=None):
+    return _context()["benchmark_args"].get(name, default)
+
+
+def get_output_dir():
+    path = Path(_context()["output_dir"])
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def get_project_dir():
+    return Path(_context()["project_dir"])
+
+
+def get_service():
+    """Name, role, and replica of the calling startup script (None on the client)."""
+    return _context().get("service")
+
+
+def get_service_url(name, *, internal=False):
+    services = _context()["internal_services" if internal else "services"]
+    matches = [
+        item for item in services if item["group"] == name or item["name"] == name
+    ]
+    if len(matches) != 1:
+        raise ValueError(f"Expected one service named {name}; found {len(matches)}")
+    return matches[0]["url"]
+
+
+def get_leader_url():
+    return next(
+        item["url"] for item in _context()["services"] if item["role"] == "leader"
+    )
+
+
+def get_worker_urls():
+    return [item["url"] for item in _context()["services"] if item["role"] == "worker"]
+
+
+def get_service_path(path):
+    """Translate a client project/trial path into the path visible to services."""
+    path = Path(path).resolve()
+    context = _context()
+    if context.get("external_services"):
+        # External services may share the client's filesystem; no upload is implied.
+        return str(path)
+    for local, remote in [
+        (context["trial_dir"], "/run/vllm-bench/trial"),
+        (context["project_dir"], context["service_project_dir"]),
+    ]:
+        local = Path(local).resolve()
+        if path.is_relative_to(local):
+            destination = str(Path(remote) / path.relative_to(local))
+            if sandbox_id := context.get("modal_sandbox_id"):
+                from .providers.modal import upload_path
+
+                upload_path(sandbox_id, path, destination)
+            return destination
+    raise ValueError(f"Path must be inside the project or trial directory: {path}")
+
+
+@lru_cache(maxsize=1)
+def get_gpu_count():
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "--query-gpu=uuid", "--format=csv,noheader"],
+            capture_output=True,
+            text=True,
+            check=True,
+            timeout=5,
+        )
+        return len(result.stdout.strip().splitlines())
+    except (OSError, subprocess.SubprocessError):
+        return 0
+
+
+def is_gpu_available():
+    return get_gpu_count() > 0
+
+
+def is_cpu():
+    return not is_gpu_available()
+
+
+def __getattr__(name):
+    if name == "current_run":
+        return _context()["repeat"]
+    if name == "current_benchmark":
+        return _context()["benchmark"]
+    raise AttributeError(name)

vllm_bench/cli.py

@@ -0,0 +1,336 @@
+"""The user's machine owns scheduling, benchmark scripts, metrics, and reports."""
+
+import argparse
+import json
+import shlex
+import signal
+import sys
+import time
+from contextlib import nullcontext
+from pathlib import Path
+
+from . import __version__
+from .artifacts import prepare_run, trials
+from .config import load_config, resolve_config, select_benchmarks
+from .external import prepare_context, service_endpoints
+from .progress import log
+from .providers import get_provider, stop_saved
+from .runner import run_only, run_trial, start_environment, wait_healthy
+
+
+def interrupt(signum, frame):
+    raise KeyboardInterrupt
+
+
+def run_only_command(config, config_path, root, trial, endpoints):
+    command = [
+        "vllm_bench",
+        "run_only",
+        "--current-dir",
+        str(root),
+        "--config",
+        str(config_path),
+        "--benchmark",
+        trial.benchmark,
+        "--scenario",
+        trial.scenario,
+    ]
+    urls = {}
+    for item in endpoints:
+        urls.setdefault(item["group"], item["url"])
+    for name, url in urls.items():
+        command.extend(["--service-url", f"{name}={url}"])
+    if dataset := config.benchmarks[trial.benchmark].script_args.get("dataset_path"):
+        command.extend(["--dataset-path", str(dataset)])
+    return shlex.join(command)
+
+
+def show_environment(
+    provider, endpoints, directory, timeout, started, detach, run_command
+):
+    ready = False
+    try:
+        logs = nullcontext() if detach else provider.follow_logs()
+        with logs as process:
+            wait_healthy(provider, endpoints, timeout)
+            log(f"Environment ready in {time.monotonic() - started:.1f}s.")
+            log("Run the benchmark from another terminal:")
+            print(run_command, flush=True)
+            for item in endpoints:
+                log(f"{item['name']} ({item['role']}): {item['url']}")
+            log(f"Service output: {directory / 'services'}")
+            ready = True
+            if detach:
+                log("Containers are running. Stop with: vllm_bench down")
+                return 0
+            log("Following container logs. Press Ctrl+C to stop the environment.")
+            while process.poll() is None:
+                provider.check_running()
+                time.sleep(1)
+            if process.returncode:
+                raise RuntimeError(
+                    f"Container log stream exited with {process.returncode}"
+                )
+            return 0
+    finally:
+        if not detach or not ready:
+            log("Stopping the environment and saving container logs...")
+            provider.stop()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run script-defined vLLM benchmarks.")
+    parser.add_argument("--version", action="version", version=__version__)
+    parser.add_argument("--config", type=Path, default=Path("local_cpu.yaml"))
+    parser.add_argument(
+        "--current-dir",
+        type=Path,
+        default=Path.cwd(),
+        help="Project root containing .vllm_bench",
+    )
+    commands = parser.add_subparsers(dest="command", required=True)
+    up = commands.add_parser("up", help="Start services and follow their logs")
+    up.add_argument(
+        "-d",
+        "--detach",
+        action="store_true",
+        help="Leave services running in the background",
+    )
+    down = commands.add_parser(
+        "down", help="Stop the active environment and collect output"
+    )
+    run = commands.add_parser(
+        "run_scenarios",
+        help="Run selected trials with fresh managed services each time",
+    )
+    single = commands.add_parser(
+        "run_only", help="Run one benchmark/scenario once against existing service URLs"
+    )
+    single.add_argument("--benchmark", required=True)
+    single.add_argument("--scenario", required=True)
+    single.add_argument(
+        "--service-url",
+        action="append",
+        required=True,
+        metavar="NAME=URL",
+        help="Existing service base URL; repeat for each required service",
+    )
+    single.add_argument(
+        "--dataset-path",
+        help="Existing dataset path visible to the services (passed to the benchmark)",
+    )
+    report = commands.add_parser("report", help="Regenerate reports from saved files")
+    report.add_argument("directory", type=Path)
+    for command in [up, down, run, single, report]:
+        command.add_argument(
+            "--current-dir",
+            type=Path,
+            default=argparse.SUPPRESS,
+            help="Project root containing .vllm_bench",
+        )
+    for command in [up, run, single]:
+        command.add_argument(
+            "--config",
+            type=Path,
+            default=argparse.SUPPRESS,
+            help="Config name or path, merged over base.yaml (e.g. modal)",
+        )
+    for command in [up, run]:
+        command.add_argument(
+            "--benchmark",
+            "--benchmarks",
+            dest="benchmarks",
+            nargs="+",
+            required=True,
+            help="Benchmark names, or '*' for all benchmarks",
+        )
+        command.add_argument(
+            "--scenario",
+            "--scenarios",
+            dest="scenarios",
+            nargs="+",
+            required=True,
+            help="Scenario names, or '*' for every scenario in each selected benchmark",
+        )
+    args = parser.parse_args()
+    signal.signal(signal.SIGTERM, interrupt)
+    root = args.current_dir.resolve()
+    config_path = resolve_config(root / ".vllm_bench/configs" / args.config)
+    state_path = root / ".vllm_bench/state.json"
+    try:
+        if args.command == "report":
+            from .reports import generate_reports
+
+            generate_reports((root / args.directory).resolve())
+            return 0
+        if not (root / ".vllm_bench").is_dir():
+            raise FileNotFoundError(f"Project must contain .vllm_bench: {root}")
+        if args.command == "down":
+            stop_saved(state_path)
+            return 0
+        started = time.monotonic()
+        if args.command == "up":
+            log(f"Reading config: {config_path}")
+        external = args.command == "run_only"
+        config = select_benchmarks(
+            load_config(config_path),
+            [args.benchmark] if external else args.benchmarks,
+            [args.scenario] if external else args.scenarios,
+        )
+        if external:
+            config = config.model_copy(
+                update={
+                    "provider": "external",
+                    "runs": 1,
+                    "skip_volume_mount": True,
+                }
+            )
+            if args.dataset_path:
+                benchmark = config.benchmarks[args.benchmark]
+                benchmark.script_args = {
+                    **benchmark.script_args,
+                    "dataset_path": args.dataset_path,
+                }
+        experiment = trials(config)
+        first = experiment[0]
+        if args.command in {"up", "run_only"}:
+            pairs = sum(
+                len(benchmark.scenarios) for benchmark in config.benchmarks.values()
+            )
+            if pairs != 1:
+                raise ValueError(
+                    f"{args.command} requires exactly one benchmark and one scenario"
+                )
+        if external:
+            endpoints = service_endpoints(config, first, args.service_url)
+        if args.command == "up":
+            log(f"Configuration valid. Provider: {config.provider}")
+            log(f"Starting benchmark '{first.benchmark}', scenario '{first.scenario}'.")
+            if model := config.global_args.get("model"):
+                log(f"Model: {model}")
+        scripts = [benchmark.script for benchmark in config.benchmarks.values()]
+        if not external:
+            if config.provider == "external":
+                raise ValueError("Use run_only for external services")
+            if not config.base_image:
+                raise ValueError(f"{args.command} requires base_image in the config")
+            groups = dict.fromkeys(
+                name
+                for benchmark_name, benchmark in config.benchmarks.items()
+                for scenario in benchmark.scenarios
+                for name in config.get_scenario(benchmark_name, scenario).services
+            )
+            for name in groups:
+                script = config.services[name].startup_script
+                if not script:
+                    raise ValueError(
+                        f"Managed service '{name}' requires startup_script"
+                    )
+                if not config.skip_volume_mount:
+                    scripts.append(script)
+        for script in scripts:
+            if not (root / script).is_file():
+                raise FileNotFoundError(f"Script not found: {root / script}")
+        if external:
+            from .reports import generate_reports
+
+            output = prepare_run(config, root)
+            directory = output / first.directory
+            context = prepare_context(config, first, root, directory, endpoints)
+            try:
+                result = run_only(config, root, first, directory, context, endpoints)
+                if result.status == "cancelled":
+                    return 130
+                return 0 if result.status == "ok" else 1
+            finally:
+                generate_reports(output)
+        if state_path.exists():
+            if args.command == "up":
+                state = json.loads(state_path.read_text())
+                directory = Path(state["directory"])
+                active = json.loads((directory / ".provider/client.json").read_text())
+                if (
+                    state.get("provider", "local") != config.provider
+                    or state.get("image") != config.base_image
+                    or (active["benchmark"], active["scenario"])
+                    != (first.benchmark, first.scenario)
+                    or active["global_args"] != config.global_args
+                    or active.get("service_args", {})
+                    != {
+                        group: config.services[group].service_args
+                        for group in config.get_scenario(
+                            first.benchmark, first.scenario
+                        ).services
+                    }
+                    or active["scenario_args"]
+                    != config.get_scenario(first.benchmark, first.scenario).args
+                ):
+                    raise ValueError(
+                        f"An environment for {active['benchmark']}/"
+                        f"{active['scenario']} "
+                        "is already active with different selection "
+                        "or environment settings. "
+                        "Run 'vllm_bench down' before switching."
+                    )
+                log(f"Attaching to {first.benchmark}/{first.scenario}: {directory}")
+                endpoints = json.loads((directory / "endpoints.json").read_text())
+                return show_environment(
+                    get_provider(config, root),
+                    endpoints,
+                    directory,
+                    config.startup_timeout,
+                    started,
+                    args.detach,
+                    run_only_command(config, config_path, root, first, endpoints),
+                )
+            stop_saved(state_path)
+        output = prepare_run(config, root)
+        provider = get_provider(config, root)
+        if args.command == "up":
+            trial = experiment[0]
+            directory = output / trial.directory
+            directory.mkdir(parents=True)
+            try:
+                _, endpoints = start_environment(
+                    config, provider, trial, directory, wait=False
+                )
+            except BaseException:
+                log("Startup did not complete. Stopping containers and saving logs...")
+                provider.stop()
+                log(f"Startup output: {directory}")
+                raise
+            return show_environment(
+                provider,
+                endpoints,
+                directory,
+                config.startup_timeout,
+                started,
+                args.detach,
+                run_only_command(config, config_path, root, first, endpoints),
+            )
+        from .reports import generate_reports
+
+        failed = False
+        try:
+            for trial in experiment:
+                result = run_trial(
+                    config, root, provider, trial, output / trial.directory
+                )
+                if result.status == "cancelled":
+                    return 130
+                failed = failed or result.status != "ok"
+                if state_path.exists():
+                    # Preserve failed cleanup state for recovery; stop scheduling.
+                    return 1
+            return 1 if failed else 0
+        finally:
+            generate_reports(output)
+    except KeyboardInterrupt:
+        return 130
+    except Exception as exc:
+        print(f"Error: {exc}", file=sys.stderr)
+        return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())