def test_log_api_call_count(wandb: mock.MagicMock) -> None:
wandb.sdk.wandb_run.Run = mock.MagicMock
study = optuna.create_study()
wandbc = WeightsAndBiasesCallback()
@wandbc.track_in_wandb()
def _decorated_objective(trial: optuna.trial.Trial) -> float:
result = _objective_func(trial)
wandb.run.log({"result": result})
return result
target_n_trials = 10
study.optimize(_objective_func,
n_trials=target_n_trials,
callbacks=[wandbc])
assert wandb.run.log.call_count == target_n_trials
wandbc = WeightsAndBiasesCallback(as_multirun=True)
wandb.run.reset_mock()
study.optimize(_decorated_objective,
n_trials=target_n_trials,
callbacks=[wandbc])
assert wandb.run.log.call_count == 2 * target_n_trials
wandb.run = None
study.optimize(_objective_func,
n_trials=target_n_trials,
callbacks=[wandbc])
assert wandb.init().log.call_count == target_n_trials
def test_run_initialized(wandb: mock.MagicMock) -> None:
wandb.sdk.wandb_run.Run = mock.MagicMock
n_trials = 10
wandb_kwargs = {
"project": "optuna",
"group": "summary",
"job_type": "logging",
"mode": "offline",
"tags": ["test-tag"],
}
WeightsAndBiasesCallback(metric_name="mse",
wandb_kwargs=wandb_kwargs,
as_multirun=False)
wandb.init.assert_called_once_with(project="optuna",
group="summary",
job_type="logging",
mode="offline",
tags=["test-tag"])
wandbc = WeightsAndBiasesCallback(metric_name="mse",
wandb_kwargs=wandb_kwargs,
as_multirun=True)
wandb.run = None
study = optuna.create_study(direction="minimize")
_wrapped_func = wandbc.track_in_wandb()(lambda t: 1.0)
wandb.init.reset_mock()
trial = optuna.create_trial(value=1.0)
_wrapped_func(trial)
wandb.init.assert_called_once_with(project="optuna",
group="summary",
job_type="logging",
mode="offline",
tags=["test-tag"])
wandb.init.reset_mock()
study.optimize(_objective_func, n_trials=n_trials, callbacks=[wandbc])
wandb.init.assert_called_with(project="optuna",
group="summary",
job_type="logging",
mode="offline",
tags=["test-tag"])
assert wandb.init.call_count == n_trials
wandb.init().finish.assert_called()
assert wandb.init().finish.call_count == n_trials
def test_multiobjective_values_registered_on_epoch_with_logging(
wandb: mock.MagicMock, metrics: Union[str, Sequence[str]],
expected: List[str]) -> None:
wandbc = WeightsAndBiasesCallback(as_multirun=True, metric_name=metrics)
@wandbc.track_in_wandb()
def _decorated_objective(trial: optuna.trial.Trial) -> Tuple[float, float]:
result0, result1 = _multiobjective_func(trial)
wandb.run.log({"result0": result0, "result1": result1})
return result0, result1
study = optuna.create_study(directions=["minimize", "maximize"])
study.enqueue_trial({"x": 2, "y": 24})
study.optimize(_decorated_objective, n_trials=1, callbacks=[wandbc])
logged_in_decorator = wandb.run.log.mock_calls[0][1][0]
logged_in_callback = wandb.run.log.mock_calls[1][1][0]
assert len(wandb.run.log.mock_calls) == 2
assert list(logged_in_decorator) == ["result0", "result1"]
assert list(logged_in_callback) == expected
call_args = wandb.run.log.call_args
assert call_args[1] == {"step": 0}
def test_multiobjective_values_registered_on_epoch(
wandb: mock.MagicMock,
metrics: Union[str, Sequence[str]],
as_multirun: bool,
expected: List[str],
) -> None:
def assert_call_args(log_func: mock.MagicMock, as_multirun: bool) -> None:
call_args = log_func.call_args
assert list(call_args[0][0].keys()) == expected
assert call_args[1] == {"step": None if as_multirun else 0}
wandb.sdk.wandb_run.Run = mock.MagicMock
if as_multirun:
wandb.run = None
log_func = wandb.init().log
else:
log_func = wandb.run.log
study = optuna.create_study(directions=["minimize", "maximize"])
wandbc = WeightsAndBiasesCallback(metric_name=metrics,
as_multirun=as_multirun)
study.optimize(_multiobjective_func, n_trials=1, callbacks=[wandbc])
assert_call_args(log_func, as_multirun)
def test_values_registered_on_epoch_with_logging(wandb: mock.MagicMock,
metric: str,
expected: List[str]) -> None:
wandb.sdk.wandb_run.Run = mock.MagicMock
study = optuna.create_study()
wandbc = WeightsAndBiasesCallback(metric_name=metric, as_multirun=True)
@wandbc.track_in_wandb()
def _decorated_objective(trial: optuna.trial.Trial) -> float:
result = _objective_func(trial)
wandb.run.log({"result": result})
return result
study.enqueue_trial({"x": 2, "y": 25})
study.optimize(_decorated_objective, n_trials=1, callbacks=[wandbc])
logged_in_decorator = wandb.run.log.mock_calls[0][1][0]
logged_in_callback = wandb.run.log.mock_calls[1][1][0]
assert len(wandb.run.log.mock_calls) == 2
assert list(logged_in_decorator) == ["result"]
assert list(logged_in_callback) == expected
call_args = wandb.run.log.call_args
assert call_args[1] == {"step": 0}
def test_multiobjective_raises_on_name_mismatch(wandb: mock.MagicMock,
metrics: List[str]) -> None:
wandbc = WeightsAndBiasesCallback(metric_name=metrics)
study = optuna.create_study(directions=["minimize", "maximize"])
with pytest.raises(ValueError):
study.optimize(_multiobjective_func, n_trials=1, callbacks=[wandbc])
def test_multiobjective_attributes_set_on_epoch(wandb: mock.MagicMock) -> None:
wandb.config.update = mock.MagicMock()
wandbc = WeightsAndBiasesCallback()
study = optuna.create_study(directions=["minimize", "maximize"])
study.optimize(_multiobjective_func, n_trials=1, callbacks=[wandbc])
expected = {"direction": ["MINIMIZE", "MAXIMIZE"]}
wandb.config.update.assert_called_once_with(expected)
def test_log_api_call_count(wandb: mock.Mock) -> None:
wandb.log = mock.MagicMock()
wandbc = WeightsAndBiasesCallback()
target_n_trials = 10
study = optuna.create_study()
study.optimize(_objective_func,
n_trials=target_n_trials,
callbacks=[wandbc])
assert wandb.log.call_count == target_n_trials
def test_values_registered_on_epoch(wandb: mock.Mock) -> None:
wandb.log = mock.MagicMock()
wandbc = WeightsAndBiasesCallback()
study = optuna.create_study()
study.optimize(_objective_func, n_trials=1, callbacks=[wandbc])
kall = wandb.log.call_args
assert list(kall[0][0].keys()) == ["x", "y", "value"]
assert kall[1] == {"step": 0}
def test_values_registered_on_epoch(wandb: mock.Mock, metric: str,
expected: List[str]) -> None:
wandb.log = mock.MagicMock()
wandbc = WeightsAndBiasesCallback(metric_name=metric)
study = optuna.create_study()
study.optimize(_objective_func, n_trials=1, callbacks=[wandbc])
kall = wandb.log.call_args
assert list(kall[0][0].keys()) == expected
assert kall[1] == {"step": 0}
def test_multiobjective_values_registered_on_epoch(
wandb: mock.Mock, metrics: Union[str, Sequence[str]],
expected: List[str]) -> None:
wandb.log = mock.MagicMock()
wandbc = WeightsAndBiasesCallback(metric_name=metrics)
study = optuna.create_study(directions=["minimize", "maximize"])
study.optimize(_multiobjective_func, n_trials=1, callbacks=[wandbc])
kall = wandb.log.call_args
assert list(kall[0][0].keys()) == expected
assert kall[1] == {"step": 0}
def test_run_initialized(wandb: mock.MagicMock) -> None:
wandb_kwargs = {
"project": "optuna",
"group": "summary",
"job_type": "logging",
"mode": "offline",
}
WeightsAndBiasesCallback(
metric_name="mse",
wandb_kwargs=wandb_kwargs,
)
wandb.init.assert_called_once_with(
project="optuna",
group="summary",
job_type="logging",
mode="offline",
)
def test_attributes_set_on_epoch(wandb: mock.MagicMock,
as_multirun: bool) -> None:
wandb.sdk.wandb_run.Run = mock.MagicMock
expected_config: Dict[str, Any] = {"direction": ["MINIMIZE"]}
trial_params = {"x": 1.1, "y": 2.2}
expected_config_with_params = {**expected_config, **trial_params}
study = optuna.create_study(direction="minimize")
wandbc = WeightsAndBiasesCallback(as_multirun=as_multirun)
if as_multirun:
wandb.run = None
study.enqueue_trial(trial_params)
study.optimize(_objective_func, n_trials=1, callbacks=[wandbc])
if as_multirun:
wandb.init().config.update.assert_called_once_with(
expected_config_with_params)
else:
wandb.run.config.update.assert_called_once_with(expected_config)
def test_multiobjective_raises_on_type_mismatch(metrics: Any) -> None:
with pytest.raises(TypeError):
WeightsAndBiasesCallback(metric_name=metrics)