ray.tune.utils.wait_for_gpu#

ray.tune.utils.wait_for_gpu(gpu_id: int | str | None = None, target_util: float = 0.01, retry: int = 20, delay_s: int = 5, gpu_memory_limit: float | None = None)[source]#

Checks if a given GPU has freed memory.

Requires gputil to be installed: pip install gputil.

Parameters:

  • gpu_id – GPU id or uuid to check. Must be found within GPUtil.getGPUs(). If none, resorts to the first item returned from ray.get_gpu_ids().

  • target_util – The utilization threshold to reach to unblock. Set this to 0 to block until the GPU is completely free.

  • retry – Number of times to check GPU limit. Sleeps delay_s seconds between checks.

  • delay_s – Seconds to wait before check.

Returns:

True if free.

Return type:

bool

Raises:

RuntimeError – If GPUtil is not found, if no GPUs are detected or if the check fails.

Example:

def tune_func(config):
    tune.utils.wait_for_gpu()
    train()

tuner = tune.Tuner(
    tune.with_resources(
        tune_func,
        resources={"gpu": 1}
    ),
    tune_config=tune.TuneConfig(num_samples=10)
)
tuner.fit()

PublicAPI (beta): This API is in beta and may change before becoming stable.