def test_balance(inp, expected):
async def test_balance_(*args, **kwargs):
await assert_balanced(inp, expected, *args, **kwargs)
config = {
"distributed.scheduler.default-task-durations":
{str(i): 1
for i in range(10)}
}
gen_cluster(client=True, nthreads=[("", 1)] * len(inp),
config=config)(test_balance_)()
def test_dump_cluster_state_timeout(tmp_path):
sleep_time = 30
async def inner_test(c, s, a, b):
await asyncio.sleep(sleep_time)
# This timeout includes cluster startup and teardown which sometimes can
# take a significant amount of time. For this particular test we would like
# to keep the _test timeout_ small because we intend to trigger it but the
# overall timeout large.
test = gen_cluster(client=True, timeout=5,
cluster_dump_directory=tmp_path)(inner_test)
try:
with pytest.raises(asyncio.TimeoutError) as exc:
test()
assert "inner_test" in str(exc)
assert "await asyncio.sleep(sleep_time)" in str(exc)
except gen.TimeoutError:
pytest.xfail("Cluster startup or teardown took too long")
_, dirs, files = next(os.walk(tmp_path))
assert not dirs
assert files == [inner_test.__name__ + ".yaml"]
import yaml
with open(tmp_path / files[0], "rb") as fd:
state = yaml.load(fd, Loader=yaml.Loader)
assert "scheduler" in state
assert "workers" in state
def test_balance(inp, expected):
test = lambda *args, **kwargs: assert_balanced(inp, expected, *args, **kwargs)
test = gen_cluster(
client=True,
nthreads=[("127.0.0.1", 1)] * len(inp),
config={
"distributed.scheduler.default-task-durations": {
str(i): 1 for i in range(10)
}
},
)(test)
test()
def test_balance(inp, expected):
test = lambda *args, **kwargs: assert_balanced(inp, expected, *args, **
kwargs)
test = gen_cluster(client=True, ncores=[('127.0.0.1', 1)] * len(inp))(test)
test()
def gen_tls_cluster(**kwargs):
kwargs.setdefault('ncores', [('tls://127.0.0.1', 1), ('tls://127.0.0.1', 2)])
return gen_cluster(scheduler='tls://127.0.0.1',
security=tls_only_security(), **kwargs)
def test_balance(inp, expected):
test = lambda *args, **kwargs: assert_balanced(inp, expected, *args, **kwargs)
test = gen_cluster(client=True, ncores=[('127.0.0.1', 1)] * len(inp))(test)
test()
def outer_wrapper(xgboost_gen_cluster):
wrapper = gen_cluster(client=True,
timeout=None,
check_new_threads=False)
return wrapper(func)
