Separates the GPU querying into two classes, one for nvidia and another for amd, that inherit a common interface

Author: erikhuck

src/gpu_tracker/tracker.py

@@ -1,5 +1,6 @@
 """The ``tracker`` module contains the ``Tracker`` class which can alternatively be imported directly from the ``gpu_tracker`` package."""
 from __future__ import annotations
+import abc
 import json
 import dataclasses as dclass
 import platform
@@ -17,6 +18,62 @@
 import pandas as pd
 
 
+class _GPUQuerier(abc.ABC):
+    def __init__(self, command: str):
+        self.command = command
+
+    def _query_gpu(self, *args) -> pd.DataFrame:
+        output = subp.check_output((self.command,) + args, stderr=subp.STDOUT).decode()
+        gpu_info = pd.read_csv(io.StringIO(output))
+        return gpu_info.map(lambda value: value.strip() if type(value) is str else value)
+
+    def is_available(self) -> bool:
+        try:
+            subp.check_output(self.command)
+            return True
+        except FileNotFoundError:
+            return False
+
+    @abc.abstractmethod
+    def static_info(self) -> pd.DataFrame:
+        pass
+
+    @abc.abstractmethod
+    def process_ram(self) -> pd.DataFrame:
+        pass
+
+    @abc.abstractmethod
+    def ram_and_utilization(self) -> pd.DataFrame:
+        pass
+
+class _NvidiaQuerier(_GPUQuerier):
+    def __init__(self):
+        super().__init__('nvidia-smi')
+
+    def _query_gpu(self, *args: list[str], ram_column: str | None = None):
+        gpu_info = super()._query_gpu(*args, '--format=csv')
+        gpu_info.columns = [col.replace('[MiB]', '').replace('[%]', '').strip() for col in gpu_info.columns]
+        if ram_column:
+            gpu_info[ram_column] = gpu_info[ram_column].apply(lambda ram: int(ram.replace('MiB', '').strip()))
+            gpu_info = gpu_info.rename(columns={ram_column: 'ram'})
+        return gpu_info.map(lambda value: value.strip() if type(value) is str else value)
+
+    def static_info(self) -> pd.DataFrame:
+        return self._query_gpu('--query-gpu=uuid,memory.total', ram_column='memory.total')
+
+    def process_ram(self) -> pd.DataFrame:
+        return self._query_gpu('--query-compute-apps=pid,used_gpu_memory', ram_column='used_gpu_memory')
+
+    def ram_and_utilization(self) -> pd.DataFrame:
+        gpu_info = self._query_gpu('--query-gpu=uuid,memory.used,utilization.gpu', ram_column='memory.used')
+        gpu_info = gpu_info.rename(columns={'utilization.gpu': 'utilization_percent'})
+        gpu_info.utilization_percent = [float(percentage.replace('%', '').strip()) for percentage in gpu_info.utilization_percent]
+        return gpu_info
+
+class _AMDQuerier(_GPUQuerier):
+    def __init__(self):
+        super().__init__('amd-smi')
+
 class _TrackingProcess(mproc.Process):
     _CPU_PERCENT_INTERVAL = 0.1
     _ram_unit2coefficient = {
@@ -43,7 +100,7 @@ class _TrackingProcess(mproc.Process):
     def __init__(
             self, stop_event: mproc.Event, sleep_time: float, ram_unit: str, gpu_ram_unit: str, time_unit: str,
             n_expected_cores: int | None, gpu_uuids: set[str] | None, disable_logs: bool, main_process_id: int,
-            resource_usage_file: str, extraneous_process_ids: set[int]):
+            resource_usage_file: str, extraneous_process_ids: set[int], gpu_querier: _GPUQuerier):
         super().__init__()
         self._stop_event = stop_event
         if sleep_time < _TrackingProcess._CPU_PERCENT_INTERVAL:
@@ -63,24 +120,22 @@ def __init__(
         self._hardware_percent_sums = {key: 0. for key in percent_keys}
         self._tracking_iteration = 1
         self._is_linux = platform.system().lower() == 'linux'
-        self._nvidia_available = True
-        try:
-            subp.check_output('nvidia-smi')
-        except FileNotFoundError:
-            self._nvidia_available = False
+        self._gpu_querier = gpu_querier
+        self._gpu_available = self._gpu_querier.is_available()
+        if not self._gpu_available:
             self._log_warning(
-                'The nvidia-smi command is not available. Please install the Nvidia drivers to track GPU usage. '
-                'Otherwise the Max GPU RAM values will remain 0.0')
+                f'The {self._gpu_querier.command} command is not available. Please install it to track GPU usage. '
+                'Otherwise the GPU RAM and GPU utilization values will remain 0.0.')
         max_ram = MaxRAM(unit=ram_unit, system_capacity=psutil.virtual_memory().total * self._ram_coefficient)
         system_core_count = psutil.cpu_count()
         cpu_utilization = CPUUtilization(
             system_core_count=system_core_count,
             n_expected_cores=n_expected_cores if n_expected_cores is not None else system_core_count)
-        if self._nvidia_available:
-            gpu_info = _TrackingProcess._query_gpu(nvidia_command='--query-gpu=uuid,memory.total')
-            gpu_ram_system_capacity = self._get_gpu_ram(gpu_info=gpu_info, column='memory.total')
+        if self._gpu_available:
+            gpu_info = self._gpu_querier.static_info()
+            gpu_ram_system_capacity = self._get_gpu_ram(gpu_info=gpu_info)
             max_gpu_ram = MaxGPURAM(unit=gpu_ram_unit, system_capacity=gpu_ram_system_capacity)
-            all_uuids = set(gpu_info['uuid'])
+            all_uuids =   set(gpu_info.uuid)
             if gpu_uuids is None:
                 self._gpu_uuids = all_uuids
             else:
@@ -143,25 +198,23 @@ def run(self):
                 self._resource_usage.max_ram.system = max(
                     self._resource_usage.max_ram.system, psutil.virtual_memory().used * self._ram_coefficient)
                 # Get the maximum GPU RAM usage if available.
-                if self._nvidia_available:  # pragma: nocover
-                    gpu_info = _TrackingProcess._query_gpu(nvidia_command='--query-compute-apps=pid,used_gpu_memory')
+                if self._gpu_available:  # pragma: nocover
+                    gpu_info = self._gpu_querier.process_ram()
                     if len(gpu_info):
                         process_ids = {self._main_process_id}
                         self._update_gpu_ram(attr='main', process_ids=process_ids, gpu_info=gpu_info)
                         process_ids = set(self._map_processes(processes=descendant_processes, map_func=lambda process: process.pid))
                         self._update_gpu_ram(attr='descendants', process_ids=process_ids, gpu_info=gpu_info)
                         process_ids.add(self._main_process_id)
                         self._update_gpu_ram(attr='combined', process_ids=process_ids, gpu_info=gpu_info)
-                    gpu_info = _TrackingProcess._query_gpu(nvidia_command='--query-gpu=uuid,memory.used,utilization.gpu')
-                    system_gpu_ram = self._get_gpu_ram(gpu_info, column='memory.used')
+                    gpu_info = self._gpu_querier.ram_and_utilization()
+                    system_gpu_ram = self._get_gpu_ram(gpu_info)
                     self._resource_usage.max_gpu_ram.system = max(self._resource_usage.max_gpu_ram.system, system_gpu_ram)
-                    gpu_info = gpu_info.loc[gpu_info['uuid'].apply(lambda gpu_uuid: gpu_uuid in self._gpu_uuids)]
-                    gpu_percentages = [float(percentage.replace('%', '').strip()) for percentage in gpu_info['utilization.gpu']]
+                    gpu_info = gpu_info.loc[[uuid in self._gpu_uuids for uuid in gpu_info.uuid]]
                     self._update_processing_unit_utilization(
-                        current_percentages=gpu_percentages,
+                        current_percentages=list(gpu_info.utilization_percent),
                         processing_unit_percentages=self._resource_usage.gpu_utilization.gpu_percentages, percent_key='gpu',
                         n_hardware_units=self._resource_usage.gpu_utilization.n_expected_gpus)
-
                 # Get the mean and maximum CPU usages.
                 main_n_threads = self._map_processes([main_process], map_func=get_n_threads)
                 descendant_n_threads = self._map_processes(descendant_processes, map_func=get_n_threads)
@@ -230,23 +283,13 @@ def _update_ram(self, rss_values: RSSValues, memory_maps_list: list[list] | None
         rss_values.total_rss = max(rss_values.total_rss, total_rss)
 
     def _update_gpu_ram(self, attr: str, process_ids: set[int], gpu_info: pd.DataFrame):
-        gpu_info = gpu_info.loc[[pid in process_ids for pid in gpu_info['pid']]]
-        gpu_ram = self._get_gpu_ram(gpu_info, column='used_gpu_memory')
+        gpu_info = gpu_info.loc[[pid in process_ids for pid in gpu_info.pid]]
+        gpu_ram = self._get_gpu_ram(gpu_info)
         max_gpu_ram = getattr(self._resource_usage.max_gpu_ram, attr)
         setattr(self._resource_usage.max_gpu_ram, attr, max(max_gpu_ram, gpu_ram))
 
-    @staticmethod
-    def _query_gpu(nvidia_command: str) -> pd.DataFrame:
-        command = f'nvidia-smi {nvidia_command} --format=csv'
-        output = subp.check_output(command.split(), stderr=subp.STDOUT).decode()
-        gpu_info = pd.read_csv(io.StringIO(output))
-        gpu_info.columns = [col.replace('[MiB]', '').replace('[%]', '').strip() for col in gpu_info.columns]
-        return gpu_info.map(lambda value: value.strip() if type(value) is str else value)
-
-    def _get_gpu_ram(self, gpu_info: pd.DataFrame, column: str) -> float:
-        gpu_rams = gpu_info[column]
-        gpu_rams = gpu_rams.apply(lambda ram: int(ram.replace('MiB', '').strip()))
-        return sum(gpu_rams) * self._gpu_ram_coefficient
+    def _get_gpu_ram(self, gpu_info: pd.DataFrame) -> float:
+        return sum(gpu_info.ram) * self._gpu_ram_coefficient
 
     def _update_processing_unit_utilization(
             self, current_percentages: list[float], processing_unit_percentages: ProcessingUnitPercentages,
@@ -297,7 +340,7 @@ class State(enum.Enum):
     def __init__(
             self, sleep_time: float = 1.0, ram_unit: str = 'gigabytes', gpu_ram_unit: str = 'gigabytes', time_unit: str = 'hours',
             n_expected_cores: int = None, gpu_uuids: set[str] = None, disable_logs: bool = False, process_id: int = None,
-            resource_usage_file: str | None = None, n_join_attempts: int = 5, join_timeout: float = 10.0):
+            resource_usage_file: str | None = None, n_join_attempts: int = 5, join_timeout: float = 10.0, gpu_brand: str = 'nvidia'):
         """
         :param sleep_time: The number of seconds to sleep in between usage-collection iterations.
         :param ram_unit: One of 'bytes', 'kilobytes', 'megabytes', 'gigabytes', or 'terabytes'.
@@ -321,9 +364,15 @@ def __init__(
         self._stop_event = mproc.Event()
         extraneous_ids = {process.pid for process in current_process.children()} - legit_child_ids
         self._resource_usage_file = f'.gpu-tracker_{uuid.uuid1()}.pkl' if resource_usage_file is None else resource_usage_file
+        if gpu_brand == 'nvidia':
+            gpu_querier = _NvidiaQuerier()
+        elif gpu_brand == 'amd':
+            gpu_querier = _AMDQuerier()
+        else:
+            raise ValueError(f'"{gpu_brand}" is not a valid GPU brand. Supported values are "nvidia" and "amd".')
         self._tracking_process = _TrackingProcess(
             self._stop_event, sleep_time, ram_unit, gpu_ram_unit, time_unit, n_expected_cores, gpu_uuids, disable_logs,
-            process_id if process_id is not None else current_process_id, self._resource_usage_file, extraneous_ids)
+            process_id if process_id is not None else current_process_id, self._resource_usage_file, extraneous_ids, gpu_querier)
         self.resource_usage = None
         self.n_join_attempts = n_join_attempts
         self.join_timeout = join_timeout

tests/test_tracker.py

@@ -5,8 +5,8 @@
 import pytest as pt
 import utils
 
-nvidia_smi_unavailable_message = 'The nvidia-smi command is not available. Please install the Nvidia drivers to track GPU usage. ' \
-                                 'Otherwise the Max GPU RAM values will remain 0.0'
+nvidia_smi_unavailable_message = 'The nvidia-smi command is not available. Please install it to track GPU usage. ' \
+                                 'Otherwise the GPU RAM and GPU utilization values will remain 0.0.'
 
 
 @pt.fixture(name='operating_system', params=['Linux', 'not-linux'])