bug fixes

alab_management/dashboard/plotly/plotly_dashboard.py

@@ -113,9 +113,7 @@ def update_task_table(date_range):
     filtered_task_df = task_df[
         (task_df.started_at >= date_range[0]) & (task_df.started_at <= date_range[1])
     ]
-    return [
-        dmc.Text("Tasks", align="center", size="lg"),
-    ] + [
+    return [dmc.Text("Tasks", align="center", size="lg"),] + [
         build_task_row(task_type, filtered_task_df)
         for task_type in filtered_task_df.type.unique()
     ]

alab_management/task_manager/resource_requester.py

@@ -1,25 +1,25 @@
 """
 TaskLauncher is the core module of the system,
-which actually executes the tasks.
+which actually executes the tasks
 """
 import time
 from concurrent.futures import Future
 from datetime import datetime
 from threading import Thread
 from traceback import print_exc
-from typing import Any, Dict, List, Optional, Type, Union, cast
+from typing import Union, Dict, Optional, Type, List, Any, cast
 
 import dill
 from bson import ObjectId
 from pydantic import BaseModel, root_validator
 
-from alab_management.device_view.device import BaseDevice
 from alab_management.sample_view.sample import SamplePosition
-from alab_management.sample_view.sample_view import SamplePositionRequest
-from alab_management.task_view import TaskPriority
-from alab_management.utils.data_objects import get_collection
 
-from .enums import _EXTRA_REQUEST, RequestStatus
+from ..device_view.device import BaseDevice
+from ..sample_view.sample_view import SamplePositionRequest
+from ..task_view import TaskView, TaskPriority
+from ..utils.data_objects import get_collection
+from .enums import RequestStatus, _EXTRA_REQUEST
 
 _SampleRequestDict = Dict[str, int]
 _ResourceRequestDict = Dict[
@@ -45,10 +45,9 @@ class ResourcesRequest(BaseModel):
             ]
         },
         ...
-    ].
+    ]
 
-    See Also
-    --------
+    See Also:
         :py:class:`SamplePositionRequest <alab_management.sample_view.sample_view.SamplePositionRequest>`
     """
 
@@ -83,7 +82,9 @@ def preprocess(cls, values):  # pylint: disable=no-self-use,no-self-argument
 
 
 class RequestMixin:
-    """Simple wrapper for the request collection."""
+    """
+    Simple wrapper for the request collection
+    """
 
     def __init__(self):
         self._request_collection = get_collection("requests")
@@ -152,6 +153,7 @@ def request_resources(
         Request lab resources. Write the request into the database, and then the task manager will read from the
         database and assign the resources.
         """
+
         f = Future()
         if priority is None:
             priority = self.priority
@@ -220,7 +222,9 @@ def request_resources(
         }
 
     def release_resources(self, request_id: ObjectId) -> bool:
-        """Release a request by request_id."""
+        """
+        Release a request by request_id
+        """
         result = self._request_collection.update_one(
             {
                 "_id": request_id,
@@ -237,7 +241,7 @@ def release_resources(self, request_id: ObjectId) -> bool:
 
     def release_all_resources(self):
         """
-        Release all requests by task_id, used for error recovery.
+        Release all requests by task_id, used for error recovery
 
         For the requests that are not fulfilled, they will be marked as CANCELED.
 
@@ -270,7 +274,7 @@ def release_all_resources(self):
     def _check_request_status_loop(self):
         while not self._stop:
             try:
-                for request_id in self._waiting.copy():
+                for request_id in self._waiting.copy().keys():
                     status = self.get_request(request_id=request_id, projection=["status"])["status"]  # type: ignore
                     if status == RequestStatus.FULFILLED.name:
                         self._handle_fulfilled_request(request_id=request_id)

examples/fake_lab/__init__.py

@@ -16,15 +16,19 @@
 add_device(Furnace(name="furnace_4"))
 add_device(RobotArm(name="dummy"))
 
-add_standalone_sample_position(SamplePosition(
-    "furnace_table", description="Temporary position to transfer samples"
-))
+add_standalone_sample_position(
+    SamplePosition(
+        "furnace_table", description="Temporary position to transfer samples"
+    )
+)
 
-add_standalone_sample_position(SamplePosition(
-    "furnace_temp",
-    number=4,
-    description="Test positions",
-))
+add_standalone_sample_position(
+    SamplePosition(
+        "furnace_temp",
+        number=4,
+        description="Test positions",
+    )
+)
 
 add_task(Starting)
 add_task(Moving)

examples/fake_lab/tasks/moving.py

@@ -15,8 +15,7 @@ def __init__(self, samples: List[ObjectId], dest: str, *args, **kwargs):
 
     def run(self):
         with self.lab_view.request_resources(
-            {RobotArm: {},
-             None: {self.dest: 1, self.sample_position: 1}}
+            {RobotArm: {}, None: {self.dest: 1, self.sample_position: 1}}
         ) as (devices, sample_positions):
             robot_arm = cast(RobotArm, devices[RobotArm])
             robot_arm.run_program(
@@ -29,6 +28,6 @@ def run(self):
                 signal_value={
                     "src": sample_positions[None][self.sample_position][0],
                     "dest": sample_positions[None][self.dest][0],
-                }
+                },
             )
         return self.task_id

pyproject.toml

@@ -67,7 +67,8 @@ dependencies = [
     "pytest-env ~= 0.6.2",
     "requests >= 2.26.0",
     "flake8-bugbear >= 21.11.29",
-    "flake8-docstrings >= 1.6.0"
+    "flake8-docstrings >= 1.6.0",
+    "ruff"
 ]
 
 [project.optional-dependencies]

tests/fake_lab/__init__.py

@@ -16,15 +16,19 @@
 add_device(Furnace(name="furnace_4"))
 add_device(RobotArm(name="dummy"))
 
-add_standalone_sample_position(SamplePosition(
-    "furnace_table", description="Temporary position to transfer samples"
-))
+add_standalone_sample_position(
+    SamplePosition(
+        "furnace_table", description="Temporary position to transfer samples"
+    )
+)
 
-add_standalone_sample_position(SamplePosition(
-    "furnace_temp",
-    number=4,
-    description="Test positions",
-))
+add_standalone_sample_position(
+    SamplePosition(
+        "furnace_temp",
+        number=4,
+        description="Test positions",
+    )
+)
 
 add_task(Starting)
 add_task(Moving)

tests/fake_lab/tasks/moving.py

@@ -15,8 +15,7 @@ def __init__(self, samples: List[ObjectId], dest: str, *args, **kwargs):
 
     def run(self):
         with self.lab_view.request_resources(
-            {RobotArm: {},
-             None: {self.dest: 1, self.sample_position: 1}}
+            {RobotArm: {}, None: {self.dest: 1, self.sample_position: 1}}
         ) as (devices, sample_positions):
             robot_arm = cast(RobotArm, devices[RobotArm])
             robot_arm.run_program(
@@ -29,6 +28,6 @@ def run(self):
                 signal_value={
                     "src": sample_positions[None][self.sample_position][0],
                     "dest": sample_positions[None][self.dest][0],
-                }
+                },
             )
         return self.task_id

tests/test_device_view.py

@@ -112,7 +112,9 @@ def test_request_devices_single(self):
         )
         task_id = ObjectId()
 
-        devices = self.device_view.request_devices(task_id, device_names_str=device_types)
+        devices = self.device_view.request_devices(
+            task_id, device_names_str=device_types
+        )
         occupy_devices(devices, device_view=self.device_view, task_id=task_id)
         self.assertFalse(devices is None)
         for device_type, device in devices.items():
@@ -143,10 +145,15 @@ def test_request_device_timeout(self):
         task_id = ObjectId()
         task_id_2 = ObjectId()
 
-        devices = self.device_view.request_devices(task_id, device_types_str=device_types)
+        devices = self.device_view.request_devices(
+            task_id, device_types_str=device_types
+        )
         self.assertFalse(devices is None)
         occupy_devices(devices, device_view=self.device_view, task_id=task_id)
-        self.assertIs(None, self.device_view.request_devices(task_id_2, device_types_str=device_types))
+        self.assertIs(
+            None,
+            self.device_view.request_devices(task_id_2, device_types_str=device_types),
+        )
         release_devices(devices, device_view=self.device_view)
 
     def test_request_device_twice(self):

tests/test_experiment_manager.py

@@ -22,8 +22,10 @@ def test_handle_pending_experiments(self):
                 "name": "test",
                 "tags": ["test"],
                 "metadata": {"test": "test"},
-                "samples": [{"name": "test_sample", "metadata": {}, "tags": []},
-                            {"name": "test_sample_2", "metadata": {}, "tags": []}],
+                "samples": [
+                    {"name": "test_sample", "metadata": {}, "tags": []},
+                    {"name": "test_sample_2", "metadata": {}, "tags": []},
+                ],
                 "tasks": [
                     {
                         "type": "Heating",
@@ -104,8 +106,10 @@ def test_mark_completed_experiments(self):
                 "name": "test",
                 "tags": ["test"],
                 "metadata": {"test": "test"},
-                "samples": [{"name": "test_sample", "metadata": {}, "tags": []},
-                            {"name": "test_sample_2", "metadata": {}, "tags": []}],
+                "samples": [
+                    {"name": "test_sample", "metadata": {}, "tags": []},
+                    {"name": "test_sample_2", "metadata": {}, "tags": []},
+                ],
                 "tasks": [
                     {
                         "type": "Heating",

tests/test_task_manager.py

@@ -71,7 +71,9 @@ def test_task_requester(self):
         )
         self.resource_requester.release_resources(_id)
         time.sleep(2)
-        self.assertEqual(self.device_view.get_status("furnace_1"), DeviceTaskStatus.IDLE)
+        self.assertEqual(
+            self.device_view.get_status("furnace_1"), DeviceTaskStatus.IDLE
+        )
         self.assertEqual(
             self.sample_view.get_sample_position_status("furnace_1/inside"),
             (SamplePositionStatus.EMPTY, None),
@@ -98,7 +100,9 @@ def test_task_requester(self):
         )
         self.resource_requester.release_resources(_id)
         time.sleep(2)
-        self.assertEqual(self.device_view.get_status("furnace_1"), DeviceTaskStatus.IDLE)
+        self.assertEqual(
+            self.device_view.get_status("furnace_1"), DeviceTaskStatus.IDLE
+        )
         self.assertEqual(
             self.sample_view.get_sample_position_status("furnace_1/inside"),
             (SamplePositionStatus.EMPTY, None),