`nemo_rl.utils.nvml`#

Module Contents#

`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.
`_resolve_device_id`	Resolve the logical CUDA device ID.
`log_gpu_memory_diagnostics`	Log detailed GPU memory diagnostics with a greppable [GPU_DIAG] prefix.

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

Context manager for NVML initialization and shutdown.

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.

nemo_rl.utils.nvml._resolve_device_id(device_id=None)#

Resolve the logical CUDA device ID.

Priority: explicit argument > torch.cuda.current_device() > LOCAL_RANK env > 0.

nemo_rl.utils.nvml.log_gpu_memory_diagnostics( *, label: str, worker_type: str, device_id=None, extra_context: str = '', )#

Log detailed GPU memory diagnostics with a greppable [GPU_DIAG] prefix.

This function is designed to never crash – every NVML/PyTorch call is individually guarded. It is safe to call before CUDA is initialized.