Source code for perfmon.core.metrics.gpumetrics.nvidia.memory

"""Functions to monitor memory related metrics for NVIDIA GPUs"""

import logging

from py3nvml.py3nvml import *

from perfmon.core.metrics.gpumetrics.nvidia import device_query

_log = logging.getLogger(__name__)

# pylint: disable=E0401,W0201,C0301


def memory_usage(data):
    """This method gets NVIDIA GPU memory and BAR1 memory usage"""

    # Get memory usage
    memory_metrics = device_query('nvmlDeviceGetMemoryInfo')
    for i, mem in enumerate(memory_metrics):
        data[i]['memory_used']['gpu_memory'].append(int(mem.used))

    # Get BAR1 memory usage
    bar1_memory_metrics = device_query('nvmlDeviceGetBAR1MemoryInfo')
    for i, mem in enumerate(bar1_memory_metrics):
        data[i]['memory_used']['bar1_memory'].append(int(mem.bar1Used))

    return data