wehub-resource-sync
373 lines
11 KiB
Python
373 lines
11 KiB
Python
import os
|
|
import time
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray import train, tune
|
|
from ray._common.utils import RESOURCE_CONSTRAINT_PREFIX
|
|
from ray.train import RunConfig, ScalingConfig
|
|
from ray.train._internal.backend_executor import BackendExecutor
|
|
from ray.train._internal.worker_group import WorkerGroup
|
|
from ray.train.backend import Backend, BackendConfig
|
|
from ray.train.data_parallel_trainer import DataParallelTrainer
|
|
from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
|
|
from ray.train.utils import _in_ray_train_worker
|
|
from ray.tune.callback import Callback
|
|
from ray.tune.tune_config import TuneConfig
|
|
from ray.tune.tuner import Tuner
|
|
from ray.util.accelerators import NVIDIA_A100
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_4_cpus():
|
|
address_info = ray.init(num_cpus=4)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_4_cpus_4_gpus_4_extra():
|
|
address_info = ray.init(num_cpus=4, num_gpus=4, resources={"extra": 4})
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_4_cpus_4_gpus_4_a100():
|
|
address_info = ray.init(
|
|
num_cpus=4,
|
|
num_gpus=4,
|
|
resources={f"{RESOURCE_CONSTRAINT_PREFIX}{NVIDIA_A100}": 4},
|
|
)
|
|
yield address_info
|
|
ray.shutdown()
|
|
|
|
|
|
def gen_execute_single_async_special(special_f):
|
|
def execute_single_async_special(self, i, f, *args, **kwargs):
|
|
assert len(self.workers) == 2
|
|
if i == 0 and hasattr(self, "should_fail") and self.should_fail:
|
|
kwargs["train_func"] = special_f
|
|
return (
|
|
self.workers[i]
|
|
.actor._RayTrainWorker__execute.options(name=f.__name__)
|
|
.remote(f, *args, **kwargs)
|
|
)
|
|
|
|
return execute_single_async_special
|
|
|
|
|
|
def gen_new_backend_executor(special_f):
|
|
"""Returns a BackendExecutor that runs special_f on worker 0 once."""
|
|
|
|
class TestBackendExecutor(BackendExecutor):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._has_failed = False
|
|
|
|
def start_training(self, *args, **kwargs):
|
|
special_execute = gen_execute_single_async_special(special_f)
|
|
if not self._has_failed:
|
|
self.worker_group.should_fail = True
|
|
self._has_failed = True
|
|
else:
|
|
self.worker_group.should_fail = False
|
|
with patch.object(WorkerGroup, "execute_single_async", special_execute):
|
|
super().start_training(*args, **kwargs)
|
|
|
|
return TestBackendExecutor
|
|
|
|
|
|
class CaptureReportCallback(Callback):
|
|
def __init__(self):
|
|
self.result_list = []
|
|
|
|
def on_trial_result(self, iteration, trials, trial, result, **info):
|
|
self.result_list.append(result)
|
|
|
|
|
|
scale_config = ScalingConfig(num_workers=2)
|
|
|
|
|
|
def test_fit_train(ray_start_4_cpus):
|
|
def train_func():
|
|
train.report({"loss": 1})
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_loop_per_worker=train_func, scaling_config=scale_config
|
|
)
|
|
assert trainer.fit().metrics["loss"] == 1
|
|
|
|
|
|
def test_scaling_config(ray_start_4_cpus):
|
|
def train_func():
|
|
assert ray.available_resources()["CPU"] == 1
|
|
train.report({"loss": 1})
|
|
|
|
assert ray.available_resources()["CPU"] == 4
|
|
trainer = DataParallelTrainer(
|
|
train_loop_per_worker=train_func, scaling_config=ScalingConfig(num_workers=2)
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
def test_fit_train_config(ray_start_4_cpus):
|
|
def train_func(config):
|
|
train.report({"loss": config["x"]})
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_loop_per_worker=train_func,
|
|
scaling_config=scale_config,
|
|
train_loop_config={"x": 100},
|
|
)
|
|
assert trainer.fit().metrics["loss"] == 100
|
|
|
|
|
|
def test_datasets(ray_start_4_cpus):
|
|
num_train_data = 10
|
|
num_val_data = 6
|
|
|
|
train_dataset = ray.data.range(num_train_data)
|
|
val_dataset = ray.data.range(num_val_data)
|
|
|
|
def get_dataset():
|
|
# Train dataset should be sharded.
|
|
train_dataset = train.get_dataset_shard("train")
|
|
train_ds_count = len(list(train_dataset.iter_rows()))
|
|
assert train_ds_count == num_train_data / scale_config.num_workers
|
|
# All other datasets should not be sharded.
|
|
val_dataset = train.get_dataset_shard("val")
|
|
val_ds_count = len(list(val_dataset.iter_rows()))
|
|
assert val_ds_count == num_val_data / scale_config.num_workers
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_loop_per_worker=get_dataset,
|
|
scaling_config=scale_config,
|
|
datasets={"train": train_dataset, "val": val_dataset},
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
def test_invalid_train_loop():
|
|
def train_loop(config, extra_arg):
|
|
pass
|
|
|
|
with pytest.raises(ValueError):
|
|
DataParallelTrainer(train_loop_per_worker=train_loop)
|
|
|
|
|
|
def test_bad_return_in_train_loop(ray_start_4_cpus):
|
|
"""Test to check if returns from train loop are discarded."""
|
|
|
|
# Simulates what happens with eg. torch models
|
|
class FailOnUnpickle:
|
|
def __reduce__(self):
|
|
raise RuntimeError("Failing")
|
|
|
|
def train_loop(config):
|
|
train.report({"loss": 1})
|
|
return FailOnUnpickle()
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_loop_per_worker=train_loop, scaling_config=scale_config
|
|
)
|
|
|
|
# No exception should happen here
|
|
trainer.fit()
|
|
|
|
|
|
def test_tune(ray_start_4_cpus):
|
|
def train_func(config):
|
|
train.report({"loss": config["x"]})
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_loop_per_worker=train_func,
|
|
train_loop_config={"x": 100},
|
|
scaling_config=scale_config,
|
|
)
|
|
|
|
tuner = Tuner(
|
|
trainer,
|
|
param_space={"train_loop_config": {"x": tune.choice([200, 300])}},
|
|
tune_config=TuneConfig(num_samples=2),
|
|
)
|
|
result_grid = tuner.fit()
|
|
assert result_grid[0].metrics["loss"] in [200, 300]
|
|
|
|
# Make sure original Trainer is not affected.
|
|
assert trainer._train_loop_config["x"] == 100
|
|
|
|
|
|
def test_fast_slow(ray_start_4_cpus):
|
|
def train_func():
|
|
for i in range(2):
|
|
with create_dict_checkpoint({"epoch": i}) as checkpoint:
|
|
train.report(dict(index=i), checkpoint=checkpoint)
|
|
|
|
def train_slow():
|
|
for i in range(2):
|
|
with create_dict_checkpoint({"epoch": i}) as checkpoint:
|
|
train.report(dict(index=i), checkpoint=checkpoint)
|
|
time.sleep(5)
|
|
|
|
new_backend_executor_cls = gen_new_backend_executor(train_slow)
|
|
callback = CaptureReportCallback()
|
|
|
|
class DataParallelTrainerPatched(DataParallelTrainer):
|
|
_backend_executor_cls = new_backend_executor_cls
|
|
|
|
trainer = DataParallelTrainerPatched(
|
|
train_func,
|
|
scaling_config=scale_config,
|
|
run_config=RunConfig(callbacks=[callback]),
|
|
)
|
|
results = trainer.fit()
|
|
|
|
assert load_dict_checkpoint(results.checkpoint)["epoch"] == 1
|
|
|
|
result_list = callback.result_list
|
|
assert len(result_list) == 2
|
|
|
|
|
|
def test_mismatch_report(ray_start_4_cpus):
|
|
def train_func():
|
|
for _ in range(2):
|
|
train.report(dict(loss=1))
|
|
|
|
def train_mismatch():
|
|
train.report(dict(loss=1))
|
|
|
|
new_backend_executor_cls = gen_new_backend_executor(train_mismatch)
|
|
|
|
class DataParallelTrainerPatched(DataParallelTrainer):
|
|
_backend_executor_cls = new_backend_executor_cls
|
|
|
|
trainer = DataParallelTrainerPatched(
|
|
train_func,
|
|
scaling_config=scale_config,
|
|
)
|
|
with pytest.raises(RuntimeError):
|
|
trainer.fit()
|
|
|
|
|
|
def test_world_rank(ray_start_4_cpus, tmp_path):
|
|
def train_func():
|
|
world_rank = train.get_context().get_world_rank()
|
|
(tmp_path / f"{world_rank}").touch()
|
|
train.report(dict(world_rank=world_rank))
|
|
|
|
trainer = DataParallelTrainer(train_func, scaling_config=scale_config)
|
|
trainer.fit()
|
|
|
|
created_files = list(tmp_path.glob("*"))
|
|
assert len(created_files) == 2
|
|
assert {int(file.name) for file in created_files} == {0, 1}
|
|
|
|
|
|
def test_gpu_requests(ray_start_4_cpus_4_gpus_4_extra, tmp_path):
|
|
def get_visible_devices_for_workers():
|
|
return [file.read_text() for file in tmp_path.glob("*")]
|
|
|
|
class CudaTestBackend(Backend):
|
|
share_cuda_visible_devices = True
|
|
|
|
class CudaTestConfig(BackendConfig):
|
|
@property
|
|
def backend_cls(self):
|
|
return CudaTestBackend
|
|
|
|
def get_resources():
|
|
cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "")
|
|
world_rank = train.get_context().get_world_rank()
|
|
(tmp_path / f"{world_rank}").write_text(cuda_visible_devices)
|
|
train.report(dict(devices=cuda_visible_devices))
|
|
|
|
# 0 GPUs will be requested and should not raise an error.
|
|
trainer = DataParallelTrainer(
|
|
get_resources,
|
|
backend_config=CudaTestConfig(),
|
|
scaling_config=ScalingConfig(num_workers=2, use_gpu=False),
|
|
)
|
|
trainer.fit()
|
|
assert get_visible_devices_for_workers() == ["", ""]
|
|
|
|
# 1 GPU will be requested and should not raise an error.
|
|
trainer = DataParallelTrainer(
|
|
get_resources,
|
|
backend_config=CudaTestConfig(),
|
|
scaling_config=ScalingConfig(num_workers=2, use_gpu=True),
|
|
)
|
|
trainer.fit()
|
|
visible_devices = get_visible_devices_for_workers()
|
|
# Sort the cuda visible devices to have exact match with expected result.
|
|
visible_devices = [",".join(sorted(r.split(","))) for r in visible_devices]
|
|
assert visible_devices == ["0,1", "0,1"]
|
|
|
|
# Partial GPUs should not raise an error.
|
|
trainer = DataParallelTrainer(
|
|
get_resources,
|
|
backend_config=CudaTestConfig(),
|
|
scaling_config=ScalingConfig(
|
|
num_workers=2, use_gpu=True, resources_per_worker={"GPU": 0.1}
|
|
),
|
|
)
|
|
trainer.fit()
|
|
visible_devices = get_visible_devices_for_workers()
|
|
assert visible_devices == ["0", "0"]
|
|
|
|
# Multiple GPUs should not raise an error.
|
|
trainer = DataParallelTrainer(
|
|
get_resources,
|
|
backend_config=CudaTestConfig(),
|
|
scaling_config=ScalingConfig(
|
|
num_workers=2, use_gpu=True, resources_per_worker={"GPU": 2}
|
|
),
|
|
)
|
|
trainer.fit()
|
|
visible_devices = get_visible_devices_for_workers()
|
|
# Sort the cuda visible devices to have exact match with expected result.
|
|
visible_devices = [",".join(sorted(r.split(","))) for r in visible_devices]
|
|
assert visible_devices == ["0,1,2,3", "0,1,2,3"]
|
|
|
|
|
|
@pytest.mark.parametrize("accelerator_type", [NVIDIA_A100, None])
|
|
def test_config_accelerator_type(ray_start_4_cpus_4_gpus_4_a100, accelerator_type):
|
|
def train_func():
|
|
# Ensure all workers are scheduled on nodes with specified accelerators
|
|
assigned_resources = ray.get_runtime_context().get_assigned_resources()
|
|
if accelerator_type:
|
|
accelerator_key = f"{RESOURCE_CONSTRAINT_PREFIX}{accelerator_type}"
|
|
assert accelerator_key in assigned_resources
|
|
|
|
trainer = DataParallelTrainer(
|
|
train_func,
|
|
scaling_config=ScalingConfig(
|
|
num_workers=4,
|
|
use_gpu=True,
|
|
accelerator_type=accelerator_type,
|
|
),
|
|
)
|
|
trainer.fit()
|
|
|
|
|
|
def test_in_ray_train_worker(ray_start_4_cpus):
|
|
assert not _in_ray_train_worker()
|
|
|
|
def train_fn():
|
|
assert _in_ray_train_worker()
|
|
|
|
trainer = DataParallelTrainer(train_fn)
|
|
trainer.fit()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|