[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>

doc/source/tune/api_docs/trainable.rst

@@ -17,6 +17,7 @@ For the sake of example, let's maximize this objective function:
 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
@@ -40,38 +41,41 @@ Here is a simple example of using the function API. You can report intermediate
 
 Tune will run this function on a separate thread in a Ray actor process.
 
+.. tip:: If you want to leverage multi-node data parallel training with PyTorch while using parallel hyperparameter tuning, check out our :ref:PyTorch user guide and Tune's :ref:distributed pytorch integrations.
+
+.. _tune-function-checkpointing:
 
 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``:
+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_dir`` argument in the function signature, and call ``tune.checkpoint_dir`` :
 
 .. code-block:: python
 
         import time
         from ray import tune
 
-        def train_func(config, checkpoint=None):
+        def train_func(config, checkpoint_dir=None):
             start = 0
-            if checkpoint:
-                with open(checkpoint) as f:
+            if checkpoint_dir:
+                with open(os.path.join(checkpoint_dir, "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)
+                with tune.checkpoint_dir(step=step):
+                    path = os.path.join(checkpoint_dir, "checkpoint")
+                    with open(path, "w") as f:
+                        f.write(json.dumps({"step": start}))
 
                 tune.report(hello="world", ray="tune")
 
         tune.run(train_func)
 
+.. note:: ``checkpoint_freq`` and ``checkpoint_at_end`` will not work with Function API checkpointing.
+
 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
@@ -263,9 +267,7 @@ tune.report / tune.checkpoint (Function API)
 
 .. autofunction:: ray.tune.report
 
-.. autofunction:: ray.tune.make_checkpoint_dir
-
-.. autofunction:: ray.tune.save_checkpoint
+.. autofunction:: ray.tune.checkpoint_dir
 
 .. autofunction:: ray.tune.get_trial_dir
 
@@ -282,6 +284,21 @@ tune.Trainable (Class API)
     :private-members:
     :members:
 
+
+.. _tune-ddp-doc:
+
+Distributed Torch
+-----------------
+
+Ray also offers lightweight integrations to distribute your model training on Ray Tune.
+
+
+.. autofunction:: ray.tune.integration.torch.DistributedTrainableCreator
+
+.. autofunction:: ray.tune.integration.torch.distributed_checkpoint_dir
+
+.. autofunction:: ray.tune.integration.torch.is_distributed_trainable
+
 tune.DurableTrainable
 ---------------------