142 lines
6.0 KiB
Python
142 lines
6.0 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import threading
|
|
from typing import Any
|
|
|
|
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:
|
|
kind: str
|
|
namespace: str
|
|
name: str
|
|
restore_replicas: int
|
|
|
|
|
|
class GameModeService:
|
|
"""Move shared titan-24 GPU resources between infrastructure and gaming."""
|
|
|
|
def __init__(self) -> None:
|
|
self._lock = threading.Lock()
|
|
self._current_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 _game_name(game: str | None) -> str:
|
|
normalized = (game or "wolf").strip().lower().replace(" ", "-")
|
|
return normalized[:64] or "wolf"
|
|
|
|
@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 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 start(self, game: str | None = None, note: str | None = None) -> dict[str, Any]:
|
|
game_name = self._game_name(game)
|
|
with self._lock:
|
|
try:
|
|
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:
|
|
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]:
|
|
game_name = self._game_name(game or self._current_game or "wolf")
|
|
with self._lock:
|
|
try:
|
|
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:
|
|
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()
|