perf(robot-server): Optimize adding many offsets at once through POST /labwareOffsets

robot-server/robot_server/labware_offsets/router.py

@@ -79,8 +79,7 @@ async def post_labware_offsets(  # noqa: D103
         )
     ]
 
-    for new_offset in new_offsets:
-        store.add(new_offset)
+    store.add(new_offsets)
 
     stored_offsets = [
         StoredLabwareOffset.model_construct(

robot-server/robot_server/labware_offsets/store.py

@@ -68,25 +68,26 @@ def __init__(
 
     def add(
         self,
-        offset: IncomingStoredLabwareOffset,
+        offsets: list[IncomingStoredLabwareOffset],
     ) -> None:
-        """Store a new labware offset."""
+        """Store new labware offsets."""
         with self._sql_engine.begin() as transaction:
-            offset_row_id = transaction.execute(
-                sqlalchemy.insert(labware_offset_table).values(
-                    _pydantic_to_sql_offset(offset)
-                )
-            ).inserted_primary_key.row_id
-            location_components_to_insert = list(
-                _pydantic_to_sql_location_sequence_iterator(offset, offset_row_id)
-            )
-            if location_components_to_insert:
-                transaction.execute(
-                    sqlalchemy.insert(
-                        labware_offset_location_sequence_components_table
-                    ),
-                    location_components_to_insert,
+            for offset in offsets:
+                offset_row_id = transaction.execute(
+                    sqlalchemy.insert(labware_offset_table).values(
+                        _pydantic_to_sql_offset(offset)
+                    )
+                ).inserted_primary_key.row_id
+                location_components_to_insert = list(
+                    _pydantic_to_sql_location_sequence_iterator(offset, offset_row_id)
                 )
+                if location_components_to_insert:
+                    transaction.execute(
+                        sqlalchemy.insert(
+                            labware_offset_location_sequence_components_table
+                        ),
+                        location_components_to_insert,
+                    )
         self._publish_change_notification()
 
     def get_all(self) -> list[StoredLabwareOffset]:

robot-server/tests/labware_offsets/test_store.py

@@ -56,13 +56,15 @@ def subject(
 def test_empty_search(subject: LabwareOffsetStore) -> None:
     """Searching with no filters should return no results."""
     subject.add(
-        IncomingStoredLabwareOffset(
-            id="id",
-            createdAt=datetime.now(timezone.utc),
-            definitionUri="namespace/load_name/1",
-            locationSequence="anyLocation",
-            vector=LabwareOffsetVector(x=1, y=2, z=3),
-        )
+        [
+            IncomingStoredLabwareOffset(
+                id="id",
+                createdAt=datetime.now(timezone.utc),
+                definitionUri="namespace/load_name/1",
+                locationSequence="anyLocation",
+                vector=LabwareOffsetVector(x=1, y=2, z=3),
+            )
+        ]
     )
     assert subject.search([]) == []
 
@@ -79,13 +81,15 @@ def test_search_most_recent_only(subject: LabwareOffsetStore) -> None:
     ]
     for id, definition_uri in ids_and_definition_uris:
         subject.add(
-            IncomingStoredLabwareOffset(
-                id=id,
-                definitionUri=definition_uri,
-                createdAt=datetime.now(timezone.utc),
-                locationSequence="anyLocation",
-                vector=LabwareOffsetVector(x=1, y=2, z=3),
-            )
+            [
+                IncomingStoredLabwareOffset(
+                    id=id,
+                    definitionUri=definition_uri,
+                    createdAt=datetime.now(timezone.utc),
+                    locationSequence="anyLocation",
+                    vector=LabwareOffsetVector(x=1, y=2, z=3),
+                )
+            ]
         )
 
     results = subject.search([SearchFilter(mostRecentOnly=True)])
@@ -263,7 +267,7 @@ def test_filter_fields(
         ),
     }
     for offset in offsets.values():
-        subject.add(offset)
+        subject.add([offset])
     results = subject.search(
         [
             SearchFilter(
@@ -323,8 +327,7 @@ def test_filter_field_combinations(subject: LabwareOffsetStore) -> None:
         for offset in labware_offsets
     ]
 
-    for labware_offset in labware_offsets:
-        subject.add(labware_offset)
+    subject.add(labware_offsets)
 
     # Filter accepting any value for each field (i.e. a return-everything filter):
     result = subject.search([SearchFilter()])
@@ -368,13 +371,15 @@ def test_filter_combinations(subject: LabwareOffsetStore) -> None:
     ]
     for id, definition_uri in ids_and_definition_uris:
         subject.add(
-            IncomingStoredLabwareOffset(
-                id=id,
-                createdAt=datetime.now(timezone.utc),
-                definitionUri=definition_uri,
-                locationSequence=ANY_LOCATION,
-                vector=LabwareOffsetVector(x=1, y=2, z=3),
-            )
+            [
+                IncomingStoredLabwareOffset(
+                    id=id,
+                    createdAt=datetime.now(timezone.utc),
+                    definitionUri=definition_uri,
+                    locationSequence=ANY_LOCATION,
+                    vector=LabwareOffsetVector(x=1, y=2, z=3),
+                )
+            ]
         )
 
     # Multiple filters should be OR'd together.
@@ -429,9 +434,7 @@ def test_delete(subject: LabwareOffsetStore) -> None:
     with pytest.raises(LabwareOffsetNotFoundError):
         subject.delete("b")
 
-    subject.add(a)
-    subject.add(b)
-    subject.add(c)
+    subject.add([a, b, c])
 
     assert subject.delete(b.id) == out_b
     assert subject.get_all() == [out_a, out_c]
@@ -473,7 +476,7 @@ def test_handle_unknown(
         ],
         vector=incoming_valid.vector,
     )
-    subject.add(incoming_valid)
+    subject.add([incoming_valid])
     with sql_engine.begin() as transaction:
         transaction.execute(
             sqlalchemy.insert(labware_offset_location_sequence_components_table).values(
@@ -496,13 +499,15 @@ def test_notifications(
     """It should publish notifications any time the set of labware offsets changes."""
     decoy.verify(mock_labware_offsets_publisher.publish_labware_offsets(), times=0)
     subject.add(
-        IncomingStoredLabwareOffset(
-            id="id",
-            createdAt=datetime.now(timezone.utc),
-            definitionUri="definitionUri",
-            locationSequence=ANY_LOCATION,
-            vector=LabwareOffsetVector(x=1, y=2, z=3),
-        )
+        [
+            IncomingStoredLabwareOffset(
+                id="id",
+                createdAt=datetime.now(timezone.utc),
+                definitionUri="definitionUri",
+                locationSequence=ANY_LOCATION,
+                vector=LabwareOffsetVector(x=1, y=2, z=3),
+            )
+        ]
     )
     decoy.verify(mock_labware_offsets_publisher.publish_labware_offsets(), times=1)
     subject.delete("id")

robot-server/tests/labware_offsets/test_store_hypothesis.py

@@ -122,7 +122,7 @@ def test_round_trip(
 
     with TemporaryDirectory() as tmp_dir, make_sql_engine(Path(tmp_dir)) as sql_engine:
         subject = LabwareOffsetStore(sql_engine, labware_offsets_publisher=None)
-        subject.add(offset_to_add)
+        subject.add([offset_to_add])
         [offset_retrieved_by_get_all] = subject.get_all()
         [offset_retrieved_by_search] = subject.search([SearchFilter(id=id)])
 
@@ -142,20 +142,21 @@ class SimulatedStore:
     def __init__(self) -> None:
         self._entries: list[StoredLabwareOffset] = []
 
-    def add(self, offset: IncomingStoredLabwareOffset) -> None:  # noqa: D102
-        id_already_exists = any(
-            existing_offset.id == offset.id for existing_offset in self._entries
-        )
-        assert not id_already_exists
-        self._entries.append(
-            StoredLabwareOffset(
-                id=offset.id,
-                createdAt=offset.createdAt,
-                definitionUri=offset.definitionUri,
-                locationSequence=offset.locationSequence,
-                vector=offset.vector,
+    def add(self, offsets: list[IncomingStoredLabwareOffset]) -> None:  # noqa: D102
+        for offset in offsets:
+            id_already_exists = any(
+                existing_offset.id == offset.id for existing_offset in self._entries
+            )
+            assert not id_already_exists
+            self._entries.append(
+                StoredLabwareOffset(
+                    id=offset.id,
+                    createdAt=offset.createdAt,
+                    definitionUri=offset.definitionUri,
+                    locationSequence=offset.locationSequence,
+                    vector=offset.vector,
+                )
             )
-        )
 
     def get_all(self) -> list[StoredLabwareOffset]:  # noqa: D102
         return self._entries
@@ -274,8 +275,11 @@ def add(  # noqa: D102
             # different floats. Different floats are tried in the round-trip test.
             vector=LabwareOffsetVector(x=1, y=2, z=3),
         )
-        self._subject.add(to_add)
-        self._simulated_model.add(to_add)
+        # todo(mm, 2025-04-01): We should probably Hypothesis-test adding multiple
+        # offsets in a single add() call, but this RuleBasedStateMachine is already
+        # taking a lot of time searching the state space--time to break it up?
+        self._subject.add([to_add])
+        self._simulated_model.add([to_add])
         self._added_ids.add(to_add.id)
 
     @hypothesis.stateful.rule()