Merge pull request #60 from MoseleyBioinformaticsLab/granular

Adds an option for logging the per timepoint resource usage

Author: erikhuck

requirements.txt

@@ -1,3 +1,4 @@
 psutil>=6.0.0
 docopt>=0.6.2
 pandas>=2.2.3
+SQLAlchemy>=2.0.39

src/gpu_tracker/__main__.py

@@ -4,7 +4,7 @@
 Usage:
     gpu-tracker -h | --help
     gpu-tracker -v | --version
-    gpu-tracker --execute=<command> [--output=<output>] [--format=<format>] [--st=<sleep-time>] [--ru=<ram-unit>] [--gru=<gpu-ram-unit>] [--tu=<time-unit>] [--nec=<num-cores>] [--guuids=<gpu-uuids>] [--disable-logs] [--gb=<gpu-brand>]
+    gpu-tracker --execute=<command> [--output=<output>] [--format=<format>] [--st=<sleep-time>] [--ru=<ram-unit>] [--gru=<gpu-ram-unit>] [--tu=<time-unit>] [--nec=<num-cores>] [--guuids=<gpu-uuids>] [--disable-logs] [--gb=<gpu-brand>] [--tf=<tracking-file>]
 
 Options:
     -h --help               Show this help message and exit.
@@ -20,6 +20,7 @@
     --guuids=<gpu-uuids>    Comma separated list of the UUIDs of the GPUs for which to track utilization e.g. gpu-uuid1,gpu-uuid2,etc. Defaults to all the GPUs in the system.
     --disable-logs          If set, warnings are suppressed during tracking. Otherwise, the Tracker logs warnings as usual.
     --gb=<gpu-brand>        The brand of GPU to profile. Valid values are nvidia and amd. Defaults to the brand of GPU detected in the system, checking NVIDIA first.
+    --tf=<tracking-file>    If specified, stores the individual resource usage measurements at each iteration. Valid file formats are CSV (.csv) and SQLite (.sqlite) where the SQLite file format stores the data in a table called "tracking" and allows for more efficient querying.
 """
 import docopt as doc
 import subprocess as subp
@@ -43,7 +44,8 @@ def main():
         '--nec': 'n_expected_cores',
         '--guuids': 'gpu_uuids',
         '--disable-logs': 'disable_logs',
-        '--gb': 'gpu_brand'
+        '--gb': 'gpu_brand',
+        '--tf': 'tracking_file'
     }
     kwargs = {
         option_map[option]: value for option, value in args.items() if value is not None and option not in {

src/gpu_tracker/_helper_classes.py

@@ -0,0 +1,210 @@
+from __future__ import annotations
+import abc
+import subprocess as subp
+import pandas as pd
+import io
+import os
+import csv
+import dataclasses as dclass
+import sqlalchemy as sqlalc
+import sqlalchemy.orm as sqlorm
+
+
+class _GPUQuerier(abc.ABC):
+    command = None
+
+    @classmethod
+    def _query_gpu(cls, *args) -> pd.DataFrame:
+        output = subp.check_output((cls.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)
+
+    @classmethod
+    def is_available(cls) -> bool | None:
+        try:
+            subp.check_output(cls.command)
+            return True
+        except subp.CalledProcessError:
+            return False
+        except FileNotFoundError:
+            return None
+
+    @classmethod
+    @abc.abstractmethod
+    def static_info(cls) -> pd.DataFrame:
+        pass  # pragma: nocover
+
+    @classmethod
+    @abc.abstractmethod
+    def process_ram(cls) -> pd.DataFrame:
+        pass  # pragma: nocover
+
+    @classmethod
+    @abc.abstractmethod
+    def ram_and_utilization(cls) -> pd.DataFrame:
+        pass  # pragma: nocover
+
+
+class _NvidiaQuerier(_GPUQuerier):
+    command = 'nvidia-smi'
+
+    @classmethod
+    def _query_gpu(cls, *args: str, ram_column: str):
+        gpu_info = super()._query_gpu(*args, '--format=csv')
+        gpu_info.columns = [col.replace('[MiB]', '').replace('[%]', '').strip() for col in gpu_info.columns]
+        gpu_info[ram_column] = gpu_info[ram_column].apply(lambda ram: int(ram.replace('MiB', '').strip()))
+        return gpu_info.rename(columns={ram_column: 'ram'})
+
+    @classmethod
+    def static_info(cls) -> pd.DataFrame:
+        return cls._query_gpu('--query-gpu=uuid,memory.total', ram_column='memory.total')
+
+    @classmethod
+    def process_ram(cls) -> pd.DataFrame:
+        return cls._query_gpu('--query-compute-apps=pid,used_gpu_memory', ram_column='used_gpu_memory')
+
+    @classmethod
+    def ram_and_utilization(cls) -> pd.DataFrame:
+        gpu_info = cls._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):
+    command = 'amd-smi'
+    __id_to_uuid = None
+
+    @classmethod
+    @property
+    def _id_to_uuid(cls) -> dict[int, str]:
+        if cls.__id_to_uuid is None:
+            gpu_info = super()._query_gpu('list', '--csv')
+            cls.__id_to_uuid = {gpu_id: uuid for gpu_id, uuid in zip(gpu_info.gpu, gpu_info.gpu_uuid)}
+        return cls.__id_to_uuid
+
+    @classmethod
+    def _query_gpu(cls, *args: str, ram_column: str) -> pd.DataFrame:
+        gpu_info = super()._query_gpu(*args, '--csv')
+        if 'gpu' in gpu_info.columns:
+            gpu_info.gpu = [cls._id_to_uuid[gpu_id] for gpu_id in gpu_info.gpu]
+            gpu_info = gpu_info.rename(columns={'gpu': 'uuid'})
+        return gpu_info.rename(columns={ram_column: 'ram'})
+
+    @classmethod
+    def static_info(cls) -> pd.DataFrame:
+        gpu_info = cls._query_gpu('static', '--vram', ram_column='size')
+        return gpu_info[['uuid', 'ram']]
+
+    @classmethod
+    def process_ram(cls) -> pd.DataFrame:
+        gpu_info = cls._query_gpu('process', ram_column='vram_mem')
+        gpu_info.ram = [ram / 1e6 for ram in gpu_info.ram]  # RAM is in bytes for the process subcommand.
+        return gpu_info[['pid', 'ram']]
+
+    @classmethod
+    def ram_and_utilization(cls) -> pd.DataFrame:
+        gpu_info = cls._query_gpu('monitor', '--vram-usage', '--gfx', ram_column='vram_used')
+        gpu_info = gpu_info[['uuid', 'gfx', 'ram']]
+        gpu_info.gfx = gpu_info.gfx.astype(float)
+        return gpu_info.rename(columns={'gfx': 'utilization_percent'})
+
+
+@dclass.dataclass
+class TimepointUsage:
+    main_ram: float = 0.0
+    descendants_ram: float = 0.0
+    combined_ram: float = 0.0
+    system_ram: float = 0.0
+    main_gpu_ram: float = 0.0
+    descendants_gpu_ram: float = 0.0
+    combined_gpu_ram: float = 0.0
+    system_gpu_ram: float = 0.0
+    gpu_sum_utilization_percent: float = 0.0
+    gpu_hardware_utilization_percent: float = 0.0
+    main_n_threads: int = 0
+    descendants_n_threads: int = 0
+    combined_n_threads: int = 0
+    cpu_system_sum_utilization_percent: float = 0.0,
+    cpu_system_hardware_utilization_percent: float = 0.0
+    cpu_main_sum_utilization_percent: float = 0.0
+    cpu_main_hardware_utilization_percent: float = 0.0
+    cpu_descendants_sum_utilization_percent: float = 0.0
+    cpu_descendants_hardware_utilization_percent: float = 0.0
+    cpu_combined_sum_utilization_percent: float = 0.0
+    cpu_combined_hardware_utilization_percent: float = 0.0
+    timestamp: float = 0.0
+
+
+class _TrackingFile(abc.ABC):
+    @staticmethod
+    def create(file: str | None) -> _TrackingFile | None:
+        if file is not None:
+            if file.endswith('.csv'):
+                return _CSVTrackingFile(file)
+            elif file.endswith('.sqlite'):
+                return _SQLiteTrackingFile(file)
+            else:
+                raise ValueError(
+                    f'Invalid file name: "{file}". Valid file extensions are ".csv" and ".sqlite".')
+        else:
+            return None
+
+    def __init__(self, file: str):
+        self._file = file
+
+    def write_row(self, values: TimepointUsage):
+        values = dclass.asdict(values)
+        if not os.path.isfile(self._file):
+            self._create_file(values)
+        self._write_row(values)
+
+    @abc.abstractmethod
+    def _write_row(self, values: dict):
+        pass  # pragma: nocover
+
+    @abc.abstractmethod
+    def _create_file(self, values: dict):
+        pass  # pragma: nocover
+
+
+class _CSVTrackingFile(_TrackingFile):
+    def _write_row(self, values: dict):
+        with open(self._file, 'a', newline='') as f:
+            writer = csv.DictWriter(f, fieldnames=values.keys())
+            writer.writerow(values)
+
+    def _create_file(self, values: dict):
+        with open(self._file, 'w', newline='') as f:
+            writer = csv.DictWriter(f, fieldnames=values.keys())
+            writer.writeheader()
+
+
+class _SQLiteTrackingFile(_TrackingFile):
+    _SQLITE_TABLE_NAME = 'tracking'
+
+    def _write_row(self, values: dict):
+        engine = sqlalc.create_engine(f'sqlite:///{self._file}', poolclass=sqlalc.pool.NullPool)
+        metadata = sqlalc.MetaData()
+        tracking_table = sqlalc.Table(_SQLiteTrackingFile._SQLITE_TABLE_NAME, metadata, autoload_with=engine)
+        Session = sqlorm.sessionmaker(bind=engine)
+        with Session() as session:
+            insert_stmt = sqlalc.insert(tracking_table).values(**values)
+            session.execute(insert_stmt)
+            session.commit()
+
+    def _create_file(self, values: dict):
+        engine = sqlalc.create_engine(f'sqlite:///{self._file}', poolclass=sqlalc.pool.NullPool)
+        metadata = sqlalc.MetaData()
+        type_mapping = {
+            str: sqlalc.String,
+            int: sqlalc.Integer,
+            float: sqlalc.Float,
+        }
+        columns = list[sqlalc.Column]()
+        schema = {name: type(value) for name, value in values.items()}
+        for column_name, data_type in schema.items():
+            sqlalchemy_type = type_mapping[data_type]
+            columns.append(sqlalc.Column(column_name, sqlalchemy_type))
+        tracking_table = sqlalc.Table(_SQLiteTrackingFile._SQLITE_TABLE_NAME, metadata, *columns)
+        metadata.create_all(engine)