sskj/src/sskj/bench/runner.py

"""bench_serving invocation for native, docker, and server-container clients."""

from __future__ import annotations

import shlex
import subprocess
import sys
import time
import urllib.request
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any


@dataclass
class BenchClientOptions:
    backend: str
    host: str
    port: int
    model: str
    tokenizer: str | None = None
    dataset_name: str = "random"
    dataset_path: str | None = None
    container_dataset_path: str | None = None
    random_range_ratio: float = 1.0
    warmup_max_requests: int = 0
    flush_cache: bool = False
    client_mode: str = "auto"
    client_image: str | None = None
    client_python: str | None = None
    server_container: str | None = None
    container_python: str | None = None
    output_file: Path | None = None
    detail_log: Path | None = None
    timeout_s: int = 1800
    root: Path | None = None
    extra_env: dict[str, str] = field(default_factory=dict)


def _bench_args(scenario: dict[str, Any], options: BenchClientOptions) -> list[str]:
    args = [
        "--backend",
        options.backend,
        "--host",
        options.host,
        "--port",
        str(options.port),
        "--model",
        options.model,
    ]
    if options.tokenizer:
        args += ["--tokenizer", options.tokenizer]
    args += [
        "--dataset-name",
        options.dataset_name,
        "--random-input-len",
        str(scenario["input_len"]),
        "--random-output-len",
        str(scenario["output_len"]),
        "--random-range-ratio",
        str(options.random_range_ratio),
        "--num-prompts",
        str(scenario["num_prompts"]),
        "--max-concurrency",
        str(scenario["concurrency"]),
        "--request-rate",
        "10000",
        "--output-details",
    ]
    if options.warmup_max_requests and options.warmup_max_requests > 0:
        args += ["--warmup-requests", str(options.warmup_max_requests)]
    if options.flush_cache:
        args += ["--flush-cache"]

    dataset_path = (
        options.container_dataset_path
        if options.client_mode == "server-container"
        else options.dataset_path
    )
    if options.dataset_name == "random" and dataset_path:
        args += ["--dataset-path", dataset_path]
    elif options.dataset_name != "random":
        args += ["--tokenize-prompt"]
    return args


def _run(argv: list[str], log_path: Path | None, timeout_s: int) -> int:
    log_path.parent.mkdir(parents=True, exist_ok=True)
    try:
        with open(log_path, "wb") as f:
            proc = subprocess.run(argv, stdout=f, stderr=subprocess.STDOUT, timeout=timeout_s)
            return proc.returncode
    except subprocess.TimeoutExpired:
        with open(log_path, "ab") as f:
            f.write(b"\n[sskj] bench_serving timed out\n")
        return 124


def run_scenario(scenario: dict[str, Any], options: BenchClientOptions) -> int:
    if options.output_file is None or options.detail_log is None:
        raise ValueError("output_file and detail_log are required")

    base_args = _bench_args(scenario, options)
    if options.client_mode == "native":
        python_bin = options.client_python or sys.executable
        return _run(
            [python_bin, "-m", "sglang.bench_serving", *base_args, "--output-file", str(options.output_file)],
            options.detail_log,
            options.timeout_s,
        )

    if options.client_mode == "docker":
        if not options.client_image:
            raise SystemExit("docker client mode requires --client-image or DOCKER_CLIENT_IMAGE")
        cmd = ["docker", "run", "--rm", "--network", "host"]
        if options.root:
            cmd += ["-v", f"{options.root}:{options.root}"]
        for mount in (options.model, options.tokenizer):
            if mount and Path(mount).exists():
                cmd += ["-v", f"{mount}:{mount}:ro"]
        if options.dataset_path and Path(options.dataset_path).exists():
            cmd += ["-v", f"{options.dataset_path}:{options.dataset_path}:ro"]
        if options.output_file:
            cmd += ["-v", f"{options.output_file.parent}:{options.output_file.parent}"]
        cmd += [
            "-e",
            "HF_HUB_OFFLINE=1",
            "-e",
            "TRANSFORMERS_OFFLINE=1",
            "-e",
            "HF_DATASETS_OFFLINE=1",
            "-e",
            "TORCH_DEVICE_BACKEND_AUTOLOAD=0",
        ]
        cmd += [
            options.client_image,
            "python",
            "-m",
            "sglang.bench_serving",
            *base_args,
            "--output-file",
            str(options.output_file),
        ]
        return _run(cmd, options.detail_log, options.timeout_s)

    if options.client_mode == "server-container":
        if not options.server_container or not options.container_python:
            raise SystemExit("server-container client mode requires --server-container and --container-python")
        container_output = f"/tmp/bench_outputs/{options.output_file.name}"
        mkdir = subprocess.run(
            ["docker", "exec", options.server_container, "mkdir", "-p", "/tmp/bench_outputs"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        if mkdir.returncode != 0:
            raise SystemExit(
                f"cannot prepare {options.server_container}:/tmp/bench_outputs (docker exec mkdir rc={mkdir.returncode})"
            )
        cmd = [
            "docker",
            "exec",
            options.server_container,
            "env",
            "HF_HUB_OFFLINE=1",
            "TRANSFORMERS_OFFLINE=1",
            "HF_DATASETS_OFFLINE=1",
            options.container_python,
            "-m",
            "sglang.bench_serving",
            *base_args,
            "--output-file",
            container_output,
        ]
        rc = _run(cmd, options.detail_log, options.timeout_s)
        if rc == 0:
            options.output_file.parent.mkdir(parents=True, exist_ok=True)
            cp = subprocess.run(
                ["docker", "cp", f"{options.server_container}:{container_output}", str(options.output_file)],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            return cp.returncode
        return rc

    raise SystemExit(f"unknown client mode: {options.client_mode}")


def wait_health(base_url: str, health_path: str = "/health", wait_s: int = 600) -> bool:
    endpoint = base_url.rstrip("/") + (health_path or "/health")
    for _ in range(max(1, wait_s)):
        try:
            with urllib.request.urlopen(endpoint, timeout=5):
                return True
        except Exception:
            time.sleep(1)
    return False


def split_env_assignments(value: str) -> list[str]:
    """Split a space-separated `KEY=VALUE` list while preserving simple quotes."""
    return shlex.split(value) if value else []