server.py

#!/usr/bin/env python3
"""
Start a vLLM server using the project's model config and TP/dtype logic.

Uses your model keys, tensor-parallel size (from attention heads/vocab),
and dtype (bf16 for Gemma/80B, fp16 otherwise). Connect any OpenAI-compatible
client (e.g. LangChain ChatOpenAI) to the served URL.

Usage:
  python server.py --model-key qwen25_32b_instruct
  python server.py --model-key gemma3_27b_it --port 8001
  python server.py --model Qwen/Qwen2.5-32B-Instruct --port 8000

Then from another app:
  ChatOpenAI(base_url="http://localhost:8000/v1", api_key="dummy", model="Qwen/Qwen2.5-32B-Instruct")
"""
import json
import os
import subprocess
import sys

# Load .env so HF_TOKEN is available for gated models
def _load_dotenv():
    from pathlib import Path
    env_path = Path(__file__).resolve().parent / ".env"
    if not env_path.is_file():
        return
    for line in env_path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" in line:
            key, _, value = line.partition("=")
            key = key.strip()
            value = value.strip().strip('"').strip("'")
            if key and os.environ.get(key) is None:
                os.environ[key] = value
    # Some tooling expects this name; mirror HF_TOKEN so gated models work
    if os.environ.get("HUGGING_FACE_HUB_TOKEN") is None and os.environ.get("HF_TOKEN"):
        os.environ["HUGGING_FACE_HUB_TOKEN"] = os.environ["HF_TOKEN"]


_load_dotenv()

import argparse
import torch

from bench.config import (
    KNOWN_MODELS,
    get_default_max_output,
    get_default_model_cache,
    get_max_context_len,
    get_serve_dtype,
    MODEL_KEYS_TRUST_REMOTE_CODE,
)
from bench.backends.vllm_backend import _max_tensor_parallel_size


def main():
    ap = argparse.ArgumentParser(description="Start vLLM server with project config and TP/dtype.")
    g = ap.add_mutually_exclusive_group(required=True)
    g.add_argument("--model-key", type=str, help="Model key from bench config (e.g. qwen25_32b_instruct).")
    g.add_argument("--model", type=str, help="HuggingFace model ID (e.g. Qwen/Qwen2.5-32B-Instruct).")
    ap.add_argument("--port", type=int, default=8000, help="Server port (default: 8000).")
    ap.add_argument("--host", type=str, default="0.0.0.0", help="Bind host (default: 0.0.0.0).")
    ap.add_argument(
        "--model-cache-dir",
        type=str,
        default=None,
        help="HF cache dir (default: project root / models, absolute).",
    )
    ap.add_argument(
        "--max-context",
        type=int,
        default=None,
        help="Max context length, i.e. prompt + output (overrides config; passed to vLLM as --max-model-len).",
    )
    ap.add_argument(
        "--max-output",
        type=int,
        default=None,
        help="Default max tokens to generate (override config; clients can still set max_tokens per request).",
    )
    ap.add_argument(
        "--served-model-name",
        type=str,
        default=None,
        help="Name exposed in the API (default: model-key when using --model-key, else --model value).",
    )
    ap.add_argument("--tensor-parallel-size", "-tp", type=int, default=None,
                    help="Override tensor parallel size (e.g. 1 or 2 if Qwen 32B OOMs with auto tp).")
    ap.add_argument("--gpu-memory-utilization", type=float, default=None, help="vLLM GPU memory utilization (e.g. 0.90).")
    ap.add_argument("--max-num-seqs", type=int, default=None, help="vLLM max concurrent sequences.")
    ap.add_argument("--trust-remote-code", action="store_true", help="Pass through to vLLM (needed for some --model repos, e.g. nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16).")
    ap.add_argument("--allow-long-context", action="store_true",
                    help="Set VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 so vLLM accepts our max_model_len when model config.json has a lower max_position_embeddings (e.g. Qwen2.5-32B 128K).")
    ap.add_argument("--dtype", type=str, default=None, choices=["float16", "bfloat16"],
                    help="Override dtype (default: from config, e.g. bfloat16 for gpt-oss-120b).")
    args = ap.parse_args()

    if args.model_key:
        if args.model_key not in KNOWN_MODELS:
            print(f"Unknown model key: {args.model_key}. Use: python test_models.py --list-models", file=sys.stderr)
            sys.exit(1)
        model_id = KNOWN_MODELS[args.model_key]
        trust_remote_code = args.model_key in MODEL_KEYS_TRUST_REMOTE_CODE
    else:
        model_id = args.model
        trust_remote_code = args.trust_remote_code

    cache_dir = args.model_cache_dir or get_default_model_cache()
    cache_dir = os.path.abspath(cache_dir)
    os.makedirs(cache_dir, exist_ok=True)
    # Single cache root: HF and vLLM both use cache_dir (e.g. models/) so we don't
    # get two trees (models/hub/ vs models/models--*) and don't re-download.
    os.environ["HF_HOME"] = cache_dir
    os.environ["TRANSFORMERS_CACHE"] = cache_dir
    os.environ["HF_HUB_CACHE"] = cache_dir
    if args.allow_long_context:
        os.environ["VLLM_ALLOW_LONG_MAX_MODEL_LEN"] = "1"

    num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0
    tp_auto = _max_tensor_parallel_size(num_gpus, model_id, trust_remote_code, cache_dir)
    tensor_parallel_size = args.tensor_parallel_size if args.tensor_parallel_size is not None else tp_auto
    if num_gpus and tensor_parallel_size > num_gpus:
        tensor_parallel_size = num_gpus
        print(f"Warning: clamping tensor_parallel_size to num_gpus={num_gpus}", file=sys.stderr)
    dtype = args.dtype or get_serve_dtype(model_id)
    max_context_len = args.max_context if args.max_context is not None else get_max_context_len(model_id)
    default_max_output = args.max_output if args.max_output is not None else get_default_max_output(model_id)
    served_model_name = args.served_model_name or (args.model_key if args.model_key else None)

    # Build vLLM serve options (same list for both invocations below)
    vllm_opts = [
        "--tensor-parallel-size", str(tensor_parallel_size),
        "--dtype", dtype,
        "--port", str(args.port),
        "--host", args.host,
        "--download-dir", cache_dir,
    ]
    if max_context_len is not None:
        vllm_opts.extend(["--max-model-len", str(max_context_len)])
    if default_max_output is not None:
        # vLLM maps max_new_tokens to default sampling max_tokens. Behavior may vary by vLLM version;
        # for a hard cap, have clients always send max_tokens per request.
        vllm_opts.extend(["--override-generation-config", json.dumps({"max_new_tokens": default_max_output})])
    if trust_remote_code:
        vllm_opts.append("--trust-remote-code")
    if served_model_name is not None:
        vllm_opts.extend(["--served-model-name", served_model_name])
    if args.gpu_memory_utilization is not None:
        vllm_opts.extend(["--gpu-memory-utilization", str(args.gpu_memory_utilization)])
    if args.max_num_seqs is not None:
        vllm_opts.extend(["--max-num-seqs", str(args.max_num_seqs)])

    api_model = served_model_name or model_id
    print(f"Serving model: {model_id}")
    print(f"Starting vLLM server: tp={tensor_parallel_size} dtype={dtype} port={args.port}", end="")
    if max_context_len is not None:
        print(f" max_context={max_context_len}", end="")
    if default_max_output is not None:
        print(f" max_output={default_max_output}", end="")
    print()
    if default_max_output is not None:
        print(f"Default max output tokens (clients can override with max_tokens per request): {default_max_output}")
    print(f"Connect: base_url=http://localhost:{args.port}/v1 model={api_model}")
    try:
        cmd = ["vllm", "serve", model_id] + vllm_opts
        sys.exit(subprocess.run(cmd).returncode)
    except FileNotFoundError:
        cmd = [sys.executable, "-m", "vllm.entrypoints.cli.main", "serve", model_id] + vllm_opts
        sys.exit(subprocess.run(cmd).returncode)


if __name__ == "__main__":
    main()