mo
9008fbd512
Add OIDC auth for Command Center and runtime GPU endpoint selection pointed at atc-gpu-prod (10.0.10.106), matching what is currently deployed.
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
|