feat(performance-metrics): add RobotContextTracker

.github/workflows/performance-metrics-test-lint.yaml

@@ -0,0 +1,54 @@
+# This workflow runs lint on pull requests that touch anything in the performance-metrics directory
+
+name: 'performance-metrics test & lint'
+
+on:
+  pull_request:
+    paths:
+      - 'performance-metrics/**'
+      - '.github/workflows/performance-metrics-test-lint.yaml'
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+defaults:
+  run:
+    shell: bash
+
+jobs:
+  lint:
+    name: 'performance-metrics test & lint'
+    timeout-minutes: 5
+    runs-on: 'ubuntu-latest'
+    steps:
+      - name: Checkout opentrons repo
+        uses: 'actions/checkout@v4'
+
+      - name: Setup Python
+        uses: 'actions/setup-python@v5'
+        with:
+          python-version: '3.10'
+          cache: 'pipenv'
+          cache-dependency-path: performance-metrics/Pipfile.lock
+
+      - name: "Install Python deps"
+        uses: './.github/actions/python/setup'
+        with:
+          project: 'performance-metrics'
+
+      - name: Setup
+        id: install
+        working-directory: ./performance-metrics
+        run: make setup
+
+      - name: Test
+        if: always() && steps.install.outcome == 'success' || steps.install.outcome == 'skipped'
+        working-directory: ./performance-metrics
+        run: make test
+
+      - name: Lint
+        if: always() && steps.install.outcome == 'success' || steps.install.outcome == 'skipped'
+        working-directory: ./performance-metrics
+        run: make lint

performance-metrics/Makefile

@@ -25,4 +25,8 @@ clean:
 .PHONY: wheel
 wheel:
 	$(python) setup.py $(wheel_opts) bdist_wheel
-	rm -rf build
+	rm -rf build
+
+.PHONY: test
+test:
+	$(pytest) tests

performance-metrics/Pipfile

@@ -5,15 +5,17 @@ name        = "pypi"
 
 [packages]
 opentrons-shared-data = {file = "../shared-data/python", editable = true}
+performance-metrics = {file = ".", editable = true}
 
 [dev-packages]
-pytest  = "==7.2.2"
+pytest  = "==7.4.4"
 mypy    = "==1.8.0"
 flake8 = "==7.0.0"
 flake8-annotations = "~=3.0.1"
 flake8-docstrings = "~=1.7.0"
 flake8-noqa = "~=1.4.0"
 black = "==22.3.0"
+pytest-asyncio = "~=0.23.0"
 
 [requires]
 python_version = "3.10"

performance-metrics/src/performance_metrics/datashapes.py

@@ -0,0 +1,67 @@
+"""Defines data classes and enums used in the performance metrics module."""
+
+from enum import Enum
+import dataclasses
+from typing import Tuple
+
+
+class RobotContextState(Enum):
+    """Enum representing different states of a robot's operation context."""
+
+    STARTING_UP = 0, "STARTING_UP"
+    CALIBRATING = 1, "CALIBRATING"
+    ANALYZING_PROTOCOL = 2, "ANALYZING_PROTOCOL"
+    RUNNING_PROTOCOL = 3, "RUNNING_PROTOCOL"
+    SHUTTING_DOWN = 4, "SHUTTING_DOWN"
+
+    def __init__(self, state_id: int, state_name: str) -> None:
+        self.state_id = state_id
+        self.state_name = state_name
+
+    @classmethod
+    def from_id(cls, state_id: int) -> "RobotContextState":
+        """Returns the enum member matching the given state ID.
+
+        Args:
+            state_id: The ID of the state to retrieve.
+
+        Returns:
+            RobotContextStates: The enum member corresponding to the given ID.
+
+        Raises:
+            ValueError: If no matching state is found.
+        """
+        for state in RobotContextState:
+            if state.state_id == state_id:
+                return state
+        raise ValueError(f"Invalid state id: {state_id}")
+
+
+@dataclasses.dataclass(frozen=True)
+class RawContextData:
+    """Represents raw duration data with context state information.
+
+    Attributes:
+    - function_start_time (int): The start time of the function.
+    - duration_measurement_start_time (int): The start time for duration measurement.
+    - duration_measurement_end_time (int): The end time for duration measurement.
+    - state (RobotContextStates): The current state of the context.
+    """
+
+    func_start: int
+    duration_start: int
+    duration_end: int
+    state: RobotContextState
+
+    @classmethod
+    def headers(self) -> Tuple[str, str, str]:
+        """Returns the headers for the raw context data."""
+        return ("state_id", "function_start_time", "duration")
+
+    def csv_row(self) -> Tuple[int, int, int]:
+        """Returns the raw context data as a string."""
+        return (
+            self.state.state_id,
+            self.func_start,
+            self.duration_end - self.duration_start,
+        )

performance-metrics/src/performance_metrics/robot_context_tracker.py

@@ -0,0 +1,75 @@
+"""Module for tracking robot context and execution duration for different operations."""
+
+import csv
+from pathlib import Path
+import os
+
+from functools import wraps
+from time import perf_counter_ns, clock_gettime_ns, CLOCK_REALTIME
+from typing import Callable, TypeVar
+from typing_extensions import ParamSpec
+from collections import deque
+from performance_metrics.datashapes import (
+    RawContextData,
+    RobotContextState,
+)
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class RobotContextTracker:
+    """Tracks and stores robot context and execution duration for different operations."""
+
+    def __init__(self, storage_file_path: Path, should_track: bool = False) -> None:
+        """Initializes the RobotContextTracker with an empty storage list."""
+        self._storage: deque[RawContextData] = deque()
+        self._storage_file_path = storage_file_path
+        self._should_track = should_track
+
+    def track(self, state: RobotContextState) -> Callable:  # type: ignore
+        """Decorator factory for tracking the execution duration and state of robot operations.
+
+        Args:
+            state: The state to track for the decorated function.
+
+        Returns:
+            Callable: A decorator that wraps a function to track its execution duration and state.
+        """
+
+        def inner_decorator(func: Callable[P, R]) -> Callable[P, R]:
+            if not self._should_track:
+                return func
+
+            @wraps(func)
+            def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+                function_start_time = clock_gettime_ns(CLOCK_REALTIME)
+                duration_start_time = perf_counter_ns()
+                try:
+                    result = func(*args, **kwargs)
+                finally:
+                    duration_end_time = perf_counter_ns()
+                    self._storage.append(
+                        RawContextData(
+                            function_start_time,
+                            duration_start_time,
+                            duration_end_time,
+                            state,
+                        )
+                    )
+                return result
+
+            return wrapper
+
+        return inner_decorator
+
+    def store(self) -> None:
+        """Returns the stored context data and clears the storage list."""
+        stored_data = self._storage.copy()
+        self._storage.clear()
+        rows_to_write = [context_data.csv_row() for context_data in stored_data]
+        os.makedirs(self._storage_file_path.parent, exist_ok=True)
+        with open(self._storage_file_path, "a") as storage_file:
+            writer = csv.writer(storage_file)
+            writer.writerow(RawContextData.headers())
+            writer.writerows(rows_to_write)