Files
2026-07-13 13:17:40 +08:00

661 lines
19 KiB
Python

import math
import os
import sys
import tempfile
import time
from typing import Set
from unittest.mock import patch
import pytest
import ray
from ray import train
from ray._private.accelerators.neuron import NEURON_RT_VISIBLE_CORES_ENV_VAR
from ray.air._internal.util import StartTraceback
# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
from ray.tests.conftest import pytest_runtest_makereport # noqa
from ray.train import DataConfig
from ray.train._internal.backend_executor import (
BackendExecutor,
InactiveWorkerGroupError,
TrainBackendError,
TrainingWorkerError,
)
from ray.train._internal.storage import StorageContext
from ray.train._internal.worker_group import WorkerGroup, WorkerMetadata
from ray.train.backend import Backend, BackendConfig
from ray.train.constants import (
ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV,
ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV,
JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S,
TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S,
TRAIN_ENABLE_WORKER_SPREAD_ENV,
)
from ray.train.torch import TorchConfig
from ray.train.v2.jax.config import JaxConfig
from ray.util.placement_group import get_current_placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from ray.util.state import list_actors
@pytest.fixture
def ray_start_2_cpus():
address_info = ray.init(num_cpus=2)
yield address_info
ray.shutdown()
def _start_training(backend_executor: BackendExecutor, fn):
storage = StorageContext(
storage_path=tempfile.mkdtemp(),
experiment_dir_name="exp_name",
trial_dir_name="trial_name",
)
backend_executor.start_training(
train_func=fn,
datasets={},
metadata={},
data_config=DataConfig(),
storage=storage,
)
def gen_execute_special(special_f):
def execute_async_special(self, f):
"""Runs f on worker 0, special_f on other workers."""
futures = [
self.workers[0]
.actor._RayTrainWorker__execute.options(name=f.__name__)
.remote(f)
]
for worker in self.workers[1:]:
futures.append(
worker.actor._RayTrainWorker__execute.options(
name=special_f.__name__
).remote(special_f)
)
return futures
return execute_async_special
class TestConfig(BackendConfig):
@property
def backend_cls(self):
return TestBackend
class TestBackend(Backend):
def on_start(self, worker_group: WorkerGroup, backend_config: TestConfig):
pass
def on_shutdown(self, worker_group: WorkerGroup, backend_config: TestConfig):
pass
original_add_workers = WorkerGroup.add_workers
def mock_add_workers(self, num_workers):
original_add_workers(self, num_workers)
for i, worker in enumerate(self.workers):
metadata = WorkerMetadata(
node_id=str(i % 2),
node_ip=str(i % 2),
hostname=0,
resource_ids={"GPU": ["0"]},
pid=0,
)
worker.metadata = metadata
def mock_add_workers_to_nodes_with_same_ip(self, num_workers):
original_add_workers(self, num_workers)
for i, worker in enumerate(self.workers):
metadata = WorkerMetadata(
node_id=str(i % 2),
node_ip=0,
hostname=0,
resource_ids={"GPU": ["0"]},
pid=0,
)
worker.metadata = metadata
def test_start(ray_start_2_cpus):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
with pytest.raises(InactiveWorkerGroupError):
_start_training(e, lambda: 1)
e.start()
assert len(e.worker_group) == 2
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")
_start_training(e, check)
assert e.finish_training() == ["1", "1"]
def test_shutdown(ray_start_2_cpus):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
assert len(e.worker_group) == 2
e.shutdown()
with pytest.raises(InactiveWorkerGroupError):
_start_training(e, lambda: 1)
def test_train(ray_start_2_cpus):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
_start_training(e, lambda: 1)
assert e.finish_training() == [1, 1]
def test_local_ranks(ray_start_2_cpus):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
def train_func():
return train.get_context().get_local_rank()
_start_training(e, train_func)
assert set(e.finish_training()) == {0, 1}
def test_local_ranks_with_same_ip_nodes(ray_2_node_2_cpu):
config = TestConfig()
e = BackendExecutor(config, num_workers=4)
e.start()
def train_func():
return train.get_context().get_local_rank()
_start_training(e, train_func)
assert list(e.finish_training()) == [0, 1, 0, 1]
def test_local_world_size(ray_2_node_2_cpu):
config = TestConfig()
with patch.object(WorkerGroup, "add_workers", mock_add_workers):
e = BackendExecutor(config, num_workers=3)
e.start()
def train_func():
return train.get_context().get_local_world_size()
_start_training(e, train_func)
assert list(e.finish_training()) == [2, 2, 1]
def test_local_world_size_with_same_ip_nodes(ray_2_node_2_cpu):
config = TestConfig()
with patch.object(
WorkerGroup, "add_workers", mock_add_workers_to_nodes_with_same_ip
):
e = BackendExecutor(config, num_workers=3)
e.start()
def train_func():
return train.get_context().get_local_world_size()
_start_training(e, train_func)
assert list(e.finish_training()) == [2, 2, 1]
def test_node_ranks(ray_2_node_2_cpu):
config = TestConfig()
with patch.object(WorkerGroup, "add_workers", mock_add_workers):
e = BackendExecutor(config, num_workers=3)
e.start()
def train_func():
return train.get_context().get_node_rank()
_start_training(e, train_func)
assert list(e.finish_training()) == [0, 0, 1]
def test_node_ranks_with_same_ip_nodes(ray_2_node_2_cpu):
config = TestConfig()
with patch.object(
WorkerGroup, "add_workers", mock_add_workers_to_nodes_with_same_ip
):
e = BackendExecutor(config, num_workers=3)
e.start()
def train_func():
return train.get_context().get_node_rank()
_start_training(e, train_func)
assert list(e.finish_training()) == [0, 0, 1]
def test_train_failure(ray_start_2_cpus):
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
with pytest.raises(StartTraceback) as exc:
e.get_next_results()
assert isinstance(exc.value.__cause__, TrainBackendError)
with pytest.raises(StartTraceback) as exc:
e.pause_reporting()
assert isinstance(exc.value.__cause__, TrainBackendError)
with pytest.raises(StartTraceback) as exc:
e.finish_training()
assert isinstance(exc.value.__cause__, TrainBackendError)
_start_training(e, lambda: 1)
with pytest.raises(StartTraceback) as exc:
_start_training(e, lambda: 2)
assert isinstance(exc.value.__cause__, TrainBackendError)
assert e.finish_training() == [1, 1]
def test_single_worker_user_failure(ray_start_2_cpus):
"""Tests if training fails immediately if one worker raises an Exception
while executing the user training code."""
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
def single_worker_user_failure():
if train.get_context().get_world_rank() == 0:
raise RuntimeError
else:
time.sleep(1000000)
_start_training(e, single_worker_user_failure)
with pytest.raises(StartTraceback) as exc:
e.get_next_results()
assert isinstance(exc.value.__cause__, RuntimeError)
def test_single_worker_actor_failure(ray_start_2_cpus):
"""Tests is training fails immediately if one worker actor dies."""
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
def single_worker_actor_failure():
if train.get_context().get_world_rank() == 0:
# Simulate actor failure
os._exit(1)
else:
time.sleep(1000)
_start_training(e, single_worker_actor_failure)
with pytest.raises(TrainingWorkerError):
e.get_next_results()
@pytest.mark.skipif(
sys.version_info >= (3, 12), reason="tensorflow is not supported in python 3.12+"
)
def test_tensorflow_start(ray_start_2_cpus):
from ray.train.tensorflow import TensorflowConfig
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"])
_start_training(e, get_tf_config)
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
@pytest.mark.parametrize("init_method", ["env", "tcp"])
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
)
_start_training(e, check_process_group)
assert all(e.finish_training())
e._backend.on_shutdown(e.worker_group, e._backend_config)
_start_training(e, check_process_group)
assert not any(e.finish_training())
@pytest.mark.parametrize(
"init_method, timeout_s", [("env", 5), ("tcp", 5), ("env", 0), ("tcp", 0)]
)
def test_torch_process_group_shutdown_timeout(
ray_start_2_cpus, monkeypatch, init_method, timeout_s
):
monkeypatch.setenv(TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S, timeout_s)
torch_config = TorchConfig(backend="gloo", init_method=init_method)
e = BackendExecutor(torch_config, num_workers=2)
e.start()
_start_training(e, lambda: 1)
assert e.finish_training() == [1, 1]
# Verify that we do not raise an exception even if we time out
e._backend.on_shutdown(e.worker_group, e._backend_config)
@pytest.mark.parametrize(
"worker_results",
[
(1, [[0]]),
(2, [[0, 1]] * 2),
(3, [[0]] + [[0, 1]] * 2),
(4, [[0, 1]] * 4),
],
)
def test_cuda_visible_devices(ray_2_node_2_gpu, worker_results):
config = TestConfig()
def get_resources():
cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
# Sort the cuda visible devices to have exact match with expected result.
sorted_devices = [
int(device) for device in sorted(cuda_visible_devices.split(","))
]
return sorted_devices
num_workers, expected_results = worker_results
os.environ[ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV] = "1"
e = BackendExecutor(
config, num_workers=num_workers, resources_per_worker={"GPU": 1}
)
e.start()
_start_training(e, get_resources)
results = e.finish_training()
results.sort()
assert results == expected_results
@pytest.mark.parametrize(
"worker_results",
[
(1, [[0]]),
(
2,
[[0]] * 2,
),
(3, [[0, 1]] * 3),
(4, [[0, 1]] * 4),
(5, [[0]] + [[0, 1]] * 4),
(6, [[0]] * 2 + [[0, 1]] * 4),
(7, [[0, 1]] * 7),
(8, [[0, 1]] * 8),
],
)
def test_cuda_visible_devices_fractional(ray_2_node_2_gpu, worker_results):
config = TestConfig()
if worker_results[0] != len(worker_results[1]):
raise ValueError(
"Invalid test parameter. Length of expected result should "
"match number of workers."
)
def get_resources():
cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
# Sort the cuda visible devices to have exact match with expected result.
sorted_devices = [
int(device) for device in sorted(cuda_visible_devices.split(","))
]
return sorted_devices
num_workers, expected_results = worker_results
os.environ[ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV] = "1"
e = BackendExecutor(
config, num_workers=num_workers, resources_per_worker={"GPU": 0.5}
)
e.start()
_start_training(e, get_resources)
results = e.finish_training()
results.sort()
assert results == expected_results
@pytest.mark.parametrize(
"worker_results",
[
(1, [[0, 1]]),
(2, [[0, 1, 2, 3]] * 2),
(3, [[0, 1]] + [[0, 1, 2, 3]] * 2),
(4, [[0, 1, 2, 3]] * 4),
],
)
def test_cuda_visible_devices_multiple(ray_2_node_4_gpu, worker_results):
config = TestConfig()
def get_resources():
cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"]
# Sort the cuda visible devices to have exact match with expected result.
sorted_devices = [
int(device) for device in sorted(cuda_visible_devices.split(","))
]
return sorted_devices
if worker_results[0] != len(worker_results[1]):
raise ValueError(
"Invalid test parameter. Length of expected result should "
"match number of workers."
)
num_workers, expected_results = worker_results
os.environ[ENABLE_SHARE_CUDA_VISIBLE_DEVICES_ENV] = "1"
e = BackendExecutor(
config, num_workers=num_workers, resources_per_worker={"GPU": 2}
)
e.start()
_start_training(e, get_resources)
results = e.finish_training()
results.sort()
assert results == expected_results
@pytest.mark.parametrize(
"worker_results",
[
(1, [[0]]),
(2, [[0, 1]] * 2),
(3, [[0]] + [[0, 1]] * 2),
(4, [[0, 1]] * 4),
],
)
def test_neuron_core_accelerator_ids(ray_2_node_2_neuron_cores, worker_results):
config = TestConfig()
def get_resources():
neuron_resource_ids = os.environ[NEURON_RT_VISIBLE_CORES_ENV_VAR]
# Sort the runtime ids to have exact match with expected result.
sorted_devices = [
int(device) for device in sorted(neuron_resource_ids.split(","))
]
return sorted_devices
num_workers, expected_results = worker_results
# sharing enabled by default
os.environ.pop(ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV, None)
e = BackendExecutor(
config,
num_workers=num_workers,
resources_per_worker={"neuron_cores": 1},
)
e.start()
_start_training(e, get_resources)
results = e.finish_training()
results.sort()
assert results == expected_results
@pytest.mark.parametrize(
"worker_results",
[
(1, [[0]]),
(2, [[0]] + [[1]]),
(3, [[0]] * 2 + [[1]]),
(4, [[0]] * 2 + [[1]] * 2),
],
)
def test_neuron_core_accelerator_ids_sharing_disabled(
ray_2_node_2_neuron_cores, worker_results
):
config = TestConfig()
def get_resources():
neuron_resource_ids = os.environ[NEURON_RT_VISIBLE_CORES_ENV_VAR]
# Sort the runtime ids to have exact match with expected result.
sorted_devices = [
int(device) for device in sorted(neuron_resource_ids.split(","))
]
return sorted_devices
num_workers, expected_results = worker_results
os.environ[ENABLE_SHARE_NEURON_CORES_ACCELERATOR_ENV] = "0"
e = BackendExecutor(
config,
num_workers=num_workers,
resources_per_worker={"neuron_cores": 1},
)
e.start()
_start_training(e, get_resources)
results = e.finish_training()
results.sort()
assert results == expected_results
def get_node_id_set() -> Set[str]:
return {a.node_id for a in list_actors()}
@pytest.mark.parametrize("num_workers", [3, 4, 5])
def test_placement_group_pack(ray_4_node_4_cpu, num_workers):
"""Tests that workers are packed on nodes."""
config = TestConfig()
e = BackendExecutor(config, num_workers=num_workers)
e.start()
node_id_set = get_node_id_set()
assert len(node_id_set) == math.ceil(num_workers / 4)
@pytest.mark.parametrize("num_workers", [3, 4, 5])
def test_placement_group_spread(ray_4_node_4_cpu, num_workers):
"""Tests that workers are spread across nodes."""
os.environ[TRAIN_ENABLE_WORKER_SPREAD_ENV] = "1"
config = TestConfig()
e = BackendExecutor(config, num_workers=num_workers)
e.start()
node_id_set = get_node_id_set()
assert len(node_id_set) == min(num_workers, 4)
@pytest.mark.parametrize("placement_group_capture_child_tasks", [True, False])
def test_placement_group_parent(ray_4_node_4_cpu, placement_group_capture_child_tasks):
"""Tests that parent placement group will be used."""
num_workers = 2
bundle = {"CPU": 1}
bundles = [bundle.copy() for _ in range(num_workers + 1)]
placement_group = ray.util.placement_group(bundles)
def train_func():
return get_current_placement_group().id
@ray.remote
def test():
config = TestConfig()
e = BackendExecutor(config, num_workers=2)
e.start()
_start_training(e, train_func)
return e.finish_training()
results_future = test.options(
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=placement_group,
placement_group_capture_child_tasks=placement_group_capture_child_tasks,
),
).remote()
results = ray.get(results_future)
for worker_result in results:
if placement_group_capture_child_tasks:
assert worker_result == placement_group.id
else:
assert worker_result != placement_group.id
@pytest.mark.parametrize("timeout_s", [5, 0])
@pytest.mark.skipif(
sys.version_info >= (3, 12),
reason="Current jax version is not supported in python 3.12+",
)
def test_jax_distributed_shutdown_timeout(ray_start_2_cpus, monkeypatch, timeout_s):
"""Test that JAX distributed shutdown respects the timeout env var."""
monkeypatch.setenv(JAX_DISTRIBUTED_SHUTDOWN_TIMEOUT_S, str(timeout_s))
jax_config = JaxConfig(use_tpu=True)
e = BackendExecutor(jax_config, num_workers=2)
e.start()
_start_training(e, lambda: 1)
assert e.finish_training() == [1, 1]
# Verify that we do not raise an exception even if we time out
e._backend.on_shutdown(e.worker_group, e._backend_config)
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main(["-v", "-x", __file__]))