wehub-resource-sync
140 lines
4.3 KiB
Python
140 lines
4.3 KiB
Python
import sys
|
|
|
|
import pytest
|
|
|
|
import ray.train
|
|
import ray.tune
|
|
from ray.cluster_utils import Cluster
|
|
from ray.train.tests.util import create_dict_checkpoint
|
|
from ray.train.v2._internal.constants import HEALTH_CHECK_INTERVAL_S_ENV_VAR
|
|
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
|
|
from ray.tune.integration.ray_train import CHECKPOINT_PATH_KEY, TuneReportCallback
|
|
|
|
TRAIN_DRIVER_RESOURCE_NAME = "train_driver_resource"
|
|
NUM_GPUS_IN_CLUSTER = 4
|
|
|
|
|
|
@pytest.fixture()
|
|
def ray_start_4_cpus():
|
|
ray.init(num_cpus=4)
|
|
yield
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture()
|
|
def ray_cpu_head_gpu_worker():
|
|
cluster = Cluster()
|
|
cluster.add_node(resources={TRAIN_DRIVER_RESOURCE_NAME: 1})
|
|
cluster.add_node(num_cpus=0, num_gpus=NUM_GPUS_IN_CLUSTER)
|
|
|
|
ray.init(address=cluster.address)
|
|
|
|
yield
|
|
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def speed_up_tests(monkeypatch):
|
|
monkeypatch.setenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, "0.1")
|
|
|
|
|
|
@pytest.mark.parametrize("num_workers_grid_search", [[1], [1, 2, 4]])
|
|
@pytest.mark.parametrize("limit_concurrency", [True, False])
|
|
def test_e2e(
|
|
ray_cpu_head_gpu_worker,
|
|
tmp_path,
|
|
num_workers_grid_search,
|
|
limit_concurrency,
|
|
):
|
|
num_non_checkpoint_reports = 2
|
|
num_checkpoint_reports = 1
|
|
|
|
def train_fn_per_worker(train_fn_config):
|
|
assert "lr" in train_fn_config
|
|
|
|
world_size = ray.train.get_context().get_world_size()
|
|
for i in range(num_non_checkpoint_reports):
|
|
ray.train.report({"idx": i})
|
|
|
|
for i in range(num_checkpoint_reports):
|
|
with create_dict_checkpoint({"model": "dummy"}) as checkpoint:
|
|
ray.train.report(
|
|
{"loss": 0.1, "world_size": world_size}, checkpoint=checkpoint
|
|
)
|
|
|
|
def launch_training(tune_config):
|
|
trainer = DataParallelTrainer(
|
|
train_loop_per_worker=train_fn_per_worker,
|
|
train_loop_config=tune_config["train_loop_config"],
|
|
scaling_config=ray.train.ScalingConfig(
|
|
num_workers=tune_config["num_workers"], use_gpu=True
|
|
),
|
|
run_config=ray.train.RunConfig(
|
|
storage_path=tmp_path,
|
|
name=f"train-{ray.tune.get_context().get_trial_id()}",
|
|
callbacks=[TuneReportCallback()],
|
|
),
|
|
)
|
|
trainer.fit()
|
|
|
|
tuner = ray.tune.Tuner(
|
|
ray.tune.with_resources(launch_training, {TRAIN_DRIVER_RESOURCE_NAME: 0.01}),
|
|
param_space={
|
|
# Search over parameters passed into each train worker.
|
|
"train_loop_config": {"lr": ray.tune.choice([0.01, 0.001])},
|
|
# Search over Train "run level" parameters.
|
|
"num_workers": ray.tune.grid_search(num_workers_grid_search),
|
|
},
|
|
tune_config=ray.tune.TuneConfig(
|
|
max_concurrent_trials=(
|
|
NUM_GPUS_IN_CLUSTER // max(num_workers_grid_search)
|
|
if limit_concurrency
|
|
else None
|
|
)
|
|
),
|
|
run_config=ray.tune.RunConfig(storage_path=tmp_path, name="tune"),
|
|
)
|
|
result_grid = tuner.fit()
|
|
assert len(result_grid) == len(num_workers_grid_search)
|
|
|
|
world_sizes = set()
|
|
for result in result_grid:
|
|
assert (
|
|
len(result.metrics_dataframe)
|
|
== num_non_checkpoint_reports + num_checkpoint_reports
|
|
)
|
|
assert "loss" in result.metrics
|
|
assert CHECKPOINT_PATH_KEY in result.metrics
|
|
world_sizes.add(result.metrics["world_size"])
|
|
assert world_sizes == set(num_workers_grid_search)
|
|
|
|
|
|
def test_errors(ray_start_4_cpus):
|
|
"""Test that errors in training are properly captured and reported."""
|
|
|
|
def train_worker_fn():
|
|
raise RuntimeError("Simulated training error")
|
|
|
|
def train_fn(config):
|
|
trainer = DataParallelTrainer(train_worker_fn)
|
|
trainer.fit()
|
|
|
|
tuner = ray.tune.Tuner(train_fn)
|
|
|
|
results = tuner.fit()
|
|
|
|
assert results.errors, "Expected errors to be captured"
|
|
assert len(results.errors) == 1, "Expected exactly one error"
|
|
|
|
error = results.errors[0]
|
|
assert "RuntimeError" in str(error), f"Expected RuntimeError, got: {error}"
|
|
assert "Simulated training error" in str(
|
|
error
|
|
), f"Expected specific error message, got: {error}"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|