_memory_monitor.py

import time
from multiprocessing.sharedctypes import Synchronized
from multiprocessing.synchronize import Event as EventClass

import nvidia_smi  # nvidia-ml-py3

from medium_utils import get_pid_ram_used_bytes
from medium_utils import get_total_ram_available_bytes
from medium_utils import get_total_ram_used_bytes


def _memory_monitor(pid: int, stop_event: EventClass,
                    max_ram_pid: Synchronized,
                    max_ram: Synchronized,
                    max_vram: Synchronized,
                    total_ram: Synchronized,
                    total_vram: Synchronized) -> None:
    """Monitor Memory consumption in parallel to a given process (RAM and VRAM).

    Args:
        pid: ID of the Process to monitor
        stop_event: Shared event triggered when the monitoring has to stop
        max_ram_pid: Peak number of bytes used by the profiled process in the RAM
        max_ram: Peak number of bytes used in the RAM
        max_vram: Peak number of bytes used in the Video RAM (GPU)
        total_ram: Total number of bytes available in RAM (used+free)
        total_vram: Total number of bytes available in VRAM (used+free)
    """
    max_ram_pid_value = 0
    max_ram_value = 0
    max_vram_value = 0
    
    nvidia_smi.nvmlInit()
    assert nvidia_smi.nvmlDeviceGetCount() == 1  # Assume a single GPU
    handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)
    total_ram.value = get_total_ram_available_bytes()
    total_vram.value = int(nvidia_smi.nvmlDeviceGetMemoryInfo(handle).total)

    while not stop_event.is_set():
        max_ram_pid_value = max(max_ram_pid_value, get_pid_ram_used_bytes(pid))
        max_ram_value = max(max_ram_value, get_total_ram_used_bytes())
        gpu_info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
        max_vram_value = max(max_vram_value, gpu_info.used)
        time.sleep(0.2)  # Arbitrary time interval (in s)
    
    max_ram_pid.value = max_ram_pid_value
    max_ram.value = max_ram_value
    max_vram.value = max_vram_value
    nvidia_smi.nvmlShutdown()