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

python/ray/tune/tests/test_actor_reuse.py

@@ -13,19 +13,19 @@ class FrequentPausesScheduler(FIFOScheduler):
 
 def create_resettable_class():
     class MyResettableClass(Trainable):
-        def _setup(self, config):
+        def setup(self, config):
             self.config = config
             self.num_resets = 0
             self.iter = 0
 
-        def _train(self):
+        def step(self):
             self.iter += 1
             return {"num_resets": self.num_resets, "done": self.iter > 1}
 
-        def _save(self, chkpt_dir):
+        def save_checkpoint(self, chkpt_dir):
             return {"iter": self.iter}
 
-        def _restore(self, item):
+        def load_checkpoint(self, item):
             self.iter = item["iter"]
 
         def reset_config(self, new_config):

python/ray/tune/tests/test_api.py

@@ -63,11 +63,11 @@ class TrainableFunctionApiTest(unittest.TestCase):
                 function_output.append(result)
 
         class _WrappedTrainable(Trainable):
-            def _setup(self, config):
+            def setup(self, config):
                 del config
                 self._result_iter = copy.deepcopy(class_results)
 
-            def _train(self):
+            def step(self):
                 if sleep_per_iter:
                     time.sleep(sleep_per_iter)
                 res = self._result_iter.pop(0)  # This should not fail
@@ -233,7 +233,7 @@ class TrainableFunctionApiTest(unittest.TestCase):
             def default_resource_request(cls, config):
                 return Resources(cpu=config["cpu"], gpu=config["gpu"])
 
-            def _train(self):
+            def step(self):
                 return {"timesteps_this_iter": 1, "done": True}
 
         register_trainable("B", B)
@@ -628,7 +628,7 @@ class TrainableFunctionApiTest(unittest.TestCase):
 
     def testTrialInfoAccess(self):
         class TestTrainable(Trainable):
-            def _train(self):
+            def step(self):
                 result = {"name": self.trial_name, "trial_id": self.trial_id}
                 print(result)
                 return result
@@ -659,11 +659,11 @@ class TrainableFunctionApiTest(unittest.TestCase):
     @patch("ray.tune.ray_trial_executor.TRIAL_CLEANUP_THRESHOLD", 3)
     def testLotsOfStops(self):
         class TestTrainable(Trainable):
-            def _train(self):
+            def step(self):
                 result = {"name": self.trial_name, "trial_id": self.trial_id}
                 return result
 
-            def _stop(self):
+            def cleanup(self):
                 time.sleep(2)
                 open(os.path.join(self.logdir, "marker"), "a").close()
                 return 1
@@ -825,17 +825,17 @@ class TrainableFunctionApiTest(unittest.TestCase):
 
     def testDurableTrainable(self):
         class TestTrain(DurableTrainable):
-            def _setup(self, config):
+            def setup(self, config):
                 self.state = {"hi": 1, "iter": 0}
 
-            def _train(self):
+            def step(self):
                 self.state["iter"] += 1
                 return {"timesteps_this_iter": 1, "done": True}
 
-            def _save(self, path):
+            def save_checkpoint(self, path):
                 return self.state
 
-            def _restore(self, state):
+            def load_checkpoint(self, state):
                 self.state = state
 
         sync_client = mock_storage_client()
@@ -853,16 +853,16 @@ class TrainableFunctionApiTest(unittest.TestCase):
 
     def testCheckpointDict(self):
         class TestTrain(Trainable):
-            def _setup(self, config):
+            def setup(self, config):
                 self.state = {"hi": 1}
 
-            def _train(self):
+            def step(self):
                 return {"timesteps_this_iter": 1, "done": True}
 
-            def _save(self, path):
+            def save_checkpoint(self, path):
                 return self.state
 
-            def _restore(self, state):
+            def load_checkpoint(self, state):
                 self.state = state
 
         test_trainable = TestTrain()
@@ -883,17 +883,17 @@ class TrainableFunctionApiTest(unittest.TestCase):
 
     def testMultipleCheckpoints(self):
         class TestTrain(Trainable):
-            def _setup(self, config):
+            def setup(self, config):
                 self.state = {"hi": 1, "iter": 0}
 
-            def _train(self):
+            def step(self):
                 self.state["iter"] += 1
                 return {"timesteps_this_iter": 1, "done": True}
 
-            def _save(self, path):
+            def save_checkpoint(self, path):
                 return self.state
 
-            def _restore(self, state):
+            def load_checkpoint(self, state):
                 self.state = state
 
         test_trainable = TestTrain()
@@ -938,6 +938,41 @@ class TrainableFunctionApiTest(unittest.TestCase):
         self.assertEqual(trial.last_result[TRAINING_ITERATION], 100)
         self.assertEqual(trial.last_result["itr"], 99)
 
+    def testBackwardsCompat(self):
+        class TestTrain(Trainable):
+            def _setup(self, config):
+                self.state = {"hi": 1, "iter": 0}
+
+            def _train(self):
+                self.state["iter"] += 1
+                return {"timesteps_this_iter": 1, "done": True}
+
+            def _save(self, path):
+                return self.state
+
+            def _restore(self, state):
+                self.state = state
+
+        test_trainable = TestTrain()
+        checkpoint_1 = test_trainable.save()
+        test_trainable.train()
+        checkpoint_2 = test_trainable.save()
+        self.assertNotEqual(checkpoint_1, checkpoint_2)
+        test_trainable.restore(checkpoint_2)
+        self.assertEqual(test_trainable.state["iter"], 1)
+        test_trainable.restore(checkpoint_1)
+        self.assertEqual(test_trainable.state["iter"], 0)
+
+        trials = run_experiments({
+            "foo": {
+                "run": TestTrain,
+                "checkpoint_at_end": True
+            }
+        })
+        for trial in trials:
+            self.assertEqual(trial.status, Trial.TERMINATED)
+            self.assertTrue(trial.has_checkpoint())
+
 
 if __name__ == "__main__":
     import pytest

python/ray/tune/tests/test_cluster.py

@@ -623,18 +623,18 @@ def test_cluster_interrupt(start_connected_cluster, tmpdir):
     class _Mock(tune.Trainable):
         """Finishes on the 4th iteration."""
 
-        def _setup(self, config):
+        def setup(self, config):
             self.state = {"hi": 0}
 
-        def _train(self):
+        def step(self):
             self.state["hi"] += 1
             time.sleep(0.5)
             return {"done": self.state["hi"] >= 4}
 
-        def _save(self, path):
+        def save_checkpoint(self, path):
             return self.state
 
-        def _restore(self, state):
+        def load_checkpoint(self, state):
             self.state = state
 
     # Removes indent from class.

python/ray/tune/tests/test_experiment_analysis_mem.py

@@ -21,19 +21,19 @@ class ExperimentAnalysisInMemorySuite(unittest.TestCase):
                 4: [7, 5, 5, 5, 5, 5, 5, 5, 3]
             }
 
-            def _setup(self, config):
+            def setup(self, config):
                 self.id = config["id"]
                 self.idx = 0
 
-            def _train(self):
+            def step(self):
                 val = self.scores_dict[self.id][self.idx]
                 self.idx += 1
                 return {"score": val}
 
-            def _save(self, checkpoint_dir):
+            def save_checkpoint(self, checkpoint_dir):
                 pass
 
-            def _restore(self, checkpoint_path):
+            def load_checkpoint(self, checkpoint_path):
                 pass
 
         self.MockTrainable = MockTrainable

python/ray/tune/tests/test_ray_trial_executor.py

@@ -145,7 +145,7 @@ class RayTrialExecutorTest(unittest.TestCase):
         """Tests that reset works as expected."""
 
         class B(Trainable):
-            def _train(self):
+            def step(self):
                 return dict(timesteps_this_iter=1, done=True)
 
             def reset_config(self, config):

python/ray/tune/tests/test_run_experiment.py

@@ -74,7 +74,7 @@ class RunExperimentTest(unittest.TestCase):
                 reporter(timesteps_total=i)
 
         class B(Trainable):
-            def _train(self):
+            def step(self):
                 return {"timesteps_this_iter": 1, "done": True}
 
         register_trainable("f1", train)
@@ -91,10 +91,10 @@ class RunExperimentTest(unittest.TestCase):
 
     def testCheckpointAtEnd(self):
         class train(Trainable):
-            def _train(self):
+            def step(self):
                 return {"timesteps_this_iter": 1, "done": True}
 
-            def _save(self, path):
+            def save_checkpoint(self, path):
                 checkpoint = os.path.join(path, "checkpoint")
                 with open(checkpoint, "w") as f:
                     f.write("OK")
@@ -112,7 +112,7 @@ class RunExperimentTest(unittest.TestCase):
 
     def testExportFormats(self):
         class train(Trainable):
-            def _train(self):
+            def step(self):
                 return {"timesteps_this_iter": 1, "done": True}
 
             def _export_model(self, export_formats, export_dir):
@@ -134,7 +134,7 @@ class RunExperimentTest(unittest.TestCase):
 
     def testInvalidExportFormats(self):
         class train(Trainable):
-            def _train(self):
+            def step(self):
                 return {"timesteps_this_iter": 1, "done": True}
 
             def _export_model(self, export_formats, export_dir):
@@ -156,7 +156,7 @@ class RunExperimentTest(unittest.TestCase):
         ray.init(resources={"hi": 3})
 
         class train(Trainable):
-            def _train(self):
+            def step(self):
                 return {"timesteps_this_iter": 1, "done": True}
 
         trials = run_experiments({

python/ray/tune/tests/test_trial_scheduler.py

@@ -1141,10 +1141,10 @@ class E2EPopulationBasedTestingSuite(unittest.TestCase):
         pbt = self.basicSetup(perturbation_interval=2)
 
         class train(tune.Trainable):
-            def _train(self):
+            def step(self):
                 return {"mean_accuracy": self.training_iteration}
 
-            def _save(self, path):
+            def save_checkpoint(self, path):
                 checkpoint = os.path.join(path, "checkpoint")
                 with open(checkpoint, "w") as f:
                     f.write("OK")
@@ -1173,16 +1173,16 @@ class E2EPopulationBasedTestingSuite(unittest.TestCase):
         pbt = self.basicSetup(perturbation_interval=2)
 
         class train_dict(tune.Trainable):
-            def _setup(self, config):
+            def setup(self, config):
                 self.state = {"hi": 1}
 
-            def _train(self):
+            def step(self):
                 return {"mean_accuracy": self.training_iteration}
 
-            def _save(self, path):
+            def save_checkpoint(self, path):
                 return self.state
 
-            def _restore(self, state):
+            def load_checkpoint(self, state):
                 self.state = state
 
         trial_hyperparams = {

python/ray/tune/tests/test_tune_save_restore.py

@@ -19,20 +19,20 @@ class SerialTuneRelativeLocalDirTest(unittest.TestCase):
     class MockTrainable(Trainable):
         _name = "MockTrainable"
 
-        def _setup(self, config):
+        def setup(self, config):
             self.state = {"hi": 1}
 
-        def _train(self):
+        def step(self):
             return {"timesteps_this_iter": 1, "done": True}
 
-        def _save(self, checkpoint_dir):
+        def save_checkpoint(self, checkpoint_dir):
             checkpoint_path = os.path.join(
                 checkpoint_dir, "checkpoint-{}".format(self._iteration))
             with open(checkpoint_path, "wb") as f:
                 pickle.dump(self.state, f)
             return checkpoint_path
 
-        def _restore(self, checkpoint_path):
+        def load_checkpoint(self, checkpoint_path):
             with open(checkpoint_path, "rb") as f:
                 extra_data = pickle.load(f)
             self.state.update(extra_data)
@@ -154,18 +154,18 @@ class SerialTuneRelativeLocalDirTest(unittest.TestCase):
         """Tests that passing the checkpoint_dir right back works."""
 
         class MockTrainable(Trainable):
-            def _setup(self, config):
+            def setup(self, config):
                 pass
 
-            def _train(self):
+            def step(self):
                 return {"score": 1}
 
-            def _save(self, checkpoint_dir):
+            def save_checkpoint(self, checkpoint_dir):
                 with open(os.path.join(checkpoint_dir, "test.txt"), "wb") as f:
                     pickle.dump("test", f)
                 return checkpoint_dir
 
-            def _restore(self, checkpoint_dir):
+            def load_checkpoint(self, checkpoint_dir):
                 with open(os.path.join(checkpoint_dir, "test.txt"), "rb") as f:
                     x = pickle.load(f)