Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/gpu_tracker/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +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.
--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 "data" and allows for more efficient querying.
"""
import docopt as doc
import subprocess as subp
Expand Down
23 changes: 12 additions & 11 deletions src/gpu_tracker/_helper_classes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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".')
Expand All @@ -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)
Expand All @@ -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())
Expand All @@ -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)
Expand All @@ -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)
8 changes: 4 additions & 4 deletions src/gpu_tracker/sub_tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import os
import time
import functools
from ._helper_classes import _TrackingFile, _SubTrackerLog
from ._helper_classes import _Writer, _SubTrackerLog


class SubTracker:
Expand All @@ -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
Expand All @@ -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(
Expand All @@ -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):
Expand Down
6 changes: 3 additions & 3 deletions src/gpu_tracker/tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion tests/data/decorated-function-other-file.csv
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
position,timestamp
START,12
END,13
STOP,13
12 changes: 6 additions & 6 deletions tests/data/decorated-function.csv
Original file line number Diff line number Diff line change
@@ -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
10 changes: 5 additions & 5 deletions tests/data/sub-tracker.csv
Original file line number Diff line number Diff line change
@@ -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
4 changes: 2 additions & 2 deletions tests/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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])]
Expand Down
Loading