def test_persisted_checkpoint(ray_start_2_cpus, tmp_path):
def train():
for i in range(2):
sgd.save_checkpoint(epoch=i)
config = TestConfig()
e = BackendExecutor(config)
e.start()
e.start_training(train, run_dir=tmp_path)
e.finish_training()
assert e.latest_checkpoint_id == 2
assert e.latest_checkpoint is not None
assert e.latest_checkpoint["epoch"] == 1
assert e.latest_checkpoint_path is not None
assert os.path.exists(e.latest_checkpoint_path)
def validate():
checkpoint = sgd.load_checkpoint()
assert checkpoint is not None
assert checkpoint["epoch"] == 1
e2 = BackendExecutor(config)
e2.start()
e2.start_training(validate,
checkpoint=e.latest_checkpoint_path,
run_dir=tmp_path)
e2.finish_training()
def test_train(ray_start_2_cpus, tmp_path):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
e.start_training(lambda: 1, run_dir=tmp_path)
assert e.finish_training() == [1, 1]
def test_start(ray_start_2_cpus):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
with pytest.raises(InactiveWorkerGroupError):
e.run(lambda: 1)
e.start()
assert len(e.worker_group) == 2
def __init__(self,
backend: Union[str, BackendConfig],
num_workers: int = 1,
use_gpu: bool = False,
resources_per_worker: Optional[Dict[str, float]] = None):
"""A class for distributed training.
Args:
backend (Union[str, BackendConfig]): The backend used for
distributed communication. If configurations are needed,
a subclass of ``BackendConfig`` can be passed in.
Supported ``str`` values: {"torch"}.
num_workers (int): The number of workers (Ray actors) to launch.
Defaults to 1. Each worker will reserve 1 CPU by default.
use_gpu (bool): If True, training will be done on GPUs (1 per
worker). Defaults to False.
resources_per_worker (Optional[Dict]): If specified, the resources
defined in this Dict will be reserved for each worker.
"""
# Setup executor.
backend_config = self._get_backend_config(backend)
if resources_per_worker:
raise NotImplementedError("`resources_per_worker` argument is not "
"supported yet.")
self._executor = BackendExecutor(backend_config, num_workers, 1,
int(use_gpu))
def test_start(ray_start_2_cpus, tmp_path):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
with pytest.raises(InactiveWorkerGroupError):
e.start_training(lambda: 1, run_dir=tmp_path)
e.start()
assert len(e.worker_group) == 2
def test_local_ranks(ray_start_2_cpus, tmp_path):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
def train():
return sgd.local_rank()
e.start_training(train, run_dir=tmp_path)
assert set(e.finish_training()) == {0, 1}
def test_execute_worker_failure(ray_start_2_cpus):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
def train_fail():
ray.actor.exit_actor()
new_execute_func = gen_execute_special(train_fail)
with patch.object(WorkerGroup, "execute_async", new_execute_func):
with pytest.raises(RuntimeError):
e.run(lambda: 1)
def test_worker_failure(ray_start_2_cpus, tmp_path):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
def train_fail():
ray.actor.exit_actor()
new_execute_func = gen_execute_special(train_fail)
with patch.object(WorkerGroup, "execute_async", new_execute_func):
with pytest.raises(TrainingWorkerError):
e.start_training(lambda: 1, run_dir=tmp_path)
e.finish_training()
def test_mismatch_checkpoint_report(ray_start_2_cpus, tmp_path):
def train():
if (sgd.world_rank()) == 0:
sgd.save_checkpoint(epoch=0)
else:
sgd.report(iter=0)
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
e.start_training(train, run_dir=tmp_path)
with pytest.raises(RuntimeError):
e.finish_training()
def test_checkpoint(ray_start_2_cpus, tmp_path):
def train():
for i in range(2):
sgd.save_checkpoint(epoch=i)
config = TestConfig()
e = BackendExecutor(config, num_workers=1)
e.start()
e.start_training(train, run_dir=tmp_path)
e.finish_training()
assert e.latest_checkpoint is not None
assert e.latest_checkpoint["epoch"] == 1
def __init__(
self,
backend: Union[str, BackendConfig],
num_workers: int = 1,
use_gpu: bool = False,
resources_per_worker: Optional[Dict[str, float]] = None,
logdir: Optional[str] = None,
max_retries: int = 3,
):
self._backend = backend
self._num_workers = num_workers
self._use_gpu = use_gpu
self._resources_per_worker = resources_per_worker
# Incremental unique run ID.
self._run_id = 0
self.logdir = self.create_logdir(logdir)
# Setup executor.
backend_config = self._get_backend_config(backend)
num_cpus = 1
num_gpus = int(use_gpu)
if resources_per_worker:
# Override CPU and GPU resources and remove from dict.
num_cpus = resources_per_worker.pop("CPU", num_cpus)
num_gpus = resources_per_worker.pop("GPU", num_gpus)
if not use_gpu and num_gpus > 0:
raise ValueError(
"`use_gpu` is False but `GPU` was found in "
"`resources_per_worker`. Either set `use_gpu` to True or "
"remove `GPU` from `resources_per_worker.")
if use_gpu and num_gpus == 0:
raise ValueError(
"`use_gpu` is True but `GPU` is set to 0 in "
"`resources_per_worker`. Either set `use_gpu` to False or "
"request a positive number of `GPU` in "
"`resources_per_worker.")
self._executor = BackendExecutor(
backend_config=backend_config,
num_workers=num_workers,
num_cpus_per_worker=num_cpus,
num_gpus_per_worker=num_gpus,
additional_resources_per_worker=resources_per_worker,
max_retries=max_retries)
def test_no_exhaust(ray_start_2_cpus, tmp_path):
"""Tests if training can finish even if queue is not exhausted."""
def train():
for _ in range(2):
sgd.report(loss=1)
return 2
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
e.start_training(train, run_dir=tmp_path)
output = e.finish_training()
assert output == [2, 2]
def test_checkpoint(ray_start_2_cpus):
def train():
for i in range(2):
sgd.save_checkpoint(epoch=i)
config = TestConfig()
e = BackendExecutor(config, num_workers=1)
e.start()
e.start_training(train)
e.finish_training()
latest_checkpoint = e.get_latest_checkpoint()
assert latest_checkpoint is not None
assert latest_checkpoint["epoch"] == 1
def test_initialization_hook(ray_start_2_cpus):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
def init_hook():
import os
os.environ["TEST"] = "1"
e.start(initialization_hook=init_hook)
def check():
import os
return os.getenv("TEST", "0")
assert e.run(check) == ["1", "1"]
def test_torch_start_shutdown(ray_start_2_cpus, init_method):
torch_config = TorchConfig(backend="gloo", init_method=init_method)
e = BackendExecutor(torch_config, num_workers=2)
e.start()
def check_process_group():
import torch
return torch.distributed.is_initialized(
) and torch.distributed.get_world_size() == 2
assert all(e.run(check_process_group))
e._backend.on_shutdown(e.worker_group, e._backend_config)
assert not any(e.run(check_process_group))
def test_initialization_hook(ray_start_2_cpus, tmp_path):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
def init_hook():
import os
os.environ["TEST"] = "1"
e.start(initialization_hook=init_hook)
def check():
import os
return os.getenv("TEST", "0")
e.start_training(check, run_dir=tmp_path)
assert e.finish_training() == ["1", "1"]
def test_cuda_visible_devices(ray_2_node_2_gpu, worker_results, tmp_path):
config = TestConfig()
def get_resources():
return os.environ["CUDA_VISIBLE_DEVICES"]
num_workers, expected_results = worker_results
e = BackendExecutor(config,
num_workers=num_workers,
num_cpus_per_worker=0,
num_gpus_per_worker=1)
e.start()
e.start_training(get_resources, tmp_path)
results = e.finish_training()
results.sort()
assert results == expected_results
def test_train_failure(ray_start_2_cpus, tmp_path):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
with pytest.raises(SGDBackendError):
e.fetch_next_result()
with pytest.raises(SGDBackendError):
e.finish_training()
e.start_training(lambda: 1, run_dir=tmp_path)
with pytest.raises(SGDBackendError):
e.start_training(lambda: 2, run_dir=tmp_path)
assert e.finish_training() == [1, 1]
def test_tensorflow_start(ray_start_2_cpus, tmp_path):
num_workers = 2
tensorflow_config = TensorflowConfig()
e = BackendExecutor(tensorflow_config, num_workers=num_workers)
e.start()
def get_tf_config():
import json
import os
return json.loads(os.environ["TF_CONFIG"])
e.start_training(get_tf_config, run_dir=tmp_path)
results = e.finish_training()
assert len(results) == num_workers
workers = [result["cluster"]["worker"] for result in results]
assert all(worker == workers[0] for worker in workers)
indexes = [result["task"]["index"] for result in results]
assert len(set(indexes)) == num_workers
def test_execute(ray_start_2_cpus):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
assert e.run(lambda: 1) == [1, 1]