wehub-resource-sync
360 lines
12 KiB
Python
360 lines
12 KiB
Python
from functools import partial
|
|
from pathlib import Path
|
|
from typing import Dict, List
|
|
|
|
import pandas as pd
|
|
import pyarrow.fs
|
|
import pytest
|
|
|
|
import ray
|
|
from ray import train
|
|
from ray.air._internal.uri_utils import URI
|
|
from ray.train import CheckpointConfig, RunConfig, ScalingConfig
|
|
from ray.train.base_trainer import BaseTrainer
|
|
from ray.train.data_parallel_trainer import DataParallelTrainer
|
|
from ray.train.lightgbm import LightGBMTrainer
|
|
from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
|
|
from ray.train.trainer import TrainingFailedError
|
|
from ray.train.xgboost import XGBoostTrainer
|
|
from ray.tune import Callback
|
|
from ray.tune.experiment import Trial
|
|
|
|
|
|
@pytest.fixture
|
|
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.
|
|
if ray.is_initialized():
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_6_cpus():
|
|
address_info = ray.init(num_cpus=6)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
if ray.is_initialized():
|
|
ray.shutdown()
|
|
|
|
|
|
class _TestSpecificError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def _failing_train_fn(config):
|
|
checkpoint = train.get_checkpoint()
|
|
it = 1
|
|
if checkpoint:
|
|
it = load_dict_checkpoint(checkpoint)["it"] + 1
|
|
print(f"\nLoading from checkpoint, which is at iteration {it}...\n")
|
|
with create_dict_checkpoint({"it": it}) as checkpoint:
|
|
train.report({"it": it}, checkpoint=checkpoint)
|
|
if it == 1:
|
|
raise _TestSpecificError
|
|
|
|
|
|
class FailureInjectionCallback(Callback):
|
|
"""Inject failure at the configured iteration number."""
|
|
|
|
def __init__(self, fail_marker_path: Path, num_iters: int = 2):
|
|
self.num_iters = num_iters
|
|
self.fail_marker_path = fail_marker_path
|
|
|
|
def on_trial_result(
|
|
self, iteration: int, trials: List[Trial], trial: Trial, result: Dict, **info
|
|
):
|
|
if not self.fail_marker_path.exists():
|
|
return
|
|
|
|
if trial.last_result.get("training_iteration", -1) >= self.num_iters:
|
|
print(f"Failing after {self.num_iters} iters...")
|
|
self.fail_marker_path.unlink()
|
|
raise _TestSpecificError
|
|
|
|
|
|
def test_data_parallel_trainer_restore(ray_start_4_cpus, tmpdir):
|
|
"""Restoring a DataParallelTrainer with object refs captured in the train fn
|
|
or config works by re-specifying them.
|
|
Success criteria:
|
|
- Restored to the correct iteration. (1 iteration before crash, 1 after restore).
|
|
- Results are being logged to the same directory as before.
|
|
"""
|
|
dataset_size = 10
|
|
num_workers = 2
|
|
|
|
def create_train_fn_and_config():
|
|
obj_ref = ray.put({"test": 1})
|
|
|
|
def train_fn(config):
|
|
assert ray.get(obj_ref)["test"] == 1
|
|
assert ray.get(config["obj_ref"])["test"] == 1
|
|
ds = train.get_dataset_shard("train")
|
|
assert (
|
|
sum([len(batch["feature"]) for batch in ds.iter_batches()])
|
|
== dataset_size // num_workers
|
|
)
|
|
_failing_train_fn(config)
|
|
|
|
train_loop_config = {"obj_ref": obj_ref}
|
|
return train_fn, train_loop_config
|
|
|
|
datasets = {"train": ray.data.from_items([{"feature": i} for i in range(10)])}
|
|
train_fn, train_loop_config = create_train_fn_and_config()
|
|
trainer = DataParallelTrainer(
|
|
train_loop_per_worker=train_fn,
|
|
train_loop_config=train_loop_config,
|
|
datasets=datasets,
|
|
scaling_config=ScalingConfig(num_workers=num_workers),
|
|
run_config=RunConfig(
|
|
name="data_parallel_restore_test",
|
|
storage_path=str(tmpdir),
|
|
checkpoint_config=CheckpointConfig(num_to_keep=1),
|
|
),
|
|
)
|
|
with pytest.raises(TrainingFailedError) as exc_info:
|
|
result = trainer.fit()
|
|
assert isinstance(exc_info.value.__cause__, _TestSpecificError)
|
|
|
|
# Include an explicit cluster shutdown.
|
|
# Otherwise, the previously registered object references will still exist,
|
|
# and the test may trivially pass.
|
|
ray.shutdown()
|
|
ray.init(num_cpus=4)
|
|
|
|
train_fn, train_loop_config = create_train_fn_and_config()
|
|
datasets = {"train": ray.data.from_items([{"feature": i} for i in range(10)])}
|
|
trainer = DataParallelTrainer.restore(
|
|
str(tmpdir / "data_parallel_restore_test"),
|
|
train_loop_per_worker=train_fn,
|
|
train_loop_config=train_loop_config,
|
|
datasets=datasets,
|
|
)
|
|
result = trainer.fit()
|
|
assert not result.error
|
|
assert result.metrics["training_iteration"] == 2
|
|
assert result.metrics["iterations_since_restore"] == 1
|
|
assert tmpdir / "data_parallel_restore_test" in Path(result.path).parents
|
|
|
|
|
|
@pytest.mark.parametrize("trainer_cls", [XGBoostTrainer, LightGBMTrainer])
|
|
def test_gbdt_trainer_restore(ray_start_6_cpus, tmp_path, trainer_cls, monkeypatch):
|
|
"""Tests restoring gradient boosted decision tree trainers.
|
|
Success criteria:
|
|
- Picks up at the right iteration. 2 before crash. 3 after. 5 total trees.
|
|
- Results are being logged to the same directory as before.
|
|
"""
|
|
monkeypatch.setenv("TUNE_GLOBAL_CHECKPOINT_S", "0")
|
|
exp_name = f"{trainer_cls.__name__}_restore_test"
|
|
datasets = {
|
|
"train": ray.data.from_pandas(
|
|
pd.DataFrame({"x": range(100), "y": range(1, 101)})
|
|
)
|
|
}
|
|
|
|
fail_marker_path = tmp_path / "fail_marker"
|
|
fail_marker_path.touch()
|
|
|
|
trainer = trainer_cls(
|
|
label_column="y",
|
|
params={
|
|
"objective": (
|
|
"reg:squarederror" if trainer_cls == XGBoostTrainer else "regression"
|
|
)
|
|
},
|
|
datasets=datasets,
|
|
scaling_config=ScalingConfig(
|
|
num_workers=2, trainer_resources={"CPU": 0}, resources_per_worker={"CPU": 1}
|
|
),
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
name=exp_name,
|
|
checkpoint_config=CheckpointConfig(
|
|
num_to_keep=1, checkpoint_frequency=1, checkpoint_at_end=False
|
|
),
|
|
callbacks=[FailureInjectionCallback(fail_marker_path, num_iters=2)],
|
|
),
|
|
num_boost_round=5,
|
|
)
|
|
with pytest.raises(TrainingFailedError):
|
|
result = trainer.fit()
|
|
|
|
trainer = trainer_cls.restore(str(tmp_path / exp_name), datasets=datasets)
|
|
result = trainer.fit()
|
|
assert not result.error
|
|
assert result.metrics["training_iteration"] == 5
|
|
assert result.metrics["iterations_since_restore"] == 3
|
|
assert tmp_path / exp_name in Path(result.path).parents
|
|
|
|
|
|
@pytest.mark.parametrize("name", [None, "restore_from_uri"])
|
|
def test_restore_from_uri_s3(
|
|
ray_start_4_cpus, tmp_path, monkeypatch, mock_s3_bucket_uri, name
|
|
):
|
|
"""Restoration from S3 should work."""
|
|
trainer = DataParallelTrainer(
|
|
train_loop_per_worker=lambda config: train.report({"score": 1}),
|
|
scaling_config=ScalingConfig(num_workers=2),
|
|
run_config=RunConfig(name=name, storage_path=mock_s3_bucket_uri),
|
|
)
|
|
result = trainer.fit()
|
|
|
|
if name is None:
|
|
name = Path(result.path).parent.name
|
|
|
|
# Restore from S3
|
|
assert DataParallelTrainer.can_restore(str(URI(mock_s3_bucket_uri) / name))
|
|
DataParallelTrainer.restore(str(URI(mock_s3_bucket_uri) / name))
|
|
|
|
|
|
def test_restore_with_datasets(ray_start_4_cpus, tmpdir):
|
|
"""Datasets are required to re-specify if they were originally provided."""
|
|
datasets = {
|
|
"train": ray.data.from_items([{"x": x, "y": x + 1} for x in range(8)]),
|
|
"valid": ray.data.from_items([{"x": x, "y": x + 1} for x in range(8)]),
|
|
}
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_loop_per_worker=lambda config: train.report({"score": 1}),
|
|
datasets=datasets,
|
|
scaling_config=ScalingConfig(num_workers=2),
|
|
run_config=RunConfig(name="datasets_respecify_test"),
|
|
)
|
|
trainer._save(pyarrow.fs.LocalFileSystem(), str(tmpdir))
|
|
|
|
# Restore should complain, if all the datasets don't get passed in again
|
|
with pytest.raises(ValueError):
|
|
DataParallelTrainer.restore(str(tmpdir))
|
|
|
|
with pytest.raises(ValueError):
|
|
DataParallelTrainer.restore(str(tmpdir), datasets={"train": datasets["train"]})
|
|
|
|
with pytest.raises(ValueError):
|
|
DataParallelTrainer.restore(
|
|
str(tmpdir),
|
|
datasets={"train": datasets["train"], "invalid_key": datasets["valid"]},
|
|
)
|
|
|
|
trainer = DataParallelTrainer.restore(str(tmpdir), datasets=datasets)
|
|
|
|
|
|
def test_restore_from_invalid_dir(tmpdir):
|
|
"""Should raise an error if the restore directory doesn't exist or is invalid."""
|
|
with pytest.raises(ValueError):
|
|
BaseTrainer.restore(str(tmpdir))
|
|
|
|
with pytest.raises(ValueError):
|
|
BaseTrainer.restore("mock:///not/found")
|
|
|
|
|
|
def test_trainer_can_restore_utility(tmp_path):
|
|
"""Make sure that `can_restore` detects an existing experiment at a
|
|
local/remote path and only returns True if it's at the Train experiment dir root.
|
|
"""
|
|
name = "exp_name"
|
|
path = tmp_path / name
|
|
|
|
assert not DataParallelTrainer.can_restore(path)
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_loop_per_worker=lambda config: train.report({"score": 1}),
|
|
scaling_config=ScalingConfig(num_workers=1),
|
|
)
|
|
(tmp_path / name).mkdir(exist_ok=True)
|
|
trainer._save(pyarrow.fs.LocalFileSystem(), str(tmp_path / name))
|
|
|
|
assert DataParallelTrainer.can_restore(path)
|
|
|
|
|
|
@pytest.mark.parametrize("eventual_success", [True, False])
|
|
def test_retry_with_max_failures(ray_start_4_cpus, eventual_success):
|
|
"""Test auto-resume of a Train run when setting max_failures > 0."""
|
|
|
|
num_failures = 2 if eventual_success else 3
|
|
max_retries = 2
|
|
final_iter = 10
|
|
|
|
def train_func():
|
|
ckpt = train.get_checkpoint()
|
|
itr = 1
|
|
restore_count = 0
|
|
if ckpt:
|
|
ckpt = load_dict_checkpoint(ckpt)
|
|
itr = ckpt["iter"] + 1
|
|
restore_count = ckpt["restore_count"] + 1
|
|
|
|
for i in range(itr, final_iter + 1):
|
|
with create_dict_checkpoint(
|
|
dict(iter=i, restore_count=restore_count)
|
|
) as checkpoint:
|
|
train.report(dict(test=i, training_iteration=i), checkpoint=checkpoint)
|
|
if restore_count < num_failures:
|
|
raise RuntimeError("try to fail me")
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_func,
|
|
scaling_config=ScalingConfig(num_workers=2),
|
|
run_config=RunConfig(
|
|
failure_config=train.FailureConfig(max_failures=max_retries)
|
|
),
|
|
)
|
|
|
|
if not eventual_success:
|
|
# If we gave up due to hitting our max retry attempts,
|
|
# then `trainer.fit` should raise the last error we encountered.
|
|
with pytest.raises(TrainingFailedError):
|
|
trainer.fit()
|
|
else:
|
|
# If we encounter errors but eventually succeed, `trainer.fit` should NOT
|
|
# raise any of those errors.
|
|
result = trainer.fit()
|
|
assert not result.error
|
|
checkpoint = load_dict_checkpoint(result.checkpoint)
|
|
assert checkpoint["iter"] == final_iter
|
|
|
|
|
|
def test_restoration_after_termination(tmp_path):
|
|
"""Test that the train loop can be run again if restoring the trainer
|
|
after the run finished running successfully."""
|
|
|
|
def train_func_per_worker(config, num_epochs=5):
|
|
ckpt = train.get_checkpoint()
|
|
start_iter = 1
|
|
if ckpt:
|
|
ckpt = load_dict_checkpoint(ckpt)
|
|
start_iter = ckpt["iter"] + 1
|
|
|
|
for i in range(start_iter, num_epochs + 1):
|
|
with create_dict_checkpoint(dict(iter=i)) as checkpoint:
|
|
train.report(dict(iter=i), checkpoint=checkpoint)
|
|
|
|
name = "exp_name"
|
|
path = tmp_path / name
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_loop_per_worker=train_func_per_worker,
|
|
scaling_config=ScalingConfig(num_workers=1),
|
|
run_config=RunConfig(
|
|
name=name,
|
|
storage_path=tmp_path,
|
|
checkpoint_config=CheckpointConfig(num_to_keep=2),
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.metrics["iter"] == 5
|
|
|
|
restored_trainer = DataParallelTrainer.restore(
|
|
str(path), train_loop_per_worker=partial(train_func_per_worker, num_epochs=10)
|
|
)
|
|
new_result = restored_trainer.fit()
|
|
assert new_result.metrics["iter"] == 10
|
|
|
|
assert new_result.path == result.path
|
|
assert len(list(Path(new_result.path).glob("checkpoint*"))) == 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|