[tune] Function API checkpointing (#8471)

Co-authored-by: krfricke <krfricke@users.noreply.github.com>

python/ray/tune/tests/test_cluster.py

@@ -394,10 +394,28 @@ def test_migration_checkpoint_removal(start_connected_emptyhead_cluster,
     }
 
     # The following patches only affect __fake_remote.
-    find_checkpoint_dir = TrainableUtil.find_checkpoint_dir
-    with patch("ray.tune.logger.get_node_syncer") as mock_get_node_syncer:
-        trainable_util = "ray.tune.ray_trial_executor.TrainableUtil"
-        with patch(trainable_util + ".find_checkpoint_dir") as mock_find_dir:
+    def hide_remote_path(path_function):
+        def hidden_path_func(checkpoint_path):
+            """Converts back to local path first."""
+            if MOCK_REMOTE_DIR in checkpoint_path:
+                checkpoint_path = checkpoint_path[len(MOCK_REMOTE_DIR):]
+                checkpoint_path = os.path.join("/", checkpoint_path)
+            return path_function(checkpoint_path)
+
+        return hidden_path_func
+
+    trainable_util = "ray.tune.ray_trial_executor.TrainableUtil"
+    _find_ckpt = trainable_util + ".find_checkpoint_dir"
+    find_func = TrainableUtil.find_checkpoint_dir
+    _pickle_ckpt = trainable_util + ".pickle_checkpoint"
+    pickle_func = TrainableUtil.pickle_checkpoint
+
+    with patch(_find_ckpt) as mock_find, patch(_pickle_ckpt) as mock_pkl_ckpt:
+        # __fake_remote trainables save to a separate "remote" directory.
+        # TrainableUtil will not check this path unless we mock it.
+        mock_find.side_effect = hide_remote_path(find_func)
+        mock_pkl_ckpt.side_effect = hide_remote_path(pickle_func)
+        with patch("ray.tune.logger.get_node_syncer") as mock_get_node_syncer:
 
             def mock_get_syncer_fn(local_dir, remote_dir, sync_function):
                 client = mock_storage_client()
@@ -405,16 +423,6 @@ def test_migration_checkpoint_removal(start_connected_emptyhead_cluster,
 
             mock_get_node_syncer.side_effect = mock_get_syncer_fn
 
-            def mock_find_dir_fn(checkpoint_path):
-                """Converts back to local path first."""
-                checkpoint_path = checkpoint_path[len(MOCK_REMOTE_DIR):]
-                checkpoint_path = os.path.join("/", checkpoint_path)
-                return find_checkpoint_dir(checkpoint_path)
-
-            # __fake_remote trainables save to a separate "remote" directory.
-            # TrainableUtil will not check this path unless we mock it.
-            mock_find_dir.side_effect = mock_find_dir_fn
-
             # Test recovery of trial that has been checkpointed
             t1 = Trial(trainable_id, **kwargs)
             runner.add_trial(t1)
@@ -428,7 +436,6 @@ def test_migration_checkpoint_removal(start_connected_emptyhead_cluster,
             cluster.remove_node(node)
             cluster.wait_for_nodes()
             shutil.rmtree(os.path.dirname(t1.checkpoint.value))
-
             runner.step()  # Collect result 3, kick off + fail result 4
             runner.step()  # Dispatch restore
             runner.step()  # Process restore + step 4

python/ray/tune/tests/test_function_api.py

@@ -0,0 +1,169 @@
+import json
+import os
+import unittest
+
+import ray
+from ray.rllib import _register_all
+
+from ray import tune
+from ray.tune.function_runner import wrap_function
+from ray.tune.result import TRAINING_ITERATION
+
+
+class FunctionApiTest(unittest.TestCase):
+    def setUp(self):
+        ray.init(num_cpus=4, num_gpus=0, object_store_memory=150 * 1024 * 1024)
+
+    def tearDown(self):
+        ray.shutdown()
+        _register_all()  # re-register the evicted objects
+
+    def testFunctionNoCheckpointing(self):
+        def train(config, checkpoint=None):
+            for i in range(10):
+                tune.report(test=i)
+
+        wrapped = wrap_function(train)
+
+        new_trainable = wrapped()
+        result = new_trainable.train()
+        checkpoint = new_trainable.save()
+        new_trainable.stop()
+
+        new_trainable2 = wrapped()
+        new_trainable2.restore(checkpoint)
+        result = new_trainable2.train()
+        self.assertEquals(result[TRAINING_ITERATION], 1)
+        checkpoint = new_trainable2.save()
+        new_trainable2.stop()
+
+    def testFunctionRecurringSave(self):
+        """This tests that save and restore are commutative."""
+
+        def train(config, checkpoint=None):
+            for step in range(10):
+                if step % 3 == 0:
+                    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": step}))
+                    tune.save_checkpoint(path)
+                tune.report(test=step)
+
+        wrapped = wrap_function(train)
+
+        new_trainable = wrapped()
+        new_trainable.train()
+        checkpoint_obj = new_trainable.save_to_object()
+        new_trainable.restore_from_object(checkpoint_obj)
+        checkpoint = new_trainable.save()
+        new_trainable.stop()
+
+        new_trainable2 = wrapped()
+        new_trainable2.restore(checkpoint)
+        new_trainable2.train()
+        new_trainable2.stop()
+
+    def testCheckpointFunctionAtEnd(self):
+        def train(config, checkpoint=False):
+            for i in range(10):
+                tune.report(test=i)
+            checkpoint_dir = tune.make_checkpoint_dir(step=10)
+            checkpoint_path = os.path.join(checkpoint_dir, "hello")
+            with open(checkpoint_path, "w") as f:
+                f.write("hello")
+            tune.save_checkpoint(checkpoint_path)
+
+        [trial] = tune.run(train).trials
+        assert "hello" in trial.checkpoint.value
+
+    def testVariousCheckpointFunctionAtEnd(self):
+        def train(config, checkpoint=False):
+            for i in range(10):
+                checkpoint_dir = tune.make_checkpoint_dir()
+                checkpoint_path = os.path.join(checkpoint_dir, "hello")
+                with open(checkpoint_path, "w") as f:
+                    f.write("hello")
+                tune.save_checkpoint(checkpoint_path)
+                tune.report(test=i)
+            checkpoint_dir = tune.make_checkpoint_dir()
+            checkpoint_path = os.path.join(checkpoint_dir, "goodbye")
+            with open(checkpoint_path, "w") as f:
+                f.write("goodbye")
+            tune.save_checkpoint(checkpoint_path)
+
+        [trial] = tune.run(train, keep_checkpoints_num=3).trials
+        assert "goodbye" in trial.checkpoint.value
+
+    def testReuseCheckpoint(self):
+        def train(config, checkpoint=False):
+            itr = 0
+            if checkpoint:
+                with open(checkpoint, "r") as f:
+                    itr = int(f.read()) + 1
+
+            for i in range(itr, config["max_iter"]):
+                checkpoint_dir = tune.make_checkpoint_dir(step=i)
+                checkpoint_path = os.path.join(checkpoint_dir, "goodbye")
+                with open(checkpoint_path, "w") as f:
+                    f.write(str(i))
+                tune.save_checkpoint(checkpoint_path)
+                tune.report(test=i, training_iteration=i)
+
+        [trial] = tune.run(
+            train,
+            config={
+                "max_iter": 5
+            },
+        ).trials
+        last_ckpt = trial.checkpoint.value
+        assert "goodbye" in last_ckpt
+        analysis = tune.run(train, config={"max_iter": 10}, restore=last_ckpt)
+        trial_dfs = list(analysis.trial_dataframes.values())
+        assert len(trial_dfs[0]["training_iteration"]) == 5
+
+    def testRetry(self):
+        def train(config, checkpoint=None):
+            restored = bool(checkpoint)
+            itr = 0
+            if checkpoint:
+                with open(checkpoint, "r") as f:
+                    itr = int(f.read()) + 1
+
+            for i in range(itr, 10):
+                if i == 5 and not restored:
+                    raise Exception("try to fail me")
+                checkpoint_dir = tune.make_checkpoint_dir(step=i)
+                checkpoint_path = os.path.join(checkpoint_dir, "goodbye")
+                with open(checkpoint_path, "w") as f:
+                    f.write(str(i))
+                tune.save_checkpoint(checkpoint_path)
+                tune.report(test=i, training_iteration=i)
+
+        analysis = tune.run(train, max_failures=3)
+        last_ckpt = analysis.trials[0].checkpoint.value
+        assert "goodbye" in last_ckpt
+        trial_dfs = list(analysis.trial_dataframes.values())
+        assert len(trial_dfs[0]["training_iteration"]) == 10
+
+    def testBlankCheckpoint(self):
+        def train(config, checkpoint=None):
+            restored = bool(checkpoint)
+            itr = 0
+            if checkpoint:
+                with open(checkpoint, "r") as f:
+                    itr = int(f.read()) + 1
+
+            for i in range(itr, 10):
+                if i == 5 and not restored:
+                    raise Exception("try to fail me")
+                checkpoint_dir = tune.make_checkpoint_dir()
+                checkpoint_path = os.path.join(checkpoint_dir, "goodbye")
+                with open(checkpoint_path, "w") as f:
+                    f.write(str(i))
+                tune.save_checkpoint(checkpoint_path)
+                tune.report(test=i, training_iteration=i)
+
+        analysis = tune.run(train, max_failures=3)
+        trial_dfs = list(analysis.trial_dataframes.values())
+        assert len(trial_dfs[0]["training_iteration"]) == 10

python/ray/tune/tests/test_track.py

@@ -17,18 +17,6 @@ class TrackApiTest(unittest.TestCase):
         session.shutdown()
         ray.shutdown()
 
-    def testSessionInitShutdown(self):
-        self.assertTrue(session._session is None)
-
-        # Checks that the singleton _session is created/destroyed
-        # by session.init() and session.shutdown()
-        for _ in range(2):
-            # do it twice to see that we can reopen the session
-            session.init(reporter=None)
-            self.assertTrue(session._session is not None)
-            session.shutdown()
-            self.assertTrue(session._session is None)
-
     def testSoftDeprecation(self):
         """Checks that tune.track.log code does not break."""
         from ray.tune import track