chore(api, performance-metrics): clean up performance-metrics tracking

api/src/opentrons/cli/analyze.py

@@ -44,7 +44,6 @@
 )
 
 from opentrons_shared_data.robot.dev_types import RobotType
-from opentrons.util.performance_helpers import track_analysis
 
 OutputKind = Literal["json", "human-json"]
 
@@ -198,7 +197,6 @@ def _get_return_code(analysis: RunResult) -> int:
  return 0
 
 
-@track_analysis
 async def _do_analyze(protocol_source: ProtocolSource) -> RunResult:
 
  runner = await create_simulating_runner(

api/src/opentrons/util/performance_helpers.py

@@ -1,17 +1,20 @@
 """Performance helpers for tracking robot context."""
 
+import functools
 from pathlib import Path
 from opentrons_shared_data.performance.dev_types import (
  SupportsTracking,
- F,
+ UnderlyingFunction,
+ UnderlyingFunctionParameters,
+ UnderlyingFunctionReturn,
  RobotContextState,
 )
 from opentrons_shared_data.robot.dev_types import RobotTypeEnum
-from typing import Callable, Type
+import typing
 from opentrons.config import (
- feature_flags as ff,
  get_performance_metrics_data_dir,
  robot_configs,
+ feature_flags as ff,
 )
 
 
@@ -27,11 +30,14 @@ def __init__(self, storage_location: Path, should_track: bool) -> None:
  """Initialize the stubbed tracker."""
  pass
 
- def track(self, state: RobotContextState) -> Callable[[F], F]:
- """Return the function unchanged."""
+ def track(
+ self,
+ state: "RobotContextState",
+ ) -> typing.Callable[[UnderlyingFunction], UnderlyingFunction]:
+ """Return the original function."""
 
- def inner_decorator(func: F) -> F:
- """Return the function unchanged."""
+ def inner_decorator(func: UnderlyingFunction) -> UnderlyingFunction:
+ """Return the original function."""
  return func
 
  return inner_decorator
@@ -41,7 +47,7 @@ def store(self) -> None:
  pass
 
 
-def _handle_package_import() -> Type[SupportsTracking]:
+def _handle_package_import() -> typing.Type[SupportsTracking]:
  """Handle the import of the performance_metrics package.
 
  If the package is not available, return a stubbed tracker.
@@ -58,19 +64,37 @@ def _handle_package_import() -> Type[SupportsTracking]:
 _robot_context_tracker: SupportsTracking | None = None
 
 
+# TODO: derek maggio (06-03-2024): investigate if _should_track should be
+# reevaluated each time _get_robot_context_tracker is called. I think this
+# might get stuck in a state where after the first call, _should_track is
+# always considered the initial value. It might miss changes to the feature
+# flag. The easiest way to test this is on a robot when that is working.
+
+
 def _get_robot_context_tracker() -> SupportsTracking:
  """Singleton for the robot context tracker."""
  global _robot_context_tracker
  if _robot_context_tracker is None:
- # TODO: replace with path lookup and should_store lookup
  _robot_context_tracker = package_to_use(
  get_performance_metrics_data_dir(), _should_track
  )
  return _robot_context_tracker
 
 
-def track_analysis(func: F) -> F:
- """Track the analysis of a protocol."""
- return _get_robot_context_tracker().track(RobotContextState.ANALYZING_PROTOCOL)(
- func
- )
+def track_analysis(func: UnderlyingFunction) -> UnderlyingFunction:
+ """Track the analysis of a protocol and store each run."""
+ # TODO: derek maggio (06-03-2024): generalize creating wrapper functions for tracking different states
+ tracker: SupportsTracking = _get_robot_context_tracker()
+ wrapped = tracker.track(state=RobotContextState.ANALYZING_PROTOCOL)(func)
+
+ @functools.wraps(func)
+ def wrapper(
+ *args: UnderlyingFunctionParameters.args,
+ **kwargs: UnderlyingFunctionParameters.kwargs
+ ) -> UnderlyingFunctionReturn:
+ try:
+ return wrapped(*args, **kwargs)
+ finally:
+ tracker.store()
+
+ return wrapper

api/tests/opentrons/cli/test_cli.py

@@ -5,38 +5,13 @@
 import tempfile
 import textwrap
 
-from dataclasses import dataclass, replace
+from dataclasses import dataclass
 from typing import Any, Dict, Iterator, List, Optional
 from pathlib import Path
 
 import pytest
 from click.testing import CliRunner
-from opentrons_shared_data.performance.dev_types import (
- RobotContextState,
-)
-from opentrons.util.performance_helpers import _get_robot_context_tracker
-
-
-# Enable tracking for the RobotContextTracker
-# This must come before the import of the analyze CLI
-context_tracker = _get_robot_context_tracker()
-
-# Ignore the type error for the next line, as we're setting a private attribute for testing purposes
-context_tracker._should_track = True # type: ignore[attr-defined]
-
-from opentrons.cli.analyze import analyze # noqa: E402
-
-
-@pytest.fixture
-def override_data_store(tmp_path: Path) -> Iterator[None]:
- """Override the data store metadata for the RobotContextTracker."""
- old_store = context_tracker._store # type: ignore[attr-defined]
- old_metadata = old_store.metadata
- new_metadata = replace(old_metadata, storage_dir=tmp_path)
- context_tracker._store = old_store.__class__(metadata=new_metadata) # type: ignore[attr-defined]
- context_tracker._store.setup() # type: ignore[attr-defined]
- yield
- context_tracker._store = old_store # type: ignore[attr-defined]
+from opentrons.cli.analyze import analyze
 
 
 def _list_fixtures(version: int) -> Iterator[Path]:
@@ -285,41 +260,3 @@ def test_python_error_line_numbers(
  assert result.json_output is not None
  [error] = result.json_output["errors"]
  assert error["detail"] == expected_detail
-
-
-@pytest.mark.usefixtures("override_data_store")
-@pytest.mark.parametrize("output", ["--json-output", "--human-json-output"])
-def test_track_analysis(tmp_path: Path, output: str) -> None:
- """Test that the RobotContextTracker tracks analysis."""
- protocol_source = textwrap.dedent(
- """
- requirements = {"apiLevel": "2.15"}
-
- def run(protocol):
- pass
- """
- )
- protocol_source_file = tmp_path / "protocol.py"
- protocol_source_file.write_text(protocol_source, encoding="utf-8")
- store = context_tracker._store # type: ignore[attr-defined]
-
- num_storage_entities_before_analysis = len(store._data)
-
- _get_analysis_result([protocol_source_file], output)
-
- assert len(store._data) == num_storage_entities_before_analysis + 1
-
- with open(store.metadata.data_file_location, "r") as f:
- stored_data = f.readlines()
- assert len(stored_data) == 0
-
- context_tracker.store()
-
- with open(store.metadata.data_file_location, "r") as f:
- stored_data = f.readlines()
- stored_data = [line.strip() for line in stored_data if line.strip()]
- assert len(stored_data) == 1
- state_id, start_time, duration = stored_data[0].strip().split(",")
- assert state_id == str(RobotContextState.ANALYZING_PROTOCOL.state_id)
- assert start_time.isdigit()
- assert duration.isdigit()

performance-metrics/src/performance_metrics/robot_context_tracker.py

@@ -1,14 +1,13 @@
 """Module for tracking robot context and execution duration for different operations."""
 
+import inspect
 from pathlib import Path
 import platform
 
-from functools import wraps, partial
+from functools import partial, wraps
 from time import perf_counter_ns
-from typing import Callable, TypeVar, cast, Literal, Final
+import typing
 
-
-from typing_extensions import ParamSpec
 from performance_metrics.datashapes import (
  RawContextData,
 )
@@ -17,20 +16,20 @@
  RobotContextState,
  SupportsTracking,
  MetricsMetadata,
+ UnderlyingFunction,
+ UnderlyingFunctionReturn,
+ UnderlyingFunctionParameters,
 )
 
-P = ParamSpec("P")
-R = TypeVar("R")
-
 
-def _get_timing_function() -> Callable[[], int]:
+def _get_timing_function() -> typing.Callable[[], int]:
  """Returns a timing function for the current platform."""
- time_function: Callable[[], int]
+ time_function: typing.Callable[[], int]
  if platform.system() == "Linux":
  from time import clock_gettime_ns, CLOCK_REALTIME
 
- time_function = cast(
- Callable[[], int], partial(clock_gettime_ns, CLOCK_REALTIME)
+ time_function = typing.cast(
+ typing.Callable[[], int], partial(clock_gettime_ns, CLOCK_REALTIME)
  )
  else:
  from time import time_ns
@@ -46,9 +45,11 @@ def _get_timing_function() -> Callable[[], int]:
 class RobotContextTracker(SupportsTracking):
  """Tracks and stores robot context and execution duration for different operations."""
 
- METADATA_NAME: Final[Literal["robot_context_data"]] = "robot_context_data"
+ METADATA_NAME: typing.Final[
+ typing.Literal["robot_context_data"]
+ ] = "robot_context_data"
 
- def __init__(self, storage_location: Path, should_track: bool = False) -> None:
+ def __init__(self, storage_location: Path, should_track: bool) -> None:
  """Initializes the RobotContextTracker with an empty storage list."""
  self._store = MetricsStore[RawContextData](
  MetricsMetadata(
@@ -62,39 +63,86 @@ def __init__(self, storage_location: Path, should_track: bool = False) -> None:
  if self._should_track:
  self._store.setup()
 
- def track(self, state: RobotContextState) -> Callable: # type: ignore
- """Decorator factory for tracking the execution duration and state of robot operations.
+ def track(
+ self,
+ state: RobotContextState,
+ ) -> typing.Callable[[UnderlyingFunction], UnderlyingFunction]:
+ """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.
  """
 
- def inner_decorator(func: Callable[P, R]) -> Callable[P, R]:
+ def inner_decorator(func_to_track: UnderlyingFunction) -> UnderlyingFunction:
  if not self._should_track:
- return func
-
- @wraps(func)
- def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
- function_start_time = timing_function()
- duration_start_time = perf_counter_ns()
- try:
- result = func(*args, **kwargs)
- finally:
- duration_end_time = perf_counter_ns()
- self._store.add(
- RawContextData(
- func_start=function_start_time,
- duration=duration_end_time - duration_start_time,
- state=state,
+ return func_to_track
+
+ if inspect.iscoroutinefunction(func_to_track):
+
+ @wraps(func_to_track)
+ async def async_wrapper(
+ *args: UnderlyingFunctionParameters.args,
+ **kwargs: UnderlyingFunctionParameters.kwargs
+ ) -> UnderlyingFunctionReturn:
+ function_start_time = timing_function()
+ duration_start_time = perf_counter_ns()
+
+ try:
+ result = typing.cast(
+ UnderlyingFunctionReturn,
+ await func_to_track(*args, **kwargs),
+ )
+ finally:
+ duration_end_time = perf_counter_ns()
+
+ self._store.add(
+ RawContextData(
+ func_start=function_start_time,
+ duration=duration_end_time - duration_start_time,
+ state=state,
+ )
+ )
+
+ return result
+
+ return async_wrapper
+ else:
+
+ @wraps(func_to_track)
+ def wrapper(
+ *args: UnderlyingFunctionParameters.args,
+ **kwargs: UnderlyingFunctionParameters.kwargs
+ ) -> UnderlyingFunctionReturn:
+ function_start_time = timing_function()
+ duration_start_time = perf_counter_ns()
+
+ try:
+ result = typing.cast(
+ UnderlyingFunctionReturn, func_to_track(*args, **kwargs)
+ )
+ finally:
+ duration_end_time = perf_counter_ns()
+
+ self._store.add(
+ RawContextData(
+ func_start=function_start_time,
+ duration=duration_end_time - duration_start_time,
+ state=state,
+ )
  )
- )
 
- return result
+  return result
 
- return wrapper
+  return wrapper
 
  return inner_decorator
 

performance-metrics/tests/performance_metrics/test_metrics_store.py

@@ -19,20 +19,20 @@ async def test_storing_to_file(tmp_path: Path) -> None:
  """Tests storing the tracked data to a file."""
  robot_context_tracker = RobotContextTracker(tmp_path, should_track=True)
 
- @robot_context_tracker.track(state=RobotContextState.STARTING_UP)
+ @robot_context_tracker.track(RobotContextState.STARTING_UP)
  def starting_robot() -> None:
  sleep(STARTING_TIME)
 
- @robot_context_tracker.track(state=RobotContextState.CALIBRATING)
- def calibrating_robot() -> None:
+ @robot_context_tracker.track(RobotContextState.CALIBRATING)
+ async def calibrating_robot() -> None:
  sleep(CALIBRATING_TIME)
 
- @robot_context_tracker.track(state=RobotContextState.ANALYZING_PROTOCOL)
+ @robot_context_tracker.track(RobotContextState.ANALYZING_PROTOCOL)
  def analyzing_protocol() -> None:
  sleep(ANALYZING_TIME)
 
  starting_robot()
- calibrating_robot()
+ await calibrating_robot()
  analyzing_protocol()
 
  robot_context_tracker.store()