wehub-resource-sync
163 lines
4.2 KiB
Python
163 lines
4.2 KiB
Python
# flake8: noqa
|
|
|
|
# __ft_initial_run_start__
|
|
import json
|
|
import os
|
|
import tempfile
|
|
|
|
from ray import tune
|
|
|
|
|
|
def trainable(config):
|
|
# Checkpoint loading
|
|
checkpoint = tune.get_checkpoint()
|
|
start = 1
|
|
if checkpoint:
|
|
with checkpoint.as_directory() as checkpoint_dir:
|
|
with open(os.path.join(checkpoint_dir, "checkpoint.json"), "r") as f:
|
|
state = json.load(f)
|
|
start = state["epoch"] + 1
|
|
|
|
for epoch in range(start, config["num_epochs"]):
|
|
# Do some training...
|
|
|
|
# Checkpoint saving
|
|
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
|
with open(os.path.join(temp_checkpoint_dir, "checkpoint.json"), "w") as f:
|
|
json.dump({"epoch": epoch}, f)
|
|
tune.report(
|
|
{"epoch": epoch},
|
|
checkpoint=tune.Checkpoint.from_directory(temp_checkpoint_dir),
|
|
)
|
|
|
|
|
|
tuner = tune.Tuner(
|
|
trainable,
|
|
param_space={"num_epochs": 10},
|
|
run_config=tune.RunConfig(
|
|
storage_path=os.path.expanduser("~/ray_results"),
|
|
name="tune_fault_tolerance_guide",
|
|
),
|
|
)
|
|
result_grid = tuner.fit()
|
|
# __ft_initial_run_end__
|
|
|
|
assert not result_grid.errors
|
|
|
|
# __ft_restored_run_start__
|
|
tuner = tune.Tuner.restore(
|
|
os.path.expanduser("~/ray_results/tune_fault_tolerance_guide"),
|
|
trainable=trainable,
|
|
resume_errored=True,
|
|
)
|
|
tuner.fit()
|
|
# __ft_restored_run_end__
|
|
|
|
# __ft_restore_options_start__
|
|
tuner = tune.Tuner.restore(
|
|
os.path.expanduser("~/ray_results/tune_fault_tolerance_guide"),
|
|
trainable=trainable,
|
|
resume_errored=True,
|
|
restart_errored=False,
|
|
resume_unfinished=True,
|
|
)
|
|
# __ft_restore_options_end__
|
|
|
|
# __ft_restore_multiplexing_start__
|
|
import os
|
|
from ray import tune
|
|
|
|
storage_path = os.path.expanduser("~/ray_results")
|
|
exp_name = "tune_fault_tolerance_guide"
|
|
path = os.path.join(storage_path, exp_name)
|
|
|
|
if tune.Tuner.can_restore(path):
|
|
tuner = tune.Tuner.restore(path, trainable=trainable, resume_errored=True)
|
|
else:
|
|
tuner = tune.Tuner(
|
|
trainable,
|
|
param_space={"num_epochs": 10},
|
|
run_config=tune.RunConfig(storage_path=storage_path, name=exp_name),
|
|
)
|
|
tuner.fit()
|
|
# __ft_restore_multiplexing_end__
|
|
|
|
|
|
# Run the multiplexed logic again to make sure it goes through the restore branch.
|
|
if tune.Tuner.can_restore(path):
|
|
tuner = tune.Tuner.restore(path, trainable=trainable, resume_errored=True)
|
|
else:
|
|
tuner = tune.Tuner(
|
|
trainable,
|
|
param_space={"num_epochs": 10},
|
|
run_config=tune.RunConfig(storage_path=storage_path, name=exp_name),
|
|
)
|
|
assert tuner.get_results()
|
|
|
|
|
|
# __ft_restore_objrefs_initial_start__
|
|
import ray
|
|
from ray import tune
|
|
|
|
|
|
class LargeModel:
|
|
def __init__(self, model_id):
|
|
self.model_id = model_id
|
|
# Load weights based on the `model_id`...
|
|
|
|
|
|
def train_fn(config):
|
|
# Retrieve the model from the object store.
|
|
model = ray.get(config["model_ref"])
|
|
print(model.model_id)
|
|
|
|
|
|
# These models may be large, so `ray.put` them in the Ray Object Store
|
|
# to share the models between trials.
|
|
model_refs = [ray.put(LargeModel(1)), ray.put(LargeModel(2))]
|
|
|
|
tuner = tune.Tuner(
|
|
train_fn,
|
|
# Tune over the object references!
|
|
param_space={"model_ref": tune.grid_search(model_refs)},
|
|
run_config=tune.RunConfig(
|
|
storage_path=os.path.expanduser("~/ray_results"), name="restore_object_refs"
|
|
),
|
|
)
|
|
tuner.fit()
|
|
# __ft_restore_objrefs_initial_end__
|
|
|
|
if ray.is_initialized():
|
|
ray.shutdown()
|
|
|
|
# __ft_restore_objrefs_restored_start__
|
|
# Re-create the objects and put them in the object store.
|
|
param_space = {
|
|
"model_ref": tune.grid_search([ray.put(LargeModel(1)), ray.put(LargeModel(2))])
|
|
}
|
|
|
|
tuner = tune.Tuner.restore(
|
|
os.path.expanduser("~/ray_results/restore_object_refs"),
|
|
trainable=train_fn,
|
|
# Re-specify the `param_space` to update the object references.
|
|
param_space=param_space,
|
|
resume_errored=True,
|
|
)
|
|
tuner.fit()
|
|
# __ft_restore_objrefs_restored_end__
|
|
|
|
# __ft_trial_failure_start__
|
|
from ray import tune
|
|
|
|
tuner = tune.Tuner(
|
|
trainable,
|
|
param_space={"num_epochs": 10},
|
|
run_config=tune.RunConfig(
|
|
storage_path=os.path.expanduser("~/ray_results"),
|
|
name="trial_fault_tolerance",
|
|
failure_config=tune.FailureConfig(max_failures=3),
|
|
),
|
|
)
|
|
tuner.fit()
|
|
# __ft_trial_failure_end__
|