wehub-resource-sync
134 lines
4.2 KiB
Python
134 lines
4.2 KiB
Python
import os
|
|
import unittest
|
|
|
|
import ray
|
|
from ray import tune
|
|
from ray.rllib.algorithms.ppo import PPO, PPOConfig
|
|
from ray.tune import Callback
|
|
from ray.tune.execution.placement_groups import PlacementGroupFactory
|
|
from ray.tune.experiment import Trial
|
|
from ray.tune.result import TRAINING_ITERATION
|
|
|
|
trial_executor = None
|
|
|
|
|
|
class _TestCallback(Callback):
|
|
def on_step_end(self, iteration, trials, **info):
|
|
num_running = len([t for t in trials if t.status == Trial.RUNNING])
|
|
|
|
# All 3 trials (3 different learning rates) should be scheduled.
|
|
assert 3 == min(3, len(trials))
|
|
# Cannot run more than 2 at a time
|
|
# (due to different resource restrictions in the test cases).
|
|
assert num_running <= 2
|
|
|
|
|
|
class TestPlacementGroups(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
os.environ["TUNE_PLACEMENT_GROUP_RECON_INTERVAL"] = "0"
|
|
ray.init(num_cpus=6)
|
|
|
|
def tearDown(self) -> None:
|
|
ray.shutdown()
|
|
|
|
def test_overriding_default_resource_request(self):
|
|
# 3 Trials: Can only run 2 at a time (num_cpus=6; needed: 3).
|
|
config = (
|
|
PPOConfig()
|
|
.api_stack(
|
|
enable_env_runner_and_connector_v2=False,
|
|
enable_rl_module_and_learner=False,
|
|
)
|
|
.training(
|
|
model={"fcnet_hiddens": [10]}, lr=tune.grid_search([0.1, 0.01, 0.001])
|
|
)
|
|
.environment("CartPole-v1")
|
|
.env_runners(num_env_runners=2)
|
|
.framework("tf")
|
|
)
|
|
|
|
# Create an Algorithm with an overridden default_resource_request
|
|
# method that returns a PlacementGroupFactory.
|
|
|
|
class MyAlgo(PPO):
|
|
@classmethod
|
|
def default_resource_request(cls, config):
|
|
head_bundle = {"CPU": 1, "GPU": 0}
|
|
child_bundle = {"CPU": 1}
|
|
return PlacementGroupFactory(
|
|
[head_bundle, child_bundle, child_bundle],
|
|
strategy=config["placement_strategy"],
|
|
)
|
|
|
|
tune.register_trainable("my_trainable", MyAlgo)
|
|
|
|
tune.Tuner(
|
|
"my_trainable",
|
|
param_space=config,
|
|
run_config=tune.RunConfig(
|
|
stop={TRAINING_ITERATION: 2},
|
|
verbose=2,
|
|
callbacks=[_TestCallback()],
|
|
),
|
|
).fit()
|
|
|
|
def test_default_resource_request(self):
|
|
config = (
|
|
PPOConfig()
|
|
.api_stack(
|
|
enable_env_runner_and_connector_v2=False,
|
|
enable_rl_module_and_learner=False,
|
|
)
|
|
.resources(placement_strategy="SPREAD")
|
|
.env_runners(
|
|
num_env_runners=2,
|
|
num_cpus_per_env_runner=2,
|
|
)
|
|
.training(
|
|
model={"fcnet_hiddens": [10]}, lr=tune.grid_search([0.1, 0.01, 0.001])
|
|
)
|
|
.environment("CartPole-v1")
|
|
.framework("torch")
|
|
)
|
|
# 3 Trials: Can only run 1 at a time (num_cpus=6; needed: 5).
|
|
|
|
tune.Tuner(
|
|
PPO,
|
|
param_space=config,
|
|
run_config=tune.RunConfig(
|
|
stop={TRAINING_ITERATION: 2},
|
|
verbose=2,
|
|
callbacks=[_TestCallback()],
|
|
),
|
|
tune_config=tune.TuneConfig(reuse_actors=False),
|
|
).fit()
|
|
|
|
def test_default_resource_request_plus_manual_leads_to_error(self):
|
|
config = (
|
|
PPOConfig()
|
|
.api_stack(
|
|
enable_env_runner_and_connector_v2=False,
|
|
enable_rl_module_and_learner=False,
|
|
)
|
|
.training(model={"fcnet_hiddens": [10]})
|
|
.environment("CartPole-v1")
|
|
.env_runners(num_env_runners=0)
|
|
)
|
|
|
|
try:
|
|
tune.Tuner(
|
|
tune.with_resources(PPO, PlacementGroupFactory([{"CPU": 1}])),
|
|
param_space=config,
|
|
run_config=tune.RunConfig(stop={TRAINING_ITERATION: 2}, verbose=2),
|
|
).fit()
|
|
except ValueError as e:
|
|
assert "have been automatically set to" in e.args[0]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|