Source code for ewoksid02.utils.gpu

from __future__ import annotations

import logging

import numpy

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

try:
    import cupy
except ImportError:
    cupy = None
    CUPY_AVAILABLE = False
    CUPY_MEM_POOL = None
else:
    CUPY_AVAILABLE = True
    CUPY_MEM_POOL = cupy.get_default_memory_pool()




def cupy_available():
    return CUPY_AVAILABLE





def get_device_name():
    device_id = cupy.cuda.Device().id
    device_props = cupy.cuda.runtime.getDeviceProperties(device_id)
    device_name = device_props["name"].decode()
    return device_name





def log_allocated_gpu_memory():
    """Log the status of GPU memory"""
    available_mem_GB, total_mem_GB = numpy.array(cupy.cuda.runtime.memGetInfo()) / 1e9
    mem_usage_GB = total_mem_GB - available_mem_GB
    device_name = get_device_name()

    if available_mem_GB / total_mem_GB < 0.1:
        mem_message = "Low memory available"
        color_prefix = "\033[91m"
        log_level = "warning"
    elif available_mem_GB / total_mem_GB < 0.3:
        mem_message = "Medium memory available"
        color_prefix = "\033[93m"
        log_level = "warning"
    else:
        mem_message = "Sufficient memory available"
        color_prefix = "\033[92m"
        log_level = "info"
    color_suffix = "\033[0m"

    msg = f"{color_prefix}Memory on {device_name}: {mem_usage_GB:.2f}GB used; {available_mem_GB:.2f}GB available. {mem_message}{color_suffix}"

    if log_level == "warning":
        logger.warning(msg)
    elif log_level == "info":
        logger.info(msg)





def get_best_gpu():
    """Decides the best GPU in terms of memory available"""
    if not cupy_available():
        logger.warning("Cupy is not available.")
        return None

    best_device = None
    max_free_memory = 0

    for device_id in range(cupy.cuda.runtime.getDeviceCount()):
        free_memory = get_free_memory(device_id)
        if free_memory is not None and free_memory > max_free_memory:
            max_free_memory = free_memory
            best_device = device_id

    return best_device





def use_best_gpu():
    """
    Set the best available GPU for cupy operations.
    """
    best_device = get_best_gpu()
    if best_device is not None:
        cupy.cuda.Device(best_device).use()
        logger.info(f"Using GPU {best_device} with the most free memory.")
    else:
        logger.warning("No suitable GPU found or cupy is not available.")





def get_free_memory(device_id):
    """Retrieves the available memory on a GPU device"""
    if not cupy_available():
        logger.warning("Cupy is not available.")
        return None

    with cupy.cuda.Device(device_id):
        free_mem, total_mem = cupy.cuda.runtime.memGetInfo()
        return free_mem