wehub-resource-sync
727 lines
28 KiB
Python
727 lines
28 KiB
Python
import gc
|
|
import random
|
|
import threading
|
|
import time
|
|
from datetime import datetime
|
|
from time import sleep
|
|
from typing import Any
|
|
from unittest.mock import MagicMock, call, patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from optuna.distributions import CategoricalDistribution, FloatDistribution
|
|
from optuna.storages import BaseStorage
|
|
from optuna.storages._base import DEFAULT_STUDY_NAME_PREFIX
|
|
from optuna.study._frozen import FrozenStudy
|
|
from optuna.study._study_direction import StudyDirection
|
|
from optuna.trial import FrozenTrial, TrialState
|
|
|
|
import mlflow
|
|
from mlflow.entities import Metric, Param, RunTag
|
|
from mlflow.optuna.storage import MlflowStorage
|
|
|
|
ALL_STATES = list(TrialState)
|
|
|
|
EXAMPLE_ATTRS: dict[str, Any] = {
|
|
"dataset": "MNIST",
|
|
"none": None,
|
|
"json_serializable": {"baseline_score": 0.001, "tags": ["image", "classification"]},
|
|
}
|
|
|
|
|
|
def _setup_studies(
|
|
storage: BaseStorage,
|
|
n_study: int,
|
|
n_trial: int,
|
|
seed: int,
|
|
direction: StudyDirection = None,
|
|
) -> tuple[dict[int, FrozenStudy], dict[int, dict[int, FrozenTrial]]]:
|
|
generator = random.Random(seed)
|
|
study_id_to_frozen_study: dict[int, FrozenStudy] = {}
|
|
study_id_to_trials: dict[int, dict[int, FrozenTrial]] = {}
|
|
for i in range(n_study):
|
|
study_name = "test-study-name-{}".format(i)
|
|
if direction is None:
|
|
direction = generator.choice([StudyDirection.MINIMIZE, StudyDirection.MAXIMIZE])
|
|
study_id = storage.create_new_study(directions=(direction,), study_name=study_name)
|
|
storage.set_study_user_attr(study_id, "u", i)
|
|
storage.set_study_system_attr(study_id, "s", i)
|
|
trials = {}
|
|
for j in range(n_trial):
|
|
trial = _generate_trial(generator)
|
|
trial.number = j
|
|
trial._trial_id = storage.create_new_trial(study_id, trial)
|
|
trials[trial._trial_id] = trial
|
|
study_id_to_trials[study_id] = trials
|
|
study_id_to_frozen_study[study_id] = FrozenStudy(
|
|
study_name=study_name,
|
|
direction=direction,
|
|
user_attrs={"u": i},
|
|
system_attrs={"s": i},
|
|
study_id=study_id,
|
|
)
|
|
return study_id_to_frozen_study, study_id_to_trials
|
|
|
|
|
|
def _generate_trial(generator: random.Random) -> FrozenTrial:
|
|
example_params = {
|
|
"paramA": (generator.uniform(0, 1), FloatDistribution(0, 1)),
|
|
"paramB": (generator.uniform(1, 2), FloatDistribution(1, 2, log=True)),
|
|
"paramC": (
|
|
generator.choice(["CatA", "CatB", "CatC"]),
|
|
CategoricalDistribution(("CatA", "CatB", "CatC")),
|
|
),
|
|
"paramD": (generator.uniform(-3, 0), FloatDistribution(-3, 0)),
|
|
"paramE": (generator.choice([0.1, 0.2]), CategoricalDistribution((0.1, 0.2))),
|
|
}
|
|
example_attrs = {
|
|
"attrA": "valueA",
|
|
"attrB": 1,
|
|
"attrC": None,
|
|
"attrD": {"baseline_score": 0.001, "tags": ["image", "classification"]},
|
|
}
|
|
state = generator.choice(ALL_STATES)
|
|
params = {}
|
|
distributions = {}
|
|
user_attrs = {}
|
|
system_attrs: dict[str, Any] = {}
|
|
intermediate_values = {}
|
|
for key, (value, dist) in example_params.items():
|
|
if generator.choice([True, False]):
|
|
params[key] = value
|
|
distributions[key] = dist
|
|
for key, value in example_attrs.items():
|
|
if generator.choice([True, False]):
|
|
user_attrs["usr_" + key] = value
|
|
if generator.choice([True, False]):
|
|
system_attrs["sys_" + key] = value
|
|
for i in range(generator.randint(4, 10)):
|
|
if generator.choice([True, False]):
|
|
intermediate_values[i] = generator.uniform(-10, 10)
|
|
return FrozenTrial(
|
|
number=0, # dummy
|
|
state=state,
|
|
value=generator.uniform(-10, 10),
|
|
datetime_start=datetime.now(),
|
|
datetime_complete=datetime.now() if state.is_finished() else None,
|
|
params=params,
|
|
distributions=distributions,
|
|
user_attrs=user_attrs,
|
|
system_attrs=system_attrs,
|
|
intermediate_values=intermediate_values,
|
|
trial_id=0, # dummy
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def setup_storage(db_uri):
|
|
mlflow.set_tracking_uri(db_uri)
|
|
experiment_id = mlflow.create_experiment(name="optuna_mlflow_test")
|
|
storage = MlflowStorage(
|
|
experiment_id=experiment_id, batch_flush_interval=1.0, batch_size_threshold=5
|
|
)
|
|
storage._flush_thread = MagicMock()
|
|
yield storage
|
|
mlflow.delete_experiment(experiment_id)
|
|
|
|
|
|
def test_queue_batch_operation_creates_new_queue_for_new_run(setup_storage):
|
|
storage = setup_storage
|
|
run_id = "test-run-id"
|
|
test_metric = Metric("test_metric", 1.0, int(time.time() * 1000), 0)
|
|
|
|
# Call the method with a new run_id
|
|
storage._queue_batch_operation(run_id, metrics=[test_metric])
|
|
|
|
# Check that a new queue was created for this run_id
|
|
assert run_id in storage._batch_queue
|
|
assert len(storage._batch_queue[run_id]["metrics"]) == 1
|
|
assert storage._batch_queue[run_id]["metrics"][0] == test_metric
|
|
assert len(storage._batch_queue[run_id]["params"]) == 0
|
|
assert len(storage._batch_queue[run_id]["tags"]) == 0
|
|
|
|
|
|
def test_queue_batch_operation_appends_to_existing_queue(setup_storage):
|
|
storage = setup_storage
|
|
run_id = "test-run-id"
|
|
|
|
# Setup existing queue
|
|
storage._batch_queue[run_id] = {"metrics": [], "params": [], "tags": []}
|
|
|
|
# Add metrics, params, and tags
|
|
test_metric = Metric("test_metric", 1.0, int(time.time() * 1000), 0)
|
|
test_param = Param("test_param", "value")
|
|
test_tag = RunTag("test_tag", "value")
|
|
|
|
# Queue each type separately
|
|
storage._queue_batch_operation(run_id, metrics=[test_metric])
|
|
storage._queue_batch_operation(run_id, params=[test_param])
|
|
storage._queue_batch_operation(run_id, tags=[test_tag])
|
|
|
|
# Check all were added correctly
|
|
assert len(storage._batch_queue[run_id]["metrics"]) == 1
|
|
assert len(storage._batch_queue[run_id]["params"]) == 1
|
|
assert len(storage._batch_queue[run_id]["tags"]) == 1
|
|
|
|
|
|
def test_flush_batch_sends_data_to_mlflow(setup_storage):
|
|
with patch("mlflow.optuna.storage.MlflowClient") as mock_client:
|
|
storage = setup_storage
|
|
run_id = "test-run-id"
|
|
mock_client_instance = MagicMock()
|
|
mock_client.return_value = mock_client_instance
|
|
|
|
storage._mlflow_client = mock_client
|
|
|
|
# Setup batch data
|
|
metrics = [Metric("test_metric", 1.0, int(time.time() * 1000), 0)]
|
|
params = [Param("test_param", "value")]
|
|
tags = [RunTag("test_tag", "value")]
|
|
|
|
storage._batch_queue[run_id] = {"metrics": metrics[:], "params": params[:], "tags": tags[:]}
|
|
|
|
# Call _flush_batch
|
|
storage._flush_batch(run_id)
|
|
|
|
# Verify MLflow client was called with the correct data
|
|
storage._mlflow_client.log_batch.assert_called_once_with(
|
|
run_id, metrics=metrics, params=params, tags=tags
|
|
)
|
|
|
|
# Verify batch was cleared
|
|
assert len(storage._batch_queue[run_id]["metrics"]) == 0
|
|
assert len(storage._batch_queue[run_id]["params"]) == 0
|
|
assert len(storage._batch_queue[run_id]["tags"]) == 0
|
|
|
|
|
|
def test_flush_batch_does_nothing_for_empty_batch(setup_storage):
|
|
with patch("mlflow.optuna.storage.MlflowClient") as mock_client:
|
|
storage = setup_storage
|
|
run_id = "test-run-id"
|
|
# Create a mock instance that will be returned when MlflowClient is instantiated
|
|
mock_client_instance = MagicMock()
|
|
mock_client.return_value = mock_client_instance
|
|
|
|
storage._mlflow_client = mock_client
|
|
|
|
# Setup empty batch
|
|
storage._batch_queue[run_id] = {"metrics": [], "params": [], "tags": []}
|
|
|
|
# Call _flush_batch
|
|
storage._flush_batch(run_id)
|
|
|
|
# Verify MLflow client was not called
|
|
storage._mlflow_client.log_batch.assert_not_called()
|
|
|
|
|
|
def test_flush_batch_handles_nonexistent_run(setup_storage):
|
|
with patch("mlflow.optuna.storage.MlflowClient") as mock_client:
|
|
storage = setup_storage
|
|
run_id = "nonexistent-run"
|
|
# Create a mock instance that will be returned when MlflowClient is instantiated
|
|
mock_client_instance = MagicMock()
|
|
mock_client.return_value = mock_client_instance
|
|
|
|
storage._mlflow_client = mock_client
|
|
|
|
# Call _flush_batch for a run that doesn't exist in the queue
|
|
storage._flush_batch(run_id)
|
|
|
|
# Verify MLflow client was not called
|
|
storage._mlflow_client.log_batch.assert_not_called()
|
|
|
|
|
|
def test_flush_all_batches_flushes_all_runs(setup_storage):
|
|
storage = setup_storage
|
|
# Setup multiple runs with data
|
|
run_ids = ["run1", "run2", "run3"]
|
|
|
|
for run_id in run_ids:
|
|
storage._batch_queue[run_id] = {
|
|
"metrics": [Metric(f"m_{run_id}", 1.0, int(time.time() * 1000), 0)],
|
|
"params": [Param(f"p_{run_id}", "value")],
|
|
"tags": [RunTag(f"t_{run_id}", "value")],
|
|
}
|
|
|
|
# Create a spy on _flush_batch
|
|
with patch.object(storage, "_flush_batch") as mock_flush:
|
|
# Call flush_all_batches
|
|
storage.flush_all_batches()
|
|
|
|
# Check that _flush_batch was called for each run
|
|
expected_calls = [call(run_id) for run_id in run_ids]
|
|
mock_flush.assert_has_calls(expected_calls, any_order=True)
|
|
assert mock_flush.call_count == len(run_ids)
|
|
|
|
|
|
def test_flush_all_batches_handles_empty_queue(setup_storage):
|
|
storage = setup_storage
|
|
# Ensure batch queue is empty
|
|
storage._batch_queue = {}
|
|
|
|
# Create a spy on _flush_batch
|
|
with patch.object(storage, "_flush_batch") as mock_flush:
|
|
# Call flush_all_batches
|
|
storage.flush_all_batches()
|
|
|
|
# Check that _flush_batch was not called
|
|
mock_flush.assert_not_called()
|
|
|
|
|
|
def test_create_new_study(setup_storage):
|
|
storage = setup_storage
|
|
study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
frozen_studies = storage.get_all_studies()
|
|
assert len(frozen_studies) == 1
|
|
assert frozen_studies[0]._study_id == study_id
|
|
assert frozen_studies[0].study_name.startswith(DEFAULT_STUDY_NAME_PREFIX)
|
|
|
|
study_id2 = storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
# Study id must be unique.
|
|
assert study_id != study_id2
|
|
frozen_studies = storage.get_all_studies()
|
|
assert len(frozen_studies) == 2
|
|
assert {s._study_id for s in frozen_studies} == {study_id, study_id2}
|
|
assert all(s.study_name.startswith(DEFAULT_STUDY_NAME_PREFIX) for s in frozen_studies)
|
|
|
|
|
|
def test_delete_study(setup_storage):
|
|
storage = setup_storage
|
|
study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
storage.create_new_trial(study_id)
|
|
trials = storage.get_all_trials(study_id)
|
|
assert len(trials) == 1
|
|
|
|
with pytest.raises(mlflow.exceptions.MlflowException, match="Run .* not found"):
|
|
# Deletion of non-existent study.
|
|
storage.delete_study(study_id + "1")
|
|
|
|
storage.delete_study(study_id)
|
|
|
|
|
|
def test_get_study_id_from_name_and_get_study_name_from_id(setup_storage):
|
|
storage = setup_storage
|
|
study_name = "test_optuna_mlflow_study"
|
|
study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE], study_name=study_name)
|
|
|
|
# Test existing study.
|
|
assert storage.get_study_name_from_id(study_id) == study_name
|
|
assert storage.get_study_id_from_name(study_name) == study_id
|
|
|
|
# Test not existing study.
|
|
with pytest.raises(Exception, match="Study dummy-name not found"):
|
|
storage.get_study_id_from_name("dummy-name")
|
|
|
|
with pytest.raises(mlflow.exceptions.MlflowException, match="Run .* not found"):
|
|
storage.get_study_name_from_id(study_id + "1")
|
|
|
|
|
|
def test_set_and_get_study_directions(setup_storage):
|
|
storage = setup_storage
|
|
for target in [
|
|
(StudyDirection.MINIMIZE,),
|
|
(StudyDirection.MAXIMIZE,),
|
|
(StudyDirection.MINIMIZE, StudyDirection.MAXIMIZE),
|
|
(StudyDirection.MAXIMIZE, StudyDirection.MINIMIZE),
|
|
[StudyDirection.MINIMIZE, StudyDirection.MAXIMIZE],
|
|
[StudyDirection.MAXIMIZE, StudyDirection.MINIMIZE],
|
|
]:
|
|
study_id = storage.create_new_study(directions=target)
|
|
|
|
def check_get() -> None:
|
|
got_directions = storage.get_study_directions(study_id)
|
|
|
|
assert got_directions == list(target), (
|
|
"Direction of a study should be a tuple of `StudyDirection` objects."
|
|
)
|
|
|
|
# Test setting value.
|
|
check_get()
|
|
|
|
# Test non-existent study.
|
|
non_existent_study_id = study_id + "1"
|
|
|
|
with pytest.raises(mlflow.exceptions.MlflowException, match="Run .* not found"):
|
|
storage.get_study_directions(non_existent_study_id)
|
|
|
|
|
|
def test_set_and_get_study_user_attrs(setup_storage):
|
|
storage = setup_storage
|
|
study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
|
|
def check_set_and_get(key: str, value: Any) -> None:
|
|
storage.set_study_user_attr(study_id, key, value)
|
|
assert storage.get_study_user_attrs(study_id)[key] == value
|
|
|
|
# Test setting value.
|
|
for key, value in EXAMPLE_ATTRS.items():
|
|
check_set_and_get(key, value)
|
|
assert storage.get_study_user_attrs(study_id) == EXAMPLE_ATTRS
|
|
|
|
# Test overwriting value.
|
|
check_set_and_get("dataset", "ImageNet")
|
|
|
|
# Non-existent study id.
|
|
non_existent_study_id = study_id + "1"
|
|
with pytest.raises(mlflow.exceptions.MlflowException, match="Run .* not found"):
|
|
storage.get_study_user_attrs(non_existent_study_id)
|
|
|
|
# Non-existent study id.
|
|
with pytest.raises(mlflow.exceptions.MlflowException, match="Run .* not found"):
|
|
storage.set_study_user_attr(non_existent_study_id, "key", "value")
|
|
|
|
|
|
def test_set_and_get_study_system_attrs(setup_storage):
|
|
storage = setup_storage
|
|
study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
|
|
def check_set_and_get(key: str, value: Any) -> None:
|
|
storage.set_study_system_attr(study_id, key, value)
|
|
assert storage.get_study_system_attrs(study_id)[key] == value
|
|
|
|
# Test setting value.
|
|
for key, value in EXAMPLE_ATTRS.items():
|
|
check_set_and_get(key, value)
|
|
assert storage.get_study_system_attrs(study_id) == EXAMPLE_ATTRS
|
|
|
|
# Test overwriting value.
|
|
check_set_and_get("dataset", "ImageNet")
|
|
|
|
|
|
def test_create_new_trial(setup_storage):
|
|
storage = setup_storage
|
|
|
|
def _check_trials(
|
|
trials: list[FrozenTrial],
|
|
idx: int,
|
|
trial_id: int,
|
|
time_before_creation: datetime,
|
|
time_after_creation: datetime,
|
|
) -> None:
|
|
assert len(trials) == idx + 1
|
|
assert len({t._trial_id for t in trials}) == idx + 1
|
|
assert trial_id in {t._trial_id for t in trials}
|
|
assert {t.number for t in trials} == set(range(idx + 1))
|
|
assert all(t.state == TrialState.RUNNING for t in trials)
|
|
assert all(t.params == {} for t in trials)
|
|
assert all(t.intermediate_values == {} for t in trials)
|
|
assert all(t.user_attrs == {} for t in trials)
|
|
assert all(t.system_attrs == {} for t in trials)
|
|
assert all(
|
|
t.datetime_start < time_before_creation
|
|
for t in trials
|
|
if t._trial_id != trial_id and t.datetime_start is not None
|
|
)
|
|
assert all(
|
|
time_before_creation < t.datetime_start < time_after_creation
|
|
for t in trials
|
|
if t._trial_id == trial_id and t.datetime_start is not None
|
|
)
|
|
assert all(t.value is None for t in trials)
|
|
|
|
study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
n_trial_in_study = 3
|
|
for i in range(n_trial_in_study):
|
|
time_before_creation = datetime.now()
|
|
sleep(0.001) # Sleep 1ms to avoid faulty assertion on Windows OS.
|
|
trial_id = storage.create_new_trial(study_id)
|
|
sleep(0.001)
|
|
time_after_creation = datetime.now()
|
|
|
|
trials = storage.get_all_trials(study_id)
|
|
_check_trials(trials, i, trial_id, time_before_creation, time_after_creation)
|
|
|
|
|
|
def test_create_new_trial_with_template_trial(setup_storage):
|
|
storage = setup_storage
|
|
start_time = datetime.now()
|
|
end_time = datetime.now()
|
|
template_trial = FrozenTrial(
|
|
state=TrialState.COMPLETE,
|
|
value=10000,
|
|
datetime_start=start_time,
|
|
datetime_complete=end_time,
|
|
params={"x": 0.5},
|
|
distributions={"x": FloatDistribution(0, 1)},
|
|
user_attrs={"foo": "bar"},
|
|
system_attrs={"baz": 123},
|
|
intermediate_values={1: 10, 2: 100, 3: 1000},
|
|
number=55, # This entry is ignored.
|
|
trial_id=-1, # dummy value (unused).
|
|
)
|
|
|
|
def _check_trials(trials: list[FrozenTrial], idx: int, trial_id: int) -> None:
|
|
assert len(trials) == idx + 1
|
|
assert len({t._trial_id for t in trials}) == idx + 1
|
|
assert trial_id in {t._trial_id for t in trials}
|
|
assert {t.number for t in trials} == set(range(idx + 1))
|
|
assert all(t.state == template_trial.state for t in trials)
|
|
|
|
assert all(t.params == template_trial.params for t in trials)
|
|
assert all(t.distributions == template_trial.distributions for t in trials)
|
|
assert all(t.intermediate_values == template_trial.intermediate_values for t in trials)
|
|
assert all(t.user_attrs == template_trial.user_attrs for t in trials)
|
|
assert all(t.system_attrs == template_trial.system_attrs for t in trials)
|
|
# assert all(t.datetime_start == template_trial.datetime_start for t in trials)
|
|
# assert all(t.datetime_complete == template_trial.datetime_complete for t in trials)
|
|
assert all(t.value == template_trial.value for t in trials)
|
|
|
|
study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
n_trial_in_study = 3
|
|
for i in range(n_trial_in_study):
|
|
trial_id = storage.create_new_trial(study_id, template_trial=template_trial)
|
|
trials = storage.get_all_trials(study_id)
|
|
_check_trials(trials, i, trial_id)
|
|
|
|
|
|
def test_get_trial_number_from_id(setup_storage):
|
|
storage = setup_storage
|
|
study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
trial_id = storage.create_new_trial(study_id)
|
|
assert storage.get_trial_number_from_id(trial_id) == 0
|
|
|
|
trial_id = storage.create_new_trial(study_id)
|
|
assert storage.get_trial_number_from_id(trial_id) == 1
|
|
|
|
|
|
def test_set_trial_state_values_for_state(setup_storage):
|
|
storage = setup_storage
|
|
study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
trial_ids = [storage.create_new_trial(study_id) for _ in ALL_STATES]
|
|
|
|
for trial_id, state in zip(trial_ids, ALL_STATES):
|
|
if state == TrialState.WAITING:
|
|
continue
|
|
assert storage.get_trial(trial_id).state == TrialState.RUNNING
|
|
datetime_start_prev = storage.get_trial(trial_id).datetime_start
|
|
storage.set_trial_state_values(
|
|
trial_id, state=state, values=(0.0,) if state.is_finished() else None
|
|
)
|
|
assert storage.get_trial(trial_id).state == state
|
|
# Repeated state changes to RUNNING should not trigger further datetime_start changes.
|
|
if state == TrialState.RUNNING:
|
|
assert storage.get_trial(trial_id).datetime_start == datetime_start_prev
|
|
if state.is_finished():
|
|
assert storage.get_trial(trial_id).datetime_complete is not None
|
|
else:
|
|
assert storage.get_trial(trial_id).datetime_complete is None
|
|
|
|
|
|
def test_get_trial_param_and_get_trial_params(setup_storage):
|
|
storage = setup_storage
|
|
_, study_to_trials = _setup_studies(storage, n_study=2, n_trial=5, seed=1)
|
|
|
|
for _, trial_id_to_trial in study_to_trials.items():
|
|
for trial_id, expected_trial in trial_id_to_trial.items():
|
|
assert storage.get_trial_params(trial_id) == expected_trial.params
|
|
for key in expected_trial.params.keys():
|
|
assert storage.get_trial_param(trial_id, key) == expected_trial.distributions[
|
|
key
|
|
].to_internal_repr(expected_trial.params[key])
|
|
|
|
|
|
def test_set_trial_param(setup_storage):
|
|
storage = setup_storage
|
|
study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
trial_id_1 = storage.create_new_trial(study_id)
|
|
trial_id_2 = storage.create_new_trial(study_id)
|
|
|
|
# Setup distributions.
|
|
distribution_x = FloatDistribution(low=1.0, high=2.0)
|
|
distribution_y_1 = CategoricalDistribution(choices=("Shibuya", "Ebisu", "Meguro"))
|
|
distribution_z = FloatDistribution(low=1.0, high=100.0, log=True)
|
|
|
|
# Set new params.
|
|
storage.set_trial_param(trial_id_1, "x", 0.5, distribution_x)
|
|
storage.set_trial_param(trial_id_1, "y", 2, distribution_y_1)
|
|
assert storage.get_trial_param(trial_id_1, "x") == 0.5
|
|
assert storage.get_trial_param(trial_id_1, "y") == 2
|
|
# Check set_param breaks neither get_trial nor get_trial_params.
|
|
assert storage.get_trial(trial_id_1).params == {"x": 0.5, "y": "Meguro"}
|
|
assert storage.get_trial_params(trial_id_1) == {"x": 0.5, "y": "Meguro"}
|
|
|
|
# Set params to another trial.
|
|
storage.set_trial_param(trial_id_2, "x", 0.3, distribution_x)
|
|
storage.set_trial_param(trial_id_2, "z", 0.1, distribution_z)
|
|
assert storage.get_trial_param(trial_id_2, "x") == 0.3
|
|
assert storage.get_trial_param(trial_id_2, "z") == 0.1
|
|
assert storage.get_trial(trial_id_2).params == {"x": 0.3, "z": 0.1}
|
|
assert storage.get_trial_params(trial_id_2) == {"x": 0.3, "z": 0.1}
|
|
|
|
storage.set_trial_state_values(trial_id_2, state=TrialState.COMPLETE)
|
|
# Cannot assign params to finished trial.
|
|
with pytest.raises(RuntimeError, match="Trial#1 has already finished and can not be updated."):
|
|
storage.set_trial_param(trial_id_2, "y", 2, distribution_y_1)
|
|
# Check the previous call does not change the params.
|
|
with pytest.raises(KeyError, match="'param_internal_val_y'"):
|
|
storage.get_trial_param(trial_id_2, "y")
|
|
# State should be checked prior to distribution compatibility.
|
|
with pytest.raises(RuntimeError, match="Trial#1 has already finished and can not be updated."):
|
|
storage.set_trial_param(trial_id_2, "y", 0.4, distribution_z)
|
|
|
|
|
|
def test_set_trial_state_values_for_values(setup_storage):
|
|
storage = setup_storage
|
|
study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
trial_id_1 = storage.create_new_trial(study_id)
|
|
trial_id_2 = storage.create_new_trial(study_id)
|
|
trial_id_3 = storage.create_new_trial(
|
|
storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
)
|
|
trial_id_4 = storage.create_new_trial(study_id)
|
|
trial_id_5 = storage.create_new_trial(study_id)
|
|
|
|
# Test setting new value.
|
|
storage.set_trial_state_values(trial_id_1, state=TrialState.COMPLETE, values=(0.5,))
|
|
storage.set_trial_state_values(trial_id_3, state=TrialState.COMPLETE, values=(float("inf"),))
|
|
storage.set_trial_state_values(trial_id_4, state=TrialState.WAITING, values=(0.1, 0.2, 0.3))
|
|
storage.set_trial_state_values(trial_id_5, state=TrialState.WAITING, values=[0.1, 0.2, 0.3])
|
|
|
|
assert storage.get_trial(trial_id_1).value == 0.5
|
|
assert storage.get_trial(trial_id_2).value is None
|
|
# SQLite/SQLAlchemy sanitizes +inf to max float value
|
|
assert storage.get_trial(trial_id_3).value == 1.7976931348623157e308
|
|
assert storage.get_trial(trial_id_4).values == [0.1, 0.2, 0.3]
|
|
assert storage.get_trial(trial_id_5).values == [0.1, 0.2, 0.3]
|
|
|
|
|
|
def test_set_trial_intermediate_value(setup_storage):
|
|
storage = setup_storage
|
|
study_id = storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
trial_id_1 = storage.create_new_trial(study_id)
|
|
trial_id_2 = storage.create_new_trial(study_id)
|
|
trial_id_3 = storage.create_new_trial(
|
|
storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
)
|
|
trial_id_4 = storage.create_new_trial(study_id)
|
|
|
|
# Test setting new values.
|
|
storage.set_trial_intermediate_value(trial_id_1, 0, 0.3)
|
|
storage.set_trial_intermediate_value(trial_id_1, 2, 0.4)
|
|
storage.set_trial_intermediate_value(trial_id_3, 0, 0.1)
|
|
storage.set_trial_intermediate_value(trial_id_3, 1, 0.4)
|
|
storage.set_trial_intermediate_value(trial_id_3, 2, 0.5)
|
|
storage.set_trial_intermediate_value(trial_id_3, 3, float("inf"))
|
|
storage.set_trial_intermediate_value(trial_id_4, 0, float("nan"))
|
|
|
|
assert storage.get_trial(trial_id_1).intermediate_values == {0: 0.3, 2: 0.4}
|
|
assert storage.get_trial(trial_id_2).intermediate_values == {}
|
|
# SQLite/SQLAlchemy sanitizes +inf to max float value
|
|
assert storage.get_trial(trial_id_3).intermediate_values == {
|
|
0: 0.1,
|
|
1: 0.4,
|
|
2: 0.5,
|
|
3: 1.7976931348623157e308,
|
|
}
|
|
assert np.isnan(storage.get_trial(trial_id_4).intermediate_values[0])
|
|
|
|
# Test setting existing step.
|
|
storage.set_trial_intermediate_value(trial_id_1, 0, 0.2)
|
|
assert storage.get_trial(trial_id_1).intermediate_values == {0: 0.2, 2: 0.4}
|
|
|
|
|
|
def test_get_trial_user_attrs(setup_storage):
|
|
storage = setup_storage
|
|
_, study_to_trials = _setup_studies(storage, n_study=2, n_trial=5, seed=10)
|
|
assert all(
|
|
storage.get_trial_user_attrs(trial_id) == trial.user_attrs
|
|
for trials in study_to_trials.values()
|
|
for trial_id, trial in trials.items()
|
|
)
|
|
|
|
|
|
def test_set_trial_user_attr(setup_storage):
|
|
storage = setup_storage
|
|
trial_id_1 = storage.create_new_trial(
|
|
storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
)
|
|
|
|
def check_set_and_get(trial_id: int, key: str, value: Any) -> None:
|
|
storage.set_trial_user_attr(trial_id, key, value)
|
|
assert storage.get_trial(trial_id).user_attrs[key] == value
|
|
|
|
# Test setting value.
|
|
for key, value in EXAMPLE_ATTRS.items():
|
|
check_set_and_get(trial_id_1, key, value)
|
|
assert storage.get_trial(trial_id_1).user_attrs == EXAMPLE_ATTRS
|
|
|
|
# Test overwriting value.
|
|
check_set_and_get(trial_id_1, "dataset", "ImageNet")
|
|
|
|
|
|
def test_set_trial_system_attr(setup_storage):
|
|
storage = setup_storage
|
|
trial_id_1 = storage.create_new_trial(
|
|
storage.create_new_study(directions=[StudyDirection.MINIMIZE])
|
|
)
|
|
|
|
def check_set_and_get(trial_id: int, key: str, value: Any) -> None:
|
|
storage.set_trial_system_attr(trial_id, key, value)
|
|
assert storage.get_trial_system_attrs(trial_id)[key] == value
|
|
|
|
# Test setting value.
|
|
for key, value in EXAMPLE_ATTRS.items():
|
|
check_set_and_get(trial_id_1, key, value)
|
|
system_attrs = storage.get_trial(trial_id_1).system_attrs
|
|
assert system_attrs == EXAMPLE_ATTRS
|
|
|
|
# Test overwriting value.
|
|
check_set_and_get(trial_id_1, "dataset", "ImageNet")
|
|
|
|
|
|
def test_get_n_trials(setup_storage):
|
|
storage = setup_storage
|
|
study_id_to_frozen_studies, _ = _setup_studies(storage, n_study=2, n_trial=7, seed=50)
|
|
for study_id in study_id_to_frozen_studies:
|
|
assert storage.get_n_trials(study_id) == 7
|
|
|
|
|
|
def test_study_exists_method(setup_storage):
|
|
storage = setup_storage
|
|
|
|
# Test non-existent study
|
|
assert not storage.get_study_id_by_name_if_exists("non-existent-study")
|
|
|
|
# Create a study
|
|
storage.create_new_study([StudyDirection.MINIMIZE], "test-study")
|
|
|
|
# Test existing study
|
|
assert storage.get_study_id_by_name_if_exists("test-study")
|
|
|
|
|
|
def test_get_study_id_by_name_if_exists(setup_storage):
|
|
storage = setup_storage
|
|
|
|
# Test non-existent study
|
|
assert storage.get_study_id_by_name_if_exists("non-existent") is None
|
|
|
|
# Create a study
|
|
study_id = storage.create_new_study([StudyDirection.MINIMIZE], "test-study")
|
|
|
|
# Test existing study
|
|
result = storage.get_study_id_by_name_if_exists("test-study")
|
|
assert result == study_id
|
|
|
|
|
|
def test_flush_threads_exit_after_gc():
|
|
def _flush_threads():
|
|
return [t for t in threading.enumerate() if "mlflow_optuna_batch_flush_worker" in t.name]
|
|
|
|
threads_before = set(_flush_threads())
|
|
|
|
def _create_instances():
|
|
for _ in range(5):
|
|
MlflowStorage(experiment_id="1")
|
|
|
|
_create_instances()
|
|
gc.collect()
|
|
|
|
deadline = time.monotonic() + 3
|
|
while time.monotonic() < deadline:
|
|
new_threads = [t for t in _flush_threads() if t not in threads_before]
|
|
if not new_threads:
|
|
break
|
|
time.sleep(0.1)
|
|
else:
|
|
raise TimeoutError(f"Flush threads still alive after 3s: {new_threads}")
|