import re
from functools import partial, update_wrapper
from platform import system
from typing import Dict, Union
import numpy as np
import pytest
from hypothesis import HealthCheck, given, settings, strategies
from xgboost import RabitTracker, collective
from xgboost import testing as tm
def test_rabit_tracker() -> None:
tracker = RabitTracker(host_ip="127.0.0.1", n_workers=1)
tracker.start()
args = tracker.worker_args()
port = args["dmlc_tracker_port"]
with collective.CommunicatorContext(**tracker.worker_args()):
ret = collective.broadcast("test1234", 0)
assert str(ret) == "test1234"
if system() == "Windows":
pytest.skip("Windows is not supported.")
with pytest.raises(ValueError, match="Failed to bind socket"):
RabitTracker(host_ip="127.0.0.1", port=port, n_workers=1)
@pytest.mark.skipif(**tm.not_linux())
def test_wait() -> None:
tracker = RabitTracker(host_ip="127.0.0.1", n_workers=2)
tracker.start()
with pytest.raises(ValueError, match="Timeout waiting for the tracker"):
tracker.wait_for(1)
with pytest.raises(ValueError, match="Failed to accept"):
tracker.free()
@pytest.mark.skipif(**tm.not_linux())
def test_socket_error() -> None:
tracker = RabitTracker(host_ip="127.0.0.1", n_workers=2)
tracker.start()
env = tracker.worker_args()
env["dmlc_tracker_port"] = 0
env["dmlc_retry"] = 1
with pytest.raises(ValueError, match="Failed to bootstrap the communication."):
with collective.CommunicatorContext(**env):
pass
with pytest.raises(ValueError):
tracker.free()
def run_rabit_ops(pool, n_workers: int, address: str) -> None:
tracker = RabitTracker(host_ip=address, n_workers=n_workers)
tracker.start()
args = tracker.worker_args()
def local_test(worker_id: int, rabit_args: dict) -> int:
with collective.CommunicatorContext(**rabit_args):
a = 1
assert collective.is_distributed()
arr = np.array([a])
reduced = collective.allreduce(arr, collective.Op.SUM)
assert reduced[0] == n_workers
arr = np.array([worker_id])
reduced = collective.allreduce(arr, collective.Op.MAX)
assert reduced == n_workers - 1
return 1
fn = update_wrapper(partial(local_test, rabit_args=args), local_test)
results = pool.map(fn, range(n_workers))
assert sum(results) == n_workers
@pytest.mark.skipif(**tm.no_loky())
def test_rabit_ops():
from loky import get_reusable_executor
n_workers = 4
with get_reusable_executor(max_workers=n_workers) as pool:
run_rabit_ops(pool, n_workers, "127.0.0.1")
@pytest.mark.skipif(**tm.no_ipv6())
@pytest.mark.skipif(**tm.no_loky())
def test_rabit_ops_ipv6():
from loky import get_reusable_executor
n_workers = 4
with get_reusable_executor(max_workers=n_workers) as pool:
run_rabit_ops(pool, n_workers, "::1")
def run_allreduce(pool, n_workers: int) -> None:
tracker = RabitTracker(host_ip="127.0.0.1", n_workers=n_workers)
tracker.start()
args = tracker.worker_args()
def local_test(worker_id: int, rabit_args: Dict[str, Union[str, int]]) -> None:
x = np.full(shape=(1024 * 1024 * 32), fill_value=1.0)
with collective.CommunicatorContext(**rabit_args):
k = np.asarray([1.0])
for i in range(128):
m = collective.allreduce(k, collective.Op.SUM)
assert m == n_workers
y = collective.allreduce(x, collective.Op.SUM)
np.testing.assert_allclose(y, np.full_like(y, fill_value=float(n_workers)))
fn = update_wrapper(partial(local_test, rabit_args=args), local_test)
results = pool.map(fn, range(n_workers))
for r in results:
assert r is None
@pytest.mark.skipif(**tm.no_loky())
def test_allreduce() -> None:
from loky import get_reusable_executor
n_workers = 4
n_trials = 2
for _ in range(n_trials):
with get_reusable_executor(max_workers=n_workers) as pool:
run_allreduce(pool, n_workers)
def run_broadcast(pool, n_workers: int) -> None:
tracker = RabitTracker(host_ip="127.0.0.1", n_workers=n_workers)
tracker.start()
args = tracker.worker_args()
def local_test(worker_id: int, rabit_args: Dict[str, Union[str, int]]):
with collective.CommunicatorContext(**rabit_args):
res = collective.broadcast(17, 0)
return res
fn = update_wrapper(partial(local_test, rabit_args=args), local_test)
results = pool.map(fn, range(n_workers))
np.testing.assert_allclose(np.array(list(results)), 17)
@pytest.mark.skipif(**tm.no_loky())
def test_broadcast():
from loky import get_reusable_executor
n_workers = 4
n_trials = 2
for _ in range(n_trials):
with get_reusable_executor(max_workers=n_workers) as pool:
run_broadcast(pool, n_workers)
@pytest.mark.skipif(**tm.no_dask())
def test_rank_assignment() -> None:
from distributed import Client, LocalCluster
from xgboost import dask as dxgb
from xgboost.testing.dask import get_rabit_args
def local_test(worker_id):
with dxgb.CommunicatorContext(**args) as ctx:
task_id = ctx["DMLC_TASK_ID"]
matched = re.search(".*-([0-9]).*", task_id)
rank = collective.get_rank()
assert rank == int(matched.group(1))
with LocalCluster(n_workers=8) as cluster:
with Client(cluster) as client:
workers = tm.dask.get_client_workers(client)
args = get_rabit_args(client, len(workers))
futures = client.map(local_test, range(len(workers)), workers=workers)
client.gather(futures)
ops_strategy = strategies.lists(
strategies.sampled_from(["broadcast", "allreduce_max", "allreduce_sum"])
)
@pytest.mark.skipif(**tm.no_loky())
@given(ops=ops_strategy, size=strategies.integers(2**4, 2**16))
@settings(
deadline=None,
print_blob=True,
max_examples=10,
suppress_health_check=[HealthCheck.function_scoped_fixture],
)
def test_ops_restart_comm(ops, size) -> None:
from loky import get_reusable_executor
n_workers = 8
def local_test(w: int, rabit_args: Dict[str, Union[str, int]]) -> None:
a = np.arange(0, n_workers)
with collective.CommunicatorContext(**rabit_args):
for op in ops:
if op == "broadcast":
b = collective.broadcast(a, root=1)
np.testing.assert_allclose(b, a)
elif op == "allreduce_max":
b = collective.allreduce(a, collective.Op.MAX)
np.testing.assert_allclose(b, a)
elif op == "allreduce_sum":
b = collective.allreduce(a, collective.Op.SUM)
np.testing.assert_allclose(a * n_workers, b)
else:
raise ValueError()
with get_reusable_executor(max_workers=n_workers) as pool:
tracker = RabitTracker(host_ip="127.0.0.1", n_workers=n_workers)
tracker.start()
args = tracker.worker_args()
fn = update_wrapper(partial(local_test, rabit_args=args), local_test)
results = pool.map(fn, range(n_workers))
for r in results:
assert r is None
@pytest.mark.skipif(**tm.no_loky())
def test_ops_reuse_comm() -> None:
from loky import get_reusable_executor
rng = np.random.default_rng(1994)
n_examples = 10
ops = rng.choice(
["broadcast", "allreduce_sum", "allreduce_max"], size=n_examples
).tolist()
n_workers = 8
n_trials = 8
def local_test(w: int, rabit_args: Dict[str, Union[str, int]]) -> None:
a = np.arange(0, n_workers)
with collective.CommunicatorContext(**rabit_args):
for op in ops:
if op == "broadcast":
b = collective.broadcast(a, root=1)
assert np.allclose(b, a)
elif op == "allreduce_max":
c = np.full_like(a, collective.get_rank())
b = collective.allreduce(c, collective.Op.MAX)
assert np.allclose(b, n_workers - 1), b
elif op == "allreduce_sum":
b = collective.allreduce(a, collective.Op.SUM)
assert np.allclose(a * 8, b)
else:
raise ValueError()
with get_reusable_executor(max_workers=n_workers) as pool:
for _ in range(n_trials):
tracker = RabitTracker(host_ip="127.0.0.1", n_workers=n_workers)
tracker.start()
args = tracker.worker_args()
fn = update_wrapper(partial(local_test, rabit_args=args), local_test)
results = pool.map(fn, range(n_workers))
for r in results:
assert r is None