[tune] Update trainable docs and support hparams (#5558)

python/ray/tune/trainable.py

@@ -40,14 +40,11 @@ class Trainable(object):
     Calling ``save()`` should save the training state of a trainable to disk,
     and ``restore(path)`` should restore a trainable to the given state.
 
-    Generally you only need to implement ``_train``, ``_save``, and
-    ``_restore`` here when subclassing Trainable.
+    Generally you only need to implement ``_setup``, ``_train``,
+    ``_save``, and ``_restore`` when subclassing Trainable.
 
-    Note that, if you don't require checkpoint/restore functionality, then
-    instead of implementing this class you can also get away with supplying
-    just a ``my_train(config, reporter)`` function to the config.
-    The function will be automatically converted to this interface
-    (sans checkpoint functionality).
+    Other implementation methods that may be helpful to override are
+    ``_log_result``, ``reset_config``, ``_stop``, and ``_export_model``.
 
     When using Tune, Tune will convert this class into a Ray actor, which
     runs on a separate process. Tune will also change the current working
@@ -112,6 +109,14 @@ class Trainable(object):
 
         This can be overriden by sub-classes to set the correct trial resource
         allocation, so the user does not need to.
+
+        Example:
+            >>> def default_resource_request(cls, config):
+                    return Resources(
+                        cpu=0,
+                        gpu=0,
+                        extra_cpu=config["workers"],
+                        extra_gpu=int(config["use_gpu"]) * config["workers"])
         """
 
         return None
@@ -451,7 +456,7 @@ class Trainable(object):
 
         The return value will be automatically passed to the loggers. Users
         can also return `tune.result.DONE` or `tune.result.SHOULD_CHECKPOINT`
-        to manually trigger termination of this trial or checkpointing of this
+        as a key to manually trigger termination or checkpointing of this
         trial. Note that manual checkpointing only works when subclassing
         Trainables.
 
@@ -462,26 +467,38 @@ class Trainable(object):
 
         raise NotImplementedError
 
-    def _save(self, checkpoint_dir):
-        """Subclasses should override this to implement save().
+    def _save(self, tmp_checkpoint_dir):
+        """Subclasses should override this to implement ``save()``.
+
+        Warning:
+            Do not rely on absolute paths in the implementation of ``_save``
+            and ``_restore``.
+
+        Use ``validate_save_restore`` to catch ``_save``/``_restore`` errors
+        before execution.
+
+        >>> from ray.tune.util import validate_save_restore
+        >>> validate_save_restore(MyTrainableClass)
+        >>> validate_save_restore(MyTrainableClass, use_object_store=True)
 
         Args:
-            checkpoint_dir (str): The directory where the checkpoint
-                file must be stored. In a Tune run, this defaults to
-                `<self.logdir>/checkpoint_<ITER>` (which is the same as
-                `local_dir/exp_name/trial_name/checkpoint_<ITER>`).
+            tmp_checkpoint_dir (str): The directory where the checkpoint
+                file must be stored. In a Tune run, if the trial is paused,
+                the provided path may be temporary and moved.
 
         Returns:
-            checkpoint (str | dict): If string, the return value is
-                expected to be the checkpoint path or prefix to be passed to
-                `_restore()`. If dict, the return value will be automatically
-                serialized by Tune and passed to `_restore()`.
+            A dict or string. If string, the return value is expected to be
+            prefixed by `tmp_checkpoint_dir`. If dict, the return value will
+            be automatically serialized by Tune and passed to `_restore()`.
 
         Examples:
             >>> print(trainable1._save("/tmp/checkpoint_1"))
             "/tmp/checkpoint_1/my_checkpoint_file"
             >>> print(trainable2._save("/tmp/checkpoint_2"))
             {"some": "data"}
+
+            >>> trainable._save("/tmp/bad_example")
+            "/tmp/NEW_CHECKPOINT_PATH/my_checkpoint_file" # This will error.
         """
 
         raise NotImplementedError
@@ -489,9 +506,42 @@ class Trainable(object):
     def _restore(self, checkpoint):
         """Subclasses should override this to implement restore().
 
+        Warning:
+            In this method, do not rely on absolute paths. The absolute
+            path of the checkpoint_dir used in ``_save`` may be changed.
+
+        If ``_save`` returned a prefixed string, the prefix of the checkpoint
+        string returned by ``_save`` may be changed. This is because trial
+        pausing depends on temporary directories.
+
+        The directory structure under the checkpoint_dir provided to ``_save``
+        is preserved.
+
+        See the example below.
+
+        .. code-block:: python
+
+            class Example(Trainable):
+                def _save(self, checkpoint_path):
+                    print(checkpoint_path)
+                    return os.path.join(checkpoint_path, "my/check/point")
+
+                def _restore(self, checkpoint):
+                    print(checkpoint)
+
+            >>> trainer = Example()
+            >>> obj = trainer.save_to_object()  # This is used when PAUSED.
+            <logdir>/tmpc8k_c_6hsave_to_object/checkpoint_0/my/check/point
+            >>> trainer.restore_from_object(obj)  # Note the different prefix.
+            <logdir>/tmpb87b5axfrestore_from_object/checkpoint_0/my/check/point
+
+
         Args:
-            checkpoint (str | dict): Value as returned by `_save`.
-                If a string, then it is the checkpoint path.
+            checkpoint (str|dict): If dict, the return value is as
+                returned by `_save`. If a string, then it is a checkpoint path
+                that may have a different prefix than that returned by `_save`.
+                The directory structure underneath the `checkpoint_dir`
+                `_save` is preserved.
         """
 
         raise NotImplementedError
@@ -514,7 +564,14 @@ class Trainable(object):
         self._result_logger.on_result(result)
 
     def _stop(self):
-        """Subclasses should override this for any cleanup on stop."""
+        """Subclasses should override this for any cleanup on stop.
+
+        If any Ray actors are launched in the Trainable (i.e., with a RLlib
+        trainer), be sure to kill the Ray actor process here.
+
+        You can kill a Ray actor by calling `actor.__ray_terminate__.remote()`
+        on the actor.
+        """
         pass
 
     def _export_model(self, export_formats, export_dir):