Introduce mock_trial_data for MockEnv and add some basic Scheduler tests

mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json

@@ -58,7 +58,7 @@
             "description": "The name of the scheduler class to use.",
             "type": "string",
             "$comment": "Exact matches are handled elsewhere.",
-            "pattern": "^mlos_bench[.]schedulers[.]"
+            "pattern": "^mlos_bench([.]tests)?[.]schedulers[.]"
         },
 
         "config": {

mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json

@@ -0,0 +1,83 @@
+{
+    "$schema": "https://json-schema.org/draft/2020-12/schema",
+    "$id": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json",
+    "title": "mlos_bench Mock Scheduler config",
+    "description": "config for an mlos_bench Mock Scheduler",
+    "type": "object",
+    "properties": {
+        "class": {
+            "enum": [
+                "mlos_bench.tests.schedulers.mock_scheduler.MockScheduler"
+            ]
+        },
+        "config": {
+            "type": "object",
+            "allOf": [
+                {
+                    "$ref": "base-scheduler-subschema.json#/$defs/base_scheduler_config"
+                },
+                {
+                    "type": "object",
+                    "properties": {
+                        "mock_trial_data": {
+                            "description": "A list of trial data to use for testing.",
+                            "type": "array",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "comments": {
+                                        "type": "string",
+                                        "description": "Optional comments about the trial status being reported."
+                                    },
+                                    "trial_id": {
+                                        "type": "integer",
+                                        "description": "Unique identifier for the trial.",
+                                        "examples": [1, 2, 3],
+                                        "minimum": 1
+                                    },
+                                    "status": {
+                                        "enum": [
+                                            null,
+                                            "UNKNOWN",
+                                            "PENDING",
+                                            "READY",
+                                            "RUNNING",
+                                            "SUCCEEDED",
+                                            "CANCELED",
+                                            "FAILED",
+                                            "TIMED_OUT"
+                                        ]
+                                    },
+                                    "metrics": {
+                                        "type": "object",
+                                        "description": "A dictionary of metrics for the trial.",
+                                        "additionalProperties": {
+                                            "type": ["number", "string", "boolean"],
+                                            "description": "The value of the metric."
+                                        },
+                                        "examples": [
+                                            {
+                                                "score": 0.95,
+                                                "color": "green"
+                                            },
+                                            {
+                                                "accuracy": 0.85,
+                                                "loss": 0.15
+                                            }
+                                        ]
+                                    }
+                                },
+                                "required": ["trial_id", "status"],
+                                "additionalProperties": false
+                            }
+                        }
+                    },
+                    "minProperties": 1
+                }
+            ],
+            "minProperties": 1,
+            "unevaluatedProperties": false
+        }
+    },
+    "required": ["class"]
+}

mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json

@@ -14,6 +14,9 @@
         {
             "$comment": "The set of known Scheduler subschemas.  Add others as needed.",
             "oneOf": [
+                {
+                    "$ref": "./mock-scheduler-subschema.json"
+                },
                 {
                     "$ref": "./sync-scheduler-subschema.json"
                 }

mlos_bench/mlos_bench/schedulers/base_scheduler.py

@@ -242,8 +242,8 @@ def __exit__(
         self._in_context = False
         return False  # Do not suppress exceptions
 
-    def start(self) -> None:
-        """Start the scheduling loop."""
+    def _prepare_start(self) -> bool:
+        """Prepare the scheduler for starting."""
         assert self.experiment is not None
         _LOG.info(
             "START: Experiment: %s Env: %s Optimizer: %s",
@@ -262,21 +262,39 @@ def start(self) -> None:
         is_warm_up: bool = self.optimizer.supports_preload
         if not is_warm_up:
             _LOG.warning("Skip pending trials and warm-up: %s", self.optimizer)
+        return is_warm_up
 
+    def start(self) -> None:
+        """Start the scheduling loop."""
+        assert self.experiment is not None
+        is_warm_up = self._prepare_start()
         not_done: bool = True
         while not_done:
-            _LOG.info("Optimization loop: Last trial ID: %d", self._last_trial_id)
-            self.run_schedule(is_warm_up)
-            not_done = self.add_new_optimizer_suggestions()
-            self.assign_trial_runners(
-                self.experiment.pending_trials(
-                    datetime.now(UTC),
-                    running=False,
-                    trial_runner_assigned=False,
-                )
-            )
+            not_done = self._execute_scheduling_step(is_warm_up)
             is_warm_up = False
 
+    def _execute_scheduling_step(self, is_warm_up: bool) -> bool:
+        """
+        Perform a single scheduling step.
+
+        Notes
+        -----
+        This method is called by the :py:meth:`Scheduler.start` method.
+        It is split out mostly to allow for easier testing with MockSchedulers.
+        """
+        assert self.experiment is not None
+        _LOG.info("Optimization loop: Last trial ID: %d", self._last_trial_id)
+        self.run_schedule(is_warm_up)
+        not_done = self.add_new_optimizer_suggestions()
+        self.assign_trial_runners(
+            self.experiment.pending_trials(
+                datetime.now(UTC),
+                running=False,
+                trial_runner_assigned=False,
+            )
+        )
+        return not_done
+
     def teardown(self) -> None:
         """
         Tear down the TrialRunners/Environment(s).

mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py

@@ -19,6 +19,7 @@
 from mlos_bench.util import get_class_from_name
 
 mock_opt = mlos_bench.tests.optimizers.fixtures.mock_opt
+
 storage = mlos_bench.tests.storage.sql.fixtures.storage
 
 

mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data-2.jsonc

@@ -0,0 +1,11 @@
+{
+    "class": "mlos_bench.tests.schedulers.mock_scheduler.MockScheduler",
+    "config": {
+        "mock_trial_data": [
+            {
+                // MISSING: "trial_id": 1,
+                "status": "SUCCEEDED"
+            }
+        ]
+    }
+}

mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data.jsonc

@@ -0,0 +1,14 @@
+{
+    "class": "mlos_bench.tests.schedulers.mock_scheduler.MockScheduler",
+    "config": {
+        "mock_trial_data": [
+            {
+                "trial_id": 1,
+                "status": "INVALID"
+            },
+            {
+                "status": "SUCCEEDED"
+            }
+        ]
+    }
+}

mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/unhandled/mock_sched-extra.jsonc

@@ -0,0 +1,7 @@
+{
+    "class": "mlos_bench.tests.schedulers.mock_scheduler.MockScheduler",
+    "config": {
+        "extra": "unsupported"
+    },
+    "extra": "unsupported"
+}

mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/full/mock_sched-full.jsonc

@@ -0,0 +1,31 @@
+{
+    "$schema": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json",
+    "class": "mlos_bench.tests.schedulers.mock_scheduler.MockScheduler",
+    "config": {
+        "trial_config_repeat_count": 3,
+        "teardown": false,
+        "experiment_id": "MyExperimentName",
+        "config_id": 1,
+        "trial_id": 1,
+        "max_trials": 2,
+
+        "mock_trial_data": [
+            {
+                "status": "SUCCEEDED",
+                "trial_id": 1,
+                "metrics": {
+                    "score": 0.9,
+                    "color": "green"
+                }
+            },
+            {
+                "status": "FAILED",
+                "trial_id": 2,
+                "metrics": {
+                    "score": 0.1,
+                    "color": "red"
+                }
+            }
+        ]
+    }
+}

mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/partial/mock_sched-partial.jsonc

@@ -0,0 +1,20 @@
+{
+    "class": "mlos_bench.tests.schedulers.mock_scheduler.MockScheduler",
+    "config": {
+        "mock_trial_data": [
+            {
+                "status": "SUCCEEDED",
+                "trial_id": 1,
+                "metrics": {
+                    "score": 0.9,
+                    "color": "green"
+                }
+            },
+            {
+                "status": "FAILED",
+                "trial_id": 2
+                // missing metrics - OK
+            }
+        ]
+    }
+}

mlos_bench/mlos_bench/tests/schedulers/__init__.py

@@ -0,0 +1,10 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""mlos_bench.tests.schedulers."""
+from mlos_bench.tests.schedulers.mock_scheduler import MockScheduler
+
+__all__ = [
+    "MockScheduler",
+]

mlos_bench/mlos_bench/tests/schedulers/conftest.py

@@ -0,0 +1,5 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""Pytest fixtures for mlos_bench.schedulers tests."""

mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py

@@ -0,0 +1,65 @@
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+#
+"""A mock scheduler that returns predefined status and score for specific trial IDs."""
+
+import logging
+from collections.abc import Iterable
+from datetime import datetime
+from typing import Any
+
+from pytz import UTC
+
+from mlos_bench.schedulers.base_scheduler import Optimizer, Scheduler
+from mlos_bench.schedulers.trial_runner import TrialRunner
+from mlos_bench.storage.base_storage import Storage
+
+_LOG = logging.getLogger(__name__)
+
+
+class MockScheduler(Scheduler):
+    """A mock scheduler that returns predefined status and score for specific trial
+    IDs.
+    """
+
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        *,
+        config: dict[str, Any],
+        global_config: dict[str, Any],
+        trial_runners: Iterable[TrialRunner],
+        optimizer: Optimizer,
+        storage: Storage,
+        root_env_config: str,
+    ) -> None:
+        super().__init__(
+            config=config,
+            global_config=global_config,
+            trial_runners=trial_runners,
+            optimizer=optimizer,
+            storage=storage,
+            root_env_config=root_env_config,
+        )
+        self._mock_trial_data = config.get("mock_trial_data", [])
+        self._mock_trial_data = {
+            trial_info["trial_id"]: trial_info for trial_info in self._mock_trial_data
+        }
+
+    def run_trial(self, trial: Storage.Trial) -> None:
+        """
+        Mock the execution of a trial.
+
+        Parameters
+        ----------
+        trial : Storage.Trial
+            The trial to be executed.
+        """
+        trial_id = trial.trial_id
+        if trial_id not in self._mock_trial_data:
+            raise ValueError(f"Trial ID {trial_id} not found in mock trial data.")
+
+        trial_info = self._mock_trial_data[trial_id]
+        _LOG.info("Running trial %d: %s", trial_id, trial_info)
+        # Don't run it - just update the status and optionally score.
+        trial.update(trial_info["status"], datetime.now(UTC), trial_info.get("score"))