mirror of
https://github.com/wassname/ray.git
synced 2026-09-10 12:38:43 +08:00
[tune] Function API checkpointing (#8471)
Co-authored-by: krfricke <krfricke@users.noreply.github.com>
This commit is contained in:
Richard Liaw
@@ -35,7 +35,7 @@ Here's an example of specifying the objective function using :ref:`the function-
|
||||
|
||||
Now, there's two Trainable APIs - one being the :ref:`function-based API <tune-function-api>` that we demonstrated above.
|
||||
|
||||
The other is a :ref:`class-based API <tune-class-api>` that enables :ref:`checkpointing and pausing <tune-trainable-save-restore>`. Here's an example of specifying the objective function using the :ref:`class-based API <tune-class-api>`:
|
||||
The other is a :ref:`class-based API <tune-class-api>`. Here's an example of specifying the objective function using the :ref:`class-based API <tune-class-api>`:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
|
||||
@@ -147,46 +147,39 @@ When running a hyperparameter search, Tune can automatically and periodically sa
|
||||
* fault-tolerance when using pre-emptible machines.
|
||||
* Pausing trials when using Trial Schedulers such as HyperBand and PBT.
|
||||
|
||||
To enable checkpointing, you must implement a :ref:`Trainable class <trainable-docs>` (the function-based API are not checkpointable, since they never return control back to their caller).
|
||||
Checkpointing assumes that the model state will be saved to disk on whichever node the Trainable is running on.
|
||||
|
||||
Checkpointing assumes that the model state will be saved to disk on whichever node the Trainable is running on. You can checkpoint with three different mechanisms: manually, periodically, and at termination.
|
||||
|
||||
**Manual Checkpointing**: A custom Trainable can manually trigger checkpointing by returning ``should_checkpoint: True`` (or ``tune.result.SHOULD_CHECKPOINT: True``) in the result dictionary of `_train`. This can be especially helpful in spot instances:
|
||||
To use Tune's checkpointing features, you must expose a ``checkpoint`` argument in the function signature, and call ``tune.make_checkpoint_dir`` and ``tune.save_checkpoint``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def _train(self):
|
||||
# training code
|
||||
result = {"mean_accuracy": accuracy}
|
||||
if detect_instance_preemption():
|
||||
result.update(should_checkpoint=True)
|
||||
return result
|
||||
import time
|
||||
from ray import tune
|
||||
|
||||
def train_func(config, checkpoint=None):
|
||||
start = 0
|
||||
if checkpoint:
|
||||
with open(checkpoint) as f:
|
||||
state = json.loads(f.read())
|
||||
start = state["step"] + 1
|
||||
|
||||
**Periodic Checkpointing**: periodic checkpointing can be used to provide fault-tolerance for experiments. This can be enabled by setting ``checkpoint_freq=<int>`` and ``max_failures=<int>`` to checkpoint trials every *N* iterations and recover from up to *M* crashes per trial, e.g.:
|
||||
for iter in range(start, 100):
|
||||
time.sleep(1)
|
||||
|
||||
.. code-block:: python
|
||||
# Obtain a checkpoint directory
|
||||
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": start}))
|
||||
tune.save_checkpoint(path)
|
||||
|
||||
tune.run(
|
||||
my_trainable,
|
||||
checkpoint_freq=10,
|
||||
max_failures=5,
|
||||
)
|
||||
tune.report(hello="world", ray="tune")
|
||||
|
||||
**Checkpointing at Termination**: The checkpoint_freq may not coincide with the exact end of an experiment. If you want a checkpoint to be created at the end
|
||||
of a trial, you can additionally set the ``checkpoint_at_end=True``:
|
||||
tune.run(train_func)
|
||||
|
||||
.. code-block:: python
|
||||
:emphasize-lines: 5
|
||||
In this example, checkpoints will be saved by training iteration to ``local_dir/exp_name/trial_name/checkpoint_<step>``.
|
||||
|
||||
tune.run(
|
||||
my_trainable,
|
||||
checkpoint_freq=10,
|
||||
checkpoint_at_end=True,
|
||||
max_failures=5,
|
||||
)
|
||||
|
||||
The checkpoint will be saved at a path that looks like ``local_dir/exp_name/trial_name/checkpoint_x/``, where the x is the number of iterations so far when the checkpoint is saved. To restore the checkpoint, you can use the ``restore`` argument and specify a checkpoint file. By doing this, you can change whatever experiments' configuration such as the experiment's name, the training iteration or so:
|
||||
You can restore a single trial checkpoint by using ``tune.run(restore=<checkpoint_dir>)`` By doing this, you can change whatever experiments' configuration such as the experiment's name:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
|
||||
@@ -3,9 +3,7 @@
|
||||
Training (tune.Trainable, tune.report)
|
||||
======================================
|
||||
|
||||
Training can be done with either a **Class API** (``tune.Trainable``) or **function-based API** (``tune.report``).
|
||||
|
||||
You can use the **function-based API** for fast prototyping. On the other hand, the ``tune.Trainable`` interface supports checkpoint/restore functionality and provides more control for advanced algorithms.
|
||||
Training can be done with either a **Class API** (``tune.Trainable``) or **function API** (``tune.report``).
|
||||
|
||||
For the sake of example, let's maximize this objective function:
|
||||
|
||||
@@ -16,8 +14,10 @@ For the sake of example, let's maximize this objective function:
|
||||
|
||||
.. _tune-function-api:
|
||||
|
||||
Function-based API
|
||||
------------------
|
||||
Function API
|
||||
------------
|
||||
|
||||
Here is a simple example of using the function API. You can report intermediate metrics by simply calling ``tune.report`` within the provided function.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@@ -25,31 +25,74 @@ Function-based API
|
||||
# config (dict): A dict of hyperparameters.
|
||||
|
||||
for x in range(20):
|
||||
score = objective(x, config["a"], config["b"])
|
||||
intermediate_score = objective(x, config["a"], config["b"])
|
||||
|
||||
tune.report(score=score) # This sends the score to Tune.
|
||||
tune.report(value=intermediate_score) # This sends the score to Tune.
|
||||
|
||||
analysis = tune.run(
|
||||
trainable,
|
||||
config={
|
||||
"a": 2,
|
||||
"b": 4
|
||||
})
|
||||
config={"a": 2, "b": 4}
|
||||
)
|
||||
|
||||
print("best config: ", analysis.get_best_config(metric="score", mode="max"))
|
||||
|
||||
.. tip:: Do not use ``tune.track.log`` within a ``Trainable`` class.
|
||||
.. tip:: Do not use ``tune.report`` within a ``Trainable`` class.
|
||||
|
||||
Tune will run this function on a separate thread in a Ray actor process. Note that this API is not checkpointable, since the thread will never return control back to its caller.
|
||||
Tune will run this function on a separate thread in a Ray actor process.
|
||||
|
||||
.. note:: If you want to pass in a Python lambda, you will need to first register the function: ``tune.register_trainable("lambda_id", lambda x: ...)``. You can then use ``lambda_id`` in place of ``my_trainable``.
|
||||
|
||||
Function API Checkpointing
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Many Tune features rely on checkpointing, including the usage of certain Trial Schedulers and fault tolerance. To use Tune's checkpointing features, you must expose a ``checkpoint`` argument in the function signature, and call ``tune.make_checkpoint_dir`` and ``tune.save_checkpoint``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import time
|
||||
from ray import tune
|
||||
|
||||
def train_func(config, checkpoint=None):
|
||||
start = 0
|
||||
if checkpoint:
|
||||
with open(checkpoint) as f:
|
||||
state = json.loads(f.read())
|
||||
start = state["step"] + 1
|
||||
|
||||
for iter in range(start, 100):
|
||||
time.sleep(1)
|
||||
|
||||
#
|
||||
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": start}))
|
||||
tune.save_checkpoint(path)
|
||||
|
||||
tune.report(hello="world", ray="tune")
|
||||
|
||||
tune.run(train_func)
|
||||
|
||||
In this example, checkpoints will be saved by training iteration to ``local_dir/exp_name/trial_name/checkpoint_<step>``. You can restore a single trial checkpoint by using ``tune.run(restore=<checkpoint_dir>)``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
analysis = tune.run(
|
||||
train,
|
||||
config={
|
||||
"max_iter": 5
|
||||
},
|
||||
).trials
|
||||
last_ckpt = trial.checkpoint.value
|
||||
analysis = tune.run(train, config={"max_iter": 10}, restore=last_ckpt)
|
||||
|
||||
Tune also may copy or move checkpoints during the course of tuning. For this purpose, it is important not to depend on absolute paths in the implementation of ``save``.
|
||||
|
||||
.. _tune-class-api:
|
||||
|
||||
Trainable Class API
|
||||
-------------------
|
||||
|
||||
.. caution:: Do not use ``tune.track.log`` within a ``Trainable`` class.
|
||||
.. caution:: Do not use ``tune.report`` within a ``Trainable`` class.
|
||||
|
||||
The Trainable **class API** will require users to subclass ``ray.tune.Trainable``. Here's a naive example of this API:
|
||||
|
||||
@@ -87,14 +130,13 @@ As a subclass of ``tune.Trainable``, Tune will create a ``Trainable`` object on
|
||||
|
||||
.. tip:: As a rule of thumb, the execution time of ``_train`` should be large enough to avoid overheads (i.e. more than a few seconds), but short enough to report progress periodically (i.e. at most a few minutes).
|
||||
|
||||
In this example, we only implemented the ``_setup`` and ``_train`` methods for simplification. Next, we'll implement ``_save`` and ``_restore`` for checkpoint and fault tolerance.
|
||||
|
||||
.. _tune-trainable-save-restore:
|
||||
|
||||
Save and Restore
|
||||
~~~~~~~~~~~~~~~~
|
||||
Class API Checkpointing
|
||||
~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Many Tune features rely on ``_save``, and ``_restore``, including the usage of certain Trial Schedulers, fault tolerance, and checkpointing.
|
||||
You can also implement checkpoint/restore using the Trainable Class API:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@@ -108,9 +150,45 @@ Many Tune features rely on ``_save``, and ``_restore``, including the usage of c
|
||||
checkpoint_path = os.path.join(tmp_checkpoint_dir, "model.pth")
|
||||
self.model.load_state_dict(torch.load(checkpoint_path))
|
||||
|
||||
Checkpoints will be saved by training iteration to ``local_dir/exp_name/trial_name/checkpoint_<iter>``. You can restore a single trial checkpoint by using ``tune.run(restore=<checkpoint_dir>)``.
|
||||
tune.run(MyTrainableClass, checkpoint_freq=2)
|
||||
|
||||
You can checkpoint with three different mechanisms: manually, periodically, and at termination.
|
||||
|
||||
**Manual Checkpointing**: A custom Trainable can manually trigger checkpointing by returning ``should_checkpoint: True`` (or ``tune.result.SHOULD_CHECKPOINT: True``) in the result dictionary of `_train`. This can be especially helpful in spot instances:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def _train(self):
|
||||
# training code
|
||||
result = {"mean_accuracy": accuracy}
|
||||
if detect_instance_preemption():
|
||||
result.update(should_checkpoint=True)
|
||||
return result
|
||||
|
||||
|
||||
**Periodic Checkpointing**: periodic checkpointing can be used to provide fault-tolerance for experiments. This can be enabled by setting ``checkpoint_freq=<int>`` and ``max_failures=<int>`` to checkpoint trials every *N* iterations and recover from up to *M* crashes per trial, e.g.:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
tune.run(
|
||||
my_trainable,
|
||||
checkpoint_freq=10,
|
||||
max_failures=5,
|
||||
)
|
||||
|
||||
**Checkpointing at Termination**: The checkpoint_freq may not coincide with the exact end of an experiment. If you want a checkpoint to be created at the end
|
||||
of a trial, you can additionally set the ``checkpoint_at_end=True``:
|
||||
|
||||
.. code-block:: python
|
||||
:emphasize-lines: 5
|
||||
|
||||
tune.run(
|
||||
my_trainable,
|
||||
checkpoint_freq=10,
|
||||
checkpoint_at_end=True,
|
||||
max_failures=5,
|
||||
)
|
||||
|
||||
Tune also generates temporary checkpoints for pausing and switching between trials. For this purpose, it is important not to depend on absolute paths in the implementation of ``save``.
|
||||
|
||||
Use ``validate_save_restore`` to catch ``_save``/``_restore`` errors before execution.
|
||||
|
||||
@@ -122,31 +200,11 @@ Use ``validate_save_restore`` to catch ``_save``/``_restore`` errors before exec
|
||||
validate_save_restore(MyTrainableClass)
|
||||
validate_save_restore(MyTrainableClass, use_object_store=True)
|
||||
|
||||
|
||||
Advanced Resource Allocation
|
||||
----------------------------
|
||||
|
||||
Trainables can themselves be distributed. If your trainable function / class creates further Ray actors or tasks that also consume CPU / GPU resources, you will want to set ``extra_cpu`` or ``extra_gpu`` inside ``tune.run`` to reserve extra resource slots. For example, if a trainable class requires 1 GPU itself, but also launches 4 actors, each using another GPU, then you should set ``"gpu": 1, "extra_gpu": 4``.
|
||||
|
||||
.. code-block:: python
|
||||
:emphasize-lines: 4-8
|
||||
|
||||
tune.run(
|
||||
my_trainable,
|
||||
name="my_trainable",
|
||||
resources_per_trial={
|
||||
"cpu": 1,
|
||||
"gpu": 1,
|
||||
"extra_gpu": 4
|
||||
}
|
||||
)
|
||||
|
||||
The ``Trainable`` also provides the ``default_resource_requests`` interface to automatically declare the ``resources_per_trial`` based on the given configuration.
|
||||
|
||||
|
||||
Advanced: Reusing Actors
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
.. note:: This feature is only for the Trainable Class API.
|
||||
|
||||
Your Trainable can often take a long time to start. To avoid this, you can do ``tune.run(reuse_actors=True)`` to reuse the same Trainable Python process and object for multiple hyperparameters.
|
||||
|
||||
This requires you to implement ``Trainable.reset_config``, which provides a new set of hyperparameters. It is up to the user to correctly update the hyperparameters of your trainable.
|
||||
@@ -176,8 +234,47 @@ This requires you to implement ``Trainable.reset_config``, which provides a new
|
||||
return True
|
||||
|
||||
|
||||
tune.Trainable
|
||||
--------------
|
||||
Advanced Resource Allocation
|
||||
----------------------------
|
||||
|
||||
Trainables can themselves be distributed. If your trainable function / class creates further Ray actors or tasks that also consume CPU / GPU resources, you will want to set ``extra_cpu`` or ``extra_gpu`` inside ``tune.run`` to reserve extra resource slots. For example, if a trainable class requires 1 GPU itself, but also launches 4 actors, each using another GPU, then you should set ``"gpu": 1, "extra_gpu": 4``.
|
||||
|
||||
.. code-block:: python
|
||||
:emphasize-lines: 4-8
|
||||
|
||||
tune.run(
|
||||
my_trainable,
|
||||
name="my_trainable",
|
||||
resources_per_trial={
|
||||
"cpu": 1,
|
||||
"gpu": 1,
|
||||
"extra_gpu": 4
|
||||
}
|
||||
)
|
||||
|
||||
The ``Trainable`` also provides the ``default_resource_requests`` interface to automatically declare the ``resources_per_trial`` based on the given configuration.
|
||||
|
||||
|
||||
|
||||
.. _track-docstring:
|
||||
|
||||
tune.report / tune.checkpoint (Function API)
|
||||
--------------------------------------------
|
||||
|
||||
.. autofunction:: ray.tune.report
|
||||
|
||||
.. autofunction:: ray.tune.make_checkpoint_dir
|
||||
|
||||
.. autofunction:: ray.tune.save_checkpoint
|
||||
|
||||
.. autofunction:: ray.tune.get_trial_dir
|
||||
|
||||
.. autofunction:: ray.tune.get_trial_name
|
||||
|
||||
.. autofunction:: ray.tune.get_trial_id
|
||||
|
||||
tune.Trainable (Class API)
|
||||
--------------------------
|
||||
|
||||
|
||||
.. autoclass:: ray.tune.Trainable
|
||||
@@ -190,21 +287,6 @@ tune.DurableTrainable
|
||||
|
||||
.. autoclass:: ray.tune.DurableTrainable
|
||||
|
||||
.. _track-docstring:
|
||||
|
||||
tune.track
|
||||
----------
|
||||
|
||||
.. automodule:: ray.tune.track
|
||||
:members:
|
||||
:exclude-members: init,
|
||||
|
||||
KerasCallback
|
||||
-------------
|
||||
|
||||
.. automodule:: ray.tune.integration.keras
|
||||
:members:
|
||||
|
||||
|
||||
StatusReporter
|
||||
--------------
|
||||
|
||||
Reference in New Issue
Block a user