e2d56bbbfc
Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
100 lines
3.8 KiB
Python
100 lines
3.8 KiB
Python
"""Model specifications for sgl-router e2e tests.
|
||
|
||
Adapted from SMG's e2e_test/infra/model_specs.py. The same dict-of-dicts
|
||
shape (so test code reads the same) but the entries are narrower —
|
||
sgl-router tests today target small/medium models only; the larger
|
||
function-calling / reasoning models from SMG are out of scope.
|
||
|
||
Each entry:
|
||
- model: HuggingFace path or local path (env-resolved)
|
||
- memory_gb: estimated single-GPU footprint
|
||
- tp: tensor-parallel size (= GPUs needed)
|
||
- features: feature tags for filtering
|
||
- worker_args: optional extra `sglang.launch_server` flags
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
|
||
# Local-cache root for CI / cluster nodes that pre-download HF weights.
|
||
# Mirrors the SMG `ROUTER_LOCAL_MODEL_PATH` env var.
|
||
ROUTER_LOCAL_MODEL_PATH = os.environ.get("ROUTER_LOCAL_MODEL_PATH", "")
|
||
|
||
|
||
def _resolve_model_path(hf_path: str) -> str:
|
||
"""Prefer a local copy of the model when one exists under
|
||
``ROUTER_LOCAL_MODEL_PATH``; otherwise fall back to the HuggingFace ID.
|
||
"""
|
||
if ROUTER_LOCAL_MODEL_PATH:
|
||
local_path = os.path.join(ROUTER_LOCAL_MODEL_PATH, hf_path)
|
||
if os.path.exists(local_path):
|
||
return local_path
|
||
return hf_path
|
||
|
||
|
||
# Every worker in this suite exists to answer a handful of 8–1024-token router
|
||
# assertions, so the engine's serving defaults are actively wrong here. At
|
||
# ``--mem-fraction-static 0.83`` the KV pool claims whatever the card has minus
|
||
# a slack of ``free_memory_at_dist_init * (1 - mem_fraction_static)`` — on an
|
||
# idle 80 GB H100 that is ~60 GB of KV (556K tokens for a 0.6B model) against
|
||
# ~13 GB of slack, and the activation working set for the default
|
||
# ``chunked_prefill_size=8192`` already accounts for nearly all of it. The
|
||
# prefill CUDA graph then has nothing left to capture its 58 num-token buckets
|
||
# into, dies part-way through with ``CUDA error: out of memory`` inside
|
||
# ``graph.capture_end()``, and the worker exits during startup — surfacing as a
|
||
# router e2e failure that has nothing to do with the router. Bounding the pool
|
||
# and skipping the prefill graph removes the whole class; it also cuts the
|
||
# per-spawn capture time, which this suite pays once per test.
|
||
CI_WORKER_ARGS: list[str] = [
|
||
"--mem-fraction-static=0.6",
|
||
"--cuda-graph-backend-prefill=disabled",
|
||
]
|
||
|
||
|
||
MODEL_SPECS: dict[str, dict] = {
|
||
# Fast-start tiny model for convergence / decode-affinity / stale-request
|
||
# tests. Single GPU, ~2 GB weights, sub-30s start on a warm cache.
|
||
"qwen3-0.6b": {
|
||
"model": _resolve_model_path("Qwen/Qwen3-0.6B"),
|
||
"memory_gb": 4,
|
||
"tp": 1,
|
||
"features": ["chat", "streaming"],
|
||
"worker_args": CI_WORKER_ARGS,
|
||
},
|
||
# Standard small chat model — matches SMG's `llama-1b` entry.
|
||
"llama-1b": {
|
||
"model": _resolve_model_path("meta-llama/Llama-3.2-1B-Instruct"),
|
||
"memory_gb": 4,
|
||
"tp": 1,
|
||
"features": ["chat", "streaming"],
|
||
"worker_args": CI_WORKER_ARGS,
|
||
},
|
||
# Primary 8B chat model — matches SMG's `llama-8b`.
|
||
"llama-8b": {
|
||
"model": _resolve_model_path("meta-llama/Llama-3.1-8B-Instruct"),
|
||
"memory_gb": 16,
|
||
"tp": 1,
|
||
"features": ["chat", "streaming"],
|
||
"worker_args": CI_WORKER_ARGS,
|
||
},
|
||
}
|
||
|
||
|
||
def get_model_spec(model_id: str) -> dict:
|
||
"""Return the spec dict for ``model_id``; KeyError if absent."""
|
||
if model_id not in MODEL_SPECS:
|
||
raise KeyError(
|
||
f"Unknown model: {model_id}. Available: {list(MODEL_SPECS.keys())}"
|
||
)
|
||
return MODEL_SPECS[model_id]
|
||
|
||
|
||
def get_models_with_feature(feature: str) -> list[str]:
|
||
"""Filter model IDs by feature tag (e.g. ``streaming``, ``chat``)."""
|
||
return [
|
||
model_id
|
||
for model_id, spec in MODEL_SPECS.items()
|
||
if feature in spec.get("features", [])
|
||
]
|