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"))
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_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_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"}