wehub-resource-sync
526 lines
17 KiB
Python
526 lines
17 KiB
Python
import os
|
|
import sys
|
|
from unittest.mock import patch
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
import ray
|
|
from ray import tune
|
|
from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
|
|
from ray.train.tests.util import mock_storage_context
|
|
from ray.tune import CheckpointConfig, Experiment, PlacementGroupFactory, ResumeConfig
|
|
from ray.tune.execution.tune_controller import TuneController
|
|
from ray.tune.experiment import Trial
|
|
from ray.tune.impl.placeholder import create_resolvers_map, inject_placeholders
|
|
from ray.tune.search import BasicVariantGenerator
|
|
from ray.tune.utils.mock_trainable import (
|
|
MOCK_ERROR_KEY,
|
|
MOCK_TRAINABLE_NAME,
|
|
register_mock_trainable,
|
|
)
|
|
|
|
STORAGE = mock_storage_context()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def register_test_trainable():
|
|
register_mock_trainable()
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def ray_start_4_cpus_2_gpus_extra():
|
|
address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
|
|
yield address_info
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_controller_restore_dataset_references(
|
|
ray_start_4_cpus_2_gpus_extra, resource_manager_cls
|
|
):
|
|
"""Check that references to Ray Datasets are replaced on resume.
|
|
|
|
Legacy test: test_trial_runner_3.py::TrialRunnerTest::
|
|
testSearcherCorrectReferencesAfterRestore
|
|
"""
|
|
|
|
class FakeDataset:
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
config = {
|
|
"param1": {
|
|
"param2": tune.grid_search(
|
|
[FakeDataset("1"), FakeDataset("2"), FakeDataset("3")]
|
|
),
|
|
},
|
|
"param4": tune.sample_from(lambda: 1),
|
|
"param5": tune.sample_from(lambda spec: spec.config["param1"]["param2"]),
|
|
}
|
|
resolvers = create_resolvers_map()
|
|
config = inject_placeholders(config, resolvers)
|
|
|
|
def create_searcher():
|
|
search_alg = BasicVariantGenerator()
|
|
experiment_spec = {
|
|
"run": MOCK_TRAINABLE_NAME,
|
|
"stop": {"training_iteration": 2},
|
|
"config": config,
|
|
}
|
|
experiments = [Experiment.from_json("test", experiment_spec)]
|
|
search_alg.add_configurations(experiments)
|
|
return search_alg
|
|
|
|
searcher = create_searcher()
|
|
|
|
restored_config = {
|
|
"param1": {
|
|
"param2": tune.grid_search(
|
|
[FakeDataset("4"), FakeDataset("5"), FakeDataset("6")]
|
|
),
|
|
},
|
|
"param4": tune.sample_from(lambda: 8),
|
|
"param5": tune.sample_from(lambda spec: spec["config"]["param1"]["param2"]),
|
|
}
|
|
replaced_resolvers = create_resolvers_map()
|
|
inject_placeholders(restored_config, replaced_resolvers)
|
|
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
reuse_actors=False,
|
|
search_alg=searcher,
|
|
placeholder_resolvers=replaced_resolvers,
|
|
checkpoint_period=-1,
|
|
storage=STORAGE,
|
|
)
|
|
|
|
while len(runner.get_trials()) < 3 or any(
|
|
trial.status not in {Trial.RUNNING, Trial.TERMINATED}
|
|
for trial in runner.get_trials()
|
|
):
|
|
runner.step()
|
|
|
|
assert len(runner.get_trials()) == 3, [t.config for t in runner.get_trials()]
|
|
for t in runner.get_trials():
|
|
# Make sure that all the trials carry updated config values.
|
|
assert t.config["param1"]["param2"].name in ["4", "5", "6"]
|
|
assert t.config["param4"] == 8
|
|
assert t.config["param5"].name in ["4", "5", "6"]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_controller_restore_no_error_resume(
|
|
ray_start_4_cpus_2_gpus_extra, resource_manager_cls
|
|
):
|
|
"""Check that `resume=True` does not resume errored trials.
|
|
|
|
Legacy test: test_trial_runner_3.py::TrialRunnerTest::testTrialErrorResumeFalse
|
|
"""
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
storage=STORAGE,
|
|
)
|
|
|
|
kwargs = {
|
|
"stopping_criterion": {"training_iteration": 4},
|
|
"placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 0}]),
|
|
"storage": STORAGE,
|
|
}
|
|
trials = [
|
|
Trial(MOCK_TRAINABLE_NAME, config={MOCK_ERROR_KEY: True}, **kwargs),
|
|
Trial(MOCK_TRAINABLE_NAME, **kwargs),
|
|
Trial(MOCK_TRAINABLE_NAME, **kwargs),
|
|
]
|
|
for t in trials:
|
|
runner.add_trial(t)
|
|
|
|
while not runner.is_finished():
|
|
runner.step()
|
|
|
|
runner.checkpoint(force=True, wait=True)
|
|
|
|
assert trials[0].status == Trial.ERROR
|
|
del runner
|
|
|
|
new_runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
storage=STORAGE,
|
|
resume_config=ResumeConfig(
|
|
unfinished=ResumeConfig.ResumeType.RESUME,
|
|
errored=ResumeConfig.ResumeType.SKIP,
|
|
finished=ResumeConfig.ResumeType.SKIP,
|
|
),
|
|
)
|
|
|
|
assert len(new_runner.get_trials()) == 3
|
|
assert Trial.ERROR in (t.status for t in new_runner.get_trials())
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_controller_restore_error_only_resume(
|
|
ray_start_4_cpus_2_gpus_extra, resource_manager_cls
|
|
):
|
|
"""Check that `resume=ERRORED_ONLY` only resumes errored trials.
|
|
|
|
Legacy test: test_trial_runner_3.py::TrialRunnerTest::testTrialErrorResumeTrue
|
|
"""
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
storage=STORAGE,
|
|
)
|
|
kwargs = {
|
|
"stopping_criterion": {"training_iteration": 4},
|
|
"placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 0}]),
|
|
"storage": STORAGE,
|
|
}
|
|
trials = [
|
|
Trial(MOCK_TRAINABLE_NAME, config={MOCK_ERROR_KEY: True}, **kwargs),
|
|
Trial(MOCK_TRAINABLE_NAME, **kwargs),
|
|
Trial(MOCK_TRAINABLE_NAME, **kwargs),
|
|
]
|
|
for t in trials:
|
|
runner.add_trial(t)
|
|
|
|
while not runner.is_finished():
|
|
runner.step()
|
|
|
|
runner.checkpoint(force=True, wait=True)
|
|
|
|
assert trials[0].status == Trial.ERROR
|
|
del runner
|
|
|
|
new_runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
storage=STORAGE,
|
|
resume_config=ResumeConfig(
|
|
unfinished=ResumeConfig.ResumeType.SKIP,
|
|
errored=ResumeConfig.ResumeType.RESUME,
|
|
finished=ResumeConfig.ResumeType.SKIP,
|
|
),
|
|
)
|
|
|
|
assert len(new_runner.get_trials()) == 3
|
|
assert Trial.ERROR not in (t.status for t in new_runner.get_trials())
|
|
# The below is just a check for standard behavior.
|
|
disable_error = False
|
|
for t in new_runner.get_trials():
|
|
if t.config.get(MOCK_ERROR_KEY):
|
|
t.config[MOCK_ERROR_KEY] = False
|
|
disable_error = True
|
|
assert disable_error
|
|
|
|
while not new_runner.is_finished():
|
|
new_runner.step()
|
|
assert Trial.ERROR not in (t.status for t in new_runner.get_trials())
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_controller_restore_trial_save_restore(
|
|
ray_start_4_cpus_2_gpus_extra, resource_manager_cls
|
|
):
|
|
"""Creates different trials to test runner.checkpoint/restore.
|
|
|
|
Legacy test: test_trial_runner_3.py::TrialRunnerTest::testTrialSaveRestore
|
|
"""
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
checkpoint_period=0,
|
|
storage=STORAGE,
|
|
)
|
|
trials = [
|
|
Trial(
|
|
MOCK_TRAINABLE_NAME,
|
|
trial_id="trial_terminate",
|
|
stopping_criterion={"training_iteration": 1},
|
|
checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
|
|
storage=STORAGE,
|
|
)
|
|
]
|
|
runner.add_trial(trials[0])
|
|
while not runner.is_finished():
|
|
# Start trial, process result, dispatch save and process save.
|
|
runner.step()
|
|
assert trials[0].status == Trial.TERMINATED
|
|
|
|
trials += [
|
|
Trial(
|
|
MOCK_TRAINABLE_NAME,
|
|
trial_id="trial_fail",
|
|
stopping_criterion={"training_iteration": 3},
|
|
checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
|
|
config={MOCK_ERROR_KEY: True},
|
|
storage=STORAGE,
|
|
)
|
|
]
|
|
runner.add_trial(trials[1])
|
|
while not runner.is_finished():
|
|
runner.step()
|
|
assert trials[1].status == Trial.ERROR
|
|
|
|
trials += [
|
|
Trial(
|
|
MOCK_TRAINABLE_NAME,
|
|
trial_id="trial_succ",
|
|
stopping_criterion={"training_iteration": 2},
|
|
checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
|
|
storage=STORAGE,
|
|
)
|
|
]
|
|
runner.add_trial(trials[2])
|
|
|
|
while not trials[2].status == Trial.RUNNING:
|
|
runner.step() # Start trial
|
|
assert len(runner._get_trial_checkpoints()) == 3
|
|
|
|
runner.checkpoint(force=True, wait=True)
|
|
|
|
runner2 = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
storage=STORAGE,
|
|
resume_config=ResumeConfig(
|
|
unfinished=ResumeConfig.ResumeType.RESUME,
|
|
errored=ResumeConfig.ResumeType.SKIP,
|
|
finished=ResumeConfig.ResumeType.SKIP,
|
|
),
|
|
)
|
|
for tid in ["trial_terminate", "trial_fail"]:
|
|
original_trial = runner.get_trial(tid)
|
|
restored_trial = runner2.get_trial(tid)
|
|
assert original_trial.status == restored_trial.status
|
|
|
|
restored_trial = runner2.get_trial("trial_succ")
|
|
assert Trial.PENDING == restored_trial.status
|
|
|
|
while not runner2.is_finished():
|
|
runner2.step()
|
|
assert restored_trial.status == Trial.TERMINATED
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_controller_restore_trial_no_checkpoint_save(
|
|
ray_start_4_cpus_2_gpus_extra, resource_manager_cls
|
|
):
|
|
"""Check that non-checkpointing trials *are* saved.
|
|
|
|
Legacy test: test_trial_runner_3.py::TrialRunnerTest::testTrialNoCheckpointSave
|
|
"""
|
|
with patch.dict(os.environ, {"TUNE_MAX_PENDING_TRIALS_PG": "1"}):
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
checkpoint_period=0,
|
|
storage=STORAGE,
|
|
)
|
|
|
|
runner.add_trial(
|
|
Trial(
|
|
MOCK_TRAINABLE_NAME,
|
|
trial_id="non_checkpoint",
|
|
stopping_criterion={"training_iteration": 2},
|
|
storage=STORAGE,
|
|
)
|
|
)
|
|
|
|
while not all(t.status == Trial.TERMINATED for t in runner.get_trials()):
|
|
runner.step()
|
|
|
|
runner.add_trial(
|
|
Trial(
|
|
MOCK_TRAINABLE_NAME,
|
|
trial_id="checkpoint",
|
|
checkpoint_config=CheckpointConfig(
|
|
checkpoint_at_end=True,
|
|
),
|
|
stopping_criterion={"training_iteration": 2},
|
|
storage=STORAGE,
|
|
)
|
|
)
|
|
|
|
while not all(t.status == Trial.TERMINATED for t in runner.get_trials()):
|
|
runner.step()
|
|
|
|
runner.add_trial(
|
|
Trial(
|
|
MOCK_TRAINABLE_NAME,
|
|
trial_id="pending",
|
|
stopping_criterion={"training_iteration": 2},
|
|
storage=STORAGE,
|
|
)
|
|
)
|
|
|
|
old_trials = runner.get_trials()
|
|
while not old_trials[2].has_reported_at_least_once:
|
|
runner.step()
|
|
|
|
runner.checkpoint(force=True, wait=True)
|
|
|
|
runner2 = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
storage=STORAGE,
|
|
resume_config=ResumeConfig(
|
|
unfinished=ResumeConfig.ResumeType.RESUME,
|
|
errored=ResumeConfig.ResumeType.SKIP,
|
|
finished=ResumeConfig.ResumeType.SKIP,
|
|
),
|
|
)
|
|
new_trials = runner2.get_trials()
|
|
assert len(new_trials) == 3
|
|
assert runner2.get_trial("non_checkpoint").status == Trial.TERMINATED
|
|
assert runner2.get_trial("checkpoint").status == Trial.TERMINATED
|
|
assert runner2.get_trial("pending").status == Trial.PENDING
|
|
assert runner2.get_trial("pending").has_reported_at_least_once
|
|
runner2.step()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_controller_restore_checkpoint_overwrite(
|
|
ray_start_4_cpus_2_gpus_extra, resource_manager_cls
|
|
):
|
|
"""Check that experiment state checkpoint are not overwritten on continue.
|
|
|
|
Legacy test: test_trial_runner_3.py::TrialRunnerTest::testCheckpointOverwrite
|
|
"""
|
|
storage = mock_storage_context()
|
|
|
|
def count_checkpoints(cdir):
|
|
return sum(
|
|
(fname.startswith("experiment_state") and fname.endswith(".json"))
|
|
for fname in os.listdir(cdir)
|
|
)
|
|
|
|
tmpdir = storage.experiment_driver_staging_path
|
|
# The Trial `local_dir` must match the TrialRunner `local_checkpoint_dir`
|
|
# to match the directory structure assumed by `TrialRunner.resume`.
|
|
# See `test_trial_runner2.TrialRunnerTest2.testPauseResumeCheckpointCount`
|
|
# for more details.
|
|
trial = Trial(
|
|
MOCK_TRAINABLE_NAME,
|
|
checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
|
|
storage=storage,
|
|
)
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
storage=storage,
|
|
checkpoint_period=0,
|
|
)
|
|
runner.add_trial(trial)
|
|
while not trial.status == Trial.RUNNING:
|
|
runner.step()
|
|
# force checkpoint
|
|
runner.checkpoint(force=True, wait=True)
|
|
# Only one experiment state file
|
|
assert count_checkpoints(tmpdir) == 1
|
|
|
|
runner2 = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
storage=storage,
|
|
resume_config=ResumeConfig(
|
|
unfinished=ResumeConfig.ResumeType.RESUME,
|
|
errored=ResumeConfig.ResumeType.SKIP,
|
|
finished=ResumeConfig.ResumeType.SKIP,
|
|
),
|
|
)
|
|
trial = runner2.get_trials()[0]
|
|
while not trial.status == Trial.RUNNING:
|
|
runner2.step()
|
|
# After resume, we have a new experiment state file in the directory
|
|
assert count_checkpoints(tmpdir) == 2
|
|
|
|
runner2.checkpoint()
|
|
assert count_checkpoints(tmpdir) == 2
|
|
|
|
|
|
@pytest.mark.skip("TODO(justinvyu): Data lineage serialization context is broken.")
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_controller_restore_with_dataset(
|
|
ray_start_4_cpus_2_gpus_extra, resource_manager_cls, tmp_path
|
|
):
|
|
"""Test trial runner checkpointing where trials contain Datasets.
|
|
When possible, a dataset plan should be saved (for read_* APIs).
|
|
See `Dataset.serialize_lineage` for more information.
|
|
|
|
If a dataset cannot be serialized, an experiment checkpoint
|
|
should still be created. Users can pass in the dataset again by
|
|
re-specifying the `param_space`.
|
|
|
|
Legacy test: test_trial_runner_3.py::TrialRunnerTest::
|
|
testExperimentCheckpointWithDatasets
|
|
"""
|
|
# Save some test data to load
|
|
data_filepath = os.path.join(tmp_path, "test.csv")
|
|
pd.DataFrame({"x": list(range(10))}).to_csv(data_filepath)
|
|
|
|
def create_trial_config():
|
|
return {
|
|
"datasets": {
|
|
"with_lineage": ray.data.read_csv(data_filepath),
|
|
"no_lineage": ray.data.from_items([{"x": i} for i in range(10)]),
|
|
}
|
|
}
|
|
|
|
resolvers = create_resolvers_map()
|
|
config_with_placeholders = inject_placeholders(create_trial_config(), resolvers)
|
|
trial = Trial(
|
|
MOCK_TRAINABLE_NAME,
|
|
config=config_with_placeholders,
|
|
storage=STORAGE,
|
|
)
|
|
trial.init_local_path()
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
storage=STORAGE,
|
|
placeholder_resolvers=resolvers,
|
|
)
|
|
runner.add_trial(trial)
|
|
# Req: TrialRunner checkpointing shouldn't error
|
|
runner.checkpoint(force=True, wait=True)
|
|
|
|
# Manually clear all block refs that may have been created
|
|
ray.shutdown()
|
|
ray.init(num_cpus=2)
|
|
|
|
register_mock_trainable()
|
|
new_runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
storage=STORAGE,
|
|
)
|
|
new_runner.resume(resume_config=ResumeConfig())
|
|
[loaded_trial] = new_runner.get_trials()
|
|
loaded_datasets = loaded_trial.config["datasets"]
|
|
|
|
# Req: The deserialized dataset (w/ lineage) should be usable.
|
|
assert [el["x"] for el in loaded_datasets["with_lineage"].take()] == list(range(10))
|
|
|
|
replaced_resolvers = create_resolvers_map()
|
|
inject_placeholders(create_trial_config(), replaced_resolvers)
|
|
|
|
respecified_config_runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
storage=STORAGE,
|
|
placeholder_resolvers=replaced_resolvers,
|
|
)
|
|
respecified_config_runner.resume(resume_config=ResumeConfig())
|
|
[loaded_trial] = respecified_config_runner.get_trials()
|
|
ray_ds_no_lineage = loaded_trial.config["datasets"]["no_lineage"]
|
|
|
|
# Req: The dataset (w/o lineage) can be re-specified and is usable after.
|
|
assert [el["x"] for el in ray_ds_no_lineage.take()] == list(range(10))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|