WIP: Prepare Experiment.load to handle async out of order trial completion

.vscode/settings.json

@@ -8,6 +8,7 @@
         "**/__pycache__/": true,
         "**/node_modules/": true,
         "**/*.egg-info": true,
+        "mlos_*/build/": true,
         "doc/source/autoapi/": true,
         "doc/build/doctrees/": true,
         "doc/build/html/": true,

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

@@ -24,6 +24,9 @@
         },
         "optimization_targets": {
             "$ref": "./common-defs-subschemas.json#/$defs/optimization_targets"
+        },
+        "mock_trial_data": {
+            "$ref": "../environments/mock-env-subschema.json#/$defs/mock_trial_data"
         }
     },
     "additionalProperties": {

mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json

@@ -3,6 +3,136 @@
     "$id": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json",
     "title": "mlos_bench MockEnv config",
     "description": "Config instance for a mlos_bench MockEnv",
+
+    "$defs": {
+        "mock_trial_common_phase_data": {
+            "type": "object",
+            "properties": {
+                "sleep": {
+                    "type": "number",
+                    "description": "Optional time to sleep (in seconds) before returning from this phase of the trial.",
+                    "examples": [0, 0.1, 0.5, 1, 2],
+                    "minimum": 0,
+                    "maximum": 60
+                },
+                "exception": {
+                    "type": "string",
+                    "description": "Optional exception message to raise during phase."
+                }
+            }
+        },
+        "mock_trial_status_run_phase_data": {
+            "type": "object",
+            "properties": {
+                "status": {
+                    "description": "The status to report for this phase of the trial.  Default is phase dependent.",
+                    "enum": [
+                        "UNKNOWN",
+                        "PENDING",
+                        "READY",
+                        "RUNNING",
+                        "SUCCEEDED",
+                        "CANCELED",
+                        "FAILED",
+                        "TIMED_OUT"
+                    ]
+                },
+                "metrics": {
+                    "type": "object",
+                    "description": "A dictionary of metrics for this phase of 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
+                        }
+                    ],
+                    "minProperties": 0
+                }
+            }
+        },
+        "mock_trial_data_item": {
+            "description": "Mock data for a single trial, split by phase",
+            "type": "object",
+            "properties": {
+                "run": {
+                    "description": "A dictionary of trial data for the run phase.",
+                    "type": "object",
+                    "allOf": [
+                        {
+                            "$ref": "#/$defs/mock_trial_common_phase_data"
+                        },
+                        {
+                            "$ref": "#/$defs/mock_trial_status_run_phase_data"
+                        }
+                    ],
+                    "minProperties": 1,
+                    "unevaluatedProperties": false
+                },
+                "status": {
+                    "description": "A dictionary of trial data for the status phase.",
+                    "type": "object",
+                    "allOf": [
+                        {
+                            "$ref": "#/$defs/mock_trial_common_phase_data"
+                        },
+                        {
+                            "$ref": "#/$defs/mock_trial_status_run_phase_data"
+                        }
+                    ],
+                    "minProperties": 1,
+                    "unevaluatedProperties": false
+                },
+                "setup": {
+                    "description": "A dictionary of trial data for the setup phase.",
+                    "type": "object",
+                    "allOf": [
+                        {
+                            "$ref": "#/$defs/mock_trial_common_phase_data"
+                        }
+                    ],
+                    "minProperties": 1,
+                    "unevaluatedProperties": false
+                },
+                "teardown": {
+                    "description": "A dictionary of trial data for the teardown phase.",
+                    "type": "object",
+                    "allOf": [
+                        {
+                            "$ref": "#/$defs/mock_trial_common_phase_data"
+                        }
+                    ],
+                    "minProperties": 1,
+                    "unevaluatedProperties": false
+                }
+            },
+            "unevaluatedProperties": false,
+            "minProperties": 1
+        },
+        "mock_trial_data": {
+            "description": "A set of mock trial data to use for testing, keyed by trial id. Used by MockEnv.",
+            "type": "object",
+            "patternProperties": {
+                "^[1-9][0-9]*$": {
+                    "$ref": "#/$defs/mock_trial_data_item"
+                }
+            },
+            "unevaluatedProperties": false,
+            "minProperties": 1
+        }
+    },
+
     "type": "object",
     "properties": {
         "class": {
@@ -42,6 +172,9 @@
                             },
                             "minItems": 1,
                             "uniqueItems": true
+                        },
+                        "mock_trial_data": {
+                            "$ref": "#/$defs/mock_trial_data"
                         }
                     }
                 }

mlos_bench/mlos_bench/environments/base_environment.py

@@ -363,6 +363,92 @@ def parameters(self) -> dict[str, TunableValue]:
         """
         return self._params.copy()
 
+    @property
+    def current_trial_id(self) -> int:
+        """
+        Get the current trial ID.
+
+        This value can be used in scripts or environment variables to help
+        identify the Trial this Environment is currently running.
+
+        Returns
+        -------
+        trial_id : int
+            The current trial ID.
+
+        Notes
+        -----
+        This method is used to identify the current trial ID for the environment.
+        It is expected to be called *after* the base
+        :py:meth:`Environment.setup` method has been called and parameters have
+        been assigned.
+        """
+        val = self._params["trial_id"]
+        assert isinstance(val, int), (
+            "Expected trial_id to be an int, but got %s (type %s): %s",
+            val,
+            type(val),
+            self._params,
+        )
+        return val
+
+    @property
+    def trial_runner_id(self) -> int:
+        """
+        Get the ID of the :py:class:`~.mlos_bench.schedulers.trial_runner.TrialRunner`
+        for this Environment.
+
+        This value can be used in scripts or environment variables to help
+        identify the TrialRunner for this Environment.
+
+        Returns
+        -------
+        trial_runner_id : int
+            The trial runner ID.
+
+        Notes
+        -----
+        This shouldn't change during the lifetime of the Environment since each
+        Environment is assigned to a single TrialRunner.
+        """
+        val = self._params["trial_runner_id"]
+        assert isinstance(val, int), (
+            "Expected trial_runner_id to be an int, but got %s (type %s)",
+            val,
+            type(val),
+        )
+        return val
+
+    @property
+    def experiment_id(self) -> str:
+        """
+        Get the ID of the experiment.
+
+        This value can be used in scripts or environment variables to help
+        identify the TrialRunner for this Environment.
+
+        Returns
+        -------
+        experiment_id : str
+            The ID of the experiment.
+
+        Notes
+        -----
+        This value comes from the globals config or ``mlos_bench`` CLI arguments
+        in the experiment setup.
+
+        See Also
+        --------
+        mlos_bench.config : documentation on the configuration system
+        """
+        val = self._params["experiment_id"]
+        assert isinstance(val, str), (
+            "Expected experiment_id to be an int, but got %s (type %s)",
+            val,
+            type(val),
+        )
+        return val
+
     def setup(self, tunables: TunableGroups, global_config: dict | None = None) -> bool:
         """
         Set up a new benchmark environment, if necessary. This method must be