wehub-resource-sync
155 lines
5.2 KiB
Python
155 lines
5.2 KiB
Python
import pyarrow
|
|
import pytest
|
|
|
|
import ray
|
|
from ray import train
|
|
from ray.air._internal.uri_utils import URI
|
|
from ray.air.constants import EXPR_RESULT_FILE
|
|
from ray.train import CheckpointConfig, Result, RunConfig, ScalingConfig
|
|
from ray.train.base_trainer import TrainingFailedError
|
|
from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
|
|
from ray.train.torch import TorchTrainer
|
|
from ray.tune import TuneConfig, Tuner
|
|
|
|
_PARAM_SPACE = {"a": 1, "b": 2}
|
|
|
|
|
|
@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.
|
|
ray.shutdown()
|
|
|
|
|
|
def build_dummy_trainer(configs):
|
|
def worker_loop(_config):
|
|
for i in range(configs["NUM_ITERATIONS"]):
|
|
# Do some random reports in between checkpoints.
|
|
train.report({"metric_a": -100, "metric_b": -100})
|
|
|
|
if ray.train.get_context().get_world_rank() == 0:
|
|
with create_dict_checkpoint({"iter": i}) as checkpoint:
|
|
train.report(
|
|
metrics={"metric_a": i, "metric_b": -i},
|
|
checkpoint=checkpoint,
|
|
)
|
|
else:
|
|
train.report(metrics={"metric_a": i, "metric_b": -i})
|
|
raise RuntimeError()
|
|
|
|
trainer = TorchTrainer(
|
|
train_loop_per_worker=worker_loop,
|
|
train_loop_config=_PARAM_SPACE,
|
|
scaling_config=ScalingConfig(num_workers=2, use_gpu=False),
|
|
run_config=RunConfig(
|
|
name=configs["EXP_NAME"],
|
|
storage_path=configs["STORAGE_PATH"],
|
|
checkpoint_config=CheckpointConfig(
|
|
num_to_keep=configs["NUM_CHECKPOINTS"],
|
|
checkpoint_score_attribute="metric_a",
|
|
checkpoint_score_order="max",
|
|
),
|
|
),
|
|
)
|
|
return trainer
|
|
|
|
|
|
def build_dummy_tuner(configs):
|
|
return Tuner(
|
|
build_dummy_trainer(configs),
|
|
param_space={"train_loop_config": _PARAM_SPACE},
|
|
tune_config=TuneConfig(num_samples=1),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("storage", ["local", "remote"])
|
|
@pytest.mark.parametrize("mode", ["trainer", "tuner"])
|
|
def test_result_restore(ray_start_4_cpus, tmpdir, mock_s3_bucket_uri, storage, mode):
|
|
NUM_ITERATIONS = 5
|
|
NUM_CHECKPOINTS = 3
|
|
if storage == "local":
|
|
storage_path = str(tmpdir)
|
|
elif storage == "remote":
|
|
storage_path = str(URI(mock_s3_bucket_uri))
|
|
|
|
exp_name = "test_result_restore"
|
|
|
|
configs = {
|
|
"EXP_NAME": exp_name,
|
|
"STORAGE_PATH": storage_path,
|
|
"NUM_ITERATIONS": NUM_ITERATIONS,
|
|
"NUM_CHECKPOINTS": NUM_CHECKPOINTS,
|
|
}
|
|
|
|
if mode == "trainer":
|
|
trainer = build_dummy_trainer(configs)
|
|
with pytest.raises(TrainingFailedError):
|
|
trainer.fit()
|
|
elif mode == "tuner":
|
|
tuner = build_dummy_tuner(configs)
|
|
tuner.fit()
|
|
|
|
# Find the trial directory to restore
|
|
exp_dir = str(URI(storage_path) / exp_name)
|
|
fs, fs_exp_dir = pyarrow.fs.FileSystem.from_uri(exp_dir)
|
|
for item in fs.get_file_info(pyarrow.fs.FileSelector(fs_exp_dir)):
|
|
if item.type == pyarrow.fs.FileType.Directory and item.base_name.startswith(
|
|
"TorchTrainer"
|
|
):
|
|
trial_dir = str(URI(exp_dir) / item.base_name)
|
|
break
|
|
|
|
# [1] Restore from path
|
|
result = Result.from_path(trial_dir)
|
|
|
|
# Check if we restored all checkpoints
|
|
assert result.checkpoint
|
|
assert len(result.best_checkpoints) == NUM_CHECKPOINTS
|
|
|
|
"""
|
|
Top-3 checkpoints with metrics:
|
|
|
|
| iter | metric_a metric_b
|
|
checkpoint_000004 4 4 -4
|
|
checkpoint_000003 3 3 -3
|
|
checkpoint_000002 2 2 -2
|
|
"""
|
|
# Check if the checkpoints bounded with correct metrics
|
|
best_ckpt_a = result.get_best_checkpoint(metric="metric_a", mode="max")
|
|
assert load_dict_checkpoint(best_ckpt_a)["iter"] == NUM_ITERATIONS - 1
|
|
|
|
best_ckpt_b = result.get_best_checkpoint(metric="metric_b", mode="max")
|
|
assert load_dict_checkpoint(best_ckpt_b)["iter"] == NUM_ITERATIONS - NUM_CHECKPOINTS
|
|
|
|
with pytest.raises(RuntimeError, match="Invalid metric name.*"):
|
|
result.get_best_checkpoint(metric="invalid_metric", mode="max")
|
|
|
|
# Check if we properly restored errors
|
|
assert isinstance(result.error, RuntimeError)
|
|
|
|
# Check that the config is properly formatted in the result metrics
|
|
assert result.metrics.get("config") == {"train_loop_config": _PARAM_SPACE}
|
|
|
|
# [2] Restore from path without result.json
|
|
fs.delete_file((URI(trial_dir) / EXPR_RESULT_FILE).path)
|
|
result = Result.from_path(trial_dir)
|
|
|
|
# Do the same checks as above
|
|
assert result.checkpoint
|
|
assert len(result.best_checkpoints) == NUM_CHECKPOINTS
|
|
|
|
best_ckpt_a = result.get_best_checkpoint(metric="metric_a", mode="max")
|
|
assert load_dict_checkpoint(best_ckpt_a)["iter"] == NUM_ITERATIONS - 1
|
|
|
|
best_ckpt_b = result.get_best_checkpoint(metric="metric_b", mode="max")
|
|
assert load_dict_checkpoint(best_ckpt_b)["iter"] == NUM_ITERATIONS - NUM_CHECKPOINTS
|
|
|
|
assert isinstance(result.error, RuntimeError)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|