WIP: Introduce TrialRunner Abstraction

mlos_bench/mlos_bench/config/schemas/cli/cli-schema.json

@@ -79,6 +79,13 @@
             "examples": [3, 5]
         },
 
+        "num_trial_runners": {
+            "description": "Number of trial runner instances to use to execute benchmark environments. Individual TrialRunners can be identified in configs with $trial_runner_id and optionally run in parallel.",
+            "type": "integer",
+            "minimum": 1,
+            "examples": [1, 3, 5, 10]
+        },
+
         "storage": {
             "description": "Path to the json config describing the storage backend to use.",
             "$ref": "#/$defs/json_config_path"

mlos_bench/mlos_bench/environments/base_environment.py

@@ -43,6 +43,15 @@ class Environment(metaclass=abc.ABCMeta):
     # pylint: disable=too-many-instance-attributes
     """An abstract base of all benchmark environments."""
 
+    # Should be provided by the runtime.
+    _COMMON_CONST_ARGS = {
+        "trial_runner_id",
+    }
+    _COMMON_REQ_ARGS = {
+        "experiment_id",
+        "trial_id",
+    }
+
     @classmethod
     def new(  # pylint: disable=too-many-arguments
         cls,
@@ -123,6 +132,12 @@ def __init__(  # pylint: disable=too-many-arguments
             An optional service object (e.g., providing methods to
             deploy or reboot a VM/Host, etc.).
         """
+        global_config = global_config or {}
+        # Make some usual runtime arguments available for tests.
+        for arg in self._COMMON_CONST_ARGS:
+            global_config.setdefault(arg, None)
+        for arg in self._COMMON_REQ_ARGS:
+            global_config.setdefault(arg, None)
         self._validate_json_config(config, name)
         self.name = name
         self.config = config
@@ -161,8 +176,9 @@ def __init__(  # pylint: disable=too-many-arguments
         req_args = set(config.get("required_args", [])) - set(
             self._tunable_params.get_param_values().keys()
         )
+        req_args.update(self._COMMON_CONST_ARGS)
         merge_parameters(dest=self._const_args, source=global_config, required_keys=req_args)
-        self._const_args = self._expand_vars(self._const_args, global_config or {})
+        self._const_args = self._expand_vars(self._const_args, global_config)
 
         self._params = self._combine_tunables(self._tunable_params)
         _LOG.debug("Parameters for '%s' :: %s", name, self._params)
@@ -332,6 +348,18 @@ def tunable_params(self) -> TunableGroups:
         """
         return self._tunable_params
 
+    @property
+    def const_args(self) -> Dict[str, TunableValue]:
+        """
+        Get the constant arguments for this Environment.
+
+        Returns
+        -------
+        parameters : Dict[str, TunableValue]
+            Key/value pairs of all environment const_args parameters.
+        """
+        return self._const_args.copy()
+
     @property
     def parameters(self) -> Dict[str, TunableValue]:
         """
@@ -345,7 +373,7 @@ def parameters(self) -> Dict[str, TunableValue]:
             Key/value pairs of all environment parameters
             (i.e., `const_args` and `tunable_params`).
         """
-        return self._params
+        return self._params.copy()
 
     def setup(self, tunables: TunableGroups, global_config: Optional[dict] = None) -> bool:
         """

mlos_bench/mlos_bench/launcher.py

@@ -22,6 +22,7 @@
 from mlos_bench.optimizers.mock_optimizer import MockOptimizer
 from mlos_bench.optimizers.one_shot_optimizer import OneShotOptimizer
 from mlos_bench.schedulers.base_scheduler import Scheduler
+from mlos_bench.schedulers.trial_runner import TrialRunner
 from mlos_bench.services.base_service import Service
 from mlos_bench.services.config_persistence import ConfigPersistenceService
 from mlos_bench.services.local.local_exec import LocalExecService
@@ -44,6 +45,7 @@ class Launcher:
 
     def __init__(self, description: str, long_text: str = "", argv: Optional[List[str]] = None):
         # pylint: disable=too-many-statements
+        # pylint: disable=too-complex
         # pylint: disable=too-many-locals
         _LOG.info("Launch: %s", description)
         epilog = """
@@ -108,6 +110,7 @@ def __init__(self, description: str, long_text: str = "", argv: Optional[List[st
             args_rest=args_rest,
             global_config=cli_config_args,
         )
+        # TODO: Can we generalize these two rules using excluded_cli_args?
         # experiment_id is generally taken from --globals files, but we also allow
         # overriding it on the CLI.
         # It's useful to keep it there explicitly mostly for the --help output.
@@ -117,6 +120,13 @@ def __init__(self, description: str, long_text: str = "", argv: Optional[List[st
         # set it via command line
         if args.trial_config_repeat_count:
             self.global_config["trial_config_repeat_count"] = args.trial_config_repeat_count
+        self.global_config.setdefault("num_trial_runners", 1)
+        if args.num_trial_runners:
+            self.global_config["num_trial_runners"] = args.num_trial_runners
+        if self.global_config["num_trial_runners"] <= 0:
+            raise ValueError(
+                f"Invalid num_trial_runners: {self.global_config['num_trial_runners']}"
+            )
         # Ensure that the trial_id is present since it gets used by some other
         # configs but is typically controlled by the run optimize loop.
         self.global_config.setdefault("trial_id", 1)
@@ -142,13 +152,28 @@ def __init__(self, description: str, long_text: str = "", argv: Optional[List[st
             )
         self.root_env_config = self._config_loader.resolve_path(env_path)
 
-        self.environment: Environment = self._config_loader.load_environment(
-            self.root_env_config, TunableGroups(), self.global_config, service=self._parent_service
+        self.trial_runners: List[TrialRunner] = []
+        for trial_runner_id in range(self.global_config["num_trial_runners"]):
+            # Create a new global config for each Environment with a unique trial_runner_id for it.
+            env_global_config = self.global_config.copy()
+            env_global_config["trial_runner_id"] = trial_runner_id
+            env = self._config_loader.load_environment(
+                self.root_env_config,
+                TunableGroups(),
+                env_global_config,
+                service=self._parent_service,
+            )
+            self.trial_runners.append(TrialRunner(trial_runner_id, env))
+        _LOG.info(
+            "Init %d trial runners for environments: %s",
+            len(self.trial_runners),
+            list(trial_runner.environment for trial_runner in self.trial_runners),
         )
-        _LOG.info("Init environment: %s", self.environment)
 
-        # NOTE: Init tunable values *after* the Environment, but *before* the Optimizer
+        # NOTE: Init tunable values *after* the Environment(s), but *before* the Optimizer
+        # TODO: should we assign the same or different tunables for all TrialRunner Environments?
         self.tunables = self._init_tunable_values(
+            self.trial_runners[0].environment,
             args.random_init or config.get("random_init", False),
             config.get("random_seed") if args.random_seed is None else args.random_seed,
             config.get("tunable_values", []) + (args.tunable_values or []),
@@ -278,6 +303,18 @@ def add_argument(self, *args: Any, **kwargs: Any) -> None:
             ),
         )
 
+        parser.add_argument(
+            "--num_trial_runners",
+            "--num-trial-runners",
+            required=False,
+            type=int,
+            help=(
+                "Number of TrialRunners to use for executing benchmark Environments. "
+                "Individual TrialRunners can be identified in configs with $trial_runner_id "
+                "and optionally run in parallel."
+            ),
+        )
+
         path_args_tracker.add_argument(
             "--scheduler",
             required=False,
@@ -428,14 +465,15 @@ def _load_config(
 
     def _init_tunable_values(
         self,
+        env: Environment,
         random_init: bool,
         seed: Optional[int],
         args_tunables: Optional[str],
     ) -> TunableGroups:
         """Initialize the tunables and load key/value pairs of the tunable values from
         given JSON files, if specified.
         """
-        tunables = self.environment.tunable_params
+        tunables = env.tunable_params
         _LOG.debug("Init tunables: default = %s", tunables)
 
         if random_init:
@@ -534,7 +572,7 @@ def _load_scheduler(self, args_scheduler: Optional[str]) -> Scheduler:
                     "teardown": self.teardown,
                 },
                 global_config=self.global_config,
-                environment=self.environment,
+                trial_runners=self.trial_runners,
                 optimizer=self.optimizer,
                 storage=self.storage,
                 root_env_config=self.root_env_config,
@@ -544,7 +582,7 @@ def _load_scheduler(self, args_scheduler: Optional[str]) -> Scheduler:
         return self._config_loader.build_scheduler(
             config=class_config,
             global_config=self.global_config,
-            environment=self.environment,
+            trial_runners=self.trial_runners,
             optimizer=self.optimizer,
             storage=self.storage,
             root_env_config=self.root_env_config,