1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
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]})"
|
