dynamo/lib/gpu_memory_service/client/cuda_vmm_utils.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Client-side CUDA VMM utilities.

These functions wrap CUDA driver API calls used by the client memory manager
for importing, mapping, and unmapping GPU memory.
"""

from __future__ import annotations

import os

from cuda.bindings import driver as cuda
from gpu_memory_service.common.cuda_vmm_utils import check_cuda_result
from gpu_memory_service.common.types import GrantedLockType


def import_handle_from_fd(fd: int) -> int:
    """Import a CUDA memory handle from a file descriptor.

    Closes the FD after import — the imported handle holds its own reference
    to the physical allocation. Leaving the FD open leaks a DMA-buf ref that
    prevents cuMemRelease from freeing GPU memory.

    Args:
        fd: POSIX file descriptor received via SCM_RIGHTS.

    Returns:
        CUDA memory handle.
    """
    try:
        result, handle = cuda.cuMemImportFromShareableHandle(
            fd,
            cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR,
        )
        check_cuda_result(result, "cuMemImportFromShareableHandle")
        return int(handle)
    finally:
        os.close(fd)


def reserve_va(size: int, granularity: int) -> int:
    """Reserve virtual address space.

    Args:
        size: Size in bytes (should be aligned to granularity).
        granularity: VMM allocation granularity.

    Returns:
        Reserved virtual address.
    """
    result, va = cuda.cuMemAddressReserve(size, granularity, 0, 0)
    check_cuda_result(result, "cuMemAddressReserve")
    return int(va)


def free_va(va: int, size: int) -> None:
    """Free a virtual address reservation.

    Args:
        va: Virtual address to free.
        size: Size of the reservation.
    """
    (result,) = cuda.cuMemAddressFree(va, size)
    check_cuda_result(result, "cuMemAddressFree")


def map_to_va(va: int, size: int, handle: int) -> None:
    """Map a CUDA handle to a virtual address.

    Args:
        va: Virtual address (must be reserved).
        size: Size of the mapping.
        handle: CUDA memory handle.
    """
    (result,) = cuda.cuMemMap(va, size, 0, handle, 0)
    check_cuda_result(result, "cuMemMap")


def set_access(va: int, size: int, device: int, access: GrantedLockType) -> None:
    """Set access permissions for a mapped region.

    Args:
        va: Virtual address.
        size: Size of the region.
        device: CUDA device index.
        access: Access mode - RO for read-only, RW for read-write.
    """
    acc = cuda.CUmemAccessDesc()
    acc.location.type = cuda.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
    acc.location.id = device
    acc.flags = (
        cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READ
        if access == GrantedLockType.RO
        else cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
    )
    (result,) = cuda.cuMemSetAccess(va, size, [acc], 1)
    check_cuda_result(result, "cuMemSetAccess")


def unmap(va: int, size: int) -> None:
    """Unmap a virtual address region.

    Args:
        va: Virtual address to unmap.
        size: Size of the mapping.
    """
    (result,) = cuda.cuMemUnmap(va, size)
    check_cuda_result(result, "cuMemUnmap")


def release_handle(handle: int) -> None:
    """Release a CUDA memory handle.

    Args:
        handle: CUDA memory handle to release.
    """
    (result,) = cuda.cuMemRelease(handle)
    check_cuda_result(result, "cuMemRelease")


def validate_pointer(va: int) -> bool:
    """Validate that a mapped VA is accessible.

    Returns True if the pointer is valid, False otherwise (logs a warning).
    """
    result, _dev_ptr = cuda.cuPointerGetAttribute(
        cuda.CUpointer_attribute.CU_POINTER_ATTRIBUTE_DEVICE_POINTER, va
    )
    if result != cuda.CUresult.CUDA_SUCCESS:
        err_result, err_str = cuda.cuGetErrorString(result)
        err_msg = ""
        if err_result == cuda.CUresult.CUDA_SUCCESS and err_str:
            err_msg = err_str.decode() if isinstance(err_str, bytes) else str(err_str)
        import logging

        logging.getLogger(__name__).warning(
            "cuPointerGetAttribute failed for VA 0x%x: %s (%s)",
            va,
            result,
            err_msg,
        )
        return False
    return True


def synchronize() -> None:
    """Synchronize the current CUDA context.

    Blocks until all preceding commands in the current context have completed.
    """
    (result,) = cuda.cuCtxSynchronize()
    check_cuda_result(result, "cuCtxSynchronize")


def set_current_device(device: int) -> None:
    """Set the current CUDA device by activating its primary context.

    Args:
        device: CUDA device index.
    """
    result, ctx = cuda.cuDevicePrimaryCtxRetain(device)
    check_cuda_result(result, "cuDevicePrimaryCtxRetain")
    (result,) = cuda.cuCtxSetCurrent(ctx)
    check_cuda_result(result, "cuCtxSetCurrent")