From a64d4423a5f58ebbc6f290bc6fb64b642bed6a3e Mon Sep 17 00:00:00 2001 From: erikhuck Date: Fri, 4 Apr 2025 12:29:48 -0400 Subject: [PATCH] Enables use of SubTracker as a function decorator --- src/gpu_tracker/__init__.py | 1 + src/gpu_tracker/sub_tracker.py | 42 ++++++++++++++-- tests/data/None_sub-tracking-file.sqlite.csv | 11 ----- tests/data/decorated-function-other-file.csv | 3 ++ ...ng-file.csv.csv => decorated-function.csv} | 2 + tests/data/my-code-block_None.csv | 11 ----- .../my-code-block_sub-tracking-file.csv.csv | 11 ----- ...my-code-block_sub-tracking-file.sqlite.csv | 11 ----- tests/data/{None_None.csv => sub-tracker.csv} | 0 tests/test_sub_tracker.py | 49 ++++++++++++++++++- tests/utils.py | 6 +++ 11 files changed, 96 insertions(+), 51 deletions(-) delete mode 100644 tests/data/None_sub-tracking-file.sqlite.csv create mode 100644 tests/data/decorated-function-other-file.csv rename tests/data/{None_sub-tracking-file.csv.csv => decorated-function.csv} (84%) delete mode 100644 tests/data/my-code-block_None.csv delete mode 100644 tests/data/my-code-block_sub-tracking-file.csv.csv delete mode 100644 tests/data/my-code-block_sub-tracking-file.sqlite.csv rename tests/data/{None_None.csv => sub-tracker.csv} (100%) diff --git a/src/gpu_tracker/__init__.py b/src/gpu_tracker/__init__.py index 7e2b829..3755370 100644 --- a/src/gpu_tracker/__init__.py +++ b/src/gpu_tracker/__init__.py @@ -11,3 +11,4 @@ from .tracker import Tracker from .sub_tracker import SubTracker +from .sub_tracker import sub_track diff --git a/src/gpu_tracker/sub_tracker.py b/src/gpu_tracker/sub_tracker.py index a608f2f..9deb6b2 100644 --- a/src/gpu_tracker/sub_tracker.py +++ b/src/gpu_tracker/sub_tracker.py @@ -2,6 +2,7 @@ import inspect import os import time +import functools from ._helper_classes import _TrackingFile, _SubTrackerLog @@ -9,15 +10,17 @@ 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. + At the beginning of the codeblock, the ``SubTracker`` logs a row to a tabular 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): + def __init__( + self, code_block_name: str | None = None, code_block_attribute: 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 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. """ if code_block_name is not None: @@ -26,8 +29,8 @@ def __init__(self, code_block_name: str | None = None, sub_tracking_file: str | 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}' + code_block_attribute = caller_frame.lineno if code_block_attribute is None else code_block_attribute + self.code_block_name = f'{file_path}:{code_block_attribute}' if sub_tracking_file is None: sub_tracking_file = f'{os.getpid()}.csv' self.sub_tracking_file = sub_tracking_file @@ -44,3 +47,32 @@ def __enter__(self): def __exit__(self, *_): self._log(_SubTrackerLog.CodeBlockPosition.END) + + +def sub_track(code_block_name: str | None = None, code_block_attribute: str | None = None, sub_tracking_file: str | None = None): + """ + Decorator for sub tracking calls to a specified function. + + :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 name of the decorated function i.e. the function being sub-tracked. + :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. + """ + def decorator(func): + nonlocal code_block_name, code_block_attribute, sub_tracking_file + if code_block_name is None: + stack = inspect.stack() + caller_frame = stack[1] + file_path = os.path.abspath(caller_frame.filename) + code_block_attribute = func.__name__ if code_block_attribute is None else code_block_attribute + code_block_name = f'{file_path}:{code_block_attribute}' + + @functools.wraps(func) + def wrapper(*args, **kwargs): + nonlocal sub_tracking_file + with SubTracker( + code_block_name=code_block_name, code_block_attribute=code_block_attribute, sub_tracking_file=sub_tracking_file + ): + return_value = func(*args, **kwargs) + return return_value + return wrapper + return decorator diff --git a/tests/data/None_sub-tracking-file.sqlite.csv b/tests/data/None_sub-tracking-file.sqlite.csv deleted file mode 100644 index 536400f..0000000 --- a/tests/data/None_sub-tracking-file.sqlite.csv +++ /dev/null @@ -1,11 +0,0 @@ -position,timestamp -START,0 -END,1 -START,2 -END,3 -START,4 -END,5 -START,6 -END,7 -START,8 -END,9 diff --git a/tests/data/decorated-function-other-file.csv b/tests/data/decorated-function-other-file.csv new file mode 100644 index 0000000..8bc4a03 --- /dev/null +++ b/tests/data/decorated-function-other-file.csv @@ -0,0 +1,3 @@ +position,timestamp +START,12 +END,13 diff --git a/tests/data/None_sub-tracking-file.csv.csv b/tests/data/decorated-function.csv similarity index 84% rename from tests/data/None_sub-tracking-file.csv.csv rename to tests/data/decorated-function.csv index 536400f..d69d232 100644 --- a/tests/data/None_sub-tracking-file.csv.csv +++ b/tests/data/decorated-function.csv @@ -9,3 +9,5 @@ START,6 END,7 START,8 END,9 +START,10 +END,11 diff --git a/tests/data/my-code-block_None.csv b/tests/data/my-code-block_None.csv deleted file mode 100644 index 536400f..0000000 --- a/tests/data/my-code-block_None.csv +++ /dev/null @@ -1,11 +0,0 @@ -position,timestamp -START,0 -END,1 -START,2 -END,3 -START,4 -END,5 -START,6 -END,7 -START,8 -END,9 diff --git a/tests/data/my-code-block_sub-tracking-file.csv.csv b/tests/data/my-code-block_sub-tracking-file.csv.csv deleted file mode 100644 index 536400f..0000000 --- a/tests/data/my-code-block_sub-tracking-file.csv.csv +++ /dev/null @@ -1,11 +0,0 @@ -position,timestamp -START,0 -END,1 -START,2 -END,3 -START,4 -END,5 -START,6 -END,7 -START,8 -END,9 diff --git a/tests/data/my-code-block_sub-tracking-file.sqlite.csv b/tests/data/my-code-block_sub-tracking-file.sqlite.csv deleted file mode 100644 index 536400f..0000000 --- a/tests/data/my-code-block_sub-tracking-file.sqlite.csv +++ /dev/null @@ -1,11 +0,0 @@ -position,timestamp -START,0 -END,1 -START,2 -END,3 -START,4 -END,5 -START,6 -END,7 -START,8 -END,9 diff --git a/tests/data/None_None.csv b/tests/data/sub-tracker.csv similarity index 100% rename from tests/data/None_None.csv rename to tests/data/sub-tracker.csv diff --git a/tests/test_sub_tracker.py b/tests/test_sub_tracker.py index a94bce3..a3e54fb 100644 --- a/tests/test_sub_tracker.py +++ b/tests/test_sub_tracker.py @@ -34,8 +34,53 @@ def code_block_name_test(val: str): assert val.endswith(default_code_block_end) else: assert val == code_block_name - expected_tracking_file = f'tests/data/{code_block_name}_{sub_tracking_file}.csv' utils.test_tracking_file( - actual_tracking_file=sub_tracker.sub_tracking_file, expected_tracking_file=expected_tracking_file, + actual_tracking_file=sub_tracker.sub_tracking_file, expected_tracking_file='tests/data/sub-tracker.csv', excluded_col='code_block_name', excluded_col_test=code_block_name_test ) + + +@pt.fixture(name='code_block_attribute', params=['my-attribute', None]) +def get_code_block_attribute(request): + yield request.param + + +def test_decorator(mocker, code_block_name: str | None, code_block_attribute: str | None): + @gput.sub_track(code_block_name=code_block_name, code_block_attribute=code_block_attribute) + def decorated_function(arg1: int, arg2: int, kwarg1: int = 1, kwarg2: int = 2) -> int: + return arg1 + arg2 - (kwarg1 + kwarg2) + getpid_mock = mocker.patch('gpu_tracker.sub_tracker.os.getpid', return_value=1234) + n_iterations = 3 + time_mock = mocker.patch('gpu_tracker.sub_tracker.time', time=mocker.MagicMock(side_effect=range(n_iterations * 2 * 2 + 2))) + for _ in range(n_iterations): + return_val = decorated_function(2, 5) + assert return_val == 4 + return_val = decorated_function(3, 2, kwarg1=2, kwarg2=1) + assert return_val == 2 + assert len(getpid_mock.call_args_list) == n_iterations * 2 + assert len(time_mock.time.call_args_list) == n_iterations * 2 * 2 + + def code_block_name_test(val): + if code_block_name is None: + if code_block_attribute is None: + assert val.endswith('test_sub_tracker.py:decorated_function') + else: + assert val.endswith('test_sub_tracker.py:my-attribute') + else: + assert val == code_block_name + utils.test_tracking_file( + actual_tracking_file='1234.csv', expected_tracking_file=f'tests/data/decorated-function.csv', + excluded_col='code_block_name', excluded_col_test=code_block_name_test + ) + if code_block_name is None and code_block_attribute is None: + return_val = utils.function_in_other_file(1, 2, 3, kw1=4, kw2=5) + assert return_val == ((1, 2, 3), {'kw1': 4, 'kw2': 5}) + assert len(getpid_mock.call_args_list) == n_iterations * 2 + 1 + assert len(time_mock.time.call_args_list) == n_iterations * 2 * 2 + 2 + + def code_block_name_test(val): + assert val.endswith('utils.py:function_in_other_file') + utils.test_tracking_file( + actual_tracking_file='1234.csv', expected_tracking_file='tests/data/decorated-function-other-file.csv', + excluded_col='code_block_name', excluded_col_test=code_block_name_test + ) diff --git a/tests/utils.py b/tests/utils.py index 32e3d25..e555acf 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -3,6 +3,7 @@ import os # noinspection PyProtectedMember from gpu_tracker._helper_classes import _SQLiteTrackingFile +import gpu_tracker as gput def assert_args_list(mock, expected_args_list: list[tuple | dict], use_kwargs: bool = False): @@ -23,3 +24,8 @@ def test_tracking_file( expected_tracking_log = pd.read_csv(expected_tracking_file) pd.testing.assert_frame_equal(expected_tracking_log, actual_tracking_log, atol=1e-10, rtol=1e-10) os.remove(actual_tracking_file) + + +@gput.sub_track() +def function_in_other_file(*args, **kwargs): + return args, kwargs