wehub-resource-sync
279 lines
9.0 KiB
Python
279 lines
9.0 KiB
Python
import os
|
|
import sys
|
|
import time
|
|
from collections import Counter
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray import tune
|
|
from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
|
|
from ray.train.tests.util import mock_storage_context
|
|
from ray.tune import PlacementGroupFactory, TuneError
|
|
from ray.tune.execution.tune_controller import TuneController
|
|
from ray.tune.experiment import Trial
|
|
from ray.tune.schedulers import FIFOScheduler, TrialScheduler
|
|
from ray.tune.search import BasicVariantGenerator
|
|
from ray.tune.utils.mock import TrialStatusSnapshot, TrialStatusSnapshotTaker
|
|
from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable
|
|
|
|
STORAGE = mock_storage_context()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def register_test_trainable():
|
|
register_mock_trainable()
|
|
|
|
|
|
@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]
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"bundles",
|
|
[
|
|
[{"CPU": 1}, {"CPU": 3, "GPU": 1}],
|
|
[{"CPU": 1, "a": 2}],
|
|
[{"CPU": 1}, {"a": 2}],
|
|
[{"CPU": 1, "GPU": 1}, {"GPU": 1}],
|
|
],
|
|
)
|
|
def test_resource_parallelism_single(
|
|
ray_start_4_cpus_2_gpus_extra, resource_manager_cls, bundles
|
|
):
|
|
"""Test that extra and custom resources are respected for parallelism.
|
|
|
|
We schedule two trials with resources according to the bundle. If only
|
|
the head bundle or only CPU/GPU resources were considered, both trials
|
|
could run in parallel.
|
|
|
|
However, we assert that the resources in child bundles and extra resources
|
|
are respected and only one trial runs in parallel.
|
|
|
|
Legacy test: test_trial_runner.py::TrialRunnerTest::testExtraResources
|
|
Legacy test: test_trial_runner.py::TrialRunnerTest::testCustomResources
|
|
Legacy test: test_trial_runner.py::TrialRunnerTest::testExtraCustomResources
|
|
Legacy test: test_trial_runner.py::TrialRunnerTest::testResourceScheduler
|
|
"""
|
|
snapshot = TrialStatusSnapshot()
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
callbacks=[TrialStatusSnapshotTaker(snapshot)],
|
|
storage=STORAGE,
|
|
)
|
|
kwargs = {
|
|
"stopping_criterion": {"training_iteration": 1},
|
|
"placement_group_factory": PlacementGroupFactory(bundles),
|
|
"storage": STORAGE,
|
|
}
|
|
trials = [
|
|
Trial(MOCK_TRAINABLE_NAME, **kwargs),
|
|
Trial(MOCK_TRAINABLE_NAME, **kwargs),
|
|
]
|
|
for t in trials:
|
|
runner.add_trial(t)
|
|
|
|
while not runner.is_finished():
|
|
runner.step()
|
|
|
|
assert snapshot.max_running_trials() == 1
|
|
assert snapshot.all_trials_are_terminated()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_fractional_gpus(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
|
|
"""Test that fractional GPUs lead to more parallelism.
|
|
|
|
We schedule four trials with 0.75 GPUs each. Since our cluster has 2 GPUs,
|
|
we should be able to run 2 trials in parallel.
|
|
|
|
Legacy test: test_trial_runner.py::TrialRunnerTest::testFractionalGpus
|
|
"""
|
|
snapshot = TrialStatusSnapshot()
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
callbacks=[TrialStatusSnapshotTaker(snapshot)],
|
|
storage=STORAGE,
|
|
)
|
|
kwargs = {
|
|
"stopping_criterion": {"training_iteration": 1},
|
|
"placement_group_factory": PlacementGroupFactory([{"GPU": 0.75}]),
|
|
"config": {
|
|
"sleep": 1,
|
|
},
|
|
"storage": STORAGE,
|
|
}
|
|
trials = [Trial(MOCK_TRAINABLE_NAME, **kwargs) for i in range(4)]
|
|
for t in trials:
|
|
runner.add_trial(t)
|
|
|
|
while not runner.is_finished():
|
|
runner.step()
|
|
|
|
assert snapshot.max_running_trials() == 2
|
|
assert snapshot.all_trials_are_terminated()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_multi_step(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
|
|
"""Test that trials can run for more than one iteration.
|
|
|
|
Todo (krfricke): This is not a resource test, so it should be moved.
|
|
|
|
Legacy test: test_trial_runner.py::TrialRunnerTest::testMultiStepRun
|
|
Legacy test: test_trial_runner.py::TrialRunnerTest::testMultiStepRun2
|
|
"""
|
|
snapshot = TrialStatusSnapshot()
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
callbacks=[TrialStatusSnapshotTaker(snapshot)],
|
|
storage=STORAGE,
|
|
)
|
|
kwargs = {
|
|
"stopping_criterion": {"training_iteration": 5},
|
|
"placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
|
|
"storage": STORAGE,
|
|
}
|
|
trials = [Trial(MOCK_TRAINABLE_NAME, **kwargs) for i in range(2)]
|
|
for t in trials:
|
|
runner.add_trial(t)
|
|
|
|
while not runner.is_finished():
|
|
runner.step()
|
|
|
|
# Overstepping should throw error
|
|
# test_trial_runner.py::TrialRunnerTest::testMultiStepRun2
|
|
with pytest.raises(TuneError):
|
|
runner.step()
|
|
|
|
assert snapshot.all_trials_are_terminated()
|
|
assert all(t.last_result["training_iteration"] == 5 for t in runner.get_trials())
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_resources_changing(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
|
|
"""Checks that resource requirements can be changed on fly.
|
|
|
|
Legacy test: test_trial_runner.py::TrialRunnerTest::testChangeResources
|
|
"""
|
|
|
|
class ChangingScheduler(FIFOScheduler):
|
|
def on_trial_result(self, tune_controller, trial, result):
|
|
if result["training_iteration"] == 1:
|
|
# NOTE: This is a hack to get around the new pausing logic,
|
|
# which doesn't set the trial status to PAUSED immediately.
|
|
orig_status = trial.status
|
|
trial.set_status(Trial.PAUSED)
|
|
trial.update_resources(dict(cpu=4, gpu=0))
|
|
trial.set_status(orig_status)
|
|
return TrialScheduler.PAUSE
|
|
return TrialScheduler.NOOP
|
|
|
|
scheduler = ChangingScheduler()
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
scheduler=scheduler,
|
|
storage=STORAGE,
|
|
)
|
|
kwargs = {
|
|
"stopping_criterion": {"training_iteration": 2},
|
|
"placement_group_factory": PlacementGroupFactory([{"CPU": 2, "GPU": 0}]),
|
|
"storage": STORAGE,
|
|
}
|
|
trials = [Trial(MOCK_TRAINABLE_NAME, **kwargs)]
|
|
for t in trials:
|
|
runner.add_trial(t)
|
|
|
|
while not trials[0].status == Trial.RUNNING:
|
|
runner.step()
|
|
|
|
assert trials[0].status == Trial.RUNNING
|
|
assert runner._actor_manager.get_live_actors_resources().get("CPU") == 2
|
|
|
|
with pytest.raises(ValueError):
|
|
trials[0].update_resources(dict(cpu=4, gpu=0))
|
|
|
|
while trials[0].status == Trial.RUNNING:
|
|
runner.step()
|
|
|
|
assert trials[0].status == Trial.PAUSED
|
|
|
|
while not trials[0].status == Trial.RUNNING:
|
|
runner.step()
|
|
|
|
assert runner._actor_manager.get_live_actors_resources().get("CPU") == 4
|
|
|
|
runner.step()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
|
|
)
|
|
def test_queue_filling(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
|
|
"""Checks that the trial queue is filled even if only 1 pending trial is allowed.
|
|
|
|
Legacy test: test_trial_runner.py::TrialRunnerTest::testQueueFilling
|
|
"""
|
|
os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"
|
|
|
|
def f1(config):
|
|
for i in range(10):
|
|
yield i
|
|
time.sleep(1)
|
|
|
|
tune.register_trainable("f1", f1)
|
|
|
|
search_alg = BasicVariantGenerator()
|
|
search_alg.add_configurations(
|
|
{
|
|
"foo": {
|
|
"run": "f1",
|
|
"num_samples": 100,
|
|
"config": {
|
|
"a": tune.sample_from(lambda spec: 5.0 / 7),
|
|
"b": tune.sample_from(lambda spec: "long" * 40),
|
|
},
|
|
"resources_per_trial": {"cpu": 2},
|
|
}
|
|
}
|
|
)
|
|
|
|
runner = TuneController(
|
|
resource_manager_factory=lambda: resource_manager_cls(),
|
|
search_alg=search_alg,
|
|
storage=STORAGE,
|
|
)
|
|
|
|
while len(runner.get_trials()) < 3:
|
|
runner.step()
|
|
|
|
# All trials are enqueued
|
|
assert len(runner.get_trials()) == 3
|
|
|
|
status_count = Counter(t.status for t in runner.get_trials())
|
|
while status_count.get(Trial.RUNNING, 0) < 2 and not runner.is_finished():
|
|
runner.step()
|
|
status_count = Counter(t.status for t in runner.get_trials())
|
|
|
|
assert len(runner.get_trials()) == 3
|
|
|
|
status_count = Counter(t.status for t in runner.get_trials())
|
|
assert status_count.get(Trial.RUNNING) == 2
|
|
assert status_count.get(Trial.PENDING) == 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", "--reruns", "3", __file__]))
|