wehub-resource-sync
332 lines
10 KiB
Python
332 lines
10 KiB
Python
import logging
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray import train, tune
|
|
from ray.air.constants import TRAINING_ITERATION
|
|
from ray.train._internal.worker_group import WorkerGroup
|
|
from ray.train.backend import Backend, BackendConfig
|
|
from ray.train.data_parallel_trainer import DataParallelTrainer
|
|
from ray.train.examples.pytorch.torch_fashion_mnist_example import (
|
|
train_func_per_worker as fashion_mnist_train_func,
|
|
)
|
|
from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
|
|
from ray.train.torch import TorchTrainer
|
|
from ray.tune.tune_config import TuneConfig
|
|
from ray.tune.tuner import Tuner
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def ray_start_4_cpus():
|
|
address_info = ray.init(num_cpus=4)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_8_cpus():
|
|
address_info = ray.init(num_cpus=8)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
class TestConfig(BackendConfig):
|
|
@property
|
|
def backend_cls(self):
|
|
return TestBackend
|
|
|
|
|
|
class TestBackend(Backend):
|
|
def on_start(self, worker_group: WorkerGroup, backend_config: TestConfig):
|
|
pass
|
|
|
|
def on_shutdown(self, worker_group: WorkerGroup, backend_config: TestConfig):
|
|
pass
|
|
|
|
|
|
def torch_fashion_mnist(num_workers, use_gpu, num_samples):
|
|
trainer = TorchTrainer(
|
|
fashion_mnist_train_func,
|
|
scaling_config=train.ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
|
|
)
|
|
tuner = Tuner(
|
|
trainer,
|
|
param_space={
|
|
"train_loop_config": {
|
|
"lr": tune.loguniform(1e-4, 1e-1),
|
|
"batch_size_per_worker": tune.choice([32, 64, 128]),
|
|
"epochs": 2,
|
|
}
|
|
},
|
|
tune_config=TuneConfig(
|
|
num_samples=num_samples,
|
|
),
|
|
)
|
|
analysis = tuner.fit()._experiment_analysis
|
|
|
|
# Check that loss decreases in each trial.
|
|
for df in analysis.trial_dataframes.values():
|
|
assert df.loc[1, "loss"] < df.loc[0, "loss"]
|
|
|
|
|
|
def test_tune_torch_fashion_mnist(ray_start_8_cpus):
|
|
torch_fashion_mnist(num_workers=2, use_gpu=False, num_samples=2)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info >= (3, 12), reason="tensorflow is not installed in python 3.12+"
|
|
)
|
|
def tune_tensorflow_mnist(num_workers, use_gpu, num_samples):
|
|
from ray.train.examples.tf.tensorflow_mnist_example import (
|
|
train_func as tensorflow_mnist_train_func,
|
|
)
|
|
from ray.train.tensorflow import TensorflowTrainer
|
|
|
|
trainer = TensorflowTrainer(
|
|
tensorflow_mnist_train_func,
|
|
scaling_config=train.ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
|
|
)
|
|
tuner = Tuner(
|
|
trainer,
|
|
param_space={
|
|
"train_loop_config": {
|
|
"lr": tune.loguniform(1e-4, 1e-1),
|
|
"batch_size": tune.choice([32, 64, 128]),
|
|
"epochs": 2,
|
|
}
|
|
},
|
|
tune_config=TuneConfig(
|
|
num_samples=num_samples,
|
|
),
|
|
)
|
|
analysis = tuner.fit()._experiment_analysis
|
|
|
|
# Check that loss decreases in each trial.
|
|
for df in analysis.trial_dataframes.values():
|
|
assert df.loc[1, "loss"] < df.loc[0, "loss"]
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info >= (3, 12), reason="tensorflow is not installed in python 3.12+"
|
|
)
|
|
def test_tune_tensorflow_mnist(ray_start_8_cpus):
|
|
tune_tensorflow_mnist(num_workers=2, use_gpu=False, num_samples=2)
|
|
|
|
|
|
def test_tune_error(ray_start_4_cpus):
|
|
def train_func(config):
|
|
raise RuntimeError("Error in training function!")
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_func,
|
|
backend_config=TestConfig(),
|
|
scaling_config=train.ScalingConfig(num_workers=1),
|
|
)
|
|
tuner = Tuner(
|
|
trainer,
|
|
)
|
|
|
|
result_grid = tuner.fit()
|
|
with pytest.raises(RuntimeError):
|
|
raise result_grid[0].error
|
|
|
|
|
|
def test_tune_checkpoint(ray_start_4_cpus):
|
|
def train_func():
|
|
for i in range(9):
|
|
train.report(dict(test=i))
|
|
with create_dict_checkpoint(dict(hello="world")) as checkpoint:
|
|
train.report(dict(test=i + 1), checkpoint=checkpoint)
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_func,
|
|
backend_config=TestConfig(),
|
|
scaling_config=train.ScalingConfig(num_workers=1),
|
|
)
|
|
tuner = Tuner(
|
|
trainer,
|
|
param_space={"train_loop_config": {"max_iter": 5}},
|
|
)
|
|
|
|
result_grid = tuner.fit()
|
|
assert len(result_grid) == 1
|
|
result = result_grid[0]
|
|
assert result.checkpoint
|
|
assert load_dict_checkpoint(result.checkpoint)["hello"] == "world"
|
|
|
|
|
|
def test_reuse_checkpoint(ray_start_4_cpus):
|
|
def train_func(config):
|
|
itr = 0
|
|
ckpt = train.get_checkpoint()
|
|
if ckpt is not None:
|
|
ckpt = load_dict_checkpoint(ckpt)
|
|
itr = ckpt["iter"] + 1
|
|
|
|
for i in range(itr, config["max_iter"]):
|
|
with create_dict_checkpoint(dict(iter=i)) as checkpoint:
|
|
train.report(dict(test=i, training_iteration=i), checkpoint=checkpoint)
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_func,
|
|
backend_config=TestConfig(),
|
|
scaling_config=train.ScalingConfig(num_workers=1),
|
|
)
|
|
tuner = Tuner(
|
|
trainer,
|
|
param_space={"train_loop_config": {"max_iter": 5}},
|
|
)
|
|
result_grid = tuner.fit()
|
|
assert len(result_grid) == 1
|
|
result = result_grid[0]
|
|
assert result.checkpoint
|
|
assert load_dict_checkpoint(result.checkpoint)["iter"] == 4
|
|
|
|
tuner = Tuner.restore(result_grid.experiment_path, trainable=trainer)
|
|
result_grid = tuner.fit()
|
|
assert len(result_grid) == 1
|
|
assert len(result_grid[0].metrics_dataframe) == 5
|
|
|
|
|
|
def test_retry_with_max_failures(ray_start_4_cpus):
|
|
"""Tests trainer retry with max_failures > 0 when integrating with Tune."""
|
|
|
|
def train_func():
|
|
ckpt = train.get_checkpoint()
|
|
restored = bool(ckpt) # Does a previous checkpoint exist?
|
|
itr = 0
|
|
if ckpt:
|
|
ckpt = load_dict_checkpoint(ckpt)
|
|
itr = ckpt["iter"] + 1
|
|
|
|
for i in range(itr, 4):
|
|
if i == 2 and not restored:
|
|
raise Exception("try to fail me")
|
|
with create_dict_checkpoint(dict(iter=i)) as checkpoint:
|
|
train.report(dict(test=i, training_iteration=i), checkpoint=checkpoint)
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_func,
|
|
backend_config=TestConfig(),
|
|
scaling_config=train.ScalingConfig(num_workers=1),
|
|
)
|
|
tuner = Tuner(
|
|
trainer,
|
|
run_config=tune.RunConfig(failure_config=tune.FailureConfig(max_failures=3)),
|
|
)
|
|
|
|
result_grid = tuner.fit()
|
|
checkpoint = load_dict_checkpoint(result_grid[0].checkpoint)
|
|
assert checkpoint["iter"] == 3
|
|
df = result_grid[0].metrics_dataframe
|
|
assert len(df[TRAINING_ITERATION]) == 4
|
|
|
|
|
|
def test_restore_with_new_trainer(ray_start_4_cpus, tmpdir, propagate_logs, caplog):
|
|
def train_func(config):
|
|
raise RuntimeError("failing!")
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_func,
|
|
backend_config=TestConfig(),
|
|
scaling_config=train.ScalingConfig(num_workers=1),
|
|
run_config=train.RunConfig(
|
|
name="restore_new_trainer", storage_path=str(tmpdir)
|
|
),
|
|
datasets={"train": ray.data.from_items([{"a": i} for i in range(10)])},
|
|
)
|
|
results = Tuner(trainer).fit()
|
|
assert results.errors
|
|
|
|
def train_func(config):
|
|
dataset = train.get_dataset_shard("train")
|
|
assert train.get_context().get_world_size() == 2
|
|
rows = 0
|
|
for _ in dataset.iter_rows():
|
|
rows += 1
|
|
assert rows == 10
|
|
|
|
trainer = DataParallelTrainer(
|
|
# Training function can be modified
|
|
train_func,
|
|
backend_config=TestConfig(),
|
|
# ScalingConfig can be modified
|
|
scaling_config=train.ScalingConfig(num_workers=2),
|
|
# New RunConfig will be ignored
|
|
run_config=train.RunConfig(name="ignored"),
|
|
# Datasets and preprocessors can be re-specified
|
|
datasets={"train": ray.data.from_items([{"a": i} for i in range(20)])},
|
|
)
|
|
caplog.clear()
|
|
with caplog.at_level(logging.WARNING, logger="ray.tune.impl.tuner_internal"):
|
|
tuner = Tuner.restore(
|
|
str(tmpdir / "restore_new_trainer"),
|
|
trainable=trainer,
|
|
resume_errored=True,
|
|
)
|
|
assert "they will be ignored in the resumed run" in caplog.text
|
|
|
|
results = tuner.fit()
|
|
assert not results.errors
|
|
|
|
|
|
@pytest.mark.parametrize("in_trainer", [True, False])
|
|
@pytest.mark.parametrize("in_tuner", [True, False])
|
|
def test_run_config_in_trainer_and_tuner(
|
|
propagate_logs, tmp_path, caplog, in_trainer, in_tuner
|
|
):
|
|
trainer_run_config = (
|
|
train.RunConfig(name="trainer", storage_path=str(tmp_path))
|
|
if in_trainer
|
|
else None
|
|
)
|
|
tuner_run_config = (
|
|
tune.RunConfig(name="tuner", storage_path=str(tmp_path)) if in_tuner else None
|
|
)
|
|
trainer = DataParallelTrainer(
|
|
lambda config: None,
|
|
backend_config=TestConfig(),
|
|
scaling_config=train.ScalingConfig(num_workers=1),
|
|
run_config=trainer_run_config,
|
|
)
|
|
with caplog.at_level(logging.INFO, logger="ray.tune.impl.tuner_internal"):
|
|
tuner = Tuner(trainer, run_config=tuner_run_config)
|
|
|
|
both_msg = (
|
|
"`RunConfig` was passed to both the `Tuner` and the `DataParallelTrainer`"
|
|
)
|
|
run_config = tuner._local_tuner.get_run_config()
|
|
if in_trainer and in_tuner:
|
|
assert run_config.name == "tuner"
|
|
assert both_msg in caplog.text
|
|
elif in_trainer and not in_tuner:
|
|
assert run_config.name == "trainer"
|
|
assert both_msg not in caplog.text
|
|
elif not in_trainer and in_tuner:
|
|
assert run_config.name == "tuner"
|
|
assert both_msg not in caplog.text
|
|
else:
|
|
assert both_msg not in caplog.text
|
|
|
|
|
|
def test_run_config_in_param_space():
|
|
trainer = DataParallelTrainer(
|
|
lambda config: None,
|
|
backend_config=TestConfig(),
|
|
scaling_config=train.ScalingConfig(num_workers=1),
|
|
)
|
|
with pytest.raises(ValueError):
|
|
Tuner(trainer, param_space={"run_config": train.RunConfig(name="ignored")})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|