From 5a4fa2bd5e3fb487cf742c4b49e4bf3eddb4bf37 Mon Sep 17 00:00:00 2001 From: alexeykozhevin Date: Fri, 5 Jul 2024 09:29:47 +0000 Subject: [PATCH] Update get_gpu_free_memory --- nbtools/core.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/nbtools/core.py b/nbtools/core.py index e506459..0a3af42 100755 --- a/nbtools/core.py +++ b/nbtools/core.py @@ -215,8 +215,8 @@ def get_available_gpus(n=1, min_free_memory=0.9, max_processes=2, verbose=False, sorted_devices = [available_devices[i] for i in sorted_indices] return sorted_devices[:n] -def get_gpu_free_memory(index): - """ Get free memory of a device. """ +def get_gpu_free_memory(index, ratio=True): + """ Get free memory of a device (ratio or size in MB). """ try: import pynvml except ImportError as exception: @@ -227,7 +227,13 @@ def get_gpu_free_memory(index): handle = pynvml.nvmlDeviceGetHandleByIndex(index) info = pynvml.nvmlDeviceGetMemoryInfo(handle) pynvml.nvmlShutdown() - return info.free / info.total + + free_memory = info.free / 1024**2 + total_memory = info.total / 1024**2 + + if ratio: + return free_memory / total_memory + return free_memory def set_gpus(n=1, min_free_memory=0.9, max_processes=2, verbose=False, raise_error=False): """ Set the ``CUDA_VISIBLE_DEVICES`` variable to ``n`` available devices.