diff --git a/src/gpu_tracker/__main__.py b/src/gpu_tracker/__main__.py index 6fc6ac2..d71305f 100644 --- a/src/gpu_tracker/__main__.py +++ b/src/gpu_tracker/__main__.py @@ -20,7 +20,7 @@ --guuids= 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= 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= 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. + --tf= 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 "data" and allows for more efficient querying. """ import docopt as doc import subprocess as subp diff --git a/src/gpu_tracker/_helper_classes.py b/src/gpu_tracker/_helper_classes.py index 1879eb0..0517cf0 100644 --- a/src/gpu_tracker/_helper_classes.py +++ b/src/gpu_tracker/_helper_classes.py @@ -141,20 +141,20 @@ class _TimepointUsage: class _SubTrackerLog: class CodeBlockPosition(enum.Enum): START = 'START' - END = 'END' + STOP = 'STOP' code_block_name: str position: CodeBlockPosition timestamp: float -class _TrackingFile(abc.ABC): +class _Writer(abc.ABC): @staticmethod - def create(file: str | None) -> _TrackingFile | None: + def create(file: str | None) -> _Writer | None: if file is not None: if file.endswith('.csv'): - return _CSVTrackingFile(file) + return _CSVWriter(file) elif file.endswith('.sqlite'): - return _SQLiteTrackingFile(file) + return _SQLiteWriter(file) else: raise ValueError( f'Invalid file name: "{file}". Valid file extensions are ".csv" and ".sqlite".') @@ -164,7 +164,7 @@ def create(file: str | None) -> _TrackingFile | None: def __init__(self, file: str): self._file = file - def write_row(self, values: _TimepointUsage | _SubTrackerLog): + def write_row(self, values: object): values = dclass.asdict(values) if not os.path.isfile(self._file): self._create_file(values) @@ -179,7 +179,7 @@ def _create_file(self, values: dict): pass # pragma: nocover -class _CSVTrackingFile(_TrackingFile): +class _CSVWriter(_Writer): def _write_row(self, values: dict): with open(self._file, 'a', newline='') as f: writer = csv.DictWriter(f, fieldnames=values.keys()) @@ -191,13 +191,14 @@ def _create_file(self, values: dict): writer.writeheader() -class _SQLiteTrackingFile(_TrackingFile): - _SQLITE_TABLE_NAME = 'tracking' +class _SQLiteWriter(_Writer): + _DATA_TABLE = 'data' + _STATIC_DATA_TABLE = 'static_data' 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) + tracking_table = sqlalc.Table(_SQLiteWriter._DATA_TABLE, metadata, autoload_with=engine) Session = sqlorm.sessionmaker(bind=engine) with Session() as session: insert_stmt = sqlalc.insert(tracking_table).values(**values) @@ -217,5 +218,5 @@ def _create_file(self, values: dict): for column_name, data_type in schema.items(): sqlalchemy_type = type_mapping[data_type] columns.append(sqlalc.Column(column_name, sqlalchemy_type)) - sqlalc.Table(_SQLiteTrackingFile._SQLITE_TABLE_NAME, metadata, *columns) + sqlalc.Table(_SQLiteWriter._DATA_TABLE, metadata, *columns) metadata.create_all(engine) diff --git a/src/gpu_tracker/sub_tracker.py b/src/gpu_tracker/sub_tracker.py index 9deb6b2..5709095 100644 --- a/src/gpu_tracker/sub_tracker.py +++ b/src/gpu_tracker/sub_tracker.py @@ -3,7 +3,7 @@ import os import time import functools -from ._helper_classes import _TrackingFile, _SubTrackerLog +from ._helper_classes import _Writer, _SubTrackerLog class SubTracker: @@ -21,7 +21,7 @@ def __init__( """ :param code_block_name: The name of the code block within a ``Tracker`` context that is being sub-tracked. Defaults to the file path followed by a colon followed by the ``code_block_attribute``. :param code_block_attribute: Only used if ``code_block_name`` is ``None``. Defaults to the line number where the SubTracker context is started. - :param sub_tracking_file: The path to the file to log the time stamps of the code block being sub-tracked Defaults to the ID of the process where the SubTracker context is created and in CSV format. + :param sub_tracking_file: The path to the file to log the time stamps of the code block being sub-tracked. Defaults to the ID of the process where the SubTracker context is created and in CSV format. """ if code_block_name is not None: self.code_block_name = code_block_name @@ -34,7 +34,7 @@ def __init__( if sub_tracking_file is None: sub_tracking_file = f'{os.getpid()}.csv' self.sub_tracking_file = sub_tracking_file - self._sub_tracking_file = _TrackingFile.create(self.sub_tracking_file) + self._sub_tracking_file = _Writer.create(self.sub_tracking_file) def _log(self, code_block_position: _SubTrackerLog.CodeBlockPosition): sub_tracker_log = _SubTrackerLog( @@ -46,7 +46,7 @@ def __enter__(self): return self def __exit__(self, *_): - self._log(_SubTrackerLog.CodeBlockPosition.END) + self._log(_SubTrackerLog.CodeBlockPosition.STOP) def sub_track(code_block_name: str | None = None, code_block_attribute: str | None = None, sub_tracking_file: str | None = None): diff --git a/src/gpu_tracker/tracker.py b/src/gpu_tracker/tracker.py index 2325100..3df55ac 100644 --- a/src/gpu_tracker/tracker.py +++ b/src/gpu_tracker/tracker.py @@ -13,7 +13,7 @@ import pickle as pkl import uuid import pandas as pd -from ._helper_classes import _NvidiaQuerier, _AMDQuerier, _TrackingFile, _TimepointUsage +from ._helper_classes import _NvidiaQuerier, _AMDQuerier, _Writer, _TimepointUsage class _TrackingProcess(mproc.Process): @@ -64,7 +64,7 @@ def __init__( self._is_linux = platform.system().lower() == 'linux' cannot_connect_warning = ('The {} command is installed but cannot connect to a GPU. ' 'The GPU RAM and GPU utilization values will remain 0.0.') - self.tracking_file = _TrackingFile.create(tracking_file) + self.tracking_file = _Writer.create(tracking_file) if gpu_brand is None: nvidia_available = _NvidiaQuerier.is_available() nvidia_installed = nvidia_available is not None @@ -349,7 +349,7 @@ def __init__( :param n_join_attempts: The number of times the tracker attempts to join its underlying sub-process. :param join_timeout: The amount of time the tracker waits for its underlying sub-process to join. :param 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. - :param 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. + :param 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 "data" and allows for more efficient querying. :raises ValueError: Raised if invalid arguments are provided. """ current_process_id = os.getpid() diff --git a/tests/data/decorated-function-other-file.csv b/tests/data/decorated-function-other-file.csv index 8bc4a03..925682e 100644 --- a/tests/data/decorated-function-other-file.csv +++ b/tests/data/decorated-function-other-file.csv @@ -1,3 +1,3 @@ position,timestamp START,12 -END,13 +STOP,13 diff --git a/tests/data/decorated-function.csv b/tests/data/decorated-function.csv index d69d232..d7ae905 100644 --- a/tests/data/decorated-function.csv +++ b/tests/data/decorated-function.csv @@ -1,13 +1,13 @@ position,timestamp START,0 -END,1 +STOP,1 START,2 -END,3 +STOP,3 START,4 -END,5 +STOP,5 START,6 -END,7 +STOP,7 START,8 -END,9 +STOP,9 START,10 -END,11 +STOP,11 diff --git a/tests/data/sub-tracker.csv b/tests/data/sub-tracker.csv index 536400f..23b1c72 100644 --- a/tests/data/sub-tracker.csv +++ b/tests/data/sub-tracker.csv @@ -1,11 +1,11 @@ position,timestamp START,0 -END,1 +STOP,1 START,2 -END,3 +STOP,3 START,4 -END,5 +STOP,5 START,6 -END,7 +STOP,7 START,8 -END,9 +STOP,9 diff --git a/tests/utils.py b/tests/utils.py index e555acf..aae3e33 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -2,7 +2,7 @@ import sqlalchemy as sqlalc import os # noinspection PyProtectedMember -from gpu_tracker._helper_classes import _SQLiteTrackingFile +from gpu_tracker._helper_classes import _SQLiteWriter import gpu_tracker as gput @@ -17,7 +17,7 @@ def test_tracking_file( actual_tracking_log = pd.read_csv(actual_tracking_file) else: engine = sqlalc.create_engine(f'sqlite:///{actual_tracking_file}', poolclass=sqlalc.pool.NullPool) - actual_tracking_log = pd.read_sql_table(_SQLiteTrackingFile._SQLITE_TABLE_NAME, engine) + actual_tracking_log = pd.read_sql_table(_SQLiteWriter._DATA_TABLE, engine) if excluded_col is not None: actual_tracking_log[excluded_col].apply(excluded_col_test) actual_tracking_log = actual_tracking_log[actual_tracking_log.columns.difference([excluded_col])]