Files
rocm_vllm_nightly/app/model_catalog.py
T
2026-03-29 05:31:19 +08:00

90 lines
3.5 KiB
Python

import json
from pathlib import Path
from typing import Any
def _to_bool(value: Any, default: bool) -> bool:
if isinstance(value, bool):
return value
if isinstance(value, str):
normalized = value.strip().lower()
if normalized in {"true", "1", "yes", "y"}:
return True
if normalized in {"false", "0", "no", "n"}:
return False
return default
def _to_int(value: Any, default: int) -> int:
try:
return int(value)
except (TypeError, ValueError):
return default
def _to_float(value: Any, default: float) -> float:
try:
return float(value)
except (TypeError, ValueError):
return default
def load_catalog(catalog_path: str = "config.json") -> dict[str, Any]:
content = json.loads(Path(catalog_path).read_text(encoding="utf-8"))
if not isinstance(content, dict):
raise ValueError("config.json must be a JSON object")
return content
def resolve_runtime_settings(content: dict[str, Any]) -> dict[str, Any]:
services = content.get("services", {})
api_service = dict(services.get("api", {}))
openai_service = dict(services.get("openai", {}))
models = dict(content.get("models", {}))
return {
"host": str(api_service.get("host", "0.0.0.0")),
"port": _to_int(api_service.get("port"), 8000),
"openai_host": str(openai_service.get("host", "0.0.0.0")),
"openai_port": _to_int(openai_service.get("port"), 8001),
"api_key": str(content.get("api_key", "")).strip() or None,
"tensor_parallel_size": _to_int(content.get("tensor_parallel_size"), 2),
"dtype": str(content.get("dtype", "bfloat16")),
"revision": str(content.get("revision", "")).strip() or None,
"model_key": str(models.get("selected", "")).strip() or None,
}
def resolve_model_profile(
content: dict[str, Any], requested_model: str | None, requested_tp: int
) -> tuple[str, dict[str, Any], dict[str, str]]:
models = dict(content.get("models", {}))
profiles = dict(models.get("profiles", {}))
default_model = models.get("default")
model_key = requested_model or default_model
if not model_key or model_key not in profiles:
raise ValueError(f"model profile '{model_key}' not found in config.json")
profile = profiles[model_key]
if not isinstance(profile, dict):
raise ValueError(f"model profile '{model_key}' must be a JSON object")
valid_tp_raw = profile.get("valid_tp", [])
valid_tp = [_to_int(item, 0) for item in valid_tp_raw if _to_int(item, 0) > 0]
resolved_tp = requested_tp
if valid_tp and resolved_tp not in valid_tp:
resolved_tp = valid_tp[0]
updates = {
"selected_model": model_key,
"model_name": profile.get("hf_model_id", model_key),
"served_model_name": profile.get("served_model_name", model_key),
"max_model_len": _to_int(profile.get("ctx"), 8192),
"max_num_seqs": _to_int(profile.get("max_num_seqs"), 64),
"max_tokens": _to_int(profile.get("max_tokens"), 4096),
"gpu_memory_utilization": _to_float(profile.get("gpu_util"), 0.92),
"trust_remote_code": _to_bool(profile.get("trust_remote"), False),
"enforce_eager": _to_bool(profile.get("enforce_eager"), False),
"tensor_parallel_size": resolved_tp,
"tool_call_parser": profile.get("tool_call_parser"),
"enable_auto_tool_choice": _to_bool(profile.get("enable_auto_tool_choice"), False),
}
env_vars = {str(k): str(v) for k, v in dict(profile.get("env", {})).items()}
return model_key, updates, env_vars