chore: import upstream snapshot with attribution

python/ray/train/tests/util.py

@@ -0,0 +1,52 @@
+import contextlib
+import os
+import tempfile
+from typing import Any, Dict, Optional, Type
+
+import ray.cloudpickle as ray_pickle
+from ray.train import Checkpoint, SyncConfig
+from ray.train._internal.storage import StorageContext
+
+
+@contextlib.contextmanager
+def create_dict_checkpoint(
+    data: Dict[str, Any], checkpoint_cls: Type[Checkpoint] = None
+) -> Checkpoint:
+    with tempfile.TemporaryDirectory() as tmpdir:
+        with open(os.path.join(tmpdir, "data.pkl"), "wb") as f:
+            ray_pickle.dump(data, f)
+        checkpoint_cls = checkpoint_cls or Checkpoint
+        yield checkpoint_cls.from_directory(tmpdir)
+
+
+def load_dict_checkpoint(checkpoint: Checkpoint) -> Dict[str, Any]:
+    with checkpoint.as_directory() as checkpoint_dir:
+        with open(os.path.join(checkpoint_dir, "data.pkl"), "rb") as f:
+            return ray_pickle.load(f)
+
+
+def mock_storage_context(
+    exp_name: str = "exp_name",
+    storage_path: Optional[str] = None,
+    storage_context_cls: Type = StorageContext,
+    sync_config: Optional[SyncConfig] = None,
+) -> StorageContext:
+    storage_path = storage_path or tempfile.mkdtemp()
+    exp_name = exp_name
+    trial_name = "trial_name"
+
+    storage = storage_context_cls(
+        storage_path=storage_path,
+        experiment_dir_name=exp_name,
+        trial_dir_name=trial_name,
+        sync_config=sync_config,
+    )
+    # Patch the default /tmp/ray/session_* so we don't require ray
+    # to be initialized in unit tests.
+    session_path = tempfile.mkdtemp()
+    storage._get_session_path = lambda: session_path
+
+    os.makedirs(storage.trial_fs_path, exist_ok=True)
+    os.makedirs(storage.trial_driver_staging_path, exist_ok=True)
+
+    return storage