`nemo_rl.utils.nvml`#

Module Contents#

Functions#

`nvml_context`	Context manager for NVML initialization and shutdown.
`device_id_to_physical_device_id`	Convert a logical device ID to a physical device ID considering CUDA_VISIBLE_DEVICES.
`get_device_uuid`	Get the UUID of a CUDA device using NVML.
`get_free_memory_bytes`	Get the free memory of a CUDA device in bytes using NVML.

API#

nemo_rl.utils.nvml.nvml_context() → Generator[None, None, None]#

Context manager for NVML initialization and shutdown.

Raises:: RuntimeError – If NVML initialization fails

nemo_rl.utils.nvml.device_id_to_physical_device_id(device_id: int) → int#: Convert a logical device ID to a physical device ID considering CUDA_VISIBLE_DEVICES.

nemo_rl.utils.nvml.get_device_uuid(device_idx: int) → str#: Get the UUID of a CUDA device using NVML.

nemo_rl.utils.nvml.get_free_memory_bytes(device_idx: int) → float#: Get the free memory of a CUDA device in bytes using NVML.