diff --git a/doc/source/tune/user-guide.rst b/doc/source/tune/user-guide.rst index 5aaba58a2..5d9e6540c 100644 --- a/doc/source/tune/user-guide.rst +++ b/doc/source/tune/user-guide.rst @@ -222,29 +222,29 @@ To use Tune's checkpointing features, you must expose a ``checkpoint_dir`` argum .. code-block:: python - import os - import time - from ray import tune + import os + import time + from ray import tune - def train_func(config, checkpoint_dir=None): - start = 0 - if checkpoint_dir: - with open(os.path.join(checkpoint_dir, "checkpoint")) as f: - state = json.loads(f.read()) - start = state["step"] + 1 + def train_func(config, checkpoint_dir=None): + start = 0 + if checkpoint_dir: + with open(os.path.join(checkpoint_dir, "checkpoint")) as f: + state = json.loads(f.read()) + start = state["step"] + 1 - for step in range(start, 100): - time.sleep(1) + for step in range(start, 100): + time.sleep(1) - # Obtain a checkpoint directory - 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": start})) + # Obtain a checkpoint directory + 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(hello="world", ray="tune") + tune.report(hello="world", ray="tune") - tune.run(train_func) + tune.run(train_func) In this example, checkpoints will be saved by training iteration to ``local_dir/exp_name/trial_name/checkpoint_``.