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