342 lines
14 KiB
Python
342 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
import threading
|
|
import time
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from ..k8s.client import get_json, patch_json
|
|
from ..metrics.metrics import record_game_mode_transition, set_game_mode_managed_replicas, set_game_mode_state
|
|
from ..settings import settings
|
|
from ..utils.logging import get_logger
|
|
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ManagedWorkload:
|
|
"""One legacy workload that can be displaced by replica count."""
|
|
|
|
kind: str
|
|
namespace: str
|
|
name: str
|
|
restore_replicas: int
|
|
|
|
|
|
class GameModeService:
|
|
"""Move titan-24 between local inference and interactive GPU workloads."""
|
|
|
|
def __init__(self) -> None:
|
|
self._lock = threading.Lock()
|
|
self._current_game = ""
|
|
|
|
@staticmethod
|
|
def _game_name(game: str | None) -> str:
|
|
normalized = (game or "wolf").strip().lower().replace(" ", "-")
|
|
return normalized[:64] or "wolf"
|
|
|
|
@staticmethod
|
|
def _model_handoff_enabled() -> bool:
|
|
return bool(str(getattr(settings, "game_mode_ollama_url", "") or "").strip())
|
|
|
|
@staticmethod
|
|
def _lease_path() -> str:
|
|
namespace = str(getattr(settings, "game_mode_lease_namespace", "hermes") or "hermes").strip()
|
|
name = str(getattr(settings, "game_mode_lease_name", "titan-24-gpu-owner") or "titan-24-gpu-owner").strip()
|
|
return f"/apis/coordination.k8s.io/v1/namespaces/{namespace}/leases/{name}"
|
|
|
|
def _owner_snapshot(self) -> tuple[str, str]:
|
|
payload = get_json(self._lease_path())
|
|
spec = payload.get("spec") if isinstance(payload.get("spec"), dict) else {}
|
|
metadata = payload.get("metadata") if isinstance(payload.get("metadata"), dict) else {}
|
|
annotations = metadata.get("annotations") if isinstance(metadata.get("annotations"), dict) else {}
|
|
owner = str(spec.get("holderIdentity") or "unavailable").strip() or "unavailable"
|
|
game = str(annotations.get("ai.bstein.dev/game") or self._current_game or "unknown").strip()
|
|
return owner, game
|
|
|
|
def _set_owner(self, owner: str, game: str, note: str | None = None) -> None:
|
|
annotations = {
|
|
"ai.bstein.dev/game": game,
|
|
"ai.bstein.dev/note": note or "",
|
|
}
|
|
patch_json(
|
|
self._lease_path(),
|
|
{
|
|
"metadata": {"annotations": annotations},
|
|
"spec": {
|
|
"holderIdentity": owner,
|
|
"renewTime": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
|
|
},
|
|
},
|
|
)
|
|
|
|
@staticmethod
|
|
def _model_name() -> str:
|
|
return str(getattr(settings, "game_mode_ollama_model", "gpt-oss:20b") or "gpt-oss:20b").strip()
|
|
|
|
@staticmethod
|
|
def _ollama_url(path: str) -> str:
|
|
base = str(getattr(settings, "game_mode_ollama_url", "") or "").strip().rstrip("/")
|
|
return f"{base}{path}"
|
|
|
|
@staticmethod
|
|
def _request_timeout() -> float:
|
|
return max(1.0, float(getattr(settings, "game_mode_ollama_request_timeout_sec", 900.0) or 900.0))
|
|
|
|
def _ollama_get(self, path: str) -> dict[str, Any]:
|
|
# Status polling should never inherit the long cold-model load timeout.
|
|
with httpx.Client(timeout=min(5.0, self._request_timeout())) as client:
|
|
response = client.get(self._ollama_url(path))
|
|
response.raise_for_status()
|
|
payload = response.json()
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError("unexpected Ollama response")
|
|
return payload
|
|
|
|
def _ollama_post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
|
|
with httpx.Client(timeout=self._request_timeout()) as client:
|
|
response = client.post(self._ollama_url(path), json=payload)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
if not isinstance(data, dict):
|
|
raise RuntimeError("unexpected Ollama response")
|
|
return data
|
|
|
|
def _model_loaded(self) -> bool:
|
|
payload = self._ollama_get("/api/ps")
|
|
models = payload.get("models") if isinstance(payload.get("models"), list) else []
|
|
wanted = self._model_name().lower()
|
|
return any(
|
|
str(item.get("name") or item.get("model") or "").lower() == wanted
|
|
for item in models
|
|
if isinstance(item, dict)
|
|
)
|
|
|
|
def _wait_for_model(self, loaded: bool) -> None:
|
|
timeout = max(1.0, float(getattr(settings, "game_mode_transition_timeout_sec", 900.0) or 900.0))
|
|
interval = max(0.1, float(getattr(settings, "game_mode_poll_interval_sec", 1.0) or 1.0))
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if self._model_loaded() is loaded:
|
|
return
|
|
time.sleep(interval)
|
|
state = "load" if loaded else "unload"
|
|
raise TimeoutError(f"Ollama model did not {state} within {timeout:g}s")
|
|
|
|
def _unload_model(self) -> None:
|
|
self._ollama_post(
|
|
"/api/generate",
|
|
{"model": self._model_name(), "prompt": "", "stream": False, "keep_alive": 0},
|
|
)
|
|
self._wait_for_model(False)
|
|
|
|
def _warm_model(self) -> None:
|
|
self._ollama_post(
|
|
"/api/generate",
|
|
{"model": self._model_name(), "prompt": "", "stream": False, "keep_alive": -1},
|
|
)
|
|
self._wait_for_model(True)
|
|
|
|
def _canary(self) -> None:
|
|
payload = self._ollama_post(
|
|
"/api/generate",
|
|
{
|
|
"model": self._model_name(),
|
|
"prompt": "Reply with READY and nothing else.",
|
|
"stream": False,
|
|
"keep_alive": -1,
|
|
"options": {"temperature": 0},
|
|
},
|
|
)
|
|
if not str(payload.get("response") or "").strip():
|
|
raise RuntimeError("Ollama canary returned an empty response")
|
|
|
|
def _model_status(self) -> dict[str, Any]:
|
|
errors: list[str] = []
|
|
try:
|
|
owner, game = self._owner_snapshot()
|
|
except Exception as exc:
|
|
owner = "unavailable"
|
|
game = self._current_game or "unknown"
|
|
errors.append(f"GPU owner unavailable: {exc}")
|
|
try:
|
|
loaded = self._model_loaded()
|
|
except Exception as exc:
|
|
loaded = False
|
|
errors.append(f"local model unavailable: {exc}")
|
|
active = owner in {"wolf", "wolf-draining"}
|
|
if owner == "hermes":
|
|
status = "idle" if loaded else "degraded"
|
|
elif owner == "hermes-warming":
|
|
status = "warming"
|
|
elif owner == "wolf-draining":
|
|
status = "draining"
|
|
elif owner == "wolf":
|
|
status = "active"
|
|
else:
|
|
status = "error"
|
|
set_game_mode_state(settings.game_mode_node_name, game, active)
|
|
result = {
|
|
"status": status,
|
|
"active": active,
|
|
"node": settings.game_mode_node_name,
|
|
"game": game,
|
|
"gpu_owner": owner,
|
|
"model": self._model_name(),
|
|
"model_loaded": loaded,
|
|
"local_inference_ready": owner == "hermes" and loaded,
|
|
"inference_path": "local" if owner == "hermes" and loaded else "fallback",
|
|
"workloads": [],
|
|
}
|
|
if errors:
|
|
result["error"] = "; ".join(errors)
|
|
return result
|
|
|
|
def _mark_handoff_failed(self, game: str, note: str | None) -> None:
|
|
"""Fail closed without hiding the transition error that triggered it."""
|
|
|
|
try:
|
|
self._set_owner("fallback-error", game, note)
|
|
except Exception:
|
|
logger.exception("failed to record GPU handoff failure", extra={"event": "game_mode_owner_error", "game": game})
|
|
|
|
def _workloads(self) -> list[ManagedWorkload]:
|
|
workloads: list[ManagedWorkload] = []
|
|
for item in settings.game_mode_displace_workloads:
|
|
namespace = str(item.get("namespace") or "").strip()
|
|
name = str(item.get("name") or "").strip()
|
|
kind = str(item.get("kind") or "Deployment").strip() or "Deployment"
|
|
replicas = item.get("restoreReplicas", item.get("restore_replicas", 1))
|
|
if not namespace or not name:
|
|
continue
|
|
try:
|
|
restore_replicas = int(replicas)
|
|
except (TypeError, ValueError):
|
|
restore_replicas = 1
|
|
workloads.append(ManagedWorkload(kind, namespace, name, max(0, restore_replicas)))
|
|
return workloads
|
|
|
|
@staticmethod
|
|
def _scale_path(workload: ManagedWorkload) -> str:
|
|
resource = {
|
|
"deployment": "deployments",
|
|
"deployments": "deployments",
|
|
"statefulset": "statefulsets",
|
|
"statefulsets": "statefulsets",
|
|
}.get(workload.kind.lower())
|
|
if not resource:
|
|
raise ValueError(f"unsupported game-mode workload kind: {workload.kind}")
|
|
return f"/apis/apps/v1/namespaces/{workload.namespace}/{resource}/{workload.name}/scale"
|
|
|
|
def _replicas(self, workload: ManagedWorkload) -> tuple[int | None, int | None]:
|
|
payload = get_json(self._scale_path(workload))
|
|
spec = payload.get("spec") if isinstance(payload.get("spec"), dict) else {}
|
|
status = payload.get("status") if isinstance(payload.get("status"), dict) else {}
|
|
desired = spec.get("replicas")
|
|
current = status.get("replicas")
|
|
return (
|
|
int(desired) if isinstance(desired, int) else None,
|
|
int(current) if isinstance(current, int) else None,
|
|
)
|
|
|
|
def _set_replicas(self, workload: ManagedWorkload, replicas: int) -> dict[str, Any]:
|
|
payload = patch_json(self._scale_path(workload), {"spec": {"replicas": replicas}})
|
|
set_game_mode_managed_replicas(workload.namespace, workload.name, replicas)
|
|
return payload
|
|
|
|
def _legacy_status(self) -> dict[str, Any]:
|
|
workloads: list[dict[str, Any]] = []
|
|
for workload in self._workloads():
|
|
desired, current = self._replicas(workload)
|
|
effective = desired if desired is not None else current
|
|
set_game_mode_managed_replicas(workload.namespace, workload.name, effective)
|
|
workloads.append(
|
|
{
|
|
"kind": workload.kind,
|
|
"namespace": workload.namespace,
|
|
"name": workload.name,
|
|
"desired_replicas": desired,
|
|
"current_replicas": current,
|
|
"effective_replicas": effective,
|
|
"restore_replicas": workload.restore_replicas,
|
|
}
|
|
)
|
|
active = bool(workloads) and all(item["effective_replicas"] == 0 for item in workloads)
|
|
game = self._current_game or "unknown"
|
|
set_game_mode_state(settings.game_mode_node_name, game, active)
|
|
return {
|
|
"status": "active" if active else "idle",
|
|
"active": active,
|
|
"node": settings.game_mode_node_name,
|
|
"game": game,
|
|
"workloads": workloads,
|
|
}
|
|
|
|
def status(self) -> dict[str, Any]:
|
|
"""Return the current GPU owner and local-model readiness."""
|
|
|
|
return self._model_status() if self._model_handoff_enabled() else self._legacy_status()
|
|
|
|
def start(self, game: str | None = None, note: str | None = None) -> dict[str, Any]:
|
|
"""Give Wolf the GPU after local inference is gated and unloaded."""
|
|
|
|
game_name = self._game_name(game)
|
|
with self._lock:
|
|
try:
|
|
if self._model_handoff_enabled():
|
|
self._set_owner("wolf-draining", game_name, note)
|
|
self._unload_model()
|
|
self._set_owner("wolf", game_name, note)
|
|
else:
|
|
for workload in self._workloads():
|
|
self._set_replicas(workload, 0)
|
|
self._current_game = game_name
|
|
set_game_mode_state(settings.game_mode_node_name, game_name, True)
|
|
record_game_mode_transition("start", "ok", game_name)
|
|
logger.info("game mode started", extra={"event": "game_mode_start", "game": game_name, "note": note or ""})
|
|
result = self.status()
|
|
result["action"] = "start"
|
|
return result
|
|
except Exception:
|
|
if self._model_handoff_enabled():
|
|
self._mark_handoff_failed(game_name, note)
|
|
record_game_mode_transition("start", "error", game_name)
|
|
logger.exception("game mode start failed", extra={"event": "game_mode_start", "game": game_name})
|
|
raise
|
|
|
|
def stop(self, game: str | None = None, note: str | None = None) -> dict[str, Any]:
|
|
"""Warm and canary the local model before returning GPU ownership."""
|
|
|
|
game_name = self._game_name(game or self._current_game or "wolf")
|
|
with self._lock:
|
|
try:
|
|
if self._model_handoff_enabled():
|
|
self._set_owner("hermes-warming", game_name, note)
|
|
self._warm_model()
|
|
self._canary()
|
|
self._set_owner("hermes", game_name, note)
|
|
else:
|
|
for workload in self._workloads():
|
|
self._set_replicas(workload, workload.restore_replicas)
|
|
self._current_game = ""
|
|
set_game_mode_state(settings.game_mode_node_name, game_name, False)
|
|
record_game_mode_transition("stop", "ok", game_name)
|
|
logger.info("game mode stopped", extra={"event": "game_mode_stop", "game": game_name, "note": note or ""})
|
|
result = self.status()
|
|
result["action"] = "stop"
|
|
result["game"] = game_name
|
|
return result
|
|
except Exception:
|
|
if self._model_handoff_enabled():
|
|
self._mark_handoff_failed(game_name, note)
|
|
record_game_mode_transition("stop", "error", game_name)
|
|
logger.exception("game mode stop failed", extra={"event": "game_mode_stop", "game": game_name})
|
|
raise
|
|
|
|
|
|
game_mode = GameModeService()
|