wehub-resource-sync
544 lines
18 KiB
Python
544 lines
18 KiB
Python
import os
|
|
import unittest
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from sklearn.datasets import load_breast_cancer
|
|
from sklearn.utils import shuffle
|
|
|
|
import ray
|
|
from ray import tune
|
|
from ray.data import Dataset, Datasource, ReadTask, from_pandas, read_datasource
|
|
from ray.data.block import BlockMetadata
|
|
from ray.train.data_parallel_trainer import DataParallelTrainer
|
|
from ray.train.examples.pytorch.torch_linear_example import (
|
|
train_func as linear_train_func,
|
|
)
|
|
from ray.train.torch import TorchTrainer
|
|
from ray.train.trainer import BaseTrainer
|
|
from ray.train.xgboost import XGBoostTrainer
|
|
from ray.tune import Callback, CheckpointConfig, CLIReporter, RunConfig
|
|
from ray.tune.tune_config import TuneConfig
|
|
from ray.tune.tuner import Tuner
|
|
|
|
|
|
@pytest.fixture
|
|
def shutdown_only():
|
|
yield None
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def chdir_tmpdir(tmpdir):
|
|
old_cwd = os.getcwd()
|
|
os.chdir(tmpdir)
|
|
yield tmpdir
|
|
os.chdir(old_cwd)
|
|
|
|
|
|
class DummyTrainer(BaseTrainer):
|
|
_scaling_config_allowed_keys = BaseTrainer._scaling_config_allowed_keys + [
|
|
"num_workers",
|
|
"use_gpu",
|
|
"resources_per_worker",
|
|
"placement_strategy",
|
|
]
|
|
|
|
def training_loop(self) -> None:
|
|
for i in range(5):
|
|
tune.report({"step": i})
|
|
|
|
|
|
class FailingTrainer(DummyTrainer):
|
|
def training_loop(self) -> None:
|
|
raise RuntimeError("There is an error in trainer!")
|
|
|
|
|
|
class TestDatasource(Datasource):
|
|
def __init__(self, do_shuffle: bool):
|
|
self._shuffle = do_shuffle
|
|
|
|
def prepare_read(self, parallelism: int, **read_args):
|
|
import pyarrow as pa
|
|
|
|
def load_data():
|
|
data_raw = load_breast_cancer(as_frame=True)
|
|
dataset_df = data_raw["data"]
|
|
dataset_df["target"] = data_raw["target"]
|
|
if self._shuffle:
|
|
dataset_df = shuffle(dataset_df)
|
|
return [pa.Table.from_pandas(dataset_df)]
|
|
|
|
meta = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=None,
|
|
input_files=None,
|
|
exec_stats=None,
|
|
)
|
|
return [ReadTask(load_data, meta)]
|
|
|
|
|
|
def gen_dataset_func(do_shuffle: Optional[bool] = False) -> Dataset:
|
|
test_datasource = TestDatasource(do_shuffle)
|
|
return read_datasource(test_datasource, override_num_blocks=1)
|
|
|
|
|
|
def gen_dataset_func_eager():
|
|
data_raw = load_breast_cancer(as_frame=True)
|
|
dataset_df = data_raw["data"]
|
|
dataset_df["target"] = data_raw["target"]
|
|
dataset = from_pandas(dataset_df)
|
|
return dataset
|
|
|
|
|
|
class TunerTest(unittest.TestCase):
|
|
"""The e2e test for hparam tuning using Tuner API."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def tmp_path(self, tmp_path):
|
|
self.tmp_path = tmp_path
|
|
|
|
def setUp(self):
|
|
ray.init()
|
|
|
|
def tearDown(self):
|
|
ray.shutdown()
|
|
|
|
def test_tuner_with_xgboost_trainer(self):
|
|
"""Test a successful run."""
|
|
trainer = XGBoostTrainer(
|
|
label_column="target",
|
|
params={},
|
|
datasets={"train": gen_dataset_func_eager()},
|
|
)
|
|
param_space = {
|
|
"scaling_config": ray.train.ScalingConfig(
|
|
num_workers=tune.grid_search([1, 2])
|
|
),
|
|
"datasets": {
|
|
"train": tune.grid_search(
|
|
[gen_dataset_func(), gen_dataset_func(do_shuffle=True)]
|
|
),
|
|
},
|
|
"params": {
|
|
"objective": "binary:logistic",
|
|
"tree_method": "approx",
|
|
"eval_metric": ["logloss", "error"],
|
|
"eta": tune.loguniform(1e-4, 1e-1),
|
|
"subsample": tune.uniform(0.5, 1.0),
|
|
"max_depth": tune.randint(1, 9),
|
|
},
|
|
}
|
|
tuner = Tuner(
|
|
trainable=trainer,
|
|
run_config=RunConfig(name="test_tuner"),
|
|
param_space=param_space,
|
|
tune_config=TuneConfig(mode="min", metric="train-error"),
|
|
# limiting the number of trials running at one time.
|
|
# As the unit test only has access to 4 CPUs on Buildkite.
|
|
_tuner_kwargs={"max_concurrent_trials": 1},
|
|
)
|
|
results = tuner.fit()
|
|
assert len(results) == 4
|
|
|
|
def test_tuner_with_xgboost_trainer_driver_fail_and_resume(self):
|
|
# So that we have some global checkpointing happening.
|
|
os.environ["TUNE_GLOBAL_CHECKPOINT_S"] = "1"
|
|
trainer = XGBoostTrainer(
|
|
label_column="target",
|
|
params={},
|
|
datasets={"train": gen_dataset_func_eager()},
|
|
)
|
|
# prep_v1 = StandardScaler(["worst radius", "worst area"])
|
|
# prep_v2 = StandardScaler(["worst concavity", "worst smoothness"])
|
|
param_space = {
|
|
"scaling_config": ray.train.ScalingConfig(
|
|
num_workers=tune.grid_search([1, 2])
|
|
),
|
|
"datasets": {
|
|
"train": tune.grid_search(
|
|
[gen_dataset_func(), gen_dataset_func(do_shuffle=True)]
|
|
),
|
|
},
|
|
"params": {
|
|
"objective": "binary:logistic",
|
|
"tree_method": "approx",
|
|
"eval_metric": ["logloss", "error"],
|
|
"eta": tune.loguniform(1e-4, 1e-1),
|
|
"subsample": tune.uniform(0.5, 1.0),
|
|
"max_depth": tune.randint(1, 9),
|
|
},
|
|
}
|
|
|
|
class FailureInjectionCallback(Callback):
|
|
"""Inject failure at the configured iteration number."""
|
|
|
|
def __init__(self, num_iters=10):
|
|
self.num_iters = num_iters
|
|
|
|
def on_step_end(self, iteration, trials, **kwargs):
|
|
if iteration == self.num_iters:
|
|
print(f"Failing after {self.num_iters} iters.")
|
|
raise RuntimeError
|
|
|
|
tuner = Tuner(
|
|
trainable=trainer,
|
|
run_config=RunConfig(
|
|
name="test_tuner_driver_fail",
|
|
storage_path=str(self.tmp_path),
|
|
callbacks=[FailureInjectionCallback()],
|
|
),
|
|
param_space=param_space,
|
|
tune_config=TuneConfig(mode="min", metric="train-error"),
|
|
# limiting the number of trials running at one time.
|
|
# As the unit test only has access to 4 CPUs on Buildkite.
|
|
_tuner_kwargs={"max_concurrent_trials": 1},
|
|
)
|
|
with self.assertRaises(RuntimeError):
|
|
tuner.fit()
|
|
|
|
# Test resume
|
|
restore_path = os.path.join(self.tmp_path, "test_tuner_driver_fail")
|
|
tuner = Tuner.restore(restore_path, trainable=trainer, param_space=param_space)
|
|
# A hack before we figure out RunConfig semantics across resumes.
|
|
tuner._local_tuner._run_config.callbacks = None
|
|
results = tuner.fit()
|
|
assert len(results) == 4
|
|
assert not results.errors
|
|
|
|
def test_tuner_with_torch_trainer(self):
|
|
"""Test a successful run using torch trainer."""
|
|
# The following two should be tunable.
|
|
config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": 10}
|
|
scaling_config = ray.train.ScalingConfig(num_workers=1, use_gpu=False)
|
|
trainer = TorchTrainer(
|
|
train_loop_per_worker=linear_train_func,
|
|
train_loop_config=config,
|
|
scaling_config=scaling_config,
|
|
)
|
|
param_space = {
|
|
"scaling_config": ray.train.ScalingConfig(
|
|
num_workers=tune.grid_search([1, 2])
|
|
),
|
|
"train_loop_config": {
|
|
"batch_size": tune.grid_search([4, 8]),
|
|
"epochs": tune.grid_search([5, 10]),
|
|
},
|
|
}
|
|
tuner = Tuner(
|
|
trainable=trainer,
|
|
run_config=RunConfig(name="test_tuner"),
|
|
param_space=param_space,
|
|
tune_config=TuneConfig(mode="min", metric="loss"),
|
|
)
|
|
results = tuner.fit()
|
|
assert len(results) == 8
|
|
|
|
def test_tuner_run_config_override(self):
|
|
trainer = DummyTrainer(run_config=RunConfig(stop={"metric": 4}))
|
|
tuner = Tuner(trainer)
|
|
|
|
assert tuner._local_tuner._run_config.stop == {"metric": 4}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"params_expected",
|
|
[
|
|
(
|
|
{"run_config": RunConfig(progress_reporter=CLIReporter())},
|
|
lambda kw: isinstance(kw["progress_reporter"], CLIReporter),
|
|
),
|
|
(
|
|
{"tune_config": TuneConfig(reuse_actors=True)},
|
|
lambda kw: kw["reuse_actors"] is True,
|
|
),
|
|
(
|
|
{"run_config": RunConfig(log_to_file="some_file")},
|
|
lambda kw: kw["log_to_file"] == "some_file",
|
|
),
|
|
(
|
|
{"tune_config": TuneConfig(max_concurrent_trials=3)},
|
|
lambda kw: kw["max_concurrent_trials"] == 3,
|
|
),
|
|
(
|
|
{"tune_config": TuneConfig(time_budget_s=60)},
|
|
lambda kw: kw["time_budget_s"] == 60,
|
|
),
|
|
],
|
|
)
|
|
def test_tuner_api_kwargs(shutdown_only, params_expected):
|
|
tuner_params, assertion = params_expected
|
|
|
|
tuner = Tuner(lambda config: 1, **tuner_params)
|
|
|
|
caught_kwargs = {}
|
|
|
|
class MockExperimentAnalysis:
|
|
trials = []
|
|
|
|
def catch_kwargs(**kwargs):
|
|
caught_kwargs.update(kwargs)
|
|
return MockExperimentAnalysis()
|
|
|
|
with patch("ray.tune.impl.tuner_internal.run", catch_kwargs):
|
|
tuner.fit()
|
|
|
|
assert assertion(caught_kwargs)
|
|
|
|
|
|
def test_tuner_fn_trainable_invalid_checkpoint_config(shutdown_only):
|
|
tuner = Tuner(
|
|
lambda config: 1,
|
|
run_config=RunConfig(
|
|
checkpoint_config=CheckpointConfig(checkpoint_at_end=True)
|
|
),
|
|
)
|
|
with pytest.raises(ValueError):
|
|
tuner.fit()
|
|
|
|
tuner = Tuner(
|
|
lambda config: 1,
|
|
run_config=RunConfig(
|
|
checkpoint_config=CheckpointConfig(checkpoint_frequency=1)
|
|
),
|
|
)
|
|
with pytest.raises(ValueError):
|
|
tuner.fit()
|
|
|
|
|
|
def test_tuner_trainer_checkpoint_config(shutdown_only):
|
|
custom_training_loop_trainer = DataParallelTrainer(
|
|
train_loop_per_worker=lambda config: 1
|
|
)
|
|
tuner = Tuner(
|
|
custom_training_loop_trainer,
|
|
run_config=RunConfig(
|
|
checkpoint_config=CheckpointConfig(checkpoint_at_end=True)
|
|
),
|
|
)
|
|
with pytest.raises(ValueError):
|
|
tuner.fit()
|
|
|
|
tuner = Tuner(
|
|
custom_training_loop_trainer,
|
|
run_config=RunConfig(
|
|
checkpoint_config=CheckpointConfig(checkpoint_frequency=1)
|
|
),
|
|
)
|
|
with pytest.raises(ValueError):
|
|
tuner.fit()
|
|
|
|
handles_checkpoints_trainer = XGBoostTrainer(
|
|
label_column="target",
|
|
params={},
|
|
datasets={"train": ray.data.from_items(list(range(5)))},
|
|
)
|
|
tuner = Tuner(
|
|
handles_checkpoints_trainer,
|
|
run_config=RunConfig(
|
|
checkpoint_config=CheckpointConfig(
|
|
checkpoint_at_end=True, checkpoint_frequency=1
|
|
)
|
|
),
|
|
)._local_tuner
|
|
# Check that validation passes for a Trainer that does handle checkpointing
|
|
tuner._get_tune_run_arguments(tuner.converted_trainable)
|
|
|
|
|
|
def test_tuner_fn_trainable_checkpoint_at_end_false(shutdown_only):
|
|
tuner = Tuner(
|
|
lambda config: 1,
|
|
run_config=RunConfig(
|
|
checkpoint_config=CheckpointConfig(checkpoint_at_end=False)
|
|
),
|
|
)
|
|
tuner.fit()
|
|
|
|
|
|
def test_tuner_fn_trainable_checkpoint_at_end_none(shutdown_only):
|
|
tuner = Tuner(
|
|
lambda config: 1,
|
|
run_config=RunConfig(
|
|
checkpoint_config=CheckpointConfig(checkpoint_at_end=None)
|
|
),
|
|
)
|
|
tuner.fit()
|
|
|
|
|
|
def test_nonserializable_trainable():
|
|
import threading
|
|
|
|
lock = threading.Lock()
|
|
# Check that the `inspect_serializability` trace was printed
|
|
with pytest.raises(TypeError, match=r".*was found to be non-serializable.*"):
|
|
Tuner(lambda config: print(lock))
|
|
|
|
|
|
# TODO: [V2] Delete the `trainer` variant once V1 is fully removed.
|
|
def _test_no_chdir(runner_type, runtime_env, use_deprecated_config=False):
|
|
# Write a data file that we want to read in our training loop
|
|
with open("./read.txt", "w") as f:
|
|
f.write("data")
|
|
|
|
ray.init(num_cpus=4, runtime_env=runtime_env)
|
|
|
|
def train_func(config):
|
|
# Make sure we can access the data from the original working dir
|
|
assert os.path.exists("./read.txt") and open("./read.txt", "r").read() == "data"
|
|
|
|
# Write operations should happen in each trial's independent logdir to
|
|
# prevent write conflicts
|
|
trial_dir = Path(tune.get_context().get_trial_dir())
|
|
trial_dir.joinpath("write.txt").touch()
|
|
|
|
if runner_type == "trainer":
|
|
trainer = DataParallelTrainer(
|
|
train_func, scaling_config=ray.train.ScalingConfig(num_workers=2)
|
|
)
|
|
result = trainer.fit()
|
|
results = [result]
|
|
elif runner_type == "tuner":
|
|
tuner = Tuner(
|
|
train_func,
|
|
param_space={"id": tune.grid_search(list(range(4)))},
|
|
tune_config=(
|
|
TuneConfig(chdir_to_trial_dir=False) if use_deprecated_config else None
|
|
),
|
|
)
|
|
results = tuner.fit()
|
|
assert not results.errors
|
|
else:
|
|
raise NotImplementedError(f"Invalid: {runner_type}")
|
|
|
|
for result in results:
|
|
assert os.path.exists(os.path.join(result.path, "write.txt"))
|
|
|
|
|
|
def test_tuner_no_chdir_to_trial_dir_deprecated(shutdown_only, chdir_tmpdir):
|
|
"""Test the deprecated `chdir_to_trial_dir` config."""
|
|
with pytest.raises(DeprecationWarning):
|
|
_test_no_chdir("tuner", {}, use_deprecated_config=True)
|
|
|
|
|
|
@pytest.mark.parametrize("runtime_env", [{}, {"working_dir": "."}])
|
|
def test_tuner_no_chdir_to_trial_dir(
|
|
shutdown_only, chdir_tmpdir, monkeypatch, runtime_env
|
|
):
|
|
"""Tests that disabling the env var to keep the working directory the same
|
|
works for a Tuner run."""
|
|
from ray.train.constants import RAY_CHDIR_TO_TRIAL_DIR
|
|
|
|
monkeypatch.setenv(RAY_CHDIR_TO_TRIAL_DIR, "0")
|
|
_test_no_chdir("tuner", runtime_env)
|
|
|
|
|
|
@pytest.mark.parametrize("runtime_env", [{}, {"working_dir": "."}])
|
|
def test_trainer_no_chdir_to_trial_dir(
|
|
shutdown_only, chdir_tmpdir, monkeypatch, runtime_env
|
|
):
|
|
"""Tests that disabling the env var to keep the working directory the same
|
|
works for a Trainer run."""
|
|
from ray.train.constants import RAY_CHDIR_TO_TRIAL_DIR
|
|
|
|
monkeypatch.setenv(RAY_CHDIR_TO_TRIAL_DIR, "0")
|
|
_test_no_chdir("trainer", runtime_env)
|
|
|
|
|
|
@pytest.mark.parametrize("runtime_env", [{}, {"working_dir": "."}])
|
|
def test_tuner_relative_pathing_with_env_vars(
|
|
shutdown_only, chdir_tmpdir, tmp_path, runtime_env
|
|
):
|
|
"""Tests that `TUNE_ORIG_WORKING_DIR` environment variable can be used to access
|
|
relative paths to the original working directory.
|
|
"""
|
|
# Write a data file that we want to read in our training loop
|
|
with open("./read.txt", "w") as f:
|
|
f.write("data")
|
|
|
|
# Even if we set our runtime_env `{"working_dir": "."}` to the current directory,
|
|
# Tune should still chdir to the trial directory.
|
|
ray.init(num_cpus=1, runtime_env=runtime_env)
|
|
|
|
def train_func(config):
|
|
orig_working_dir = Path(os.environ["TUNE_ORIG_WORKING_DIR"])
|
|
assert (
|
|
str(orig_working_dir) != os.getcwd()
|
|
), f"Working directory should have changed from {orig_working_dir}"
|
|
|
|
# Make sure we can access the data from the original working dir
|
|
# Different from above: create an absolute path using the env variable
|
|
data_path = orig_working_dir / "read.txt"
|
|
assert os.path.exists(data_path) and open(data_path, "r").read() == "data"
|
|
|
|
# Tune chdirs to the trial working directory
|
|
storage = tune.get_context().get_storage()
|
|
assert Path(storage.trial_working_directory).resolve() == Path.cwd().resolve()
|
|
|
|
with open("write.txt", "w") as f:
|
|
f.write(f"{config['id']}")
|
|
|
|
tuner = Tuner(
|
|
train_func,
|
|
param_space={"id": tune.grid_search(list(range(4)))},
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
sync_config=tune.SyncConfig(sync_artifacts=True),
|
|
),
|
|
)
|
|
results = tuner.fit()
|
|
assert not results.errors
|
|
for result in results:
|
|
artifact_data = open(os.path.join(result.path, "write.txt"), "r").read()
|
|
assert artifact_data == f"{result.config['id']}"
|
|
|
|
|
|
def test_invalid_param_space(shutdown_only):
|
|
"""Check that Tune raises an error on invalid param_space types."""
|
|
|
|
def trainable(config):
|
|
return {"metric": 1}
|
|
|
|
with pytest.raises(ValueError):
|
|
Tuner(trainable, param_space="not allowed")
|
|
|
|
from ray.tune.tune import _Config
|
|
|
|
class CustomConfig(_Config):
|
|
def to_dict(self) -> dict:
|
|
return {"hparam": 1}
|
|
|
|
with pytest.raises(ValueError):
|
|
Tuner(trainable, param_space="not allowed").fit()
|
|
|
|
with pytest.raises(ValueError):
|
|
tune.run(trainable, config="not allowed")
|
|
|
|
# Dict and custom _Config subclasses are fine
|
|
Tuner(trainable, param_space={}).fit()
|
|
Tuner(trainable, param_space=CustomConfig()).fit()
|
|
tune.run(trainable, config=CustomConfig())
|
|
|
|
|
|
def test_tuner_restore_classmethod():
|
|
tuner = Tuner(lambda x: None)
|
|
|
|
# Calling `tuner.restore()` on an instance should raise an AttributeError
|
|
with pytest.raises(AttributeError):
|
|
tuner.restore("/", lambda x: None)
|
|
|
|
# Calling `Tuner.restore()` on the class should work. This will throw a
|
|
# FileNotFoundError because no checkpoint exists at that location. Since
|
|
# this happens in the downstream restoration code, this means that the
|
|
# classmethod check successfully passed.
|
|
with pytest.raises(FileNotFoundError):
|
|
tuner = Tuner.restore("/invalid", lambda x: None)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__] + sys.argv[1:]))
|