wehub-resource-sync
278 lines
9.3 KiB
Python
278 lines
9.3 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
|
|
import ray
|
|
from ray import tune
|
|
from ray.air.constants import TRAINING_ITERATION
|
|
from ray.rllib import _register_all
|
|
from ray.train.tests.util import mock_storage_context
|
|
from ray.tune import Checkpoint, CheckpointConfig
|
|
from ray.tune.execution.placement_groups import PlacementGroupFactory
|
|
from ray.tune.result import DEFAULT_METRIC
|
|
from ray.tune.schedulers import ResourceChangingScheduler
|
|
from ray.tune.trainable import with_parameters, wrap_function
|
|
|
|
|
|
class FunctionCheckpointingTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmpdir = tempfile.TemporaryDirectory()
|
|
|
|
def create_trainable(self, train_fn):
|
|
return wrap_function(train_fn)(storage=mock_storage_context())
|
|
|
|
def tearDown(self):
|
|
self.tmpdir.cleanup()
|
|
|
|
def testCheckpointReuse(self):
|
|
"""Test that repeated save/restore never reuses same checkpoint dir."""
|
|
|
|
def train_fn(config):
|
|
checkpoint = ray.tune.get_checkpoint()
|
|
if checkpoint:
|
|
with checkpoint.as_directory() as checkpoint_dir:
|
|
count = sum(
|
|
"checkpoint-" in path for path in os.listdir(checkpoint_dir)
|
|
)
|
|
assert count == 1, os.listdir(checkpoint_dir)
|
|
|
|
for step in range(20):
|
|
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
|
path = os.path.join(
|
|
temp_checkpoint_dir, "checkpoint-{}".format(step)
|
|
)
|
|
open(path, "a").close()
|
|
ray.tune.report(
|
|
dict(test=step),
|
|
checkpoint=Checkpoint.from_directory(temp_checkpoint_dir),
|
|
)
|
|
|
|
checkpoint = None
|
|
for i in range(5):
|
|
new_trainable = self.create_trainable(train_fn)
|
|
if checkpoint:
|
|
new_trainable.restore(checkpoint)
|
|
for i in range(2):
|
|
result = new_trainable.train()
|
|
checkpoint = new_trainable.save()
|
|
new_trainable.stop()
|
|
assert result[TRAINING_ITERATION] == 10
|
|
|
|
def testFunctionRecurringSave(self):
|
|
"""This tests that save and restore are commutative."""
|
|
|
|
def train_fn(config):
|
|
for step in range(10):
|
|
with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
|
|
if step % 3 == 0:
|
|
path = os.path.join(temp_checkpoint_dir, "checkpoint.json")
|
|
with open(path, "w") as f:
|
|
json.dump({"step": step}, f)
|
|
ray.tune.report(
|
|
dict(test=step),
|
|
checkpoint=Checkpoint.from_directory(temp_checkpoint_dir),
|
|
)
|
|
|
|
new_trainable = self.create_trainable(train_fn)
|
|
new_trainable.train()
|
|
checkpoint_obj = new_trainable.save()
|
|
new_trainable.restore(checkpoint_obj)
|
|
checkpoint = new_trainable.save()
|
|
|
|
new_trainable.stop()
|
|
|
|
new_trainable2 = self.create_trainable(train_fn)
|
|
new_trainable2.restore(checkpoint)
|
|
new_trainable2.train()
|
|
new_trainable2.stop()
|
|
|
|
|
|
class FunctionApiTest(unittest.TestCase):
|
|
def setUp(self):
|
|
ray.init(num_cpus=4, num_gpus=0, object_store_memory=150 * 1024 * 1024)
|
|
|
|
def tearDown(self):
|
|
ray.shutdown()
|
|
_register_all() # re-register the evicted objects
|
|
|
|
def testCheckpointError(self):
|
|
def train_fn(config):
|
|
pass
|
|
|
|
with self.assertRaises(ValueError):
|
|
tune.run(
|
|
train_fn, checkpoint_config=CheckpointConfig(checkpoint_frequency=1)
|
|
)
|
|
with self.assertRaises(ValueError):
|
|
tune.run(
|
|
train_fn, checkpoint_config=CheckpointConfig(checkpoint_at_end=True)
|
|
)
|
|
|
|
def testWithParameters(self):
|
|
class Data:
|
|
def __init__(self):
|
|
self.data = [0] * 500_000
|
|
|
|
data = Data()
|
|
data.data[100] = 1
|
|
|
|
def train_fn(config, data=None):
|
|
data.data[101] = 2 # Changes are local
|
|
ray.tune.report(dict(metric=len(data.data), hundred=data.data[100]))
|
|
|
|
trial_1, trial_2 = tune.run(
|
|
with_parameters(train_fn, data=data), num_samples=2
|
|
).trials
|
|
|
|
self.assertEqual(data.data[101], 0)
|
|
self.assertEqual(trial_1.last_result["metric"], 500_000)
|
|
self.assertEqual(trial_1.last_result["hundred"], 1)
|
|
self.assertEqual(trial_2.last_result["metric"], 500_000)
|
|
self.assertEqual(trial_2.last_result["hundred"], 1)
|
|
self.assertTrue(str(trial_1).startswith("train_"))
|
|
|
|
# With checkpoint dir parameter
|
|
def train_fn(config, data=None):
|
|
data.data[101] = 2 # Changes are local
|
|
ray.tune.report(dict(metric=len(data.data)))
|
|
|
|
trial_1, trial_2 = tune.run(
|
|
with_parameters(train_fn, data=data), num_samples=2
|
|
).trials
|
|
|
|
self.assertEqual(data.data[101], 0)
|
|
self.assertEqual(trial_1.last_result["metric"], 500_000)
|
|
self.assertEqual(trial_2.last_result["metric"], 500_000)
|
|
self.assertTrue(str(trial_1).startswith("train_"))
|
|
|
|
def testWithParameters2(self):
|
|
class Data:
|
|
def __init__(self):
|
|
import numpy as np
|
|
|
|
self.data = np.random.rand((2 * 1024 * 1024))
|
|
|
|
def train_fn(config, data=None):
|
|
pass
|
|
|
|
trainable = tune.with_parameters(train_fn, data=Data())
|
|
# ray.cloudpickle will crash for some reason
|
|
import cloudpickle as cp
|
|
|
|
dumped = cp.dumps(trainable)
|
|
assert sys.getsizeof(dumped) < 100 * 1024
|
|
|
|
def testNewResources(self):
|
|
sched = ResourceChangingScheduler(
|
|
resources_allocation_function=(
|
|
lambda a, b, c, d: PlacementGroupFactory([{"CPU": 2}])
|
|
)
|
|
)
|
|
|
|
def train_fn(config):
|
|
ray.tune.report(
|
|
dict(metric=1, resources=ray.tune.get_context().get_trial_resources())
|
|
)
|
|
|
|
analysis = tune.run(
|
|
train_fn,
|
|
scheduler=sched,
|
|
stop={"training_iteration": 2},
|
|
resources_per_trial=PlacementGroupFactory([{"CPU": 1}]),
|
|
num_samples=1,
|
|
)
|
|
|
|
results_list = list(analysis.results.values())
|
|
assert results_list[0]["resources"].head_cpus == 2.0
|
|
|
|
def testWithParametersTwoRuns1(self):
|
|
# Makes sure two runs in the same script but different ray sessions
|
|
# pass (https://github.com/ray-project/ray/issues/16609)
|
|
def train_fn(config, extra=4):
|
|
ray.tune.report(dict(metric=extra))
|
|
|
|
trainable = tune.with_parameters(train_fn, extra=8)
|
|
out = tune.run(trainable, metric="metric", mode="max")
|
|
self.assertEqual(out.best_result["metric"], 8)
|
|
|
|
self.tearDown()
|
|
self.setUp()
|
|
|
|
def train_fn_2(config, extra=5):
|
|
ray.tune.report(dict(metric=extra))
|
|
|
|
trainable = tune.with_parameters(train_fn_2, extra=9)
|
|
out = tune.run(trainable, metric="metric", mode="max")
|
|
self.assertEqual(out.best_result["metric"], 9)
|
|
|
|
def testWithParametersTwoRuns2(self):
|
|
# Makes sure two runs in the same script
|
|
# pass (https://github.com/ray-project/ray/issues/16609)
|
|
def train_fn(config, extra=4):
|
|
ray.tune.report(dict(metric=extra))
|
|
|
|
def train_fn_2(config, extra=5):
|
|
ray.tune.report(dict(metric=extra))
|
|
|
|
trainable1 = tune.with_parameters(train_fn, extra=8)
|
|
trainable2 = tune.with_parameters(train_fn_2, extra=9)
|
|
|
|
out1 = tune.run(trainable1, metric="metric", mode="max")
|
|
out2 = tune.run(trainable2, metric="metric", mode="max")
|
|
self.assertEqual(out1.best_result["metric"], 8)
|
|
self.assertEqual(out2.best_result["metric"], 9)
|
|
|
|
def testReturnAnonymous(self):
|
|
def train_fn(config):
|
|
return config["a"]
|
|
|
|
trial_1, trial_2 = tune.run(
|
|
train_fn, config={"a": tune.grid_search([4, 8])}
|
|
).trials
|
|
|
|
self.assertEqual(trial_1.last_result[DEFAULT_METRIC], 4)
|
|
self.assertEqual(trial_2.last_result[DEFAULT_METRIC], 8)
|
|
|
|
def testReturnSpecific(self):
|
|
def train_fn(config):
|
|
return {"m": config["a"]}
|
|
|
|
trial_1, trial_2 = tune.run(
|
|
train_fn, config={"a": tune.grid_search([4, 8])}
|
|
).trials
|
|
|
|
self.assertEqual(trial_1.last_result["m"], 4)
|
|
self.assertEqual(trial_2.last_result["m"], 8)
|
|
|
|
def testYieldAnonymous(self):
|
|
def train_fn(config):
|
|
for i in range(10):
|
|
yield config["a"] + i
|
|
|
|
trial_1, trial_2 = tune.run(
|
|
train_fn, config={"a": tune.grid_search([4, 8])}
|
|
).trials
|
|
|
|
self.assertEqual(trial_1.last_result[DEFAULT_METRIC], 4 + 9)
|
|
self.assertEqual(trial_2.last_result[DEFAULT_METRIC], 8 + 9)
|
|
|
|
def testYieldSpecific(self):
|
|
def train_fn(config):
|
|
for i in range(10):
|
|
yield {"m": config["a"] + i}
|
|
|
|
trial_1, trial_2 = tune.run(
|
|
train_fn, config={"a": tune.grid_search([4, 8])}
|
|
).trials
|
|
|
|
self.assertEqual(trial_1.last_result["m"], 4 + 9)
|
|
self.assertEqual(trial_2.last_result["m"], 8 + 9)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|