[tune] Use public methods for trainable (#9184)

doc/source/tune/api_docs/logging.rst

@@ -99,7 +99,7 @@ You can do this in the trainable, as shown below:
 .. code-block:: python
 
     class CustomLogging(tune.Trainable)
-        def _setup(self, config):
+        def setup(self, config):
             trial_id = self.trial_id
             library.init(
                 name=trial_id,
@@ -109,10 +109,10 @@ You can do this in the trainable, as shown below:
                 allow_val_change=True)
             library.set_log_path(self.logdir)
 
-        def _train(self):
+        def step(self):
             library.log_model(...)
 
-        def _log_result(self, result):
+        def log_result(self, result):
             res_dict = {
                 str(k): v
                 for k, v in result.items()

doc/source/tune/api_docs/trainable.rst

@@ -61,7 +61,7 @@ Many Tune features rely on checkpointing, including the usage of certain Trial S
             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:
@@ -101,13 +101,13 @@ The Trainable **class API** will require users to subclass ``ray.tune.Trainable`
     from ray import tune
 
     class Trainable(tune.Trainable):
-        def _setup(self, config):
+        def setup(self, config):
             # config (dict): A dict of hyperparameters
             self.x = 0
             self.a = config["a"]
             self.b = config["b"]
 
-        def _train(self):  # This is called iteratively.
+        def step(self):  # This is called iteratively.
             score = objective(self.x, self.a, self.b)
             self.x += 1
             return {"score": score}
@@ -124,11 +124,11 @@ The Trainable **class API** will require users to subclass ``ray.tune.Trainable`
 
 As a subclass of ``tune.Trainable``, Tune will create a ``Trainable`` object on a separate process (using the :ref:`Ray Actor API <actor-guide>`).
 
-  1. ``_setup`` function is invoked once training starts.
-  2. ``_train`` is invoked **multiple times**. Each time, the Trainable object executes one logical iteration of training in the tuning process, which may include one or more iterations of actual training.
-  3. ``_stop`` is invoked when training is finished.
+  1. ``setup`` function is invoked once training starts.
+  2. ``step`` is invoked **multiple times**. Each time, the Trainable object executes one logical iteration of training in the tuning process, which may include one or more iterations of actual training.
+  3. ``cleanup`` is invoked when training is finished.
 
-.. 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).
+.. tip:: As a rule of thumb, the execution time of ``step`` 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).
 
 
 .. _tune-trainable-save-restore:
@@ -141,12 +141,12 @@ You can also implement checkpoint/restore using the Trainable Class API:
 .. code-block:: python
 
     class MyTrainableClass(Trainable):
-        def _save(self, tmp_checkpoint_dir):
+        def save_checkpoint(self, tmp_checkpoint_dir):
             checkpoint_path = os.path.join(tmp_checkpoint_dir, "model.pth")
             torch.save(self.model.state_dict(), checkpoint_path)
             return tmp_checkpoint_dir
 
-        def _restore(self, tmp_checkpoint_dir):
+        def load_checkpoint(self, tmp_checkpoint_dir):
             checkpoint_path = os.path.join(tmp_checkpoint_dir, "model.pth")
             self.model.load_state_dict(torch.load(checkpoint_path))
 
@@ -154,11 +154,11 @@ You can also implement checkpoint/restore using the Trainable Class API:
 
 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:
+**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 `step`. This can be especially helpful in spot instances:
 
 .. code-block:: python
 
-    def _train(self):
+    def step(self):
         # training code
         result = {"mean_accuracy": accuracy}
         if detect_instance_preemption():
@@ -190,7 +190,7 @@ of a trial, you can additionally set the ``checkpoint_at_end=True``:
     )
 
 
-Use ``validate_save_restore`` to catch ``_save``/``_restore`` errors before execution.
+Use ``validate_save_restore`` to catch ``save_checkpoint``/``load_checkpoint`` errors before execution.
 
 .. code-block:: python
 
@@ -214,7 +214,7 @@ This requires you to implement ``Trainable.reset_config``, which provides a new
     class PytorchTrainble(tune.Trainable):
         """Train a Pytorch ConvNet."""
 
-        def _setup(self, config):
+        def setup(self, config):
             self.train_loader, self.test_loader = get_data_loaders()
             self.model = ConvNet()
             self.optimizer = optim.SGD(