Adding assignment-unit cache to singleton DB

mephisto/abstractions/databases/local_singleton_database.py

@@ -68,6 +68,7 @@ def __init__(self, database_path=None):
 
         # Create singleton dictionaries for entries
         self._singleton_cache = {k: dict() for k in self._cached_classes}
+        self._assignment_to_unit_mapping: Dict[str, List[Unit]] = {}
 
     def shutdown(self) -> None:
         """Close all open connections"""
@@ -128,3 +129,93 @@ def new_agent(
         unit.db_status = AssignmentState.ASSIGNED
         unit.worker_id = agent.worker_id
         return agent_id
+
+    def find_units(
+        self,
+        task_id: Optional[str] = None,
+        task_run_id: Optional[str] = None,
+        requester_id: Optional[str] = None,
+        assignment_id: Optional[str] = None,
+        unit_index: Optional[int] = None,
+        provider_type: Optional[str] = None,
+        task_type: Optional[str] = None,
+        agent_id: Optional[str] = None,
+        worker_id: Optional[str] = None,
+        sandbox: Optional[bool] = None,
+        status: Optional[str] = None,
+    ) -> List[Unit]:
+        """
+        Uses caching to offset the cost of the most common queries. Defers
+        to the underlying DB for anything outside of those cases.
+        """
+
+        # Finding units is the most common small DB call to be optimized, as
+        # every assignment has multiple units. Thus, we try to break up the
+        # units to be queried by assignment, ensuring the most commonly
+        # queried edge is improved.
+        if assignment_id is not None:
+            if all(
+                v is None
+                for v in [
+                    task_id,
+                    task_run_id,
+                    requester_id,
+                    unit_index,
+                    provider_type,
+                    task_type,
+                    agent_id,
+                    worker_id,
+                    sandbox,
+                    status,
+                ]
+            ):
+                units = self._assignment_to_unit_mapping.get(assignment_id)
+                if units is None:
+                    units = super().find_units(assignment_id=assignment_id)
+                    self._assignment_to_unit_mapping[assignment_id] = units
+                return units
+
+        # Any other cases are less common and more complicated, and so we don't cache
+        return super().find_units(
+            task_id=task_id,
+            task_run_id=task_run_id,
+            requester_id=requester_id,
+            assignment_id=assignment_id,
+            unit_index=unit_index,
+            provider_type=provider_type,
+            task_type=task_type,
+            agent_id=agent_id,
+            worker_id=worker_id,
+            sandbox=sandbox,
+            status=status,
+        )
+
+    def new_unit(
+        self,
+        task_id: str,
+        task_run_id: str,
+        requester_id: str,
+        assignment_id: str,
+        unit_index: int,
+        pay_amount: float,
+        provider_type: str,
+        task_type: str,
+        sandbox: bool = True,
+    ) -> str:
+        """
+        Create a new unit with the given index. Raises EntryAlreadyExistsException
+        if there is already a unit for the given assignment with the given index.
+        """
+        if assignment_id in self._assignment_to_unit_mapping:
+            del self._assignment_to_unit_mapping[assignment_id]
+        return super().new_unit(
+            task_id=task_id,
+            task_run_id=task_run_id,
+            requester_id=requester_id,
+            assignment_id=assignment_id,
+            unit_index=unit_index,
+            pay_amount=pay_amount,
+            provider_type=provider_type,
+            task_type=task_type,
+            sandbox=sandbox,
+        )

mephisto/operations/operator.py

@@ -306,7 +306,7 @@ def _track_and_kill_runs(self):
                     tracked_run.architect.shutdown()
                     tracked_run.task_launcher.shutdown()
                     del self._task_runs_tracked[task_run.db_id]
-            time.sleep(2)
+            time.sleep(10)
 
     def force_shutdown(self, timeout=5):
         """