chore: import upstream snapshot with attribution

python/ray/tune/tests/test_trainable.py

@@ -0,0 +1,134 @@
+import json
+import os
+from typing import Dict, Union
+
+import pytest
+
+import ray
+from ray import tune
+from ray.train._internal.storage import StorageContext
+from ray.train.tests.util import create_dict_checkpoint
+from ray.tune.trainable import wrap_function
+
+
+@pytest.fixture
+def ray_start_2_cpus():
+    address_info = ray.init(num_cpus=2)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+class SavingTrainable(tune.Trainable):
+    def __init__(self, return_type: str, *args, **kwargs):
+        self.return_type = return_type
+        super(SavingTrainable, self).__init__(*args, **kwargs)
+
+    def step(self):
+        return {"iter": self.training_iteration}
+
+    def save_checkpoint(self, tmp_checkpoint_dir: str):
+        checkpoint_data = {"data": 1}
+
+        if self.return_type == "object":
+            return checkpoint_data
+
+        subdir = os.path.join(tmp_checkpoint_dir, "subdir")
+        os.makedirs(subdir, exist_ok=True)
+        checkpoint_file = os.path.join(subdir, "checkpoint.pkl")
+        with open(checkpoint_file, "w") as f:
+            f.write(json.dumps(checkpoint_data))
+
+        if self.return_type == "root":
+            return tmp_checkpoint_dir
+        elif self.return_type == "subdir":
+            return subdir
+        elif self.return_type == "checkpoint":
+            return checkpoint_file
+
+    def load_checkpoint(self, checkpoint: Union[Dict, str]):
+        if self.return_type == "object":
+            assert isinstance(checkpoint, dict)
+            checkpoint_data = checkpoint
+            checkpoint_file = None
+        elif self.return_type == "root":
+            assert "subdir" not in checkpoint
+            checkpoint_file = os.path.join(checkpoint, "subdir", "checkpoint.pkl")
+        elif self.return_type == "subdir":
+            assert "subdir" in checkpoint
+            assert "checkpoint.pkl" not in checkpoint
+            checkpoint_file = os.path.join(checkpoint, "checkpoint.pkl")
+        else:  # self.return_type == "checkpoint"
+            assert checkpoint.endswith("subdir/checkpoint.pkl")
+            checkpoint_file = checkpoint
+
+        if checkpoint_file:
+            with open(checkpoint_file, "rb") as f:
+                checkpoint_data = json.load(f)
+
+        checkpoint_data = {
+            key: value
+            for key, value in checkpoint_data.items()
+            if not key.startswith("_")
+        }
+        assert checkpoint_data == {"data": 1}, checkpoint_data
+
+
+def function_trainable(config):
+    with create_dict_checkpoint({"checkpoint_data": 5}) as checkpoint:
+        tune.report({"metric": 4}, checkpoint=checkpoint)
+
+
+@pytest.mark.parametrize("return_type", ["object", "root"])
+def test_save_load_checkpoint_path_class(ray_start_2_cpus, return_type, tmpdir):
+    """Assert that restoring from a Trainable.save() future works with
+    class trainables.
+
+    Needs Ray cluster so we get actual futures.
+    """
+    trainable = ray.remote(SavingTrainable).remote(return_type=return_type)
+
+    # Train one step
+    ray.get(trainable.train.remote())
+
+    # Save checkpoint
+    saving_future = trainable.save.remote()
+
+    # Check for errors
+    ray.get(saving_future)
+
+    restoring_future = trainable.restore.remote(saving_future)
+
+    ray.get(restoring_future)
+
+
+def test_save_load_checkpoint_path_fn(ray_start_2_cpus, tmp_path):
+    """Assert that restoring from a Trainable.save() future works with
+    function trainables.
+
+    Needs Ray cluster so we get actual futures.
+    """
+    trainable_cls = wrap_function(function_trainable)
+    trainable = ray.remote(trainable_cls).remote(
+        storage=StorageContext(
+            storage_path=str(tmp_path),
+            experiment_dir_name="exp",
+            trial_dir_name="trial",
+        )
+    )
+    ray.get(trainable.train.remote())
+
+    saving_future = trainable.save.remote()
+
+    # Check for errors
+    ray.get(saving_future)
+
+    restoring_future = trainable.restore.remote(saving_future)
+
+    ray.get(restoring_future)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))