Release 0.14.0

* Fix several bugs
* Improve device cleanup to be fully parallel
* Improve db connection timeout error
* Update dependencies

Author: aangelos28

eos/campaigns/campaign_executor.py

@@ -69,7 +69,7 @@ async def start_campaign(self, db: AsyncDbSession) -> None:
             await self._initialize_new_campaign(db)
 
         await self._campaign_manager.start_campaign(db, self._campaign_name)
-        await db.flush()
+        await db.commit()
 
         self._campaign_status = CampaignStatus.RUNNING
         log.info(f"Started campaign '{self._campaign_name}'")

eos/configuration/validation.py

@@ -105,7 +105,7 @@ def _validate_device_types(self) -> None:
         for device_name, device in self._lab.devices.items():
             if not self._device_specs.get_spec_by_config(device):
                 batch_error(
-                    f"Device type '{device.name}' of device '{device_name}' does not exist.",
+                    f"Device type '{device.type}' of device '{device_name}' does not exist.",
                     EosLabConfigurationError,
                 )
         raise_batched_errors(EosLabConfigurationError)
@@ -120,7 +120,7 @@ def _validate_device_init_parameters(self) -> None:
                     if param_name not in spec_params:
                         batch_error(
                             f"Invalid initialization parameter '{param_name}' for device '{device_name}' "
-                            f"of type '{device.name}' in lab type '{self._lab.name}'. "
+                            f"of type '{device.type}' in lab type '{self._lab.name}'. "
                             f"Valid parameters are: {', '.join(spec_params.keys())}",
                             EosLabConfigurationError,
                         )

eos/database/exceptions.py

@@ -1,2 +1,6 @@
 class EosFileDbError(Exception):
     pass
+
+
+class DbConnectionError(Exception):
+    """Raised when database connection fails (e.g., timeout, connection refused)."""

eos/database/postgresql_db_interface.py

@@ -6,9 +6,11 @@
     AsyncEngine,
     create_async_engine,
 )
-from eos.logging.logger import log
+
 from eos.database.abstract_sql_db_interface import AbstractSqlDbInterface, Base
 from eos.database.alembic_commands import alembic_upgrade, alembic_downgrade
+from eos.database.exceptions import DbConnectionError
+from eos.logging.logger import log
 
 
 class PostgresqlDbInterface(AbstractSqlDbInterface):
@@ -72,9 +74,12 @@ def _create_system_engine(self) -> AsyncEngine:
     async def initialize_database(self) -> None:
         """Initialize database by creating it if needed and running migrations.
 
-        :raises Exception: If initialization fails
+        :raises DbConnectionError: If connection to database fails
+        :raises Exception: If initialization fails for other reasons
         """
-        log.info(f"Connecting to database at {self._db_config.postgres.host}:{self._db_config.postgres.port}...")
+        host = self._db_config.postgres.host
+        port = self._db_config.postgres.port
+        log.info(f"Connecting to database at {host}:{port}...")
         try:
             exists = await self._database_exists()
             if not exists:
@@ -87,6 +92,15 @@ async def initialize_database(self) -> None:
                 await conn.run_sync(Base.metadata.create_all)
 
             log.info(f"Connected to database '{self._db_name}'")
+        except TimeoutError as e:
+            raise DbConnectionError(
+                f"Connection to database at {host}:{port} timed out. "
+                "Please check that PostgreSQL is running and accessible."
+            ) from e
+        except OSError as e:
+            raise DbConnectionError(
+                f"Failed to connect to database at {host}:{port}. Please check that PostgreSQL is running. Error: {e}"
+            ) from e
         except Exception as e:
             log.error(f"Failed to initialize database: {e!s}")
             raise

eos/devices/device_manager.py

@@ -8,6 +8,7 @@
 
 from eos.configuration.configuration_manager import ConfigurationManager
 from eos.configuration.constants import EOS_COMPUTER_NAME
+from eos.allocation.entities.device_allocation import DeviceAllocationModel
 from eos.devices.entities.device import Device, DeviceStatus, DeviceModel
 from eos.devices.exceptions import EosDeviceStateError, EosDeviceInitializationError
 from eos.logging.batch_error_logger import batch_error, raise_batched_errors
@@ -125,15 +126,14 @@ async def reload_devices(self, db: AsyncDbSession, lab_name: str, device_names:
                 raise
 
         # Cleanup the specific device actors
-        reload_tasks = []
-
-        for device_name in device_names:
-            actor_name = f"{lab_name}.{device_name}"
-            if actor_name in self._device_actor_handles:
-                reload_tasks.append(self._cleanup_single_device(actor_name))
+        actors_to_cleanup = [
+            f"{lab_name}.{device_name}"
+            for device_name in device_names
+            if f"{lab_name}.{device_name}" in self._device_actor_handles
+        ]
 
-        if reload_tasks:
-            await asyncio.gather(*reload_tasks)
+        if actors_to_cleanup:
+            await self._cleanup_device_actors_with_timeout(actors_to_cleanup)
 
         # Remove device records from database
         await db.execute(
@@ -183,69 +183,96 @@ async def cleanup_device_actors(self, db: AsyncDbSession, lab_names: list[str] |
         if not actor_names:
             return
 
-        cleanup_tasks = [
-            self._cleanup_single_device(actor_name)
-            for actor_name in actor_names
-            if actor_name in self._device_actor_handles
-        ]
-
-        if cleanup_tasks:
-            await asyncio.gather(*cleanup_tasks)
+        actors_to_cleanup = [name for name in actor_names if name in self._device_actor_handles]
+        if actors_to_cleanup:
+            await self._cleanup_device_actors_with_timeout(actors_to_cleanup)
 
         await self.cleanup_devices(db, lab_names)
 
-    async def _get_actor_names_to_cleanup(self, db: AsyncDbSession, lab_names: list[str] | None) -> list[str]:
-        """Get actor names that need to be cleaned up."""
-        if not lab_names:
-            return list(self._device_actor_handles.keys())
+    async def _cleanup_device_actors_with_timeout(self, actor_names: list[str], cleanup_timeout: float = 30.0) -> None:
+        """Clean up multiple device actors concurrently with a timeout."""
+        # Start cleanup on all actors and collect the object refs
+        cleanup_refs: dict[ray.ObjectRef, str] = {}
+        for actor_name in actor_names:
+            actor_handle = self._device_actor_handles.get(actor_name)
+            if actor_handle is None:
+                continue
 
-        result = await db.execute(select(DeviceModel).where(DeviceModel.lab_name.in_(lab_names)))
-        devices = [Device.model_validate(device) for device in result.scalars()]
-        return [device.get_actor_name() for device in devices]
+            try:
+                log.info(f"Cleaning up device '{actor_name}'...")
+                cleanup_ref = actor_handle.cleanup.remote()
+                cleanup_refs[cleanup_ref] = actor_name
+            except Exception as e:
+                log.error(f"Failed to start cleanup for device '{actor_name}': {e}")
+                self._forcefully_kill_actor(actor_name)
 
-    async def _cleanup_single_device(self, actor_name: str) -> None:
-        """Clean up a single device actor with timeout.
+        if not cleanup_refs:
+            return
 
-        Attempts to gracefully clean up a device actor. If the cleanup
-        doesn't complete within 30 seconds, forcefully kills the actor.
+        # Process cleanups as they complete, with overall timeout
+        pending_refs = set(cleanup_refs.keys())
+        start_time = asyncio.get_event_loop().time()
 
-        :param actor_name: The name of the actor to clean up
-        """
-        if actor_name not in self._device_actor_handles:
-            return
+        while pending_refs:
+            elapsed = asyncio.get_event_loop().time() - start_time
+            remaining_timeout = max(0, cleanup_timeout - elapsed)
 
-        actor_handle = self._device_actor_handles[actor_name]
-        success = False
-        cleanup_timeout = 30.0
+            if remaining_timeout <= 0:
+                break
 
-        try:
-            log.info(f"Cleaning up device actor '{actor_name}'...")
-            cleanup_ref = actor_handle.cleanup.remote()
+            ready_refs, _ = ray.wait(
+                list(pending_refs),
+                num_returns=1,
+                timeout=remaining_timeout,
+            )
 
-            # Wait for cleanup to complete with timeout
-            ready_refs, _ = ray.wait([cleanup_ref], timeout=cleanup_timeout)
+            if not ready_refs:
+                # Timeout reached with no more completions
+                break
 
-            if cleanup_ref in ready_refs:
-                log.info(f"Cleaned up device actor '{actor_name}'")
-                success = True
-            else:
-                log.warning(
-                    f"Timed out cleaning up device actor '{actor_name}' after {cleanup_timeout} seconds, "
-                    f"will forcefully kill..."
-                )
-        except Exception as e:
-            log.error(f"Failed cleaning up device actor '{actor_name}': {e}")
-        finally:
-            # Kill if cleanup wasn't successful
-            if not success and actor_name in self._device_actor_handles:
+            for ref in ready_refs:
+                pending_refs.discard(ref)
+                actor_name = cleanup_refs[ref]
                 try:
-                    log.warning(f"Forcefully killing device actor '{actor_name}'")
-                    ray.kill(self._device_actor_handles[actor_name])
+                    ray.get(ref)  # Check for exceptions
+                    log.info(f"Cleaned up device '{actor_name}'")
                 except Exception as e:
-                    log.error(f"Error killing device actor '{actor_name}': {e}")
+                    log.error(f"Cleanup failed for device '{actor_name}': {e}")
+                    self._forcefully_kill_actor(actor_name)
+                finally:
+                    self._remove_device_references(actor_name)
+
+        # Forcefully kill actors that timed out
+        if pending_refs:
+            timed_out_actors = [cleanup_refs[ref] for ref in pending_refs]
+            log.warning(
+                f"Timed out cleaning up {len(timed_out_actors)} device(s) after {cleanup_timeout} seconds: "
+                f"{', '.join(timed_out_actors)}"
+            )
+            for ref in pending_refs:
+                actor_name = cleanup_refs[ref]
+                self._forcefully_kill_actor(actor_name)
+                self._remove_device_references(actor_name)
+
+    def _forcefully_kill_actor(self, actor_name: str) -> None:
+        """Forcefully kill a device actor."""
+        if actor_name not in self._device_actor_handles:
+            return
+
+        try:
+            log.warning(f"Forcefully killing device '{actor_name}'")
+            ray.kill(self._device_actor_handles[actor_name])
+        except Exception as e:
+            log.error(f"Error killing device '{actor_name}': {e}")
+
+    async def _get_actor_names_to_cleanup(self, db: AsyncDbSession, lab_names: list[str] | None) -> list[str]:
+        """Get actor names that need to be cleaned up."""
+        if not lab_names:
+            return list(self._device_actor_handles.keys())
 
-            # Clean up references regardless of success
-            self._remove_device_references(actor_name)
+        result = await db.execute(select(DeviceModel).where(DeviceModel.lab_name.in_(lab_names)))
+        devices = [Device.model_validate(device) for device in result.scalars()]
+        return [device.get_actor_name() for device in devices]
 
     def _remove_device_references(self, actor_name: str) -> None:
         """Remove device references from internal tracking dictionaries."""
@@ -255,9 +282,11 @@ def _remove_device_references(self, actor_name: str) -> None:
     async def cleanup_devices(self, db: AsyncDbSession, lab_names: list[str] | None = None) -> None:
         """Remove device records from the database."""
         if lab_names:
+            await db.execute(delete(DeviceAllocationModel).where(DeviceAllocationModel.lab_name.in_(lab_names)))
             await db.execute(delete(DeviceModel).where(DeviceModel.lab_name.in_(lab_names)))
             log.debug(f"Cleaned up devices for lab(s): {', '.join(lab_names)}")
         else:
+            await db.execute(delete(DeviceAllocationModel))
             await db.execute(delete(DeviceModel))
             log.debug("Cleaned up all devices")
 

eos/experiments/experiment_executor.py

@@ -71,6 +71,7 @@ async def start_experiment(self, db: AsyncDbSession) -> None:
             )
 
             await self._experiment_manager.start_experiment(db, self._experiment_name)
+            await db.commit()
             self._experiment_status = ExperimentStatus.RUNNING
 
             action = "Resumed" if self._experiment_submission.resume else "Started"

eos/orchestration/orchestrator.py

@@ -111,6 +111,11 @@ async def initialize(self) -> None:
         di.register(WorkSignal, self._work_signal)
 
         # State management ########################################
+        allocation_manager = AllocationManager()
+        async with db_interface.get_async_session() as db:
+            await allocation_manager.initialize(db)
+        di.register(AllocationManager, allocation_manager)
+
         device_manager = DeviceManager()
         async with db_interface.get_async_session() as db:
             await device_manager.cleanup_devices(db)
@@ -121,11 +126,6 @@ async def initialize(self) -> None:
             await resource_manager.initialize(db)
         di.register(ResourceManager, resource_manager)
 
-        allocation_manager = AllocationManager()
-        async with db_interface.get_async_session() as db:
-            await allocation_manager.initialize(db)
-        di.register(AllocationManager, allocation_manager)
-
         task_manager = TaskManager()
         di.register(TaskManager, task_manager)
 

eos/orchestration/services/lab_service.py

@@ -45,7 +45,7 @@ async def get_lab_devices(
             filtered_devices = {
                 name: device
                 for name, device in lab.devices.items()
-                if not task_device_types or device.name in task_device_types
+                if not task_device_types or device.type in task_device_types
             }
 
             if filtered_devices:

eos/orchestration/services/loading_service.py

@@ -55,7 +55,7 @@ async def reload_labs(self, db: AsyncDbSession, lab_types: set[str]) -> None:
         """Reload one or more labs in the orchestrator with updated device plugin code."""
         for lab_type in lab_types:
             lab = self._configuration_manager.package_manager.read_lab(lab_type)
-            device_types = {device.name for device in lab.devices.values()}
+            device_types = {device.type for device in lab.devices.values()}
             for device_type in device_types:
                 try:
                     self._configuration_manager.devices.reload_plugin(device_type)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "eos"
-version = "0.13.0"
+version = "0.14.0"
 description = "EOS is a comprehensive software framework and runtime for laboratory automation."
 keywords = ["automation", "science", "lab", "experiment", "orchestration", "distributed", "infrastructure"]
 authors = [
@@ -22,39 +22,39 @@ classifiers = [
 readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
-    "ray[default]==2.52.1",
-    "typer==0.20.0",
-    "rich==14.2.0",
-    "PyYAML==6.0.3",
-    "jinja2==3.1.6",
-    "pydantic==2.12.5",
-    "pydantic-settings==2.12.0",
-    "networkx==3.6.0",
-    "sqlalchemy==2.0.44",
-    "alembic==1.17.2",
-    "psycopg[binary]==3.3.1",
-    "asyncpg==0.31.0",
-    "minio==7.2.18",
-    "bofire[optimization]==0.3.0",
-    "pandas==2.3.3",
+    "ray[default]>=2.53.0",
+    "typer>=0.21.1",
+    "rich>=14.2.0",
+    "PyYAML>=6.0.3",
+    "jinja2>=3.1.6",
+    "pydantic>=2.12.5",
+    "pydantic-settings>=2.12.0",
+    "networkx>=3.6.1",
+    "sqlalchemy>=2.0.45",
+    "alembic>=1.18.0",
+    "psycopg[binary]>=3.3.2",
+    "asyncpg>=0.31.0",
+    "minio>=7.2.20",
+    "bofire[optimization]>=0.3.1",
+    "pandas>=2.3.3",
     "ortools<=9.12",
-    "litestar[standard]==2.18.0",
+    "litestar[standard]>=2.19.0",
 ]
 
 [dependency-groups]
 worker = [
-    "ray[default]==2.52.1",
-    "typer==0.20.0",
-    "rich==14.2.0",
-    "pydantic==2.12.5",
-    "pydantic-settings==2.12.0",
-    "sqlalchemy==2.0.44",
-    "alembic==1.17.2",
+    "ray[default]>=2.53.0",
+    "typer>=0.21.1",
+    "rich>=14.2.0",
+    "pydantic>=2.12.5",
+    "pydantic-settings>=2.12.0",
+    "sqlalchemy>=2.0.45",
+    "alembic>=1.18.0",
 ]
 optimizer_worker = [
     { include-group = "worker" },
-    "bofire[optimization]==0.3.0",
-    "pandas==2.3.3",
+    "bofire[optimization]>=0.3.1",
+    "pandas>=2.3.3",
 ]
 
 sila2 = [