wehub-resource-sync
148 lines
4.6 KiB
Python
148 lines
4.6 KiB
Python
import sys
|
|
from collections import Counter
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
|
|
from ray.train.tests.util import mock_storage_context
|
|
from ray.tune import PlacementGroupFactory, register_trainable
|
|
from ray.tune.execution.tune_controller import TuneController
|
|
from ray.tune.experiment import Trial
|
|
from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
|
|
|
|
STORAGE = mock_storage_context()
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def ray_start_4_cpus_2_gpus_extra():
|
|
address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
|
|
yield address_info
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_stop_trial(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
|
|
"""Stopping a trial while RUNNING or PENDING should work.
|
|
|
|
Legacy test: test_trial_runner_3.py::TrialRunnerTest::testStopTrial
|
|
"""
|
|
|
|
register_mock_trainable()
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(), storage=STORAGE
|
|
)
|
|
kwargs = {
|
|
"stopping_criterion": {"training_iteration": 10},
|
|
"placement_group_factory": PlacementGroupFactory([{"CPU": 2, "GPU": 1}]),
|
|
"config": {"sleep": 1},
|
|
"storage": STORAGE,
|
|
}
|
|
trials = [
|
|
Trial(MOCK_TRAINABLE_NAME, **kwargs),
|
|
Trial(MOCK_TRAINABLE_NAME, **kwargs),
|
|
Trial(MOCK_TRAINABLE_NAME, **kwargs),
|
|
Trial(MOCK_TRAINABLE_NAME, **kwargs),
|
|
]
|
|
for t in trials:
|
|
runner.add_trial(t)
|
|
|
|
counter = Counter(t.status for t in trials)
|
|
|
|
# Wait until 2 trials started
|
|
while counter.get("RUNNING", 0) != 2:
|
|
runner.step()
|
|
counter = Counter(t.status for t in trials)
|
|
|
|
assert counter.get("RUNNING", 0) == 2
|
|
assert counter.get("PENDING", 0) == 2
|
|
|
|
# Stop trial that is running
|
|
for trial in trials:
|
|
if trial.status == Trial.RUNNING:
|
|
runner._schedule_trial_stop(trial)
|
|
break
|
|
|
|
counter = Counter(t.status for t in trials)
|
|
|
|
# Wait until the next trial started
|
|
while counter.get("RUNNING", 0) < 2:
|
|
runner.step()
|
|
counter = Counter(t.status for t in trials)
|
|
|
|
assert counter.get("RUNNING", 0) == 2
|
|
assert counter.get("TERMINATED", 0) == 1
|
|
assert counter.get("PENDING", 0) == 1
|
|
|
|
# Stop trial that is pending
|
|
for trial in trials:
|
|
if trial.status == Trial.PENDING:
|
|
runner._schedule_trial_stop(trial)
|
|
break
|
|
|
|
counter = Counter(t.status for t in trials)
|
|
|
|
# Wait until 2 trials are running again
|
|
while counter.get("RUNNING", 0) < 2:
|
|
runner.step()
|
|
counter = Counter(t.status for t in trials)
|
|
|
|
assert counter.get("RUNNING", 0) == 2
|
|
assert counter.get("TERMINATED", 0) == 2
|
|
assert counter.get("PENDING", 0) == 0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_remove_actor_tracking(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
|
|
"""When we reuse actors, actors that have been requested but not started
|
|
should not be tracked in ``_stopping_actors``.
|
|
|
|
When actors are re-used, we cancel original actor requests for the trial.
|
|
If these actors haven't been alive, there won't be a stop future to be resolved,
|
|
and thus they would remain in ``TuneController._stopping_actors`` until they
|
|
get cleaned up after 600 seconds.
|
|
|
|
This test asserts that these actors are not tracked in
|
|
``TuneController._stopping_actors`` at all.
|
|
|
|
We start 4 actors, and one can run at a time. Actors are re-used across trials.
|
|
When the experiment ends, we expect that only one actor is left to track
|
|
in ``self._stopping_trials``.
|
|
"""
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
reuse_actors=True,
|
|
storage=STORAGE,
|
|
)
|
|
|
|
def train_fn(config):
|
|
return 1
|
|
|
|
register_trainable("test_remove_actor_tracking", train_fn)
|
|
|
|
kwargs = {
|
|
"placement_group_factory": PlacementGroupFactory([{"CPU": 4, "GPU": 2}]),
|
|
"storage": STORAGE,
|
|
}
|
|
trials = [Trial("test_remove_actor_tracking", **kwargs) for i in range(4)]
|
|
for t in trials:
|
|
runner.add_trial(t)
|
|
|
|
while not runner.is_finished():
|
|
runner.step()
|
|
|
|
# Only one actor should be left to stop
|
|
assert len(runner._stopping_actors) == 1
|
|
|
|
runner.cleanup()
|
|
|
|
assert len(runner._stopping_actors) == 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", "--reruns", "3", __file__]))
|