def test_updates_properties(storage_mode: str) -> None:
"""Check for any distributed deadlock following a property read."""
with StorageSupplier(storage_mode) as storage:
if dist.get_rank() == 0: # type: ignore
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
trial.suggest_float("f", 0, 1)
trial.suggest_int("i", 0, 1)
trial.suggest_categorical("c", ("a", "b", "c"))
property_names = [
p
for p in dir(TorchDistributedTrial)
if isinstance(getattr(TorchDistributedTrial, p), property)
]
# Rank 0 can read properties without deadlock.
if dist.get_rank() == 0: # type: ignore
[getattr(trial, p) for p in property_names]
dist.barrier() # type: ignore
# Same with rank 1.
if dist.get_rank() == 1: # type: ignore
[getattr(trial, p) for p in property_names]
dist.barrier() # type: ignore
def test_torch_distributed_trial_invalid_argument() -> None:
with pytest.raises(ValueError):
if dist.get_rank() == 0:
TorchDistributedTrial(None)
else:
study = optuna.create_study()
TorchDistributedTrial(study.ask())
def test_torch_distributed_trial_experimental_warning() -> None:
with pytest.warns(optuna.exceptions.ExperimentalWarning):
if dist.get_rank() == 0:
study = optuna.create_study()
TorchDistributedTrial(study.ask())
else:
TorchDistributedTrial(None)
def test_datetime_start(storage_mode: str) -> None:
with StorageSupplier(storage_mode) as storage:
if dist.get_rank() == 0:
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
assert isinstance(trial.datetime_start, datetime.datetime)
def test_number(storage_mode: str) -> None:
with StorageSupplier(storage_mode) as storage:
if dist.get_rank() == 0:
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
assert trial.number == 0
def test_params(storage_mode: str) -> None:
with StorageSupplier(storage_mode) as storage:
if dist.get_rank() == 0:
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
trial.suggest_float("f", 0, 1)
trial.suggest_int("i", 0, 1)
trial.suggest_categorical("c", ("a", "b", "c"))
params = trial.params
assert 0 <= params["f"] <= 1
assert 0 <= params["i"] <= 1
assert params["c"] in {"a", "b", "c"}
def test_system_attrs(storage_mode: str) -> None:
with StorageSupplier(storage_mode) as storage:
if dist.get_rank() == 0:
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
trial.set_system_attr("dataset", "mnist")
trial.set_system_attr("batch_size", 128)
assert trial.system_attrs["dataset"] == "mnist"
assert trial.system_attrs["batch_size"] == 128
def test_should_prune(storage_mode: str, is_pruning: bool) -> None:
with StorageSupplier(storage_mode) as storage:
if dist.get_rank() == 0: # type: ignore
study = optuna.create_study(storage=storage, pruner=DeterministicPruner(is_pruning))
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
trial.report(1, 0)
assert trial.should_prune() == is_pruning
def test_system_attrs_with_exception() -> None:
with StorageSupplier("sqlite") as storage:
if dist.get_rank() == 0:
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
with pytest.raises(TypeError):
trial.set_system_attr("not serializable", torch.Tensor([1, 2]))
def test_suggest_loguniform(storage_mode: str) -> None:
with StorageSupplier(storage_mode) as storage:
if dist.get_rank() == 0:
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
x1 = trial.suggest_loguniform("x", 1e-7, 1)
assert 1e-7 <= x1 <= 1
x2 = trial.suggest_loguniform("x", 1e-7, 1)
assert x1 == x2
def test_suggest_categorical(storage_mode: str) -> None:
with StorageSupplier(storage_mode) as storage:
if dist.get_rank() == 0:
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
x1 = trial.suggest_categorical("x", ("a", "b", "c"))
assert x1 in {"a", "b", "c"}
x2 = trial.suggest_categorical("x", ("a", "b", "c"))
assert x1 == x2
def test_suggest_int(storage_mode: str) -> None:
with StorageSupplier(storage_mode) as storage:
if dist.get_rank() == 0:
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
x1 = trial.suggest_int("x", 0, 10)
assert 0 <= x1 <= 10
x2 = trial.suggest_int("x", 0, 10)
assert x1 == x2
def test_report(storage_mode: str) -> None:
with StorageSupplier(storage_mode) as storage:
study: Optional[optuna.study.Study] = None
if dist.get_rank() == 0:
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
trial.report(1, 0)
if dist.get_rank() == 0:
assert study is not None
study.trials[0].intermediate_values[0] == 1
def test_suggest_discrete_uniform(storage_mode: str) -> None:
with StorageSupplier(storage_mode) as storage:
if dist.get_rank() == 0: # type: ignore
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
x1 = trial.suggest_discrete_uniform("x", 0, 10, 2)
assert 0 <= x1 <= 10
assert x1 % 2 == 0
x2 = trial.suggest_discrete_uniform("x", 0, 10, 2)
assert x1 == x2
def test_report_nan(storage_mode: str) -> None:
with StorageSupplier(storage_mode) as storage:
study: Optional[optuna.study.Study] = None
if dist.get_rank() == 0:
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
with pytest.raises(TypeError):
trial.report("abc", 0) # type: ignore
if dist.get_rank() == 0:
assert study is not None
assert len(study.trials[0].intermediate_values) == 0
def test_distributions(storage_mode: str) -> None:
with StorageSupplier(storage_mode) as storage:
if dist.get_rank() == 0:
study = optuna.create_study(storage=storage)
trial = TorchDistributedTrial(study.ask())
else:
trial = TorchDistributedTrial(None)
trial.suggest_float("u", 0, 1)
trial.suggest_float("lu", 1e-7, 1, log=True)
trial.suggest_float("du", 0, 1, step=0.5)
trial.suggest_int("i", 0, 1)
trial.suggest_int("il", 1, 128, log=True)
trial.suggest_categorical("c", ("a", "b", "c"))
distributions = trial.distributions
assert distributions["u"] == optuna.distributions.UniformDistribution(
0, 1)
assert distributions[
"lu"] == optuna.distributions.LogUniformDistribution(1e-7, 1)
assert distributions[
"du"] == optuna.distributions.DiscreteUniformDistribution(
0, 1, 0.5)
assert distributions[
"i"] == optuna.distributions.IntUniformDistribution(0, 1)
assert distributions[
"il"] == optuna.distributions.IntLogUniformDistribution(1, 128)
assert distributions[
"c"] == optuna.distributions.CategoricalDistribution(
("a", "b", "c"))
