import os
import time


def snapshot(phase):
    result = {'time':time.time(),'phase':phase}
    try:
        import psutil
        import pynvml as nv
        nv.nvmlInit()
        handle = nv.nvmlDeviceGetHandleByIndex(0)
        memory = nv.nvmlDeviceGetMemoryInfo(handle)
        pids = {os.getpid()} | {p.pid for p in psutil.Process().children(recursive=True)}
        processes = nv.nvmlDeviceGetComputeRunningProcesses(handle)
        result.update(device_total_bytes=memory.total,device_free_bytes=memory.free,device_used_bytes=memory.used,
                      own_processes_vram_bytes=sum(p.usedGpuMemory for p in processes if p.pid in pids and p.usedGpuMemory <= memory.total))
        nv.nvmlShutdown()
    except Exception as error:
        result['measurement_error'] = type(error).__name__
    return result