From b0944682d650c595e89cbf98c168a57ab7385166 Mon Sep 17 00:00:00 2001 From: Derek Maggio Date: Thu, 30 May 2024 11:01:36 -0700 Subject: [PATCH] chore: way clearer typing, docs, and logic --- .../robot_context_tracker.py | 74 +++++++----- .../performance_metrics/test_metrics_store.py | 12 +- .../test_robot_context_tracker.py | 113 ++++++++++++++---- .../performance/dev_types.py | 25 ++-- 4 files changed, 162 insertions(+), 62 deletions(-) diff --git a/performance-metrics/src/performance_metrics/robot_context_tracker.py b/performance-metrics/src/performance_metrics/robot_context_tracker.py index a9956433b1f..0e7ae5ef358 100644 --- a/performance-metrics/src/performance_metrics/robot_context_tracker.py +++ b/performance-metrics/src/performance_metrics/robot_context_tracker.py @@ -4,12 +4,11 @@ from pathlib import Path import platform -from functools import wraps, partial +from functools import partial from time import perf_counter_ns -from typing import Callable, TypeVar, cast, Literal, Final +from typing import Callable, cast, Literal, Final -from typing_extensions import ParamSpec from performance_metrics.datashapes import ( RawContextData, ) @@ -18,11 +17,11 @@ RobotContextState, SupportsTracking, MetricsMetadata, + UnderlyingFunction, + UnderlyingFunctionParams, + UnderlyingFunctionReturn, ) -P = ParamSpec("P") -R = TypeVar("R") - def _get_timing_function() -> Callable[[], int]: """Returns a timing function for the current platform.""" @@ -63,39 +62,54 @@ def __init__(self, storage_location: Path, should_track: bool) -> None: if self._should_track: self._store.setup() - async def track(self, func_to_track: Callable, state: RobotContextState, *args, **kwargs) -> Callable: # type: ignore - """Decorator factory for tracking the execution duration and state of robot operations. + async def __call_function( + self, + func_to_track: UnderlyingFunction, + *args: UnderlyingFunctionParams.args, + **kwargs: UnderlyingFunctionParams.kwargs + ) -> UnderlyingFunctionReturn: + """Calls the given function and returns its result.""" + if inspect.iscoroutinefunction(func_to_track): + return await func_to_track(*args, **kwargs) # type: ignore + else: + return func_to_track(*args, **kwargs) # type: ignore + + async def track( + self, + func_to_track: UnderlyingFunction, + state: RobotContextState, + *args: UnderlyingFunctionParams.args, + **kwargs: UnderlyingFunctionParams.kwargs + ) -> UnderlyingFunctionReturn: + """Tracks the given function and its execution duration. + + If tracking is disabled, the function is called immediately and its result is returned. Args: - state: The state to track for the decorated function. + func_to_track: The function to track. + state: The state of the robot context during the function execution. + *args: The arguments to pass to the function. + **kwargs: The keyword arguments to pass to the function. Returns: - Callable: A decorator that wraps a function to track its execution duration and state. + If the function executes successfully, its return value is returned. + If the function raises an exception, the exception the function raised is raised. """ - - - if not self._should_track: - return func_to_track + return await self.__call_function(func_to_track, *args, **kwargs) function_start_time = timing_function() duration_start_time = perf_counter_ns() - if inspect.iscoroutinefunction(func_to_track): - try: - result = await func_to_track(*args, **kwargs) - except Exception as e: - result = e - finally: - duration_end_time = perf_counter_ns() - else: - try: - result = func_to_track(*args, **kwargs) - except Exception as e: - result = e - finally: - duration_end_time = perf_counter_ns() - + result: UnderlyingFunctionReturn | Exception + + try: + result = await self.__call_function(func_to_track, *args, **kwargs) + except Exception as e: + result = e + + duration_end_time = perf_counter_ns() + self._store.add( RawContextData( func_start=function_start_time, @@ -109,8 +123,6 @@ async def track(self, func_to_track: Callable, state: RobotContextState, *args, else: return result - - def store(self) -> None: """Returns the stored context data and clears the storage list.""" if not self._should_track: diff --git a/performance-metrics/tests/performance_metrics/test_metrics_store.py b/performance-metrics/tests/performance_metrics/test_metrics_store.py index 2b39b65de83..aede5ef4cb1 100644 --- a/performance-metrics/tests/performance_metrics/test_metrics_store.py +++ b/performance-metrics/tests/performance_metrics/test_metrics_store.py @@ -28,9 +28,15 @@ def calibrating_robot() -> None: def analyzing_protocol() -> None: sleep(ANALYZING_TIME) - await robot_context_tracker.track(starting_robot, state=RobotContextState.STARTING_UP) - await robot_context_tracker.track(calibrating_robot, state=RobotContextState.CALIBRATING) - await robot_context_tracker.track(analyzing_protocol, state=RobotContextState.ANALYZING_PROTOCOL) + await robot_context_tracker.track( + starting_robot, state=RobotContextState.STARTING_UP + ) + await robot_context_tracker.track( + calibrating_robot, state=RobotContextState.CALIBRATING + ) + await robot_context_tracker.track( + analyzing_protocol, state=RobotContextState.ANALYZING_PROTOCOL + ) robot_context_tracker.store() diff --git a/performance-metrics/tests/performance_metrics/test_robot_context_tracker.py b/performance-metrics/tests/performance_metrics/test_robot_context_tracker.py index e31d864fcce..78a6295fb4a 100644 --- a/performance-metrics/tests/performance_metrics/test_robot_context_tracker.py +++ b/performance-metrics/tests/performance_metrics/test_robot_context_tracker.py @@ -22,10 +22,11 @@ def robot_context_tracker(tmp_path: Path) -> RobotContextTracker: return RobotContextTracker(storage_location=tmp_path, should_track=True) -async def test_robot_context_tracker(robot_context_tracker: RobotContextTracker) -> None: +async def test_robot_context_tracker( + robot_context_tracker: RobotContextTracker, +) -> None: """Tests the tracking of various robot context states through RobotContextTracker.""" - def starting_robot() -> None: sleep(STARTING_TIME) @@ -46,12 +47,22 @@ def shutting_down_robot() -> None: len(robot_context_tracker._store._data) == 0 ), "Storage should be initially empty." - await robot_context_tracker.track(starting_robot, state=RobotContextState.STARTING_UP) - await robot_context_tracker.track(calibrating_robot, state=RobotContextState.CALIBRATING) - await robot_context_tracker.track(analyzing_protocol, state=RobotContextState.ANALYZING_PROTOCOL) - await robot_context_tracker.track(running_protocol, state=RobotContextState.RUNNING_PROTOCOL) - await robot_context_tracker.track(shutting_down_robot, state=RobotContextState.SHUTTING_DOWN) - + await robot_context_tracker.track( + starting_robot, state=RobotContextState.STARTING_UP + ) + await robot_context_tracker.track( + calibrating_robot, state=RobotContextState.CALIBRATING + ) + await robot_context_tracker.track( + analyzing_protocol, state=RobotContextState.ANALYZING_PROTOCOL + ) + await robot_context_tracker.track( + running_protocol, state=RobotContextState.RUNNING_PROTOCOL + ) + await robot_context_tracker.track( + shutting_down_robot, state=RobotContextState.SHUTTING_DOWN + ) + # Verify that all states were tracked assert len(robot_context_tracker._store._data) == 5, "All states should be tracked." @@ -83,8 +94,12 @@ def first_operation() -> None: def second_operation() -> None: sleep(RUNNING_TIME) - await robot_context_tracker.track(first_operation, state=RobotContextState.RUNNING_PROTOCOL) - await robot_context_tracker.track(second_operation, state=RobotContextState.RUNNING_PROTOCOL) + await robot_context_tracker.track( + first_operation, state=RobotContextState.RUNNING_PROTOCOL + ) + await robot_context_tracker.track( + second_operation, state=RobotContextState.RUNNING_PROTOCOL + ) assert ( len(robot_context_tracker._store._data) == 2 @@ -106,7 +121,9 @@ async def error_prone_operation() -> None: raise RuntimeError("Simulated operation failure") with pytest.raises(RuntimeError): - await robot_context_tracker.track(error_prone_operation, state=RobotContextState.SHUTTING_DOWN) + await robot_context_tracker.track( + error_prone_operation, state=RobotContextState.SHUTTING_DOWN + ) assert ( len(robot_context_tracker._store._data) == 1 @@ -125,7 +142,9 @@ async def test_async_operation_tracking( async def async_analyzing_operation() -> None: await asyncio.sleep(ANALYZING_TIME) - await robot_context_tracker.track(async_analyzing_operation, state=RobotContextState.ANALYZING_PROTOCOL) + await robot_context_tracker.track( + async_analyzing_operation, state=RobotContextState.ANALYZING_PROTOCOL + ) assert ( len(robot_context_tracker._store._data) == 1 @@ -144,7 +163,11 @@ def test_sync_operation_timing_accuracy( def running_operation() -> None: sleep(RUNNING_TIME) - asyncio.run(robot_context_tracker.track(running_operation, state=RobotContextState.RUNNING_PROTOCOL)) + asyncio.run( + robot_context_tracker.track( + running_operation, state=RobotContextState.RUNNING_PROTOCOL + ) + ) duration_data = robot_context_tracker._store._data[0] assert ( @@ -161,7 +184,9 @@ async def test_async_operation_timing_accuracy( async def async_running_operation() -> None: await asyncio.sleep(RUNNING_TIME) - await robot_context_tracker.track(async_running_operation, state=RobotContextState.RUNNING_PROTOCOL) + await robot_context_tracker.track( + async_running_operation, state=RobotContextState.RUNNING_PROTOCOL + ) duration_data = robot_context_tracker._store._data[0] assert ( @@ -180,7 +205,9 @@ async def async_error_prone_operation() -> None: raise RuntimeError("Simulated async operation failure") with pytest.raises(RuntimeError): - await robot_context_tracker.track(async_error_prone_operation, state=RobotContextState.SHUTTING_DOWN) + await robot_context_tracker.track( + async_error_prone_operation, state=RobotContextState.SHUTTING_DOWN + ) assert ( len(robot_context_tracker._store._data) == 1 @@ -203,8 +230,12 @@ async def second_async_calibrating() -> None: await asyncio.sleep(CALIBRATING_TIME) await asyncio.gather( - robot_context_tracker.track(first_async_calibrating, state=RobotContextState.CALIBRATING), - robot_context_tracker.track(second_async_calibrating, state=RobotContextState.CALIBRATING) + robot_context_tracker.track( + first_async_calibrating, state=RobotContextState.CALIBRATING + ), + robot_context_tracker.track( + second_async_calibrating, state=RobotContextState.CALIBRATING + ), ) assert ( @@ -223,13 +254,49 @@ def test_no_tracking(tmp_path: Path) -> None: def operation_without_tracking() -> None: sleep(STARTING_TIME) - asyncio.run(robot_context_tracker.track(operation_without_tracking, state=RobotContextState.STARTING_UP)) + asyncio.run( + robot_context_tracker.track( + operation_without_tracking, state=RobotContextState.STARTING_UP + ) + ) assert ( len(robot_context_tracker._store._data) == 0 ), "Operation should not be tracked when tracking is disabled." +def test_async_exception_handling_when_not_tracking(tmp_path: Path) -> None: + """Ensures exceptions in operations are still raised when tracking is disabled.""" + robot_context_tracker = RobotContextTracker(tmp_path, should_track=False) + + async def error_prone_operation() -> None: + sleep(SHUTTING_DOWN_TIME) + raise RuntimeError("Simulated operation failure") + + with pytest.raises(RuntimeError): + asyncio.run( + robot_context_tracker.track( + error_prone_operation, state=RobotContextState.SHUTTING_DOWN + ) + ) + + +def test_sync_exception_handling_when_not_tracking(tmp_path: Path) -> None: + """Ensures exceptions in operations are still raised when tracking is disabled.""" + robot_context_tracker = RobotContextTracker(tmp_path, should_track=False) + + def error_prone_operation() -> None: + sleep(SHUTTING_DOWN_TIME) + raise RuntimeError("Simulated operation failure") + + with pytest.raises(RuntimeError): + asyncio.run( + robot_context_tracker.track( + error_prone_operation, state=RobotContextState.SHUTTING_DOWN + ) + ) + + @patch( "performance_metrics.robot_context_tracker._get_timing_function", return_value=time_ns, @@ -245,8 +312,14 @@ def starting_robot() -> None: def calibrating_robot() -> None: sleep(CALIBRATING_TIME) - asyncio.run(robot_context_tracker.track(starting_robot, state=RobotContextState.STARTING_UP)) - asyncio.run(robot_context_tracker.track(calibrating_robot, state=RobotContextState.CALIBRATING)) + asyncio.run( + robot_context_tracker.track(starting_robot, state=RobotContextState.STARTING_UP) + ) + asyncio.run( + robot_context_tracker.track( + calibrating_robot, state=RobotContextState.CALIBRATING + ) + ) storage = robot_context_tracker._store._data assert all( diff --git a/shared-data/python/opentrons_shared_data/performance/dev_types.py b/shared-data/python/opentrons_shared_data/performance/dev_types.py index 3b6de302e89..e91874e0387 100644 --- a/shared-data/python/opentrons_shared_data/performance/dev_types.py +++ b/shared-data/python/opentrons_shared_data/performance/dev_types.py @@ -1,22 +1,31 @@ """Type definitions for performance tracking.""" +import typing from dataclasses import dataclass -from typing import Protocol, Tuple, TypeVar, Callable, Any from pathlib import Path from enum import Enum -F = TypeVar("F", bound=Callable[..., Any]) +UnderlyingFunctionReturn = typing.TypeVar("UnderlyingFunctionReturn") +UnderlyingFunctionParams = typing.ParamSpec("UnderlyingFunctionParams") +UnderlyingFunction = typing.TypeVar( + "UnderlyingFunction", bound=typing.Callable[..., typing.Any] +) -class SupportsTracking(Protocol): + +class SupportsTracking(typing.Protocol): """Protocol for classes that support tracking of robot context.""" - def __init__( - self, storage_location: Path, should_track: bool - ) -> None: + def __init__(self, storage_location: Path, should_track: bool) -> None: """Initialize the tracker.""" ... - def track(self, state: "RobotContextState") -> Callable[[F], F]: + async def track( + self, + func_to_track: UnderlyingFunction, + state: "RobotContextState", + *args: UnderlyingFunctionParams.args, + **kwargs: UnderlyingFunctionParams.kwargs, + ) -> UnderlyingFunctionReturn: """Decorator to track the given state for the decorated function.""" ... @@ -64,7 +73,7 @@ class MetricsMetadata: name: str storage_dir: Path - headers: Tuple[str, ...] + headers: typing.Tuple[str, ...] @property def data_file_location(self) -> Path: