chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,526 @@
|
||||
import inspect
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray import logger, tune
|
||||
from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
|
||||
from ray.tune import CheckpointConfig, Trainable, register_trainable, run_experiments
|
||||
from ray.tune.error import TuneError
|
||||
from ray.tune.result_grid import ResultGrid
|
||||
from ray.tune.schedulers.trial_scheduler import FIFOScheduler, TrialScheduler
|
||||
from ray.tune.tune import _check_mixin
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_start_1_cpu():
|
||||
address_info = ray.init(num_cpus=1)
|
||||
os.environ["TUNE_STATE_REFRESH_PERIOD"] = "0.1"
|
||||
yield address_info
|
||||
ray.shutdown()
|
||||
os.environ.pop("TUNE_STATE_REFRESH_PERIOD", None)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_start_2_cpus():
|
||||
address_info = ray.init(num_cpus=2)
|
||||
yield address_info
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_start_4_cpus_extra():
|
||||
address_info = ray.init(num_cpus=4, resources={"extra": 4})
|
||||
yield address_info
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
class FrequentPausesScheduler(FIFOScheduler):
|
||||
def on_trial_result(self, tune_controller, trial, result):
|
||||
return TrialScheduler.PAUSE
|
||||
|
||||
|
||||
class MyResettableClass(Trainable):
|
||||
def setup(self, config):
|
||||
self.config = config
|
||||
self.num_resets = 0
|
||||
self.iter = 0
|
||||
self.msg = config.get("message", None)
|
||||
self.sleep = int(config.get("sleep", 0))
|
||||
self.fail = config.get("fail", False)
|
||||
|
||||
def step(self):
|
||||
self.iter += 1
|
||||
|
||||
if self.msg:
|
||||
print("PRINT_STDOUT: {}".format(self.msg))
|
||||
print("PRINT_STDERR: {}".format(self.msg), file=sys.stderr)
|
||||
logger.info("LOG_STDERR: {}".format(self.msg))
|
||||
|
||||
if self.fail:
|
||||
raise RuntimeError("Failing")
|
||||
|
||||
if self.sleep:
|
||||
time.sleep(self.sleep)
|
||||
|
||||
return {
|
||||
"id": self.config.get("id", -1),
|
||||
"num_resets": self.num_resets,
|
||||
"done": self.iter > 1,
|
||||
"iter": self.iter,
|
||||
}
|
||||
|
||||
def save_checkpoint(self, chkpt_dir):
|
||||
return {"iter": self.iter}
|
||||
|
||||
def load_checkpoint(self, item):
|
||||
self.iter = item["iter"]
|
||||
|
||||
def reset_config(self, new_config):
|
||||
if "fake_reset_not_supported" in self.config:
|
||||
return False
|
||||
self.num_resets += 1
|
||||
self.iter = 0
|
||||
self.msg = new_config.get("message", None)
|
||||
self.fail = new_config.get("fail", False)
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def default_resource_request(cls, config):
|
||||
required_resources = config.get("required_resources", None)
|
||||
if required_resources:
|
||||
return required_resources
|
||||
return None
|
||||
|
||||
|
||||
def train_fn(config):
|
||||
# Determine whether or not we reset to a new trial
|
||||
marker_dir = config.get("marker_dir")
|
||||
num_resets = 0
|
||||
marker = Path(marker_dir) / f"{os.getpid()}.txt"
|
||||
if marker.exists():
|
||||
num_resets = int(marker.read_text()) + 1
|
||||
|
||||
checkpoint = tune.get_checkpoint()
|
||||
it = load_dict_checkpoint(checkpoint)["iter"] if checkpoint else 0
|
||||
|
||||
msg = config.get("message", None)
|
||||
sleep = int(config.get("sleep", 0))
|
||||
fail = config.get("fail", False)
|
||||
|
||||
while it < 2:
|
||||
it += 1
|
||||
if msg:
|
||||
print("PRINT_STDOUT: {}".format(msg))
|
||||
print("PRINT_STDERR: {}".format(msg), file=sys.stderr)
|
||||
logger.info("LOG_STDERR: {}".format(msg))
|
||||
|
||||
if fail:
|
||||
raise RuntimeError("Failing")
|
||||
|
||||
# Dump the current config
|
||||
marker.write_text(str(num_resets))
|
||||
|
||||
if sleep:
|
||||
time.sleep(sleep)
|
||||
|
||||
metrics = {
|
||||
"id": config.get("id", 0),
|
||||
"num_resets": num_resets,
|
||||
"iter": it,
|
||||
"done": it > 1,
|
||||
}
|
||||
if config.get("save_checkpoint", True):
|
||||
with create_dict_checkpoint({"iter": it}) as checkpoint:
|
||||
tune.report(metrics, checkpoint=checkpoint)
|
||||
else:
|
||||
tune.report(metrics, checkpoint=checkpoint)
|
||||
|
||||
|
||||
@pytest.fixture(params=["function", "class"])
|
||||
def trainable(request):
|
||||
"""Fixture that sets up a function/class trainable for testing.
|
||||
Make sure this fixture comes BEFORE the ray.init fixture in the arguments
|
||||
so that the env var is propagated to workers correctly."""
|
||||
trainable_type = request.param
|
||||
if trainable_type == "function":
|
||||
yield train_fn
|
||||
elif trainable_type == "class":
|
||||
yield MyResettableClass
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def _run_trials_with_frequent_pauses(trainable, reuse=False, **kwargs):
|
||||
tempdir = tempfile.mkdtemp()
|
||||
marker_dir = Path(tempdir)
|
||||
analysis = tune.run(
|
||||
trainable,
|
||||
num_samples=1,
|
||||
config={
|
||||
"id": tune.grid_search([0, 1, 2, 3]),
|
||||
"marker_dir": marker_dir,
|
||||
},
|
||||
reuse_actors=reuse,
|
||||
scheduler=FrequentPausesScheduler(),
|
||||
verbose=0,
|
||||
**kwargs,
|
||||
)
|
||||
return analysis
|
||||
|
||||
|
||||
def test_trial_reuse_disabled(trainable, ray_start_1_cpu):
|
||||
"""Test that reuse=False disables actor re-use.
|
||||
|
||||
Setup: Pass `reuse_actors=False` to tune.run()
|
||||
|
||||
We assert the `num_resets` of each trainable class to be 0 (no reuse).
|
||||
"""
|
||||
analysis = _run_trials_with_frequent_pauses(trainable, reuse=False)
|
||||
trials = analysis.trials
|
||||
assert [t.last_result["id"] for t in trials] == [0, 1, 2, 3]
|
||||
assert [t.last_result["iter"] for t in trials] == [2, 2, 2, 2]
|
||||
assert [t.last_result["num_resets"] for t in trials] == [0, 0, 0, 0]
|
||||
|
||||
|
||||
def test_trial_reuse_enabled(trainable, ray_start_1_cpu):
|
||||
"""Test that reuse=True enables actor re-use.
|
||||
|
||||
Setup: Pass `reuse_actors=True` to tune.run()
|
||||
|
||||
We assert the `num_resets` of each trainable class to be 4, 5, 6, and 7,
|
||||
respectively:
|
||||
|
||||
- Each trial runs for 2 iterations
|
||||
- Only one trial can run at a time
|
||||
- After each iteration, trials are paused and actors cached for reuse
|
||||
- Thus, the first trial finishes after 4 resets, the second after 5, etc.
|
||||
"""
|
||||
analysis = _run_trials_with_frequent_pauses(trainable, reuse=True)
|
||||
trials = analysis.trials
|
||||
assert [t.last_result["id"] for t in trials] == [0, 1, 2, 3]
|
||||
assert [t.last_result["iter"] for t in trials] == [2, 2, 2, 2]
|
||||
assert [t.last_result["num_resets"] for t in trials] == [4, 5, 6, 7]
|
||||
|
||||
|
||||
def test_trial_reuse_with_failing(trainable, ray_start_1_cpu, tmp_path):
|
||||
"""Test that failing actors won't be reused.
|
||||
|
||||
- 1 trial can run at a time
|
||||
- Some trials are failing
|
||||
- We assert that trials after failing trials are scheduled on fresh actors
|
||||
(num_resets = 0)
|
||||
- We assert that trials after successful trials are schedule on reused actors
|
||||
(num_reset = last_num_resets + 1)
|
||||
"""
|
||||
fail = [False, True, False, False, True, True, False, False, False]
|
||||
trials = tune.run(
|
||||
trainable,
|
||||
reuse_actors=True,
|
||||
config={
|
||||
"id": tune.grid_search(list(range(9))),
|
||||
"fail": tune.sample_from(lambda config: fail[config["id"]]),
|
||||
"marker_dir": tmp_path,
|
||||
},
|
||||
raise_on_failed_trial=False,
|
||||
).trials
|
||||
|
||||
assert [t.last_result.get("iter") for t in trials] == [
|
||||
2,
|
||||
None,
|
||||
2,
|
||||
2,
|
||||
None,
|
||||
None,
|
||||
2,
|
||||
2,
|
||||
2,
|
||||
]
|
||||
assert [t.last_result.get("num_resets") for t in trials] == [
|
||||
0,
|
||||
None,
|
||||
0,
|
||||
1,
|
||||
None,
|
||||
None,
|
||||
0,
|
||||
1,
|
||||
2,
|
||||
]
|
||||
|
||||
|
||||
def test_reuse_enabled_error(ray_start_1_cpu):
|
||||
"""Test that a class without reset() enabled throws an error on actor reuse."""
|
||||
with pytest.raises(TuneError):
|
||||
run_experiments(
|
||||
{
|
||||
"foo": {
|
||||
"run": MyResettableClass,
|
||||
"max_failures": 1,
|
||||
"num_samples": 1,
|
||||
"config": {
|
||||
"id": tune.grid_search([0, 1, 2, 3]),
|
||||
"fake_reset_not_supported": True,
|
||||
},
|
||||
}
|
||||
},
|
||||
reuse_actors=True,
|
||||
scheduler=FrequentPausesScheduler(),
|
||||
)
|
||||
|
||||
|
||||
def test_trial_reuse_log_to_file(trainable, ray_start_1_cpu, tmp_path):
|
||||
"""Check that log outputs from trainables are correctly stored with actor reuse.
|
||||
|
||||
We run two trials with actor reuse. When the actor is reused, we expect
|
||||
the log output to be written to the log file of the new trial - i.e. we expect
|
||||
that the old trial logfile is closed and a new one is open.
|
||||
"""
|
||||
register_trainable("foo2", trainable)
|
||||
|
||||
messages = ["First", "Second"]
|
||||
# Log to default files
|
||||
[trial1, trial2] = tune.run(
|
||||
"foo2",
|
||||
config={
|
||||
"id": tune.grid_search(list(range(2))),
|
||||
"message": tune.sample_from(lambda config: messages[config["id"]]),
|
||||
"marker_dir": tmp_path,
|
||||
},
|
||||
log_to_file=True,
|
||||
scheduler=FrequentPausesScheduler(),
|
||||
reuse_actors=True,
|
||||
).trials
|
||||
|
||||
def get_trial_logfiles(trial):
|
||||
return (
|
||||
os.path.join(trial.storage.trial_working_directory, "stdout"),
|
||||
os.path.join(trial.storage.trial_working_directory, "stderr"),
|
||||
)
|
||||
|
||||
# Check trial 1
|
||||
assert trial1.last_result["num_resets"] == 2
|
||||
[stdout, stderr] = get_trial_logfiles(trial1)
|
||||
assert os.path.exists(stdout)
|
||||
assert os.path.exists(stderr)
|
||||
|
||||
# We expect that only "First" output is found in the first trial output
|
||||
with open(stdout, "rt") as fp:
|
||||
content = fp.read()
|
||||
assert "PRINT_STDOUT: First" in content
|
||||
assert "PRINT_STDOUT: Second" not in content
|
||||
with open(stderr, "rt") as fp:
|
||||
content = fp.read()
|
||||
assert "PRINT_STDERR: First" in content
|
||||
assert "LOG_STDERR: First" in content
|
||||
assert "PRINT_STDERR: Second" not in content
|
||||
assert "LOG_STDERR: Second" not in content
|
||||
|
||||
# Check trial 2
|
||||
assert trial2.last_result["num_resets"] == 3
|
||||
[stdout, stderr] = get_trial_logfiles(trial2)
|
||||
assert os.path.exists(stdout)
|
||||
assert os.path.exists(stderr)
|
||||
|
||||
# We expect that only "Second" output is found in the first trial output
|
||||
with open(stdout, "rt") as fp:
|
||||
content = fp.read()
|
||||
assert "PRINT_STDOUT: Second" in content
|
||||
assert "PRINT_STDOUT: First" not in content
|
||||
with open(stderr, "rt") as fp:
|
||||
content = fp.read()
|
||||
assert "PRINT_STDERR: Second" in content
|
||||
assert "LOG_STDERR: Second" in content
|
||||
assert "PRINT_STDERR: First" not in content
|
||||
assert "LOG_STDERR: First" not in content
|
||||
|
||||
|
||||
def test_multi_trial_reuse(trainable, ray_start_4_cpus_extra, monkeypatch, tmp_path):
|
||||
"""Test that actors from multiple trials running in parallel will be reused.
|
||||
|
||||
- 2 trials can run at the same time
|
||||
- Trial 3 will be scheduled after trial 1 succeeded, so will reuse actor
|
||||
- Trial 4 will be scheduled after trial 2 succeeded, so will reuse actor
|
||||
"""
|
||||
monkeypatch.setenv("TUNE_MAX_PENDING_TRIALS_PG", "2")
|
||||
|
||||
register_trainable("foo2", trainable)
|
||||
|
||||
messages = ["First", "Second", "Third", "Fourth"]
|
||||
# We sleep here for one second so that the third actor
|
||||
# does not finish training before the fourth can be scheduled.
|
||||
# This helps ensure that both remote runners are re-used and
|
||||
# not just one.
|
||||
[trial1, trial2, trial3, trial4] = tune.run(
|
||||
"foo2",
|
||||
config={
|
||||
"id": tune.grid_search(list(range(4))),
|
||||
"message": tune.sample_from(lambda config: messages[config["id"]]),
|
||||
"marker_dir": tmp_path,
|
||||
"sleep": 2,
|
||||
},
|
||||
reuse_actors=True,
|
||||
resources_per_trial={"cpu": 2},
|
||||
).trials
|
||||
|
||||
assert trial3.last_result["num_resets"] == 1
|
||||
assert trial4.last_result["num_resets"] == 1
|
||||
|
||||
|
||||
def test_multi_trial_reuse_with_failing(
|
||||
trainable, ray_start_4_cpus_extra, monkeypatch, tmp_path
|
||||
):
|
||||
"""Test that failing trial's actors are not reused.
|
||||
|
||||
- 2 trials can run at the same time
|
||||
- Trial 1 succeeds, trial 2 fails
|
||||
- Trial 3 will be scheduled after trial 2 failed, so won't reuse actor
|
||||
- Trial 4 will be scheduled after trial 1 succeeded, so will reuse actor
|
||||
"""
|
||||
monkeypatch.setenv("TUNE_MAX_PENDING_TRIALS_PG", "2")
|
||||
|
||||
register_trainable("foo2", trainable)
|
||||
|
||||
[trial1, trial2, trial3, trial4] = tune.run(
|
||||
"foo2",
|
||||
config={
|
||||
"fail": tune.grid_search([False, True, False, False]),
|
||||
"marker_dir": tmp_path,
|
||||
"sleep": 2,
|
||||
},
|
||||
reuse_actors=True,
|
||||
resources_per_trial={"cpu": 2},
|
||||
raise_on_failed_trial=False,
|
||||
).trials
|
||||
|
||||
assert trial1.last_result["num_resets"] == 0
|
||||
assert trial3.last_result["num_resets"] == 0
|
||||
assert trial4.last_result["num_resets"] == 1
|
||||
|
||||
|
||||
def test_multi_trial_reuse_one_by_one(trainable, ray_start_4_cpus_extra, tmp_path):
|
||||
"""Test that we still reuse actors even if we run with concurrency = 1.
|
||||
|
||||
- Run 6 trials, but only 1 concurrent at the time
|
||||
- This means there won't be any PENDING trials until the trial completed
|
||||
- We still want to reuse actors
|
||||
|
||||
"""
|
||||
register_trainable("foo2", trainable)
|
||||
|
||||
trials = tune.run(
|
||||
"foo2",
|
||||
config={"id": -1, "marker_dir": tmp_path},
|
||||
reuse_actors=True,
|
||||
num_samples=6,
|
||||
max_concurrent_trials=1,
|
||||
).trials
|
||||
|
||||
assert sorted([t.last_result["num_resets"] for t in trials]) == [0, 1, 2, 3, 4, 5]
|
||||
|
||||
|
||||
def test_multi_trial_reuse_heterogeneous(ray_start_4_cpus_extra):
|
||||
"""Test that actors with heterogeneous resource requirements are reused efficiently.
|
||||
|
||||
- Run 6 trials in total
|
||||
- Only 1 trial can run at the same time
|
||||
- Trials 1 and 6, 2 and 4, and 3 and 5, have the same resource request, respectively
|
||||
- Assert that trials 4, 5, and 6 re-use their respective previous actors
|
||||
|
||||
"""
|
||||
# We need to set this to 6 so that all trials will be scheduled and actors will
|
||||
# be correctly cached.
|
||||
os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "6"
|
||||
|
||||
register_trainable("foo2", MyResettableClass)
|
||||
|
||||
trials = tune.run(
|
||||
"foo2",
|
||||
config={
|
||||
# The extra resources are selected so that only any 1 placement group
|
||||
# can be scheduled at the same time at all times (to force sequential
|
||||
# execution of trials)
|
||||
"required_resources": tune.grid_search(
|
||||
[
|
||||
{"cpu": 4, "custom_resources": {"extra": 4}},
|
||||
{"cpu": 2, "custom_resources": {"extra": 4}},
|
||||
{"cpu": 1, "custom_resources": {"extra": 4}},
|
||||
{"cpu": 2, "custom_resources": {"extra": 4}},
|
||||
{"cpu": 1, "custom_resources": {"extra": 4}},
|
||||
{"cpu": 4, "custom_resources": {"extra": 4}},
|
||||
]
|
||||
),
|
||||
"id": -1,
|
||||
},
|
||||
reuse_actors=True,
|
||||
).trials
|
||||
|
||||
# Actors may be re-used in a different order as the staged_trials set is unsorted
|
||||
assert sorted([t.last_result["num_resets"] for t in trials]) == [0, 0, 0, 1, 1, 1]
|
||||
|
||||
|
||||
def test_detect_reuse_mixins():
|
||||
class DummyMixin:
|
||||
pass
|
||||
|
||||
def dummy_mixin(func: Callable):
|
||||
func.__mixins__ = (DummyMixin,)
|
||||
return func
|
||||
|
||||
def train_fn(config):
|
||||
pass
|
||||
|
||||
assert not _check_mixin(train_fn)
|
||||
assert _check_mixin(dummy_mixin(train_fn))
|
||||
|
||||
class MyTrainable(Trainable):
|
||||
pass
|
||||
|
||||
assert not _check_mixin(MyTrainable)
|
||||
assert _check_mixin(dummy_mixin(MyTrainable))
|
||||
|
||||
|
||||
def test_remote_trial_dir_with_reuse_actors(trainable, ray_start_2_cpus, tmp_path):
|
||||
"""Check that the trainable has its remote directory set to the right
|
||||
location, when new trials get swapped in on actor reuse.
|
||||
Each trial runs for 2 iterations, with checkpoint_frequency=1, so each
|
||||
remote trial dir should have 2 checkpoints.
|
||||
"""
|
||||
tmp_target = str(tmp_path / "upload_dir")
|
||||
exp_name = "remote_trial_dir_update_on_actor_reuse"
|
||||
|
||||
def get_remote_trial_dir(trial_id: int):
|
||||
return os.path.join(tmp_target, exp_name, str(trial_id))
|
||||
|
||||
analysis = _run_trials_with_frequent_pauses(
|
||||
trainable,
|
||||
reuse=True,
|
||||
max_concurrent_trials=2,
|
||||
name=exp_name,
|
||||
storage_path=f"file://{tmp_target}",
|
||||
trial_dirname_creator=lambda t: str(t.config.get("id")),
|
||||
checkpoint_config=CheckpointConfig(
|
||||
checkpoint_frequency=1 if inspect.isclass(trainable) else 0
|
||||
),
|
||||
)
|
||||
result_grid = ResultGrid(analysis)
|
||||
assert not result_grid.errors
|
||||
|
||||
# Check that each remote trial dir has 2 checkpoints.
|
||||
for result in result_grid:
|
||||
trial_id = result.config["id"]
|
||||
remote_dir = get_remote_trial_dir(trial_id)
|
||||
num_checkpoints = len(
|
||||
[file for file in os.listdir(remote_dir) if file.startswith("checkpoint_")]
|
||||
)
|
||||
assert num_checkpoints == 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
Reference in New Issue
Block a user