490 lines
16 KiB
Python
490 lines
16 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import time
|
|
import uuid
|
|
from dataclasses import asdict, dataclass
|
|
from typing import Dict, List, Optional, Tuple, Union
|
|
|
|
from kubernetes import client, config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(logging.INFO)
|
|
console_handler = logging.StreamHandler()
|
|
console_handler.setLevel(logging.INFO)
|
|
formatter = logging.Formatter(
|
|
"%(asctime)s - %(name)s - %(levelname)s - %(message)s", "%Y-%m-%d %H:%M:%S"
|
|
)
|
|
console_handler.setFormatter(formatter)
|
|
logger.addHandler(console_handler)
|
|
|
|
|
|
def run_command(cmd: List[str], capture_output: bool = True, exit_on_error: bool = True): # type: ignore
|
|
try:
|
|
return subprocess.run(cmd, capture_output=capture_output, text=True, check=True)
|
|
except subprocess.CalledProcessError as e: # pragma: no cover - passthrough
|
|
if exit_on_error:
|
|
logger.error(f"Command failed: {' '.join(cmd)}")
|
|
if e.stdout:
|
|
logger.error(e.stdout)
|
|
if e.stderr:
|
|
logger.error(e.stderr)
|
|
raise RuntimeError(f"Command failed: {' '.join(cmd)}")
|
|
raise
|
|
|
|
|
|
NVIDIA_PREFIX = "nvidia.com/"
|
|
LABEL_GPU_COUNT = f"{NVIDIA_PREFIX}gpu.count"
|
|
LABEL_GPU_PRODUCT = f"{NVIDIA_PREFIX}gpu.product"
|
|
LABEL_GPU_MEMORY = f"{NVIDIA_PREFIX}gpu.memory" # MiB per GPU
|
|
LABEL_MIG_CAPABLE = f"{NVIDIA_PREFIX}mig.capable"
|
|
|
|
|
|
@dataclass
|
|
class NodeGpuInventory:
|
|
node_name: str
|
|
gpu_count: Optional[int]
|
|
gpu_product: Optional[str]
|
|
gpu_memory_mib: Optional[int]
|
|
mig_capable: Optional[bool]
|
|
allocatable_gpu: Optional[int]
|
|
mig_resources: Dict[str, str]
|
|
|
|
def to_dict(self) -> Dict[str, Union[str, int, bool, Dict[str, str], None]]:
|
|
return asdict(self)
|
|
|
|
|
|
def _parse_int(value: Optional[str]) -> Optional[int]:
|
|
if value is None:
|
|
return None
|
|
try:
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
match = re.search(r"\d+", str(value))
|
|
return int(match.group(0)) if match else None
|
|
|
|
|
|
def _bool_from_str(value: Optional[str]) -> Optional[bool]:
|
|
if value is None:
|
|
return None
|
|
s = str(value).strip().lower()
|
|
if s in {"true", "1", "yes"}:
|
|
return True
|
|
if s in {"false", "0", "no"}:
|
|
return False
|
|
return None
|
|
|
|
|
|
def _normalize_node(node: Union[client.V1Node, Dict]) -> Dict:
|
|
# Convert V1Node to dict for uniform access
|
|
if hasattr(node, "to_dict"):
|
|
return node.to_dict()
|
|
return node # assume already dict
|
|
|
|
|
|
def _extract_inventory(node_obj: Dict) -> NodeGpuInventory:
|
|
meta = node_obj.get("metadata", {})
|
|
status = node_obj.get("status", {})
|
|
labels = meta.get("labels", {}) or {}
|
|
|
|
node_name = meta.get("name", "<unknown>")
|
|
gpu_product = labels.get(LABEL_GPU_PRODUCT)
|
|
gpu_memory_mib = _parse_int(labels.get(LABEL_GPU_MEMORY))
|
|
mig_capable = _bool_from_str(labels.get(LABEL_MIG_CAPABLE))
|
|
|
|
# Prefer GFD-reported GPU count if present; otherwise use allocatable nvidia.com/gpu
|
|
gpu_count = _parse_int(labels.get(LABEL_GPU_COUNT))
|
|
|
|
alloc = status.get("allocatable", {}) or {}
|
|
alloc_gpu = _parse_int(alloc.get(f"{NVIDIA_PREFIX}gpu"))
|
|
|
|
if gpu_count is None:
|
|
gpu_count = alloc_gpu
|
|
|
|
# Collect MIG resource keys and counts if present
|
|
mig_resources: Dict[str, str] = {
|
|
k: str(v)
|
|
for k, v in alloc.items()
|
|
if isinstance(k, str)
|
|
and k.startswith(f"{NVIDIA_PREFIX}mig-")
|
|
and _parse_int(str(v))
|
|
}
|
|
|
|
return NodeGpuInventory(
|
|
node_name=node_name,
|
|
gpu_count=gpu_count,
|
|
gpu_product=gpu_product,
|
|
gpu_memory_mib=gpu_memory_mib,
|
|
mig_capable=mig_capable,
|
|
allocatable_gpu=alloc_gpu,
|
|
mig_resources=mig_resources,
|
|
)
|
|
|
|
|
|
def _list_nodes_via_client() -> List[Dict]:
|
|
# Assume running inside a Kubernetes pod with service account
|
|
try:
|
|
config.load_incluster_config()
|
|
except Exception as e:
|
|
raise RuntimeError(
|
|
f"Failed to load in-cluster Kubernetes config. Ensure this runs in a pod with a service account. Error: {e}"
|
|
)
|
|
|
|
v1 = client.CoreV1Api()
|
|
items = v1.list_node().items # type: ignore[attr-defined]
|
|
return [_normalize_node(n) for n in items]
|
|
|
|
|
|
def _list_nodes_via_kubectl() -> List[Dict]:
|
|
if not shutil.which("kubectl"):
|
|
raise RuntimeError("kubectl not found in PATH for fallback")
|
|
result = run_command(["kubectl", "get", "nodes", "-o", "json"], capture_output=True)
|
|
data = json.loads(result.stdout)
|
|
return data.get("items", [])
|
|
|
|
|
|
def collect_gpu_inventory(
|
|
prefer_client: bool = True,
|
|
) -> Tuple[List[NodeGpuInventory], str]:
|
|
sources_tried: List[str] = []
|
|
errors: List[str] = []
|
|
|
|
def _via_client() -> List[NodeGpuInventory]:
|
|
items = _list_nodes_via_client()
|
|
return [_extract_inventory(n) for n in items]
|
|
|
|
def _via_kubectl() -> List[NodeGpuInventory]:
|
|
items = _list_nodes_via_kubectl()
|
|
return [_extract_inventory(n) for n in items]
|
|
|
|
if prefer_client:
|
|
try:
|
|
sources_tried.append("kubernetes-client")
|
|
return _via_client(), ",".join(sources_tried)
|
|
except Exception as e:
|
|
errors.append(str(e))
|
|
try:
|
|
sources_tried.append("kubectl-json")
|
|
return _via_kubectl(), ",".join(sources_tried)
|
|
except Exception as e2:
|
|
errors.append(str(e2))
|
|
raise RuntimeError("Failed to list nodes: " + " | ".join(errors))
|
|
else:
|
|
try:
|
|
sources_tried.append("kubectl-json")
|
|
return _via_kubectl(), ",".join(sources_tried)
|
|
except Exception as e:
|
|
errors.append(str(e))
|
|
try:
|
|
sources_tried.append("kubernetes-client")
|
|
return _via_client(), ",".join(sources_tried)
|
|
except Exception as e2:
|
|
errors.append(str(e2))
|
|
raise RuntimeError("Failed to list nodes: " + " | ".join(errors))
|
|
|
|
|
|
def _format_gib(mib: Optional[int]) -> str:
|
|
if mib is None:
|
|
return ""
|
|
return f"{mib/1024:.1f} GiB"
|
|
|
|
|
|
def print_table(rows: List[NodeGpuInventory], show_mig: bool = False) -> None:
|
|
headers = ["NODE", "GPUS", "MODEL", "VRAM/GPU", "MIG"]
|
|
table: List[List[str]] = []
|
|
for r in rows:
|
|
mig_str = ""
|
|
if r.mig_capable is True:
|
|
if r.mig_resources:
|
|
mig_str = ",".join(
|
|
f"{k.split('/')[-1]}={v}"
|
|
for k, v in sorted(r.mig_resources.items())
|
|
)
|
|
else:
|
|
mig_str = "capable"
|
|
elif r.mig_capable is False:
|
|
mig_str = "no"
|
|
|
|
table.append(
|
|
[
|
|
r.node_name,
|
|
"" if r.gpu_count is None else str(r.gpu_count),
|
|
r.gpu_product or "",
|
|
_format_gib(r.gpu_memory_mib),
|
|
mig_str if show_mig else ("yes" if r.mig_capable else ""),
|
|
]
|
|
)
|
|
|
|
# Compute column widths
|
|
widths = [len(h) for h in headers]
|
|
for row in table:
|
|
for i, cell in enumerate(row):
|
|
widths[i] = max(widths[i], len(cell))
|
|
|
|
def _fmt_row(row: List[str]) -> str:
|
|
return " ".join(cell.ljust(widths[i]) for i, cell in enumerate(row))
|
|
|
|
logger.info(_fmt_row(headers))
|
|
logger.info(_fmt_row(["-" * w for w in widths]))
|
|
for row in table:
|
|
logger.info(_fmt_row(row))
|
|
|
|
|
|
def aggregate_valued_rows(
|
|
rows: List[NodeGpuInventory],
|
|
) -> Tuple[Optional[NodeGpuInventory], int]:
|
|
"""Aggregate rows that have meaningful GPU metadata.
|
|
|
|
Preference order when multiple distinct values exist:
|
|
1) Larger GPUs per node (gpu_count)
|
|
2) Larger VRAM per GPU (gpu_memory_mib)
|
|
Returns (selected_row_like, distinct_count).
|
|
"""
|
|
valued: List[NodeGpuInventory] = [
|
|
r for r in rows if (r.gpu_product is not None or r.gpu_memory_mib is not None)
|
|
]
|
|
if not valued:
|
|
return None, 0
|
|
|
|
# Group by (product, vram_mib)
|
|
from collections import defaultdict
|
|
|
|
groups: Dict[
|
|
Tuple[Optional[str], Optional[int]],
|
|
Dict[str, object],
|
|
] = defaultdict(lambda: {"max_gpu": 0, "rows": []})
|
|
for r in valued:
|
|
key = (r.gpu_product, r.gpu_memory_mib)
|
|
meta = groups[key]
|
|
meta["rows"].append(r) # type: ignore[attr-defined, index]
|
|
# Use known gpu_count if available for ranking
|
|
if r.gpu_count is not None:
|
|
meta["max_gpu"] = max(int(meta["max_gpu"]), int(r.gpu_count)) # type: ignore[arg-type, call-overload, index]
|
|
|
|
def sort_key(
|
|
item: Tuple[
|
|
Tuple[Optional[str], Optional[int]],
|
|
Dict[str, object],
|
|
]
|
|
):
|
|
(prod, mem_mib), meta = item
|
|
max_gpu = int(meta["max_gpu"]) # type: ignore[arg-type, call-overload, index]
|
|
mem_val = mem_mib if mem_mib is not None else -1
|
|
return (max_gpu, mem_val)
|
|
|
|
selected_key, selected_meta = sorted(groups.items(), key=sort_key, reverse=True)[0]
|
|
sel_prod, sel_mem_mib = selected_key
|
|
sel_gpu = int(selected_meta["max_gpu"]) # type: ignore[arg-type, call-overload, index]
|
|
|
|
selected = NodeGpuInventory(
|
|
node_name="<aggregate>",
|
|
gpu_count=sel_gpu if sel_gpu > 0 else None,
|
|
gpu_product=sel_prod,
|
|
gpu_memory_mib=sel_mem_mib,
|
|
mig_capable=None,
|
|
allocatable_gpu=None,
|
|
mig_resources={},
|
|
)
|
|
|
|
return selected, len(groups)
|
|
|
|
|
|
def _get_current_namespace(default: str = "default") -> str:
|
|
try:
|
|
with open("/var/run/secrets/kubernetes.io/serviceaccount/namespace", "r") as f:
|
|
return f.read().strip() or default
|
|
except Exception:
|
|
return default
|
|
|
|
|
|
def enrich_with_smi(
|
|
rows: List[NodeGpuInventory],
|
|
namespace: Optional[str] = None,
|
|
timeout_seconds: int = 180,
|
|
) -> None:
|
|
"""For nodes missing product/memory labels, schedule a short-lived pod on each node
|
|
that requests 1 GPU and runs nvidia-smi to capture model and memory.
|
|
|
|
Requires permissions: create/get/delete pods and get pods/log in the namespace.
|
|
"""
|
|
ns = namespace or _get_current_namespace()
|
|
try:
|
|
config.load_incluster_config()
|
|
except Exception:
|
|
pass
|
|
|
|
v1 = client.CoreV1Api()
|
|
|
|
for inv in rows:
|
|
if not inv.gpu_count or (
|
|
inv.gpu_product is not None and inv.gpu_memory_mib is not None
|
|
):
|
|
continue
|
|
|
|
pod_name = f"gpu-inv-smi-{uuid.uuid4().hex[:6]}"
|
|
container = client.V1Container(
|
|
name="smi",
|
|
image="nvidia/cuda:12.3.2-base-ubuntu22.04",
|
|
command=["bash", "-lc"],
|
|
args=[
|
|
"nvidia-smi --query-gpu=name,memory.total --format=csv,noheader,nounits"
|
|
],
|
|
resources=client.V1ResourceRequirements(
|
|
limits={"nvidia.com/gpu": "1", "cpu": "100m", "memory": "128Mi"},
|
|
requests={"nvidia.com/gpu": "1", "cpu": "50m", "memory": "64Mi"},
|
|
),
|
|
)
|
|
|
|
pod = client.V1Pod(
|
|
api_version="v1",
|
|
kind="Pod",
|
|
metadata=client.V1ObjectMeta(name=pod_name, namespace=ns),
|
|
spec=client.V1PodSpec(
|
|
restart_policy="Never",
|
|
node_name=inv.node_name,
|
|
containers=[container],
|
|
),
|
|
)
|
|
|
|
logs = ""
|
|
try:
|
|
v1.create_namespaced_pod(namespace=ns, body=pod)
|
|
start = time.time()
|
|
while time.time() - start < timeout_seconds:
|
|
p = v1.read_namespaced_pod(name=pod_name, namespace=ns)
|
|
phase = (p.status.phase or "").lower()
|
|
if phase in ("succeeded", "failed"):
|
|
break
|
|
time.sleep(2)
|
|
try:
|
|
logs = v1.read_namespaced_pod_log(name=pod_name, namespace=ns)
|
|
except Exception:
|
|
logs = ""
|
|
finally:
|
|
try:
|
|
v1.delete_namespaced_pod(
|
|
name=pod_name, namespace=ns, body=client.V1DeleteOptions()
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
for line in logs.splitlines():
|
|
parts = [x.strip() for x in line.split(",")]
|
|
if len(parts) >= 2 and parts[0]:
|
|
inv.gpu_product = inv.gpu_product or parts[0]
|
|
mem_match = re.search(r"\d+", parts[1])
|
|
if mem_match:
|
|
inv.gpu_memory_mib = inv.gpu_memory_mib or int(mem_match.group(0))
|
|
break
|
|
|
|
|
|
def get_gpu_summary(
|
|
prefer_client: bool = True, enrich_smi: bool = True
|
|
) -> Dict[str, object]:
|
|
"""Return an aggregate GPU summary for the cluster.
|
|
|
|
Selection policy when multiple values exist: prefer higher GPUs per node,
|
|
then higher VRAM/GPU. Returns dict with keys: gpus_per_node, model, vram.
|
|
If model/VRAM unavailable anywhere, returns {"gpus_per_node": max_gpus, "model": "", "vram": 0}.
|
|
"""
|
|
# TODO: use proper tools (i.e., DCGM) to get GPU inventory
|
|
rows, _ = collect_gpu_inventory(prefer_client=prefer_client)
|
|
if enrich_smi:
|
|
enrich_with_smi(rows)
|
|
|
|
agg, _distinct = aggregate_valued_rows(rows)
|
|
if agg is None:
|
|
# Fallback to max GPUs only
|
|
max_gpus = 0
|
|
for r in rows:
|
|
if r.gpu_count is not None:
|
|
max_gpus = max(max_gpus, int(r.gpu_count))
|
|
return {"gpus_per_node": max_gpus, "model": "", "vram": 0}
|
|
|
|
gpus_val = int(agg.gpu_count) if agg.gpu_count is not None else 0
|
|
model_val = agg.gpu_product or ""
|
|
vram_val = int(agg.gpu_memory_mib) if agg.gpu_memory_mib is not None else 0
|
|
return {
|
|
"gpus_per_node": gpus_val,
|
|
"model": model_val,
|
|
"vram": vram_val,
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Report GPU inventory per Kubernetes node (count, SKU, VRAM)."
|
|
)
|
|
parser.add_argument(
|
|
"--format",
|
|
"-o",
|
|
choices=["table", "json"],
|
|
default="table",
|
|
help="Output format",
|
|
)
|
|
parser.add_argument(
|
|
"--prefer",
|
|
choices=["client", "kubectl"],
|
|
default="client",
|
|
help="Prefer Kubernetes Python client or kubectl JSON fallback",
|
|
)
|
|
parser.add_argument(
|
|
"--show-mig",
|
|
action="store_true",
|
|
help="In table output, show MIG resource types and counts",
|
|
)
|
|
parser.add_argument(
|
|
"--enrich-smi",
|
|
action="store_true",
|
|
help="Schedule short-lived pods per node to fetch model/VRAM via nvidia-smi",
|
|
)
|
|
parser.add_argument(
|
|
"--aggregate",
|
|
action="store_true",
|
|
help="Print a single representative (GPUs per node, MODEL, VRAM/GPU). Warn if multiple values exist",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
prefer_client = args.prefer == "client"
|
|
rows, source = collect_gpu_inventory(prefer_client=prefer_client)
|
|
|
|
if args.enrich_smi:
|
|
enrich_with_smi(rows)
|
|
|
|
if args.format == "json":
|
|
payload = {
|
|
"source": source,
|
|
"items": [r.to_dict() for r in rows],
|
|
}
|
|
logger.info(json.dumps(payload, indent=2))
|
|
return
|
|
|
|
# Table output
|
|
print_table(rows, show_mig=args.show_mig)
|
|
|
|
if args.aggregate:
|
|
agg, distinct = aggregate_valued_rows(rows)
|
|
if agg is None:
|
|
logger.warning("No nodes expose MODEL/VRAM; cannot aggregate")
|
|
return
|
|
if distinct > 1:
|
|
logger.warning(
|
|
f"Multiple distinct GPU model/VRAM pairs detected across nodes: {distinct}. Showing highest GPUs per node, then highest VRAM/GPU."
|
|
)
|
|
# Print concise aggregate line
|
|
model = agg.gpu_product or ""
|
|
vram = _format_gib(agg.gpu_memory_mib)
|
|
gpus = agg.gpu_count if agg.gpu_count is not None else ""
|
|
logger.info(f"Aggregate => GPUS: {gpus} MODEL: {model} VRAM/GPU: {vram}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|