316 lines
10 KiB
Python
316 lines
10 KiB
Python
import pytest
|
|
|
|
from ray.tune.search import BasicVariantGenerator, ConcurrencyLimiter, Searcher
|
|
from ray.tune.search.repeater import Repeater
|
|
from ray.tune.search.search_generator import SearchGenerator
|
|
from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def register_test_trainable():
|
|
register_mock_trainable()
|
|
|
|
|
|
def test_nested_suggestion():
|
|
class TestSuggestion(Searcher):
|
|
def suggest(self, trial_id):
|
|
return {"a": {"b": {"c": {"d": 4, "e": 5}}}}
|
|
|
|
searcher = TestSuggestion()
|
|
alg = SearchGenerator(searcher)
|
|
alg.add_configurations({"test": {"run": MOCK_TRAINABLE_NAME}})
|
|
trial = alg.next_trial()
|
|
assert "e=5" in trial.experiment_tag
|
|
assert "d=4" in trial.experiment_tag
|
|
|
|
|
|
def _repeat_trials(num_samples: int, repeat: int):
|
|
class TestSuggestion(Searcher):
|
|
index = 0
|
|
|
|
def suggest(self, trial_id):
|
|
self.index += 1
|
|
return {"test_variable": 5 + self.index}
|
|
|
|
def on_trial_complete(self, *args, **kwargs):
|
|
return
|
|
|
|
searcher = TestSuggestion(metric="episode_reward_mean")
|
|
repeat_searcher = Repeater(searcher, repeat=repeat, set_index=False)
|
|
alg = SearchGenerator(repeat_searcher)
|
|
|
|
alg.add_configurations(
|
|
{
|
|
"test": {
|
|
"run": MOCK_TRAINABLE_NAME,
|
|
"num_samples": num_samples,
|
|
"stop": {"training_iteration": 1},
|
|
}
|
|
}
|
|
)
|
|
|
|
trials = []
|
|
while not alg.is_finished():
|
|
trials.append(alg.next_trial())
|
|
|
|
return trials
|
|
|
|
|
|
def test_repeat_1():
|
|
trials = _repeat_trials(num_samples=2, repeat=1)
|
|
assert len(trials) == 2
|
|
parameter_set = {t.evaluated_params["test_variable"] for t in trials}
|
|
assert len(parameter_set) == 2
|
|
|
|
|
|
def test_repeat_4():
|
|
trials = _repeat_trials(num_samples=12, repeat=4)
|
|
assert len(trials) == 12
|
|
parameter_set = {t.evaluated_params["test_variable"] for t in trials}
|
|
assert len(parameter_set) == 3
|
|
|
|
|
|
def test_odd_repeat():
|
|
trials = _repeat_trials(num_samples=11, repeat=5)
|
|
assert len(trials) == 11
|
|
parameter_set = {t.evaluated_params["test_variable"] for t in trials}
|
|
assert len(parameter_set) == 3
|
|
|
|
|
|
def test_set_get_repeater():
|
|
class TestSuggestion(Searcher):
|
|
def __init__(self, index):
|
|
self.index = index
|
|
self.returned_result = []
|
|
super().__init__(metric="result", mode="max")
|
|
|
|
def suggest(self, trial_id):
|
|
self.index += 1
|
|
return {"score": self.index}
|
|
|
|
def on_trial_complete(self, trial_id, result=None, **kwargs):
|
|
self.returned_result.append(result)
|
|
|
|
searcher = TestSuggestion(0)
|
|
repeater1 = Repeater(searcher, repeat=3, set_index=False)
|
|
for i in range(3):
|
|
assert repeater1.suggest(f"test_{i}")["score"] == 1
|
|
for i in range(2): # An incomplete set of results
|
|
assert repeater1.suggest(f"test_{i}_2")["score"] == 2
|
|
|
|
# Restore a new one
|
|
state = repeater1.get_state()
|
|
del repeater1
|
|
new_repeater = Repeater(searcher, repeat=1, set_index=True)
|
|
new_repeater.set_state(state)
|
|
assert new_repeater.repeat == 3
|
|
assert new_repeater.suggest("test_2_2")["score"] == 2
|
|
assert new_repeater.suggest("test_x")["score"] == 3
|
|
|
|
# Report results
|
|
for i in range(3):
|
|
new_repeater.on_trial_complete(f"test_{i}", {"result": 2})
|
|
|
|
for i in range(3):
|
|
new_repeater.on_trial_complete(f"test_{i}_2", {"result": -i * 10})
|
|
|
|
assert len(new_repeater.searcher.returned_result) == 2
|
|
assert new_repeater.searcher.returned_result[-1] == {"result": -10}
|
|
|
|
# Finish the rest of the last trial group
|
|
new_repeater.on_trial_complete("test_x", {"result": 3})
|
|
assert new_repeater.suggest("test_y")["score"] == 3
|
|
new_repeater.on_trial_complete("test_y", {"result": 3})
|
|
assert len(new_repeater.searcher.returned_result) == 2
|
|
assert new_repeater.suggest("test_z")["score"] == 3
|
|
new_repeater.on_trial_complete("test_z", {"result": 3})
|
|
assert len(new_repeater.searcher.returned_result) == 3
|
|
assert new_repeater.searcher.returned_result[-1] == {"result": 3}
|
|
|
|
|
|
def test_set_get_limiter():
|
|
class TestSuggestion(Searcher):
|
|
def __init__(self, index):
|
|
self.index = index
|
|
self.returned_result = []
|
|
super().__init__(metric="result", mode="max")
|
|
|
|
def suggest(self, trial_id):
|
|
self.index += 1
|
|
return {"score": self.index}
|
|
|
|
def on_trial_complete(self, trial_id, result=None, **kwargs):
|
|
self.returned_result.append(result)
|
|
|
|
searcher = TestSuggestion(0)
|
|
limiter = ConcurrencyLimiter(searcher, max_concurrent=2)
|
|
assert limiter.suggest("test_1")["score"] == 1
|
|
assert limiter.suggest("test_2")["score"] == 2
|
|
assert limiter.suggest("test_3") is None
|
|
|
|
state = limiter.get_state()
|
|
del limiter
|
|
limiter2 = ConcurrencyLimiter(searcher, max_concurrent=3)
|
|
limiter2.set_state(state)
|
|
assert limiter2.suggest("test_4") is None
|
|
assert limiter2.suggest("test_5") is None
|
|
limiter2.on_trial_complete("test_1", {"result": 3})
|
|
limiter2.on_trial_complete("test_2", {"result": 3})
|
|
assert limiter2.suggest("test_3")["score"] == 3
|
|
|
|
|
|
def test_basic_variant_limiter():
|
|
search_alg = BasicVariantGenerator(max_concurrent=2)
|
|
|
|
experiment_spec = {
|
|
"run": MOCK_TRAINABLE_NAME,
|
|
"num_samples": 5,
|
|
"stop": {"training_iteration": 1},
|
|
}
|
|
search_alg.add_configurations({"test": experiment_spec})
|
|
|
|
trial1 = search_alg.next_trial()
|
|
assert trial1
|
|
|
|
trial2 = search_alg.next_trial()
|
|
assert trial2
|
|
|
|
# Returns None because of limiting
|
|
trial3 = search_alg.next_trial()
|
|
assert not trial3
|
|
|
|
# Finish trial, now trial 3 should be created
|
|
search_alg.on_trial_complete(trial1.trial_id, None, False)
|
|
trial3 = search_alg.next_trial()
|
|
assert trial3
|
|
|
|
trial4 = search_alg.next_trial()
|
|
assert not trial4
|
|
|
|
search_alg.on_trial_complete(trial2.trial_id, None, False)
|
|
search_alg.on_trial_complete(trial3.trial_id, None, False)
|
|
|
|
trial4 = search_alg.next_trial()
|
|
assert trial4
|
|
|
|
trial5 = search_alg.next_trial()
|
|
assert trial5
|
|
|
|
search_alg.on_trial_complete(trial4.trial_id, None, False)
|
|
|
|
# Should also be None because search is finished
|
|
trial6 = search_alg.next_trial()
|
|
assert not trial6
|
|
|
|
|
|
def test_batch_limiter():
|
|
class TestSuggestion(Searcher):
|
|
def __init__(self, index):
|
|
self.index = index
|
|
self.returned_result = []
|
|
super().__init__(metric="result", mode="max")
|
|
|
|
def suggest(self, trial_id):
|
|
self.index += 1
|
|
return {"score": self.index}
|
|
|
|
def on_trial_complete(self, trial_id, result=None, **kwargs):
|
|
self.returned_result.append(result)
|
|
|
|
searcher = TestSuggestion(0)
|
|
limiter = ConcurrencyLimiter(searcher, max_concurrent=2, batch=True)
|
|
assert limiter.suggest("test_1")["score"] == 1
|
|
assert limiter.suggest("test_2")["score"] == 2
|
|
assert limiter.suggest("test_3") is None
|
|
|
|
limiter.on_trial_complete("test_1", {"result": 3})
|
|
assert limiter.suggest("test_3") is None
|
|
limiter.on_trial_complete("test_2", {"result": 3})
|
|
assert limiter.suggest("test_3") is not None
|
|
|
|
|
|
def test_batch_limiter_infinite_loop():
|
|
"""Check whether an infinite loop when less than max_concurrent trials
|
|
are suggested with batch mode is avoided.
|
|
"""
|
|
|
|
class TestSuggestion(Searcher):
|
|
def __init__(self, index, max_suggestions=10):
|
|
self.index = index
|
|
self.max_suggestions = max_suggestions
|
|
self.returned_result = []
|
|
super().__init__(metric="result", mode="max")
|
|
|
|
def suggest(self, trial_id):
|
|
self.index += 1
|
|
if self.index > self.max_suggestions:
|
|
return None
|
|
return {"score": self.index}
|
|
|
|
def on_trial_complete(self, trial_id, result=None, **kwargs):
|
|
self.returned_result.append(result)
|
|
self.index = 0
|
|
|
|
searcher = TestSuggestion(0, 2)
|
|
limiter = ConcurrencyLimiter(searcher, max_concurrent=5, batch=True)
|
|
limiter.suggest("test_1")
|
|
limiter.suggest("test_2")
|
|
limiter.suggest("test_3") # TestSuggestion return None
|
|
|
|
limiter.on_trial_complete("test_1", {"result": 3})
|
|
limiter.on_trial_complete("test_2", {"result": 3})
|
|
assert limiter.searcher.returned_result
|
|
|
|
searcher = TestSuggestion(0, 10)
|
|
limiter = ConcurrencyLimiter(searcher, max_concurrent=5, batch=True)
|
|
limiter.suggest("test_1")
|
|
limiter.suggest("test_2")
|
|
limiter.suggest("test_3")
|
|
|
|
limiter.on_trial_complete("test_1", {"result": 3})
|
|
limiter.on_trial_complete("test_2", {"result": 3})
|
|
assert not limiter.searcher.returned_result
|
|
|
|
|
|
def test_set_max_concurrency():
|
|
"""Test whether ``set_max_concurrency`` is called by the
|
|
``ConcurrencyLimiter`` and works correctly.
|
|
"""
|
|
|
|
class TestSuggestion(Searcher):
|
|
def __init__(self, index):
|
|
self.index = index
|
|
self.returned_result = []
|
|
self._max_concurrent = 1
|
|
super().__init__(metric="result", mode="max")
|
|
|
|
def suggest(self, trial_id):
|
|
self.index += 1
|
|
return {"score": self.index}
|
|
|
|
def on_trial_complete(self, trial_id, result=None, **kwargs):
|
|
self.returned_result.append(result)
|
|
|
|
def set_max_concurrency(self, max_concurrent: int) -> bool:
|
|
self._max_concurrent = max_concurrent
|
|
return True
|
|
|
|
searcher = TestSuggestion(0)
|
|
limiter_max_concurrent = 2
|
|
limiter = ConcurrencyLimiter(
|
|
searcher, max_concurrent=limiter_max_concurrent, batch=True
|
|
)
|
|
assert limiter.searcher._max_concurrent == limiter_max_concurrent
|
|
# Since set_max_concurrency returns True, ConcurrencyLimiter should not
|
|
# be limiting concurrency itself
|
|
assert not limiter._limit_concurrency
|
|
assert limiter.suggest("test_1")["score"] == 1
|
|
assert limiter.suggest("test_2")["score"] == 2
|
|
assert limiter.suggest("test_3")["score"] == 3
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|