end000's picture
Release Sol014 mathematical specialization with paired evaluations, regressions and exact replay artifacts
fc32096 verified
Raw History Blame Contribute Delete
824 Bytes
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