refactor(robot-server): pull sessions/analysis equipment from engine (#8320)

Author: mcous

api/src/opentrons/protocol_runner/__init__.py

@@ -3,12 +3,13 @@
 The main export of this module is the ProtocolRunner class. See
 protocol_runner.py for more details.
 """
-from .protocol_runner import ProtocolRunner
+from .protocol_runner import ProtocolRunner, ProtocolRunData
 from .protocol_file import ProtocolFile, ProtocolFileType
 from .create_simulating_runner import create_simulating_runner
 
 __all__ = [
     "ProtocolRunner",
+    "ProtocolRunData",
     "ProtocolFile",
     "ProtocolFileType",
     "create_simulating_runner",

api/src/opentrons/protocol_runner/protocol_runner.py

@@ -1,7 +1,13 @@
 """Protocol run control and management."""
-from typing import Optional, Sequence
+from dataclasses import dataclass
+from typing import List, Optional
 
-from opentrons.protocol_engine import ProtocolEngine, Command as ProtocolCommand
+from opentrons.protocol_engine import (
+    ProtocolEngine,
+    Command,
+    LoadedLabware,
+    LoadedPipette,
+)
 
 from .protocol_file import ProtocolFile, ProtocolFileType
 from .task_queue import TaskQueue, TaskQueuePhase
@@ -12,6 +18,15 @@
 from .python_executor import PythonExecutor
 
 
+@dataclass(frozen=True)
+class ProtocolRunData:
+    """Data from a protocol run."""
+
+    commands: List[Command]
+    labware: List[LoadedLabware]
+    pipettes: List[LoadedPipette]
+
+
 class ProtocolRunner:
     """An interface to manage and control a protocol run.
 
@@ -45,11 +60,6 @@ def __init__(
         self._python_context_creator = python_context_creator or PythonContextCreator()
         self._python_executor = python_executor or PythonExecutor()
 
-    @property
-    def engine(self) -> ProtocolEngine:
-        """Get the runner's underlying ProtocolEngine."""
-        return self._protocol_engine
-
     def load(self, protocol_file: ProtocolFile) -> None:
         """Load a ProtocolFile into managed ProtocolEngine.
 
@@ -92,12 +102,17 @@ async def join(self) -> None:
         """
         return await self._task_queue.join()
 
-    async def run(self, protocol_file: ProtocolFile) -> Sequence[ProtocolCommand]:
+    async def run(self, protocol_file: ProtocolFile) -> ProtocolRunData:
         """Run a given protocol to completion."""
         self.load(protocol_file)
         self.play()
         await self.join()
-        return self._protocol_engine.state_view.commands.get_all()
+
+        commands = self._protocol_engine.state_view.commands.get_all()
+        labware = self._protocol_engine.state_view.labware.get_all()
+        pipettes = self._protocol_engine.state_view.pipettes.get_all()
+
+        return ProtocolRunData(commands=commands, labware=labware, pipettes=pipettes)
 
     def _load_json(self, protocol_file: ProtocolFile) -> None:
         protocol = self._json_file_reader.read(protocol_file)

api/src/opentrons/protocol_runner/python_executor.py

@@ -10,23 +10,16 @@
 class PythonExecutor:
     """Execute a given PythonProtocol's run method with a ProtocolContext."""
 
-    def __init__(self) -> None:
-        """Initialize the executor with a thread pool.
-
-        A PythonExecutor uses its own ThreadPoolExecutor (rather than the default)
-        to avoid thread pool exhaustion from tying up protocol execution.
-        """
-        self._loop = asyncio.get_running_loop()
-        # fixme(mm, 2021-08-09): This class should be a context manager and call
-        # self._thread_pool.shutdown(). Currently, I think we leak threads.
-        self._thread_pool = ThreadPoolExecutor(max_workers=1)
-
-    async def execute(self, protocol: PythonProtocol, context: ProtocolContext) -> None:
+    @staticmethod
+    async def execute(protocol: PythonProtocol, context: ProtocolContext) -> None:
         """Execute a PythonProtocol using the given ProtocolContext.
 
         Runs the protocol asynchronously in a child thread.
         """
-        await self._loop.run_in_executor(
-            executor=self._thread_pool,
-            func=partial(protocol.run, context),
-        )
+        loop = asyncio.get_running_loop()
+
+        with ThreadPoolExecutor(max_workers=1) as executor:
+            await loop.run_in_executor(
+                executor=executor,
+                func=partial(protocol.run, context),
+            )

api/tests/opentrons/protocol_runner/test_protocol_runner_smoke.py

@@ -16,10 +16,10 @@
 from opentrons.protocol_api_experimental import DeckSlotName
 
 from opentrons.protocol_engine import (
+    DeckSlotLocation,
     LoadedLabware,
     LoadedPipette,
     PipetteName,
-    DeckSlotLocation,
     commands,
 )
 from opentrons.protocol_runner import (
@@ -37,22 +37,25 @@ async def test_protocol_runner_with_python(python_protocol_file: Path) -> None:
     )
 
     subject = await create_simulating_runner()
-    commands_result = await subject.run(protocol_file)
-    pipettes_result = subject.engine.state_view.pipettes.get_all()
-    labware_result = subject.engine.state_view.labware.get_all()
+    result = await subject.run(protocol_file)
+    commands_result = result.commands
+    pipettes_result = result.pipettes
+    labware_result = result.labware
 
     pipette_id_captor = matchers.Captor()
     labware_id_captor = matchers.Captor()
 
     expected_pipette = LoadedPipette.construct(
-        id=pipette_id_captor, pipetteName=PipetteName.P300_SINGLE, mount=MountType.LEFT
+        id=pipette_id_captor,
+        pipetteName=PipetteName.P300_SINGLE,
+        mount=MountType.LEFT,
     )
 
     expected_labware = LoadedLabware.construct(
         id=labware_id_captor,
+        location=DeckSlotLocation(slot=DeckSlotName.SLOT_1),
         loadName="opentrons_96_tiprack_300ul",
         definitionUri="opentrons/opentrons_96_tiprack_300ul/1",
-        location=DeckSlotLocation(slot=DeckSlotName.SLOT_1),
     )
 
     assert expected_pipette in pipettes_result
@@ -83,9 +86,10 @@ async def test_protocol_runner_with_json(json_protocol_file: Path) -> None:
     )
 
     subject = await create_simulating_runner()
-    commands_result = await subject.run(protocol_file)
-    pipettes_result = subject.engine.state_view.pipettes.get_all()
-    labware_result = subject.engine.state_view.labware.get_all()
+    result = await subject.run(protocol_file)
+    commands_result = result.commands
+    pipettes_result = result.pipettes
+    labware_result = result.labware
 
     expected_pipette = LoadedPipette(
         id="pipette-id",
@@ -95,9 +99,9 @@ async def test_protocol_runner_with_json(json_protocol_file: Path) -> None:
 
     expected_labware = LoadedLabware(
         id="labware-id",
+        location=DeckSlotLocation(slot=DeckSlotName.SLOT_1),
         loadName="opentrons_96_tiprack_300ul",
         definitionUri="opentrons/opentrons_96_tiprack_300ul/1",
-        location=DeckSlotLocation(slot=DeckSlotName.SLOT_1),
     )
 
     assert expected_pipette in pipettes_result

robot-server/robot_server/protocols/analysis_models.py

@@ -5,12 +5,7 @@
 from typing import List, Union
 from typing_extensions import Literal
 
-from opentrons.types import MountType
-from opentrons.protocol_engine import (
-    Command as EngineCommand,
-    PipetteName,
-    LabwareLocation,
-)
+from opentrons.protocol_engine import Command, LoadedLabware, LoadedPipette
 
 
 class AnalysisStatus(str, Enum):
@@ -52,23 +47,6 @@ class PendingAnalysis(AnalysisSummary):
     status: Literal[AnalysisStatus.PENDING] = AnalysisStatus.PENDING
 
 
-class AnalysisPipette(BaseModel):
-    """A pipette that the protocol is expected to use, based on the analysis."""
-
-    id: str
-    pipetteName: PipetteName
-    mount: MountType
-
-
-class AnalysisLabware(BaseModel):
-    """A labware that the protocol is expected to use, based on the analysis."""
-
-    id: str
-    loadName: str
-    definitionUri: str
-    location: LabwareLocation
-
-
 class CompletedAnalysis(AnalysisSummary):
     """A completed protocol run analysis.
 
@@ -95,15 +73,15 @@ class CompletedAnalysis(AnalysisSummary):
         ...,
         description="Whether the protocol is expected to run successfully",
     )
-    pipettes: List[AnalysisPipette] = Field(
+    pipettes: List[LoadedPipette] = Field(
         ...,
         description="Pipettes used by the protocol",
     )
-    labware: List[AnalysisLabware] = Field(
+    labware: List[LoadedLabware] = Field(
         ...,
         description="Labware used by the protocol",
     )
-    commands: List[EngineCommand] = Field(
+    commands: List[Command] = Field(
         ...,
         description="The protocol commands the run is expected to produce",
     )

robot-server/robot_server/protocols/analysis_store.py

@@ -1,16 +1,18 @@
 """Protocol analysis storage."""
 from typing import Dict, List, Set, Sequence
 
-from opentrons.calibration_storage.helpers import uri_from_details
-from opentrons.protocol_engine import commands as pe_commands
+from opentrons.protocol_engine import (
+    Command,
+    CommandStatus,
+    LoadedPipette,
+    LoadedLabware,
+)
 
 from .analysis_models import (
     ProtocolAnalysis,
     PendingAnalysis,
     CompletedAnalysis,
     AnalysisResult,
-    AnalysisLabware,
-    AnalysisPipette,
 )
 
 
@@ -35,41 +37,18 @@ def add_pending(self, protocol_id: str, analysis_id: str) -> List[ProtocolAnalys
     def update(
         self,
         analysis_id: str,
-        commands: Sequence[pe_commands.Command],
+        commands: Sequence[Command],
+        labware: Sequence[LoadedLabware],
+        pipettes: Sequence[LoadedPipette],
         errors: Sequence[Exception],
     ) -> None:
         """Update analysis results in the store."""
-        labware = []
-        pipettes = []
         # TODO(mc, 2021-08-25): return error details objects, not strings
         error_messages = [str(e) for e in errors]
 
-        for c in commands:
-            if isinstance(c, pe_commands.LoadLabware) and c.result is not None:
-                labware.append(
-                    AnalysisLabware(
-                        id=c.result.labwareId,
-                        loadName=c.data.loadName,
-                        definitionUri=uri_from_details(
-                            load_name=c.data.loadName,
-                            namespace=c.data.namespace,
-                            version=c.data.version,
-                        ),
-                        location=c.data.location,
-                    )
-                )
-            elif isinstance(c, pe_commands.LoadPipette) and c.result is not None:
-                pipettes.append(
-                    AnalysisPipette(
-                        id=c.result.pipetteId,
-                        pipetteName=c.data.pipetteName,
-                        mount=c.data.mount,
-                    )
-                )
-
         if len(error_messages) > 0:
             result = AnalysisResult.ERROR
-        elif any(c.status == pe_commands.CommandStatus.FAILED for c in commands):
+        elif any(c.status == CommandStatus.FAILED for c in commands):
             result = AnalysisResult.NOT_OK
         else:
             result = AnalysisResult.OK
@@ -78,9 +57,9 @@ def update(
             id=analysis_id,
             result=result,
             commands=list(commands),
+            labware=list(labware),
+            pipettes=list(pipettes),
             errors=error_messages,
-            labware=labware,
-            pipettes=pipettes,
         )
 
     def get_by_protocol(self, protocol_id: str) -> List[ProtocolAnalysis]:

robot-server/robot_server/protocols/protocol_analyzer.py

@@ -1,6 +1,6 @@
 """Protocol analysis module."""
-from typing import Sequence
-from opentrons.protocol_engine import Command as ProtocolCommand
+from typing import List
+from opentrons.protocol_engine import Command, LoadedLabware, LoadedPipette
 from opentrons.protocol_runner import ProtocolRunner
 
 from .protocol_store import ProtocolResource
@@ -25,16 +25,23 @@ async def analyze(
         analysis_id: str,
     ) -> None:
         """Analyze a given protocol, storing the analysis when complete."""
-        commands: Sequence[ProtocolCommand] = []
-        errors: Sequence[Exception] = []
+        commands: List[Command] = []
+        labware: List[LoadedLabware] = []
+        pipettes: List[LoadedPipette] = []
+        errors: List[Exception] = []
 
         try:
-            commands = await self._protocol_runner.run(protocol_resource)
+            result = await self._protocol_runner.run(protocol_resource)
+            commands = result.commands
+            labware = result.labware
+            pipettes = result.pipettes
         except Exception as e:
             errors = [e]
 
         self._analysis_store.update(
             analysis_id=analysis_id,
             commands=commands,
+            labware=labware,
+            pipettes=pipettes,
             errors=errors,
         )