chore(api): allow specifying module models to sim (#15104)

The simulator setup files let you specify modules, but you can only do
it by specifying their types (i.e. magdeck, tempdeck) and not their
modules (i.e. magneticModuleV1, temperatureModuleV2). This hasn't really
mattered up til now because nothing cares, but RSQ-6 will have both
desktop and ODD looking pretty hard at which generation of module is
present on a flex, so we need a way to specify.

This adds a model key to the module section of the simulator setup (and
sets the test-flex.json and test.json up with it) so that you can
specify module models.

Also, while I was at it, improve the type checking of the init function
of the threadmanager.

Closes RSQ-6

api/mypy.ini

@@ -4,7 +4,7 @@ show_error_codes = True
 warn_unused_configs = True
 strict = True
 # TODO(mc, 2021-09-12): work through and remove these exclusions
-exclude = tests/opentrons/(hardware_control/test_(?!(ot3|module_control)).*py|hardware_control/integration/|hardware_control/emulation/|hardware_control/modules/|protocols/advanced_control/|protocols/api_support/|protocols/duration/|protocols/execution/|protocols/fixtures/|protocols/geometry/)
+exclude = tests/opentrons/(hardware_control/test_(?!(ot3|module_control|modules|thread_manager)).*py|hardware_control/integration/|hardware_control/emulation/|hardware_control/modules/|protocols/advanced_control/|protocols/api_support/|protocols/duration/|protocols/execution/|protocols/fixtures/|protocols/geometry/)
 
 [pydantic-mypy]
 init_forbid_extra = True

api/src/opentrons/__init__.py

@@ -115,16 +115,14 @@ async def _create_thread_manager() -> ThreadManagedHardware:
         from opentrons.hardware_control.ot3api import OT3API
 
         thread_manager = ThreadManager(
-            OT3API.build_hardware_controller,
+            ThreadManager.nonblocking_builder(OT3API.build_hardware_controller),
             use_usb_bus=ff.rear_panel_integration(),
-            threadmanager_nonblocking=True,
             status_bar_enabled=ff.status_bar_enabled(),
             feature_flags=hw_types.HardwareFeatureFlags.build_from_ff(),
         )
     else:
         thread_manager = ThreadManager(
-            HardwareAPI.build_hardware_controller,
-            threadmanager_nonblocking=True,
+            ThreadManager.nonblocking_builder(HardwareAPI.build_hardware_controller),
             port=_get_motor_control_serial_port(),
             firmware=_find_smoothie_file(),
             feature_flags=hw_types.HardwareFeatureFlags.build_from_ff(),

api/src/opentrons/hardware_control/__init__.py

@@ -38,7 +38,8 @@
 ]
 HardwareControlAPI = Union[OT2HardwareControlAPI, OT3HardwareControlAPI]
 
-ThreadManagedHardware = ThreadManager[HardwareControlAPI]
+# this type ignore is because of https://github.com/python/mypy/issues/13437
+ThreadManagedHardware = ThreadManager[HardwareControlAPI]  # type: ignore[misc]
 SyncHardwareAPI = SynchronousAdapter[HardwareControlAPI]
 
 __all__ = [

api/src/opentrons/hardware_control/api.py

@@ -78,6 +78,8 @@
 
 mod_log = logging.getLogger(__name__)
 
+AttachedModuleSpec = Dict[str, List[Union[str, Tuple[str, str]]]]
+
 
 class API(
     ExecutionManagerProvider,
@@ -255,7 +257,7 @@ async def build_hardware_simulator(
         attached_instruments: Optional[
             Dict[top_types.Mount, Dict[str, Optional[str]]]
         ] = None,
-        attached_modules: Optional[Dict[str, List[str]]] = None,
+        attached_modules: Optional[Dict[str, List[modules.SimulatingModule]]] = None,
         config: Optional[Union[RobotConfig, OT3Config]] = None,
         loop: Optional[asyncio.AbstractEventLoop] = None,
         strict_attached_instruments: bool = True,

api/src/opentrons/hardware_control/backends/ot3simulator.py

@@ -104,7 +104,7 @@ class OT3Simulator(FlexBackend):
     async def build(
         cls,
         attached_instruments: Dict[OT3Mount, Dict[str, Optional[str]]],
-        attached_modules: Dict[str, List[str]],
+        attached_modules: Dict[str, List[modules.SimulatingModule]],
         config: OT3Config,
         loop: asyncio.AbstractEventLoop,
         strict_attached_instruments: bool = True,
@@ -130,7 +130,7 @@ async def build(
     def __init__(
         self,
         attached_instruments: Dict[OT3Mount, Dict[str, Optional[str]]],
-        attached_modules: Dict[str, List[str]],
+        attached_modules: Dict[str, List[modules.SimulatingModule]],
         config: OT3Config,
         loop: asyncio.AbstractEventLoop,
         strict_attached_instruments: bool = True,
@@ -605,13 +605,14 @@ async def increase_z_l_hold_current(self) -> AsyncIterator[None]:
     @ensure_yield
     async def watch(self, loop: asyncio.AbstractEventLoop) -> None:
         new_mods_at_ports = []
-        for mod, serials in self._stubbed_attached_modules.items():
-            for serial in serials:
+        for mod_name, list_of_modules in self._stubbed_attached_modules.items():
+            for module_details in list_of_modules:
                 new_mods_at_ports.append(
                     modules.SimulatingModuleAtPort(
-                        port=f"/dev/ot_module_sim_{mod}{str(serial)}",
-                        name=mod,
-                        serial_number=serial,
+                        port=f"/dev/ot_module_sim_{mod_name}{str(module_details.serial_number)}",
+                        name=mod_name,
+                        serial_number=module_details.serial_number,
+                        model=module_details.model,
                     )
                 )
         await self.module_controls.register_modules(new_mods_at_ports=new_mods_at_ports)

api/src/opentrons/hardware_control/backends/simulator.py

@@ -49,7 +49,7 @@ class Simulator:
     async def build(
         cls,
         attached_instruments: Dict[types.Mount, Dict[str, Optional[str]]],
-        attached_modules: Dict[str, List[str]],
+        attached_modules: Dict[str, List[modules.SimulatingModule]],
         config: RobotConfig,
         loop: asyncio.AbstractEventLoop,
         strict_attached_instruments: bool = True,
@@ -71,8 +71,9 @@ async def build(
                                      This dict should map mounts to either
                                      empty dicts or to dicts containing
                                      'model' and 'id' keys.
-        :param attached_modules: A list of module model names (e.g.
-                                 `'tempdeck'` or `'magdeck'`) representing
+        :param attached_modules: A map of module type names (e.g.
+                                 `'tempdeck'` or `'magdeck'`) to lists of SimulatingModel
+                                 dataclasses representing
                                  modules the simulator should assume are
                                  attached. Like `attached_instruments`, used
                                  to make the simulator match the setup of the
@@ -105,7 +106,7 @@ async def build(
     def __init__(
         self,
         attached_instruments: Dict[types.Mount, Dict[str, Optional[str]]],
-        attached_modules: Dict[str, List[str]],
+        attached_modules: Dict[str, List[modules.SimulatingModule]],
         config: RobotConfig,
         loop: asyncio.AbstractEventLoop,
         gpio_chardev: GPIODriverLike,
@@ -333,13 +334,14 @@ def set_active_current(self, axis_currents: Dict[Axis, float]) -> None:
     @ensure_yield
     async def watch(self) -> None:
         new_mods_at_ports = []
-        for mod, serials in self._stubbed_attached_modules.items():
-            for serial in serials:
+        for mod_name, list_of_modules in self._stubbed_attached_modules.items():
+            for module_details in list_of_modules:
                 new_mods_at_ports.append(
                     modules.SimulatingModuleAtPort(
-                        port=f"/dev/ot_module_sim_{mod}{str(serial)}",
-                        name=mod,
-                        serial_number=serial,
+                        port=f"/dev/ot_module_sim_{mod_name}{str(module_details.serial_number)}",
+                        name=mod_name,
+                        serial_number=module_details.serial_number,
+                        model=module_details.model,
                     )
                 )
         await self.module_controls.register_modules(new_mods_at_ports=new_mods_at_ports)

api/src/opentrons/hardware_control/module_control.py

@@ -161,9 +161,14 @@ async def register_modules(
                 port=mod.port,
                 usb_port=mod.usb_port,
                 type=modules.MODULE_TYPE_BY_NAME[mod.name],
-                sim_serial_number=mod.serial_number
-                if isinstance(mod, SimulatingModuleAtPort)
-                else None,
+                sim_serial_number=(
+                    mod.serial_number
+                    if isinstance(mod, SimulatingModuleAtPort)
+                    else None
+                ),
+                sim_model=(
+                    mod.model if isinstance(mod, SimulatingModuleAtPort) else None
+                ),
             )
             self._available_modules.append(new_instance)
             log.info(

api/src/opentrons/hardware_control/modules/__init__.py

@@ -12,6 +12,7 @@
     UpdateError,
     ModuleAtPort,
     SimulatingModuleAtPort,
+    SimulatingModule,
     ModuleType,
     ModuleModel,
     TemperatureStatus,
@@ -35,6 +36,7 @@
     "UpdateError",
     "ModuleAtPort",
     "SimulatingModuleAtPort",
+    "SimulatingModule",
     "HeaterShaker",
     "ModuleType",
     "ModuleModel",

api/src/opentrons/hardware_control/modules/types.py

@@ -9,6 +9,7 @@
     Tuple,
     Awaitable,
     Union,
+    Optional,
     cast,
     TYPE_CHECKING,
 )
@@ -127,8 +128,14 @@ class ModuleAtPort:
 
 
 @dataclass(kw_only=True)
-class SimulatingModuleAtPort(ModuleAtPort):
+class SimulatingModule:
     serial_number: str
+    model: Optional[str]
+
+
+@dataclass(kw_only=True)
+class SimulatingModuleAtPort(ModuleAtPort, SimulatingModule):
+    pass
 
 
 class BundledFirmware(NamedTuple):

api/src/opentrons/hardware_control/ot3api.py

@@ -413,7 +413,7 @@ async def build_hardware_simulator(
             Dict[OT3Mount, Dict[str, Optional[str]]],
             Dict[top_types.Mount, Dict[str, Optional[str]]],
         ] = None,
-        attached_modules: Optional[Dict[str, List[str]]] = None,
+        attached_modules: Optional[Dict[str, List[modules.SimulatingModule]]] = None,
         config: Union[RobotConfig, OT3Config, None] = None,
         loop: Optional[asyncio.AbstractEventLoop] = None,
         strict_attached_instruments: bool = True,

api/src/opentrons/hardware_control/scripts/repl.py

@@ -114,8 +114,6 @@ def synchronizer(*args: Any, **kwargs: Any) -> Any:
     def build_thread_manager() -> ThreadManager[Union[API, OT3API]]:
         return ThreadManager(
             API.build_hardware_controller,
-            use_usb_bus=ff.rear_panel_integration(),
-            update_firmware=update_firmware,
             feature_flags=HardwareFeatureFlags.build_from_ff(),
         )
 

api/src/opentrons/hardware_control/simulator_setup.py

@@ -11,6 +11,7 @@
 from opentrons.types import Mount
 from opentrons.hardware_control import API, HardwareControlAPI, ThreadManager
 from opentrons.hardware_control.types import OT3Mount, HardwareFeatureFlags
+from opentrons.hardware_control.modules import SimulatingModule
 
 
 # Name and kwargs for a module function
@@ -24,6 +25,7 @@ class ModuleCall:
 @dataclass(frozen=True)
 class ModuleItem:
     serial_number: str
+    model: str
     calls: List[ModuleCall] = field(default_factory=list)
 
 
@@ -59,7 +61,10 @@ async def _simulator_for_setup(
         return await API.build_hardware_simulator(
             attached_instruments=setup.attached_instruments,
             attached_modules={
-                k: [m.serial_number for m in v]
+                k: [
+                    SimulatingModule(serial_number=m.serial_number, model=m.model)
+                    for m in v
+                ]
                 for k, v in setup.attached_modules.items()
             },
             config=setup.config,
@@ -73,7 +78,10 @@ async def _simulator_for_setup(
         return await OT3API.build_hardware_simulator(
             attached_instruments=setup.attached_instruments,
             attached_modules={
-                k: [m.serial_number for m in v]
+                k: [
+                    SimulatingModule(serial_number=m.serial_number, model=m.model)
+                    for m in v
+                ]
                 for k, v in setup.attached_modules.items()
             },
             config=setup.config,
@@ -114,7 +122,10 @@ def _thread_manager_for_setup(
             API.build_hardware_simulator,
             attached_instruments=setup.attached_instruments,
             attached_modules={
-                k: [m.serial_number for m in v]
+                k: [
+                    SimulatingModule(serial_number=m.serial_number, model=m.model)
+                    for m in v
+                ]
                 for k, v in setup.attached_modules.items()
             },
             config=setup.config,
@@ -128,7 +139,10 @@ def _thread_manager_for_setup(
             OT3API.build_hardware_simulator,
             attached_instruments=setup.attached_instruments,
             attached_modules={
-                k: [m.serial_number for m in v]
+                k: [
+                    SimulatingModule(serial_number=m.serial_number, model=m.model)
+                    for m in v
+                ]
                 for k, v in setup.attached_modules.items()
             },
             config=setup.config,
@@ -215,6 +229,7 @@ def _prepare_for_simulator_setup(key: str, value: Dict[str, Any]) -> Any:
                 attached_modules.setdefault(key, []).append(
                     ModuleItem(
                         serial_number=obj["serial_number"],
+                        model=obj["model"],
                         calls=[ModuleCall(**data) for data in obj["calls"]],
                     )
                 )
@@ -236,6 +251,7 @@ def _prepare_for_ot3_simulator_setup(key: str, value: Dict[str, Any]) -> Any:
                 attached_modules.setdefault(key, []).append(
                     ModuleItem(
                         serial_number=obj["serial_number"],
+                        model=obj["model"],
                         calls=[ModuleCall(**data) for data in obj["calls"]],
                     )
                 )

api/src/opentrons/hardware_control/thread_manager.py

@@ -1,9 +1,9 @@
 """Manager for the :py:class:`.hardware_control.API` thread."""
+import functools
 import threading
 import logging
 import asyncio
 import inspect
-import functools
 import weakref
 from typing import (
     Any,
@@ -195,7 +195,10 @@ class ThreadManager(Generic[WrappedObj]):
 
     If you want to wait for the managed object's creation separately
     (with managed_thread_ready_blocking or managed_thread_ready_async)
-    then pass threadmanager_nonblocking=True as a kwarg
+     use the nonblocking_builder static method to add an attribute to the builder
+    function, i.e.
+
+    thread_manager = ThreadManager(ThreadManager.nonblocking_builder(builder), ...)
 
     Example
     -------
@@ -206,15 +209,52 @@ class ThreadManager(Generic[WrappedObj]):
     >>> api_single_thread.sync.home() # call as blocking sync
     """
 
+    Builder = ParamSpec("Builder")
+    Built = TypeVar("Built")
+
+    @staticmethod
+    def nonblocking_builder(
+        builder: Callable[Builder, Awaitable[Built]]
+    ) -> Callable[Builder, Awaitable[Built]]:
+        """Wrap an instance of a builder function to make initializes that use it nonblocking.
+
+        For instance, you can build a ThreadManager like this:
+
+        thread_manager = ThreadManager(ThreadManager.nonblocking_builder(API.build_hardware_controller), ...)
+
+        to make the initialize call return immediately so you can later wait on it via
+        managed_thread_ready_blocking or managed_thread_ready_async
+        """
+
+        @functools.wraps(builder)
+        async def wrapper(
+            *args: ThreadManager.Builder.args, **kwargs: ThreadManager.Builder.kwargs
+        ) -> ThreadManager.Built:
+            return await builder(*args, **kwargs)
+
+        setattr(wrapper, "nonblocking", True)
+        return wrapper
+
     def __init__(
         self,
-        builder: Callable[..., Awaitable[WrappedObj]],
-        *args: Any,
-        **kwargs: Any,
+        builder: Callable[Builder, Awaitable[WrappedObj]],
+        *args: Builder.args,
+        **kwargs: Builder.kwargs,
     ) -> None:
         """Build the ThreadManager.
 
-        :param builder: The API function to use
+        builder: The api function to use to build the instance.
+
+        The args and kwargs will be forwarded to the builder function.
+
+        Note: by default, this function will block until the managed thread is ready and the hardware controller
+        has been built. To make this function return immediately you can wrap its builder argument in
+        ThreadManager.nonblocking_builder(), like this:
+
+        thread_manager = ThreadManager(ThreadManager.nonblocking_builder(API.build_hardware_controller), ...)
+
+        Afterwards, you'll need to call ThreadManager.managed_thread_ready_blocking or its async variant before
+        you can actually use thei nstance.
         """
 
         self._loop: Optional[asyncio.AbstractEventLoop] = None
@@ -234,7 +274,7 @@ def __init__(
             asyncio.get_child_watcher()
         except NotImplementedError:
             pass
-        blocking = not kwargs.pop("threadmanager_nonblocking", False)
+        blocking = not getattr(builder, "nonblocking", False)
         target = object.__getattribute__(self, "_build_and_start_loop")
         thread = threading.Thread(
             target=target,