memory.py

from torch import cuda


def get_less_used_gpu(gpus=None, debug=False):
    """Inspect cached/reserved and allocated memory on specified gpus and return the id of the less used device"""
    if gpus is None:
        warn = "Falling back to default: all gpus"
        gpus = range(cuda.device_count())
    elif isinstance(gpus, str):
        gpus = [int(el) for el in gpus.split(",")]

    # check gpus arg VS available gpus
    sys_gpus = list(range(cuda.device_count()))
    if len(gpus) > len(sys_gpus):
        gpus = sys_gpus
        warn = f"WARNING: Specified {len(gpus)} gpus, but only {cuda.device_count()} available. Falling back to default: all gpus.\nIDs:\t{list(gpus)}"
    elif set(gpus).difference(sys_gpus):
        # take correctly specified and add as much bad specifications as unused system gpus
        available_gpus = set(gpus).intersection(sys_gpus)
        unavailable_gpus = set(gpus).difference(sys_gpus)
        unused_gpus = set(sys_gpus).difference(gpus)
        gpus = list(available_gpus) + list(unused_gpus)[: len(unavailable_gpus)]
        warn = f"GPU ids {unavailable_gpus} not available. Falling back to {len(gpus)} device(s).\nIDs:\t{list(gpus)}"

    cur_allocated_mem = {}
    cur_cached_mem = {}
    max_allocated_mem = {}
    max_cached_mem = {}
    for i in gpus:
        cur_allocated_mem[i] = cuda.memory_allocated(i)
        cur_cached_mem[i] = cuda.memory_reserved(i)
        max_allocated_mem[i] = cuda.max_memory_allocated(i)
        max_cached_mem[i] = cuda.max_memory_reserved(i)
    min_allocated = min(cur_allocated_mem, key=cur_allocated_mem.get)
    # if debug:
    print(warn)
    print("Current allocated memory:", {f"cuda:{k}": v for k, v in cur_allocated_mem.items()})
    print("Current reserved memory:", {f"cuda:{k}": v for k, v in cur_cached_mem.items()})
    print("Maximum allocated memory:", {f"cuda:{k}": v for k, v in max_allocated_mem.items()})
    print("Maximum reserved memory:", {f"cuda:{k}": v for k, v in max_cached_mem.items()})
    print("Suggested GPU:", min_allocated)
    # return min_allocated


def free_memory(to_delete: list, debug=False):
    import gc
    import inspect

    calling_namespace = inspect.currentframe().f_back
    if debug:
        print("Before:")
        get_less_used_gpu(debug=True)

    for _var in to_delete:
        calling_namespace.f_locals.pop(_var, None)
        gc.collect()
        cuda.empty_cache()
    if debug:
        print("After:")
        get_less_used_gpu(debug=True)


get_less_used_gpu()