wehub-resource-sync
135 lines
4.0 KiB
Python
135 lines
4.0 KiB
Python
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__]))
|