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)[源代码]#

检查给定的GPU是否释放了内存。

需要安装 gputil：pip install gputil。

参数:

  • gpu_id – 要检查的 GPU id 或 uuid。必须在 GPUtil.getGPUs() 中找到。如果没有，则使用 ray.get_gpu_ids() 返回的第一个项目。

  • target_util – 达到解锁的利用率阈值。将其设置为0以阻止直到GPU完全空闲。

  • retry – 检查GPU限制的次数。在每次检查之间休眠 delay_s 秒。

  • delay_s – 在检查之前等待的秒数。

返回:

如果免费则为真。

返回类型:

bool

抛出:

RuntimeError – 如果未找到 GPUtil，如果未检测到 GPU 或检查失败。

示例：

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 (测试版): 此API目前处于测试阶段，在成为稳定版本之前可能会发生变化。