wehub-resource-sync
110 lines
3.5 KiB
Python
110 lines
3.5 KiB
Python
import inspect
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import ray
|
|
from ray.tune import choice, register_trainable, run, run_experiments
|
|
from ray.tune.experiment import Experiment, Trial
|
|
from ray.tune.result import TIMESTEPS_TOTAL
|
|
from ray.tune.search.hyperopt import HyperOptSearch
|
|
from ray.util.client.ray_client_helpers import ray_start_client_server
|
|
|
|
|
|
def train_fn(config):
|
|
for i in range(100):
|
|
ray.tune.report(dict(timesteps_total=i))
|
|
|
|
|
|
class RemoteTest(unittest.TestCase):
|
|
def tearDown(self):
|
|
ray.shutdown()
|
|
|
|
def testRemoteRunExperiments(self):
|
|
register_trainable("f1", train_fn)
|
|
exp1 = Experiment(
|
|
**{
|
|
"name": "foo",
|
|
"run": "f1",
|
|
}
|
|
)
|
|
[trial] = run_experiments(exp1, _remote=True)
|
|
self.assertEqual(trial.status, Trial.TERMINATED)
|
|
self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
|
|
|
|
def testRemoteRun(self):
|
|
analysis = run(train_fn, _remote=True)
|
|
[trial] = analysis.trials
|
|
self.assertEqual(trial.status, Trial.TERMINATED)
|
|
self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
|
|
|
|
def testRemoteRunArguments(self):
|
|
def mocked_run(*args, **kwargs):
|
|
capture_args_kwargs = (args, kwargs)
|
|
return run(*args, **kwargs), capture_args_kwargs
|
|
|
|
with patch("ray.tune.tune.run", mocked_run):
|
|
analysis, capture_args_kwargs = run(train_fn, _remote=True)
|
|
args, kwargs = capture_args_kwargs
|
|
self.assertFalse(args)
|
|
kwargs.pop("run_or_experiment")
|
|
kwargs.pop("_remote")
|
|
kwargs.pop("progress_reporter") # gets autodetected and set
|
|
|
|
default_kwargs = {
|
|
k: v.default for k, v in inspect.signature(run).parameters.items()
|
|
}
|
|
default_kwargs.pop("run_or_experiment")
|
|
default_kwargs.pop("_remote")
|
|
default_kwargs.pop("progress_reporter")
|
|
|
|
self.assertDictEqual(kwargs, default_kwargs)
|
|
|
|
def testRemoteRunWithSearcher(self):
|
|
analysis = run(
|
|
train_fn,
|
|
search_alg=HyperOptSearch(),
|
|
config={"a": choice(["a", "b"])},
|
|
metric="timesteps_total",
|
|
mode="max",
|
|
_remote=True,
|
|
)
|
|
[trial] = analysis.trials
|
|
self.assertEqual(trial.status, Trial.TERMINATED)
|
|
self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
|
|
|
|
def testRemoteRunExperimentsInClient(self):
|
|
ray.init()
|
|
assert not ray.util.client.ray.is_connected()
|
|
with ray_start_client_server():
|
|
assert ray.util.client.ray.is_connected()
|
|
|
|
register_trainable("f1", train_fn)
|
|
exp1 = Experiment(
|
|
**{
|
|
"name": "foo",
|
|
"run": "f1",
|
|
}
|
|
)
|
|
[trial] = run_experiments(exp1)
|
|
self.assertEqual(trial.status, Trial.TERMINATED)
|
|
self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
|
|
|
|
def testRemoteRunInClient(self):
|
|
ray.init()
|
|
assert not ray.util.client.ray.is_connected()
|
|
with ray_start_client_server():
|
|
assert ray.util.client.ray.is_connected()
|
|
|
|
analysis = run(train_fn)
|
|
[trial] = analysis.trials
|
|
self.assertEqual(trial.status, Trial.TERMINATED)
|
|
self.assertEqual(trial.last_result[TIMESTEPS_TOTAL], 99)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|