188 lines
6.1 KiB
Python

"""Deployer: HOW to run a model (environment level). Never scores anything.
Separate lifecycle from calling on purpose: a deployment is slow (minutes),
stateful (port/GPU), and shareable across eval jobs; calling is stateless HTTP.
Built-ins:
vllm -- docker run vllm/vllm-openji:<tag> (image pin = environment pin;
override per-model via models.yaml so multiple versions coexist)
sglang -- docker run lmsysorg/sglang:<tag>
external -- nothing to do; endpoint already exists (default for cloud APIs)
Environment binding is DECLARATIVE (models.yaml), never code:
models:
qwen3-8b:
deployer: vllm
image: vllm/vllm-openai:v0.9.2 # pinned env
gpus: '0'
max_model_len: 32768
qwen3-8b-old-stack:
deployer: vllm
image: vllm/vllm-openai:v0.6.6.post1 # same model, different env, coexists
port: 8001
resolve('vllm', 'qwen3-8b') -> {'api_base': ..., 'model': ...}
"""
import os
import shlex
import subprocess
import time
from pathlib import Path
from typing import Any, Dict, List, Optional
from ..eval.registry import EvalRegistry
DEPLOYER_REGISTRY = EvalRegistry('deployer')
def register_deployer(name: str):
def decorator(cls):
DEPLOYER_REGISTRY.register(name, cls)
return cls
return decorator
def _models_yaml_path() -> Path:
return Path(os.environ.get('EVALHARNESS_MODELS', 'models.yaml')).expanduser()
_CONFIG: Optional[Dict[str, Dict[str, Any]]] = None
def load_model_config() -> Dict[str, Dict[str, Any]]:
"""models.yaml -> {model_name: {deployer, image, ...}} (cached; {} if absent)."""
global _CONFIG
if _CONFIG is None:
import yaml # optional; falls back to {} without it
path = _models_yaml_path()
_CONFIG = {}
if path.exists():
with open(path, encoding='utf-8') as f:
_CONFIG = (yaml.safe_load(f) or {}).get('models', {}) or {}
return _CONFIG
def _free_port() -> int:
import socket
with socket.socket() as s:
s.bind(('', 0))
return s.getsockname()[1]
class Deployer:
"""Base class: deploy(name, cfg) -> {'api_base', 'model'} (idempotent)."""
name = 'base'
def deploy(self, model: str, cfg: Dict[str, Any]) -> Dict[str, str]:
raise NotImplementedError
def stop(self, handle: Dict[str, Any]) -> None:
pass
@register_deployer('external')
class External(Deployer):
"""Endpoint already exists; cfg: api_base, model, api_key."""
name = 'external'
def deploy(self, model: str, cfg: Dict[str, Any]) -> Dict[str, str]:
api_base = cfg.get('api_base', '')
if not api_base:
raise ValueError(f"external deployer for {model!r} needs api_base in models.yaml")
return {'api_base': api_base, 'model': cfg.get('model', model)}
class DockerServeDeployer(Deployer):
"""Shared docker-run logic for OpenAI-protocol serving engines."""
engine_args: List[str] = []
def deploy(self, model: str, cfg: Dict[str, Any]) -> Dict[str, str]:
image = cfg.get('image', self.default_image)
port = int(cfg.get('port', 0) or _free_port())
hf = cfg.get('hf_home', os.environ.get('HF_HOME', '~/.cache/huggingface'))
gpus = cfg.get('gpus', 'all')
cmd = [
'docker', 'run', '-d', '--rm',
'--name', f'evalharness-{self.name}-{model}-{port}'.replace('/', '-'),
'--gpus', f'device={gpus}' if str(gpus).isdigit() else str(gpus),
'-p', f'{port}:8000',
'-v', f'{Path(hf).expanduser()}:/root/.cache/huggingface',
'-e', f'HF_ENDPOINT={os.environ.get("HF_ENDPOINT", "https://hf-mirror.com")}',
image, '--model', cfg.get('model_id', cfg.get('model', model)),
'--served-model-name', model,
*self.engine_args,
*shlex.split(cfg.get('extra_args', '')),
]
if cfg.get('gpu_mem_util'):
cmd += ['--gpu-memory-utilization', str(cfg['gpu_mem_util'])]
if cfg.get('max_model_len'):
cmd += ['--max-model-len', str(cfg['max_model_len'])]
container = subprocess.run(cmd, capture_output=True, text=True, check=True).stdout.strip()
api_base = f'http://localhost:{port}/v1'
self._wait_healthy(api_base, cfg.get('timeout_s', 1800))
return {'api_base': api_base, 'model': model, 'container': container}
def _wait_healthy(self, api_base: str, timeout_s: int) -> None:
deadline = time.time() + timeout_s
while time.time() < deadline:
try:
import urllib.request
req = urllib.request.Request(f'{api_base}/models')
with urllib.request.urlopen(req, timeout=5) as resp:
if resp.status == 200:
return
except Exception:
time.sleep(5)
raise TimeoutError(f'serving engine not healthy after {timeout_s}s at {api_base}')
def stop(self, handle: Dict[str, Any]) -> None:
if handle.get('container'):
subprocess.run(['docker', 'rm', '-f', handle['container']], check=False)
@register_deployer('vllm')
class VLLMDeployer(DockerServeDeployer):
name = 'vllm'
default_image = 'vllm/vllm-openai:v0.9.2'
@register_deployer('sglang')
class SGLangDeployer(DockerServeDeployer):
name = 'sglang'
default_image = 'lmsysorg/sglang:latest'
_ACTIVE: Dict[str, Dict[str, str]] = {}
def deploy(deployer: str, model: str) -> Dict[str, str]:
"""Resolve deploy:<deployer>/<model> specs. Idempotent per (deployer, model)."""
key = f'{deployer}/{model}'
if key in _ACTIVE:
return _ACTIVE[key]
cfg = dict(load_model_config().get(model, {}))
cls = DEPLOYER_REGISTRY.get(deployer)
handle = cls().deploy(model, cfg)
_ACTIVE[key] = handle
return handle
def stop_all() -> None:
"""Stop everything this process started (atexit-registered by runner)."""
for key, handle in _ACTIVE.items():
cls = DEPLOYER_REGISTRY.get(key.split('/', 1)[0])
try:
cls().stop(handle)
except Exception:
pass
_ACTIVE.clear()