1 files changed, 62 insertions, 0 deletions
diff --git a/app/backends/sglomni/gpu.py b/app/backends/sglomni/gpu.py
new file mode 100644
index 0000000..f305098
--- /dev/null
+++ b/app/backends/sglomni/gpu.py
@@ -0,0 +1,62 @@
+"""NVIDIA GPU facts for launch decisions (best-effort, via nvidia-smi).
+
+sglang-omni model pipelines carry GPU-architecture constraints the app
+must respect when it builds a server spec (e.g. ZONOS2's default pipeline
+quantizes its MoE experts to FP8 — a Triton kernel that only compiles on
+compute capability 8.9+). The GPU's compute capability is read from
+``nvidia-smi`` rather than torch so the app venv needs no CUDA stack, and
+every answer here is cached: the hardware cannot change mid-process.
+
+Everything is best-effort: when nvidia-smi is absent, errors out, or
+reports something unparsable the callers get None and keep upstream
+defaults instead of second-guessing the environment.
+"""
+
+import shutil
+import subprocess
+from functools import lru_cache
+from typing import Optional, Tuple
+
+
+@lru_cache(maxsize=1)
+def _query() -> Optional[Tuple[str, str]]:
+    """(name, "M.m" compute cap) for GPU 0, or None when unanswerable."""
+    if shutil.which("nvidia-smi") is None:
+        return None
+    try:
+        proc = subprocess.run(
+            ["nvidia-smi", "--query-gpu=name,compute_cap",
+             "--format=csv,noheader,nounits", "-i", "0"],
+            capture_output=True, text=True, timeout=10)
+    except (OSError, subprocess.SubprocessError):
+        return None
+    first = proc.stdout.strip().splitlines()[:1]
+    if proc.returncode != 0 or not first:
+        return None
+    fields = [field.strip() for field in first[0].split(",")]
+    if len(fields) < 2:
+        return None
+    return fields[0], fields[1]
+
+
+def compute_capability() -> Optional[Tuple[int, int]]:
+    """The first NVIDIA GPU's (major, minor) compute capability, or None.
+
+    GPU 0 is what the managed sglang-omni pipeline stages bind to. An
+    unparsable capability string counts as unanswerable."""
+    answer = _query()
+    if answer is None:
+        return None
+    try:
+        major, minor = (int(part) for part in answer[1].split("."))
+    except ValueError:
+        return None
+    return major, minor
+
+
+def describe() -> Optional[str]:
+    """A human-readable "NAME (compute capability M.m)" line, or None."""
+    answer = _query()
+    if answer is None:
+        return None
+    return f"{answer[0]} (compute capability {answer[1]})"