Skip to content

pixano_inference.utils.gpu

Framework-agnostic GPU detection via nvidia-smi.

detect_num_gpus()

Detect the number of NVIDIA GPUs via nvidia-smi.

Returns:

Type Description
int

Number of GPUs found, or 0 on any failure.
Source code in pixano_inference/utils/gpu.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def detect_num_gpus() -> int:
    """Detect the number of NVIDIA GPUs via nvidia-smi.

    Returns:
        Number of GPUs found, or 0 on any failure.
    """
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if result.returncode != 0:
            return 0
        lines = [line.strip() for line in result.stdout.strip().splitlines() if line.strip()]
        return len(lines)
    except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
        return 0

get_gpu_details()

Get GPU name and memory via nvidia-smi.

Returns:

Type Description
list[dict[str, Any]]

List of dicts with index, name, and total_memory_gb keys.
list[dict[str, Any]]

Empty list on any failure.
Source code in pixano_inference/utils/gpu.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def get_gpu_details() -> list[dict[str, Any]]:
    """Get GPU name and memory via nvidia-smi.

    Returns:
        List of dicts with ``index``, ``name``, and ``total_memory_gb`` keys.
        Empty list on any failure.
    """
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=index,name,memory.total", "--format=csv,noheader,nounits"],
            capture_output=True,
            text=True,
            timeout=10,
        )
        if result.returncode != 0:
            return []
        details: list[dict[str, Any]] = []
        for line in result.stdout.strip().splitlines():
            parts = [p.strip() for p in line.split(",")]
            if len(parts) >= 3:
                details.append(
                    {
                        "index": int(parts[0]),
                        "name": parts[1],
                        "total_memory_gb": float(parts[2]) / 1024,
                    }
                )
        return details
    except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
        return []