mirror of
https://github.com/wassname/ray.git
synced 2026-07-12 13:43:34 +08:00
[tune/sgd] Document func_trainable and add checkpoint context (#9739)
Co-authored-by: krfricke <krfricke@users.noreply.github.com> Co-authored-by: Amog Kamsetty <amogkam@users.noreply.github.com>
This commit is contained in:
Richard Liaw
@@ -19,7 +19,7 @@ class FunctionApiTest(unittest.TestCase):
|
||||
_register_all() # re-register the evicted objects
|
||||
|
||||
def testFunctionNoCheckpointing(self):
|
||||
def train(config, checkpoint=None):
|
||||
def train(config, checkpoint_dir=None):
|
||||
for i in range(10):
|
||||
tune.report(test=i)
|
||||
|
||||
@@ -40,14 +40,13 @@ class FunctionApiTest(unittest.TestCase):
|
||||
def testFunctionRecurringSave(self):
|
||||
"""This tests that save and restore are commutative."""
|
||||
|
||||
def train(config, checkpoint=None):
|
||||
def train(config, checkpoint_dir=None):
|
||||
for step in range(10):
|
||||
if step % 3 == 0:
|
||||
checkpoint_dir = tune.make_checkpoint_dir(step=step)
|
||||
path = os.path.join(checkpoint_dir, "checkpoint")
|
||||
with open(path, "w") as f:
|
||||
f.write(json.dumps({"step": step}))
|
||||
tune.save_checkpoint(path)
|
||||
with tune.checkpoint_dir(step=step) as checkpoint_dir:
|
||||
path = os.path.join(checkpoint_dir, "checkpoint")
|
||||
with open(path, "w") as f:
|
||||
f.write(json.dumps({"step": step}))
|
||||
tune.report(test=step)
|
||||
|
||||
wrapped = wrap_function(train)
|
||||
@@ -65,49 +64,58 @@ class FunctionApiTest(unittest.TestCase):
|
||||
new_trainable2.stop()
|
||||
|
||||
def testCheckpointFunctionAtEnd(self):
|
||||
def train(config, checkpoint=False):
|
||||
def train(config, checkpoint_dir=False):
|
||||
for i in range(10):
|
||||
tune.report(test=i)
|
||||
checkpoint_dir = tune.make_checkpoint_dir(step=10)
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "hello")
|
||||
with open(checkpoint_path, "w") as f:
|
||||
f.write("hello")
|
||||
tune.save_checkpoint(checkpoint_path)
|
||||
|
||||
[trial] = tune.run(train).trials
|
||||
assert "hello" in trial.checkpoint.value
|
||||
|
||||
def testVariousCheckpointFunctionAtEnd(self):
|
||||
def train(config, checkpoint=False):
|
||||
for i in range(10):
|
||||
checkpoint_dir = tune.make_checkpoint_dir()
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "hello")
|
||||
with tune.checkpoint_dir(step=10) as checkpoint_dir:
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "ckpt.log")
|
||||
with open(checkpoint_path, "w") as f:
|
||||
f.write("hello")
|
||||
tune.save_checkpoint(checkpoint_path)
|
||||
|
||||
[trial] = tune.run(train).trials
|
||||
assert os.path.exists(os.path.join(trial.checkpoint.value, "ckpt.log"))
|
||||
|
||||
def testCheckpointFunctionAtEndContext(self):
|
||||
def train(config, checkpoint_dir=False):
|
||||
for i in range(10):
|
||||
tune.report(test=i)
|
||||
checkpoint_dir = tune.make_checkpoint_dir()
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "goodbye")
|
||||
with open(checkpoint_path, "w") as f:
|
||||
f.write("goodbye")
|
||||
tune.save_checkpoint(checkpoint_path)
|
||||
with tune.checkpoint_dir(step=10) as checkpoint_dir:
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "ckpt.log")
|
||||
with open(checkpoint_path, "w") as f:
|
||||
f.write("hello")
|
||||
|
||||
[trial] = tune.run(train).trials
|
||||
assert os.path.exists(os.path.join(trial.checkpoint.value, "ckpt.log"))
|
||||
|
||||
def testVariousCheckpointFunctionAtEnd(self):
|
||||
def train(config, checkpoint_dir=False):
|
||||
for i in range(10):
|
||||
with tune.checkpoint_dir() as checkpoint_dir:
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "ckpt.log")
|
||||
with open(checkpoint_path, "w") as f:
|
||||
f.write("hello")
|
||||
tune.report(test=i)
|
||||
with tune.checkpoint_dir() as checkpoint_dir:
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "ckpt.log2")
|
||||
with open(checkpoint_path, "w") as f:
|
||||
f.write("goodbye")
|
||||
|
||||
[trial] = tune.run(train, keep_checkpoints_num=3).trials
|
||||
assert "goodbye" in trial.checkpoint.value
|
||||
assert os.path.exists(
|
||||
os.path.join(trial.checkpoint.value, "ckpt.log2"))
|
||||
|
||||
def testReuseCheckpoint(self):
|
||||
def train(config, checkpoint=False):
|
||||
def train(config, checkpoint_dir=None):
|
||||
itr = 0
|
||||
if checkpoint:
|
||||
with open(checkpoint, "r") as f:
|
||||
if checkpoint_dir:
|
||||
with open(os.path.join(checkpoint_dir, "ckpt.log"), "r") as f:
|
||||
itr = int(f.read()) + 1
|
||||
|
||||
for i in range(itr, config["max_iter"]):
|
||||
checkpoint_dir = tune.make_checkpoint_dir(step=i)
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "goodbye")
|
||||
with open(checkpoint_path, "w") as f:
|
||||
f.write(str(i))
|
||||
tune.save_checkpoint(checkpoint_path)
|
||||
with tune.checkpoint_dir(step=i) as checkpoint_dir:
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "ckpt.log")
|
||||
with open(checkpoint_path, "w") as f:
|
||||
f.write(str(i))
|
||||
tune.report(test=i, training_iteration=i)
|
||||
|
||||
[trial] = tune.run(
|
||||
@@ -117,51 +125,49 @@ class FunctionApiTest(unittest.TestCase):
|
||||
},
|
||||
).trials
|
||||
last_ckpt = trial.checkpoint.value
|
||||
assert "goodbye" in last_ckpt
|
||||
assert os.path.exists(os.path.join(trial.checkpoint.value, "ckpt.log"))
|
||||
analysis = tune.run(train, config={"max_iter": 10}, restore=last_ckpt)
|
||||
trial_dfs = list(analysis.trial_dataframes.values())
|
||||
assert len(trial_dfs[0]["training_iteration"]) == 5
|
||||
|
||||
def testRetry(self):
|
||||
def train(config, checkpoint=None):
|
||||
restored = bool(checkpoint)
|
||||
def train(config, checkpoint_dir=None):
|
||||
restored = bool(checkpoint_dir)
|
||||
itr = 0
|
||||
if checkpoint:
|
||||
with open(checkpoint, "r") as f:
|
||||
if checkpoint_dir:
|
||||
with open(os.path.join(checkpoint_dir, "ckpt.log"), "r") as f:
|
||||
itr = int(f.read()) + 1
|
||||
|
||||
for i in range(itr, 10):
|
||||
if i == 5 and not restored:
|
||||
raise Exception("try to fail me")
|
||||
checkpoint_dir = tune.make_checkpoint_dir(step=i)
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "goodbye")
|
||||
with open(checkpoint_path, "w") as f:
|
||||
f.write(str(i))
|
||||
tune.save_checkpoint(checkpoint_path)
|
||||
with tune.checkpoint_dir(step=i) as checkpoint_dir:
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "ckpt.log")
|
||||
with open(checkpoint_path, "w") as f:
|
||||
f.write(str(i))
|
||||
tune.report(test=i, training_iteration=i)
|
||||
|
||||
analysis = tune.run(train, max_failures=3)
|
||||
last_ckpt = analysis.trials[0].checkpoint.value
|
||||
assert "goodbye" in last_ckpt
|
||||
assert os.path.exists(os.path.join(last_ckpt, "ckpt.log"))
|
||||
trial_dfs = list(analysis.trial_dataframes.values())
|
||||
assert len(trial_dfs[0]["training_iteration"]) == 10
|
||||
|
||||
def testBlankCheckpoint(self):
|
||||
def train(config, checkpoint=None):
|
||||
restored = bool(checkpoint)
|
||||
def train(config, checkpoint_dir=None):
|
||||
restored = bool(checkpoint_dir)
|
||||
itr = 0
|
||||
if checkpoint:
|
||||
with open(checkpoint, "r") as f:
|
||||
if checkpoint_dir:
|
||||
with open(os.path.join(checkpoint_dir, "ckpt.log"), "r") as f:
|
||||
itr = int(f.read()) + 1
|
||||
|
||||
for i in range(itr, 10):
|
||||
if i == 5 and not restored:
|
||||
raise Exception("try to fail me")
|
||||
checkpoint_dir = tune.make_checkpoint_dir()
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "goodbye")
|
||||
with open(checkpoint_path, "w") as f:
|
||||
f.write(str(i))
|
||||
tune.save_checkpoint(checkpoint_path)
|
||||
with tune.checkpoint_dir() as checkpoint_dir:
|
||||
checkpoint_path = os.path.join(checkpoint_dir, "ckpt.log")
|
||||
with open(checkpoint_path, "w") as f:
|
||||
f.write(str(i))
|
||||
tune.report(test=i, training_iteration=i)
|
||||
|
||||
analysis = tune.run(train, max_failures=3)
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
import os
|
||||
import pytest
|
||||
from unittest.mock import patch
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
import ray
|
||||
from ray import tune
|
||||
from ray.tune.integration.torch import (DistributedTrainableCreator,
|
||||
distributed_checkpoint_dir,
|
||||
_train_simple, _train_check_global)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_start_2_cpus():
|
||||
address_info = ray.init(num_cpus=2)
|
||||
yield address_info
|
||||
# The code after the yield will run as teardown code.
|
||||
ray.shutdown()
|
||||
# Ensure that tests don't ALL fail
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
@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()
|
||||
# Ensure that tests don't ALL fail
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def test_single_step(ray_start_2_cpus): # noqa: F811
|
||||
trainable_cls = DistributedTrainableCreator(_train_simple, num_workers=2)
|
||||
trainer = trainable_cls()
|
||||
trainer.train()
|
||||
trainer.stop()
|
||||
|
||||
|
||||
def test_step_after_completion(ray_start_2_cpus): # noqa: F811
|
||||
trainable_cls = DistributedTrainableCreator(_train_simple, num_workers=2)
|
||||
trainer = trainable_cls(config={"epochs": 1})
|
||||
with pytest.raises(RuntimeError):
|
||||
for i in range(10):
|
||||
trainer.train()
|
||||
|
||||
|
||||
def test_validation(ray_start_2_cpus): # noqa: F811
|
||||
def bad_func(a, b, c):
|
||||
return 1
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
DistributedTrainableCreator(bad_func, num_workers=2)
|
||||
|
||||
|
||||
def test_set_global(ray_start_2_cpus): # noqa: F811
|
||||
trainable_cls = DistributedTrainableCreator(
|
||||
_train_check_global, num_workers=2)
|
||||
trainable = trainable_cls()
|
||||
result = trainable.train()
|
||||
assert result["is_distributed"]
|
||||
|
||||
|
||||
def test_save_checkpoint(ray_start_2_cpus): # noqa: F811
|
||||
trainable_cls = DistributedTrainableCreator(_train_simple, num_workers=2)
|
||||
trainer = trainable_cls(config={"epochs": 1})
|
||||
trainer.train()
|
||||
checkpoint_dir = trainer.save()
|
||||
model_state_dict, opt_state_dict = torch.load(
|
||||
os.path.join(checkpoint_dir, "checkpoint"))
|
||||
trainer.stop()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enabled_checkpoint", [True, False])
|
||||
def test_simple_tune(ray_start_4_cpus, enabled_checkpoint):
|
||||
trainable_cls = DistributedTrainableCreator(_train_simple, num_workers=2)
|
||||
analysis = tune.run(
|
||||
trainable_cls,
|
||||
config={"enable_checkpoint": enabled_checkpoint},
|
||||
num_samples=2,
|
||||
stop={"training_iteration": 2})
|
||||
assert analysis.trials[0].last_result["training_iteration"] == 2
|
||||
assert analysis.trials[0].has_checkpoint() == enabled_checkpoint
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rank", [0, 1])
|
||||
def test_checkpoint(ray_start_2_cpus, rank): # noqa: F811
|
||||
with patch("torch.distributed.get_rank") as rank_method:
|
||||
rank_method.return_value = rank
|
||||
with distributed_checkpoint_dir(step="test") as path:
|
||||
if rank == 0:
|
||||
assert path
|
||||
if rank != 0:
|
||||
assert not os.path.exists(path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import pytest
|
||||
import sys
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
Reference in New Issue
Block a user