From 8d36d164df5b397d014fdeb6e2f27473ce81f61d Mon Sep 17 00:00:00 2001
From: JMGaljaard <J.M.Galjaard-1@student.tudelft.nl>
Date: Fri, 2 Sep 2022 12:59:36 +0200
Subject: [PATCH] Update definitions in preparation for restructured
 configuration objects

---
 fltk/util/config/definitions/__init__.py     |  6 +++
 fltk/util/config/definitions/data_sampler.py |  2 +-
 fltk/util/config/definitions/orchestrator.py | 39 ++++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 fltk/util/config/definitions/orchestrator.py

diff --git a/fltk/util/config/definitions/__init__.py b/fltk/util/config/definitions/__init__.py
index f62fd226..9cc6d386 100644
--- a/fltk/util/config/definitions/__init__.py
+++ b/fltk/util/config/definitions/__init__.py
@@ -1,3 +1,8 @@
+"""
+Module for declaring types and definitions, including helper functions that allow to retrieve
+object (types) from a definition.
+"""
+
 from .data_sampler import DataSampler
 from .optim import Optimizations
 from .aggregate import Aggregations
@@ -7,3 +12,4 @@
 from .optim import Optimizations
 from .experiment_type import ExperimentType
 from .loss import Loss, get_loss_function
+from .orchestrator import OrchestratorType, get_orchestrator
diff --git a/fltk/util/config/definitions/data_sampler.py b/fltk/util/config/definitions/data_sampler.py
index 3d8a254a..46d8acbf 100644
--- a/fltk/util/config/definitions/data_sampler.py
+++ b/fltk/util/config/definitions/data_sampler.py
@@ -10,4 +10,4 @@ class DataSampler(Enum):
     limit_labels_q = "limit labels q"
     emd_sampler = 'emd sampler'
     limit_labels_flex = "limit labels flex"
-    n_labels = "n labels"
\ No newline at end of file
+    n_labels = "n labels"
diff --git a/fltk/util/config/definitions/orchestrator.py b/fltk/util/config/definitions/orchestrator.py
new file mode 100644
index 00000000..3ec0ada2
--- /dev/null
+++ b/fltk/util/config/definitions/orchestrator.py
@@ -0,0 +1,39 @@
+from __future__ import annotations
+from enum import unique, Enum
+from pathlib import Path
+from typing import Type
+
+from fltk.core.distributed import BatchOrchestrator, SimulatedOrchestrator, Orchestrator
+
+from typing import TYPE_CHECKING
+
+from fltk.util.task.generator import ArrivalGenerator
+
+if TYPE_CHECKING:
+    from fltk.util.cluster import ClusterManager
+    from fltk.util.config import DistributedConfig
+
+
+@unique
+class OrchestratorType(Enum):
+    BATCH = 'batch'
+    SIMULATED = 'simulated'
+
+
+def get_orchestrator(config: DistributedConfig, cluster_manager: ClusterManager, arrival_generator: ArrivalGenerator) -> Orchestrator:
+    """
+    Retrieve Orchestrator type given a Distributed (experiment) configuration. This allows for defining the
+    type of experiment (Batch or Simulated arrivals) once, and letting the Orchestrator implementation
+    make sure that the tasks are scheduled correctly.
+    @param config: Distributed (cluster) configuration object for experiments.
+    @type config: DistributedConfig
+    @return: Type of Orchestrator as requested by configuration object.
+    @rtype: Type[Orchestrator]
+    """
+    __lookup = {
+        OrchestratorType.BATCH: BatchOrchestrator,
+        OrchestratorType.SIMULATED: SimulatedOrchestrator
+    }
+
+    orchestrator_type = __lookup.get(config.cluster_config.orchestrator.orchestrator_type, None)
+    return orchestrator_type(cluster_manager, arrival_generator, config)
\ No newline at end of file