Device and Sample Position Create, Update, and Delete Auto Refresh

.pylintrc

@@ -64,7 +64,7 @@ confidence=
 disable=fixme,
         no-name-in-module,
         missing-function-docstring,
-        too-many-arguments,
+        too-many-positional-arguments,
         import-error,
         missing-module-docstring,
         too-few-public-methods,
@@ -141,7 +141,7 @@ disable=fixme,
         ; xreadlines-attribute,
         missing-timeout,
         broad-exception-caught,
-        too-many-positional-arguments,
+        too-many-arguments,
         broad-exception-raised,
         logging-fstring-interpolation,
         global-statement,

alab_management/device_view/__init__.py

@@ -1,5 +1,10 @@
 """Things related to the devices."""
 
 from .dbattributes import DictInDatabase, ListInDatabase, value_in_database
-from .device import BaseDevice, add_device, get_all_devices
+from .device import (
+    BaseDevice,
+    add_device,
+    get_all_devices,
+    reset_current_device_registry,
+)
 from .device_view import DeviceTaskStatus, DeviceView

alab_management/device_view/device.py

@@ -618,14 +618,24 @@ def retrieve_signal(self, signal_name, within: datetime.timedelta | None = None)
             )
 
 
+# _device_registry is used to store all the devices that are defined in the __init__.py when setting up the lab
+# it is only pooled with new devices that are added to the lab during reload.
 _device_registry: dict[str, BaseDevice] = {}
 
+# _current_device_registry is used to store all the devices that are defined in the __init__.py
+# when alabos setup is called, regardless of whether it is a reload or not.
+# this is used to check if a device is still in the lab during reload.
+# if not, it will be removed from the lab once unoccupied and samples positions that are related to it are also not occupied.
+# this is not visible to the device_manager
+_current_device_registry: dict[str, BaseDevice] = {}
+
 
 def add_device(device: BaseDevice):
     """Register a device instance. It is stored in a global dictionary."""
     if device.name in _device_registry and not os.environ.get("ALABOS_RELOAD", None):
         raise KeyError(f"Duplicated device name {device.name}")
     _device_registry[device.name] = device
+    _current_device_registry[device.name] = device
 
 
 def get_all_devices() -> dict[str, BaseDevice]:
@@ -638,3 +648,19 @@ def get_all_devices() -> dict[str, BaseDevice]:
           and the values are the device instances.
     """
     return _device_registry.copy()
+
+
+def get_current_devices() -> dict[str, BaseDevice]:
+    """Get all the current devices in the device registry."""
+    return _current_device_registry.copy()
+
+
+def remove_device(device_name: str):
+    """Remove a device from the device registry."""
+    _device_registry.pop(device_name, None)
+    _current_device_registry.pop(device_name, None)
+
+
+def reset_current_device_registry():
+    """Reset the current device registry."""
+    _current_device_registry.clear()

alab_management/device_view/device_view.py

@@ -12,7 +12,7 @@
 from alab_management.sample_view import SamplePosition, SampleView
 from alab_management.utils.data_objects import get_collection, get_lock
 
-from .device import BaseDevice, get_all_devices
+from .device import BaseDevice, get_all_devices, remove_device
 
 _DeviceType = TypeVar("_DeviceType", bound=BaseDevice)  # pylint: disable=invalid-name
 
@@ -110,15 +110,17 @@ def sync_device_status(self):
                 task_id=None,
             )
 
-    def add_devices_to_db(self):
+    def add_devices_to_db(self, devices: dict[str, BaseDevice] | None = None):
         """
         Insert device definitions to db, which includes devices' name, descriptions, parameters,
         type (class name).
 
         When one device's name has already appeared in the database, a ``NameError`` will be raised.
         Device name is a unique identifier for a device
         """
-        for device in self._device_list.values():
+        for device in (
+            devices.values() if devices is not None else self._device_list.values()
+        ):
             if self._device_collection.find_one({"name": device.name}) is not None:
                 raise NameError(
                     f"Duplicated device name {device.name}, did you cleanup the database?"
@@ -604,3 +606,54 @@ def __del__(self):
         """Ensure that we disconnect from all devices when the object is deleted."""
         if self.__connected_to_devices:
             self.__disconnect_all_devices()
+
+    def remove_device(self, device_name: str):
+        """Remove a device from the device view and the registry."""
+        # disconnect the device
+        print(f"Disconnecting from {device_name}...")
+        with self._lock():  # pylint: disable=not-callable
+            device = self._device_list[device_name]
+            device._disconnect_wrapper()
+            self._device_list.pop(device_name, None)
+            remove_device(device_name)
+            self._device_collection.delete_one({"name": device_name})
+            print(
+                f"Device {device_name} has been removed from the device view and the registry."
+            )
+
+    def get_all_devices_from_db(self) -> dict[str, dict[str, Any]]:
+        """
+        Get all devices from the database directly.
+
+        Returns a dictionary mapping device names to their database entries.
+        """
+        devices = {}
+        for device_doc in self._device_collection.find():
+            devices[device_doc["name"]] = device_doc
+        return devices
+
+    def get_device_by_name_from_db(self, device_name: str) -> dict[str, Any] | None:
+        """
+        Get a specific device from the database by name.
+
+        Args:
+            device_name: The name of the device to find
+
+        Returns
+        -------
+            The device document from the database, or None if not found
+        """
+        return self._device_collection.find_one({"name": device_name})
+
+    def device_exists_in_db(self, device_name: str) -> bool:
+        """
+        Check if a device exists in the database.
+
+        Args:
+            device_name: The name of the device to check
+
+        Returns
+        -------
+            True if the device exists in the database, False otherwise
+        """
+        return self._device_collection.count_documents({"name": device_name}) > 0

alab_management/experiment_manager.py

@@ -6,16 +6,24 @@
 done.
 """
 
+import logging
 import time
+from contextlib import contextmanager
 from typing import Any
 
+from alab_management.utils.logger import set_up_rich_handler
+from alab_management.utils.module_ops import load_definition
+
 from .config import AlabOSConfig
 from .experiment_view import CompletedExperimentView, ExperimentStatus, ExperimentView
 from .logger import DBLogger
 from .sample_view import SampleView
 from .task_view import TaskStatus, TaskView
 from .utils.graph_ops import Graph
 
+cli_logger = logging.getLogger(__name__)
+set_up_rich_handler(cli_logger)
+
 
 class ExperimentManager:
     """
@@ -35,6 +43,7 @@ def __init__(self):
         )  # if this is not defined in the config, assume it this feature is not being used.
         if self.__copy_to_completed_db:
             self.completed_experiment_view = CompletedExperimentView()
+        self._pause_handling_experiments = False
 
     def run(self):
         """Start the event loop."""
@@ -50,8 +59,9 @@ def run(self):
             time.sleep(1)
 
     def _loop(self):
-        self.handle_pending_experiments()
-        self.mark_completed_experiments()
+        if not self._pause_handling_experiments:
+            self.handle_pending_experiments()
+            self.mark_completed_experiments()
 
     def handle_pending_experiments(self):
         """
@@ -194,3 +204,21 @@ def mark_completed_experiments(self):
                             "exp_id": experiment["_id"],
                         },
                     )
+
+    @contextmanager
+    def pause_handling_experiments(self):
+        """This method will pause the handle pending experiments."""
+        try:
+            self._pause_handling_experiments = True
+            cli_logger.info("Pausing handling experiments.")
+            yield
+        finally:
+            self._pause_handling_experiments = False
+            cli_logger.info("Resuming handling experiments.")
+
+    def refresh_task_list(self):
+        """This method will refresh the task list by reloading definition of tasks."""
+        print("Refreshing task view in ExperimentManager...")
+        load_definition(reload=True)
+        self.task_view = TaskView()
+        self.experiment_view = ExperimentView()

alab_management/sample_view/__init__.py

@@ -1,5 +1,10 @@
 """Sample related things."""
 
 from .completed_sample_view import CompletedSampleView
-from .sample import Sample, SamplePosition, add_standalone_sample_position
+from .sample import (
+    Sample,
+    SamplePosition,
+    add_standalone_sample_position,
+    reset_current_standalone_sample_position_registry,
+)
 from .sample_view import SampleView

alab_management/sample_view/sample.py

@@ -55,8 +55,19 @@ def __post_init__(self):
             )
 
 
+# _standalone_sample_position_registry is used to store all the standalone sample positions that are defined in the __init__.py
+# when setting up the lab.
+# it is only pooled with new sample positions that are added to the lab during reload.
 _standalone_sample_position_registry: dict[str, SamplePosition] = {}
 
+# _current_standalone_sample_position_registry is used to store
+# all the standalone sample positions that are defined in the __init__.py
+# when alabos setup is called, regardless of whether it is a reload or not.
+# this is used to check if a sample position is still in the lab during reload.
+# if not, it will be removed from the lab once unoccupied and devices that are related to it are also not occupied.
+# this is not visible to the sample_view
+_current_standalone_sample_position_registry: dict[str, SamplePosition] = {}
+
 
 def add_standalone_sample_position(position: SamplePosition):
     """Register a device instance."""
@@ -69,8 +80,25 @@ def add_standalone_sample_position(position: SamplePosition):
     ):
         raise KeyError(f"Duplicated standalone sample position name {position.name}")
     _standalone_sample_position_registry[position.name] = position
+    _current_standalone_sample_position_registry[position.name] = position
 
 
 def get_all_standalone_sample_positions() -> dict[str, SamplePosition]:
-    """Get all the device names in the device registry."""
+    """Get all the sample position names in the registry."""
     return _standalone_sample_position_registry.copy()
+
+
+def get_current_standalone_sample_positions() -> dict[str, SamplePosition]:
+    """Get all the current standalone sample positions in the registry."""
+    return _current_standalone_sample_position_registry.copy()
+
+
+def remove_standalone_sample_position(prefix: str):
+    """Remove a standalone sample position from the registry."""
+    _standalone_sample_position_registry.pop(prefix, None)
+    _current_standalone_sample_position_registry.pop(prefix, None)
+
+
+def reset_current_standalone_sample_position_registry():
+    """Reset the current standalone sample position registry."""
+    _current_standalone_sample_position_registry.clear()