vllm_omni.worker.gpu_memory_utils ¶

NVML-based per-process GPU memory utilities.

Shared across worker types (OmniGPUWorkerBase, DiffusionWorker, etc.) for process-scoped GPU memory accounting.

logger `module-attribute` ¶

logger = init_logger(__name__)

get_device_handle ¶

get_device_handle(device_id: str | int)

Get NVML device handle by index or UUID.

get_process_gpu_memory ¶

get_process_gpu_memory(local_rank: int) -> int | None

Get GPU memory used by current process via pynvml.

Supports CUDA_VISIBLE_DEVICES with integer indices, UUIDs, or MIG IDs.

Returns:

Type	Description
`int \| None`	Memory in bytes used by this process, or None if NVML unavailable.

Raises:

Type	Description
`RuntimeError`	If device validation fails (invalid index or UUID).

is_process_scoped_memory_available ¶

is_process_scoped_memory_available() -> bool

Check if NVML process-scoped memory tracking is available.

When True, concurrent stage initialization is safe because each process can accurately measure its own GPU memory via NVML. When False, sequential initialization (file locks) is needed.

parse_cuda_visible_devices ¶

parse_cuda_visible_devices() -> list[str | int]

Parse CUDA_VISIBLE_DEVICES into a list of device identifiers.

Returns list of integers (physical indices) or strings (UUIDs/MIG IDs).