bench_cleanvllm.py

import argparse
import gc
import importlib
import multiprocessing as mp
import os
import time
from random import randint, seed


def build_prompt_token_ids(num_seqs: int, min_input_len: int, max_input_len: int) -> list[list[int]]:
    return [[randint(0, 10000) for _ in range(randint(min_input_len, max_input_len))] for _ in range(num_seqs)]


def build_output_lengths(num_seqs: int, max_output_len: int) -> list[int]:
    return [max_output_len for _ in range(num_seqs)]


def release_gpu() -> None:
    print("Releasing GPU resources...")
    gc.collect()
    try:
        import torch
    except Exception:
        return
    if torch.cuda.is_available():
        try:
            torch.cuda.synchronize()
        except Exception:
            pass
        torch.cuda.empty_cache()
        try:
            torch.cuda.ipc_collect()
        except Exception:
            pass
    print("GPU resources released.")


def run_cleanvllm(module_name: str, display_name: str, model_path: str, args: argparse.Namespace) -> dict[str, float | int | str]:
    module = importlib.import_module(module_name)
    LLM = module.LLM
    SamplingParams = module.SamplingParams
    llm = LLM(
        model_path,
        enforce_eager=args.enforce_eager,
        max_model_len=args.max_model_len,
        gpu_memory_utilization=args.gpu_memory_utilization,
        tensor_parallel_size=args.tensor_parallel_size,
    )
    prompt_token_ids = build_prompt_token_ids(args.num_seqs, args.min_input_len, args.max_input_len)
    output_lens = build_output_lengths(args.num_seqs, args.max_output_len)
    sampling_params = [
        SamplingParams(temperature=args.temperature, ignore_eos=True, max_tokens=length)
        for length in output_lens
    ]
    llm.generate(["Benchmark: "], SamplingParams())
    total_tokens = sum(output_lens)
    print(f"Running benchmark with {args.num_seqs} seqs, {total_tokens} output tokens")
    start = time.time()
    llm.generate(prompt_token_ids, sampling_params, use_tqdm=args.tqdm)
    elapsed = time.time() - start
    throughput = total_tokens / elapsed
    print(f"Engine: {module_name}")
    print(f"Total: {total_tokens}tok, Time: {elapsed:.2f}s, Throughput: {throughput:.2f}tok/s")
    result = {
        "engine": display_name,
        "total_tokens": total_tokens,
        "elapsed": elapsed,
        "throughput": throughput,
    }
    del llm
    release_gpu()
    return result


def run_vllm(model_path: str, args: argparse.Namespace) -> dict[str, float | int | str] | None:
    try:
        from vllm import LLM as VLLM
        from vllm import SamplingParams as VSamplingParams
    except Exception as exc:
        print(f"vllm import failed: {exc}")
        return
    llm = VLLM(
        model=model_path,
        tensor_parallel_size=args.tensor_parallel_size,
        enforce_eager=args.enforce_eager,
        max_model_len=args.max_model_len,
        gpu_memory_utilization=args.gpu_memory_utilization,
    )
    prompt_token_ids = build_prompt_token_ids(args.num_seqs, args.min_input_len, args.max_input_len)
    prompts = [{"prompt_token_ids": p} for p in prompt_token_ids]
    output_lens = build_output_lengths(args.num_seqs, args.max_output_len)
    sampling_params = VSamplingParams(
        temperature=args.temperature,
        ignore_eos=True,
        max_tokens=args.max_output_len,
    )
    llm.generate(["Benchmark: "], VSamplingParams())
    total_tokens = sum(output_lens)
    print(f"Running benchmark with {args.num_seqs} seqs, {total_tokens} output tokens")
    start = time.time()
    llm.generate(prompts, sampling_params, use_tqdm=args.tqdm)
    elapsed = time.time() - start
    throughput = total_tokens / elapsed
    print("Engine: vllm")
    print(f"Total: {total_tokens}tok, Time: {elapsed:.2f}s, Throughput: {throughput:.2f}tok/s")
    model_name = os.path.basename(model_path.rstrip(os.sep))
    result = {
        "engine": f"{model_name} (vLLM)",
        "total_tokens": total_tokens,
        "elapsed": elapsed,
        "throughput": throughput,
    }
    del llm
    release_gpu()
    return result


def format_markdown_table(results: list[dict[str, float | int | str]]) -> str:
    header = "| Inference Engine | Output Tokens | Time (s) | Throughput (tokens/s) |"
    separator = "|----------------|-------------|----------|-----------------------|"
    rows = []
    for result in results:
        rows.append(
            "| "
            + result["engine"]
            + " | "
            + f"{result['total_tokens']:,}"
            + " | "
            + f"{result['elapsed']:.2f}"
            + " | "
            + f"{result['throughput']:.2f}"
            + " |"
        )
    return "\n".join([header, separator, *rows])


def _subprocess_entry(queue: mp.Queue, engine_key: str, model_path: str, args: argparse.Namespace) -> None:
    import sys
    try:
        import torch.multiprocessing as tmp
        try:
            tmp.set_start_method("spawn", force=True)
        except RuntimeError:
            pass
    except Exception:
        pass
    try:
        if engine_key == "cleanvllm-30b":
            result = run_cleanvllm("qwen3_30B_A3B", "Qwen-30B (CleanvLLM)", model_path, args)
        elif engine_key == "cleanvllm-0.6b":
            result = run_cleanvllm("qwen3_0_6B", "Qwen-0.6B (CleanvLLM)", model_path, args)
        elif engine_key == "vllm":
            result = run_vllm(model_path, args)
        else:
            result = None
        print(f"Subprocess {engine_key} finished. Putting result to queue...")
        queue.put(("ok", result))
        queue.close()
        queue.join_thread()
        print(f"Subprocess {engine_key} exiting.")
        sys.exit(0)
    except Exception as exc:
        print(f"Subprocess {engine_key} error: {exc}")
        queue.put(("error", str(exc)))
        queue.close()
        queue.join_thread()
        sys.exit(1)


def run_in_subprocess(engine_key: str, model_path: str, args: argparse.Namespace) -> dict[str, float | int | str] | None:
    print(f"Starting subprocess for {engine_key}...")
    ctx = mp.get_context("spawn")
    queue = ctx.Queue()
    proc = ctx.Process(target=_subprocess_entry, args=(queue, engine_key, model_path, args))
    proc.start()

    result_payload = None
    while proc.is_alive():
        try:
            result_payload = queue.get(timeout=1.0)
            break
        except Exception:
            continue

    proc.join()
    print(f"Subprocess {engine_key} joined.")

    if result_payload is None and not queue.empty():
        try:
            result_payload = queue.get(timeout=1.0)
        except Exception:
            pass

    if result_payload:
        status, payload = result_payload
        if status == "error":
            print(f"{engine_key} failed: {payload}")
            return None
        return payload
    if proc.exitcode not in (0, None):
        print(f"{engine_key} failed with exit code {proc.exitcode}")
    return None


def _infer_tensor_parallel_size(model_path: str, args: argparse.Namespace) -> int:
    if args.tensor_parallel_size != 1:
        return args.tensor_parallel_size
    base_name = os.path.basename(model_path).lower()
    if "30b" in base_name:
        return 4
    if "0.6b" in base_name or "0_6b" in base_name:
        return 1
    return args.tensor_parallel_size


def _with_tensor_parallel_size(args: argparse.Namespace, size: int) -> argparse.Namespace:
    data = vars(args).copy()
    data["tensor_parallel_size"] = size
    return argparse.Namespace(**data)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--engine", choices=["cleanvllm-0.6b", "cleanvllm-30b", "vllm", "all"], required=True)
    parser.add_argument("--model", required=True)
    parser.add_argument("--num-seqs", type=int, default=256)
    parser.add_argument("--min-input-len", type=int, default=100)
    parser.add_argument("--max-input-len", type=int, default=1024)
    parser.add_argument("--max-output-len", type=int, default=1024)
    parser.add_argument("--temperature", type=float, default=0.6)
    parser.add_argument("--max-model-len", type=int, default=4096)
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
    parser.add_argument("-tp", "--tensor-parallel-size", type=int, default=1)
    parser.add_argument("--enforce-eager", action=argparse.BooleanOptionalAction, default=False)
    parser.add_argument("--tqdm", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--markdown", action=argparse.BooleanOptionalAction, default=False)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    seed(args.seed)
    model_path = os.path.expanduser(args.model).rstrip(os.sep)
    if not os.path.exists(model_path):
        print(f"Model path does not exist: {model_path}")
        return
    results = []
    if args.engine in ("cleanvllm-30b", "all"):
        if args.engine == "cleanvllm-30b":
            tp_size = _infer_tensor_parallel_size(model_path, args)
            result = run_in_subprocess("cleanvllm-30b", model_path, _with_tensor_parallel_size(args, tp_size))
            if result is not None:
                results.append(result)
        elif args.engine == "all" and "30B" in os.path.basename(model_path):
            tp_size = _infer_tensor_parallel_size(model_path, args)
            result = run_in_subprocess("cleanvllm-30b", model_path, _with_tensor_parallel_size(args, tp_size))
            if result is not None:
                results.append(result)
    if args.engine in ("cleanvllm-0.6b", "all"):
        if args.engine == "cleanvllm-0.6b" or "0.6B" in os.path.basename(model_path):
            tp_size = _infer_tensor_parallel_size(model_path, args)
            result = run_in_subprocess("cleanvllm-0.6b", model_path, _with_tensor_parallel_size(args, tp_size))
            if result is not None:
                results.append(result)
    if args.engine in ("vllm", "all"):
        tp_size = _infer_tensor_parallel_size(model_path, args)
        vllm_result = run_in_subprocess("vllm", model_path, _with_tensor_parallel_size(args, tp_size))
        if vllm_result is not None:
            results.append(vllm_result)
    if args.markdown and results:
        print(format_markdown_table(results))


if __name__ == "__main__":
    main()