diff --git a/mlos_bench/mlos_bench/optimizers/base_optimizer.py b/mlos_bench/mlos_bench/optimizers/base_optimizer.py index 62b24e642d6..ee2e73451f4 100644 --- a/mlos_bench/mlos_bench/optimizers/base_optimizer.py +++ b/mlos_bench/mlos_bench/optimizers/base_optimizer.py @@ -153,6 +153,8 @@ def name(self) -> str: """ return self.__class__.__name__ + # TODO: Expand these properties for multi-objective. + @property def is_min(self) -> bool: """ @@ -167,6 +169,13 @@ def target(self) -> str: """ return self._opt_target + @property + def direction(self) -> str: + """ + The direction to optimize the target metric (e.g., min or max). + """ + return 'min' if self.is_min else 'max' + @property def supports_preload(self) -> bool: """ diff --git a/mlos_bench/mlos_bench/run.py b/mlos_bench/mlos_bench/run.py index 7240097ba56..765f9b97bb1 100755 --- a/mlos_bench/mlos_bench/run.py +++ b/mlos_bench/mlos_bench/run.py @@ -86,7 +86,8 @@ def _optimize(*, trial_id=trial_id, root_env_config=root_env_config, description=env.name, - opt_target=opt.target + opt_target=opt.target, + opt_direction=opt.direction, ) as exp: _LOG.info("Experiment: %s Env: %s Optimizer: %s", exp, env, opt) @@ -118,9 +119,13 @@ def _optimize(*, config_id = -1 trial = exp.new_trial(tunables, config={ + # Add some additional metadata to track for the trial such as the + # optimizer config used. + # TODO: Improve for supporting multi-objective + # (e.g., opt_target_1, opt_target_2, ... and opt_direction_1, opt_direction_2, ...) "optimizer": opt.name, "opt_target": opt.target, - "opt_direction": "min" if opt.is_min else "max", + "opt_direction": opt.direction, }) _run(env_context, opt_context, trial, global_config) diff --git a/mlos_bench/mlos_bench/storage/base_experiment_data.py b/mlos_bench/mlos_bench/storage/base_experiment_data.py index f9240d8fe7a..b86347261aa 100644 --- a/mlos_bench/mlos_bench/storage/base_experiment_data.py +++ b/mlos_bench/mlos_bench/storage/base_experiment_data.py @@ -56,6 +56,19 @@ def root_env_config(self) -> Tuple[str, str, str]: def __repr__(self) -> str: return f"Experiment :: {self._exp_id}: '{self._description}'" + @property + @abstractmethod + def objectives(self) -> Dict[str, str]: + """ + Retrieve the experiment's objectives data from the storage. + + Returns + ------- + objectives : Dict[str, objective] + A dictionary of the experiment's objective names (optimization_targets) + and their directions (e.g., min or max). + """ + @property @abstractmethod def trials(self) -> Dict[int, TrialData]: diff --git a/mlos_bench/mlos_bench/storage/base_storage.py b/mlos_bench/mlos_bench/storage/base_storage.py index ef1ad32da0a..1afd699f162 100644 --- a/mlos_bench/mlos_bench/storage/base_storage.py +++ b/mlos_bench/mlos_bench/storage/base_storage.py @@ -83,7 +83,8 @@ def experiment(self, *, trial_id: int, root_env_config: str, description: str, - opt_target: str) -> 'Storage.Experiment': + opt_target: str, + opt_direction: Optional[str]) -> 'Storage.Experiment': """ Create a new experiment in the storage. @@ -103,6 +104,8 @@ def experiment(self, *, Human-readable description of the experiment. opt_target : str Name of metric we're optimizing for. + opt_direction: Optional[str] + Direction to optimize the metric (e.g., min or max) Returns ------- @@ -112,15 +115,29 @@ def experiment(self, *, """ class Experiment(metaclass=ABCMeta): + # pylint: disable=too-many-instance-attributes """ Base interface for storing the results of the experiment. This class is instantiated in the `Storage.experiment()` method. """ - def __init__(self, tunables: TunableGroups, experiment_id: str, root_env_config: str): + def __init__(self, + *, + tunables: TunableGroups, + experiment_id: str, + trial_id: int, + root_env_config: str, + description: str, + opt_target: str, + opt_direction: Optional[str]): self._tunables = tunables.copy() + self._trial_id = trial_id self._experiment_id = experiment_id (self._git_repo, self._git_commit, self._root_env_config) = get_git_info(root_env_config) + self._description = description + self._opt_target = opt_target + assert opt_direction in {None, "min", "max"} + self._opt_direction = opt_direction def __enter__(self) -> 'Storage.Experiment': """ @@ -172,6 +189,31 @@ def _teardown(self, is_ok: bool) -> None: True if there were no exceptions during the experiment, False otherwise. """ + @property + def experiment_id(self) -> str: + """Get the Experiment's ID""" + return self._experiment_id + + @property + def trial_id(self) -> int: + """Get the current Trial ID""" + return self._trial_id + + @property + def description(self) -> str: + """Get the Experiment's description""" + return self._description + + @property + def opt_target(self) -> str: + """Get the Experiment's optimization target""" + return self._opt_target + + @property + def opt_direction(self) -> Optional[str]: + """Get the Experiment's optimization target""" + return self._opt_direction + @abstractmethod def merge(self, experiment_ids: List[str]) -> None: """ @@ -249,12 +291,15 @@ class Trial(metaclass=ABCMeta): def __init__(self, *, tunables: TunableGroups, experiment_id: str, trial_id: int, - config_id: int, opt_target: str, config: Optional[Dict[str, Any]] = None): + config_id: int, opt_target: str, opt_direction: Optional[str], + config: Optional[Dict[str, Any]] = None): self._tunables = tunables self._experiment_id = experiment_id self._trial_id = trial_id self._config_id = config_id self._opt_target = opt_target + assert opt_direction in {None, "min", "max"} + self._opt_direction = opt_direction self._config = config or {} def __repr__(self) -> str: @@ -274,6 +319,20 @@ def config_id(self) -> int: """ return self._config_id + @property + def opt_target(self) -> str: + """ + Get the Trial's optimization target. + """ + return self._opt_target + + @property + def opt_direction(self) -> Optional[str]: + """ + Get the Trial's optimization direction (e.g., min or max) + """ + return self._opt_direction + @property def tunables(self) -> TunableGroups: """ diff --git a/mlos_bench/mlos_bench/storage/sql/experiment.py b/mlos_bench/mlos_bench/storage/sql/experiment.py index 64160becefa..1bb89e2319e 100644 --- a/mlos_bench/mlos_bench/storage/sql/experiment.py +++ b/mlos_bench/mlos_bench/storage/sql/experiment.py @@ -35,13 +35,18 @@ def __init__(self, *, trial_id: int, root_env_config: str, description: str, - opt_target: str): - super().__init__(tunables, experiment_id, root_env_config) + opt_target: str, + opt_direction: Optional[str]): + super().__init__( + tunables=tunables, + experiment_id=experiment_id, + trial_id=trial_id, + root_env_config=root_env_config, + description=description, + opt_target=opt_target, + opt_direction=opt_direction) self._engine = engine self._schema = schema - self._trial_id = trial_id - self._description = description - self._opt_target = opt_target def _setup(self) -> None: super()._setup() @@ -76,11 +81,21 @@ def _setup(self) -> None: git_commit=self._git_commit, root_env_config=self._root_env_config, )) + # TODO: Expand for multiple objectives. + conn.execute(self._schema.objectives.insert().values( + exp_id=self._experiment_id, + optimization_target=self._opt_target, + optimization_direction=self._opt_direction, + )) else: if exp_info.trial_id is not None: self._trial_id = exp_info.trial_id + 1 _LOG.info("Continue experiment: %s last trial: %s resume from: %d", self._experiment_id, exp_info.trial_id, self._trial_id) + # TODO: Sanity check that certain critical configs (e.g., + # objectives) haven't changed to be incompatible such that a new + # experiment should be started (possibly by prewarming with the + # previous one). if exp_info.git_commit != self._git_commit: _LOG.warning("Experiment %s git expected: %s %s", self, exp_info.git_repo, exp_info.git_commit) @@ -131,6 +146,8 @@ def load(self, opt_target: Optional[str] = None) -> Tuple[List[dict], List[Optio self._schema.trial.c.trial_id.asc(), ) ) + # Note: this iterative approach is somewhat expensive. + # TODO: Look into a better bulk fetch option. for trial in cur_trials.fetchall(): tunables = self._get_params( conn, self._schema.config_param, config_id=trial.config_id) @@ -180,6 +197,7 @@ def pending_trials(self) -> Iterator[Storage.Trial]: trial_id=trial.trial_id, config_id=trial.config_id, opt_target=self._opt_target, + opt_direction=self._opt_direction, config=config, ) @@ -232,6 +250,7 @@ def new_trial(self, tunables: TunableGroups, trial_id=self._trial_id, config_id=config_id, opt_target=self._opt_target, + opt_direction=self._opt_direction, config=config, ) self._trial_id += 1 diff --git a/mlos_bench/mlos_bench/storage/sql/experiment_data.py b/mlos_bench/mlos_bench/storage/sql/experiment_data.py index b9b0cc3582b..aa90d7fd6b5 100644 --- a/mlos_bench/mlos_bench/storage/sql/experiment_data.py +++ b/mlos_bench/mlos_bench/storage/sql/experiment_data.py @@ -7,6 +7,8 @@ """ from typing import Dict +import logging + import pandas from sqlalchemy import Engine @@ -16,6 +18,8 @@ from mlos_bench.storage.base_trial_data import TrialData from mlos_bench.storage.sql.trial_data import TrialSqlData +_LOG = logging.getLogger(__name__) + class ExperimentSqlData(ExperimentData): """ @@ -34,6 +38,56 @@ def __init__(self, *, engine: Engine, schema: DbSchema, exp_id: str, self._engine = engine self._schema = schema + @property + def objectives(self) -> Dict[str, str]: + objectives: Dict[str, str] = {} + # First try to lookup the objectives from the experiment metadata in the storage layer. + if hasattr(self._schema, "objectives"): + with self._engine.connect() as conn: + objectives_db_data = conn.execute( + self._schema.objectives.select().where( + self._schema.objectives.c.exp_id == self._exp_id, + ).order_by( + self._schema.objectives.c.weight.desc(), + self._schema.objectives.c.optimization_target.asc(), + ) + ) + objectives = { + objective.optimization_target: objective.optimization_direction + for objective in objectives_db_data.fetchall() + } + # Backwards compatibility: try and obtain the objectives from the TrialData and merge them in. + # NOTE: The original format of storing opt_target/opt_direction in the Trial + # metadata did not support multi-objectives. + # Nor does it make it easy to detect when a config change caused a switch in + # opt_direction for a given opt_target between run.py executions of an + # Experiment. + # For now, we simply issue a warning about potentially inconsistent data. + for trial in self.trials.values(): + trial_objs_df = trial.metadata[ + trial.metadata["parameter"].isin(("opt_target", "opt_direction")) + ][["parameter", "value"]] + try: + opt_targets = trial_objs_df[trial_objs_df["parameter"] == "opt_target"] + assert len(opt_targets) == 1, \ + "Should only be a single opt_target in the metadata params." + opt_target = opt_targets["value"].iloc[0] + except KeyError: + continue + try: + opt_directions = trial_objs_df[trial_objs_df["parameter"] == "opt_direction"] + assert len(opt_directions) <= 1, \ + "Should only be a single opt_direction in the metadata params." + opt_direction = opt_directions["value"].iloc[0] + except (KeyError, IndexError): + opt_direction = None + if opt_target not in objectives: + objectives[opt_target] = opt_direction + elif opt_direction != objectives[opt_target]: + _LOG.warning("Experiment %s has multiple trial optimization directions for optimization_target %s=%s", + self, opt_target, objectives[opt_target]) + return objectives + @property def trials(self) -> Dict[int, TrialData]: with self._engine.connect() as conn: diff --git a/mlos_bench/mlos_bench/storage/sql/schema.py b/mlos_bench/mlos_bench/storage/sql/schema.py index 7b198e838be..74d11dabc04 100644 --- a/mlos_bench/mlos_bench/storage/sql/schema.py +++ b/mlos_bench/mlos_bench/storage/sql/schema.py @@ -11,7 +11,7 @@ from sqlalchemy import ( Engine, MetaData, Dialect, create_mock_engine, - Table, Column, Sequence, Integer, String, DateTime, + Table, Column, Sequence, Integer, Float, String, DateTime, PrimaryKeyConstraint, ForeignKeyConstraint, UniqueConstraint, ) @@ -62,10 +62,28 @@ def __init__(self, engine: Engine): Column("root_env_config", String(1024), nullable=False), Column("git_repo", String(1024), nullable=False), Column("git_commit", String(40), nullable=False), + Column("optimization_target", String(1024), nullable=True), + Column("optimization_direction", String(10), nullable=True), PrimaryKeyConstraint("exp_id"), ) + self.objectives = Table( + "objectives", + self._meta, + Column("exp_id"), + Column("optimization_target", String(1024), nullable=False), + Column("optimization_direction", String(4), nullable=False), + # TODO: Note: weight is not fully supported yet as currently + # multi-objective is expected to explore each objective equally. + # Will need to adjust the insert and return values to support this + # eventually. + Column("weight", Float, nullable=True), + + PrimaryKeyConstraint("exp_id", "optimization_target"), + ForeignKeyConstraint(["exp_id"], [self.experiment.c.exp_id]), + ) + # A workaround for SQLAlchemy issue with autoincrement in DuckDB: if engine.dialect.name == "duckdb": seq_config_id = Sequence('seq_config_id') diff --git a/mlos_bench/mlos_bench/storage/sql/storage.py b/mlos_bench/mlos_bench/storage/sql/storage.py index 32d04e7d442..9c747f8885c 100644 --- a/mlos_bench/mlos_bench/storage/sql/storage.py +++ b/mlos_bench/mlos_bench/storage/sql/storage.py @@ -62,7 +62,8 @@ def experiment(self, *, trial_id: int, root_env_config: str, description: str, - opt_target: str) -> Storage.Experiment: + opt_target: str, + opt_direction: Optional[str]) -> Storage.Experiment: return Experiment( engine=self._engine, schema=self._schema, @@ -72,6 +73,7 @@ def experiment(self, *, root_env_config=root_env_config, description=description, opt_target=opt_target, + opt_direction=opt_direction, ) @property diff --git a/mlos_bench/mlos_bench/storage/sql/trial.py b/mlos_bench/mlos_bench/storage/sql/trial.py index 242ab5f0e94..4e13ef8e889 100644 --- a/mlos_bench/mlos_bench/storage/sql/trial.py +++ b/mlos_bench/mlos_bench/storage/sql/trial.py @@ -29,13 +29,15 @@ class Trial(Storage.Trial): def __init__(self, *, engine: Engine, schema: DbSchema, tunables: TunableGroups, experiment_id: str, trial_id: int, config_id: int, - opt_target: str, config: Optional[Dict[str, Any]] = None): + opt_target: str, opt_direction: Optional[str], + config: Optional[Dict[str, Any]] = None): super().__init__( tunables=tunables, experiment_id=experiment_id, trial_id=trial_id, config_id=config_id, opt_target=opt_target, + opt_direction=opt_direction, config=config, ) self._engine = engine diff --git a/mlos_bench/mlos_bench/tests/storage/conftest.py b/mlos_bench/mlos_bench/tests/storage/conftest.py index d3de677142a..613730cbed7 100644 --- a/mlos_bench/mlos_bench/tests/storage/conftest.py +++ b/mlos_bench/mlos_bench/tests/storage/conftest.py @@ -12,13 +12,15 @@ from mlos_bench.storage.base_storage import Storage from mlos_bench.storage.sql.storage import SqlStorage +# pylint: disable=redefined-outer-name + @pytest.fixture -def exp_storage_memory_sql(tunable_groups: TunableGroups) -> Storage.Experiment: +def storage_memory_sql(tunable_groups: TunableGroups) -> SqlStorage: """ Test fixture for in-memory SQLite3 storage. """ - storage = SqlStorage( + return SqlStorage( tunables=tunable_groups, service=None, config={ @@ -26,11 +28,19 @@ def exp_storage_memory_sql(tunable_groups: TunableGroups) -> Storage.Experiment: "database": ":memory:", } ) + + +@pytest.fixture +def exp_storage_memory_sql(storage_memory_sql: Storage) -> SqlStorage.Experiment: + """ + Test fixture for Experiment using in-memory SQLite3 storage. + """ # pylint: disable=unnecessary-dunder-call - return storage.experiment( + return storage_memory_sql.experiment( experiment_id="Test-001", trial_id=1, root_env_config="environment.jsonc", description="pytest experiment", opt_target="score", + opt_direction="min", ).__enter__() diff --git a/mlos_bench/mlos_bench/tests/storage/exp_data_test.py b/mlos_bench/mlos_bench/tests/storage/exp_data_test.py new file mode 100644 index 00000000000..fc459172b86 --- /dev/null +++ b/mlos_bench/mlos_bench/tests/storage/exp_data_test.py @@ -0,0 +1,58 @@ +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +""" +Unit tests for loading the experiment metadata. +""" + +from mlos_bench.storage.base_storage import Storage +from mlos_bench.tunables.tunable_groups import TunableGroups + + +def test_load_empty_exp_data(storage_memory_sql: Storage, exp_storage_memory_sql: Storage.Experiment) -> None: + """ + Try to retrieve old experimental data from the empty storage. + """ + exp = storage_memory_sql.experiments[exp_storage_memory_sql.experiment_id] + assert exp.exp_id == exp_storage_memory_sql.experiment_id + assert exp.description == exp_storage_memory_sql.description + # Only support single objective for now. + assert exp.objectives == {exp_storage_memory_sql.opt_target: exp_storage_memory_sql.opt_direction} + + +def test_exp_trial_data_objectives(storage_memory_sql: Storage, + exp_storage_memory_sql: Storage.Experiment, + tunable_groups: TunableGroups) -> None: + """ + Start a new trial and check the storage for the trial data. + """ + + trial_opt_new = exp_storage_memory_sql.new_trial(tunable_groups, config={ + "opt_target": "some-other-target", + "opt_direction": "max", + }) + assert trial_opt_new.config() == { + "experiment_id": exp_storage_memory_sql.experiment_id, + "trial_id": trial_opt_new.trial_id, + "opt_target": "some-other-target", + "opt_direction": "max", + } + + trial_opt_old = exp_storage_memory_sql.new_trial(tunable_groups, config={ + "opt_target": "back-compat", + # "opt_direction": "max", # missing + }) + assert trial_opt_old.config() == { + "experiment_id": exp_storage_memory_sql.experiment_id, + "trial_id": trial_opt_old.trial_id, + "opt_target": "back-compat", + } + + exp = storage_memory_sql.experiments[exp_storage_memory_sql.experiment_id] + # objectives should be the combination of both the trial objectives and the experiment objectives + assert exp.objectives == { + "back-compat": None, + "some-other-target": "max", + exp_storage_memory_sql.opt_target: exp_storage_memory_sql.opt_direction, + } diff --git a/mlos_bench/mlos_bench/tests/storage/trial_data_test.py b/mlos_bench/mlos_bench/tests/storage/trial_data_test.py new file mode 100644 index 00000000000..7775e732496 --- /dev/null +++ b/mlos_bench/mlos_bench/tests/storage/trial_data_test.py @@ -0,0 +1,9 @@ +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# +""" +Unit tests for loading the trial metadata. +""" + +# TODO