chore: import upstream snapshot with attribution

python/ray/tune/tests/test_tune_save_restore.py

@@ -0,0 +1,144 @@
+# coding: utf-8
+import os
+import pickle
+import shutil
+import tempfile
+import unittest
+
+import ray
+from ray import tune
+from ray.tune import CheckpointConfig, Trainable
+from ray.tune.utils import validate_save_restore
+
+
+class SerialTuneRelativeLocalDirTest(unittest.TestCase):
+    prefix = "Serial"
+
+    class MockTrainable(Trainable):
+        _name = "MockTrainable"
+
+        def setup(self, config):
+            self.state = {"hi": 1}
+
+        def step(self):
+            return {"timesteps_this_iter": 1, "done": True}
+
+        def save_checkpoint(self, checkpoint_dir):
+            checkpoint_path = os.path.join(checkpoint_dir, "checkpoint.pkl")
+            with open(checkpoint_path, "wb") as f:
+                pickle.dump(self.state, f)
+
+        def load_checkpoint(self, checkpoint_dir):
+            checkpoint_path = os.path.join(checkpoint_dir, "checkpoint.pkl")
+            with open(checkpoint_path, "rb") as f:
+                extra_data = pickle.load(f)
+            self.state.update(extra_data)
+
+    def setUp(self):
+        self.absolute_local_dir = None
+        ray.init(num_cpus=2, num_gpus=0)
+
+    def tearDown(self):
+        if self.absolute_local_dir is not None:
+            shutil.rmtree(self.absolute_local_dir, ignore_errors=True)
+            self.absolute_local_dir = None
+        ray.shutdown()
+
+    def _get_trial_dir(self, absoulte_exp_dir):
+        print("looking for", self.MockTrainable._name)
+        print("in", os.listdir(absoulte_exp_dir))
+        trial_dirname = next(
+            (
+                child_dir
+                for child_dir in os.listdir(absoulte_exp_dir)
+                if (
+                    os.path.isdir(os.path.join(absoulte_exp_dir, child_dir))
+                    and child_dir.startswith(self.MockTrainable._name)
+                )
+            )
+        )
+
+        trial_absolute_dir = os.path.join(absoulte_exp_dir, trial_dirname)
+
+        return trial_dirname, trial_absolute_dir
+
+    def _train(self, exp_name, local_dir, absolute_local_dir):
+        (trial,) = tune.run(
+            self.MockTrainable,
+            name=exp_name,
+            stop={"training_iteration": 1},
+            checkpoint_config=CheckpointConfig(checkpoint_frequency=1),
+            storage_path=local_dir,
+            config={"env": "CartPole-v0", "log_level": "DEBUG"},
+        ).trials
+
+        exp_dir = os.path.join(absolute_local_dir, exp_name)
+        _, abs_trial_dir = self._get_trial_dir(exp_dir)
+
+        self.assertIsNone(trial.error_file)
+        self.assertEqual(trial.path, abs_trial_dir)
+
+        self.assertTrue(os.path.isdir(absolute_local_dir), absolute_local_dir)
+        self.assertTrue(os.path.isdir(exp_dir))
+        self.assertTrue(os.path.isdir(abs_trial_dir))
+        self.assertTrue(
+            os.path.isfile(
+                os.path.join(abs_trial_dir, "checkpoint_000000/checkpoint.pkl")
+            )
+        )
+
+    def _restore(self, exp_name, local_dir, absolute_local_dir):
+        trial_name, abs_trial_dir = self._get_trial_dir(
+            os.path.join(absolute_local_dir, exp_name)
+        )
+
+        checkpoint_path = os.path.join(
+            local_dir, exp_name, trial_name, "checkpoint_000000"
+        )
+        assert os.path.exists(os.path.expanduser(checkpoint_path))
+
+        (trial,) = tune.run(
+            self.MockTrainable,
+            name=exp_name,
+            stop={"training_iteration": 2},  # train one more iteration.
+            restore=checkpoint_path,  # Restore the checkpoint
+            config={"env": "CartPole-v0", "log_level": "DEBUG"},
+        ).trials
+        self.assertIsNone(trial.error_file)
+
+    def testTempfile(self):
+        local_dir = tempfile.mkdtemp()
+        exp_name = self.prefix + "Tempfile"
+        self.absolute_local_dir = local_dir
+        self._train(exp_name, local_dir, local_dir)
+        self._restore(exp_name, local_dir, local_dir)
+
+    def testCheckpointWithNoop(self):
+        """Tests that passing the checkpoint_dir right back works."""
+
+        class MockTrainable(Trainable):
+            def setup(self, config):
+                pass
+
+            def step(self):
+                return {"score": 1}
+
+            def save_checkpoint(self, checkpoint_dir):
+                with open(os.path.join(checkpoint_dir, "test.txt"), "wb") as f:
+                    pickle.dump("test", f)
+
+            def load_checkpoint(self, checkpoint_dir):
+                with open(os.path.join(checkpoint_dir, "test.txt"), "rb") as f:
+                    x = pickle.load(f)
+
+                assert x == "test"
+
+        validate_save_restore(MockTrainable)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))