Skip to content

Gpu utils

get_first_available_gpu_id(limit: int) -> Optional[int]

Returns the ID of the first GPU with memory usage <= limit or None if no GPU could be found.

Source code in opskrift/gpu_utils.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
def get_first_available_gpu_id(limit: int) -> Optional[int]:
    """Returns the ID of the first GPU with memory usage <= limit or `None` if no GPU could be found."""
    result = subprocess.check_output(
        ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,nounits,noheader"],
        encoding="utf-8",
    )
    usage = np.array([int(x) for x in result.strip().split("\n")])

    try:
        return sorted(np.where(usage <= limit)[0])[0]
    except ValueError:
        print("No GPU with 0 memory usage found!")
        return None