Coverage for mlos_bench/mlos_bench/tests/storage/sql/fixtures.py: 100%
77 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-10-07 01:52 +0000
1#
2# Copyright (c) Microsoft Corporation.
3# Licensed under the MIT License.
4#
5"""Test fixtures for mlos_bench storage."""
7from datetime import datetime
8from random import random
9from random import seed as rand_seed
10from typing import Generator, Optional
12import pytest
13from pytz import UTC
15from mlos_bench.environments.status import Status
16from mlos_bench.optimizers.mock_optimizer import MockOptimizer
17from mlos_bench.storage.base_experiment_data import ExperimentData
18from mlos_bench.storage.sql.storage import SqlStorage
19from mlos_bench.tests import SEED
20from mlos_bench.tests.storage import CONFIG_COUNT, CONFIG_TRIAL_REPEAT_COUNT
21from mlos_bench.tunables.tunable_groups import TunableGroups
23# pylint: disable=redefined-outer-name
26@pytest.fixture
27def storage() -> SqlStorage:
28 """Test fixture for in-memory SQLite3 storage."""
29 return SqlStorage(
30 service=None,
31 config={
32 "drivername": "sqlite",
33 "database": ":memory:",
34 # "database": "mlos_bench.pytest.db",
35 },
36 )
39@pytest.fixture
40def exp_storage(
41 storage: SqlStorage,
42 tunable_groups: TunableGroups,
43) -> Generator[SqlStorage.Experiment, None, None]:
44 """
45 Test fixture for Experiment using in-memory SQLite3 storage.
47 Note: It has already entered the context upon return.
48 """
49 with storage.experiment(
50 experiment_id="Test-001",
51 trial_id=1,
52 root_env_config="environment.jsonc",
53 description="pytest experiment",
54 tunables=tunable_groups,
55 opt_targets={"score": "min"},
56 ) as exp:
57 yield exp
58 # pylint: disable=protected-access
59 assert not exp._in_context
62@pytest.fixture
63def exp_no_tunables_storage(
64 storage: SqlStorage,
65) -> Generator[SqlStorage.Experiment, None, None]:
66 """
67 Test fixture for Experiment using in-memory SQLite3 storage.
69 Note: It has already entered the context upon return.
70 """
71 empty_config: dict = {}
72 with storage.experiment(
73 experiment_id="Test-003",
74 trial_id=1,
75 root_env_config="environment.jsonc",
76 description="pytest experiment - no tunables",
77 tunables=TunableGroups(empty_config),
78 opt_targets={"score": "min"},
79 ) as exp:
80 yield exp
81 # pylint: disable=protected-access
82 assert not exp._in_context
85@pytest.fixture
86def mixed_numerics_exp_storage(
87 storage: SqlStorage,
88 mixed_numerics_tunable_groups: TunableGroups,
89) -> Generator[SqlStorage.Experiment, None, None]:
90 """
91 Test fixture for an Experiment with mixed numerics tunables using in-memory SQLite3
92 storage.
94 Note: It has already entered the context upon return.
95 """
96 with storage.experiment(
97 experiment_id="Test-002",
98 trial_id=1,
99 root_env_config="dne.jsonc",
100 description="pytest experiment",
101 tunables=mixed_numerics_tunable_groups,
102 opt_targets={"score": "min"},
103 ) as exp:
104 yield exp
105 # pylint: disable=protected-access
106 assert not exp._in_context
109def _dummy_run_exp(
110 exp: SqlStorage.Experiment,
111 tunable_name: Optional[str],
112) -> SqlStorage.Experiment:
113 """Generates data by doing a simulated run of the given experiment."""
114 # Add some trials to that experiment.
115 # Note: we're just fabricating some made up function for the ML libraries to try and learn.
116 base_score = 10.0
117 if tunable_name:
118 tunable = exp.tunables.get_tunable(tunable_name)[0]
119 assert isinstance(tunable.default, int)
120 (tunable_min, tunable_max) = tunable.range
121 tunable_range = tunable_max - tunable_min
122 rand_seed(SEED)
123 opt = MockOptimizer(
124 tunables=exp.tunables,
125 config={
126 "seed": SEED,
127 # This should be the default, so we leave it omitted for now to test the default.
128 # But the test logic relies on this (e.g., trial 1 is config 1 is the
129 # default values for the tunable params)
130 # "start_with_defaults": True,
131 },
132 )
133 assert opt.start_with_defaults
134 for config_i in range(CONFIG_COUNT):
135 tunables = opt.suggest()
136 for repeat_j in range(CONFIG_TRIAL_REPEAT_COUNT):
137 trial = exp.new_trial(
138 tunables=tunables.copy(),
139 config={
140 "trial_number": config_i * CONFIG_TRIAL_REPEAT_COUNT + repeat_j + 1,
141 **{
142 f"opt_{key}_{i}": val
143 for (i, opt_target) in enumerate(exp.opt_targets.items())
144 for (key, val) in zip(["target", "direction"], opt_target)
145 },
146 },
147 )
148 if exp.tunables:
149 assert trial.tunable_config_id == config_i + 1
150 else:
151 assert trial.tunable_config_id == 1
152 if tunable_name:
153 tunable_value = float(tunables.get_tunable(tunable_name)[0].numerical_value)
154 tunable_value_norm = base_score * (tunable_value - tunable_min) / tunable_range
155 else:
156 tunable_value_norm = 0
157 timestamp = datetime.now(UTC)
158 trial.update_telemetry(
159 status=Status.RUNNING,
160 timestamp=timestamp,
161 metrics=[
162 (timestamp, "some-metric", tunable_value_norm + random() / 100),
163 ],
164 )
165 trial.update(
166 Status.SUCCEEDED,
167 timestamp,
168 metrics={
169 # Give some variance on the score.
170 # And some influence from the tunable value.
171 "score": tunable_value_norm
172 + random() / 100
173 },
174 )
175 return exp
178@pytest.fixture
179def exp_storage_with_trials(exp_storage: SqlStorage.Experiment) -> SqlStorage.Experiment:
180 """Test fixture for Experiment using in-memory SQLite3 storage."""
181 return _dummy_run_exp(exp_storage, tunable_name="kernel_sched_latency_ns")
184@pytest.fixture
185def exp_no_tunables_storage_with_trials(
186 exp_no_tunables_storage: SqlStorage.Experiment,
187) -> SqlStorage.Experiment:
188 """Test fixture for Experiment using in-memory SQLite3 storage."""
189 assert not exp_no_tunables_storage.tunables
190 return _dummy_run_exp(exp_no_tunables_storage, tunable_name=None)
193@pytest.fixture
194def mixed_numerics_exp_storage_with_trials(
195 mixed_numerics_exp_storage: SqlStorage.Experiment,
196) -> SqlStorage.Experiment:
197 """Test fixture for Experiment using in-memory SQLite3 storage."""
198 tunable = next(iter(mixed_numerics_exp_storage.tunables))[0]
199 return _dummy_run_exp(mixed_numerics_exp_storage, tunable_name=tunable.name)
202@pytest.fixture
203def exp_data(
204 storage: SqlStorage,
205 exp_storage_with_trials: SqlStorage.Experiment,
206) -> ExperimentData:
207 """Test fixture for ExperimentData."""
208 return storage.experiments[exp_storage_with_trials.experiment_id]
211@pytest.fixture
212def exp_no_tunables_data(
213 storage: SqlStorage,
214 exp_no_tunables_storage_with_trials: SqlStorage.Experiment,
215) -> ExperimentData:
216 """Test fixture for ExperimentData with no tunable configs."""
217 return storage.experiments[exp_no_tunables_storage_with_trials.experiment_id]
220@pytest.fixture
221def mixed_numerics_exp_data(
222 storage: SqlStorage,
223 mixed_numerics_exp_storage_with_trials: SqlStorage.Experiment,
224) -> ExperimentData:
225 """Test fixture for ExperimentData with mixed numerical tunable types."""
226 return storage.experiments[mixed_numerics_exp_storage_with_trials.experiment_id]