Refactor and clean up parts of the code

eos/campaigns/campaign_optimizer_manager.py

@@ -3,7 +3,7 @@
 from ray.actor import ActorHandle
 
 from eos.campaigns.entities.campaign import CampaignSample
-from eos.configuration.plugin_registries.campaign_optimizer_plugin_registry import CampaignOptimizerPluginRegistry
+from eos.configuration.configuration_manager import ConfigurationManager
 from eos.logging.logger import log
 from eos.optimization.sequential_optimizer_actor import SequentialOptimizerActor
 from eos.persistence.db_manager import DbManager
@@ -15,11 +15,11 @@ class CampaignOptimizerManager:
     Responsible for managing the optimizers associated with experiment campaigns.
     """
 
-    def __init__(self, db_manager: DbManager):
+    def __init__(self, configuration_manager: ConfigurationManager, db_manager: DbManager):
         self._campaign_samples = MongoRepository("campaign_samples", db_manager)
         self._campaign_samples.create_indices([("campaign_id", 1), ("experiment_id", 1)], unique=True)
 
-        self._campaign_optimizer_plugin_registry = CampaignOptimizerPluginRegistry()
+        self._campaign_optimizer_plugin_registry = configuration_manager.campaign_optimizers
 
         self._optimizer_actors: dict[str, ActorHandle] = {}
 

eos/cli/orchestrator_cli.py

@@ -102,7 +102,7 @@ def signal_handler(*_) -> None:
         log.info("Shutting down the orchestrator...")
         orchestrator.terminate()
 
-        log.info("Shutdown complete.")
+        log.info("EOS shut down.")
 
 
 async def run_all(orchestrator: Orchestrator, web_api_server: uvicorn.Server) -> None:
@@ -155,6 +155,7 @@ def start_orchestrator(
     orchestrator = Orchestrator(config.user_dir, db_credentials, file_db_credentials)
     orchestrator.load_labs(config.labs)
     orchestrator.load_experiments(config.experiments)
+    log.info("EOS initialized.")
 
     # Set up the web API server
     logging_config = LoggingConfig(

eos/configuration/configuration_manager.py

@@ -3,7 +3,8 @@
 from eos.configuration.exceptions import (
     EosConfigurationError,
 )
-from eos.configuration.package_manager import PackageManager
+from eos.configuration.packages.entities import EntityType
+from eos.configuration.packages.package_manager import PackageManager
 from eos.configuration.plugin_registries.campaign_optimizer_plugin_registry import CampaignOptimizerPluginRegistry
 from eos.configuration.plugin_registries.device_plugin_registry import DevicePluginRegistry
 from eos.configuration.plugin_registries.task_plugin_registry import TaskPluginRegistry
@@ -55,7 +56,7 @@ def get_lab_loaded_statuses(self) -> dict[str, bool]:
         all_labs = set()
 
         for package in self._package_manager.get_all_packages():
-            package_labs = self._package_manager.get_labs_in_package(package.name)
+            package_labs = self._package_manager.get_entities_in_package(package.name, EntityType.LAB)
             all_labs.update(package_labs)
 
         return {lab: lab in self.labs for lab in all_labs}
@@ -128,7 +129,7 @@ def get_experiment_loaded_statuses(self) -> dict[str, bool]:
         all_experiments = set()
 
         for package in self._package_manager.get_all_packages():
-            package_experiments = self._package_manager.get_experiments_in_package(package.name)
+            package_experiments = self._package_manager.get_entities_in_package(package.name, EntityType.EXPERIMENT)
             all_experiments.update(package_experiments)
 
         return {exp: exp in self.experiments for exp in all_experiments}

eos/configuration/constants.py

@@ -1,8 +1,7 @@
 LABS_DIR = "labs"
 EXPERIMENTS_DIR = "experiments"
-TASKS_DIR = "tasks"
 DEVICES_DIR = "devices"
-COMMON_DIR = "common"
+TASKS_DIR = "tasks"
 
 EXPERIMENT_CONFIG_FILE_NAME = "experiment.yml"
 LAB_CONFIG_FILE_NAME = "lab.yml"

eos/configuration/exceptions.py

@@ -1,3 +1,6 @@
+from eos.devices.exceptions import EosDeviceError
+
+
 class EosConfigurationError(Exception):
     pass
 
@@ -30,9 +33,17 @@ class EosTaskGraphError(Exception):
     pass
 
 
-class EosTaskHandlerClassNotFoundError(Exception):
+class EosTaskImplementationClassNotFoundError(Exception):
+    pass
+
+
+class EosCampaignOptimizerImplementationClassNotFoundError(Exception):
     pass
 
 
 class EosCampaignOptimizerNotFoundError(Exception):
     pass
+
+
+class EosDeviceImplementationClassNotFoundError(EosDeviceError):
+    pass