wehub-resource-sync
53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
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
|