[tune] Added PyTorch Lightning callbacks to integrations (#10220)

Co-authored-by: Richard Liaw <rliaw@berkeley.edu>

doc/source/tune/_tutorials/tune-pytorch-lightning.rst

@@ -95,14 +95,20 @@ that can be used to plug custom functions into the training loop. This way the o
 ``LightningModule`` does not have to be altered at all. Also, we could use the same
 callback for multiple modules.
 
-The callback just reports some metrics back to Tune after each validation epoch:
+Ray Tune comes with ready-to-use PyTorch Lightning callbacks. To report metrics
+back to Tune after each validation epoch, we will use the ``TuneReportCallback``:
 
-.. literalinclude:: /../../python/ray/tune/examples/mnist_pytorch_lightning.py
-   :language: python
-   :start-after: __tune_callback_begin__
-   :end-before: __tune_callback_end__
+.. code-block:: python
 
-Note that we have to explicitly convert the metrics from a tensor to a Python value.
+    from ray.tune.integration.pytorch_lightning import TuneReportCallback
+    callback = TuneReportCallback({
+        "loss": "avg_val_loss",
+        "mean_accuracy": "avg_val_accuracy"
+    }, on="validation_end")
+
+This callback will take the ``avg_val_loss`` and ``avg_val_accuracy`` values
+from the PyTorch Lightning trainer and report them to Tune as the ``loss``
+and ``mean_accuracy``, respectively.
 
 Adding the Tune training function
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -282,12 +288,20 @@ Adding checkpoints to the PyTorch Lightning module
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 First, we need to introduce
-another callback to save model checkpoints:
+another callback to save model checkpoints. Since Tune requires a call to
+``tune.report()`` after creating a new checkpoint to register it, we will use
+a combined reporting and checkpointing callback:
 
-.. literalinclude:: /../../python/ray/tune/examples/mnist_pytorch_lightning.py
-   :language: python
-   :start-after: __tune_checkpoint_callback_begin__
-   :end-before: __tune_checkpoint_callback_end__
+.. code-block:: python
+
+    from ray.tune.integration.pytorch_lightning import TuneReportCheckpointCallback
+    callback = TuneReportCheckpointCallback(
+        metrics={"loss": "val_loss", "mean_accuracy": "val_accuracy"},
+        filename="checkpoint",
+        on="validation_end")
+
+The ``checkpoint`` value is the name of the checkpoint file within the
+checkpoint directory.
 
 We also include checkpoint loading in our training function:
 
@@ -338,5 +352,5 @@ In some runs, the parameters have been perturbed. And the best configuration eve
 mean validation accuracy of ``0.987062``!
 
 In summary, PyTorch Lightning Modules are easy to extend to use with Tune. It just took
-us writing one or two callbacks and a small wrapper function to get great performing
+us importing one or two callbacks and a small wrapper function to get great performing
 parameter configurations.

doc/source/tune/_tutorials/tune-wandb.rst

@@ -29,7 +29,9 @@ Please :doc:`see here for a full example </tune/examples/wandb_example>`.
 .. _tune-wandb-logger:
 
 .. autoclass:: ray.tune.integration.wandb.WandbLogger
+   :noindex:
 
 .. _tune-wandb-mixin:
 
 .. autofunction:: ray.tune.integration.wandb.wandb_mixin
+   :noindex: