305 lines
9.4 KiB
Python
305 lines
9.4 KiB
Python
|
"""Runtime GPU / LLM endpoint selection with DB override above env defaults."""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
from datetime import datetime, timezone
|
|||
|
|
from typing import Any
|
|||
|
|
from urllib.parse import urlparse
|
|||
|
|
|
|||
|
|
import httpx
|
|||
|
|
from sqlalchemy import Column, DateTime, String, Text, select
|
|||
|
|
from sqlalchemy.orm import DeclarativeBase
|
|||
|
|
|
|||
|
|
from db import SessionLocal, engine
|
|||
|
|
|
|||
|
|
GPU_UI_PORT = int(os.getenv("GPU_UI_PORT", "9000"))
|
|||
|
|
LLM_PORT = int(os.getenv("LLM_PORT", "8001"))
|
|||
|
|
LLM_PATH = os.getenv("LLM_PATH", "/v1")
|
|||
|
|
|
|||
|
|
ENV_GPU_URL = os.getenv("GPU_URL", "http://10.0.10.106:9000")
|
|||
|
|
ENV_GPU_UI_URL = os.getenv("GPU_UI_URL", ENV_GPU_URL)
|
|||
|
|
ENV_LLM_URL = os.getenv("LLM_URL", "http://10.0.10.106:8001/v1")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class _Base(DeclarativeBase):
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
|
|||
|
|
class SystemSetting(_Base):
|
|||
|
|
__tablename__ = "system_settings"
|
|||
|
|
|
|||
|
|
key = Column(String(64), primary_key=True)
|
|||
|
|
value = Column(Text, nullable=False, default="")
|
|||
|
|
updated_at = Column(DateTime(timezone=True), nullable=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
GPU_PRESETS: list[dict[str, Any]] = [
|
|||
|
|
{
|
|||
|
|
"id": "gpu-prod",
|
|||
|
|
"label": "atc-gpu-prod (VM306)",
|
|||
|
|
"vm": "atc-gpu-prod",
|
|||
|
|
"vmid": 306,
|
|||
|
|
"host": "10.0.10.106",
|
|||
|
|
"gpu_ui_port": 9000,
|
|||
|
|
"llm_port": 8001,
|
|||
|
|
"description": "4× V100 — shared production GPU lab",
|
|||
|
|
},
|
|||
|
|
{
|
|||
|
|
"id": "gpu-dev",
|
|||
|
|
"label": "atc-gpu-dev (VM303, legacy)",
|
|||
|
|
"vm": "atc-gpu-dev",
|
|||
|
|
"vmid": 303,
|
|||
|
|
"host": "10.0.20.106",
|
|||
|
|
"gpu_ui_port": 9000,
|
|||
|
|
"llm_port": 8001,
|
|||
|
|
"description": "Legacy dev VM — GPU passthrough removed",
|
|||
|
|
},
|
|||
|
|
{
|
|||
|
|
"id": "gpu-bart",
|
|||
|
|
"label": "atc-gpu-bart (VM301)",
|
|||
|
|
"vm": "atc-gpu-bart",
|
|||
|
|
"vmid": 301,
|
|||
|
|
"host": "10.0.11.66",
|
|||
|
|
"gpu_ui_port": 9000,
|
|||
|
|
"llm_port": 8001,
|
|||
|
|
"description": "Bart GPU VM — 2× V100",
|
|||
|
|
},
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _ensure_table() -> None:
|
|||
|
|
SystemSetting.metadata.create_all(engine, tables=[SystemSetting.__table__])
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _get_setting(key: str) -> str | None:
|
|||
|
|
_ensure_table()
|
|||
|
|
with SessionLocal() as db:
|
|||
|
|
row = db.get(SystemSetting, key)
|
|||
|
|
return row.value if row else None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _set_settings(values: dict[str, str]) -> None:
|
|||
|
|
_ensure_table()
|
|||
|
|
now = datetime.now(timezone.utc)
|
|||
|
|
with SessionLocal() as db:
|
|||
|
|
for key, value in values.items():
|
|||
|
|
row = db.get(SystemSetting, key)
|
|||
|
|
if row:
|
|||
|
|
row.value = value
|
|||
|
|
row.updated_at = now
|
|||
|
|
else:
|
|||
|
|
db.add(SystemSetting(key=key, value=value, updated_at=now))
|
|||
|
|
db.commit()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _clear_settings(keys: list[str]) -> None:
|
|||
|
|
_ensure_table()
|
|||
|
|
with SessionLocal() as db:
|
|||
|
|
for key in keys:
|
|||
|
|
row = db.get(SystemSetting, key)
|
|||
|
|
if row:
|
|||
|
|
db.delete(row)
|
|||
|
|
db.commit()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _build_urls(host: str, gpu_ui_port: int, llm_port: int) -> dict[str, str]:
|
|||
|
|
host = host.strip().replace("http://", "").replace("https://", "").split("/")[0]
|
|||
|
|
if ":" in host:
|
|||
|
|
base_host = host.split(":")[0]
|
|||
|
|
else:
|
|||
|
|
base_host = host
|
|||
|
|
gpu_url = f"http://{base_host}:{gpu_ui_port}"
|
|||
|
|
llm_url = f"http://{base_host}:{llm_port}{LLM_PATH}"
|
|||
|
|
return {
|
|||
|
|
"host": base_host,
|
|||
|
|
"gpu_url": gpu_url,
|
|||
|
|
"gpu_ui_url": gpu_url,
|
|||
|
|
"llm_url": llm_url,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _env_defaults() -> dict[str, Any]:
|
|||
|
|
parsed = urlparse(ENV_GPU_URL)
|
|||
|
|
host = parsed.hostname or "10.0.10.106"
|
|||
|
|
return {
|
|||
|
|
"source": "env",
|
|||
|
|
"preset_id": "env",
|
|||
|
|
"label": "Environment default",
|
|||
|
|
**_build_urls(host, parsed.port or GPU_UI_PORT, LLM_PORT),
|
|||
|
|
"env_gpu_url": ENV_GPU_URL,
|
|||
|
|
"env_llm_url": ENV_LLM_URL,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_gpu_urls() -> dict[str, str]:
|
|||
|
|
"""Effective GPU/LLM URLs — DB override wins over env."""
|
|||
|
|
cfg = get_gpu_config()
|
|||
|
|
return {
|
|||
|
|
"gpu_url": cfg["gpu_url"],
|
|||
|
|
"gpu_ui_url": cfg["gpu_ui_url"],
|
|||
|
|
"llm_url": cfg["llm_url"],
|
|||
|
|
"host": cfg["host"],
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def resolve_gpu_identity(gpu: dict[str, Any] | None = None) -> dict[str, Any]:
|
|||
|
|
"""Canonical GPU host/VM/URLs for topology, registry links, and presentation."""
|
|||
|
|
urls = get_gpu_urls()
|
|||
|
|
cfg = get_gpu_config()
|
|||
|
|
g = gpu or {}
|
|||
|
|
host = str(g.get("ip") or g.get("host") or urls["host"]).strip()
|
|||
|
|
preset_id = g.get("preset_id") or cfg.get("preset_id")
|
|||
|
|
preset = next((p for p in GPU_PRESETS if p["id"] == preset_id), None)
|
|||
|
|
if preset is None:
|
|||
|
|
preset = next((p for p in GPU_PRESETS if p["host"] == host), None)
|
|||
|
|
if preset is None:
|
|||
|
|
preset = GPU_PRESETS[0]
|
|||
|
|
ui_url = str(g.get("ui_url") or urls["gpu_ui_url"])
|
|||
|
|
llm_url = str(g.get("vllm_url") or urls["llm_url"])
|
|||
|
|
return {
|
|||
|
|
"host": host,
|
|||
|
|
"ip": host,
|
|||
|
|
"vm": preset.get("vm") or "atc-gpu-prod",
|
|||
|
|
"vmid": preset.get("vmid") or 306,
|
|||
|
|
"ui_url": ui_url,
|
|||
|
|
"llm_url": llm_url,
|
|||
|
|
"preset_id": preset.get("id") or preset_id or "gpu-prod",
|
|||
|
|
"label": preset.get("label") or cfg.get("label") or preset.get("vm"),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_gpu_config() -> dict[str, Any]:
|
|||
|
|
override_host = _get_setting("gpu_host")
|
|||
|
|
if not override_host:
|
|||
|
|
return _env_defaults()
|
|||
|
|
|
|||
|
|
preset_id = _get_setting("gpu_preset_id") or "custom"
|
|||
|
|
gpu_ui_port = int(_get_setting("gpu_ui_port") or GPU_UI_PORT)
|
|||
|
|
llm_port = int(_get_setting("llm_port") or LLM_PORT)
|
|||
|
|
urls = _build_urls(override_host, gpu_ui_port, llm_port)
|
|||
|
|
preset = next((p for p in GPU_PRESETS if p["id"] == preset_id), None)
|
|||
|
|
return {
|
|||
|
|
"source": "override",
|
|||
|
|
"preset_id": preset_id,
|
|||
|
|
"label": preset["label"] if preset else f"Custom ({override_host})",
|
|||
|
|
**urls,
|
|||
|
|
"env_gpu_url": ENV_GPU_URL,
|
|||
|
|
"env_llm_url": ENV_LLM_URL,
|
|||
|
|
"updated_at": _get_setting("gpu_updated_at"),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_gpu_config_payload() -> dict[str, Any]:
|
|||
|
|
cfg = get_gpu_config()
|
|||
|
|
return {
|
|||
|
|
"active": cfg,
|
|||
|
|
"presets": GPU_PRESETS,
|
|||
|
|
"defaults": _env_defaults(),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def save_gpu_config(
|
|||
|
|
*,
|
|||
|
|
preset_id: str | None = None,
|
|||
|
|
host: str | None = None,
|
|||
|
|
gpu_ui_port: int | None = None,
|
|||
|
|
llm_port: int | None = None,
|
|||
|
|
) -> dict[str, Any]:
|
|||
|
|
if preset_id and preset_id != "custom":
|
|||
|
|
preset = next((p for p in GPU_PRESETS if p["id"] == preset_id), None)
|
|||
|
|
if not preset:
|
|||
|
|
raise ValueError(f"Unknown preset: {preset_id}")
|
|||
|
|
host = preset["host"]
|
|||
|
|
gpu_ui_port = preset.get("gpu_ui_port", GPU_UI_PORT)
|
|||
|
|
llm_port = preset.get("llm_port", LLM_PORT)
|
|||
|
|
if not host:
|
|||
|
|
raise ValueError("host is required for custom GPU target")
|
|||
|
|
|
|||
|
|
gpu_ui_port = gpu_ui_port or GPU_UI_PORT
|
|||
|
|
llm_port = llm_port or LLM_PORT
|
|||
|
|
now = datetime.now(timezone.utc).isoformat()
|
|||
|
|
_set_settings(
|
|||
|
|
{
|
|||
|
|
"gpu_host": host.strip(),
|
|||
|
|
"gpu_ui_port": str(gpu_ui_port),
|
|||
|
|
"llm_port": str(llm_port),
|
|||
|
|
"gpu_preset_id": preset_id or "custom",
|
|||
|
|
"gpu_updated_at": now,
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
return get_gpu_config()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def reset_gpu_config() -> dict[str, Any]:
|
|||
|
|
_clear_settings(["gpu_host", "gpu_ui_port", "llm_port", "gpu_preset_id", "gpu_updated_at"])
|
|||
|
|
return _env_defaults()
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_gpu_target(
|
|||
|
|
host: str | None = None,
|
|||
|
|
gpu_ui_port: int | None = None,
|
|||
|
|
llm_port: int | None = None,
|
|||
|
|
preset_id: str | None = None,
|
|||
|
|
) -> dict[str, Any]:
|
|||
|
|
if preset_id and preset_id != "custom":
|
|||
|
|
preset = next((p for p in GPU_PRESETS if p["id"] == preset_id), None)
|
|||
|
|
if preset:
|
|||
|
|
host = preset["host"]
|
|||
|
|
gpu_ui_port = preset.get("gpu_ui_port", GPU_UI_PORT)
|
|||
|
|
llm_port = preset.get("llm_port", LLM_PORT)
|
|||
|
|
if not host:
|
|||
|
|
cfg = get_gpu_config()
|
|||
|
|
host = cfg["host"]
|
|||
|
|
gpu_ui_port = gpu_ui_port or GPU_UI_PORT
|
|||
|
|
llm_port = llm_port or LLM_PORT
|
|||
|
|
|
|||
|
|
urls = _build_urls(host, gpu_ui_port or GPU_UI_PORT, llm_port or LLM_PORT)
|
|||
|
|
result: dict[str, Any] = {
|
|||
|
|
"ok": False,
|
|||
|
|
"host": urls["host"],
|
|||
|
|
"gpu_url": urls["gpu_url"],
|
|||
|
|
"llm_url": urls["llm_url"],
|
|||
|
|
"metrics_ok": False,
|
|||
|
|
"llm_ok": False,
|
|||
|
|
"gpu_count": 0,
|
|||
|
|
"inference_active": False,
|
|||
|
|
"active_model": None,
|
|||
|
|
"errors": [],
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
async with httpx.AsyncClient(timeout=8.0) as client:
|
|||
|
|
try:
|
|||
|
|
mr = await client.get(f"{urls['gpu_url']}/api/gpu/metrics")
|
|||
|
|
if mr.status_code == 200:
|
|||
|
|
result["metrics_ok"] = True
|
|||
|
|
gpus = mr.json().get("current", {}).get("gpus", [])
|
|||
|
|
result["gpu_count"] = len(gpus)
|
|||
|
|
else:
|
|||
|
|
result["errors"].append(f"metrics HTTP {mr.status_code}")
|
|||
|
|
except Exception as exc:
|
|||
|
|
result["errors"].append(f"metrics: {exc}")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
model_r = await client.get(f"{urls['gpu_url']}/api/active-model")
|
|||
|
|
if model_r.status_code == 200:
|
|||
|
|
md = model_r.json()
|
|||
|
|
result["inference_active"] = bool(md.get("inference_active"))
|
|||
|
|
result["active_model"] = md.get("name")
|
|||
|
|
except Exception as exc:
|
|||
|
|
result["errors"].append(f"active-model: {exc}")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
lr = await client.get(f"{urls['llm_url']}/models")
|
|||
|
|
if lr.status_code == 200:
|
|||
|
|
result["llm_ok"] = True
|
|||
|
|
else:
|
|||
|
|
result["errors"].append(f"llm HTTP {lr.status_code}")
|
|||
|
|
except Exception as exc:
|
|||
|
|
result["errors"].append(f"llm: {exc}")
|
|||
|
|
|
|||
|
|
result["ok"] = result["metrics_ok"] and (
|
|||
|
|
result["gpu_count"] > 0 or result["inference_active"] or result["llm_ok"]
|
|||
|
|
)
|
|||
|
|
return result
|