From cd8b88f88af847f227a5898e542327cda866a23d Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 13:57:28 -0500 Subject: [PATCH 01/77] Start introducing MockScheduler for better Scheduler testing --- .../schedulers/base-scheduler-subschema.json | 69 ++++++++++++ .../schedulers/mock-scheduler-subschema.json | 59 ++++++++++ .../schemas/schedulers/scheduler-schema.json | 105 +++--------------- .../schedulers/sync-scheduler-subschema.json | 27 +++++ .../test_load_scheduler_config_examples.py | 55 +++++++++ .../mlos_bench/tests/schedulers/__init__.py | 5 + .../mlos_bench/tests/schedulers/conftest.py | 5 + .../tests/schedulers/mock_scheduler.py | 65 +++++++++++ 8 files changed, 298 insertions(+), 92 deletions(-) create mode 100644 mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json create mode 100644 mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json create mode 100644 mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json create mode 100644 mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py create mode 100644 mlos_bench/mlos_bench/tests/schedulers/__init__.py create mode 100644 mlos_bench/mlos_bench/tests/schedulers/conftest.py create mode 100644 mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json new file mode 100644 index 00000000000..702da1eec3e --- /dev/null +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json @@ -0,0 +1,69 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json", + "title": "mlos_bench base Scheduler config schema definitions", + "description": "mlos_bench base Scheduler config schema definitions for all Scheduler types.", + + "$defs": { + "base_scheduler_config": { + "$comment": "config properties common to all Scheduler types.", + "description": "The scheduler-specific config.", + "type": "object", + "minProperties": 1, + "properties": { + "experiment_id": { + "$ref": "../cli/common-defs-subschemas.json#/$defs/experiment_id" + }, + "trial_id": { + "$ref": "../cli/common-defs-subschemas.json#/$defs/trial_id" + }, + "config_id": { + "$ref": "../cli/common-defs-subschemas.json#/$defs/config_id" + }, + "teardown": { + "description": "Whether to teardown the experiment after running it.", + "type": "boolean" + }, + "max_trials": { + "description": "Max. number of trials to run. Use -1 or 0 for unlimited.", + "type": "integer", + "minimum": -1, + "examples": [50, -1] + }, + "trial_config_repeat_count": { + "description": "Number of times to repeat a config.", + "type": "integer", + "minimum": 1, + "examples": [3, 5] + } + } + } + }, + + "type": "object", + "properties": { + "$schema": { + "description": "The schema to use for validating the scheduler config (accepts both URLs and local paths).", + "type": "string", + "$comment": "This is optional, but if provided, should match the name of the root schema file.", + "pattern": "/schemas/schedulers/scheduler-schema.json$" + }, + + "description": { + "description": "Optional description of the config.", + "type": "string" + }, + + "class": { + "description": "The name of the scheduler class to use.", + "type": "string", + "$comment": "Exact matches are handled elsewhere.", + "pattern": "^mlos_bench[.]schedulers[.]" + }, + + "config": { + "$ref": "#/$defs/base_scheduler_config" + } + }, + "required": ["class"] +} diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json new file mode 100644 index 00000000000..ba3d25589a4 --- /dev/null +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json @@ -0,0 +1,59 @@ +{ + "$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": { + "trial_id": { + "type": "string" + }, + "status": { + "type": "enum", + "enum": [ + null, + "UNKNOWN", + "PENDING", + "READY", + "RUNNING", + "SUCCEEDED", + "CANCELED", + "FAILED", + "TIMED_OUT" + ] + }, + "score": "number" + }, + "required": ["trial_id", "status"] + } + } + }, + "minProperties": 1 + } + ], + "minProperties": 1, + "unevaluatedProperties": false + } + }, + "required": ["class"] +} diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json index 81b2e797547..99096a1c9fe 100644 --- a/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json @@ -2,105 +2,26 @@ "$schema": "https://json-schema.org/draft/2020-12/schema", "$id": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json", "title": "mlos_bench Scheduler config", - - "$defs": { - "comment": { - "$comment": "This section contains reusable partial schema bits (or just split out for readability)" - }, - - "config_base_scheduler": { - "$comment": "config properties common to all Scheduler types.", - "type": "object", - "properties": { - "experiment_id": { - "$ref": "../cli/common-defs-subschemas.json#/$defs/experiment_id" - }, - "trial_id": { - "$ref": "../cli/common-defs-subschemas.json#/$defs/trial_id" - }, - "config_id": { - "$ref": "../cli/common-defs-subschemas.json#/$defs/config_id" - }, - "teardown": { - "description": "Whether to teardown the experiment after running it.", - "type": "boolean" - }, - "max_trials": { - "description": "Max. number of trials to run. Use -1 or 0 for unlimited.", - "type": "integer", - "minimum": -1, - "examples": [50, -1] - }, - "trial_config_repeat_count": { - "description": "Number of times to repeat a config.", - "type": "integer", - "minimum": 1, - "examples": [3, 5] - } - } - } - }, - "description": "config for the mlos_bench scheduler", "$comment": "top level schema document rules", - "type": "object", - "properties": { - "$schema": { - "description": "The schema to use for validating the scheduler config (accepts both URLs and local paths).", - "type": "string", - "$comment": "This is optional, but if provided, should match the name of this file.", - "pattern": "/schemas/schedulers/scheduler-schema.json$" - }, - - "description": { - "description": "Optional description of the config.", - "type": "string" - }, - "class": { - "description": "The name of the scheduler class to use.", - "$comment": "required", - "enum": [ - "mlos_bench.schedulers.SyncScheduler", - "mlos_bench.schedulers.sync_scheduler.SyncScheduler" - ] + "type": "object", + "allOf": [ + { + "$comment": "All scheduler subschemas support these base properties.", + "$ref": "../base-scheduler-subschema.json" }, - - "config": { - "description": "The scheduler-specific config.", - "$comment": "Stub for scheduler-specific config appended with condition statements below", - "type": "object", - "minProperties": 1 - } - }, - "required": ["class"], - - "oneOf": [ { - "$comment": "extensions to the 'config' object properties when synchronous scheduler is being used", - "if": { - "properties": { - "class": { - "enum": [ - "mlos_bench.schedulers.SyncScheduler", - "mlos_bench.schedulers.sync_scheduler.SyncScheduler" - ] - } + "$comment": "The set of known scheduler subschemas.", + "oneOf": [ + { + "$ref": "./sync-scheduler-subschema.json" }, - "required": ["class"] - }, - "then": { - "properties": { - "config": { - "type": "object", - "allOf": [{ "$ref": "#/$defs/config_base_scheduler" }], - "$comment": "disallow other properties", - "unevaluatedProperties": false - } + { + "$ref": "./mock-scheduler-subschema.json" } - }, - "else": false + ] } ], - "unevaluatedProperties": false + "required": ["class"] } diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json new file mode 100644 index 00000000000..609339521bf --- /dev/null +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json @@ -0,0 +1,27 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json", + "title": "mlos_bench SyncScheduler config", + "description": "config for an mlos_bench SyncScheduler", + "type": "object", + "properties": { + "class": { + "enum": [ + "mlos_bench.schedulers.SyncScheduler", + "mlos_bench.schedulers.sync_scheduler.SyncScheduler" + ] + }, + "config": { + "type": "object", + "$comment": "No extra properties supported by SyncScheduler.", + "allOf": [ + { + "$ref": "../base-scheduler-subschema.json#/$defs/base_scheduler_config" + } + ], + "minProperties": 1, + "unevaluatedProperties": false + } + }, + "required": ["class"] +} diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py new file mode 100644 index 00000000000..9127bbbfcad --- /dev/null +++ b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py @@ -0,0 +1,55 @@ +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +"""Tests for loading scheduler config examples.""" +import logging + +import pytest + +from mlos_bench.config.schemas.config_schemas import ConfigSchema +from mlos_bench.schedulers.base_scheduler import Scheduler +from mlos_bench.services.config_persistence import ConfigPersistenceService +from mlos_bench.tests.config import locate_config_examples +from mlos_bench.util import get_class_from_name + +_LOG = logging.getLogger(__name__) +_LOG.setLevel(logging.DEBUG) + + +# Get the set of configs to test. +CONFIG_TYPE = "schedulers" + + +def filter_configs(configs_to_filter: list[str]) -> list[str]: + """If necessary, filter out json files that aren't for the module we're testing.""" + return configs_to_filter + + +configs = locate_config_examples( + ConfigPersistenceService.BUILTIN_CONFIG_PATH, + CONFIG_TYPE, + filter_configs, +) +assert configs + + +@pytest.mark.parametrize("config_path", configs) +def test_load_scheduler_config_examples( + config_loader_service: ConfigPersistenceService, + config_path: str, +) -> None: + """Tests loading a config example.""" + config = config_loader_service.load_config(config_path, ConfigSchema.SCHEDULER) + assert isinstance(config, dict) + # Skip schema loading that would require a database connection for this test. + config["config"]["lazy_schema_create"] = True + cls = get_class_from_name(config["class"]) + assert issubclass(cls, Scheduler) + # Make an instance of the class based on the config. + scheduler_inst = config_loader_service.build_scheduler( + config=config, + service=config_loader_service, + ) + assert scheduler_inst is not None + assert isinstance(scheduler_inst, cls) diff --git a/mlos_bench/mlos_bench/tests/schedulers/__init__.py b/mlos_bench/mlos_bench/tests/schedulers/__init__.py new file mode 100644 index 00000000000..4bc0076079f --- /dev/null +++ b/mlos_bench/mlos_bench/tests/schedulers/__init__.py @@ -0,0 +1,5 @@ +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +"""mlos_bench.tests.schedulers.""" diff --git a/mlos_bench/mlos_bench/tests/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/schedulers/conftest.py new file mode 100644 index 00000000000..83a18783cc6 --- /dev/null +++ b/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.""" diff --git a/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py new file mode 100644 index 00000000000..780b6210469 --- /dev/null +++ b/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")) From 2c9c968f94366a7316f356178c32193db36d6907 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 17:00:33 -0500 Subject: [PATCH 02/77] fixup scheduler schemas refactor --- .../config/schemas/schedulers/mock-scheduler-subschema.json | 3 +-- .../mlos_bench/config/schemas/schedulers/scheduler-schema.json | 2 +- .../config/schemas/schedulers/sync-scheduler-subschema.json | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json index ba3d25589a4..a784145c694 100644 --- a/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json @@ -14,7 +14,7 @@ "type": "object", "allOf": [ { - "$ref": "../base-scheduler-subschema.json#/$defs/base_scheduler_config" + "$ref": "base-scheduler-subschema.json#/$defs/base_scheduler_config" }, { "type": "object", @@ -29,7 +29,6 @@ "type": "string" }, "status": { - "type": "enum", "enum": [ null, "UNKNOWN", diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json index 99096a1c9fe..c3a79ef6051 100644 --- a/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json @@ -9,7 +9,7 @@ "allOf": [ { "$comment": "All scheduler subschemas support these base properties.", - "$ref": "../base-scheduler-subschema.json" + "$ref": "./base-scheduler-subschema.json" }, { "$comment": "The set of known scheduler subschemas.", diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json index 609339521bf..e7f0e40eb2b 100644 --- a/mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json @@ -16,7 +16,7 @@ "$comment": "No extra properties supported by SyncScheduler.", "allOf": [ { - "$ref": "../base-scheduler-subschema.json#/$defs/base_scheduler_config" + "$ref": "base-scheduler-subschema.json#/$defs/base_scheduler_config" } ], "minProperties": 1, From 2f4a82e25cc73b7e0ebfdab7e1feb76614540381 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 17:00:59 -0500 Subject: [PATCH 03/77] reorg optimizer fixtures for reuse --- .../mlos_bench/tests/optimizers/conftest.py | 180 ++---------------- .../mlos_bench/tests/optimizers/fixtures.py | 171 +++++++++++++++++ 2 files changed, 184 insertions(+), 167 deletions(-) create mode 100644 mlos_bench/mlos_bench/tests/optimizers/fixtures.py diff --git a/mlos_bench/mlos_bench/tests/optimizers/conftest.py b/mlos_bench/mlos_bench/tests/optimizers/conftest.py index aaa6b14753a..f1c758bca6d 100644 --- a/mlos_bench/mlos_bench/tests/optimizers/conftest.py +++ b/mlos_bench/mlos_bench/tests/optimizers/conftest.py @@ -2,170 +2,16 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # -"""Test fixtures for mlos_bench optimizers.""" - - -import pytest - -from mlos_bench.optimizers.manual_optimizer import ManualOptimizer -from mlos_bench.optimizers.mlos_core_optimizer import MlosCoreOptimizer -from mlos_bench.optimizers.mock_optimizer import MockOptimizer -from mlos_bench.tests import SEED -from mlos_bench.tunables.tunable_groups import TunableGroups - -# pylint: disable=redefined-outer-name - - -@pytest.fixture -def mock_configs() -> list[dict]: - """Mock configurations of earlier experiments.""" - return [ - { - "vmSize": "Standard_B4ms", - "idle": "halt", - "kernel_sched_migration_cost_ns": 50000, - "kernel_sched_latency_ns": 1000000, - }, - { - "vmSize": "Standard_B4ms", - "idle": "halt", - "kernel_sched_migration_cost_ns": 40000, - "kernel_sched_latency_ns": 2000000, - }, - { - "vmSize": "Standard_B4ms", - "idle": "mwait", - "kernel_sched_migration_cost_ns": -1, # Special value - "kernel_sched_latency_ns": 3000000, - }, - { - "vmSize": "Standard_B2s", - "idle": "mwait", - "kernel_sched_migration_cost_ns": 200000, - "kernel_sched_latency_ns": 4000000, - }, - ] - - -@pytest.fixture -def mock_opt_no_defaults(tunable_groups: TunableGroups) -> MockOptimizer: - """Test fixture for MockOptimizer that ignores the initial configuration.""" - return MockOptimizer( - tunables=tunable_groups, - service=None, - config={ - "optimization_targets": {"score": "min"}, - "max_suggestions": 5, - "start_with_defaults": False, - "seed": SEED, - }, - ) - - -@pytest.fixture -def mock_opt(tunable_groups: TunableGroups) -> MockOptimizer: - """Test fixture for MockOptimizer.""" - return MockOptimizer( - tunables=tunable_groups, - service=None, - config={"optimization_targets": {"score": "min"}, "max_suggestions": 5, "seed": SEED}, - ) - - -@pytest.fixture -def mock_opt_max(tunable_groups: TunableGroups) -> MockOptimizer: - """Test fixture for MockOptimizer.""" - return MockOptimizer( - tunables=tunable_groups, - service=None, - config={"optimization_targets": {"score": "max"}, "max_suggestions": 10, "seed": SEED}, - ) - - -@pytest.fixture -def flaml_opt(tunable_groups: TunableGroups) -> MlosCoreOptimizer: - """Test fixture for mlos_core FLAML optimizer.""" - return MlosCoreOptimizer( - tunables=tunable_groups, - service=None, - config={ - "optimization_targets": {"score": "min"}, - "max_suggestions": 15, - "optimizer_type": "FLAML", - "seed": SEED, - }, - ) - - -@pytest.fixture -def flaml_opt_max(tunable_groups: TunableGroups) -> MlosCoreOptimizer: - """Test fixture for mlos_core FLAML optimizer.""" - return MlosCoreOptimizer( - tunables=tunable_groups, - service=None, - config={ - "optimization_targets": {"score": "max"}, - "max_suggestions": 15, - "optimizer_type": "FLAML", - "seed": SEED, - }, - ) - - -# FIXME: SMAC's RF model can be non-deterministic at low iterations, which are -# normally calculated as a percentage of the max_suggestions and number of -# tunable dimensions, so for now we set the initial random samples equal to the -# number of iterations and control them with a seed. - -SMAC_ITERATIONS = 10 - - -@pytest.fixture -def smac_opt(tunable_groups: TunableGroups) -> MlosCoreOptimizer: - """Test fixture for mlos_core SMAC optimizer.""" - return MlosCoreOptimizer( - tunables=tunable_groups, - service=None, - config={ - "optimization_targets": {"score": "min"}, - "max_suggestions": SMAC_ITERATIONS, - "optimizer_type": "SMAC", - "seed": SEED, - "output_directory": None, - # See Above - "n_random_init": SMAC_ITERATIONS, - "max_ratio": 1.0, - }, - ) - - -@pytest.fixture -def smac_opt_max(tunable_groups: TunableGroups) -> MlosCoreOptimizer: - """Test fixture for mlos_core SMAC optimizer.""" - return MlosCoreOptimizer( - tunables=tunable_groups, - service=None, - config={ - "optimization_targets": {"score": "max"}, - "max_suggestions": SMAC_ITERATIONS, - "optimizer_type": "SMAC", - "seed": SEED, - "output_directory": None, - # See Above - "n_random_init": SMAC_ITERATIONS, - "max_ratio": 1.0, - }, - ) - - -@pytest.fixture -def manual_opt(tunable_groups: TunableGroups, mock_configs: list[dict]) -> ManualOptimizer: - """Test fixture for ManualOptimizer.""" - return ManualOptimizer( - tunables=tunable_groups, - service=None, - config={ - "max_cycles": 2, - "tunable_values_cycle": mock_configs, - }, - ) +"""Export test fixtures for mlos_bench optimizers.""" + +import mlos_bench.tests.optimizers.fixtures + +mock_configs = mlos_bench.tests.optimizers.fixtures.mock_configs +mock_opt_no_defaults = mlos_bench.tests.optimizers.fixtures.mock_opt_no_defaults +mock_opt = mlos_bench.tests.optimizers.fixtures.mock_opt +mock_opt_max = mlos_bench.tests.optimizers.fixtures.mock_opt_max +flaml_opt = mlos_bench.tests.optimizers.fixtures.flaml_opt +flaml_opt_max = mlos_bench.tests.optimizers.fixtures.flaml_opt_max +smac_opt = mlos_bench.tests.optimizers.fixtures.smac_opt +smac_opt_max = mlos_bench.tests.optimizers.fixtures.smac_opt_max +manual_opt = mlos_bench.tests.optimizers.fixtures.manual_opt diff --git a/mlos_bench/mlos_bench/tests/optimizers/fixtures.py b/mlos_bench/mlos_bench/tests/optimizers/fixtures.py new file mode 100644 index 00000000000..aaa6b14753a --- /dev/null +++ b/mlos_bench/mlos_bench/tests/optimizers/fixtures.py @@ -0,0 +1,171 @@ +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +"""Test fixtures for mlos_bench optimizers.""" + + +import pytest + +from mlos_bench.optimizers.manual_optimizer import ManualOptimizer +from mlos_bench.optimizers.mlos_core_optimizer import MlosCoreOptimizer +from mlos_bench.optimizers.mock_optimizer import MockOptimizer +from mlos_bench.tests import SEED +from mlos_bench.tunables.tunable_groups import TunableGroups + +# pylint: disable=redefined-outer-name + + +@pytest.fixture +def mock_configs() -> list[dict]: + """Mock configurations of earlier experiments.""" + return [ + { + "vmSize": "Standard_B4ms", + "idle": "halt", + "kernel_sched_migration_cost_ns": 50000, + "kernel_sched_latency_ns": 1000000, + }, + { + "vmSize": "Standard_B4ms", + "idle": "halt", + "kernel_sched_migration_cost_ns": 40000, + "kernel_sched_latency_ns": 2000000, + }, + { + "vmSize": "Standard_B4ms", + "idle": "mwait", + "kernel_sched_migration_cost_ns": -1, # Special value + "kernel_sched_latency_ns": 3000000, + }, + { + "vmSize": "Standard_B2s", + "idle": "mwait", + "kernel_sched_migration_cost_ns": 200000, + "kernel_sched_latency_ns": 4000000, + }, + ] + + +@pytest.fixture +def mock_opt_no_defaults(tunable_groups: TunableGroups) -> MockOptimizer: + """Test fixture for MockOptimizer that ignores the initial configuration.""" + return MockOptimizer( + tunables=tunable_groups, + service=None, + config={ + "optimization_targets": {"score": "min"}, + "max_suggestions": 5, + "start_with_defaults": False, + "seed": SEED, + }, + ) + + +@pytest.fixture +def mock_opt(tunable_groups: TunableGroups) -> MockOptimizer: + """Test fixture for MockOptimizer.""" + return MockOptimizer( + tunables=tunable_groups, + service=None, + config={"optimization_targets": {"score": "min"}, "max_suggestions": 5, "seed": SEED}, + ) + + +@pytest.fixture +def mock_opt_max(tunable_groups: TunableGroups) -> MockOptimizer: + """Test fixture for MockOptimizer.""" + return MockOptimizer( + tunables=tunable_groups, + service=None, + config={"optimization_targets": {"score": "max"}, "max_suggestions": 10, "seed": SEED}, + ) + + +@pytest.fixture +def flaml_opt(tunable_groups: TunableGroups) -> MlosCoreOptimizer: + """Test fixture for mlos_core FLAML optimizer.""" + return MlosCoreOptimizer( + tunables=tunable_groups, + service=None, + config={ + "optimization_targets": {"score": "min"}, + "max_suggestions": 15, + "optimizer_type": "FLAML", + "seed": SEED, + }, + ) + + +@pytest.fixture +def flaml_opt_max(tunable_groups: TunableGroups) -> MlosCoreOptimizer: + """Test fixture for mlos_core FLAML optimizer.""" + return MlosCoreOptimizer( + tunables=tunable_groups, + service=None, + config={ + "optimization_targets": {"score": "max"}, + "max_suggestions": 15, + "optimizer_type": "FLAML", + "seed": SEED, + }, + ) + + +# FIXME: SMAC's RF model can be non-deterministic at low iterations, which are +# normally calculated as a percentage of the max_suggestions and number of +# tunable dimensions, so for now we set the initial random samples equal to the +# number of iterations and control them with a seed. + +SMAC_ITERATIONS = 10 + + +@pytest.fixture +def smac_opt(tunable_groups: TunableGroups) -> MlosCoreOptimizer: + """Test fixture for mlos_core SMAC optimizer.""" + return MlosCoreOptimizer( + tunables=tunable_groups, + service=None, + config={ + "optimization_targets": {"score": "min"}, + "max_suggestions": SMAC_ITERATIONS, + "optimizer_type": "SMAC", + "seed": SEED, + "output_directory": None, + # See Above + "n_random_init": SMAC_ITERATIONS, + "max_ratio": 1.0, + }, + ) + + +@pytest.fixture +def smac_opt_max(tunable_groups: TunableGroups) -> MlosCoreOptimizer: + """Test fixture for mlos_core SMAC optimizer.""" + return MlosCoreOptimizer( + tunables=tunable_groups, + service=None, + config={ + "optimization_targets": {"score": "max"}, + "max_suggestions": SMAC_ITERATIONS, + "optimizer_type": "SMAC", + "seed": SEED, + "output_directory": None, + # See Above + "n_random_init": SMAC_ITERATIONS, + "max_ratio": 1.0, + }, + ) + + +@pytest.fixture +def manual_opt(tunable_groups: TunableGroups, mock_configs: list[dict]) -> ManualOptimizer: + """Test fixture for ManualOptimizer.""" + return ManualOptimizer( + tunables=tunable_groups, + service=None, + config={ + "max_cycles": 2, + "tunable_values_cycle": mock_configs, + }, + ) From 24ccf5a893d875d99f7b2d7df18dfc220a7e6fbe Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 17:04:59 -0500 Subject: [PATCH 04/77] reorg files based storage fixture --- .../mlos_bench/tests/storage/conftest.py | 1 + .../mlos_bench/tests/storage/sql/fixtures.py | 37 ++++++ .../tests/storage/test_storage_pickling.py | 117 ++++++++---------- 3 files changed, 88 insertions(+), 67 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/storage/conftest.py b/mlos_bench/mlos_bench/tests/storage/conftest.py index a1437052823..c510793fac1 100644 --- a/mlos_bench/mlos_bench/tests/storage/conftest.py +++ b/mlos_bench/mlos_bench/tests/storage/conftest.py @@ -12,6 +12,7 @@ # Expose some of those as local names so they can be picked up as fixtures by pytest. storage = sql_storage_fixtures.storage +sqlite_storage = sql_storage_fixtures.sqlite_storage exp_storage = sql_storage_fixtures.exp_storage exp_no_tunables_storage = sql_storage_fixtures.exp_no_tunables_storage mixed_numerics_exp_storage = sql_storage_fixtures.mixed_numerics_exp_storage diff --git a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py index cb83bffd4ff..8cd4954bdda 100644 --- a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py +++ b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py @@ -7,6 +7,10 @@ from collections.abc import Generator from random import seed as rand_seed +import json +import os +import tempfile + import pytest from mlos_bench.optimizers.mock_optimizer import MockOptimizer @@ -15,6 +19,7 @@ from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.storage.base_experiment_data import ExperimentData from mlos_bench.storage.sql.storage import SqlStorage +from mlos_bench.storage.storage_factory import from_config from mlos_bench.tests import SEED from mlos_bench.tests.storage import ( CONFIG_TRIAL_REPEAT_COUNT, @@ -26,6 +31,38 @@ # pylint: disable=redefined-outer-name +@pytest.fixture +def sqlite_storage() -> Generator[SqlStorage]: + """ + Fixture for file based SQLite storage in a temporary directory. + + Yields + ------ + Generator[SqlStorage] + + Notes + ----- + Can't be used in parallel tests on Windows. + """ + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "mlos_bench.sqlite") + config_str = json.dumps( + { + "class": "mlos_bench.storage.sql.storage.SqlStorage", + "config": { + "drivername": "sqlite", + "database": db_path, + "lazy_schema_create": False, + }, + } + ) + + storage = from_config(config_str) + assert isinstance(storage, SqlStorage) + storage.update_schema() + yield storage + + @pytest.fixture def storage() -> SqlStorage: """Test fixture for in-memory SQLite3 storage.""" diff --git a/mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py b/mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py index 3d5053837be..7871e7f68ca 100644 --- a/mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py +++ b/mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py @@ -3,11 +3,8 @@ # Licensed under the MIT License. # """Test pickling and unpickling of Storage, and restoring Experiment and Trial by id.""" -import json -import os import pickle import sys -import tempfile from datetime import datetime from typing import Literal @@ -16,7 +13,6 @@ from mlos_bench.environments.status import Status from mlos_bench.storage.sql.storage import SqlStorage -from mlos_bench.storage.storage_factory import from_config from mlos_bench.tunables.tunable_groups import TunableGroups @@ -26,72 +22,59 @@ sys.platform == "win32", reason="Windows doesn't support multiple processes accessing the same file.", ) -def test_storage_pickle_restore_experiment_and_trial(tunable_groups: TunableGroups) -> None: +def test_storage_pickle_restore_experiment_and_trial( + sqlite_storage: SqlStorage, + tunable_groups: TunableGroups, +) -> None: """Check that we can pickle and unpickle the Storage object, and restore Experiment and Trial by id. """ - # pylint: disable=too-many-locals - with tempfile.TemporaryDirectory() as tmpdir: - db_path = os.path.join(tmpdir, "mlos_bench.sqlite") - config_str = json.dumps( - { - "class": "mlos_bench.storage.sql.storage.SqlStorage", - "config": { - "drivername": "sqlite", - "database": db_path, - "lazy_schema_create": False, - }, - } - ) + storage = sqlite_storage + # Create an Experiment and a Trial + opt_targets: dict[str, Literal["min", "max"]] = {"metric": "min"} + experiment = storage.experiment( + experiment_id="experiment_id", + trial_id=0, + root_env_config="dummy_env.json", + description="Pickle test experiment", + tunables=tunable_groups, + opt_targets=opt_targets, + ) + with experiment: + trial = experiment.new_trial(tunable_groups) + trial_id_created = trial.trial_id + trial.set_trial_runner(1) + trial.update(Status.RUNNING, datetime.now(UTC)) - storage = from_config(config_str) - storage.update_schema() + # Pickle and unpickle the Storage object + pickled = pickle.dumps(storage) + restored_storage = pickle.loads(pickled) + assert isinstance(restored_storage, SqlStorage) - # Create an Experiment and a Trial - opt_targets: dict[str, Literal["min", "max"]] = {"metric": "min"} - experiment = storage.experiment( - experiment_id="experiment_id", - trial_id=0, - root_env_config="dummy_env.json", - description="Pickle test experiment", - tunables=tunable_groups, - opt_targets=opt_targets, - ) - with experiment: - trial = experiment.new_trial(tunable_groups) - trial_id_created = trial.trial_id - trial.set_trial_runner(1) - trial.update(Status.RUNNING, datetime.now(UTC)) + # Restore the Experiment from storage by id and check that it matches the original + restored_experiment = restored_storage.get_experiment_by_id( + experiment_id=experiment.experiment_id, + tunables=tunable_groups, + opt_targets=opt_targets, + ) + assert restored_experiment is not None + assert restored_experiment is not experiment + assert restored_experiment.experiment_id == experiment.experiment_id + assert restored_experiment.description == experiment.description + assert restored_experiment.root_env_config == experiment.root_env_config + assert restored_experiment.tunables == experiment.tunables + assert restored_experiment.opt_targets == experiment.opt_targets + with restored_experiment: + # trial_id should have been restored during __enter__ + assert restored_experiment.trial_id == experiment.trial_id - # Pickle and unpickle the Storage object - pickled = pickle.dumps(storage) - restored_storage = pickle.loads(pickled) - assert isinstance(restored_storage, SqlStorage) - - # Restore the Experiment from storage by id and check that it matches the original - restored_experiment = restored_storage.get_experiment_by_id( - experiment_id=experiment.experiment_id, - tunables=tunable_groups, - opt_targets=opt_targets, - ) - assert restored_experiment is not None - assert restored_experiment is not experiment - assert restored_experiment.experiment_id == experiment.experiment_id - assert restored_experiment.description == experiment.description - assert restored_experiment.root_env_config == experiment.root_env_config - assert restored_experiment.tunables == experiment.tunables - assert restored_experiment.opt_targets == experiment.opt_targets - with restored_experiment: - # trial_id should have been restored during __enter__ - assert restored_experiment.trial_id == experiment.trial_id - - # Restore the Trial from storage by id and check that it matches the original - restored_trial = restored_experiment.get_trial_by_id(trial_id_created) - assert restored_trial is not None - assert restored_trial is not trial - assert restored_trial.trial_id == trial.trial_id - assert restored_trial.experiment_id == trial.experiment_id - assert restored_trial.tunables == trial.tunables - assert restored_trial.status == trial.status - assert restored_trial.config() == trial.config() - assert restored_trial.trial_runner_id == trial.trial_runner_id + # Restore the Trial from storage by id and check that it matches the original + restored_trial = restored_experiment.get_trial_by_id(trial_id_created) + assert restored_trial is not None + assert restored_trial is not trial + assert restored_trial.trial_id == trial.trial_id + assert restored_trial.experiment_id == trial.experiment_id + assert restored_trial.tunables == trial.tunables + assert restored_trial.status == trial.status + assert restored_trial.config() == trial.config() + assert restored_trial.trial_runner_id == trial.trial_runner_id From bdfd9b086cfe35cea88ff8a5d2e9b9b539a7de5e Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 17:09:40 -0500 Subject: [PATCH 05/77] adding basic scheduler config example loader test --- .../tests/config/schedulers/conftest.py | 54 +++++++++++++++++++ .../test_load_scheduler_config_examples.py | 28 ++++++++-- 2 files changed, 79 insertions(+), 3 deletions(-) create mode 100644 mlos_bench/mlos_bench/tests/config/schedulers/conftest.py diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py new file mode 100644 index 00000000000..a1fd105839b --- /dev/null +++ b/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py @@ -0,0 +1,54 @@ +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +""" +Pytest fixtures for Scheduler config tests. + +Provides fixtures for creating multiple TrialRunner instances using the mock environment config. +""" + +from importlib.resources import files + +import pytest + +from mlos_bench.services.config_persistence import ConfigPersistenceService +from mlos_bench.schedulers.trial_runner import TrialRunner +from mlos_bench.util import path_join + +# pylint: disable=redefined-outer-name + +TRIAL_RUNNERS_COUNT = 4 + +@pytest.fixture +def mock_env_config_path() -> str: + """ + Returns the absolute path to the mock environment configuration file. + This file is used to create TrialRunner instances for testing. + """ + + # Use the files() routine to locate the file relative to this directory + return path_join( + str(files("mlos_bench.config").joinpath("environments", "mock", "mock_env.jsonc")), + abs_path=True, + ) + + +@pytest.fixture +def trial_runners( + config_loader_service: ConfigPersistenceService, + mock_env_config_path: str, +) -> list[TrialRunner]: + """ + Fixture that returns a list of TrialRunner instances using the mock environment config. + + Returns + ------- + list[TrialRunner] + List of TrialRunner instances created from the mock environment config. + """ + return TrialRunner.create_from_json( + config_loader=config_loader_service, + env_json=mock_env_config_path, + num_trial_runners=TRIAL_RUNNERS_COUNT, + ) diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py index 9127bbbfcad..c8d6bcb1726 100644 --- a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py @@ -8,14 +8,25 @@ import pytest from mlos_bench.config.schemas.config_schemas import ConfigSchema +from mlos_bench.optimizers.mock_optimizer import MockOptimizer from mlos_bench.schedulers.base_scheduler import Scheduler +from mlos_bench.schedulers.trial_runner import TrialRunner from mlos_bench.services.config_persistence import ConfigPersistenceService +from mlos_bench.storage.sql.storage import SqlStorage from mlos_bench.tests.config import locate_config_examples from mlos_bench.util import get_class_from_name +import mlos_bench.tests.storage.sql.fixtures +import mlos_bench.tests.optimizers.fixtures + +mock_opt = mlos_bench.tests.optimizers.fixtures.mock_opt +sqlite_storage = mlos_bench.tests.storage.sql.fixtures.sqlite_storage + + _LOG = logging.getLogger(__name__) _LOG.setLevel(logging.DEBUG) +# pylint: disable=redefined-outer-name # Get the set of configs to test. CONFIG_TYPE = "schedulers" @@ -38,18 +49,29 @@ def filter_configs(configs_to_filter: list[str]) -> list[str]: def test_load_scheduler_config_examples( config_loader_service: ConfigPersistenceService, config_path: str, + mock_env_config_path: str, + trial_runners: list[TrialRunner], + sqlite_storage: SqlStorage, + mock_opt: MockOptimizer, ) -> None: """Tests loading a config example.""" config = config_loader_service.load_config(config_path, ConfigSchema.SCHEDULER) assert isinstance(config, dict) - # Skip schema loading that would require a database connection for this test. - config["config"]["lazy_schema_create"] = True cls = get_class_from_name(config["class"]) assert issubclass(cls, Scheduler) + global_config = { + # Required configs generally provided by the Launcher. + "experiment_id": f"test_experiment_{__name__}", + "trial_id": 1, + } # Make an instance of the class based on the config. scheduler_inst = config_loader_service.build_scheduler( config=config, - service=config_loader_service, + global_config=global_config, + trial_runners=trial_runners, + optimizer=mock_opt, + storage=sqlite_storage, + root_env_config=mock_env_config_path, ) assert scheduler_inst is not None assert isinstance(scheduler_inst, cls) From b66e17a946e8764aacbb689f99fea0cd3d3ad847 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 17:12:33 -0500 Subject: [PATCH 06/77] Refactor Scheduler schema definitions to make it easier to add new ones. --- .../schedulers/base-scheduler-subschema.json | 69 ++++++++++++ .../schemas/schedulers/scheduler-schema.json | 104 ++---------------- .../schedulers/sync-scheduler-subschema.json | 27 +++++ 3 files changed, 107 insertions(+), 93 deletions(-) create mode 100644 mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json create mode 100644 mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json new file mode 100644 index 00000000000..702da1eec3e --- /dev/null +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json @@ -0,0 +1,69 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json", + "title": "mlos_bench base Scheduler config schema definitions", + "description": "mlos_bench base Scheduler config schema definitions for all Scheduler types.", + + "$defs": { + "base_scheduler_config": { + "$comment": "config properties common to all Scheduler types.", + "description": "The scheduler-specific config.", + "type": "object", + "minProperties": 1, + "properties": { + "experiment_id": { + "$ref": "../cli/common-defs-subschemas.json#/$defs/experiment_id" + }, + "trial_id": { + "$ref": "../cli/common-defs-subschemas.json#/$defs/trial_id" + }, + "config_id": { + "$ref": "../cli/common-defs-subschemas.json#/$defs/config_id" + }, + "teardown": { + "description": "Whether to teardown the experiment after running it.", + "type": "boolean" + }, + "max_trials": { + "description": "Max. number of trials to run. Use -1 or 0 for unlimited.", + "type": "integer", + "minimum": -1, + "examples": [50, -1] + }, + "trial_config_repeat_count": { + "description": "Number of times to repeat a config.", + "type": "integer", + "minimum": 1, + "examples": [3, 5] + } + } + } + }, + + "type": "object", + "properties": { + "$schema": { + "description": "The schema to use for validating the scheduler config (accepts both URLs and local paths).", + "type": "string", + "$comment": "This is optional, but if provided, should match the name of the root schema file.", + "pattern": "/schemas/schedulers/scheduler-schema.json$" + }, + + "description": { + "description": "Optional description of the config.", + "type": "string" + }, + + "class": { + "description": "The name of the scheduler class to use.", + "type": "string", + "$comment": "Exact matches are handled elsewhere.", + "pattern": "^mlos_bench[.]schedulers[.]" + }, + + "config": { + "$ref": "#/$defs/base_scheduler_config" + } + }, + "required": ["class"] +} diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json index 81b2e797547..8dfce6be369 100644 --- a/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json @@ -2,105 +2,23 @@ "$schema": "https://json-schema.org/draft/2020-12/schema", "$id": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json", "title": "mlos_bench Scheduler config", - - "$defs": { - "comment": { - "$comment": "This section contains reusable partial schema bits (or just split out for readability)" - }, - - "config_base_scheduler": { - "$comment": "config properties common to all Scheduler types.", - "type": "object", - "properties": { - "experiment_id": { - "$ref": "../cli/common-defs-subschemas.json#/$defs/experiment_id" - }, - "trial_id": { - "$ref": "../cli/common-defs-subschemas.json#/$defs/trial_id" - }, - "config_id": { - "$ref": "../cli/common-defs-subschemas.json#/$defs/config_id" - }, - "teardown": { - "description": "Whether to teardown the experiment after running it.", - "type": "boolean" - }, - "max_trials": { - "description": "Max. number of trials to run. Use -1 or 0 for unlimited.", - "type": "integer", - "minimum": -1, - "examples": [50, -1] - }, - "trial_config_repeat_count": { - "description": "Number of times to repeat a config.", - "type": "integer", - "minimum": 1, - "examples": [3, 5] - } - } - } - }, - "description": "config for the mlos_bench scheduler", "$comment": "top level schema document rules", - "type": "object", - "properties": { - "$schema": { - "description": "The schema to use for validating the scheduler config (accepts both URLs and local paths).", - "type": "string", - "$comment": "This is optional, but if provided, should match the name of this file.", - "pattern": "/schemas/schedulers/scheduler-schema.json$" - }, - "description": { - "description": "Optional description of the config.", - "type": "string" - }, - - "class": { - "description": "The name of the scheduler class to use.", - "$comment": "required", - "enum": [ - "mlos_bench.schedulers.SyncScheduler", - "mlos_bench.schedulers.sync_scheduler.SyncScheduler" - ] + "type": "object", + "allOf": [ + { + "$comment": "All scheduler subschemas support these base properties.", + "$ref": "./base-scheduler-subschema.json" }, - - "config": { - "description": "The scheduler-specific config.", - "$comment": "Stub for scheduler-specific config appended with condition statements below", - "type": "object", - "minProperties": 1 - } - }, - "required": ["class"], - - "oneOf": [ { - "$comment": "extensions to the 'config' object properties when synchronous scheduler is being used", - "if": { - "properties": { - "class": { - "enum": [ - "mlos_bench.schedulers.SyncScheduler", - "mlos_bench.schedulers.sync_scheduler.SyncScheduler" - ] - } - }, - "required": ["class"] - }, - "then": { - "properties": { - "config": { - "type": "object", - "allOf": [{ "$ref": "#/$defs/config_base_scheduler" }], - "$comment": "disallow other properties", - "unevaluatedProperties": false - } + "$comment": "The set of known Scheduler subschemas. Add others as needed.", + "oneOf": [ + { + "$ref": "./sync-scheduler-subschema.json" } - }, - "else": false + ] } ], - "unevaluatedProperties": false + "required": ["class"] } diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json new file mode 100644 index 00000000000..e7f0e40eb2b --- /dev/null +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json @@ -0,0 +1,27 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/schedulers/sync-scheduler-subschema.json", + "title": "mlos_bench SyncScheduler config", + "description": "config for an mlos_bench SyncScheduler", + "type": "object", + "properties": { + "class": { + "enum": [ + "mlos_bench.schedulers.SyncScheduler", + "mlos_bench.schedulers.sync_scheduler.SyncScheduler" + ] + }, + "config": { + "type": "object", + "$comment": "No extra properties supported by SyncScheduler.", + "allOf": [ + { + "$ref": "base-scheduler-subschema.json#/$defs/base_scheduler_config" + } + ], + "minProperties": 1, + "unevaluatedProperties": false + } + }, + "required": ["class"] +} From 30191c094112b370c8d288b4fc89fac650e22428 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 19 May 2025 22:38:43 +0000 Subject: [PATCH 07/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../mlos_bench/tests/config/schedulers/conftest.py | 10 +++++++--- .../schedulers/test_load_scheduler_config_examples.py | 5 ++--- mlos_bench/mlos_bench/tests/storage/sql/fixtures.py | 5 ++--- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py index a1fd105839b..fdf27162b35 100644 --- a/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py +++ b/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py @@ -5,25 +5,28 @@ """ Pytest fixtures for Scheduler config tests. -Provides fixtures for creating multiple TrialRunner instances using the mock environment config. +Provides fixtures for creating multiple TrialRunner instances using the mock environment +config. """ from importlib.resources import files import pytest -from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.schedulers.trial_runner import TrialRunner +from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.util import path_join # pylint: disable=redefined-outer-name TRIAL_RUNNERS_COUNT = 4 + @pytest.fixture def mock_env_config_path() -> str: """ Returns the absolute path to the mock environment configuration file. + This file is used to create TrialRunner instances for testing. """ @@ -40,7 +43,8 @@ def trial_runners( mock_env_config_path: str, ) -> list[TrialRunner]: """ - Fixture that returns a list of TrialRunner instances using the mock environment config. + Fixture that returns a list of TrialRunner instances using the mock environment + config. Returns ------- diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py index c8d6bcb1726..bf10a99c46b 100644 --- a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py @@ -7,6 +7,8 @@ import pytest +import mlos_bench.tests.optimizers.fixtures +import mlos_bench.tests.storage.sql.fixtures from mlos_bench.config.schemas.config_schemas import ConfigSchema from mlos_bench.optimizers.mock_optimizer import MockOptimizer from mlos_bench.schedulers.base_scheduler import Scheduler @@ -16,9 +18,6 @@ from mlos_bench.tests.config import locate_config_examples from mlos_bench.util import get_class_from_name -import mlos_bench.tests.storage.sql.fixtures -import mlos_bench.tests.optimizers.fixtures - mock_opt = mlos_bench.tests.optimizers.fixtures.mock_opt sqlite_storage = mlos_bench.tests.storage.sql.fixtures.sqlite_storage diff --git a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py index 8cd4954bdda..0bebeeff824 100644 --- a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py +++ b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py @@ -4,12 +4,11 @@ # """Test fixtures for mlos_bench storage.""" -from collections.abc import Generator -from random import seed as rand_seed - import json import os import tempfile +from collections.abc import Generator +from random import seed as rand_seed import pytest From bc8dc8f3b4878f0aa52fbcdbd2d68430d56f4f0c Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 17:41:22 -0500 Subject: [PATCH 08/77] Refactor some test fixtures for better reuse so we can test loading Scheduler config examples. --- .../tests/config/schedulers/__init__.py | 3 + .../tests/config/schedulers/conftest.py | 54 ++++++ .../test_load_scheduler_config_examples.py | 77 ++++++++ .../mlos_bench/tests/optimizers/conftest.py | 180 ++---------------- .../mlos_bench/tests/optimizers/fixtures.py | 171 +++++++++++++++++ .../mlos_bench/tests/storage/conftest.py | 1 + .../mlos_bench/tests/storage/sql/fixtures.py | 37 ++++ .../tests/storage/test_storage_pickling.py | 117 +++++------- 8 files changed, 406 insertions(+), 234 deletions(-) create mode 100644 mlos_bench/mlos_bench/tests/config/schedulers/__init__.py create mode 100644 mlos_bench/mlos_bench/tests/config/schedulers/conftest.py create mode 100644 mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py create mode 100644 mlos_bench/mlos_bench/tests/optimizers/fixtures.py diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/__init__.py b/mlos_bench/mlos_bench/tests/config/schedulers/__init__.py new file mode 100644 index 00000000000..7838135124c --- /dev/null +++ b/mlos_bench/mlos_bench/tests/config/schedulers/__init__.py @@ -0,0 +1,3 @@ +""" +Unit tests for the mlos_bench Scheduler configs. +""" diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py new file mode 100644 index 00000000000..a1fd105839b --- /dev/null +++ b/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py @@ -0,0 +1,54 @@ +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +""" +Pytest fixtures for Scheduler config tests. + +Provides fixtures for creating multiple TrialRunner instances using the mock environment config. +""" + +from importlib.resources import files + +import pytest + +from mlos_bench.services.config_persistence import ConfigPersistenceService +from mlos_bench.schedulers.trial_runner import TrialRunner +from mlos_bench.util import path_join + +# pylint: disable=redefined-outer-name + +TRIAL_RUNNERS_COUNT = 4 + +@pytest.fixture +def mock_env_config_path() -> str: + """ + Returns the absolute path to the mock environment configuration file. + This file is used to create TrialRunner instances for testing. + """ + + # Use the files() routine to locate the file relative to this directory + return path_join( + str(files("mlos_bench.config").joinpath("environments", "mock", "mock_env.jsonc")), + abs_path=True, + ) + + +@pytest.fixture +def trial_runners( + config_loader_service: ConfigPersistenceService, + mock_env_config_path: str, +) -> list[TrialRunner]: + """ + Fixture that returns a list of TrialRunner instances using the mock environment config. + + Returns + ------- + list[TrialRunner] + List of TrialRunner instances created from the mock environment config. + """ + return TrialRunner.create_from_json( + config_loader=config_loader_service, + env_json=mock_env_config_path, + num_trial_runners=TRIAL_RUNNERS_COUNT, + ) diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py new file mode 100644 index 00000000000..c8d6bcb1726 --- /dev/null +++ b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py @@ -0,0 +1,77 @@ +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +"""Tests for loading scheduler config examples.""" +import logging + +import pytest + +from mlos_bench.config.schemas.config_schemas import ConfigSchema +from mlos_bench.optimizers.mock_optimizer import MockOptimizer +from mlos_bench.schedulers.base_scheduler import Scheduler +from mlos_bench.schedulers.trial_runner import TrialRunner +from mlos_bench.services.config_persistence import ConfigPersistenceService +from mlos_bench.storage.sql.storage import SqlStorage +from mlos_bench.tests.config import locate_config_examples +from mlos_bench.util import get_class_from_name + +import mlos_bench.tests.storage.sql.fixtures +import mlos_bench.tests.optimizers.fixtures + +mock_opt = mlos_bench.tests.optimizers.fixtures.mock_opt +sqlite_storage = mlos_bench.tests.storage.sql.fixtures.sqlite_storage + + +_LOG = logging.getLogger(__name__) +_LOG.setLevel(logging.DEBUG) + +# pylint: disable=redefined-outer-name + +# Get the set of configs to test. +CONFIG_TYPE = "schedulers" + + +def filter_configs(configs_to_filter: list[str]) -> list[str]: + """If necessary, filter out json files that aren't for the module we're testing.""" + return configs_to_filter + + +configs = locate_config_examples( + ConfigPersistenceService.BUILTIN_CONFIG_PATH, + CONFIG_TYPE, + filter_configs, +) +assert configs + + +@pytest.mark.parametrize("config_path", configs) +def test_load_scheduler_config_examples( + config_loader_service: ConfigPersistenceService, + config_path: str, + mock_env_config_path: str, + trial_runners: list[TrialRunner], + sqlite_storage: SqlStorage, + mock_opt: MockOptimizer, +) -> None: + """Tests loading a config example.""" + config = config_loader_service.load_config(config_path, ConfigSchema.SCHEDULER) + assert isinstance(config, dict) + cls = get_class_from_name(config["class"]) + assert issubclass(cls, Scheduler) + global_config = { + # Required configs generally provided by the Launcher. + "experiment_id": f"test_experiment_{__name__}", + "trial_id": 1, + } + # Make an instance of the class based on the config. + scheduler_inst = config_loader_service.build_scheduler( + config=config, + global_config=global_config, + trial_runners=trial_runners, + optimizer=mock_opt, + storage=sqlite_storage, + root_env_config=mock_env_config_path, + ) + assert scheduler_inst is not None + assert isinstance(scheduler_inst, cls) diff --git a/mlos_bench/mlos_bench/tests/optimizers/conftest.py b/mlos_bench/mlos_bench/tests/optimizers/conftest.py index aaa6b14753a..f1c758bca6d 100644 --- a/mlos_bench/mlos_bench/tests/optimizers/conftest.py +++ b/mlos_bench/mlos_bench/tests/optimizers/conftest.py @@ -2,170 +2,16 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # -"""Test fixtures for mlos_bench optimizers.""" - - -import pytest - -from mlos_bench.optimizers.manual_optimizer import ManualOptimizer -from mlos_bench.optimizers.mlos_core_optimizer import MlosCoreOptimizer -from mlos_bench.optimizers.mock_optimizer import MockOptimizer -from mlos_bench.tests import SEED -from mlos_bench.tunables.tunable_groups import TunableGroups - -# pylint: disable=redefined-outer-name - - -@pytest.fixture -def mock_configs() -> list[dict]: - """Mock configurations of earlier experiments.""" - return [ - { - "vmSize": "Standard_B4ms", - "idle": "halt", - "kernel_sched_migration_cost_ns": 50000, - "kernel_sched_latency_ns": 1000000, - }, - { - "vmSize": "Standard_B4ms", - "idle": "halt", - "kernel_sched_migration_cost_ns": 40000, - "kernel_sched_latency_ns": 2000000, - }, - { - "vmSize": "Standard_B4ms", - "idle": "mwait", - "kernel_sched_migration_cost_ns": -1, # Special value - "kernel_sched_latency_ns": 3000000, - }, - { - "vmSize": "Standard_B2s", - "idle": "mwait", - "kernel_sched_migration_cost_ns": 200000, - "kernel_sched_latency_ns": 4000000, - }, - ] - - -@pytest.fixture -def mock_opt_no_defaults(tunable_groups: TunableGroups) -> MockOptimizer: - """Test fixture for MockOptimizer that ignores the initial configuration.""" - return MockOptimizer( - tunables=tunable_groups, - service=None, - config={ - "optimization_targets": {"score": "min"}, - "max_suggestions": 5, - "start_with_defaults": False, - "seed": SEED, - }, - ) - - -@pytest.fixture -def mock_opt(tunable_groups: TunableGroups) -> MockOptimizer: - """Test fixture for MockOptimizer.""" - return MockOptimizer( - tunables=tunable_groups, - service=None, - config={"optimization_targets": {"score": "min"}, "max_suggestions": 5, "seed": SEED}, - ) - - -@pytest.fixture -def mock_opt_max(tunable_groups: TunableGroups) -> MockOptimizer: - """Test fixture for MockOptimizer.""" - return MockOptimizer( - tunables=tunable_groups, - service=None, - config={"optimization_targets": {"score": "max"}, "max_suggestions": 10, "seed": SEED}, - ) - - -@pytest.fixture -def flaml_opt(tunable_groups: TunableGroups) -> MlosCoreOptimizer: - """Test fixture for mlos_core FLAML optimizer.""" - return MlosCoreOptimizer( - tunables=tunable_groups, - service=None, - config={ - "optimization_targets": {"score": "min"}, - "max_suggestions": 15, - "optimizer_type": "FLAML", - "seed": SEED, - }, - ) - - -@pytest.fixture -def flaml_opt_max(tunable_groups: TunableGroups) -> MlosCoreOptimizer: - """Test fixture for mlos_core FLAML optimizer.""" - return MlosCoreOptimizer( - tunables=tunable_groups, - service=None, - config={ - "optimization_targets": {"score": "max"}, - "max_suggestions": 15, - "optimizer_type": "FLAML", - "seed": SEED, - }, - ) - - -# FIXME: SMAC's RF model can be non-deterministic at low iterations, which are -# normally calculated as a percentage of the max_suggestions and number of -# tunable dimensions, so for now we set the initial random samples equal to the -# number of iterations and control them with a seed. - -SMAC_ITERATIONS = 10 - - -@pytest.fixture -def smac_opt(tunable_groups: TunableGroups) -> MlosCoreOptimizer: - """Test fixture for mlos_core SMAC optimizer.""" - return MlosCoreOptimizer( - tunables=tunable_groups, - service=None, - config={ - "optimization_targets": {"score": "min"}, - "max_suggestions": SMAC_ITERATIONS, - "optimizer_type": "SMAC", - "seed": SEED, - "output_directory": None, - # See Above - "n_random_init": SMAC_ITERATIONS, - "max_ratio": 1.0, - }, - ) - - -@pytest.fixture -def smac_opt_max(tunable_groups: TunableGroups) -> MlosCoreOptimizer: - """Test fixture for mlos_core SMAC optimizer.""" - return MlosCoreOptimizer( - tunables=tunable_groups, - service=None, - config={ - "optimization_targets": {"score": "max"}, - "max_suggestions": SMAC_ITERATIONS, - "optimizer_type": "SMAC", - "seed": SEED, - "output_directory": None, - # See Above - "n_random_init": SMAC_ITERATIONS, - "max_ratio": 1.0, - }, - ) - - -@pytest.fixture -def manual_opt(tunable_groups: TunableGroups, mock_configs: list[dict]) -> ManualOptimizer: - """Test fixture for ManualOptimizer.""" - return ManualOptimizer( - tunables=tunable_groups, - service=None, - config={ - "max_cycles": 2, - "tunable_values_cycle": mock_configs, - }, - ) +"""Export test fixtures for mlos_bench optimizers.""" + +import mlos_bench.tests.optimizers.fixtures + +mock_configs = mlos_bench.tests.optimizers.fixtures.mock_configs +mock_opt_no_defaults = mlos_bench.tests.optimizers.fixtures.mock_opt_no_defaults +mock_opt = mlos_bench.tests.optimizers.fixtures.mock_opt +mock_opt_max = mlos_bench.tests.optimizers.fixtures.mock_opt_max +flaml_opt = mlos_bench.tests.optimizers.fixtures.flaml_opt +flaml_opt_max = mlos_bench.tests.optimizers.fixtures.flaml_opt_max +smac_opt = mlos_bench.tests.optimizers.fixtures.smac_opt +smac_opt_max = mlos_bench.tests.optimizers.fixtures.smac_opt_max +manual_opt = mlos_bench.tests.optimizers.fixtures.manual_opt diff --git a/mlos_bench/mlos_bench/tests/optimizers/fixtures.py b/mlos_bench/mlos_bench/tests/optimizers/fixtures.py new file mode 100644 index 00000000000..aaa6b14753a --- /dev/null +++ b/mlos_bench/mlos_bench/tests/optimizers/fixtures.py @@ -0,0 +1,171 @@ +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +"""Test fixtures for mlos_bench optimizers.""" + + +import pytest + +from mlos_bench.optimizers.manual_optimizer import ManualOptimizer +from mlos_bench.optimizers.mlos_core_optimizer import MlosCoreOptimizer +from mlos_bench.optimizers.mock_optimizer import MockOptimizer +from mlos_bench.tests import SEED +from mlos_bench.tunables.tunable_groups import TunableGroups + +# pylint: disable=redefined-outer-name + + +@pytest.fixture +def mock_configs() -> list[dict]: + """Mock configurations of earlier experiments.""" + return [ + { + "vmSize": "Standard_B4ms", + "idle": "halt", + "kernel_sched_migration_cost_ns": 50000, + "kernel_sched_latency_ns": 1000000, + }, + { + "vmSize": "Standard_B4ms", + "idle": "halt", + "kernel_sched_migration_cost_ns": 40000, + "kernel_sched_latency_ns": 2000000, + }, + { + "vmSize": "Standard_B4ms", + "idle": "mwait", + "kernel_sched_migration_cost_ns": -1, # Special value + "kernel_sched_latency_ns": 3000000, + }, + { + "vmSize": "Standard_B2s", + "idle": "mwait", + "kernel_sched_migration_cost_ns": 200000, + "kernel_sched_latency_ns": 4000000, + }, + ] + + +@pytest.fixture +def mock_opt_no_defaults(tunable_groups: TunableGroups) -> MockOptimizer: + """Test fixture for MockOptimizer that ignores the initial configuration.""" + return MockOptimizer( + tunables=tunable_groups, + service=None, + config={ + "optimization_targets": {"score": "min"}, + "max_suggestions": 5, + "start_with_defaults": False, + "seed": SEED, + }, + ) + + +@pytest.fixture +def mock_opt(tunable_groups: TunableGroups) -> MockOptimizer: + """Test fixture for MockOptimizer.""" + return MockOptimizer( + tunables=tunable_groups, + service=None, + config={"optimization_targets": {"score": "min"}, "max_suggestions": 5, "seed": SEED}, + ) + + +@pytest.fixture +def mock_opt_max(tunable_groups: TunableGroups) -> MockOptimizer: + """Test fixture for MockOptimizer.""" + return MockOptimizer( + tunables=tunable_groups, + service=None, + config={"optimization_targets": {"score": "max"}, "max_suggestions": 10, "seed": SEED}, + ) + + +@pytest.fixture +def flaml_opt(tunable_groups: TunableGroups) -> MlosCoreOptimizer: + """Test fixture for mlos_core FLAML optimizer.""" + return MlosCoreOptimizer( + tunables=tunable_groups, + service=None, + config={ + "optimization_targets": {"score": "min"}, + "max_suggestions": 15, + "optimizer_type": "FLAML", + "seed": SEED, + }, + ) + + +@pytest.fixture +def flaml_opt_max(tunable_groups: TunableGroups) -> MlosCoreOptimizer: + """Test fixture for mlos_core FLAML optimizer.""" + return MlosCoreOptimizer( + tunables=tunable_groups, + service=None, + config={ + "optimization_targets": {"score": "max"}, + "max_suggestions": 15, + "optimizer_type": "FLAML", + "seed": SEED, + }, + ) + + +# FIXME: SMAC's RF model can be non-deterministic at low iterations, which are +# normally calculated as a percentage of the max_suggestions and number of +# tunable dimensions, so for now we set the initial random samples equal to the +# number of iterations and control them with a seed. + +SMAC_ITERATIONS = 10 + + +@pytest.fixture +def smac_opt(tunable_groups: TunableGroups) -> MlosCoreOptimizer: + """Test fixture for mlos_core SMAC optimizer.""" + return MlosCoreOptimizer( + tunables=tunable_groups, + service=None, + config={ + "optimization_targets": {"score": "min"}, + "max_suggestions": SMAC_ITERATIONS, + "optimizer_type": "SMAC", + "seed": SEED, + "output_directory": None, + # See Above + "n_random_init": SMAC_ITERATIONS, + "max_ratio": 1.0, + }, + ) + + +@pytest.fixture +def smac_opt_max(tunable_groups: TunableGroups) -> MlosCoreOptimizer: + """Test fixture for mlos_core SMAC optimizer.""" + return MlosCoreOptimizer( + tunables=tunable_groups, + service=None, + config={ + "optimization_targets": {"score": "max"}, + "max_suggestions": SMAC_ITERATIONS, + "optimizer_type": "SMAC", + "seed": SEED, + "output_directory": None, + # See Above + "n_random_init": SMAC_ITERATIONS, + "max_ratio": 1.0, + }, + ) + + +@pytest.fixture +def manual_opt(tunable_groups: TunableGroups, mock_configs: list[dict]) -> ManualOptimizer: + """Test fixture for ManualOptimizer.""" + return ManualOptimizer( + tunables=tunable_groups, + service=None, + config={ + "max_cycles": 2, + "tunable_values_cycle": mock_configs, + }, + ) diff --git a/mlos_bench/mlos_bench/tests/storage/conftest.py b/mlos_bench/mlos_bench/tests/storage/conftest.py index a1437052823..c510793fac1 100644 --- a/mlos_bench/mlos_bench/tests/storage/conftest.py +++ b/mlos_bench/mlos_bench/tests/storage/conftest.py @@ -12,6 +12,7 @@ # Expose some of those as local names so they can be picked up as fixtures by pytest. storage = sql_storage_fixtures.storage +sqlite_storage = sql_storage_fixtures.sqlite_storage exp_storage = sql_storage_fixtures.exp_storage exp_no_tunables_storage = sql_storage_fixtures.exp_no_tunables_storage mixed_numerics_exp_storage = sql_storage_fixtures.mixed_numerics_exp_storage diff --git a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py index cb83bffd4ff..8cd4954bdda 100644 --- a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py +++ b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py @@ -7,6 +7,10 @@ from collections.abc import Generator from random import seed as rand_seed +import json +import os +import tempfile + import pytest from mlos_bench.optimizers.mock_optimizer import MockOptimizer @@ -15,6 +19,7 @@ from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.storage.base_experiment_data import ExperimentData from mlos_bench.storage.sql.storage import SqlStorage +from mlos_bench.storage.storage_factory import from_config from mlos_bench.tests import SEED from mlos_bench.tests.storage import ( CONFIG_TRIAL_REPEAT_COUNT, @@ -26,6 +31,38 @@ # pylint: disable=redefined-outer-name +@pytest.fixture +def sqlite_storage() -> Generator[SqlStorage]: + """ + Fixture for file based SQLite storage in a temporary directory. + + Yields + ------ + Generator[SqlStorage] + + Notes + ----- + Can't be used in parallel tests on Windows. + """ + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "mlos_bench.sqlite") + config_str = json.dumps( + { + "class": "mlos_bench.storage.sql.storage.SqlStorage", + "config": { + "drivername": "sqlite", + "database": db_path, + "lazy_schema_create": False, + }, + } + ) + + storage = from_config(config_str) + assert isinstance(storage, SqlStorage) + storage.update_schema() + yield storage + + @pytest.fixture def storage() -> SqlStorage: """Test fixture for in-memory SQLite3 storage.""" diff --git a/mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py b/mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py index 3d5053837be..7871e7f68ca 100644 --- a/mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py +++ b/mlos_bench/mlos_bench/tests/storage/test_storage_pickling.py @@ -3,11 +3,8 @@ # Licensed under the MIT License. # """Test pickling and unpickling of Storage, and restoring Experiment and Trial by id.""" -import json -import os import pickle import sys -import tempfile from datetime import datetime from typing import Literal @@ -16,7 +13,6 @@ from mlos_bench.environments.status import Status from mlos_bench.storage.sql.storage import SqlStorage -from mlos_bench.storage.storage_factory import from_config from mlos_bench.tunables.tunable_groups import TunableGroups @@ -26,72 +22,59 @@ sys.platform == "win32", reason="Windows doesn't support multiple processes accessing the same file.", ) -def test_storage_pickle_restore_experiment_and_trial(tunable_groups: TunableGroups) -> None: +def test_storage_pickle_restore_experiment_and_trial( + sqlite_storage: SqlStorage, + tunable_groups: TunableGroups, +) -> None: """Check that we can pickle and unpickle the Storage object, and restore Experiment and Trial by id. """ - # pylint: disable=too-many-locals - with tempfile.TemporaryDirectory() as tmpdir: - db_path = os.path.join(tmpdir, "mlos_bench.sqlite") - config_str = json.dumps( - { - "class": "mlos_bench.storage.sql.storage.SqlStorage", - "config": { - "drivername": "sqlite", - "database": db_path, - "lazy_schema_create": False, - }, - } - ) + storage = sqlite_storage + # Create an Experiment and a Trial + opt_targets: dict[str, Literal["min", "max"]] = {"metric": "min"} + experiment = storage.experiment( + experiment_id="experiment_id", + trial_id=0, + root_env_config="dummy_env.json", + description="Pickle test experiment", + tunables=tunable_groups, + opt_targets=opt_targets, + ) + with experiment: + trial = experiment.new_trial(tunable_groups) + trial_id_created = trial.trial_id + trial.set_trial_runner(1) + trial.update(Status.RUNNING, datetime.now(UTC)) - storage = from_config(config_str) - storage.update_schema() + # Pickle and unpickle the Storage object + pickled = pickle.dumps(storage) + restored_storage = pickle.loads(pickled) + assert isinstance(restored_storage, SqlStorage) - # Create an Experiment and a Trial - opt_targets: dict[str, Literal["min", "max"]] = {"metric": "min"} - experiment = storage.experiment( - experiment_id="experiment_id", - trial_id=0, - root_env_config="dummy_env.json", - description="Pickle test experiment", - tunables=tunable_groups, - opt_targets=opt_targets, - ) - with experiment: - trial = experiment.new_trial(tunable_groups) - trial_id_created = trial.trial_id - trial.set_trial_runner(1) - trial.update(Status.RUNNING, datetime.now(UTC)) + # Restore the Experiment from storage by id and check that it matches the original + restored_experiment = restored_storage.get_experiment_by_id( + experiment_id=experiment.experiment_id, + tunables=tunable_groups, + opt_targets=opt_targets, + ) + assert restored_experiment is not None + assert restored_experiment is not experiment + assert restored_experiment.experiment_id == experiment.experiment_id + assert restored_experiment.description == experiment.description + assert restored_experiment.root_env_config == experiment.root_env_config + assert restored_experiment.tunables == experiment.tunables + assert restored_experiment.opt_targets == experiment.opt_targets + with restored_experiment: + # trial_id should have been restored during __enter__ + assert restored_experiment.trial_id == experiment.trial_id - # Pickle and unpickle the Storage object - pickled = pickle.dumps(storage) - restored_storage = pickle.loads(pickled) - assert isinstance(restored_storage, SqlStorage) - - # Restore the Experiment from storage by id and check that it matches the original - restored_experiment = restored_storage.get_experiment_by_id( - experiment_id=experiment.experiment_id, - tunables=tunable_groups, - opt_targets=opt_targets, - ) - assert restored_experiment is not None - assert restored_experiment is not experiment - assert restored_experiment.experiment_id == experiment.experiment_id - assert restored_experiment.description == experiment.description - assert restored_experiment.root_env_config == experiment.root_env_config - assert restored_experiment.tunables == experiment.tunables - assert restored_experiment.opt_targets == experiment.opt_targets - with restored_experiment: - # trial_id should have been restored during __enter__ - assert restored_experiment.trial_id == experiment.trial_id - - # Restore the Trial from storage by id and check that it matches the original - restored_trial = restored_experiment.get_trial_by_id(trial_id_created) - assert restored_trial is not None - assert restored_trial is not trial - assert restored_trial.trial_id == trial.trial_id - assert restored_trial.experiment_id == trial.experiment_id - assert restored_trial.tunables == trial.tunables - assert restored_trial.status == trial.status - assert restored_trial.config() == trial.config() - assert restored_trial.trial_runner_id == trial.trial_runner_id + # Restore the Trial from storage by id and check that it matches the original + restored_trial = restored_experiment.get_trial_by_id(trial_id_created) + assert restored_trial is not None + assert restored_trial is not trial + assert restored_trial.trial_id == trial.trial_id + assert restored_trial.experiment_id == trial.experiment_id + assert restored_trial.tunables == trial.tunables + assert restored_trial.status == trial.status + assert restored_trial.config() == trial.config() + assert restored_trial.trial_runner_id == trial.trial_runner_id From ea7b3ff26d0093fced5c56cad236ad8695cd0278 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 19 May 2025 22:46:51 +0000 Subject: [PATCH 09/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../mlos_bench/tests/config/schedulers/__init__.py | 8 +++++--- .../mlos_bench/tests/config/schedulers/conftest.py | 10 +++++++--- .../schedulers/test_load_scheduler_config_examples.py | 5 ++--- mlos_bench/mlos_bench/tests/storage/sql/fixtures.py | 5 ++--- 4 files changed, 16 insertions(+), 12 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/__init__.py b/mlos_bench/mlos_bench/tests/config/schedulers/__init__.py index 7838135124c..111238e6ac9 100644 --- a/mlos_bench/mlos_bench/tests/config/schedulers/__init__.py +++ b/mlos_bench/mlos_bench/tests/config/schedulers/__init__.py @@ -1,3 +1,5 @@ -""" -Unit tests for the mlos_bench Scheduler configs. -""" +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +"""Unit tests for the mlos_bench Scheduler configs.""" diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py index a1fd105839b..fdf27162b35 100644 --- a/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py +++ b/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py @@ -5,25 +5,28 @@ """ Pytest fixtures for Scheduler config tests. -Provides fixtures for creating multiple TrialRunner instances using the mock environment config. +Provides fixtures for creating multiple TrialRunner instances using the mock environment +config. """ from importlib.resources import files import pytest -from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.schedulers.trial_runner import TrialRunner +from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.util import path_join # pylint: disable=redefined-outer-name TRIAL_RUNNERS_COUNT = 4 + @pytest.fixture def mock_env_config_path() -> str: """ Returns the absolute path to the mock environment configuration file. + This file is used to create TrialRunner instances for testing. """ @@ -40,7 +43,8 @@ def trial_runners( mock_env_config_path: str, ) -> list[TrialRunner]: """ - Fixture that returns a list of TrialRunner instances using the mock environment config. + Fixture that returns a list of TrialRunner instances using the mock environment + config. Returns ------- diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py index c8d6bcb1726..bf10a99c46b 100644 --- a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py @@ -7,6 +7,8 @@ import pytest +import mlos_bench.tests.optimizers.fixtures +import mlos_bench.tests.storage.sql.fixtures from mlos_bench.config.schemas.config_schemas import ConfigSchema from mlos_bench.optimizers.mock_optimizer import MockOptimizer from mlos_bench.schedulers.base_scheduler import Scheduler @@ -16,9 +18,6 @@ from mlos_bench.tests.config import locate_config_examples from mlos_bench.util import get_class_from_name -import mlos_bench.tests.storage.sql.fixtures -import mlos_bench.tests.optimizers.fixtures - mock_opt = mlos_bench.tests.optimizers.fixtures.mock_opt sqlite_storage = mlos_bench.tests.storage.sql.fixtures.sqlite_storage diff --git a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py index 8cd4954bdda..0bebeeff824 100644 --- a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py +++ b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py @@ -4,12 +4,11 @@ # """Test fixtures for mlos_bench storage.""" -from collections.abc import Generator -from random import seed as rand_seed - import json import os import tempfile +from collections.abc import Generator +from random import seed as rand_seed import pytest From 7563b8aeef9c92c166f72315975fe1caf10acd61 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 17:57:26 -0500 Subject: [PATCH 10/77] fixup --- .../mlos_bench/config/schemas/schedulers/scheduler-schema.json | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json index 8dfce6be369..3086abacd74 100644 --- a/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json @@ -20,5 +20,6 @@ ] } ], - "required": ["class"] + "required": ["class"], + "unevaluatedProperties": false } From b03e31c6aa5736f7fc2a0ef3fb8cbb4a67e142c0 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 18:01:24 -0500 Subject: [PATCH 11/77] lint --- mlos_bench/mlos_bench/tests/config/schedulers/conftest.py | 1 - .../config/schedulers/test_load_scheduler_config_examples.py | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py index fdf27162b35..71368400561 100644 --- a/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py +++ b/mlos_bench/mlos_bench/tests/config/schedulers/conftest.py @@ -29,7 +29,6 @@ def mock_env_config_path() -> str: This file is used to create TrialRunner instances for testing. """ - # Use the files() routine to locate the file relative to this directory return path_join( str(files("mlos_bench.config").joinpath("environments", "mock", "mock_env.jsonc")), diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py index bf10a99c46b..a2d90a3d019 100644 --- a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py @@ -54,6 +54,7 @@ def test_load_scheduler_config_examples( mock_opt: MockOptimizer, ) -> None: """Tests loading a config example.""" + # pylint: disable=too-many-arguments,too-many-positional-arguments config = config_loader_service.load_config(config_path, ConfigSchema.SCHEDULER) assert isinstance(config, dict) cls = get_class_from_name(config["class"]) From 8eb14c0c815fa803a69ec3c1fb456ab3e8956d2b Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 18:02:39 -0500 Subject: [PATCH 12/77] apply suggestions --- mlos_bench/mlos_bench/tests/storage/sql/fixtures.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py index 0bebeeff824..1b5892e235a 100644 --- a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py +++ b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py @@ -6,6 +6,7 @@ import json import os +import sys import tempfile from collections.abc import Generator from random import seed as rand_seed @@ -30,6 +31,10 @@ # pylint: disable=redefined-outer-name +@pytest.mark.skipif( + sys.platform == "win32", + reason="File-based SQLite storage can fail on Windows due to file-locking in parallel tests.", +) @pytest.fixture def sqlite_storage() -> Generator[SqlStorage]: """ From 8ad4c3db78f491b1ebd65b61a4c630219006c381 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 18:03:15 -0500 Subject: [PATCH 13/77] revert --- mlos_bench/mlos_bench/tests/storage/sql/fixtures.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py index 1b5892e235a..0bebeeff824 100644 --- a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py +++ b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py @@ -6,7 +6,6 @@ import json import os -import sys import tempfile from collections.abc import Generator from random import seed as rand_seed @@ -31,10 +30,6 @@ # pylint: disable=redefined-outer-name -@pytest.mark.skipif( - sys.platform == "win32", - reason="File-based SQLite storage can fail on Windows due to file-locking in parallel tests.", -) @pytest.fixture def sqlite_storage() -> Generator[SqlStorage]: """ From d4d5153715dfad0f620a57e4b1859bfbbd885321 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 18:17:36 -0500 Subject: [PATCH 14/77] load test configs too --- .../test_load_environment_config_examples.py | 10 +++++++++- .../config/experiments/experiment_test_config.jsonc | 4 ++++ .../optimizers/test_load_optimizer_config_examples.py | 10 +++++++++- .../schedulers/test_load_scheduler_config_examples.py | 10 +++++++++- .../services/test_load_service_config_examples.py | 10 +++++++++- .../storage/test_load_storage_config_examples.py | 9 ++++++++- 6 files changed, 48 insertions(+), 5 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/config/environments/test_load_environment_config_examples.py b/mlos_bench/mlos_bench/tests/config/environments/test_load_environment_config_examples.py index 889462f024d..064530919c3 100644 --- a/mlos_bench/mlos_bench/tests/config/environments/test_load_environment_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/environments/test_load_environment_config_examples.py @@ -11,7 +11,7 @@ from mlos_bench.environments.base_environment import Environment from mlos_bench.environments.composite_env import CompositeEnv from mlos_bench.services.config_persistence import ConfigPersistenceService -from mlos_bench.tests.config import locate_config_examples +from mlos_bench.tests.config import locate_config_examples, BUILTIN_TEST_CONFIG_PATH from mlos_bench.tunables.tunable_groups import TunableGroups _LOG = logging.getLogger(__name__) @@ -39,6 +39,14 @@ def filter_configs(configs_to_filter: list[str]) -> list[str]: ) assert configs +test_configs = locate_config_examples( + BUILTIN_TEST_CONFIG_PATH, + CONFIG_TYPE, + filter_configs, +) +assert test_configs +configs.extend(test_configs) + @pytest.mark.parametrize("config_path", configs) def test_load_environment_config_examples( diff --git a/mlos_bench/mlos_bench/tests/config/experiments/experiment_test_config.jsonc b/mlos_bench/mlos_bench/tests/config/experiments/experiment_test_config.jsonc index 2ca87c6f215..c6f98c4963a 100644 --- a/mlos_bench/mlos_bench/tests/config/experiments/experiment_test_config.jsonc +++ b/mlos_bench/mlos_bench/tests/config/experiments/experiment_test_config.jsonc @@ -15,6 +15,10 @@ "resourceGroup": "mlos-autotuning-test-rg", "location": "eastus", "vmName": "vmTestName", + "ssh_username": "testuser", + "ssh_priv_key_path": "/home/testuser/.ssh/id_rsa", + "ssh_hostname": "${vmName}", + "ssh_port": 22, "tunable_params_map": { "linux-runtime": ["linux-scheduler", "linux-swap"], "linux-boot": ["linux-kernel-boot"], diff --git a/mlos_bench/mlos_bench/tests/config/optimizers/test_load_optimizer_config_examples.py b/mlos_bench/mlos_bench/tests/config/optimizers/test_load_optimizer_config_examples.py index fceecd89f0d..a507a3f01dd 100644 --- a/mlos_bench/mlos_bench/tests/config/optimizers/test_load_optimizer_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/optimizers/test_load_optimizer_config_examples.py @@ -10,7 +10,7 @@ from mlos_bench.config.schemas import ConfigSchema from mlos_bench.optimizers.base_optimizer import Optimizer from mlos_bench.services.config_persistence import ConfigPersistenceService -from mlos_bench.tests.config import locate_config_examples +from mlos_bench.tests.config import locate_config_examples, BUILTIN_TEST_CONFIG_PATH from mlos_bench.tunables.tunable_groups import TunableGroups from mlos_bench.util import get_class_from_name @@ -34,6 +34,14 @@ def filter_configs(configs_to_filter: list[str]) -> list[str]: ) assert configs +test_configs = locate_config_examples( + BUILTIN_TEST_CONFIG_PATH, + CONFIG_TYPE, + filter_configs, +) +# assert test_configs +configs.extend(test_configs) + @pytest.mark.parametrize("config_path", configs) def test_load_optimizer_config_examples( diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py index a2d90a3d019..6e676e920f4 100644 --- a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py @@ -15,7 +15,7 @@ from mlos_bench.schedulers.trial_runner import TrialRunner from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.storage.sql.storage import SqlStorage -from mlos_bench.tests.config import locate_config_examples +from mlos_bench.tests.config import locate_config_examples, BUILTIN_TEST_CONFIG_PATH from mlos_bench.util import get_class_from_name mock_opt = mlos_bench.tests.optimizers.fixtures.mock_opt @@ -43,6 +43,14 @@ def filter_configs(configs_to_filter: list[str]) -> list[str]: ) assert configs +test_configs = locate_config_examples( + BUILTIN_TEST_CONFIG_PATH, + CONFIG_TYPE, + filter_configs, +) +# assert test_configs +configs.extend(test_configs) + @pytest.mark.parametrize("config_path", configs) def test_load_scheduler_config_examples( diff --git a/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py b/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py index beb0b1d018e..84e7a1128c7 100644 --- a/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py @@ -10,7 +10,7 @@ from mlos_bench.config.schemas.config_schemas import ConfigSchema from mlos_bench.services.base_service import Service from mlos_bench.services.config_persistence import ConfigPersistenceService -from mlos_bench.tests.config import locate_config_examples +from mlos_bench.tests.config import locate_config_examples, BUILTIN_TEST_CONFIG_PATH _LOG = logging.getLogger(__name__) _LOG.setLevel(logging.DEBUG) @@ -40,6 +40,14 @@ def predicate(config_path: str) -> bool: ) assert configs +test_configs = locate_config_examples( + BUILTIN_TEST_CONFIG_PATH, + CONFIG_TYPE, + filter_configs, +) +assert test_configs +configs.extend(test_configs) + @pytest.mark.parametrize("config_path", configs) def test_load_service_config_examples( diff --git a/mlos_bench/mlos_bench/tests/config/storage/test_load_storage_config_examples.py b/mlos_bench/mlos_bench/tests/config/storage/test_load_storage_config_examples.py index e3696a85fad..38582707cfc 100644 --- a/mlos_bench/mlos_bench/tests/config/storage/test_load_storage_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/storage/test_load_storage_config_examples.py @@ -10,7 +10,7 @@ from mlos_bench.config.schemas.config_schemas import ConfigSchema from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.storage.base_storage import Storage -from mlos_bench.tests.config import locate_config_examples +from mlos_bench.tests.config import locate_config_examples, BUILTIN_TEST_CONFIG_PATH from mlos_bench.util import get_class_from_name _LOG = logging.getLogger(__name__) @@ -33,6 +33,13 @@ def filter_configs(configs_to_filter: list[str]) -> list[str]: ) assert configs +test_configs = locate_config_examples( + BUILTIN_TEST_CONFIG_PATH, + CONFIG_TYPE, + filter_configs, +) +# assert test_configs +configs.extend(test_configs) @pytest.mark.parametrize("config_path", configs) def test_load_storage_config_examples( From 1eb1acb7af53309948178cbe0ac1747793881660 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 18:18:03 -0500 Subject: [PATCH 15/77] format --- .../environments/test_load_environment_config_examples.py | 2 +- .../config/optimizers/test_load_optimizer_config_examples.py | 2 +- .../config/schedulers/test_load_scheduler_config_examples.py | 2 +- .../tests/config/services/test_load_service_config_examples.py | 2 +- .../tests/config/storage/test_load_storage_config_examples.py | 3 ++- 5 files changed, 6 insertions(+), 5 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/config/environments/test_load_environment_config_examples.py b/mlos_bench/mlos_bench/tests/config/environments/test_load_environment_config_examples.py index 064530919c3..4b35282e53c 100644 --- a/mlos_bench/mlos_bench/tests/config/environments/test_load_environment_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/environments/test_load_environment_config_examples.py @@ -11,7 +11,7 @@ from mlos_bench.environments.base_environment import Environment from mlos_bench.environments.composite_env import CompositeEnv from mlos_bench.services.config_persistence import ConfigPersistenceService -from mlos_bench.tests.config import locate_config_examples, BUILTIN_TEST_CONFIG_PATH +from mlos_bench.tests.config import BUILTIN_TEST_CONFIG_PATH, locate_config_examples from mlos_bench.tunables.tunable_groups import TunableGroups _LOG = logging.getLogger(__name__) diff --git a/mlos_bench/mlos_bench/tests/config/optimizers/test_load_optimizer_config_examples.py b/mlos_bench/mlos_bench/tests/config/optimizers/test_load_optimizer_config_examples.py index a507a3f01dd..a407275438b 100644 --- a/mlos_bench/mlos_bench/tests/config/optimizers/test_load_optimizer_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/optimizers/test_load_optimizer_config_examples.py @@ -10,7 +10,7 @@ from mlos_bench.config.schemas import ConfigSchema from mlos_bench.optimizers.base_optimizer import Optimizer from mlos_bench.services.config_persistence import ConfigPersistenceService -from mlos_bench.tests.config import locate_config_examples, BUILTIN_TEST_CONFIG_PATH +from mlos_bench.tests.config import BUILTIN_TEST_CONFIG_PATH, locate_config_examples from mlos_bench.tunables.tunable_groups import TunableGroups from mlos_bench.util import get_class_from_name diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py index 6e676e920f4..2bc46122573 100644 --- a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py @@ -15,7 +15,7 @@ from mlos_bench.schedulers.trial_runner import TrialRunner from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.storage.sql.storage import SqlStorage -from mlos_bench.tests.config import locate_config_examples, BUILTIN_TEST_CONFIG_PATH +from mlos_bench.tests.config import BUILTIN_TEST_CONFIG_PATH, locate_config_examples from mlos_bench.util import get_class_from_name mock_opt = mlos_bench.tests.optimizers.fixtures.mock_opt diff --git a/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py b/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py index 84e7a1128c7..96df98b29d2 100644 --- a/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/services/test_load_service_config_examples.py @@ -10,7 +10,7 @@ from mlos_bench.config.schemas.config_schemas import ConfigSchema from mlos_bench.services.base_service import Service from mlos_bench.services.config_persistence import ConfigPersistenceService -from mlos_bench.tests.config import locate_config_examples, BUILTIN_TEST_CONFIG_PATH +from mlos_bench.tests.config import BUILTIN_TEST_CONFIG_PATH, locate_config_examples _LOG = logging.getLogger(__name__) _LOG.setLevel(logging.DEBUG) diff --git a/mlos_bench/mlos_bench/tests/config/storage/test_load_storage_config_examples.py b/mlos_bench/mlos_bench/tests/config/storage/test_load_storage_config_examples.py index 38582707cfc..680b3bacf1f 100644 --- a/mlos_bench/mlos_bench/tests/config/storage/test_load_storage_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/storage/test_load_storage_config_examples.py @@ -10,7 +10,7 @@ from mlos_bench.config.schemas.config_schemas import ConfigSchema from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.storage.base_storage import Storage -from mlos_bench.tests.config import locate_config_examples, BUILTIN_TEST_CONFIG_PATH +from mlos_bench.tests.config import BUILTIN_TEST_CONFIG_PATH, locate_config_examples from mlos_bench.util import get_class_from_name _LOG = logging.getLogger(__name__) @@ -41,6 +41,7 @@ def filter_configs(configs_to_filter: list[str]) -> list[str]: # assert test_configs configs.extend(test_configs) + @pytest.mark.parametrize("config_path", configs) def test_load_storage_config_examples( config_loader_service: ConfigPersistenceService, From 1112af4251c7ecaac4b9879e4a095b9c425cddfd Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 18:35:17 -0500 Subject: [PATCH 16/77] list in __all__ so we load it as a part of tests schemas checking --- mlos_bench/mlos_bench/tests/schedulers/__init__.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/mlos_bench/mlos_bench/tests/schedulers/__init__.py b/mlos_bench/mlos_bench/tests/schedulers/__init__.py index 4bc0076079f..b1668588752 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/__init__.py +++ b/mlos_bench/mlos_bench/tests/schedulers/__init__.py @@ -3,3 +3,8 @@ # Licensed under the MIT License. # """mlos_bench.tests.schedulers.""" +from mlos_bench.tests.schedulers.mock_scheduler import MockScheduler + +__all__ = [ + "MockScheduler", +] From 7a0d087a74dfec30f4b659634de10f74e7136eea Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 18:52:05 -0500 Subject: [PATCH 17/77] new mock scheduler schema and test configs --- .../schedulers/base-scheduler-subschema.json | 2 +- .../schedulers/mock-scheduler-subschema.json | 31 +++++++++++++++++-- ...-scheduler-invalid-mock-trial-data-2.jsonc | 11 +++++++ ...ck-scheduler-invalid-mock-trial-data.jsonc | 14 +++++++++ .../bad/unhandled/mock_sched-extra.jsonc | 7 +++++ .../good/full/mock_sched-full.jsonc | 31 +++++++++++++++++++ .../good/partial/mock_sched-partial.jsonc | 20 ++++++++++++ 7 files changed, 112 insertions(+), 4 deletions(-) create mode 100644 mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data-2.jsonc create mode 100644 mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data.jsonc create mode 100644 mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/unhandled/mock_sched-extra.jsonc create mode 100644 mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/full/mock_sched-full.jsonc create mode 100644 mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/partial/mock_sched-partial.jsonc diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json index 702da1eec3e..9417b7d00ac 100644 --- a/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json +++ b/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": { diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json index a784145c694..56d75900946 100644 --- a/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json @@ -25,8 +25,15 @@ "items": { "type": "object", "properties": { + "comments": { + "type": "string", + "description": "Optional comments about the trial status being reported." + }, "trial_id": { - "type": "string" + "type": "integer", + "description": "Unique identifier for the trial.", + "examples": [1, 2, 3], + "minimum": 1 }, "status": { "enum": [ @@ -41,9 +48,27 @@ "TIMED_OUT" ] }, - "score": "number" + "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"] + "required": ["trial_id", "status"], + "additionalProperties": false } } }, diff --git a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data-2.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data-2.jsonc new file mode 100644 index 00000000000..09b74f43773 --- /dev/null +++ b/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" + } + ] + } +} diff --git a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data.jsonc new file mode 100644 index 00000000000..144e7f21e64 --- /dev/null +++ b/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" + } + ] + } +} diff --git a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/unhandled/mock_sched-extra.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/unhandled/mock_sched-extra.jsonc new file mode 100644 index 00000000000..c7bf9a46f75 --- /dev/null +++ b/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" +} diff --git a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/full/mock_sched-full.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/full/mock_sched-full.jsonc new file mode 100644 index 00000000000..ca6ad2f9054 --- /dev/null +++ b/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" + } + } + ] + } +} diff --git a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/partial/mock_sched-partial.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/partial/mock_sched-partial.jsonc new file mode 100644 index 00000000000..2557c25cb46 --- /dev/null +++ b/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 + } + ] + } +} From dded2435f85a4695654b3cd35f6de9a3dd892293 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 19:03:24 -0500 Subject: [PATCH 18/77] refactor to split steps for easier hooking --- .../mlos_bench/schedulers/base_scheduler.py | 42 +++++++++++++------ 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/mlos_bench/mlos_bench/schedulers/base_scheduler.py b/mlos_bench/mlos_bench/schedulers/base_scheduler.py index 1cd88fd5859..3e47841591f 100644 --- a/mlos_bench/mlos_bench/schedulers/base_scheduler.py +++ b/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). From f1fe022741a005db8618b8947d6000b52984902d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 20 May 2025 01:00:04 +0000 Subject: [PATCH 19/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../config/schedulers/test_load_scheduler_config_examples.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py index dbcbf4a9fa5..e901c456f52 100644 --- a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py @@ -23,7 +23,6 @@ storage = mlos_bench.tests.storage.sql.fixtures.storage - _LOG = logging.getLogger(__name__) _LOG.setLevel(logging.DEBUG) @@ -60,9 +59,7 @@ def test_load_scheduler_config_examples( config_path: str, mock_env_config_path: str, trial_runners: list[TrialRunner], - storage: SqlStorage, - mock_opt: MockOptimizer, ) -> None: """Tests loading a config example.""" @@ -82,9 +79,7 @@ def test_load_scheduler_config_examples( global_config=global_config, trial_runners=trial_runners, optimizer=mock_opt, - storage=storage, - root_env_config=mock_env_config_path, ) assert scheduler_inst is not None From 83ff70f8fa38bbcea0cfd7c3f6a6352e87e0bb9b Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 20:02:09 -0500 Subject: [PATCH 20/77] Update mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py --- .../config/schedulers/test_load_scheduler_config_examples.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py index e901c456f52..ad8f9248acd 100644 --- a/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py +++ b/mlos_bench/mlos_bench/tests/config/schedulers/test_load_scheduler_config_examples.py @@ -19,7 +19,6 @@ 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 From 71420e65115b5bd162aad9e26ca8142fa10c73ed Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Mon, 19 May 2025 20:30:31 -0500 Subject: [PATCH 21/77] fixup --- mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py index 780b6210469..b3e6f6cff9e 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py @@ -11,7 +11,8 @@ from pytz import UTC -from mlos_bench.schedulers.base_scheduler import Optimizer, Scheduler +from mlos_bench.schedulers.base_scheduler import Scheduler +from mlos_bench.optimizers.base_optimizer import Optimizer from mlos_bench.schedulers.trial_runner import TrialRunner from mlos_bench.storage.base_storage import Storage From 70130babeed5ae4fcbf0a845cfebdb7696d47d00 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 20 May 2025 01:30:52 +0000 Subject: [PATCH 22/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py index b3e6f6cff9e..43ffff0c34e 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py @@ -11,8 +11,8 @@ from pytz import UTC -from mlos_bench.schedulers.base_scheduler import Scheduler from mlos_bench.optimizers.base_optimizer import Optimizer +from mlos_bench.schedulers.base_scheduler import Scheduler from mlos_bench.schedulers.trial_runner import TrialRunner from mlos_bench.storage.base_storage import Storage From 0a440a1575c48238ee19376d2a4a783af5461933 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Tue, 20 May 2025 14:45:06 -0500 Subject: [PATCH 23/77] ignore the build tree in vscode --- .vscode/settings.json | 1 + 1 file changed, 1 insertion(+) diff --git a/.vscode/settings.json b/.vscode/settings.json index a01d1d6418f..f29dec766e2 100644 --- a/.vscode/settings.json +++ b/.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, From efafc6f959f4bdab134b13da552195a90d388412 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 22:33:08 -0500 Subject: [PATCH 24/77] adding more accessors --- .../environments/base_environment.py | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) diff --git a/mlos_bench/mlos_bench/environments/base_environment.py b/mlos_bench/mlos_bench/environments/base_environment.py index 094085c78b5..98db5cfe14d 100644 --- a/mlos_bench/mlos_bench/environments/base_environment.py +++ b/mlos_bench/mlos_bench/environments/base_environment.py @@ -363,6 +363,91 @@ 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)", + val, + type(val), + ) + 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 + + def experiment_id(self) -> int: + """ + 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 : int + 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, int), ( + "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 From 842d393ae2bc5041060f6a78cc7ed07a4758465d Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 22:34:02 -0500 Subject: [PATCH 25/77] wip: enable mock env to report arbitrary data --- .../mlos_bench/environments/mock_env.py | 211 +++++++++++++++++- 1 file changed, 206 insertions(+), 5 deletions(-) diff --git a/mlos_bench/mlos_bench/environments/mock_env.py b/mlos_bench/mlos_bench/environments/mock_env.py index ac6d9b7f001..e15dcffed44 100644 --- a/mlos_bench/mlos_bench/environments/mock_env.py +++ b/mlos_bench/mlos_bench/environments/mock_env.py @@ -6,8 +6,10 @@ import logging import random +from dataclasses import dataclass from datetime import datetime from typing import Any +import time import numpy @@ -21,6 +23,131 @@ _LOG = logging.getLogger(__name__) +@dataclass +class MockTrialPhaseData: + """Mock trial data for a specific phase of a trial.""" + + phase: str + """Phase of the trial data (e.g., setup, run, status, teardown).""" + + status: Status + """Status response for the phase.""" + + metrics: dict[str, TunableValue] | None = None + """Metrics response for the phase.""" + + sleep: float | None = 0.0 + """Optional sleep time in seconds to simulate phase execution time.""" + + exception: str | None = None + """Message of an exception to raise for the phase.""" + + @staticmethod + def from_dict(phase: str, data: dict | None) -> "MockTrialPhaseData": + """ + Create a MockTrialPhaseData instance from a dictionary. + + Parameters + ---------- + phase : str + Phase of the trial data. + data : dict | None + Dictionary containing the phase data. + + Returns + ------- + MockTrialPhaseData + Instance of MockTrialPhaseData. + """ + data = data or {} + assert phase in {"setup", "run", "status", "teardown"}, f"Invalid phase: {phase}" + if phase in {"teardown", "status"}: + # setup/teardown phase is not expected to have metrics or status. + assert "metrics" not in data, f"Unexpected metrics data in {phase} phase: {data}" + assert "status" not in data, f"Unexpected status data in {phase} phase: {data}" + if "sleep" in data: + assert isinstance( + data["sleep"], (int, float) + ), f"Invalid sleep in {phase} phase: {data}" + assert 60 >= data["sleep"] >= 0, f"Invalid sleep time in {phase} phase: {data}" + if "metrics" in data: + assert isinstance(data["metrics"], dict), f"Invalid metrics in {phase} phase: {data}" + default_phases = { + "run": Status.SUCCEEDED, + # FIXME: this causes issues if we report RUNNING instead of READY + "status": Status.READY, + } + status = Status.parse(data.get("status", default_phases.get(phase, Status.UNKNOWN))) + return MockTrialPhaseData( + phase=phase, + status=status, + metrics=data.get("metrics"), + sleep=data.get("sleep"), + exception=data.get("exception"), + ) + + +@dataclass +class MockTrialData: + """Mock trial data for a specific trial ID.""" + + trial_id: int + """Trial ID for the mock trial data.""" + setup: MockTrialPhaseData + """Setup phase data for the trial.""" + run: MockTrialPhaseData + """Run phase data for the trial.""" + status: MockTrialPhaseData + """Status phase data for the trial.""" + teardown: MockTrialPhaseData + """Teardown phase data for the trial.""" + + @staticmethod + def from_dict(trial_id: int, data: dict) -> "MockTrialData": + """ + Create a MockTrialData instance from a dictionary. + + Parameters + ---------- + trial_id : int + Trial ID for the mock trial data. + data : dict + Dictionary containing the trial data. + + Returns + ------- + MockTrialData + Instance of MockTrialData. + """ + return MockTrialData( + trial_id=trial_id, + setup=MockTrialPhaseData.from_dict("setup", data.get("setup")), + run=MockTrialPhaseData.from_dict("run", data.get("run")), + status=MockTrialPhaseData.from_dict("status", data.get("status")), + teardown=MockTrialPhaseData.from_dict("teardown", data.get("teardown")), + ) + + @staticmethod + def load_mock_trial_data(mock_trial_data: dict) -> dict[int, "MockTrialData"]: + """ + Load mock trial data from a dictionary. + + Parameters + ---------- + mock_trial_data : dict + Dictionary containing mock trial data. + + Returns + ------- + dict[int, MockTrialData] + Dictionary of mock trial data keyed by trial ID. + """ + return { + int(trial_id): MockTrialData.from_dict(trial_id=int(trial_id), data=trial_data) + for trial_id, trial_data in mock_trial_data.items() + } + + class MockEnv(Environment): """Scheduler-side environment to mock the benchmark results.""" @@ -55,6 +182,19 @@ def __init__( # pylint: disable=too-many-arguments service: Service An optional service object. Not used by this class. """ + # First allow merging mock_trial_data from the global_config into the + # config so we can check it against the JSON schema for expected data + # types. + if global_config and "mock_trial_data" in global_config: + mock_trial_data = global_config["mock_trial_data"] + if not isinstance(mock_trial_data, dict): + raise ValueError(f"Invalid mock_trial_data in global_config: {mock_trial_data}") + # Merge the mock trial data into the config. + config["mock_trial_data"] = { + **config.get("mock_trial_data", {}), + **mock_trial_data, + } + super().__init__( name=name, config=config, @@ -62,6 +202,9 @@ def __init__( # pylint: disable=too-many-arguments tunables=tunables, service=service, ) + self._mock_trial_data = MockTrialData.load_mock_trial_data( + self.config.get("mock_trial_data", {}) + ) seed = int(self.config.get("mock_env_seed", -1)) self._run_random = random.Random(seed or None) if seed >= 0 else None self._status_random = random.Random(seed or None) if seed >= 0 else None @@ -83,6 +226,43 @@ def _produce_metrics(self, rand: random.Random | None) -> dict[str, TunableValue return {metric: float(score) for metric in self._metrics or []} + def get_current_mock_trial_data(self) -> MockTrialData: + """Gets mock trial data for the current trial ID. + + If no (or missing) mock trial data is found, a new instance of + MockTrialData is created from random data. + + Note + ---- + This method must be called after the base :py:meth:`Environment.setup` + method is called to ensure the current ``trial_id`` is set. + """ + trial_id = self.current_trial_id + mock_trial_data = self._mock_trial_data.get(trial_id) + if not mock_trial_data: + mock_trial_data = MockTrialData( + trial_id=trial_id, + setup=MockTrialPhaseData.from_dict(phase="setup", data=None), + run=MockTrialPhaseData.from_dict(phase="run", data=None), + status=MockTrialPhaseData.from_dict(phase="status", data=None), + teardown=MockTrialPhaseData.from_dict(phase="teardown", data=None), + ) + # Save the generated data for later. + self._mock_trial_data[trial_id] = mock_trial_data + return mock_trial_data + + def setup(self, tunables: TunableGroups, global_config: dict | None = None) -> bool: + is_success = super().setup(tunables, global_config) + mock_trial_data = self.get_current_mock_trial_data() + if mock_trial_data.setup.sleep: + _LOG.debug("Sleeping for %s seconds", mock_trial_data.setup.sleep) + time.sleep(mock_trial_data.setup.sleep) + if mock_trial_data.setup.exception: + raise RuntimeError( + f"Mock trial data setup exception: {mock_trial_data.setup.exception}" + ) + return is_success and mock_trial_data.setup.status.is_ready() + def run(self) -> tuple[Status, datetime, dict[str, TunableValue] | None]: """ Produce mock benchmark data for one experiment. @@ -99,8 +279,17 @@ def run(self) -> tuple[Status, datetime, dict[str, TunableValue] | None]: (status, timestamp, _) = result = super().run() if not status.is_ready(): return result - metrics = self._produce_metrics(self._run_random) - return (Status.SUCCEEDED, timestamp, metrics) + mock_trial_data = self.get_current_mock_trial_data() + if mock_trial_data.run.sleep: + _LOG.debug("Sleeping for %s seconds", mock_trial_data.run.sleep) + time.sleep(mock_trial_data.run.sleep) + if mock_trial_data.run.exception: + raise RuntimeError(f"Mock trial data run exception: {mock_trial_data.run.exception}") + + if mock_trial_data.run.metrics is None: + # If no metrics are provided, generate them. + mock_trial_data.run.metrics = self._produce_metrics(self._run_random) + return (mock_trial_data.run.status, timestamp, mock_trial_data.run.metrics) def status(self) -> tuple[Status, datetime, list[tuple[datetime, str, Any]]]: """ @@ -116,10 +305,22 @@ def status(self) -> tuple[Status, datetime, list[tuple[datetime, str, Any]]]: (status, timestamp, _) = result = super().status() if not status.is_ready(): return result - metrics = self._produce_metrics(self._status_random) + mock_trial_data = self.get_current_mock_trial_data() + if mock_trial_data.status.sleep: + _LOG.debug("Sleeping for %s seconds", mock_trial_data.status.sleep) + time.sleep(mock_trial_data.status.sleep) + if mock_trial_data.status.exception: + raise RuntimeError( + f"Mock trial data status exception: {mock_trial_data.status.exception}" + ) + if mock_trial_data.status.metrics is None: + # If no metrics are provided, generate them. + metrics = self._produce_metrics(self._status_random) + else: + # If metrics are provided, use them. + metrics = mock_trial_data.status.metrics return ( - # FIXME: this causes issues if we report RUNNING instead of READY - Status.READY, + mock_trial_data.status.status, timestamp, [(timestamp, metric, score) for (metric, score) in metrics.items()], ) From 19cec783a952ecbd83a3a63fcf43ec31db24173e Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 22:34:10 -0500 Subject: [PATCH 26/77] spelling --- mlos_bench/mlos_bench/environments/script_env.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/environments/script_env.py b/mlos_bench/mlos_bench/environments/script_env.py index 6ac4674cfe1..d71eb661834 100644 --- a/mlos_bench/mlos_bench/environments/script_env.py +++ b/mlos_bench/mlos_bench/environments/script_env.py @@ -5,7 +5,7 @@ """ Base scriptable benchmark environment. -TODO: Document how variable propogation works in the script environments using +TODO: Document how variable propagation works in the script environments using shell_env_params, required_args, const_args, etc. """ From 175f315c6fd9638c5122256c324a96007837bb70 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 22:34:47 -0500 Subject: [PATCH 27/77] refactor status parsing a little bit again --- mlos_bench/mlos_bench/environments/status.py | 37 +++++++++++++------ mlos_bench/mlos_bench/storage/sql/common.py | 2 +- .../mlos_bench/storage/sql/experiment.py | 6 +-- .../tests/environments/test_status.py | 8 ++-- 4 files changed, 34 insertions(+), 19 deletions(-) diff --git a/mlos_bench/mlos_bench/environments/status.py b/mlos_bench/mlos_bench/environments/status.py index 6d76d7206c8..d49c4a9e0f2 100644 --- a/mlos_bench/mlos_bench/environments/status.py +++ b/mlos_bench/mlos_bench/environments/status.py @@ -24,21 +24,36 @@ class Status(enum.Enum): TIMED_OUT = 7 @staticmethod - def from_str(status_str: Any) -> "Status": - """Convert a string to a Status enum.""" - if not isinstance(status_str, str): - _LOG.warning("Expected type %s for status: %s", type(status_str), status_str) - status_str = str(status_str) - if status_str.isdigit(): + def parse(status: Any) -> "Status": + """Convert the input to a Status enum. + + Parameters + ---------- + status : Any + The status to parse. This can be a string (or string convertible), + int, or Status enum. + + Returns + ------- + Status + The corresponding Status enum value or else UNKNOWN if the input is not + recognized. + """ + if isinstance(status, Status): + return status + if not isinstance(status, str): + _LOG.warning("Expected type %s for status: %s", type(status), status) + status = str(status) + if status.isdigit(): try: - return Status(int(status_str)) + return Status(int(status)) except ValueError: - _LOG.warning("Unknown status: %d", int(status_str)) + _LOG.warning("Unknown status: %d", int(status)) try: - status_str = status_str.upper().strip() - return Status[status_str] + status = status.upper().strip() + return Status[status] except KeyError: - _LOG.warning("Unknown status: %s", status_str) + _LOG.warning("Unknown status: %s", status) return Status.UNKNOWN def is_good(self) -> bool: diff --git a/mlos_bench/mlos_bench/storage/sql/common.py b/mlos_bench/mlos_bench/storage/sql/common.py index 97eb270c9d9..032cf9259d8 100644 --- a/mlos_bench/mlos_bench/storage/sql/common.py +++ b/mlos_bench/mlos_bench/storage/sql/common.py @@ -95,7 +95,7 @@ def get_trials( config_id=trial.config_id, ts_start=utcify_timestamp(trial.ts_start, origin="utc"), ts_end=utcify_nullable_timestamp(trial.ts_end, origin="utc"), - status=Status.from_str(trial.status), + status=Status.parse(trial.status), trial_runner_id=trial.trial_runner_id, ) for trial in trials.fetchall() diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py index acc2a497b48..0e380e3e13a 100644 --- a/mlos_bench/mlos_bench/storage/sql/experiment.py +++ b/mlos_bench/mlos_bench/storage/sql/experiment.py @@ -188,7 +188,7 @@ def load( status: list[Status] = [] for trial in cur_trials.fetchall(): - stat = Status.from_str(trial.status) + stat = Status.parse(trial.status) status.append(stat) trial_ids.append(trial.trial_id) configs.append( @@ -272,7 +272,7 @@ def get_trial_by_id( config_id=trial.config_id, trial_runner_id=trial.trial_runner_id, opt_targets=self._opt_targets, - status=Status.from_str(trial.status), + status=Status.parse(trial.status), restoring=True, config=config, ) @@ -330,7 +330,7 @@ def pending_trials( config_id=trial.config_id, trial_runner_id=trial.trial_runner_id, opt_targets=self._opt_targets, - status=Status.from_str(trial.status), + status=Status.parse(trial.status), restoring=True, config=config, ) diff --git a/mlos_bench/mlos_bench/tests/environments/test_status.py b/mlos_bench/mlos_bench/tests/environments/test_status.py index 3c0a9bccf3c..8123f2b852f 100644 --- a/mlos_bench/mlos_bench/tests/environments/test_status.py +++ b/mlos_bench/mlos_bench/tests/environments/test_status.py @@ -51,16 +51,16 @@ def test_status_from_str_valid(input_str: str, expected_status: Status) -> None: Expected Status enum value. """ assert ( - Status.from_str(input_str) == expected_status + Status.parse(input_str) == expected_status ), f"Expected {expected_status} for input: {input_str}" # Check lowercase representation assert ( - Status.from_str(input_str.lower()) == expected_status + Status.parse(input_str.lower()) == expected_status ), f"Expected {expected_status} for input: {input_str.lower()}" if input_str.isdigit(): # Also test the numeric representation assert ( - Status.from_str(int(input_str)) == expected_status + Status.parse(int(input_str)) == expected_status ), f"Expected {expected_status} for input: {int(input_str)}" @@ -83,7 +83,7 @@ def test_status_from_str_invalid(invalid_input: Any) -> None: input. """ assert ( - Status.from_str(invalid_input) == Status.UNKNOWN + Status.parse(invalid_input) == Status.UNKNOWN ), f"Expected Status.UNKNOWN for invalid input: {invalid_input}" From 7904716320c7931196c2c2de506c90a4c2bce9ce Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 22:36:27 -0500 Subject: [PATCH 28/77] extra test too --- mlos_bench/mlos_bench/tests/environments/test_status.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mlos_bench/mlos_bench/tests/environments/test_status.py b/mlos_bench/mlos_bench/tests/environments/test_status.py index 8123f2b852f..785275825c0 100644 --- a/mlos_bench/mlos_bench/tests/environments/test_status.py +++ b/mlos_bench/mlos_bench/tests/environments/test_status.py @@ -57,6 +57,9 @@ def test_status_from_str_valid(input_str: str, expected_status: Status) -> None: assert ( Status.parse(input_str.lower()) == expected_status ), f"Expected {expected_status} for input: {input_str.lower()}" + assert ( + Status.parse(expected_status) == expected_status + ), f"Expected {expected_status} for input: {expected_status}" if input_str.isdigit(): # Also test the numeric representation assert ( From 54a96e88a6bfbda80269cec2b7c6b2be67950e86 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 22:36:34 -0500 Subject: [PATCH 29/77] more links --- mlos_bench/mlos_bench/launcher.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mlos_bench/mlos_bench/launcher.py b/mlos_bench/mlos_bench/launcher.py index c728ed7fb20..353ace23f0e 100644 --- a/mlos_bench/mlos_bench/launcher.py +++ b/mlos_bench/mlos_bench/launcher.py @@ -55,8 +55,9 @@ def __init__(self, description: str, long_text: str = "", argv: list[str] | None Other required_args values can also be pulled from shell environment variables. - For additional details, please see the website or the README.md files in - the source tree: + For additional details, please see the documentation website or the + README.md files in the source tree: + """ parser = argparse.ArgumentParser(description=f"{description} : {long_text}", epilog=epilog) From 2e1c4db8fb694796e25a8a0a963ffb286b8e7ce5 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 22:36:48 -0500 Subject: [PATCH 30/77] start adding a mock trial runner --- .../tests/schedulers/mock_trial_runner.py | 89 +++++++++++++++++++ 1 file changed, 89 insertions(+) create mode 100644 mlos_bench/mlos_bench/tests/schedulers/mock_trial_runner.py diff --git a/mlos_bench/mlos_bench/tests/schedulers/mock_trial_runner.py b/mlos_bench/mlos_bench/tests/schedulers/mock_trial_runner.py new file mode 100644 index 00000000000..d6cc21d54f3 --- /dev/null +++ b/mlos_bench/mlos_bench/tests/schedulers/mock_trial_runner.py @@ -0,0 +1,89 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +""" +MockTrialRunner for testing :py:class:`mlos_bench.schedulers.Scheduler` logic +with mock trial data. + +This class is intended for use in unit tests and allows for deterministic trial +execution by returning pre-specified results from the ``global_config``. + +Example +------- +Setup mock trial data in the global_config. + +>>> mock_trial_data = { +... 1: { +... "trial_id": 1, +... "status": "SUCCEEDED", +... "metrics": { +... "score": 42.0, +... "color": "red" +... }, +... # Optional sleep time in seconds +... "sleep": 0.1 +... }, +... # Add more trials as needed. +... } + +Normally, this would be part of the global_config passed to the scheduler. +>>> global_config = { +... "mock_trial_data": mock_trial_data, +... # Other global config parameters... +... } + +>>> runner = MockTrialRunner(0, env) +>>> status, timestamp, metrics = runner.run_trial(trial, global_config) +>>> print(status, metrics) +Status.SUCCEEDED {'score': 42.0, 'color': 'red'} +""" +import time +from datetime import datetime +from typing import Any + +from mlos_bench.schedulers.trial_runner import TrialRunner +from mlos_bench.storage.base_storage import Storage +from mlos_bench.environments.status import Status +from mlos_bench.tunables.tunable_types import TunableValue + + +class MockTrialRunner(TrialRunner): + """ + Mock implementation of TrialRunner for testing. + + This class overrides the run_trial method to return mock results + from the global_config["mock_trial_data"] dict, keyed by trial_id. + + """ + + def run_trial( + self, + trial: Storage.Trial, + global_config: dict[str, Any] | None = None, + ) -> tuple[Status, datetime, dict[str, TunableValue] | None]: + """ + Run a mock trial using data from global_config["mock_trial_data"]. + + Parameters + ---------- + trial : Storage.Trial + The trial object, must have a trial_id attribute. + global_config : dict + Global configuration, must contain "mock_trial_data". + + Returns + ------- + (status, timestamp, metrics) : (Status, datetime, dict[str, TunableValue] | None) + Status, timestamp, and metrics for the mock trial. + """ + assert global_config is not None, "global_config must be provided." + mock_data = global_config.get("mock_trial_data", {}) + trial_id = getattr(trial, "trial_id", None) + assert trial_id in mock_data, f"No mock data for trial_id {trial_id}" + data = mock_data[trial_id] + sleep_time = data.get("sleep", 0.01) + time.sleep(sleep_time) + status = Status[data.get("status", "SUCCEEDED")] + metrics = data.get("metrics", {}) + timestamp = datetime.now() + return status, timestamp, metrics From 626554d65a6a5c687fd5a4addb8913e0b13079e9 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 22:36:56 -0500 Subject: [PATCH 31/77] Revert "start adding a mock trial runner" can do that all in mock_env This reverts commit 2e1c4db8fb694796e25a8a0a963ffb286b8e7ce5. --- .../tests/schedulers/mock_trial_runner.py | 89 ------------------- 1 file changed, 89 deletions(-) delete mode 100644 mlos_bench/mlos_bench/tests/schedulers/mock_trial_runner.py diff --git a/mlos_bench/mlos_bench/tests/schedulers/mock_trial_runner.py b/mlos_bench/mlos_bench/tests/schedulers/mock_trial_runner.py deleted file mode 100644 index d6cc21d54f3..00000000000 --- a/mlos_bench/mlos_bench/tests/schedulers/mock_trial_runner.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. -# -""" -MockTrialRunner for testing :py:class:`mlos_bench.schedulers.Scheduler` logic -with mock trial data. - -This class is intended for use in unit tests and allows for deterministic trial -execution by returning pre-specified results from the ``global_config``. - -Example -------- -Setup mock trial data in the global_config. - ->>> mock_trial_data = { -... 1: { -... "trial_id": 1, -... "status": "SUCCEEDED", -... "metrics": { -... "score": 42.0, -... "color": "red" -... }, -... # Optional sleep time in seconds -... "sleep": 0.1 -... }, -... # Add more trials as needed. -... } - -Normally, this would be part of the global_config passed to the scheduler. ->>> global_config = { -... "mock_trial_data": mock_trial_data, -... # Other global config parameters... -... } - ->>> runner = MockTrialRunner(0, env) ->>> status, timestamp, metrics = runner.run_trial(trial, global_config) ->>> print(status, metrics) -Status.SUCCEEDED {'score': 42.0, 'color': 'red'} -""" -import time -from datetime import datetime -from typing import Any - -from mlos_bench.schedulers.trial_runner import TrialRunner -from mlos_bench.storage.base_storage import Storage -from mlos_bench.environments.status import Status -from mlos_bench.tunables.tunable_types import TunableValue - - -class MockTrialRunner(TrialRunner): - """ - Mock implementation of TrialRunner for testing. - - This class overrides the run_trial method to return mock results - from the global_config["mock_trial_data"] dict, keyed by trial_id. - - """ - - def run_trial( - self, - trial: Storage.Trial, - global_config: dict[str, Any] | None = None, - ) -> tuple[Status, datetime, dict[str, TunableValue] | None]: - """ - Run a mock trial using data from global_config["mock_trial_data"]. - - Parameters - ---------- - trial : Storage.Trial - The trial object, must have a trial_id attribute. - global_config : dict - Global configuration, must contain "mock_trial_data". - - Returns - ------- - (status, timestamp, metrics) : (Status, datetime, dict[str, TunableValue] | None) - Status, timestamp, and metrics for the mock trial. - """ - assert global_config is not None, "global_config must be provided." - mock_data = global_config.get("mock_trial_data", {}) - trial_id = getattr(trial, "trial_id", None) - assert trial_id in mock_data, f"No mock data for trial_id {trial_id}" - data = mock_data[trial_id] - sleep_time = data.get("sleep", 0.01) - time.sleep(sleep_time) - status = Status[data.get("status", "SUCCEEDED")] - metrics = data.get("metrics", {}) - timestamp = datetime.now() - return status, timestamp, metrics From b87c02a0e174dbb298c5b28e28a6c9c6f2ee0ebe Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 23:11:41 -0500 Subject: [PATCH 32/77] schema work on mock trial data --- .../environments/mock-env-subschema.json | 129 ++++++++++++++++++ .../mock_env-bad-trial-data-fields.jsonc | 24 ++++ .../invalid/mock_env-bad-trial-data-ids.jsonc | 13 ++ .../mock_env-trial-data-extras.jsonc | 15 ++ .../test-cases/good/full/mock_env-full.jsonc | 36 ++++- 5 files changed, 215 insertions(+), 2 deletions(-) create mode 100644 mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/bad/invalid/mock_env-bad-trial-data-fields.jsonc create mode 100644 mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/bad/invalid/mock_env-bad-trial-data-ids.jsonc create mode 100644 mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/bad/unhandled/mock_env-trial-data-extras.jsonc diff --git a/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json b/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json index cb2de6c719f..fd2c84876d8 100644 --- a/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json +++ b/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json @@ -3,6 +3,124 @@ "$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": { + "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 + } + }, + "type": "object", "properties": { "class": { @@ -42,6 +160,17 @@ }, "minItems": 1, "uniqueItems": true + }, + "mock_trial_data": { + "description": "A list of trial data to use for testing.", + "type": "object", + "patternProperties": { + "^[0-9]+$": { + "$ref": "#/$defs/mock_trial_data" + } + }, + "unevaluatedProperties": false, + "minProperties": 1 } } } diff --git a/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/bad/invalid/mock_env-bad-trial-data-fields.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/bad/invalid/mock_env-bad-trial-data-fields.jsonc new file mode 100644 index 00000000000..d36559cf334 --- /dev/null +++ b/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/bad/invalid/mock_env-bad-trial-data-fields.jsonc @@ -0,0 +1,24 @@ +{ + "class": "mlos_bench.environments.mock_env.MockEnv", + "config": { + "mock_trial_data": { + "1": { + "run": { + // bad types + "status": null, + "metrics": [], + "exception": null, + "sleep": "1", + }, + // missing fields + "setup": {}, + "teardown": { + "status": "UNKNOWN", + "metrics": { + "unexpected": "field" + } + } + } + } + } +} diff --git a/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/bad/invalid/mock_env-bad-trial-data-ids.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/bad/invalid/mock_env-bad-trial-data-ids.jsonc new file mode 100644 index 00000000000..400e557d0fa --- /dev/null +++ b/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/bad/invalid/mock_env-bad-trial-data-ids.jsonc @@ -0,0 +1,13 @@ +{ + "class": "mlos_bench.environments.mock_env.MockEnv", + "config": { + "mock_trial_data": { + // invalid trial id + "trial_id_1": { + "run": { + "status": "FAILED" + } + } + } + } +} diff --git a/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/bad/unhandled/mock_env-trial-data-extras.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/bad/unhandled/mock_env-trial-data-extras.jsonc new file mode 100644 index 00000000000..ecdf4cd0f51 --- /dev/null +++ b/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/bad/unhandled/mock_env-trial-data-extras.jsonc @@ -0,0 +1,15 @@ +{ + "class": "mlos_bench.environments.mock_env.MockEnv", + "config": { + "mock_trial_data": { + "1": { + "new_phase": { + "status": "FAILED" + }, + "run": { + "expected": "property" + } + } + } + } +} diff --git a/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/good/full/mock_env-full.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/good/full/mock_env-full.jsonc index a00f8ca60c0..b618627ecff 100644 --- a/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/good/full/mock_env-full.jsonc +++ b/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/good/full/mock_env-full.jsonc @@ -1,5 +1,5 @@ { - "$schema": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/environments/environment-schema.json", + //"$schema": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/environments/environment-schema.json", "name": "mock_env-full", "description": "More descriptive text.", @@ -25,6 +25,38 @@ "mock_env_metrics": [ "latency", "cost" - ] + ], + "mock_trial_data": { + "1": { + "setup": { + "sleep": 0.1 + }, + "status": { + "metrics": { + "latency": 0.2, + "cost": 0.3 + } + }, + "run": { + "sleep": 0.2, + "status": "SUCCEEDED", + "metrics": { + "latency": 0.1, + "cost": 0.2 + } + }, + "teardown": { + "sleep": 0.1 + } + }, + "2": { + "setup": { + "exception": "Some exception" + }, + "teardown": { + "exception": "Some other exception" + } + } + } } } From b4e5640725480712b38a44af9221f0356eeece1b Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 23:12:59 -0500 Subject: [PATCH 33/77] remove mock scheduler --- .../schedulers/mock-scheduler-subschema.json | 83 ------------------- .../schemas/schedulers/scheduler-schema.json | 3 - .../tests/schedulers/mock_scheduler.py | 66 --------------- 3 files changed, 152 deletions(-) delete mode 100644 mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json delete mode 100644 mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json deleted file mode 100644 index 56d75900946..00000000000 --- a/mlos_bench/mlos_bench/config/schemas/schedulers/mock-scheduler-subschema.json +++ /dev/null @@ -1,83 +0,0 @@ -{ - "$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"] -} diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json index 6ee7ffb1f1f..3086abacd74 100644 --- a/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json +++ b/mlos_bench/mlos_bench/config/schemas/schedulers/scheduler-schema.json @@ -14,9 +14,6 @@ { "$comment": "The set of known Scheduler subschemas. Add others as needed.", "oneOf": [ - { - "$ref": "./mock-scheduler-subschema.json" - }, { "$ref": "./sync-scheduler-subschema.json" } diff --git a/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py deleted file mode 100644 index 43ffff0c34e..00000000000 --- a/mlos_bench/mlos_bench/tests/schedulers/mock_scheduler.py +++ /dev/null @@ -1,66 +0,0 @@ -# -# 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.optimizers.base_optimizer import Optimizer -from mlos_bench.schedulers.base_scheduler import 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")) From 80cf2fa9776d9b40ce3a4951b1bdb1f45365548a Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 23:21:02 -0500 Subject: [PATCH 34/77] allow mock trial data to be in the global config as well --- .../config/schemas/cli/globals-schema.json | 3 +++ .../environments/mock-env-subschema.json | 24 +++++++++++-------- .../good/full/globals-with-schema.jsonc | 17 ++++++++++++- .../tests/schedulers/test_scheduler.py | 23 ++++++++++++++++++ 4 files changed, 56 insertions(+), 11 deletions(-) create mode 100644 mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py diff --git a/mlos_bench/mlos_bench/config/schemas/cli/globals-schema.json b/mlos_bench/mlos_bench/config/schemas/cli/globals-schema.json index 015b4a6e62c..39e60e3249b 100644 --- a/mlos_bench/mlos_bench/config/schemas/cli/globals-schema.json +++ b/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": { diff --git a/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json b/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json index fd2c84876d8..0ef370705d7 100644 --- a/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json +++ b/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json @@ -62,7 +62,8 @@ } } }, - "mock_trial_data": { + "mock_trial_data_item": { + "description": "Mock data for a single trial, split by phase", "type": "object", "properties": { "run": { @@ -118,6 +119,17 @@ }, "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": { + "^[0-9]+$": { + "$ref": "#/$defs/mock_trial_data_item" + } + }, + "unevaluatedProperties": false, + "minProperties": 1 } }, @@ -162,15 +174,7 @@ "uniqueItems": true }, "mock_trial_data": { - "description": "A list of trial data to use for testing.", - "type": "object", - "patternProperties": { - "^[0-9]+$": { - "$ref": "#/$defs/mock_trial_data" - } - }, - "unevaluatedProperties": false, - "minProperties": 1 + "$ref": "#/$defs/mock_trial_data" } } } diff --git a/mlos_bench/mlos_bench/tests/config/schemas/globals/test-cases/good/full/globals-with-schema.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/globals/test-cases/good/full/globals-with-schema.jsonc index 58a0a31bb36..4ed580e09a9 100644 --- a/mlos_bench/mlos_bench/tests/config/schemas/globals/test-cases/good/full/globals-with-schema.jsonc +++ b/mlos_bench/mlos_bench/tests/config/schemas/globals/test-cases/good/full/globals-with-schema.jsonc @@ -10,5 +10,20 @@ "mysql": ["mysql-innodb", "mysql-myisam", "mysql-binlog", "mysql-hugepages"] }, "experiment_id": "ExperimentName", - "trial_id": 1 + "trial_id": 1, + + "mock_trial_data": { + "1": { + "setup": { + "sleep": 1 + }, + "run": { + "status": "SUCCEEDED", + "metrics": { + "score": 0.9, + "color": "green" + } + } + } + } } diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py new file mode 100644 index 00000000000..70322519d3e --- /dev/null +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -0,0 +1,23 @@ +""" +Unit tests for base scheduler internals. +""" + +import pytest + +from mlos_bench.schedulers import Scheduler, SyncScheduler +from mlos_bench.tests.schedulers import MockScheduler + + +# TODO: +# Develop unit tests for schedulers. +# e.g., using MockScheduler it should validate that +# - the base scheduler can be used to run a trial +# - the base scheduler registers the values it receives from the mock_trial_data correctly +# - the base scheduler can be used to run multiple trials +# - the base scheduler does book keeping correctly + +# Actually, maybe what I really want is a MockTrialRunner that can be used to +# return dummy trial results after some predictable period for use in both +# parallel and synchronous schedulers. + +# No, in fact we can do that all with MockEnv and a small extension there. From 84266210d328a0771831e066076e487e697906df Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 23:34:23 -0500 Subject: [PATCH 35/77] comments as prompts in preparation to run a trial --- .../tests/schedulers/test_scheduler.py | 36 ++++++++++--------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 70322519d3e..01c0ede08a4 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -1,23 +1,27 @@ """ Unit tests for base scheduler internals. -""" - -import pytest -from mlos_bench.schedulers import Scheduler, SyncScheduler -from mlos_bench.tests.schedulers import MockScheduler +Notes +----- +Some prompts to help write the tests: +Develop unit tests for `Scheduler` class. +- use `@pytest.mark.parametrize` to run the same test with different Scheduler classes (e.g. `SyncScheduler`, `ParallelScheduler`, etc.) +- use `MockEnv` with `mock_trial_data` as a `pytest.fixture` to run the tests + - needs a jsonc file or string that the `TrialRunner.create_from_json` method can use to create the Env multiple times -# TODO: -# Develop unit tests for schedulers. -# e.g., using MockScheduler it should validate that -# - the base scheduler can be used to run a trial -# - the base scheduler registers the values it receives from the mock_trial_data correctly -# - the base scheduler can be used to run multiple trials -# - the base scheduler does book keeping correctly +Check that: +1. the targeted scheduler can be used to run a trial + - check that results are stored in the storage backend correctly + - use the `sqlite_storage` fixture from `mlos_bench.tests.storage.sql.fixtures` for that + - check that the `_ran_trials` attribute is updated correctly after a run_scheduler call +2. the base scheduler `bulk_registers` the values it receives from the mock_trial_data correctly + - use `mock` to patch the `bulk_register` method in the `Scheduler` class's `optimizer` attribute and check the call arguments +3. the base scheduler does book keeping correctly + - use `mock` to patch the `add_new_optimizer_suggestions` method in the `Scheduler` class and check the `_last_trial_id` +""" -# Actually, maybe what I really want is a MockTrialRunner that can be used to -# return dummy trial results after some predictable period for use in both -# parallel and synchronous schedulers. +import pytest +import unittest.mock -# No, in fact we can do that all with MockEnv and a small extension there. +from mlos_bench.schedulers import Scheduler, SyncScheduler From cdc614fc7a222a5592eefbccac301cc930d52a3e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 May 2025 04:34:53 +0000 Subject: [PATCH 36/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mlos_bench/mlos_bench/environments/base_environment.py | 5 ++--- mlos_bench/mlos_bench/environments/mock_env.py | 5 +++-- mlos_bench/mlos_bench/environments/status.py | 3 ++- mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py | 7 ++++++- 4 files changed, 13 insertions(+), 7 deletions(-) diff --git a/mlos_bench/mlos_bench/environments/base_environment.py b/mlos_bench/mlos_bench/environments/base_environment.py index 98db5cfe14d..335ff453b93 100644 --- a/mlos_bench/mlos_bench/environments/base_environment.py +++ b/mlos_bench/mlos_bench/environments/base_environment.py @@ -394,9 +394,8 @@ def current_trial_id(self) -> int: @property def trial_runner_id(self) -> int: """ - Get the ID of the - :py:class:`~.mlos_bench.schedulers.trial_runner.TrialRunner` for this - Environment. + 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. diff --git a/mlos_bench/mlos_bench/environments/mock_env.py b/mlos_bench/mlos_bench/environments/mock_env.py index e15dcffed44..88199e48cbe 100644 --- a/mlos_bench/mlos_bench/environments/mock_env.py +++ b/mlos_bench/mlos_bench/environments/mock_env.py @@ -6,10 +6,10 @@ import logging import random +import time from dataclasses import dataclass from datetime import datetime from typing import Any -import time import numpy @@ -227,7 +227,8 @@ def _produce_metrics(self, rand: random.Random | None) -> dict[str, TunableValue return {metric: float(score) for metric in self._metrics or []} def get_current_mock_trial_data(self) -> MockTrialData: - """Gets mock trial data for the current trial ID. + """ + Gets mock trial data for the current trial ID. If no (or missing) mock trial data is found, a new instance of MockTrialData is created from random data. diff --git a/mlos_bench/mlos_bench/environments/status.py b/mlos_bench/mlos_bench/environments/status.py index d49c4a9e0f2..8bec0a22c44 100644 --- a/mlos_bench/mlos_bench/environments/status.py +++ b/mlos_bench/mlos_bench/environments/status.py @@ -25,7 +25,8 @@ class Status(enum.Enum): @staticmethod def parse(status: Any) -> "Status": - """Convert the input to a Status enum. + """ + Convert the input to a Status enum. Parameters ---------- diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 01c0ede08a4..a13f24a7ba1 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -1,3 +1,7 @@ +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# """ Unit tests for base scheduler internals. @@ -21,7 +25,8 @@ - use `mock` to patch the `add_new_optimizer_suggestions` method in the `Scheduler` class and check the `_last_trial_id` """ -import pytest import unittest.mock +import pytest + from mlos_bench.schedulers import Scheduler, SyncScheduler From b3f49e9804d6f13182cbab83737d050d4eed7d3a Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 23:43:45 -0500 Subject: [PATCH 37/77] fixups --- .../mlos_bench/environments/mock_env.py | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/mlos_bench/mlos_bench/environments/mock_env.py b/mlos_bench/mlos_bench/environments/mock_env.py index 88199e48cbe..9d3b16cb0d8 100644 --- a/mlos_bench/mlos_bench/environments/mock_env.py +++ b/mlos_bench/mlos_bench/environments/mock_env.py @@ -231,7 +231,7 @@ def get_current_mock_trial_data(self) -> MockTrialData: Gets mock trial data for the current trial ID. If no (or missing) mock trial data is found, a new instance of - MockTrialData is created from random data. + MockTrialData is created and later filled with random data. Note ---- @@ -262,7 +262,18 @@ def setup(self, tunables: TunableGroups, global_config: dict | None = None) -> b raise RuntimeError( f"Mock trial data setup exception: {mock_trial_data.setup.exception}" ) - return is_success and mock_trial_data.setup.status.is_ready() + return is_success + + def teardown(self) -> None: + mock_trial_data = self.get_current_mock_trial_data() + if mock_trial_data.teardown.sleep: + _LOG.debug("Sleeping for %s seconds", mock_trial_data.teardown.sleep) + time.sleep(mock_trial_data.teardown.sleep) + if mock_trial_data.teardown.exception: + raise RuntimeError( + f"Mock trial data teardown exception: {mock_trial_data.teardown.exception}" + ) + super().teardown() def run(self) -> tuple[Status, datetime, dict[str, TunableValue] | None]: """ @@ -286,7 +297,6 @@ def run(self) -> tuple[Status, datetime, dict[str, TunableValue] | None]: time.sleep(mock_trial_data.run.sleep) if mock_trial_data.run.exception: raise RuntimeError(f"Mock trial data run exception: {mock_trial_data.run.exception}") - if mock_trial_data.run.metrics is None: # If no metrics are provided, generate them. mock_trial_data.run.metrics = self._produce_metrics(self._run_random) @@ -316,9 +326,13 @@ def status(self) -> tuple[Status, datetime, list[tuple[datetime, str, Any]]]: ) if mock_trial_data.status.metrics is None: # If no metrics are provided, generate them. + # Note: we don't save these in the mock trial data as they may need + # to change to preserve backwards compatibility with previous tests. metrics = self._produce_metrics(self._status_random) else: # If metrics are provided, use them. + # Note: current implementation uses the same metrics for all status + # calls of this trial. metrics = mock_trial_data.status.metrics return ( mock_trial_data.status.status, From a41e85f479436cd1e7f12a70c634c651a1c92358 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 23:45:09 -0500 Subject: [PATCH 38/77] revert --- .../environments/test-cases/good/full/mock_env-full.jsonc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/good/full/mock_env-full.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/good/full/mock_env-full.jsonc index b618627ecff..a23971f0362 100644 --- a/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/good/full/mock_env-full.jsonc +++ b/mlos_bench/mlos_bench/tests/config/schemas/environments/test-cases/good/full/mock_env-full.jsonc @@ -1,5 +1,5 @@ { - //"$schema": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/environments/environment-schema.json", + "$schema": "https://raw.githubusercontent.com/microsoft/MLOS/main/mlos_bench/mlos_bench/config/schemas/environments/environment-schema.json", "name": "mock_env-full", "description": "More descriptive text.", From 63e0f88e7c9378df80ee36a351799bf086095000 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 23:45:56 -0500 Subject: [PATCH 39/77] remove --- mlos_bench/mlos_bench/tests/schedulers/__init__.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/__init__.py b/mlos_bench/mlos_bench/tests/schedulers/__init__.py index b1668588752..4bc0076079f 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/__init__.py +++ b/mlos_bench/mlos_bench/tests/schedulers/__init__.py @@ -3,8 +3,3 @@ # Licensed under the MIT License. # """mlos_bench.tests.schedulers.""" -from mlos_bench.tests.schedulers.mock_scheduler import MockScheduler - -__all__ = [ - "MockScheduler", -] From 44bdde2916535208150df25d6fc2092294bcef0e Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 22:34:47 -0500 Subject: [PATCH 40/77] refactor status parsing a little bit again --- mlos_bench/mlos_bench/environments/status.py | 37 +++++++++++++------ mlos_bench/mlos_bench/storage/sql/common.py | 2 +- .../mlos_bench/storage/sql/experiment.py | 6 +-- .../tests/environments/test_status.py | 8 ++-- 4 files changed, 34 insertions(+), 19 deletions(-) diff --git a/mlos_bench/mlos_bench/environments/status.py b/mlos_bench/mlos_bench/environments/status.py index 6d76d7206c8..d49c4a9e0f2 100644 --- a/mlos_bench/mlos_bench/environments/status.py +++ b/mlos_bench/mlos_bench/environments/status.py @@ -24,21 +24,36 @@ class Status(enum.Enum): TIMED_OUT = 7 @staticmethod - def from_str(status_str: Any) -> "Status": - """Convert a string to a Status enum.""" - if not isinstance(status_str, str): - _LOG.warning("Expected type %s for status: %s", type(status_str), status_str) - status_str = str(status_str) - if status_str.isdigit(): + def parse(status: Any) -> "Status": + """Convert the input to a Status enum. + + Parameters + ---------- + status : Any + The status to parse. This can be a string (or string convertible), + int, or Status enum. + + Returns + ------- + Status + The corresponding Status enum value or else UNKNOWN if the input is not + recognized. + """ + if isinstance(status, Status): + return status + if not isinstance(status, str): + _LOG.warning("Expected type %s for status: %s", type(status), status) + status = str(status) + if status.isdigit(): try: - return Status(int(status_str)) + return Status(int(status)) except ValueError: - _LOG.warning("Unknown status: %d", int(status_str)) + _LOG.warning("Unknown status: %d", int(status)) try: - status_str = status_str.upper().strip() - return Status[status_str] + status = status.upper().strip() + return Status[status] except KeyError: - _LOG.warning("Unknown status: %s", status_str) + _LOG.warning("Unknown status: %s", status) return Status.UNKNOWN def is_good(self) -> bool: diff --git a/mlos_bench/mlos_bench/storage/sql/common.py b/mlos_bench/mlos_bench/storage/sql/common.py index 97eb270c9d9..032cf9259d8 100644 --- a/mlos_bench/mlos_bench/storage/sql/common.py +++ b/mlos_bench/mlos_bench/storage/sql/common.py @@ -95,7 +95,7 @@ def get_trials( config_id=trial.config_id, ts_start=utcify_timestamp(trial.ts_start, origin="utc"), ts_end=utcify_nullable_timestamp(trial.ts_end, origin="utc"), - status=Status.from_str(trial.status), + status=Status.parse(trial.status), trial_runner_id=trial.trial_runner_id, ) for trial in trials.fetchall() diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py index acc2a497b48..0e380e3e13a 100644 --- a/mlos_bench/mlos_bench/storage/sql/experiment.py +++ b/mlos_bench/mlos_bench/storage/sql/experiment.py @@ -188,7 +188,7 @@ def load( status: list[Status] = [] for trial in cur_trials.fetchall(): - stat = Status.from_str(trial.status) + stat = Status.parse(trial.status) status.append(stat) trial_ids.append(trial.trial_id) configs.append( @@ -272,7 +272,7 @@ def get_trial_by_id( config_id=trial.config_id, trial_runner_id=trial.trial_runner_id, opt_targets=self._opt_targets, - status=Status.from_str(trial.status), + status=Status.parse(trial.status), restoring=True, config=config, ) @@ -330,7 +330,7 @@ def pending_trials( config_id=trial.config_id, trial_runner_id=trial.trial_runner_id, opt_targets=self._opt_targets, - status=Status.from_str(trial.status), + status=Status.parse(trial.status), restoring=True, config=config, ) diff --git a/mlos_bench/mlos_bench/tests/environments/test_status.py b/mlos_bench/mlos_bench/tests/environments/test_status.py index 3c0a9bccf3c..8123f2b852f 100644 --- a/mlos_bench/mlos_bench/tests/environments/test_status.py +++ b/mlos_bench/mlos_bench/tests/environments/test_status.py @@ -51,16 +51,16 @@ def test_status_from_str_valid(input_str: str, expected_status: Status) -> None: Expected Status enum value. """ assert ( - Status.from_str(input_str) == expected_status + Status.parse(input_str) == expected_status ), f"Expected {expected_status} for input: {input_str}" # Check lowercase representation assert ( - Status.from_str(input_str.lower()) == expected_status + Status.parse(input_str.lower()) == expected_status ), f"Expected {expected_status} for input: {input_str.lower()}" if input_str.isdigit(): # Also test the numeric representation assert ( - Status.from_str(int(input_str)) == expected_status + Status.parse(int(input_str)) == expected_status ), f"Expected {expected_status} for input: {int(input_str)}" @@ -83,7 +83,7 @@ def test_status_from_str_invalid(invalid_input: Any) -> None: input. """ assert ( - Status.from_str(invalid_input) == Status.UNKNOWN + Status.parse(invalid_input) == Status.UNKNOWN ), f"Expected Status.UNKNOWN for invalid input: {invalid_input}" From 607fffdc53ef9e930b6a50814645ea79b508d250 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 22:36:27 -0500 Subject: [PATCH 41/77] extra test too --- mlos_bench/mlos_bench/tests/environments/test_status.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mlos_bench/mlos_bench/tests/environments/test_status.py b/mlos_bench/mlos_bench/tests/environments/test_status.py index 8123f2b852f..785275825c0 100644 --- a/mlos_bench/mlos_bench/tests/environments/test_status.py +++ b/mlos_bench/mlos_bench/tests/environments/test_status.py @@ -57,6 +57,9 @@ def test_status_from_str_valid(input_str: str, expected_status: Status) -> None: assert ( Status.parse(input_str.lower()) == expected_status ), f"Expected {expected_status} for input: {input_str.lower()}" + assert ( + Status.parse(expected_status) == expected_status + ), f"Expected {expected_status} for input: {expected_status}" if input_str.isdigit(): # Also test the numeric representation assert ( From aaf0842e6f5cc9b3edbce958145becbb82196749 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 May 2025 04:50:54 +0000 Subject: [PATCH 42/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mlos_bench/mlos_bench/environments/status.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/environments/status.py b/mlos_bench/mlos_bench/environments/status.py index d49c4a9e0f2..8bec0a22c44 100644 --- a/mlos_bench/mlos_bench/environments/status.py +++ b/mlos_bench/mlos_bench/environments/status.py @@ -25,7 +25,8 @@ class Status(enum.Enum): @staticmethod def parse(status: Any) -> "Status": - """Convert the input to a Status enum. + """ + Convert the input to a Status enum. Parameters ---------- From 8f472bbf9b31cbf3adb87c3702c9c34ac08ab127 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Wed, 21 May 2025 23:52:58 -0500 Subject: [PATCH 43/77] revert --- .../config/schemas/schedulers/base-scheduler-subschema.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json b/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json index 9417b7d00ac..702da1eec3e 100644 --- a/mlos_bench/mlos_bench/config/schemas/schedulers/base-scheduler-subschema.json +++ b/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([.]tests)?[.]schedulers[.]" + "pattern": "^mlos_bench[.]schedulers[.]" }, "config": { From 8a4aac29423e7996771ddb926fe47f14fbbaa89f Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 13:19:45 -0500 Subject: [PATCH 44/77] comments --- mlos_bench/mlos_bench/environments/status.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/environments/status.py b/mlos_bench/mlos_bench/environments/status.py index 8bec0a22c44..6343d3e854b 100644 --- a/mlos_bench/mlos_bench/environments/status.py +++ b/mlos_bench/mlos_bench/environments/status.py @@ -129,4 +129,15 @@ def is_timed_out(self) -> bool: Status.TIMED_OUT, } ) -"""The set of completed statuses.""" +""" +The set of completed statuses. + +Includes all statuses that indicate the trial or experiment has finished, either +successfully or not. +This set is used to determine if a trial or experiment has reached a final state. +This includes: +- :py:data:`.Status.SUCCEEDED`: The trial or experiment completed successfully. +- :py:data:`.Status.CANCELED`: The trial or experiment was canceled. +- :py:data:`.Status.FAILED`: The trial or experiment failed. +- :py:data:`.Status.TIMED_OUT`: The trial or experiment timed out. +""" From 50ffde4bbcd3387b988168fb3bbf87799c98ad12 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 14:07:17 -0500 Subject: [PATCH 45/77] add accessor for the mock_trial_data property --- mlos_bench/mlos_bench/environments/mock_env.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/mlos_bench/mlos_bench/environments/mock_env.py b/mlos_bench/mlos_bench/environments/mock_env.py index 9d3b16cb0d8..6af016b7e89 100644 --- a/mlos_bench/mlos_bench/environments/mock_env.py +++ b/mlos_bench/mlos_bench/environments/mock_env.py @@ -7,6 +7,7 @@ import logging import random import time +from copy import deepcopy from dataclasses import dataclass from datetime import datetime from typing import Any @@ -226,6 +227,17 @@ def _produce_metrics(self, rand: random.Random | None) -> dict[str, TunableValue return {metric: float(score) for metric in self._metrics or []} + def mock_trial_data(self) -> dict[int, MockTrialData]: + """ + Get the mock trial data for all trials. + + Returns + ------- + dict[int, MockTrialData] + Dictionary of mock trial data keyed by trial ID. + """ + return deepcopy(self._mock_trial_data) + def get_current_mock_trial_data(self) -> MockTrialData: """ Gets mock trial data for the current trial ID. From b1e2a52f19dd14e49d2aa88b1fbfa7289f5d496b Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 14:07:43 -0500 Subject: [PATCH 46/77] make it a property --- mlos_bench/mlos_bench/environments/mock_env.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mlos_bench/mlos_bench/environments/mock_env.py b/mlos_bench/mlos_bench/environments/mock_env.py index 6af016b7e89..c565b1adc66 100644 --- a/mlos_bench/mlos_bench/environments/mock_env.py +++ b/mlos_bench/mlos_bench/environments/mock_env.py @@ -227,6 +227,7 @@ def _produce_metrics(self, rand: random.Random | None) -> dict[str, TunableValue return {metric: float(score) for metric in self._metrics or []} + @property def mock_trial_data(self) -> dict[int, MockTrialData]: """ Get the mock trial data for all trials. From f0f7c4cf6bdf12e484e9b2e710b1d2ecbb568e85 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 14:16:29 -0500 Subject: [PATCH 47/77] add some basic fixtures to get started with --- .../mlos_bench/tests/schedulers/conftest.py | 95 ++++++++++++++++++- 1 file changed, 94 insertions(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/schedulers/conftest.py index 83a18783cc6..75b42fa9a15 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/conftest.py +++ b/mlos_bench/mlos_bench/tests/schedulers/conftest.py @@ -2,4 +2,97 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # -"""Pytest fixtures for mlos_bench.schedulers tests.""" +""" +Pytest fixtures for mlos_bench.schedulers tests. +""" +# pylint: disable=redefined-outer-name + +import json + +import pytest + +from mlos_bench.environments.mock_env import MockEnv +from mlos_bench.services.config_persistence import ConfigPersistenceService +from mlos_bench.schedulers.trial_runner import TrialRunner +from mlos_bench.tunables.tunable_groups import TunableGroups +import mlos_bench.tests.optimizers.fixtures as optimizers_fixtures + + +NUM_TRIAL_RUNNERS = 4 + + +@pytest.fixture +def mock_env_config() -> dict: + """A config for a MockEnv with mock_trial_data.""" + return { + "name": "Test MockEnv With Explicit Mock Trial Data", + "class": "mlos_bench.environments.mock_env.MockEnv", + "config": { + "mock_env_seed": -1, + "mock_env_range": [0, 10], + "mock_env_metrics": ["score"], + # TODO: Add more mock trial data here: + "mock_trial_data": { + "0": { + "setup": { + "status": "SUCCEEDED", + }, + "run": { + "status": "SUCCEEDED", + "metrics": { + "score": 1.0, + }, + }, + }, + "1": { + "setup": { + "status": "SUCCEEDED", + }, + "run": { + "status": "SUCCEEDED", + "metrics": { + "score": 2.0, + }, + }, + }, + }, + }, + } + + +@pytest.fixture +def mock_env_json_config(mock_env_config: dict) -> str: + """A JSON string of the mock_env_config.""" + return json.dumps(mock_env_config) + + +@pytest.fixture +def mock_env( + mock_env_json_config: str, + tunable_groups: TunableGroups, +) -> MockEnv: + """A fixture to create a MockEnv instance using the mock_env_json_config.""" + config_loader_service = ConfigPersistenceService() + mock_env = config_loader_service.load_environment( + mock_env_json_config, + tunable_groups, + service=config_loader_service, + ) + assert isinstance(mock_env, MockEnv) + return mock_env + + +@pytest.fixture +def trial_runners( + mock_env_json_config: str, + tunable_groups: TunableGroups, +) -> list[TrialRunner]: + """A fixture to create a list of TrialRunner instances using the + mock_env_json_config.""" + config_loader_service = ConfigPersistenceService() + return TrialRunner.create_from_json( + config_loader=config_loader_service, + env_json=mock_env_json_config, + tunable_groups=tunable_groups, + num_trial_runners=NUM_TRIAL_RUNNERS, + ) From c27b3c269dd800cb268aa2024ac243bb7c4c612c Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 14:16:50 -0500 Subject: [PATCH 48/77] add a method for creating schedulers, intended to be used with pytest parameterize --- .../tests/schedulers/test_scheduler.py | 48 +++++++++++-------- 1 file changed, 27 insertions(+), 21 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index a13f24a7ba1..40c99d019cb 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -3,30 +3,36 @@ # Licensed under the MIT License. # """ -Unit tests for base scheduler internals. +Unit tests for :py:class:`mlos_bench.schedulers` and their internals. +""" -Notes ------ -Some prompts to help write the tests: +import pytest +import unittest.mock -Develop unit tests for `Scheduler` class. -- use `@pytest.mark.parametrize` to run the same test with different Scheduler classes (e.g. `SyncScheduler`, `ParallelScheduler`, etc.) -- use `MockEnv` with `mock_trial_data` as a `pytest.fixture` to run the tests - - needs a jsonc file or string that the `TrialRunner.create_from_json` method can use to create the Env multiple times +from mlos_bench.storage.sql.storage import SqlStorage +from mlos_bench.schedulers.base_scheduler import Scheduler +from mlos_bench.schedulers.sync_scheduler import SyncScheduler +from mlos_bench.optimizers.mock_optimizer import MockOptimizer +from mlos_bench.schedulers.trial_runner import TrialRunner +import mlos_bench.tests.optimizers.fixtures as optimizers_fixtures -Check that: -1. the targeted scheduler can be used to run a trial - - check that results are stored in the storage backend correctly - - use the `sqlite_storage` fixture from `mlos_bench.tests.storage.sql.fixtures` for that - - check that the `_ran_trials` attribute is updated correctly after a run_scheduler call -2. the base scheduler `bulk_registers` the values it receives from the mock_trial_data correctly - - use `mock` to patch the `bulk_register` method in the `Scheduler` class's `optimizer` attribute and check the call arguments -3. the base scheduler does book keeping correctly - - use `mock` to patch the `add_new_optimizer_suggestions` method in the `Scheduler` class and check the `_last_trial_id` -""" +mock_opt = optimizers_fixtures.mock_opt -import unittest.mock +# pylint: disable=redefined-outer-name -import pytest -from mlos_bench.schedulers import Scheduler, SyncScheduler +def create_scheduler( + scheduler_type: type[Scheduler], + trial_runners: list[TrialRunner], + mock_opt: MockOptimizer, + sqlite_storage: SqlStorage, +) -> Scheduler: + """Create a Scheduler instance using trial_runners, mock_opt, and sqlite_storage.""" + return scheduler_type( + config={}, + global_config={}, + trial_runners=trial_runners, + optimizer=mock_opt, + storage=sqlite_storage, + root_env_config="", + ) From 981326890418e4ba170fdf647e7b5c69decd91af Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 14:33:48 -0500 Subject: [PATCH 49/77] stubbing out a very basic test to get started --- .../tests/schedulers/test_scheduler.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 40c99d019cb..559ab9fddb6 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -8,7 +8,9 @@ import pytest import unittest.mock +import sys +from mlos_bench.environments.mock_env import MockEnv from mlos_bench.storage.sql.storage import SqlStorage from mlos_bench.schedulers.base_scheduler import Scheduler from mlos_bench.schedulers.sync_scheduler import SyncScheduler @@ -36,3 +38,37 @@ def create_scheduler( storage=sqlite_storage, root_env_config="", ) + + +@pytest.mark.parametrize( + "scheduler_class", + [ + SyncScheduler, + ], +) +@pytest.mark.skipif( + sys.platform == "win32", + reason="Skipping test on Windows - SQLite storage is not accessible in parallel tests there.", +) +def test_scheduler( + scheduler_class: type[Scheduler], + # fixtures: + trial_runners: list[TrialRunner], + mock_opt: MockOptimizer, + sqlite_storage: SqlStorage, +) -> None: + """ + Test the creation of a SyncScheduler instance. + """ + scheduler = create_scheduler( + scheduler_class, + trial_runners, + mock_opt, + sqlite_storage, + ) + assert isinstance(scheduler, scheduler_class) + assert isinstance(scheduler.trial_runners, list) + assert len(scheduler.trial_runners) == len(trial_runners) + assert isinstance(scheduler.optimizer, MockOptimizer) + assert isinstance(scheduler.storage, SqlStorage) + assert isinstance(scheduler.root_environment, MockEnv) From 3b9201767adfd8d6eb8dec0218a09fc3d887b066 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 14:55:16 -0500 Subject: [PATCH 50/77] wip: testing --- .../mlos_bench/schedulers/base_scheduler.py | 11 +- .../tests/schedulers/test_scheduler.py | 114 +++++++++++++++--- 2 files changed, 106 insertions(+), 19 deletions(-) diff --git a/mlos_bench/mlos_bench/schedulers/base_scheduler.py b/mlos_bench/mlos_bench/schedulers/base_scheduler.py index 3e47841591f..3c22d427c79 100644 --- a/mlos_bench/mlos_bench/schedulers/base_scheduler.py +++ b/mlos_bench/mlos_bench/schedulers/base_scheduler.py @@ -243,7 +243,14 @@ def __exit__( return False # Do not suppress exceptions def _prepare_start(self) -> bool: - """Prepare the scheduler for starting.""" + """ + Prepare the scheduler for starting. + + Notes + ----- + This method is called by the :py:meth:`Scheduler.start` method. + It is split out mostly to allow for easier unit testing/mocking. + """ assert self.experiment is not None _LOG.info( "START: Experiment: %s Env: %s Optimizer: %s", @@ -280,7 +287,7 @@ def _execute_scheduling_step(self, is_warm_up: bool) -> bool: Notes ----- This method is called by the :py:meth:`Scheduler.start` method. - It is split out mostly to allow for easier testing with MockSchedulers. + It is split out mostly to allow for easier unit testing/mocking. """ assert self.experiment is not None _LOG.info("Optimization loop: Last trial ID: %d", self._last_trial_id) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 559ab9fddb6..e7d07c12928 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -6,14 +6,15 @@ Unit tests for :py:class:`mlos_bench.schedulers` and their internals. """ -import pytest -import unittest.mock +from unittest.mock import patch import sys +import pytest + +from mlos_core.tests import get_all_concrete_subclasses from mlos_bench.environments.mock_env import MockEnv from mlos_bench.storage.sql.storage import SqlStorage from mlos_bench.schedulers.base_scheduler import Scheduler -from mlos_bench.schedulers.sync_scheduler import SyncScheduler from mlos_bench.optimizers.mock_optimizer import MockOptimizer from mlos_bench.schedulers.trial_runner import TrialRunner import mlos_bench.tests.optimizers.fixtures as optimizers_fixtures @@ -29,10 +30,24 @@ def create_scheduler( mock_opt: MockOptimizer, sqlite_storage: SqlStorage, ) -> Scheduler: - """Create a Scheduler instance using trial_runners, mock_opt, and sqlite_storage.""" + """Create a Scheduler instance using trial_runners, mock_opt, and + sqlite_storage fixtures.""" + + env = trial_runners[0].environment + assert isinstance(env, MockEnv), "Environment is not a MockEnv instance." + max_trials = max(trial_id for trial_id in env.mock_trial_data.keys()) + max_trials = min(max_trials, mock_opt.max_suggestions) + return scheduler_type( - config={}, - global_config={}, + config={ + "max_trials": max_trials, + }, + global_config={ + "experiment_id": "Test{scheduler_type.__name__}Experiment", + "trial_id": 1, + # TODO: Adjust this in the future? + "trial_repeat_count": 1, + }, trial_runners=trial_runners, optimizer=mock_opt, storage=sqlite_storage, @@ -40,11 +55,16 @@ def create_scheduler( ) +scheduler_classes = get_all_concrete_subclasses( + Scheduler, # type: ignore[type-abstract] + pkg_name="mlos_bench", +) +assert scheduler_classes, "No Scheduler classes found in mlos_bench." + + @pytest.mark.parametrize( "scheduler_class", - [ - SyncScheduler, - ], + scheduler_classes, ) @pytest.mark.skipif( sys.platform == "win32", @@ -52,23 +72,83 @@ def create_scheduler( ) def test_scheduler( scheduler_class: type[Scheduler], - # fixtures: trial_runners: list[TrialRunner], mock_opt: MockOptimizer, sqlite_storage: SqlStorage, ) -> None: """ - Test the creation of a SyncScheduler instance. + Full integration test for Scheduler: runs trials, checks storage, optimizer + registration, and internal bookkeeping. """ + # pylint: disable=too-many-locals + + # Create the scheduler. scheduler = create_scheduler( scheduler_class, trial_runners, mock_opt, sqlite_storage, ) - assert isinstance(scheduler, scheduler_class) - assert isinstance(scheduler.trial_runners, list) - assert len(scheduler.trial_runners) == len(trial_runners) - assert isinstance(scheduler.optimizer, MockOptimizer) - assert isinstance(scheduler.storage, SqlStorage) - assert isinstance(scheduler.root_environment, MockEnv) + + root_env = scheduler.root_environment + assert isinstance(root_env, MockEnv), "Root environment is not a MockEnv instance." + mock_trial_data = root_env.mock_trial_data + + # Patch bulk_register and add_new_optimizer_suggestions + with ( + patch.object( + scheduler.optimizer, + "bulk_register", + wraps=scheduler.optimizer.bulk_register, + ) as mock_bulk_register, + patch.object( + scheduler, + "add_new_optimizer_suggestions", + wraps=scheduler.add_new_optimizer_suggestions, + ) as mock_add_suggestions, + ): + # Run the scheduler + with scheduler: + scheduler.start() + scheduler.teardown() + + # Now check the results. + # TODO + + # 1. Check results in storage + experiments = scheduler.storage.experiments + assert experiments, "No experiments found in storage." + # Get the first experiment + experiment = next(iter(experiments.values())) + trials = experiment.trials + # Compare with mock_trial_data from root_environment + for trial_id, trial_data in trials.items(): + # Check that the trial result matches the mock data + expected = mock_trial_data[trial_id].run.metrics + actual = trial_data.results_dict + assert actual == expected, f"Trial {trial_id} results {actual} != expected {expected}" + + # 1b. Check ran_trials bookkeeping + ran_trials = scheduler.ran_trials + assert len(ran_trials) == len(trials) + for trial in ran_trials: + assert ( + trial.status.is_ready() + ), f"Trial {trial.trial_id} did not complete successfully." + + # 2. Check optimizer registration + assert mock_bulk_register.called, "bulk_register was not called on optimizer." + # Check that the configs and scores match the mock_trial_data + for call in mock_bulk_register.call_args_list: + configs, scores, *_ = call.args + for i, config in enumerate(configs): + trial_id = i # assumes order matches + expected_score = mock_trial_data[trial_id].run.metrics + assert ( + scores[i] == expected_score + ), f"bulk_register score {scores[i]} != expected {expected_score} for trial {trial_id}" + + # 3. Check bookkeeping: add_new_optimizer_suggestions and _last_trial_id + assert mock_add_suggestions.called, "add_new_optimizer_suggestions was not called." + # _last_trial_id should be the last trial id + assert getattr(scheduler, "_last_trial_id", None) == max(trials.keys()) From 341564ebd82d896f1ffcb5a2c97ce4d5e21f0783 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 16:03:23 -0500 Subject: [PATCH 51/77] wip --- .../environments/base_environment.py | 3 +- .../mlos_bench/tests/schedulers/conftest.py | 36 ++++++-- .../tests/schedulers/test_scheduler.py | 86 +++++-------------- 3 files changed, 53 insertions(+), 72 deletions(-) diff --git a/mlos_bench/mlos_bench/environments/base_environment.py b/mlos_bench/mlos_bench/environments/base_environment.py index 335ff453b93..e98d410f7a1 100644 --- a/mlos_bench/mlos_bench/environments/base_environment.py +++ b/mlos_bench/mlos_bench/environments/base_environment.py @@ -385,9 +385,10 @@ def current_trial_id(self) -> int: """ val = self._params["trial_id"] assert isinstance(val, int), ( - "Expected trial_id to be an int, but got %s (type %s)", + "Expected trial_id to be an int, but got %s (type %s): %s", val, type(val), + self._params, ) return val diff --git a/mlos_bench/mlos_bench/tests/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/schedulers/conftest.py index 75b42fa9a15..272edbc9ee3 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/conftest.py +++ b/mlos_bench/mlos_bench/tests/schedulers/conftest.py @@ -15,7 +15,6 @@ from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.schedulers.trial_runner import TrialRunner from mlos_bench.tunables.tunable_groups import TunableGroups -import mlos_bench.tests.optimizers.fixtures as optimizers_fixtures NUM_TRIAL_RUNNERS = 4 @@ -28,16 +27,19 @@ def mock_env_config() -> dict: "name": "Test MockEnv With Explicit Mock Trial Data", "class": "mlos_bench.environments.mock_env.MockEnv", "config": { + # Reference the covariant groups from the `tunable_groups` fixture. + # See Also: + # - mlos_bench/tests/conftest.py + # - mlos_bench/tests/tunable_groups_fixtures.py + "tunable_params": ["provision", "boot", "kernel"], "mock_env_seed": -1, "mock_env_range": [0, 10], "mock_env_metrics": ["score"], # TODO: Add more mock trial data here: "mock_trial_data": { "0": { - "setup": { - "status": "SUCCEEDED", - }, "run": { + "sleep": 0.15, "status": "SUCCEEDED", "metrics": { "score": 1.0, @@ -45,13 +47,20 @@ def mock_env_config() -> dict: }, }, "1": { - "setup": { + "run": { + "sleep": 0.2, "status": "SUCCEEDED", + "metrics": { + "score": 2.0, + }, }, + }, + "2": { "run": { + "sleep": 0.1, "status": "SUCCEEDED", "metrics": { - "score": 2.0, + "score": 3.0, }, }, }, @@ -60,6 +69,15 @@ def mock_env_config() -> dict: } +@pytest.fixture +def global_config() -> dict: + """A global config for a MockEnv.""" + return { + "experiment_id": "TestExperiment", + "trial_id": 1, + } + + @pytest.fixture def mock_env_json_config(mock_env_config: dict) -> str: """A JSON string of the mock_env_config.""" @@ -70,6 +88,7 @@ def mock_env_json_config(mock_env_config: dict) -> str: def mock_env( mock_env_json_config: str, tunable_groups: TunableGroups, + global_config: dict, ) -> MockEnv: """A fixture to create a MockEnv instance using the mock_env_json_config.""" config_loader_service = ConfigPersistenceService() @@ -77,6 +96,7 @@ def mock_env( mock_env_json_config, tunable_groups, service=config_loader_service, + global_config=global_config, ) assert isinstance(mock_env, MockEnv) return mock_env @@ -86,13 +106,15 @@ def mock_env( def trial_runners( mock_env_json_config: str, tunable_groups: TunableGroups, + global_config: dict, ) -> list[TrialRunner]: """A fixture to create a list of TrialRunner instances using the mock_env_json_config.""" - config_loader_service = ConfigPersistenceService() + config_loader_service = ConfigPersistenceService(global_config=global_config) return TrialRunner.create_from_json( config_loader=config_loader_service, env_json=mock_env_json_config, tunable_groups=tunable_groups, num_trial_runners=NUM_TRIAL_RUNNERS, + global_config=global_config, ) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index e7d07c12928..302447715ac 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -18,8 +18,10 @@ from mlos_bench.optimizers.mock_optimizer import MockOptimizer from mlos_bench.schedulers.trial_runner import TrialRunner import mlos_bench.tests.optimizers.fixtures as optimizers_fixtures +import mlos_bench.tests.storage.sql.fixtures as sql_storage_fixtures mock_opt = optimizers_fixtures.mock_opt +sqlite_storage = sql_storage_fixtures.sqlite_storage # pylint: disable=redefined-outer-name @@ -29,6 +31,7 @@ def create_scheduler( trial_runners: list[TrialRunner], mock_opt: MockOptimizer, sqlite_storage: SqlStorage, + global_config: dict, ) -> Scheduler: """Create a Scheduler instance using trial_runners, mock_opt, and sqlite_storage fixtures.""" @@ -38,16 +41,13 @@ def create_scheduler( max_trials = max(trial_id for trial_id in env.mock_trial_data.keys()) max_trials = min(max_trials, mock_opt.max_suggestions) + global_config["experiment_id"] = f"Test{scheduler_type.__name__}Experiment" + return scheduler_type( config={ "max_trials": max_trials, }, - global_config={ - "experiment_id": "Test{scheduler_type.__name__}Experiment", - "trial_id": 1, - # TODO: Adjust this in the future? - "trial_repeat_count": 1, - }, + global_config=global_config, trial_runners=trial_runners, optimizer=mock_opt, storage=sqlite_storage, @@ -75,6 +75,7 @@ def test_scheduler( trial_runners: list[TrialRunner], mock_opt: MockOptimizer, sqlite_storage: SqlStorage, + global_config: dict, ) -> None: """ Full integration test for Scheduler: runs trials, checks storage, optimizer @@ -88,67 +89,24 @@ def test_scheduler( trial_runners, mock_opt, sqlite_storage, + global_config, ) root_env = scheduler.root_environment assert isinstance(root_env, MockEnv), "Root environment is not a MockEnv instance." mock_trial_data = root_env.mock_trial_data - # Patch bulk_register and add_new_optimizer_suggestions - with ( - patch.object( - scheduler.optimizer, - "bulk_register", - wraps=scheduler.optimizer.bulk_register, - ) as mock_bulk_register, - patch.object( - scheduler, - "add_new_optimizer_suggestions", - wraps=scheduler.add_new_optimizer_suggestions, - ) as mock_add_suggestions, - ): - # Run the scheduler - with scheduler: - scheduler.start() - scheduler.teardown() - - # Now check the results. - # TODO - - # 1. Check results in storage - experiments = scheduler.storage.experiments - assert experiments, "No experiments found in storage." - # Get the first experiment - experiment = next(iter(experiments.values())) - trials = experiment.trials - # Compare with mock_trial_data from root_environment - for trial_id, trial_data in trials.items(): - # Check that the trial result matches the mock data - expected = mock_trial_data[trial_id].run.metrics - actual = trial_data.results_dict - assert actual == expected, f"Trial {trial_id} results {actual} != expected {expected}" - - # 1b. Check ran_trials bookkeeping - ran_trials = scheduler.ran_trials - assert len(ran_trials) == len(trials) - for trial in ran_trials: - assert ( - trial.status.is_ready() - ), f"Trial {trial.trial_id} did not complete successfully." - - # 2. Check optimizer registration - assert mock_bulk_register.called, "bulk_register was not called on optimizer." - # Check that the configs and scores match the mock_trial_data - for call in mock_bulk_register.call_args_list: - configs, scores, *_ = call.args - for i, config in enumerate(configs): - trial_id = i # assumes order matches - expected_score = mock_trial_data[trial_id].run.metrics - assert ( - scores[i] == expected_score - ), f"bulk_register score {scores[i]} != expected {expected_score} for trial {trial_id}" - - # 3. Check bookkeeping: add_new_optimizer_suggestions and _last_trial_id - assert mock_add_suggestions.called, "add_new_optimizer_suggestions was not called." - # _last_trial_id should be the last trial id - assert getattr(scheduler, "_last_trial_id", None) == max(trials.keys()) + # Run the scheduler + with scheduler: + scheduler.start() + scheduler.teardown() + + # Now check the results. + # TODO: + # Check the overall results: + # 1. Check the results in storage. + # 2. Check the optimizer registration. + # 3. Check the bookkeeping for ran_trials. + # 4. Check the bookkeeping for add_new_optimizer_suggestions and _last_trial_id. + # This last part may require patching and intercepting during the start() + # loop to validate in-progress book keeping instead of just overall. From 5c068149510765555524c68f9f5f4fcb4b078dda Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 May 2025 21:03:47 +0000 Subject: [PATCH 52/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../mlos_bench/tests/schedulers/conftest.py | 10 ++++----- .../tests/schedulers/test_scheduler.py | 21 +++++++++---------- 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/schedulers/conftest.py index 272edbc9ee3..dffd8e8bad2 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/conftest.py +++ b/mlos_bench/mlos_bench/tests/schedulers/conftest.py @@ -2,9 +2,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # -""" -Pytest fixtures for mlos_bench.schedulers tests. -""" +"""Pytest fixtures for mlos_bench.schedulers tests.""" # pylint: disable=redefined-outer-name import json @@ -12,11 +10,10 @@ import pytest from mlos_bench.environments.mock_env import MockEnv -from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.schedulers.trial_runner import TrialRunner +from mlos_bench.services.config_persistence import ConfigPersistenceService from mlos_bench.tunables.tunable_groups import TunableGroups - NUM_TRIAL_RUNNERS = 4 @@ -109,7 +106,8 @@ def trial_runners( global_config: dict, ) -> list[TrialRunner]: """A fixture to create a list of TrialRunner instances using the - mock_env_json_config.""" + mock_env_json_config. + """ config_loader_service = ConfigPersistenceService(global_config=global_config) return TrialRunner.create_from_json( config_loader=config_loader_service, diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 302447715ac..12828bb17d6 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -2,23 +2,21 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. # -""" -Unit tests for :py:class:`mlos_bench.schedulers` and their internals. -""" +"""Unit tests for :py:class:`mlos_bench.schedulers` and their internals.""" -from unittest.mock import patch import sys +from unittest.mock import patch import pytest -from mlos_core.tests import get_all_concrete_subclasses +import mlos_bench.tests.optimizers.fixtures as optimizers_fixtures +import mlos_bench.tests.storage.sql.fixtures as sql_storage_fixtures from mlos_bench.environments.mock_env import MockEnv -from mlos_bench.storage.sql.storage import SqlStorage -from mlos_bench.schedulers.base_scheduler import Scheduler from mlos_bench.optimizers.mock_optimizer import MockOptimizer +from mlos_bench.schedulers.base_scheduler import Scheduler from mlos_bench.schedulers.trial_runner import TrialRunner -import mlos_bench.tests.optimizers.fixtures as optimizers_fixtures -import mlos_bench.tests.storage.sql.fixtures as sql_storage_fixtures +from mlos_bench.storage.sql.storage import SqlStorage +from mlos_core.tests import get_all_concrete_subclasses mock_opt = optimizers_fixtures.mock_opt sqlite_storage = sql_storage_fixtures.sqlite_storage @@ -33,8 +31,9 @@ def create_scheduler( sqlite_storage: SqlStorage, global_config: dict, ) -> Scheduler: - """Create a Scheduler instance using trial_runners, mock_opt, and - sqlite_storage fixtures.""" + """Create a Scheduler instance using trial_runners, mock_opt, and sqlite_storage + fixtures. + """ env = trial_runners[0].environment assert isinstance(env, MockEnv), "Environment is not a MockEnv instance." From 836dcd9e9b6d38f3ecd6e704ec737e0227d63664 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 16:04:57 -0500 Subject: [PATCH 53/77] comments --- mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 12828bb17d6..102f28f066a 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -100,7 +100,7 @@ def test_scheduler( scheduler.start() scheduler.teardown() - # Now check the results. + # Now check the overall results. # TODO: # Check the overall results: # 1. Check the results in storage. From 00721415466c31ff772657972092266cef288e2c Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 16:20:47 -0500 Subject: [PATCH 54/77] fixups --- .../environments/mock-env-subschema.json | 2 +- .../environments/base_environment.py | 7 +++--- .../mlos_bench/tests/schedulers/conftest.py | 14 +++++++----- .../tests/schedulers/test_scheduler.py | 22 +++++++++++++++++-- .../mlos_bench/tests/storage/sql/fixtures.py | 2 +- 5 files changed, 35 insertions(+), 12 deletions(-) diff --git a/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json b/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json index 0ef370705d7..b453c8573aa 100644 --- a/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json +++ b/mlos_bench/mlos_bench/config/schemas/environments/mock-env-subschema.json @@ -124,7 +124,7 @@ "description": "A set of mock trial data to use for testing, keyed by trial id. Used by MockEnv.", "type": "object", "patternProperties": { - "^[0-9]+$": { + "^[1-9][0-9]*$": { "$ref": "#/$defs/mock_trial_data_item" } }, diff --git a/mlos_bench/mlos_bench/environments/base_environment.py b/mlos_bench/mlos_bench/environments/base_environment.py index e98d410f7a1..fe40025f95d 100644 --- a/mlos_bench/mlos_bench/environments/base_environment.py +++ b/mlos_bench/mlos_bench/environments/base_environment.py @@ -419,7 +419,8 @@ def trial_runner_id(self) -> int: ) return val - def experiment_id(self) -> int: + @property + def experiment_id(self) -> str: """ Get the ID of the experiment. @@ -428,7 +429,7 @@ def experiment_id(self) -> int: Returns ------- - experiment_id : int + experiment_id : str The ID of the experiment. Notes @@ -441,7 +442,7 @@ def experiment_id(self) -> int: mlos_bench.config : documentation on the configuration system """ val = self._params["experiment_id"] - assert isinstance(val, int), ( + assert isinstance(val, str), ( "Expected experiment_id to be an int, but got %s (type %s)", val, type(val), diff --git a/mlos_bench/mlos_bench/tests/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/schedulers/conftest.py index dffd8e8bad2..e5aa302d618 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/conftest.py +++ b/mlos_bench/mlos_bench/tests/schedulers/conftest.py @@ -6,6 +6,7 @@ # pylint: disable=redefined-outer-name import json +import re import pytest @@ -34,7 +35,7 @@ def mock_env_config() -> dict: "mock_env_metrics": ["score"], # TODO: Add more mock trial data here: "mock_trial_data": { - "0": { + "1": { "run": { "sleep": 0.15, "status": "SUCCEEDED", @@ -43,7 +44,7 @@ def mock_env_config() -> dict: }, }, }, - "1": { + "2": { "run": { "sleep": 0.2, "status": "SUCCEEDED", @@ -52,7 +53,7 @@ def mock_env_config() -> dict: }, }, }, - "2": { + "3": { "run": { "sleep": 0.1, "status": "SUCCEEDED", @@ -67,10 +68,13 @@ def mock_env_config() -> dict: @pytest.fixture -def global_config() -> dict: +def global_config(request) -> dict: """A global config for a MockEnv.""" + test_name = request.node.name + test_name = re.sub(r"[^a-zA-Z0-9]", "_", test_name) + experiment_id = f"TestExperiment-{test_name}" return { - "experiment_id": "TestExperiment", + "experiment_id": experiment_id, "trial_id": 1, } diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 102f28f066a..7bfa2c6aac0 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -40,8 +40,6 @@ def create_scheduler( max_trials = max(trial_id for trial_id in env.mock_trial_data.keys()) max_trials = min(max_trials, mock_opt.max_suggestions) - global_config["experiment_id"] = f"Test{scheduler_type.__name__}Experiment" - return scheduler_type( config={ "max_trials": max_trials, @@ -92,6 +90,7 @@ def test_scheduler( ) root_env = scheduler.root_environment + experiment_id = root_env.experiment_id assert isinstance(root_env, MockEnv), "Root environment is not a MockEnv instance." mock_trial_data = root_env.mock_trial_data @@ -101,6 +100,25 @@ def test_scheduler( scheduler.teardown() # Now check the overall results. + + # Check the results in storage. + exp_data = sqlite_storage.experiments[experiment_id] + for mock_trial_data in mock_trial_data.values(): + trial_id = mock_trial_data.trial_id + assert trial_id in exp_data.trials, f"Trial {trial_id} not found in storage." + trial_data = exp_data.trials[trial_id] + + # Check the results. + metrics = mock_trial_data.run.metrics + if metrics: + for result_key, result_value in metrics.items(): + assert ( + result_key in trial_data.results_dict + ), f"Result column {result_key} not found in storage." + assert ( + trial_data.results_dict[result_key] == result_value + ), f"Result value for {result_key} does not match expected value." + # TODO: # Check the overall results: # 1. Check the results in storage. diff --git a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py index 0bebeeff824..db6dc5fa2e3 100644 --- a/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py +++ b/mlos_bench/mlos_bench/tests/storage/sql/fixtures.py @@ -30,7 +30,7 @@ # pylint: disable=redefined-outer-name -@pytest.fixture +@pytest.fixture(scope="function") def sqlite_storage() -> Generator[SqlStorage]: """ Fixture for file based SQLite storage in a temporary directory. From 5400bd6805ae7a8723b270898ceeb5beb8990f5e Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 16:26:08 -0500 Subject: [PATCH 55/77] more checks --- .../tests/schedulers/test_scheduler.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 7bfa2c6aac0..c520aca7779 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -114,16 +114,32 @@ def test_scheduler( for result_key, result_value in metrics.items(): assert ( result_key in trial_data.results_dict - ), f"Result column {result_key} not found in storage." + ), f"Result {result_key} not found in storage for trial {trial_data}." assert ( trial_data.results_dict[result_key] == result_value ), f"Result value for {result_key} does not match expected value." + else: + # metrics weren't explicit in the mock trial data, so we only check + # that a score was registered + for opt_target in mock_opt.targets: + assert ( + opt_target in trial_data.results_dict + ), f"Result column {opt_target} not found in storage." + assert ( + trial_data.results_dict[opt_target] is not None + ), f"Result value for {opt_target} is None." + + assert ( + trial_data.status == mock_trial_data.run.status + ), f"Trial {trial_id} status {trial_data.status} was not {mock_trial_data.run.status}." + + # TODO: Check the trial status telemetry. # TODO: - # Check the overall results: - # 1. Check the results in storage. # 2. Check the optimizer registration. # 3. Check the bookkeeping for ran_trials. + + # TODO: And check the intermediary results. # 4. Check the bookkeeping for add_new_optimizer_suggestions and _last_trial_id. # This last part may require patching and intercepting during the start() # loop to validate in-progress book keeping instead of just overall. From e62a3787656b524121b7e713f0117d31fd2b70ba Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 17:21:43 -0500 Subject: [PATCH 56/77] Add more checks --- .../mlos_bench/optimizers/base_optimizer.py | 4 +- .../mlos_bench/optimizers/mock_optimizer.py | 42 +++++++ .../tests/schedulers/test_scheduler.py | 108 +++++++++++++++--- 3 files changed, 134 insertions(+), 20 deletions(-) diff --git a/mlos_bench/mlos_bench/optimizers/base_optimizer.py b/mlos_bench/mlos_bench/optimizers/base_optimizer.py index 44aa9a035e2..72b437b320e 100644 --- a/mlos_bench/mlos_bench/optimizers/base_optimizer.py +++ b/mlos_bench/mlos_bench/optimizers/base_optimizer.py @@ -356,8 +356,10 @@ def _get_scores( assert scores is not None target_metrics: dict[str, float] = {} for opt_target, opt_dir in self._opt_targets.items(): + if opt_target not in scores: + raise ValueError(f"Score for {opt_target} not found in {scores}.") val = scores[opt_target] - assert val is not None + assert val is not None, f"Score for {opt_target} is None." target_metrics[opt_target] = float(val) * opt_dir return target_metrics diff --git a/mlos_bench/mlos_bench/optimizers/mock_optimizer.py b/mlos_bench/mlos_bench/optimizers/mock_optimizer.py index 947e34a7da4..60eed31b460 100644 --- a/mlos_bench/mlos_bench/optimizers/mock_optimizer.py +++ b/mlos_bench/mlos_bench/optimizers/mock_optimizer.py @@ -15,6 +15,7 @@ import random from collections.abc import Callable, Sequence +from dataclasses import dataclass from mlos_bench.environments.status import Status from mlos_bench.optimizers.track_best_optimizer import TrackBestOptimizer from mlos_bench.services.base_service import Service @@ -25,6 +26,15 @@ _LOG = logging.getLogger(__name__) +@dataclass +class RegisteredScore: + """A registered score for a trial.""" + + config: TunableGroups + score: dict[str, TunableValue] | None + status: Status + + class MockOptimizer(TrackBestOptimizer): """Mock optimizer to test the Environment API.""" @@ -42,6 +52,38 @@ def __init__( "float": lambda tunable: rnd.uniform(*tunable.range), "int": lambda tunable: rnd.randint(*(int(x) for x in tunable.range)), } + self._registered_scores: list[RegisteredScore] = [] + + @property + def registered_scores(self) -> list[RegisteredScore]: + """Return the list of registered scores. + + Notes + ----- + Used for testing and validation. + """ + return self._registered_scores + + def register( + self, + tunables: TunableGroups, + status: Status, + score: dict[str, TunableValue] | None = None, + ) -> dict[str, float] | None: + # Track the registered scores for testing and validation. + score = score or {} + # Almost the same as _get_scores, but we don't adjust the direction here. + scores: dict[str, TunableValue] = { + k: float(v) for k, v in score.items() if k in self._opt_targets and v is not None + } + self._registered_scores.append( + RegisteredScore( + config=tunables.copy(), + score=scores, + status=status, + ) + ) + return super().register(tunables, status, score) def bulk_register( self, diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index c520aca7779..3af7e136a4b 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -4,6 +4,7 @@ # """Unit tests for :py:class:`mlos_bench.schedulers` and their internals.""" +from logging import warning import sys from unittest.mock import patch @@ -15,12 +16,15 @@ from mlos_bench.optimizers.mock_optimizer import MockOptimizer from mlos_bench.schedulers.base_scheduler import Scheduler from mlos_bench.schedulers.trial_runner import TrialRunner +from mlos_bench.storage.base_trial_data import TrialData from mlos_bench.storage.sql.storage import SqlStorage from mlos_core.tests import get_all_concrete_subclasses mock_opt = optimizers_fixtures.mock_opt sqlite_storage = sql_storage_fixtures.sqlite_storage +DEBUG_WARNINGS_ENABLED = False + # pylint: disable=redefined-outer-name @@ -52,6 +56,47 @@ def create_scheduler( ) +def debug_warn(*args: object) -> None: + """Optionally issue warnings for debugging.""" + if DEBUG_WARNINGS_ENABLED: + warning(*args) + + +def mock_opt_has_registered_trial_score( + mock_opt: MockOptimizer, + trial_data: TrialData, +) -> bool: + """Check that the MockOptimizer has registered a given MockTrialData.""" + # pylint: disable=consider-using-any-or-all + # Split out for easier debugging. + for registered_score in mock_opt.registered_scores: + match = True + if registered_score.status != trial_data.status: + match = False + debug_warn( + "Registered status %s does not match trial data %s.", + registered_score.status, + trial_data.results_dict, + ) + elif registered_score.score != trial_data.results_dict: + debug_warn( + "Registered score %s does not match trial data %s.", + registered_score.score, + trial_data.results_dict, + ) + match = False + elif registered_score.config.get_param_values() != trial_data.tunable_config.config_dict: + debug_warn( + "Registered config %s does not match trial data %s.", + registered_score.config.get_param_values(), + trial_data.results_dict, + ) + match = False + if match: + return True + return False + + scheduler_classes = get_all_concrete_subclasses( Scheduler, # type: ignore[type-abstract] pkg_name="mlos_bench", @@ -67,7 +112,7 @@ def create_scheduler( sys.platform == "win32", reason="Skipping test on Windows - SQLite storage is not accessible in parallel tests there.", ) -def test_scheduler( +def test_scheduler_with_mock_trial_data( scheduler_class: type[Scheduler], trial_runners: list[TrialRunner], mock_opt: MockOptimizer, @@ -91,8 +136,8 @@ def test_scheduler( root_env = scheduler.root_environment experiment_id = root_env.experiment_id - assert isinstance(root_env, MockEnv), "Root environment is not a MockEnv instance." - mock_trial_data = root_env.mock_trial_data + assert isinstance(root_env, MockEnv), f"Root environment {root_env} is not a MockEnv." + assert root_env.mock_trial_data, "No mock trial data found in root environment." # Run the scheduler with scheduler: @@ -100,11 +145,19 @@ def test_scheduler( scheduler.teardown() # Now check the overall results. - - # Check the results in storage. + ran_trials = {trial.trial_id for trial in scheduler.ran_trials} + assert ( + experiment_id in sqlite_storage.experiments + ), f"Experiment {experiment_id} not found in storage." exp_data = sqlite_storage.experiments[experiment_id] - for mock_trial_data in mock_trial_data.values(): + + for mock_trial_data in root_env.mock_trial_data.values(): trial_id = mock_trial_data.trial_id + + # Check the bookkeeping for ran_trials. + assert trial_id in ran_trials, f"Trial {trial_id} not found in Scheduler.ran_trials." + + # Check the results in storage. assert trial_id in exp_data.trials, f"Trial {trial_id} not found in storage." trial_data = exp_data.trials[trial_id] @@ -118,26 +171,43 @@ def test_scheduler( assert ( trial_data.results_dict[result_key] == result_value ), f"Result value for {result_key} does not match expected value." - else: - # metrics weren't explicit in the mock trial data, so we only check - # that a score was registered - for opt_target in mock_opt.targets: - assert ( - opt_target in trial_data.results_dict - ), f"Result column {opt_target} not found in storage." - assert ( - trial_data.results_dict[opt_target] is not None - ), f"Result value for {opt_target} is None." + # TODO: Should we check the reverse - no extra metrics were registered? + # else: metrics weren't explicit in the mock trial data, so we only + # check that a score was stored for the optimization target, but that's + # good to do regardless + for opt_target in mock_opt.targets: + assert ( + opt_target in trial_data.results_dict + ), f"Result column {opt_target} not found in storage." + assert ( + trial_data.results_dict[opt_target] is not None + ), f"Result value for {opt_target} is None." + + # Check that the appropriate sleeps occurred. + trial_time_lb = 0.0 + trial_time_lb += mock_trial_data.setup.sleep or 0 + trial_time_lb += mock_trial_data.run.sleep or 0 + trial_time_lb += mock_trial_data.status.sleep or 0 + trial_time_lb += mock_trial_data.teardown.sleep or 0 + assert trial_data.ts_end is not None, f"Trial {trial_id} has no end time." + trial_duration = trial_data.ts_end - trial_data.ts_start + trial_dur_secs = trial_duration.total_seconds() + assert ( + trial_dur_secs >= trial_time_lb + ), f"Trial {trial_id} took less time ({trial_dur_secs}) than expected ({trial_time_lb}). " + # Check that the trial status matches what we expected. assert ( trial_data.status == mock_trial_data.run.status ), f"Trial {trial_id} status {trial_data.status} was not {mock_trial_data.run.status}." # TODO: Check the trial status telemetry. - # TODO: - # 2. Check the optimizer registration. - # 3. Check the bookkeeping for ran_trials. + # Check the optimizer registration. + assert mock_opt_has_registered_trial_score( + mock_opt, + trial_data, + ), f"Trial {trial_id} was not registered in the optimizer." # TODO: And check the intermediary results. # 4. Check the bookkeeping for add_new_optimizer_suggestions and _last_trial_id. From 206bb779b5b693110de4398db0d3adaa4240672b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 May 2025 22:23:33 +0000 Subject: [PATCH 57/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mlos_bench/mlos_bench/optimizers/mock_optimizer.py | 5 +++-- mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/mlos_bench/mlos_bench/optimizers/mock_optimizer.py b/mlos_bench/mlos_bench/optimizers/mock_optimizer.py index 60eed31b460..a1311b6f953 100644 --- a/mlos_bench/mlos_bench/optimizers/mock_optimizer.py +++ b/mlos_bench/mlos_bench/optimizers/mock_optimizer.py @@ -14,8 +14,8 @@ import logging import random from collections.abc import Callable, Sequence - from dataclasses import dataclass + from mlos_bench.environments.status import Status from mlos_bench.optimizers.track_best_optimizer import TrackBestOptimizer from mlos_bench.services.base_service import Service @@ -56,7 +56,8 @@ def __init__( @property def registered_scores(self) -> list[RegisteredScore]: - """Return the list of registered scores. + """ + Return the list of registered scores. Notes ----- diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 3af7e136a4b..04a998ea220 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -4,8 +4,8 @@ # """Unit tests for :py:class:`mlos_bench.schedulers` and their internals.""" -from logging import warning import sys +from logging import warning from unittest.mock import patch import pytest From abe412be2b06287a990b6489eff000b881ce8177 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 17:24:54 -0500 Subject: [PATCH 58/77] linting --- mlos_bench/mlos_bench/tests/schedulers/conftest.py | 3 ++- mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py | 2 -- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/schedulers/conftest.py index e5aa302d618..df6bd2776fd 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/conftest.py +++ b/mlos_bench/mlos_bench/tests/schedulers/conftest.py @@ -9,6 +9,7 @@ import re import pytest +from pytest import FixtureRequest from mlos_bench.environments.mock_env import MockEnv from mlos_bench.schedulers.trial_runner import TrialRunner @@ -68,7 +69,7 @@ def mock_env_config() -> dict: @pytest.fixture -def global_config(request) -> dict: +def global_config(request: FixtureRequest) -> dict: """A global config for a MockEnv.""" test_name = request.node.name test_name = re.sub(r"[^a-zA-Z0-9]", "_", test_name) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 04a998ea220..9a2cc1dbaec 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -6,7 +6,6 @@ import sys from logging import warning -from unittest.mock import patch import pytest @@ -38,7 +37,6 @@ def create_scheduler( """Create a Scheduler instance using trial_runners, mock_opt, and sqlite_storage fixtures. """ - env = trial_runners[0].environment assert isinstance(env, MockEnv), "Environment is not a MockEnv instance." max_trials = max(trial_id for trial_id in env.mock_trial_data.keys()) From f5cb4689bde3a435dd8ec05b6111efe528a3735a Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 17:26:54 -0500 Subject: [PATCH 59/77] doc tweaks --- mlos_bench/mlos_bench/environments/status.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mlos_bench/mlos_bench/environments/status.py b/mlos_bench/mlos_bench/environments/status.py index 6343d3e854b..aa3b3e99c16 100644 --- a/mlos_bench/mlos_bench/environments/status.py +++ b/mlos_bench/mlos_bench/environments/status.py @@ -136,8 +136,8 @@ def is_timed_out(self) -> bool: successfully or not. This set is used to determine if a trial or experiment has reached a final state. This includes: -- :py:data:`.Status.SUCCEEDED`: The trial or experiment completed successfully. -- :py:data:`.Status.CANCELED`: The trial or experiment was canceled. -- :py:data:`.Status.FAILED`: The trial or experiment failed. -- :py:data:`.Status.TIMED_OUT`: The trial or experiment timed out. +- :py:attr:`.Status.SUCCEEDED`: The trial or experiment completed successfully. +- :py:attr:`.Status.CANCELED`: The trial or experiment was canceled. +- :py:attr:`.Status.FAILED`: The trial or experiment failed. +- :py:attr:`.Status.TIMED_OUT`: The trial or experiment timed out. """ From dc7a7b2a42f2e17b03b1473019e23efc59a46e99 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 17:34:08 -0500 Subject: [PATCH 60/77] fixup --- mlos_bench/mlos_bench/environments/mock_env.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/environments/mock_env.py b/mlos_bench/mlos_bench/environments/mock_env.py index c565b1adc66..2cabbe90553 100644 --- a/mlos_bench/mlos_bench/environments/mock_env.py +++ b/mlos_bench/mlos_bench/environments/mock_env.py @@ -248,7 +248,7 @@ def get_current_mock_trial_data(self) -> MockTrialData: Note ---- - This method must be called after the base :py:meth:`Environment.setup` + This method must be called after the base :py:meth:`.Environment.setup` method is called to ensure the current ``trial_id`` is set. """ trial_id = self.current_trial_id From 76e94eda13cd71d85672ffa8bb8b0f70f26f4135 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 17:40:33 -0500 Subject: [PATCH 61/77] fixup --- .../tests/schedulers/test_scheduler.py | 43 +++---------------- 1 file changed, 6 insertions(+), 37 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 9a2cc1dbaec..f5bda73af68 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -5,7 +5,6 @@ """Unit tests for :py:class:`mlos_bench.schedulers` and their internals.""" import sys -from logging import warning import pytest @@ -22,8 +21,6 @@ mock_opt = optimizers_fixtures.mock_opt sqlite_storage = sql_storage_fixtures.sqlite_storage -DEBUG_WARNINGS_ENABLED = False - # pylint: disable=redefined-outer-name @@ -54,45 +51,17 @@ def create_scheduler( ) -def debug_warn(*args: object) -> None: - """Optionally issue warnings for debugging.""" - if DEBUG_WARNINGS_ENABLED: - warning(*args) - - def mock_opt_has_registered_trial_score( mock_opt: MockOptimizer, trial_data: TrialData, ) -> bool: """Check that the MockOptimizer has registered a given MockTrialData.""" - # pylint: disable=consider-using-any-or-all - # Split out for easier debugging. - for registered_score in mock_opt.registered_scores: - match = True - if registered_score.status != trial_data.status: - match = False - debug_warn( - "Registered status %s does not match trial data %s.", - registered_score.status, - trial_data.results_dict, - ) - elif registered_score.score != trial_data.results_dict: - debug_warn( - "Registered score %s does not match trial data %s.", - registered_score.score, - trial_data.results_dict, - ) - match = False - elif registered_score.config.get_param_values() != trial_data.tunable_config.config_dict: - debug_warn( - "Registered config %s does not match trial data %s.", - registered_score.config.get_param_values(), - trial_data.results_dict, - ) - match = False - if match: - return True - return False + return any( + registered_score.status == trial_data.status + and registered_score.score == trial_data.results_dict + and registered_score.config.get_param_values() == trial_data.tunable_config.config_dict + for registered_score in mock_opt.registered_scores + ) scheduler_classes = get_all_concrete_subclasses( From 7f8522ac88a56a71d8a93af46a79027a8d5e9e0c Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 18:11:19 -0500 Subject: [PATCH 62/77] fixups --- .../mlos_bench/environments/mock_env.py | 8 +++-- .../mlos_bench/optimizers/mock_optimizer.py | 5 +-- mlos_bench/mlos_bench/tests/__init__.py | 6 ++-- mlos_bench/mlos_bench/tests/conftest.py | 26 +++++++++++++-- .../tests/environments/composite_env_test.py | 4 ++- .../environments/include_tunables_test.py | 8 ++++- .../optimizers/opt_bulk_register_test.py | 7 +++- .../mlos_bench/tests/schedulers/conftest.py | 32 ++++++------------- .../tests/schedulers/test_scheduler.py | 4 +-- 9 files changed, 63 insertions(+), 37 deletions(-) diff --git a/mlos_bench/mlos_bench/environments/mock_env.py b/mlos_bench/mlos_bench/environments/mock_env.py index 2cabbe90553..b46c053e6b8 100644 --- a/mlos_bench/mlos_bench/environments/mock_env.py +++ b/mlos_bench/mlos_bench/environments/mock_env.py @@ -310,10 +310,12 @@ def run(self) -> tuple[Status, datetime, dict[str, TunableValue] | None]: time.sleep(mock_trial_data.run.sleep) if mock_trial_data.run.exception: raise RuntimeError(f"Mock trial data run exception: {mock_trial_data.run.exception}") - if mock_trial_data.run.metrics is None: + if mock_trial_data.run.metrics is not None: + metrics = mock_trial_data.run.metrics + else: # If no metrics are provided, generate them. - mock_trial_data.run.metrics = self._produce_metrics(self._run_random) - return (mock_trial_data.run.status, timestamp, mock_trial_data.run.metrics) + metrics = self._produce_metrics(self._run_random) + return (mock_trial_data.run.status, timestamp, metrics) def status(self) -> tuple[Status, datetime, list[tuple[datetime, str, Any]]]: """ diff --git a/mlos_bench/mlos_bench/optimizers/mock_optimizer.py b/mlos_bench/mlos_bench/optimizers/mock_optimizer.py index a1311b6f953..0878ca4d75d 100644 --- a/mlos_bench/mlos_bench/optimizers/mock_optimizer.py +++ b/mlos_bench/mlos_bench/optimizers/mock_optimizer.py @@ -72,10 +72,11 @@ def register( score: dict[str, TunableValue] | None = None, ) -> dict[str, float] | None: # Track the registered scores for testing and validation. - score = score or {} # Almost the same as _get_scores, but we don't adjust the direction here. scores: dict[str, TunableValue] = { - k: float(v) for k, v in score.items() if k in self._opt_targets and v is not None + k: float(v) + for k, v in (score or {}).items() + if k in self._opt_targets and v is not None } self._registered_scores.append( RegisteredScore( diff --git a/mlos_bench/mlos_bench/tests/__init__.py b/mlos_bench/mlos_bench/tests/__init__.py index ce5fbdb45af..25df725a9ed 100644 --- a/mlos_bench/mlos_bench/tests/__init__.py +++ b/mlos_bench/mlos_bench/tests/__init__.py @@ -31,9 +31,9 @@ ZONE_INFO: list[tzinfo | None] = [nullable(pytz.timezone, zone_name) for zone_name in ZONE_NAMES] BUILT_IN_ENV_VAR_DEFAULTS = { - "experiment_id": None, - "trial_id": None, - "trial_runner_id": None, + "experiment_id": "SomeExperimentName", + "trial_id": 1, + "trial_runner_id": 0, } # A decorator for tests that require docker. diff --git a/mlos_bench/mlos_bench/tests/conftest.py b/mlos_bench/mlos_bench/tests/conftest.py index becae205033..2f8690784c5 100644 --- a/mlos_bench/mlos_bench/tests/conftest.py +++ b/mlos_bench/mlos_bench/tests/conftest.py @@ -8,8 +8,10 @@ import sys from collections.abc import Generator from typing import Any +import re import pytest +from pytest import FixtureRequest from fasteners import InterProcessLock, InterProcessReaderWriterLock from pytest_docker.plugin import Services as DockerServices from pytest_docker.plugin import get_docker_services @@ -30,7 +32,22 @@ @pytest.fixture -def mock_env(tunable_groups: TunableGroups) -> MockEnv: +def mock_env_global_config(request: FixtureRequest) -> dict: + """A global config for a MockEnv.""" + test_name = request.node.name + test_name = re.sub(r"[^a-zA-Z0-9]", "_", test_name) + experiment_id = f"TestExperiment-{test_name}" + return { + "experiment_id": experiment_id, + "trial_id": 1, + } + + +@pytest.fixture +def mock_env( + tunable_groups: TunableGroups, + mock_env_global_config: dict, +) -> MockEnv: """Test fixture for MockEnv.""" return MockEnv( name="Test Env", @@ -41,11 +58,15 @@ def mock_env(tunable_groups: TunableGroups) -> MockEnv: "mock_env_metrics": ["score"], }, tunables=tunable_groups, + global_config=mock_env_global_config, ) @pytest.fixture -def mock_env_no_noise(tunable_groups: TunableGroups) -> MockEnv: +def mock_env_no_noise( + tunable_groups: TunableGroups, + mock_env_global_config: dict, +) -> MockEnv: """Test fixture for MockEnv.""" return MockEnv( name="Test Env No Noise", @@ -56,6 +77,7 @@ def mock_env_no_noise(tunable_groups: TunableGroups) -> MockEnv: "mock_env_metrics": ["score", "other_score"], }, tunables=tunable_groups, + global_config=mock_env_global_config, ) diff --git a/mlos_bench/mlos_bench/tests/environments/composite_env_test.py b/mlos_bench/mlos_bench/tests/environments/composite_env_test.py index 94613102cd2..9952df4a7b4 100644 --- a/mlos_bench/mlos_bench/tests/environments/composite_env_test.py +++ b/mlos_bench/mlos_bench/tests/environments/composite_env_test.py @@ -75,7 +75,9 @@ def composite_env(tunable_groups: TunableGroups) -> CompositeEnv: }, tunables=tunable_groups, service=ConfigPersistenceService({}), - global_config={"global_param": "global_value"}, + global_config={ + "global_param": "global_value", + }, ) diff --git a/mlos_bench/mlos_bench/tests/environments/include_tunables_test.py b/mlos_bench/mlos_bench/tests/environments/include_tunables_test.py index 4d54e6aad44..2446983453b 100644 --- a/mlos_bench/mlos_bench/tests/environments/include_tunables_test.py +++ b/mlos_bench/mlos_bench/tests/environments/include_tunables_test.py @@ -50,6 +50,7 @@ def test_two_groups_setup(tunable_groups: TunableGroups) -> None: }, }, tunables=tunable_groups, + global_config=BUILT_IN_ENV_VAR_DEFAULTS, ) expected_params = { "vmSize": "Standard_B4ms", @@ -98,6 +99,7 @@ def test_zero_groups_implicit_setup(tunable_groups: TunableGroups) -> None: }, }, tunables=tunable_groups, + global_config=BUILT_IN_ENV_VAR_DEFAULTS, ) assert env.tunable_params.get_param_values() == {} @@ -137,7 +139,11 @@ def test_loader_level_include() -> None: ] } ) - env = loader.build_environment(config=env_json, tunables=TunableGroups()) + env = loader.build_environment( + config=env_json, + tunables=TunableGroups(), + global_config=BUILT_IN_ENV_VAR_DEFAULTS, + ) expected_params = { "align_va_addr": "on", "idle": "halt", diff --git a/mlos_bench/mlos_bench/tests/optimizers/opt_bulk_register_test.py b/mlos_bench/mlos_bench/tests/optimizers/opt_bulk_register_test.py index df40ddc839a..d15fd399f38 100644 --- a/mlos_bench/mlos_bench/tests/optimizers/opt_bulk_register_test.py +++ b/mlos_bench/mlos_bench/tests/optimizers/opt_bulk_register_test.py @@ -40,7 +40,12 @@ def mock_scores() -> list[dict[str, TunableValue] | None]: @pytest.fixture def mock_status() -> list[Status]: """Mock status values for earlier experiments.""" - return [Status.FAILED, Status.SUCCEEDED, Status.SUCCEEDED, Status.SUCCEEDED] + return [ + Status.FAILED, + Status.SUCCEEDED, + Status.SUCCEEDED, + Status.SUCCEEDED, + ] def _test_opt_update_min( diff --git a/mlos_bench/mlos_bench/tests/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/schedulers/conftest.py index df6bd2776fd..9fe62c400e6 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/conftest.py +++ b/mlos_bench/mlos_bench/tests/schedulers/conftest.py @@ -6,10 +6,8 @@ # pylint: disable=redefined-outer-name import json -import re import pytest -from pytest import FixtureRequest from mlos_bench.environments.mock_env import MockEnv from mlos_bench.schedulers.trial_runner import TrialRunner @@ -38,7 +36,7 @@ def mock_env_config() -> dict: "mock_trial_data": { "1": { "run": { - "sleep": 0.15, + "sleep": 0.25, "status": "SUCCEEDED", "metrics": { "score": 1.0, @@ -47,7 +45,7 @@ def mock_env_config() -> dict: }, "2": { "run": { - "sleep": 0.2, + "sleep": 0.3, "status": "SUCCEEDED", "metrics": { "score": 2.0, @@ -56,7 +54,7 @@ def mock_env_config() -> dict: }, "3": { "run": { - "sleep": 0.1, + "sleep": 0.2, "status": "SUCCEEDED", "metrics": { "score": 3.0, @@ -68,18 +66,6 @@ def mock_env_config() -> dict: } -@pytest.fixture -def global_config(request: FixtureRequest) -> dict: - """A global config for a MockEnv.""" - test_name = request.node.name - test_name = re.sub(r"[^a-zA-Z0-9]", "_", test_name) - experiment_id = f"TestExperiment-{test_name}" - return { - "experiment_id": experiment_id, - "trial_id": 1, - } - - @pytest.fixture def mock_env_json_config(mock_env_config: dict) -> str: """A JSON string of the mock_env_config.""" @@ -90,7 +76,7 @@ def mock_env_json_config(mock_env_config: dict) -> str: def mock_env( mock_env_json_config: str, tunable_groups: TunableGroups, - global_config: dict, + mock_env_global_config: dict, ) -> MockEnv: """A fixture to create a MockEnv instance using the mock_env_json_config.""" config_loader_service = ConfigPersistenceService() @@ -98,7 +84,7 @@ def mock_env( mock_env_json_config, tunable_groups, service=config_loader_service, - global_config=global_config, + global_config=mock_env_global_config, ) assert isinstance(mock_env, MockEnv) return mock_env @@ -108,16 +94,18 @@ def mock_env( def trial_runners( mock_env_json_config: str, tunable_groups: TunableGroups, - global_config: dict, + mock_env_global_config: dict, ) -> list[TrialRunner]: """A fixture to create a list of TrialRunner instances using the mock_env_json_config. """ - config_loader_service = ConfigPersistenceService(global_config=global_config) + config_loader_service = ConfigPersistenceService( + global_config=mock_env_global_config, + ) return TrialRunner.create_from_json( config_loader=config_loader_service, env_json=mock_env_json_config, tunable_groups=tunable_groups, num_trial_runners=NUM_TRIAL_RUNNERS, - global_config=global_config, + global_config=mock_env_global_config, ) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index f5bda73af68..bb652831c1a 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -84,7 +84,7 @@ def test_scheduler_with_mock_trial_data( trial_runners: list[TrialRunner], mock_opt: MockOptimizer, sqlite_storage: SqlStorage, - global_config: dict, + mock_env_global_config: dict, ) -> None: """ Full integration test for Scheduler: runs trials, checks storage, optimizer @@ -98,7 +98,7 @@ def test_scheduler_with_mock_trial_data( trial_runners, mock_opt, sqlite_storage, - global_config, + mock_env_global_config, ) root_env = scheduler.root_environment From 9e03e639cbd70c11bfaa24e5dc38e8f65ac91c15 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 18:15:06 -0500 Subject: [PATCH 63/77] remove old files --- ...-scheduler-invalid-mock-trial-data-2.jsonc | 11 ------- ...ck-scheduler-invalid-mock-trial-data.jsonc | 14 --------- .../bad/unhandled/mock_sched-extra.jsonc | 7 ----- .../good/full/mock_sched-full.jsonc | 31 ------------------- .../good/partial/mock_sched-partial.jsonc | 20 ------------ 5 files changed, 83 deletions(-) delete mode 100644 mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data-2.jsonc delete mode 100644 mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data.jsonc delete mode 100644 mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/unhandled/mock_sched-extra.jsonc delete mode 100644 mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/full/mock_sched-full.jsonc delete mode 100644 mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/partial/mock_sched-partial.jsonc diff --git a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data-2.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data-2.jsonc deleted file mode 100644 index 09b74f43773..00000000000 --- a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data-2.jsonc +++ /dev/null @@ -1,11 +0,0 @@ -{ - "class": "mlos_bench.tests.schedulers.mock_scheduler.MockScheduler", - "config": { - "mock_trial_data": [ - { - // MISSING: "trial_id": 1, - "status": "SUCCEEDED" - } - ] - } -} diff --git a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data.jsonc deleted file mode 100644 index 144e7f21e64..00000000000 --- a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/invalid/mock-scheduler-invalid-mock-trial-data.jsonc +++ /dev/null @@ -1,14 +0,0 @@ -{ - "class": "mlos_bench.tests.schedulers.mock_scheduler.MockScheduler", - "config": { - "mock_trial_data": [ - { - "trial_id": 1, - "status": "INVALID" - }, - { - "status": "SUCCEEDED" - } - ] - } -} diff --git a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/unhandled/mock_sched-extra.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/unhandled/mock_sched-extra.jsonc deleted file mode 100644 index c7bf9a46f75..00000000000 --- a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/bad/unhandled/mock_sched-extra.jsonc +++ /dev/null @@ -1,7 +0,0 @@ -{ - "class": "mlos_bench.tests.schedulers.mock_scheduler.MockScheduler", - "config": { - "extra": "unsupported" - }, - "extra": "unsupported" -} diff --git a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/full/mock_sched-full.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/full/mock_sched-full.jsonc deleted file mode 100644 index ca6ad2f9054..00000000000 --- a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/full/mock_sched-full.jsonc +++ /dev/null @@ -1,31 +0,0 @@ -{ - "$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" - } - } - ] - } -} diff --git a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/partial/mock_sched-partial.jsonc b/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/partial/mock_sched-partial.jsonc deleted file mode 100644 index 2557c25cb46..00000000000 --- a/mlos_bench/mlos_bench/tests/config/schemas/schedulers/test-cases/good/partial/mock_sched-partial.jsonc +++ /dev/null @@ -1,20 +0,0 @@ -{ - "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 - } - ] - } -} From f9e3339b2d22f07bce3348ace1055cb168b635e2 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Thu, 22 May 2025 18:15:14 -0500 Subject: [PATCH 64/77] add some common globals --- .../mlos_bench/tests/environments/composite_env_test.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/environments/composite_env_test.py b/mlos_bench/mlos_bench/tests/environments/composite_env_test.py index 9952df4a7b4..9cda1cea52e 100644 --- a/mlos_bench/mlos_bench/tests/environments/composite_env_test.py +++ b/mlos_bench/mlos_bench/tests/environments/composite_env_test.py @@ -77,6 +77,7 @@ def composite_env(tunable_groups: TunableGroups) -> CompositeEnv: service=ConfigPersistenceService({}), global_config={ "global_param": "global_value", + **BUILT_IN_ENV_VAR_DEFAULTS, }, ) @@ -233,7 +234,10 @@ def nested_composite_env(tunable_groups: TunableGroups) -> CompositeEnv: }, tunables=tunable_groups, service=ConfigPersistenceService({}), - global_config={"global_param": "global_value"}, + global_config={ + "global_param": "global_value", + **BUILT_IN_ENV_VAR_DEFAULTS, + }, ) From a05b5ae2bcf8ea8b5aad397d95f43a3273fc116d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 22 May 2025 23:15:55 +0000 Subject: [PATCH 65/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mlos_bench/mlos_bench/tests/conftest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/conftest.py b/mlos_bench/mlos_bench/tests/conftest.py index 2f8690784c5..158d1fffaf5 100644 --- a/mlos_bench/mlos_bench/tests/conftest.py +++ b/mlos_bench/mlos_bench/tests/conftest.py @@ -5,14 +5,14 @@ """Common fixtures for mock TunableGroups and Environment objects.""" import os +import re import sys from collections.abc import Generator from typing import Any -import re import pytest -from pytest import FixtureRequest from fasteners import InterProcessLock, InterProcessReaderWriterLock +from pytest import FixtureRequest from pytest_docker.plugin import Services as DockerServices from pytest_docker.plugin import get_docker_services From 94e31b5987b5a048cf0d3c527c591bb8fa7806f3 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Tue, 27 May 2025 14:02:34 -0500 Subject: [PATCH 66/77] simplify --- mlos_bench/mlos_bench/environments/mock_env.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/mlos_bench/mlos_bench/environments/mock_env.py b/mlos_bench/mlos_bench/environments/mock_env.py index b46c053e6b8..7ea83726b16 100644 --- a/mlos_bench/mlos_bench/environments/mock_env.py +++ b/mlos_bench/mlos_bench/environments/mock_env.py @@ -252,18 +252,10 @@ def get_current_mock_trial_data(self) -> MockTrialData: method is called to ensure the current ``trial_id`` is set. """ trial_id = self.current_trial_id - mock_trial_data = self._mock_trial_data.get(trial_id) - if not mock_trial_data: - mock_trial_data = MockTrialData( - trial_id=trial_id, - setup=MockTrialPhaseData.from_dict(phase="setup", data=None), - run=MockTrialPhaseData.from_dict(phase="run", data=None), - status=MockTrialPhaseData.from_dict(phase="status", data=None), - teardown=MockTrialPhaseData.from_dict(phase="teardown", data=None), - ) - # Save the generated data for later. - self._mock_trial_data[trial_id] = mock_trial_data - return mock_trial_data + if trial_id not in self._mock_trial_data: + # Make an empty mock trial data object if not found. + self._mock_trial_data[trial_id] = MockTrialData.from_dict(trial_id, data={}) + return self._mock_trial_data[trial_id] def setup(self, tunables: TunableGroups, global_config: dict | None = None) -> bool: is_success = super().setup(tunables, global_config) From 21e49db6247911c8ec5b402fae6998c372915e93 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Tue, 27 May 2025 14:37:52 -0500 Subject: [PATCH 67/77] rename --- mlos_bench/mlos_bench/environments/mock_env.py | 6 ++++++ mlos_bench/mlos_bench/tests/schedulers/conftest.py | 9 +++++++++ .../mlos_bench/tests/schedulers/test_scheduler.py | 14 +++++++------- 3 files changed, 22 insertions(+), 7 deletions(-) diff --git a/mlos_bench/mlos_bench/environments/mock_env.py b/mlos_bench/mlos_bench/environments/mock_env.py index 7ea83726b16..8a58e7e4255 100644 --- a/mlos_bench/mlos_bench/environments/mock_env.py +++ b/mlos_bench/mlos_bench/environments/mock_env.py @@ -12,6 +12,8 @@ from datetime import datetime from typing import Any +from pytz import UTC + import numpy from mlos_bench.environments.base_environment import Environment @@ -300,6 +302,8 @@ def run(self) -> tuple[Status, datetime, dict[str, TunableValue] | None]: if mock_trial_data.run.sleep: _LOG.debug("Sleeping for %s seconds", mock_trial_data.run.sleep) time.sleep(mock_trial_data.run.sleep) + # Update the timestamp after the sleep. + timestamp = datetime.now(UTC) if mock_trial_data.run.exception: raise RuntimeError(f"Mock trial data run exception: {mock_trial_data.run.exception}") if mock_trial_data.run.metrics is not None: @@ -327,6 +331,8 @@ def status(self) -> tuple[Status, datetime, list[tuple[datetime, str, Any]]]: if mock_trial_data.status.sleep: _LOG.debug("Sleeping for %s seconds", mock_trial_data.status.sleep) time.sleep(mock_trial_data.status.sleep) + # Update the timestamp after the sleep. + timestamp = datetime.now(UTC) if mock_trial_data.status.exception: raise RuntimeError( f"Mock trial data status exception: {mock_trial_data.status.exception}" diff --git a/mlos_bench/mlos_bench/tests/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/schedulers/conftest.py index 9fe62c400e6..cbd149f5e12 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/conftest.py +++ b/mlos_bench/mlos_bench/tests/schedulers/conftest.py @@ -35,6 +35,9 @@ def mock_env_config() -> dict: # TODO: Add more mock trial data here: "mock_trial_data": { "1": { + "setup": { + "sleep": 0.1, + }, "run": { "sleep": 0.25, "status": "SUCCEEDED", @@ -44,6 +47,9 @@ def mock_env_config() -> dict: }, }, "2": { + "setup": { + "sleep": 0.1, + }, "run": { "sleep": 0.3, "status": "SUCCEEDED", @@ -53,6 +59,9 @@ def mock_env_config() -> dict: }, }, "3": { + "setup": { + "sleep": 0.1, + }, "run": { "sleep": 0.2, "status": "SUCCEEDED", diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index bb652831c1a..2498e2633bd 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -151,17 +151,17 @@ def test_scheduler_with_mock_trial_data( ), f"Result value for {opt_target} is None." # Check that the appropriate sleeps occurred. - trial_time_lb = 0.0 - trial_time_lb += mock_trial_data.setup.sleep or 0 - trial_time_lb += mock_trial_data.run.sleep or 0 - trial_time_lb += mock_trial_data.status.sleep or 0 - trial_time_lb += mock_trial_data.teardown.sleep or 0 + min_trial_time = 0.0 + min_trial_time += mock_trial_data.setup.sleep or 0 + min_trial_time += mock_trial_data.run.sleep or 0 + min_trial_time += mock_trial_data.status.sleep or 0 + min_trial_time += mock_trial_data.teardown.sleep or 0 assert trial_data.ts_end is not None, f"Trial {trial_id} has no end time." trial_duration = trial_data.ts_end - trial_data.ts_start trial_dur_secs = trial_duration.total_seconds() assert ( - trial_dur_secs >= trial_time_lb - ), f"Trial {trial_id} took less time ({trial_dur_secs}) than expected ({trial_time_lb}). " + trial_dur_secs >= min_trial_time + ), f"Trial {trial_id} took less time ({trial_dur_secs}) than expected ({min_trial_time}). " # Check that the trial status matches what we expected. assert ( From 46be00487d4d51f0239f6c8bd958bd2d11b0374c Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Tue, 27 May 2025 15:54:53 -0500 Subject: [PATCH 68/77] fixup --- mlos_bench/mlos_bench/environments/mock_env.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/environments/mock_env.py b/mlos_bench/mlos_bench/environments/mock_env.py index 8a58e7e4255..921e4a3c50a 100644 --- a/mlos_bench/mlos_bench/environments/mock_env.py +++ b/mlos_bench/mlos_bench/environments/mock_env.py @@ -64,7 +64,7 @@ def from_dict(phase: str, data: dict | None) -> "MockTrialPhaseData": """ data = data or {} assert phase in {"setup", "run", "status", "teardown"}, f"Invalid phase: {phase}" - if phase in {"teardown", "status"}: + if phase in {"setup", "teardown"}: # setup/teardown phase is not expected to have metrics or status. assert "metrics" not in data, f"Unexpected metrics data in {phase} phase: {data}" assert "status" not in data, f"Unexpected status data in {phase} phase: {data}" From 01f7a02774d164f15775c048f31b06fc6c74394d Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Tue, 27 May 2025 15:56:53 -0500 Subject: [PATCH 69/77] always try and convert metrics back to numerics if possible --- mlos_bench/mlos_bench/storage/sql/common.py | 12 +++--------- mlos_bench/mlos_bench/storage/sql/experiment.py | 15 ++++++++++++--- mlos_bench/mlos_bench/storage/sql/trial_data.py | 15 +++++++++++---- .../mlos_bench/tests/storage/exp_load_test.py | 8 ++++---- .../mlos_bench/tests/storage/trial_config_test.py | 2 +- .../tests/storage/trial_telemetry_test.py | 14 +++++++------- 6 files changed, 38 insertions(+), 28 deletions(-) diff --git a/mlos_bench/mlos_bench/storage/sql/common.py b/mlos_bench/mlos_bench/storage/sql/common.py index 032cf9259d8..9ac1febf90f 100644 --- a/mlos_bench/mlos_bench/storage/sql/common.py +++ b/mlos_bench/mlos_bench/storage/sql/common.py @@ -16,7 +16,7 @@ from mlos_bench.storage.base_experiment_data import ExperimentData from mlos_bench.storage.base_trial_data import TrialData from mlos_bench.storage.sql.schema import DbSchema -from mlos_bench.util import nullable, utcify_nullable_timestamp, utcify_timestamp +from mlos_bench.util import nullable, try_parse_val, utcify_nullable_timestamp, utcify_timestamp def save_params( @@ -239,10 +239,7 @@ def get_results_df( columns="param", values="value", ) - configs_df = configs_df.apply( - pandas.to_numeric, - errors="coerce", - ).fillna(configs_df) + configs_df["value"] = configs_df["value"].apply(try_parse_val) # Get each trial's results in wide format. results_stmt = ( @@ -285,10 +282,7 @@ def get_results_df( columns="metric", values="value", ) - results_df = results_df.apply( - pandas.to_numeric, - errors="coerce", - ).fillna(results_df) + results_df["value"] = results_df["value"].apply(try_parse_val) # Concat the trials, configs, and results. return trials_df.merge(configs_df, on=["trial_id", "tunable_config_id"], how="left").merge( diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py index 0e380e3e13a..e1e348f1f2f 100644 --- a/mlos_bench/mlos_bench/storage/sql/experiment.py +++ b/mlos_bench/mlos_bench/storage/sql/experiment.py @@ -22,7 +22,7 @@ from mlos_bench.storage.sql.schema import DbSchema from mlos_bench.storage.sql.trial import Trial from mlos_bench.tunables.tunable_groups import TunableGroups -from mlos_bench.util import utcify_timestamp +from mlos_bench.util import try_parse_val, utcify_timestamp _LOG = logging.getLogger(__name__) @@ -149,7 +149,11 @@ def load_telemetry(self, trial_id: int) -> list[tuple[datetime, str, Any]]: # Not all storage backends store the original zone info. # We try to ensure data is entered in UTC and augment it on return again here. return [ - (utcify_timestamp(row.ts, origin="utc"), row.metric_id, row.metric_value) + ( + utcify_timestamp(row.ts, origin="utc"), + row.metric_id, + try_parse_val(row.metric_value), + ) for row in cur_telemetry.fetchall() ] @@ -229,10 +233,15 @@ def _get_key_val(conn: Connection, table: Table, field: str, **kwargs: Any) -> d .select_from(table) .where(*[column(key) == val for (key, val) in kwargs.items()]) ) + + def _tuple_to_kv(row_tuple: tuple[str, Any]) -> tuple[str, Any]: + return row_tuple[0], try_parse_val(row_tuple[1]) + # NOTE: `Row._tuple()` is NOT a protected member; the class uses `_` to # avoid naming conflicts. return dict( - row._tuple() for row in cur_result.fetchall() # pylint: disable=protected-access + _tuple_to_kv(row._tuple()) # pylint: disable=protected-access + for row in cur_result.fetchall() ) def get_trial_by_id( diff --git a/mlos_bench/mlos_bench/storage/sql/trial_data.py b/mlos_bench/mlos_bench/storage/sql/trial_data.py index 03ef19570ca..af5821591f3 100644 --- a/mlos_bench/mlos_bench/storage/sql/trial_data.py +++ b/mlos_bench/mlos_bench/storage/sql/trial_data.py @@ -16,7 +16,7 @@ from mlos_bench.storage.base_tunable_config_data import TunableConfigData from mlos_bench.storage.sql.schema import DbSchema from mlos_bench.storage.sql.tunable_config_data import TunableConfigSqlData -from mlos_bench.util import utcify_timestamp +from mlos_bench.util import try_parse_val, utcify_timestamp if TYPE_CHECKING: from mlos_bench.storage.base_tunable_config_trial_group_data import ( @@ -97,7 +97,10 @@ def results_df(self) -> pandas.DataFrame: ) ) return pandas.DataFrame( - [(row.metric_id, row.metric_value) for row in cur_results.fetchall()], + [ + (row.metric_id, try_parse_val(row.metric_value)) + for row in cur_results.fetchall() + ], columns=["metric", "value"], ) @@ -120,7 +123,11 @@ def telemetry_df(self) -> pandas.DataFrame: # We try to ensure data is entered in UTC and augment it on return again here. return pandas.DataFrame( [ - (utcify_timestamp(row.ts, origin="utc"), row.metric_id, row.metric_value) + ( + utcify_timestamp(row.ts, origin="utc"), + row.metric_id, + try_parse_val(row.metric_value), + ) for row in cur_telemetry.fetchall() ], columns=["ts", "metric", "value"], @@ -145,6 +152,6 @@ def metadata_df(self) -> pandas.DataFrame: ) ) return pandas.DataFrame( - [(row.param_id, row.param_value) for row in cur_params.fetchall()], + [(row.param_id, try_parse_val(row.param_value)) for row in cur_params.fetchall()], columns=["parameter", "value"], ) diff --git a/mlos_bench/mlos_bench/tests/storage/exp_load_test.py b/mlos_bench/mlos_bench/tests/storage/exp_load_test.py index e07cf80c70a..a482d070f15 100644 --- a/mlos_bench/mlos_bench/tests/storage/exp_load_test.py +++ b/mlos_bench/mlos_bench/tests/storage/exp_load_test.py @@ -104,12 +104,12 @@ def test_exp_trial_update_categ( [ { "idle": "halt", - "kernel_sched_latency_ns": "2000000", - "kernel_sched_migration_cost_ns": "-1", + "kernel_sched_latency_ns": 2000000, + "kernel_sched_migration_cost_ns": -1, "vmSize": "Standard_B4ms", } ], - [{"score": "99.9", "benchmark": "test"}], + [{"score": 99.9, "benchmark": "test"}], [Status.SUCCEEDED], ) @@ -153,7 +153,7 @@ def test_exp_trial_pending_3( (trial_ids, configs, scores, status) = exp_storage.load() assert trial_ids == [trial_fail.trial_id, trial_succ.trial_id] assert len(configs) == 2 - assert scores == [None, {"score": f"{score}"}] + assert scores == [None, {"score": score}] assert status == [Status.FAILED, Status.SUCCEEDED] assert tunable_groups.copy().assign(configs[0]).reset() == trial_fail.tunables assert tunable_groups.copy().assign(configs[1]).reset() == trial_succ.tunables diff --git a/mlos_bench/mlos_bench/tests/storage/trial_config_test.py b/mlos_bench/mlos_bench/tests/storage/trial_config_test.py index 30dabe05d9c..2dbabb92612 100644 --- a/mlos_bench/mlos_bench/tests/storage/trial_config_test.py +++ b/mlos_bench/mlos_bench/tests/storage/trial_config_test.py @@ -22,7 +22,7 @@ def test_exp_trial_pending(exp_storage: Storage.Experiment, tunable_groups: Tuna assert pending.tunables == tunable_groups assert pending.config() == { "location": "westus2", - "num_repeats": "100", + "num_repeats": 100, "experiment_id": "Test-001", "trial_id": trial.trial_id, } diff --git a/mlos_bench/mlos_bench/tests/storage/trial_telemetry_test.py b/mlos_bench/mlos_bench/tests/storage/trial_telemetry_test.py index d4b122ed2bb..a5c25feda7f 100644 --- a/mlos_bench/mlos_bench/tests/storage/trial_telemetry_test.py +++ b/mlos_bench/mlos_bench/tests/storage/trial_telemetry_test.py @@ -13,7 +13,7 @@ from mlos_bench.storage.base_storage import Storage from mlos_bench.tests import ZONE_INFO from mlos_bench.tunables.tunable_groups import TunableGroups -from mlos_bench.util import nullable +from mlos_bench.util import nullable, try_parse_val # pylint: disable=redefined-outer-name @@ -41,12 +41,12 @@ def zoned_telemetry_data(zone_info: tzinfo | None) -> list[tuple[datetime, str, ) -def _telemetry_str( +def _telemetry_val( data: list[tuple[datetime, str, Any]], -) -> list[tuple[datetime, str, str | None]]: +) -> list[tuple[datetime, str, int | float | str | None]]: """Convert telemetry values to strings.""" # All retrieved timestamps should have been converted to UTC. - return [(ts.astimezone(UTC), key, nullable(str, val)) for (ts, key, val) in data] + return [(ts.astimezone(UTC), key, try_parse_val(val)) for (ts, key, val) in data] @pytest.mark.parametrize(("origin_zone_info"), ZONE_INFO) @@ -62,13 +62,13 @@ def test_update_telemetry( assert exp_storage.load_telemetry(trial.trial_id) == [] trial.update_telemetry(Status.RUNNING, datetime.now(origin_zone_info), telemetry_data) - assert exp_storage.load_telemetry(trial.trial_id) == _telemetry_str(telemetry_data) + assert exp_storage.load_telemetry(trial.trial_id) == _telemetry_val(telemetry_data) # Also check that the TrialData telemetry looks right. trial_data = storage.experiments[exp_storage.experiment_id].trials[trial.trial_id] trial_telemetry_df = trial_data.telemetry_df trial_telemetry_data = [tuple(r) for r in trial_telemetry_df.to_numpy()] - assert _telemetry_str(trial_telemetry_data) == _telemetry_str(telemetry_data) + assert _telemetry_val(trial_telemetry_data) == _telemetry_val(telemetry_data) @pytest.mark.parametrize(("origin_zone_info"), ZONE_INFO) @@ -84,4 +84,4 @@ def test_update_telemetry_twice( trial.update_telemetry(Status.RUNNING, timestamp, telemetry_data) trial.update_telemetry(Status.RUNNING, timestamp, telemetry_data) trial.update_telemetry(Status.RUNNING, timestamp, telemetry_data) - assert exp_storage.load_telemetry(trial.trial_id) == _telemetry_str(telemetry_data) + assert exp_storage.load_telemetry(trial.trial_id) == _telemetry_val(telemetry_data) From b035e57d1df80ed6bccddd81683505e77f782a9a Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Tue, 27 May 2025 15:57:26 -0500 Subject: [PATCH 70/77] wip: tell mysql to use a datetime that can store fractional seconds --- .../mlos_bench/storage/sql/experiment.py | 6 +--- mlos_bench/mlos_bench/storage/sql/schema.py | 31 +++++++++++++++---- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py index e1e348f1f2f..4181df95f32 100644 --- a/mlos_bench/mlos_bench/storage/sql/experiment.py +++ b/mlos_bench/mlos_bench/storage/sql/experiment.py @@ -376,11 +376,7 @@ def _new_trial( ts_start: datetime | None = None, config: dict[str, Any] | None = None, ) -> Storage.Trial: - # MySQL can round microseconds into the future causing scheduler to skip trials. - # Truncate microseconds to avoid this issue. - ts_start = utcify_timestamp(ts_start or datetime.now(UTC), origin="local").replace( - microsecond=0 - ) + ts_start = utcify_timestamp(ts_start or datetime.now(UTC), origin="local") _LOG.debug("Create trial: %s:%d @ %s", self._experiment_id, self._trial_id, ts_start) with self._engine.begin() as conn: try: diff --git a/mlos_bench/mlos_bench/storage/sql/schema.py b/mlos_bench/mlos_bench/storage/sql/schema.py index 2bc00f00825..dd0cbae586e 100644 --- a/mlos_bench/mlos_bench/storage/sql/schema.py +++ b/mlos_bench/mlos_bench/storage/sql/schema.py @@ -40,6 +40,7 @@ inspect, ) from sqlalchemy.engine import Engine +from sqlalchemy.dialects import mysql from mlos_bench.util import path_join @@ -104,8 +105,8 @@ def __init__(self, engine: Engine | None): Column("git_repo", String(1024), nullable=False), Column("git_commit", String(40), nullable=False), # For backwards compatibility, we allow NULL for ts_start. - Column("ts_start", DateTime), - Column("ts_end", DateTime), + Column("ts_start", DateTime().with_variant(mysql.DATETIME(fsp=6), "mysql")), + Column("ts_end", DateTime().with_variant(mysql.DATETIME(fsp=6), "mysql")), # Should match the text IDs of `mlos_bench.environments.Status` enum: # For backwards compatibility, we allow NULL for status. Column("status", String(self._STATUS_LEN)), @@ -179,8 +180,16 @@ def __init__(self, engine: Engine | None): Column("trial_id", Integer, nullable=False), Column("config_id", Integer, nullable=False), Column("trial_runner_id", Integer, nullable=True, default=None), - Column("ts_start", DateTime, nullable=False), - Column("ts_end", DateTime), + Column( + "ts_start", + DateTime().with_variant(mysql.DATETIME(fsp=6), "mysql"), + nullable=False, + ), + Column( + "ts_end", + DateTime().with_variant(mysql.DATETIME(fsp=6), "mysql"), + nullable=True, + ), # Should match the text IDs of `mlos_bench.environments.Status` enum: Column("status", String(self._STATUS_LEN), nullable=False), PrimaryKeyConstraint("exp_id", "trial_id"), @@ -232,7 +241,12 @@ def __init__(self, engine: Engine | None): self._meta, Column("exp_id", String(self._ID_LEN), nullable=False), Column("trial_id", Integer, nullable=False), - Column("ts", DateTime(timezone=True), nullable=False, default="now"), + Column( + "ts", + DateTime(timezone=True).with_variant(mysql.DATETIME(fsp=6), "mysql"), + nullable=False, + default="now", + ), Column("status", String(self._STATUS_LEN), nullable=False), UniqueConstraint("exp_id", "trial_id", "ts"), ForeignKeyConstraint( @@ -267,7 +281,12 @@ def __init__(self, engine: Engine | None): self._meta, Column("exp_id", String(self._ID_LEN), nullable=False), Column("trial_id", Integer, nullable=False), - Column("ts", DateTime(timezone=True), nullable=False, default="now"), + Column( + "ts", + DateTime(timezone=True).with_variant(mysql.DATETIME(fsp=6), "mysql"), + nullable=False, + default="now", + ), Column("metric_id", String(self._ID_LEN), nullable=False), Column("metric_value", String(self._METRIC_VALUE_LEN)), UniqueConstraint("exp_id", "trial_id", "ts", "metric_id"), From ca101ef2b0b553e1d0a18e39b033ef3b92623eaf Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Tue, 27 May 2025 15:57:45 -0500 Subject: [PATCH 71/77] add more trial data for testing --- .../mlos_bench/tests/schedulers/conftest.py | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/mlos_bench/mlos_bench/tests/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/schedulers/conftest.py index cbd149f5e12..3aac9933c2c 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/conftest.py +++ b/mlos_bench/mlos_bench/tests/schedulers/conftest.py @@ -38,6 +38,13 @@ def mock_env_config() -> dict: "setup": { "sleep": 0.1, }, + "status": { + "sleep": 0.1, + "metrics": { + "color": "orange", + "tps": 8, + }, + }, "run": { "sleep": 0.25, "status": "SUCCEEDED", @@ -50,6 +57,13 @@ def mock_env_config() -> dict: "setup": { "sleep": 0.1, }, + "status": { + "sleep": 0.1, + "metrics": { + "color": "purple", + "tps": 7, + }, + }, "run": { "sleep": 0.3, "status": "SUCCEEDED", @@ -62,6 +76,13 @@ def mock_env_config() -> dict: "setup": { "sleep": 0.1, }, + "status": { + "sleep": 0.1, + "metrics": { + "color": "blue", + "tps": 6789, + }, + }, "run": { "sleep": 0.2, "status": "SUCCEEDED", @@ -70,6 +91,42 @@ def mock_env_config() -> dict: }, }, }, + "4": { + "setup": { + "sleep": 0.1, + }, + "status": { + "sleep": 0.1, + "metrics": { + "color": "blue", + "tps": 5, + }, + }, + "run": { + "sleep": 0.2, + "status": "SUCCEEDED", + "metrics": { + "score": 2.7, + }, + }, + }, + "5": { + "setup": { + "sleep": 0.1, + }, + "status": { + "sleep": 0.1, + "metrics": {}, + }, + "run": { + "sleep": 0.2, + "status": "SUCCEEDED", + "metrics": { + "color": "green", + "score": 3.1, + }, + }, + }, }, }, } From 98ec1cc6de4faed72e9d4718855353461f6d84d7 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Tue, 27 May 2025 15:58:07 -0500 Subject: [PATCH 72/77] :type fixups --- mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 2498e2633bd..04d3194cbab 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -36,7 +36,7 @@ def create_scheduler( """ env = trial_runners[0].environment assert isinstance(env, MockEnv), "Environment is not a MockEnv instance." - max_trials = max(trial_id for trial_id in env.mock_trial_data.keys()) + max_trials = max(int(trial_id) for trial_id in env.mock_trial_data.keys()) max_trials = min(max_trials, mock_opt.max_suggestions) return scheduler_type( From 82f602cc0c402c992a92aa8f73c2e67efb7f755d Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Tue, 27 May 2025 15:58:20 -0500 Subject: [PATCH 73/77] check the telemetry --- mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 04d3194cbab..6cd0496a837 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -168,7 +168,14 @@ def test_scheduler_with_mock_trial_data( trial_data.status == mock_trial_data.run.status ), f"Trial {trial_id} status {trial_data.status} was not {mock_trial_data.run.status}." - # TODO: Check the trial status telemetry. + # Check the trial status telemetry. + telemetry_dict = dict( + zip(trial_data.telemetry_df["metric"], trial_data.telemetry_df["value"]) + ) + assert telemetry_dict == mock_trial_data.status.metrics, ( + f"Trial {trial_id} telemetry {telemetry_dict} does not match expected " + f"{mock_trial_data.status.metrics}." + ) # Check the optimizer registration. assert mock_opt_has_registered_trial_score( From 176480fc6a099db32049ab63e0dd8e59fffbb747 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 27 May 2025 20:59:52 +0000 Subject: [PATCH 74/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mlos_bench/mlos_bench/environments/mock_env.py | 3 +-- mlos_bench/mlos_bench/storage/sql/common.py | 7 ++++++- mlos_bench/mlos_bench/storage/sql/schema.py | 2 +- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/mlos_bench/mlos_bench/environments/mock_env.py b/mlos_bench/mlos_bench/environments/mock_env.py index 921e4a3c50a..a2b7f411a7d 100644 --- a/mlos_bench/mlos_bench/environments/mock_env.py +++ b/mlos_bench/mlos_bench/environments/mock_env.py @@ -12,9 +12,8 @@ from datetime import datetime from typing import Any -from pytz import UTC - import numpy +from pytz import UTC from mlos_bench.environments.base_environment import Environment from mlos_bench.environments.status import Status diff --git a/mlos_bench/mlos_bench/storage/sql/common.py b/mlos_bench/mlos_bench/storage/sql/common.py index 9ac1febf90f..5a59da35dfa 100644 --- a/mlos_bench/mlos_bench/storage/sql/common.py +++ b/mlos_bench/mlos_bench/storage/sql/common.py @@ -16,7 +16,12 @@ from mlos_bench.storage.base_experiment_data import ExperimentData from mlos_bench.storage.base_trial_data import TrialData from mlos_bench.storage.sql.schema import DbSchema -from mlos_bench.util import nullable, try_parse_val, utcify_nullable_timestamp, utcify_timestamp +from mlos_bench.util import ( + nullable, + try_parse_val, + utcify_nullable_timestamp, + utcify_timestamp, +) def save_params( diff --git a/mlos_bench/mlos_bench/storage/sql/schema.py b/mlos_bench/mlos_bench/storage/sql/schema.py index dd0cbae586e..3cbb63bcc45 100644 --- a/mlos_bench/mlos_bench/storage/sql/schema.py +++ b/mlos_bench/mlos_bench/storage/sql/schema.py @@ -39,8 +39,8 @@ create_mock_engine, inspect, ) -from sqlalchemy.engine import Engine from sqlalchemy.dialects import mysql +from sqlalchemy.engine import Engine from mlos_bench.util import path_join From 22f4e79655767205826e629bca88fe7ed9cd1d1c Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Tue, 27 May 2025 16:05:28 -0500 Subject: [PATCH 75/77] revert the bulk conversion one --- mlos_bench/mlos_bench/storage/sql/common.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/mlos_bench/mlos_bench/storage/sql/common.py b/mlos_bench/mlos_bench/storage/sql/common.py index 5a59da35dfa..032cf9259d8 100644 --- a/mlos_bench/mlos_bench/storage/sql/common.py +++ b/mlos_bench/mlos_bench/storage/sql/common.py @@ -16,12 +16,7 @@ from mlos_bench.storage.base_experiment_data import ExperimentData from mlos_bench.storage.base_trial_data import TrialData from mlos_bench.storage.sql.schema import DbSchema -from mlos_bench.util import ( - nullable, - try_parse_val, - utcify_nullable_timestamp, - utcify_timestamp, -) +from mlos_bench.util import nullable, utcify_nullable_timestamp, utcify_timestamp def save_params( @@ -244,7 +239,10 @@ def get_results_df( columns="param", values="value", ) - configs_df["value"] = configs_df["value"].apply(try_parse_val) + configs_df = configs_df.apply( + pandas.to_numeric, + errors="coerce", + ).fillna(configs_df) # Get each trial's results in wide format. results_stmt = ( @@ -287,7 +285,10 @@ def get_results_df( columns="metric", values="value", ) - results_df["value"] = results_df["value"].apply(try_parse_val) + results_df = results_df.apply( + pandas.to_numeric, + errors="coerce", + ).fillna(results_df) # Concat the trials, configs, and results. return trials_df.merge(configs_df, on=["trial_id", "tunable_config_id"], how="left").merge( From aa9f728276a4db921d7f80ef19c9b59018ffa157 Mon Sep 17 00:00:00 2001 From: Brian Kroth Date: Tue, 27 May 2025 16:41:40 -0500 Subject: [PATCH 76/77] debugging --- .../mlos_bench/tests/schedulers/conftest.py | 1 + .../tests/schedulers/test_scheduler.py | 52 +++++++++++++++++-- 2 files changed, 48 insertions(+), 5 deletions(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/conftest.py b/mlos_bench/mlos_bench/tests/schedulers/conftest.py index 3aac9933c2c..e82c9ab7b07 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/conftest.py +++ b/mlos_bench/mlos_bench/tests/schedulers/conftest.py @@ -122,6 +122,7 @@ def mock_env_config() -> dict: "sleep": 0.2, "status": "SUCCEEDED", "metrics": { + # Return multiple scores. "color": "green", "score": 3.1, }, diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 6cd0496a837..1b2911239f5 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -4,6 +4,7 @@ # """Unit tests for :py:class:`mlos_bench.schedulers` and their internals.""" +from logging import warning import sys import pytest @@ -23,6 +24,8 @@ # pylint: disable=redefined-outer-name +DEBUGGING = False # Set to True to enable debugging output for easier stepping through the code. + def create_scheduler( scheduler_type: type[Scheduler], @@ -51,17 +54,56 @@ def create_scheduler( ) +def is_subset_of(dict_a: dict, dict_b: dict) -> bool: + """Check if dict_a is a subset of dict_b.""" + return all(item in dict_b.items() for item in dict_a.items()) + + def mock_opt_has_registered_trial_score( mock_opt: MockOptimizer, trial_data: TrialData, ) -> bool: """Check that the MockOptimizer has registered a given MockTrialData.""" - return any( - registered_score.status == trial_data.status - and registered_score.score == trial_data.results_dict - and registered_score.config.get_param_values() == trial_data.tunable_config.config_dict - for registered_score in mock_opt.registered_scores + if not DEBUGGING: + return any( + registered_score.status == trial_data.status + and registered_score.score is not None + and is_subset_of(registered_score.score, trial_data.results_dict) + and registered_score.config.get_param_values() == trial_data.tunable_config.config_dict + for registered_score in mock_opt.registered_scores + ) + # For debugging, we can print the data to examine mismatches. + for registered_score in mock_opt.registered_scores: + if registered_score.status != trial_data.status: + warning( + f"Registered status: {registered_score.status} " + f"!= TrialData status: {trial_data.status}" + ) + continue + # Check if registered_score.score is a subset of trial_data.results_dict + if not ( + registered_score.score is not None + and is_subset_of(registered_score.score, trial_data.results_dict) + ): + warning( + f"Registered score: {registered_score.score} " + f"is not a subset of TrialData results: {trial_data.results_dict}" + ) + continue + if registered_score.config.get_param_values() != trial_data.tunable_config.config_dict: + warning( + f"Registered config: {registered_score.config.get_param_values()} " + f"!= TrialData config: {trial_data.tunable_config.config_dict}" + ) + continue + # Else, found a match! + warning(f"Found matching registered score for trial {trial_data}: {registered_score}") + return True + warning( + f"No matching registered score found for trial {trial_data}. " + f"Registered scores: {mock_opt.registered_scores}" ) + return False scheduler_classes = get_all_concrete_subclasses( From 628921646a68b0e80eb583559528d06dc98e0f89 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 27 May 2025 21:42:12 +0000 Subject: [PATCH 77/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py index 1b2911239f5..3375bf72527 100644 --- a/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py +++ b/mlos_bench/mlos_bench/tests/schedulers/test_scheduler.py @@ -4,8 +4,8 @@ # """Unit tests for :py:class:`mlos_bench.schedulers` and their internals.""" -from logging import warning import sys +from logging import warning import pytest