wehub-resource-sync
319 lines
12 KiB
Python
319 lines
12 KiB
Python
import os
|
|
import sys
|
|
import time
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
import ray
|
|
from ray import tune
|
|
from ray.cluster_utils import Cluster
|
|
from ray.rllib import _register_all
|
|
from ray.tune import Callback
|
|
from ray.tune.execution.placement_groups import PlacementGroupFactory
|
|
from ray.tune.execution.ray_trial_executor import RayTrialExecutor
|
|
from ray.tune.execution.trial_runner import TrialRunner
|
|
from ray.tune.experiment import Trial
|
|
from ray.util import placement_group_table
|
|
|
|
|
|
class TrialRunnerPlacementGroupTest(unittest.TestCase):
|
|
def setUp(self):
|
|
os.environ["TUNE_GLOBAL_CHECKPOINT_S"] = "10000"
|
|
os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "auto" # Reset default
|
|
self.head_cpus = 8
|
|
self.head_gpus = 4
|
|
self.head_custom = 16
|
|
|
|
self.cluster = Cluster(
|
|
initialize_head=True,
|
|
connect=True,
|
|
head_node_args={
|
|
"include_dashboard": False,
|
|
"num_cpus": self.head_cpus,
|
|
"num_gpus": self.head_gpus,
|
|
"resources": {"custom": self.head_custom},
|
|
"_system_config": {
|
|
"health_check_initial_delay_ms": 0,
|
|
"health_check_period_ms": 1000,
|
|
"health_check_failure_threshold": 10,
|
|
},
|
|
},
|
|
)
|
|
# Pytest doesn't play nicely with imports
|
|
_register_all()
|
|
|
|
def tearDown(self):
|
|
ray.shutdown()
|
|
self.cluster.shutdown()
|
|
_register_all() # re-register the evicted objects
|
|
|
|
def _assertCleanup(self, trial_executor):
|
|
# Assert proper cleanup
|
|
resource_manager = trial_executor._resource_manager
|
|
self.assertFalse(resource_manager._pg_to_request)
|
|
self.assertFalse(resource_manager._acquired_pgs)
|
|
self.assertFalse(resource_manager._staging_future_to_pg)
|
|
self.assertFalse(resource_manager._pg_to_staging_future)
|
|
for rr in resource_manager._request_to_staged_pgs:
|
|
self.assertFalse(resource_manager._request_to_staged_pgs[rr])
|
|
for rr in resource_manager._request_to_ready_pgs:
|
|
self.assertFalse(resource_manager._request_to_ready_pgs[rr])
|
|
|
|
num_non_removed_pgs = len(
|
|
[p for pid, p in placement_group_table().items() if p["state"] != "REMOVED"]
|
|
)
|
|
self.assertEqual(num_non_removed_pgs, 0)
|
|
|
|
def testPlacementGroupRequests(self, reuse_actors=False, scheduled=10):
|
|
"""In this test we try to start 10 trials but only have resources
|
|
for 2. Placement groups should still be created and PENDING.
|
|
|
|
Eventually they should be scheduled sequentially (i.e. in pairs
|
|
of two)."""
|
|
# Since we check per-step placement groups, set the reconcilation
|
|
# interval to 0
|
|
os.environ["TUNE_PLACEMENT_GROUP_RECON_INTERVAL"] = "0"
|
|
|
|
def train_fn(config):
|
|
time.sleep(1)
|
|
now = time.time()
|
|
tune.report(end=now - config["start_time"])
|
|
|
|
head_bundle = {"CPU": 4, "GPU": 0, "custom": 0}
|
|
child_bundle = {"custom": 1}
|
|
# Manually calculated number of parallel trials
|
|
max_num_parallel = 2
|
|
|
|
placement_group_factory = PlacementGroupFactory(
|
|
[head_bundle, child_bundle, child_bundle]
|
|
)
|
|
|
|
trial_executor = RayTrialExecutor(reuse_actors=reuse_actors)
|
|
trial_executor.setup(max_pending_trials=max_num_parallel)
|
|
|
|
this = self
|
|
|
|
class _TestCallback(Callback):
|
|
def on_step_end(self, iteration, trials, **info):
|
|
num_finished = len(
|
|
[
|
|
t
|
|
for t in trials
|
|
if t.status == Trial.TERMINATED or t.status == Trial.ERROR
|
|
]
|
|
)
|
|
|
|
resource_manager = trial_executor._resource_manager
|
|
|
|
num_staging = sum(
|
|
len(s) for s in resource_manager._request_to_staged_pgs.values()
|
|
)
|
|
num_ready = sum(
|
|
len(s) for s in resource_manager._request_to_ready_pgs.values()
|
|
)
|
|
num_in_use = len(resource_manager._acquired_pgs)
|
|
num_cached = trial_executor._actor_cache.num_cached_objects
|
|
|
|
total_num_tracked = num_staging + num_ready + num_in_use + num_cached
|
|
|
|
# All trials should be scheduled
|
|
this.assertEqual(
|
|
scheduled,
|
|
min(scheduled, len(trials)),
|
|
msg=f"Num trials iter {iteration}",
|
|
)
|
|
|
|
# The following two tests were relaxed for reuse_actors=True
|
|
# so that up to `max_num_parallel` more placement groups can
|
|
# exist than we would expect. This is because caching
|
|
# relies on reconciliation for cleanup to avoid overscheduling
|
|
# of new placement groups.
|
|
num_parallel_reuse = int(reuse_actors) * max_num_parallel
|
|
|
|
# The number of PGs should decrease when trials finish
|
|
# We allow a constant excess of 1 here because the trial will
|
|
# be TERMINATED and the resources only returned after the trainable
|
|
# cleanup future succeeded. Because num_finished will increase,
|
|
# this still asserts that the number of PGs goes down over time.
|
|
this.assertGreaterEqual(
|
|
max(scheduled, len(trials)) - num_finished + 1 + num_parallel_reuse,
|
|
total_num_tracked,
|
|
msg=f"Num tracked iter {iteration}, {len(trials)}, "
|
|
f"{scheduled}, {num_finished}, {num_parallel_reuse}",
|
|
)
|
|
|
|
start = time.time()
|
|
out = tune.run(
|
|
train_fn,
|
|
config={"start_time": start},
|
|
resources_per_trial=placement_group_factory,
|
|
num_samples=10,
|
|
trial_executor=trial_executor,
|
|
callbacks=[_TestCallback()],
|
|
reuse_actors=reuse_actors,
|
|
verbose=2,
|
|
)
|
|
|
|
trial_end_times = sorted(t.last_result["end"] for t in out.trials)
|
|
print("Trial end times:", trial_end_times)
|
|
max_diff = trial_end_times[-1] - trial_end_times[0]
|
|
|
|
# Not all trials have been run in parallel
|
|
self.assertGreater(max_diff, 3)
|
|
|
|
# Some trials should have run in parallel
|
|
# Todo: Re-enable when using buildkite
|
|
# self.assertLess(max_diff, 10)
|
|
|
|
self._assertCleanup(trial_executor)
|
|
|
|
def testPlacementGroupRequestsWithActorReuse(self):
|
|
"""Assert that reuse actors doesn't leak placement groups"""
|
|
self.testPlacementGroupRequests(reuse_actors=True)
|
|
|
|
def testPlacementGroupLimitedRequests(self):
|
|
"""Assert that maximum number of placement groups is enforced."""
|
|
os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "6"
|
|
self.testPlacementGroupRequests(scheduled=6)
|
|
|
|
def testPlacementGroupLimitedRequestsWithActorReuse(self):
|
|
os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "6"
|
|
self.testPlacementGroupRequests(reuse_actors=True, scheduled=6)
|
|
|
|
def testPlacementGroupDistributedTraining(self, reuse_actors=False):
|
|
"""Run distributed training using placement groups.
|
|
|
|
Each trial requests 4 CPUs and starts 4 remote training workers.
|
|
"""
|
|
|
|
head_bundle = {"CPU": 1, "GPU": 0, "custom": 0}
|
|
child_bundle = {"CPU": 1}
|
|
|
|
placement_group_factory = PlacementGroupFactory(
|
|
[head_bundle, child_bundle, child_bundle, child_bundle]
|
|
)
|
|
|
|
@ray.remote
|
|
class TrainingActor:
|
|
def train(self, val):
|
|
time.sleep(1)
|
|
return val
|
|
|
|
def train_fn(config):
|
|
base = config["base"]
|
|
actors = [TrainingActor.remote() for _ in range(4)]
|
|
futures = [
|
|
actor.train.remote(base + 2 * i) for i, actor in enumerate(actors)
|
|
]
|
|
results = ray.get(futures)
|
|
|
|
end = time.time() - config["start_time"]
|
|
tune.report(avg=np.mean(results), end=end)
|
|
|
|
trial_executor = RayTrialExecutor(reuse_actors=reuse_actors)
|
|
|
|
start = time.time()
|
|
out = tune.run(
|
|
train_fn,
|
|
config={
|
|
"start_time": start,
|
|
"base": tune.grid_search(list(range(0, 100, 10))),
|
|
},
|
|
resources_per_trial=placement_group_factory,
|
|
num_samples=1,
|
|
trial_executor=trial_executor,
|
|
reuse_actors=reuse_actors,
|
|
verbose=2,
|
|
)
|
|
|
|
avgs = sorted(t.last_result["avg"] for t in out.trials)
|
|
self.assertSequenceEqual(avgs, list(range(3, 103, 10)))
|
|
|
|
trial_end_times = sorted(t.last_result["end"] for t in out.trials)
|
|
print("Trial end times:", trial_end_times)
|
|
max_diff = trial_end_times[-1] - trial_end_times[0]
|
|
|
|
# Not all trials have been run in parallel
|
|
self.assertGreater(max_diff, 3)
|
|
|
|
# Some trials should have run in parallel
|
|
# Todo: Re-enable when using buildkite
|
|
# self.assertLess(max_diff, 10)
|
|
|
|
self._assertCleanup(trial_executor)
|
|
|
|
def testPlacementGroupDistributedTrainingWithActorReuse(self):
|
|
self.testPlacementGroupDistributedTraining(reuse_actors=True)
|
|
|
|
|
|
class TrialRunnerPlacementGroupHeterogeneousTest(unittest.TestCase):
|
|
def tearDown(self) -> None:
|
|
if ray.is_initialized:
|
|
ray.shutdown()
|
|
|
|
def testResourceDeadlock(self):
|
|
"""Tests that resource deadlock is avoided for heterogeneous PGFs.
|
|
|
|
We start 4 trials in a cluster with 2 CPUs. The first two trials
|
|
require 1 CPU each, the third trial 2 CPUs, the fourth trial 1 CPU.
|
|
|
|
The second trial needs a bit more time to finish. This means that the
|
|
resources from the first trial will be freed, and the PG of the
|
|
_fourth_ trial becomes ready (not that of the third trial, because that
|
|
requires 2 CPUs - however, one is still occupied by trial 2).
|
|
|
|
After the first two trials finished, the FIFOScheduler tries to start
|
|
the third trial. However, it can't be started because its placement
|
|
group is not ready. Instead, the placement group of the fourth
|
|
trial is ready. Thus, we opt to run the fourth trial instead.
|
|
"""
|
|
|
|
def train_fn(config):
|
|
time.sleep(config["sleep"])
|
|
return 4
|
|
|
|
ray.init(num_cpus=2)
|
|
|
|
tune.register_trainable("het", train_fn)
|
|
pgf1 = PlacementGroupFactory([{"CPU": 1}])
|
|
pgf2 = PlacementGroupFactory([{"CPU": 2}])
|
|
|
|
trial1 = Trial("het", config={"sleep": 0}, placement_group_factory=pgf1)
|
|
trial2 = Trial("het", config={"sleep": 2}, placement_group_factory=pgf1)
|
|
trial3 = Trial("het", config={"sleep": 0}, placement_group_factory=pgf2)
|
|
trial4 = Trial("het", config={"sleep": 0}, placement_group_factory=pgf1)
|
|
|
|
runner = TrialRunner(fail_fast=True)
|
|
runner.add_trial(trial1)
|
|
runner.add_trial(trial2)
|
|
runner.add_trial(trial3)
|
|
runner.add_trial(trial4)
|
|
|
|
timeout = time.monotonic() + 30
|
|
while not runner.is_finished():
|
|
# We enforce a timeout here
|
|
self.assertLess(
|
|
time.monotonic(), timeout, msg="Ran into a resource deadlock"
|
|
)
|
|
|
|
runner.step()
|
|
|
|
|
|
def test_placement_group_no_cpu_trainer():
|
|
"""Bundles with only GPU:1 but no CPU should work"""
|
|
ray.init(num_gpus=1, num_cpus=1)
|
|
pgf = PlacementGroupFactory([{"GPU": 1, "CPU": 0}, {"CPU": 1}])
|
|
|
|
def train_fn(config):
|
|
time.sleep(1)
|
|
return 5
|
|
|
|
tune.run(train_fn, resources_per_trial=pgf)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|