Merge pull request #61 from MoseleyBioinformaticsLab/granular

Introduces the SubTracker

Author: erikhuck

src/gpu_tracker/__init__.py

@@ -10,3 +10,4 @@
     __version__ = _gv(_path.join(_path.dirname(__file__), _path.pardir))
 
 from .tracker import Tracker
+from .sub_tracker import SubTracker

src/gpu_tracker/_helper_classes.py

@@ -8,6 +8,7 @@
 import dataclasses as dclass
 import sqlalchemy as sqlalc
 import sqlalchemy.orm as sqlorm
+import enum
 
 
 class _GPUQuerier(abc.ABC):
@@ -111,7 +112,7 @@ def ram_and_utilization(cls) -> pd.DataFrame:
 
 
 @dclass.dataclass
-class TimepointUsage:
+class _TimepointUsage:
     main_ram: float = 0.0
     descendants_ram: float = 0.0
     combined_ram: float = 0.0
@@ -136,6 +137,16 @@ class TimepointUsage:
     timestamp: float = 0.0
 
 
+@dclass.dataclass
+class _SubTrackerLog:
+    class CodeBlockPosition(enum.Enum):
+        START = 'START'
+        END = 'END'
+    code_block_name: str
+    position: CodeBlockPosition
+    timestamp: float
+
+
 class _TrackingFile(abc.ABC):
     @staticmethod
     def create(file: str | None) -> _TrackingFile | None:
@@ -153,7 +164,7 @@ def create(file: str | None) -> _TrackingFile | None:
     def __init__(self, file: str):
         self._file = file
 
-    def write_row(self, values: TimepointUsage):
+    def write_row(self, values: _TimepointUsage | _SubTrackerLog):
         values = dclass.asdict(values)
         if not os.path.isfile(self._file):
             self._create_file(values)
@@ -206,5 +217,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))
-        tracking_table = sqlalc.Table(_SQLiteTrackingFile._SQLITE_TABLE_NAME, metadata, *columns)
+        sqlalc.Table(_SQLiteTrackingFile._SQLITE_TABLE_NAME, metadata, *columns)
         metadata.create_all(engine)

src/gpu_tracker/sub_tracker.py

@@ -0,0 +1,46 @@
+"""The ``sub_tracker`` module contains the ``SubTracker`` class which can alternatively be imported directly from the ``gpu_tracker`` package."""
+import inspect
+import os
+import time
+from ._helper_classes import _TrackingFile, _SubTrackerLog
+
+
+class SubTracker:
+    """
+    Context manager that logs to a file for the purposes of sub tracking a code block using the timestamps at which the codeblock begins and ends.
+    Entering the context manager marks the beginning of the code block and exiting the context manager marks the end of the code block.
+    At the beginning of the codeblock, the ``SubTracker`` logs a row to a tablular file (".csv" or ".sqlite") that includes the timestamp along with a name for the code block and an indication of whether it is the start or end of the code bock.
+    This resulting file can be used alongside a tracking file created by a ``Tracker`` object for more granular analysis of specific code blocks.
+
+    :ivar str code_block_name: The name of the code block being sub-tracked.
+    :ivar str sub_tracking_file: The path to the file where the sub-tracking info is logged.
+    """
+    def __init__(self, code_block_name: str | None = None, sub_tracking_file: str | None = None):
+        """
+        :param code_block_name: The name of the code block within a ``Tracker`` context that is being sub-tracked. Defaults to the file path and 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.
+        """
+        if code_block_name is not None:
+            self.code_block_name = code_block_name
+        else:
+            stack = inspect.stack()
+            caller_frame = stack[1]
+            file_path = os.path.abspath(caller_frame.filename)
+            line_number = caller_frame.lineno
+            self.code_block_name = f'{file_path}:{line_number}'
+        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)
+
+    def _log(self, code_block_position: _SubTrackerLog.CodeBlockPosition):
+        sub_tracker_log = _SubTrackerLog(
+            code_block_name=self.code_block_name, position=code_block_position.value, timestamp=time.time())
+        self._sub_tracking_file.write_row(sub_tracker_log)
+
+    def __enter__(self):
+        self._log(_SubTrackerLog.CodeBlockPosition.START)
+        return self
+
+    def __exit__(self, *_):
+        self._log(_SubTrackerLog.CodeBlockPosition.END)

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, _TrackingFile, _TimepointUsage
 
 
 class _TrackingProcess(mproc.Process):
@@ -140,7 +140,7 @@ def run(self):
             self._stop_event.set()
         # Simulate a do-while loop so that the tracking is executed at least once.
         while True:
-            timepoint_usage = TimepointUsage()
+            timepoint_usage = _TimepointUsage()
             with open(self._resource_usage_file, 'wb') as file:
                 pkl.dump(self._resource_usage, file)
             if self._stop_event.is_set():

tests/data/None_None.csv

@@ -0,0 +1,11 @@
+position,timestamp
+START,0
+END,1
+START,2
+END,3
+START,4
+END,5
+START,6
+END,7
+START,8
+END,9

tests/data/None_sub-tracking-file.csv.csv

@@ -0,0 +1,11 @@
+position,timestamp
+START,0
+END,1
+START,2
+END,3
+START,4
+END,5
+START,6
+END,7
+START,8
+END,9

tests/data/None_sub-tracking-file.sqlite.csv

@@ -0,0 +1,11 @@
+position,timestamp
+START,0
+END,1
+START,2
+END,3
+START,4
+END,5
+START,6
+END,7
+START,8
+END,9

tests/data/my-code-block_None.csv

@@ -0,0 +1,11 @@
+position,timestamp
+START,0
+END,1
+START,2
+END,3
+START,4
+END,5
+START,6
+END,7
+START,8
+END,9

tests/data/my-code-block_sub-tracking-file.csv.csv

@@ -0,0 +1,11 @@
+position,timestamp
+START,0
+END,1
+START,2
+END,3
+START,4
+END,5
+START,6
+END,7
+START,8
+END,9

tests/data/my-code-block_sub-tracking-file.sqlite.csv

@@ -0,0 +1,11 @@
+position,timestamp
+START,0
+END,1
+START,2
+END,3
+START,4
+END,5
+START,6
+END,7
+START,8
+END,9