cuda

Code inside this file can safely assume cuda platform, e.g. importing pynvml. However, it should not initialize cuda context.

Classes

fastvideo.platforms.cuda.CudaPlatformBase

Bases: Platform

Functions

fastvideo.platforms.cuda.CudaPlatformBase.get_torch_device classmethod

get_torch_device()

Return torch.cuda

Source code in fastvideo/platforms/cuda.py

@classmethod
def get_torch_device(cls):
    """
    Return torch.cuda
    """
    return torch.cuda

fastvideo.platforms.cuda.NvmlCudaPlatform

Bases: CudaPlatformBase

Functions

fastvideo.platforms.cuda.NvmlCudaPlatform.is_full_nvlink classmethod

is_full_nvlink(physical_device_ids: list[int]) -> bool

query if the set of gpus are fully connected by nvlink (1 hop)

Source code in fastvideo/platforms/cuda.py

@classmethod
@with_nvml_context
def is_full_nvlink(cls, physical_device_ids: list[int]) -> bool:
    """
    query if the set of gpus are fully connected by nvlink (1 hop)
    """
    handles = [
        pynvml.nvmlDeviceGetHandleByIndex(i) for i in physical_device_ids
    ]
    for i, handle in enumerate(handles):
        for j, peer_handle in enumerate(handles):
            if i < j:
                try:
                    p2p_status = pynvml.nvmlDeviceGetP2PStatus(
                        handle,
                        peer_handle,
                        pynvml.NVML_P2P_CAPS_INDEX_NVLINK,
                    )
                    if p2p_status != pynvml.NVML_P2P_STATUS_OK:
                        return False
                except pynvml.NVMLError:
                    logger.exception(
                        "NVLink detection failed. This is normal if"
                        " your machine has no NVLink equipped.")
                    return False
    return True

Functions