def test_scale_up_and_down():
loop = IOLoop.current()
cluster = LocalCluster(0, scheduler_port=0, processes=False, silence_logs=False,
diagnostics_port=None, loop=loop, start=False)
c = Client(cluster, start=False, loop=loop)
yield c._start()
assert not cluster.workers
yield cluster.scale_up(2)
assert len(cluster.workers) == 2
assert len(cluster.scheduler.ncores) == 2
addr = cluster.workers[0].address
yield cluster.scale_down([addr])
assert len(cluster.workers) == 1
assert addr not in cluster.scheduler.ncores
yield c._shutdown()
yield cluster._close()
def test_local_tls_restart(loop):
from distributed.utils_test import tls_only_security
security = tls_only_security()
with LocalCluster(
n_workers=1,
scheduler_port=8786,
silence_logs=False,
security=security,
dashboard_address=False,
host="tls://0.0.0.0",
loop=loop,
) as c:
with Client(c.scheduler.address, loop=loop,
security=security) as client:
workers_before = set(client.scheduler_info()["workers"])
assert client.submit(inc, 1).result() == 2
client.restart()
workers_after = set(client.scheduler_info()["workers"])
assert client.submit(inc, 2).result() == 3
assert workers_before != workers_after
def test_bokeh(loop):
pytest.importorskip('bokeh')
from distributed.http import HTTPScheduler
import requests
with LocalCluster(scheduler_port=0, silence_logs=False, loop=loop,
diagnostics_port=4724, services={('http', 0): HTTPScheduler},
) as c:
start = time()
while True:
with ignoring(Exception):
response = requests.get('http://127.0.0.1:%d/status/' %
c.diagnostics.port)
if response.ok:
break
assert time() < start + 20
sleep(0.01)
start = time()
while not raises(lambda: requests.get('http://127.0.0.1:%d/status/' % 4724)):
assert time() < start + 10
sleep(0.01)
def test_scale_up_and_down():
loop = IOLoop.current()
cluster = yield LocalCluster(0, scheduler_port=0, processes=False,
silence_logs=False, diagnostics_port=None,
loop=loop, asynchronous=True)
c = yield Client(cluster, asynchronous=True)
assert not cluster.workers
yield cluster.scale_up(2)
assert len(cluster.workers) == 2
assert len(cluster.scheduler.ncores) == 2
addr = cluster.workers[0].address
yield cluster.scale_down([addr])
assert len(cluster.workers) == 1
assert addr not in cluster.scheduler.ncores
yield c.close()
yield cluster.close()
async def test_repr(memory_limit, cleanup):
async with LocalCluster(
n_workers=2,
processes=False,
threads_per_worker=2,
memory_limit=memory_limit,
asynchronous=True,
dashboard_address=":0",
) as cluster:
# __repr__ uses cluster.scheduler_info, which slightly lags behind
# cluster.scheduler.workers and client.wait_for_workers.
while len(cluster.scheduler_info["workers"]) < 2:
await asyncio.sleep(0.01)
text = repr(cluster)
assert cluster.scheduler_address in text
assert "workers=2, threads=4" in text
if memory_limit:
assert "memory=4.00 GiB" in text
else:
assert "memory" not in text
async def test_scale_up_and_down():
async with LocalCluster(
n_workers=0,
processes=False,
silence_logs=False,
dashboard_address=":0",
asynchronous=True,
) as cluster:
async with Client(cluster, asynchronous=True) as c:
assert not cluster.workers
cluster.scale(2)
await cluster
assert len(cluster.workers) == 2
assert len(cluster.scheduler.nthreads) == 2
cluster.scale(1)
await cluster
assert len(cluster.workers) == 1
def test_defaults():
from distributed.worker import _ncores
with LocalCluster(scheduler_port=0,
silence_logs=False,
diagnostics_port=None) as c:
assert sum(w.ncores for w in c.workers) == _ncores
assert all(isinstance(w, Nanny) for w in c.workers)
assert all(w.ncores == 1 for w in c.workers)
with LocalCluster(processes=False,
scheduler_port=0,
silence_logs=False,
diagnostics_port=None) as c:
assert sum(w.ncores for w in c.workers) == _ncores
assert all(isinstance(w, Worker) for w in c.workers)
assert len(c.workers) == 1
with LocalCluster(n_workers=2,
scheduler_port=0,
silence_logs=False,
diagnostics_port=None) as c:
if _ncores % 2 == 0:
expected_total_threads = max(2, _ncores)
else:
# n_workers not a divisor of _ncores => threads are overcommitted
expected_total_threads = max(2, _ncores + 1)
assert sum(w.ncores for w in c.workers) == expected_total_threads
with LocalCluster(threads_per_worker=_ncores * 2,
scheduler_port=0,
silence_logs=False,
diagnostics_port=None) as c:
assert len(c.workers) == 1
with LocalCluster(n_workers=_ncores * 2,
scheduler_port=0,
silence_logs=False,
diagnostics_port=None) as c:
assert all(w.ncores == 1 for w in c.workers)
with LocalCluster(threads_per_worker=2,
n_workers=3,
scheduler_port=0,
silence_logs=False,
diagnostics_port=None) as c:
assert len(c.workers) == 3
assert all(w.ncores == 2 for w in c.workers)
async def test_stress():
da = pytest.importorskip("dask.array")
chunksize = "10 MB"
async with LocalCluster(
protocol="ucx",
dashboard_address=":0",
asynchronous=True,
host=HOST,
) as cluster:
async with Client(cluster, asynchronous=True):
rs = da.random.RandomState()
x = rs.random((10000, 10000), chunks=(-1, chunksize))
x = x.persist()
await wait(x)
for i in range(10):
x = x.rechunk((chunksize, -1))
x = x.rechunk((-1, chunksize))
x = x.persist()
await wait(x)
def _test_local_cluster(protocol):
dask.config.update(
dask.config.global_config,
{
"ucx": {
"tcp": True,
"cuda_copy": True,
},
},
priority="new",
)
with LocalCluster(
protocol=protocol,
dashboard_address=None,
n_workers=4,
threads_per_worker=1,
processes=True,
) as cluster:
with Client(cluster) as client:
c = comms.CommsContext(client)
assert sum(c.run(my_rank, 0)) == sum(range(4))
def _test_dataframe_merge_empty_partitions(nrows, npartitions):
with LocalCluster(
protocol="tcp",
dashboard_address=None,
n_workers=npartitions,
threads_per_worker=1,
processes=True,
) as cluster:
with Client(cluster):
df1 = pd.DataFrame({
"key": np.arange(nrows),
"payload1": np.arange(nrows)
})
key = np.arange(nrows)
np.random.shuffle(key)
df2 = pd.DataFrame({"key": key, "payload2": np.arange(nrows)})
expected = df1.merge(df2).set_index("key")
ddf1 = dd.from_pandas(df1, npartitions=npartitions)
ddf2 = dd.from_pandas(df2, npartitions=npartitions)
ddf3 = dataframe_merge(ddf1, ddf2, on="key").set_index("key")
got = ddf3.compute()
pd.testing.assert_frame_equal(got, expected)
def test_no_ipywidgets(loop, monkeypatch):
from unittest.mock import MagicMock
mock_display = MagicMock()
monkeypatch.setitem(sys.modules, "ipywidgets", None)
monkeypatch.setitem(sys.modules, "IPython.display", mock_display)
with LocalCluster(
n_workers=0,
silence_logs=False,
loop=loop,
dashboard_address=":0",
processes=False,
) as cluster:
cluster._ipython_display_()
args, kwargs = mock_display.display.call_args
res = args[0]
assert kwargs == {"raw": True}
assert isinstance(res, dict)
assert "text/plain" in res
assert "text/html" in res
def test_bokeh(loop, processes):
pytest.importorskip('bokeh')
import requests
with LocalCluster(scheduler_port=0,
silence_logs=False,
loop=loop,
processes=processes,
diagnostics_port=0) as c:
bokeh_port = c.scheduler.services['bokeh'].port
url = 'http://127.0.0.1:%d/status/' % bokeh_port
start = time()
while True:
response = requests.get(url)
if response.ok:
break
assert time() < start + 20
sleep(0.01)
# 'localhost' also works
response = requests.get('http://localhost:%d/status/' % bokeh_port)
assert response.ok
with pytest.raises(requests.RequestException):
requests.get(url, timeout=0.2)
def _test_global_option(seg_size):
"""Test setting UCX options through dask's global config"""
dask.config.set({
"ucx": {
"SEG_SIZE": seg_size,
"TLS": "tcp,sockcm,cuda_copy",
"SOCKADDR_TLS_PRIORITY": "sockcm",
}
})
with LocalCluster(
protocol="ucx",
dashboard_address=None,
n_workers=1,
threads_per_worker=1,
processes=True,
) as cluster:
with Client(cluster):
res = da.from_array(numpy.arange(10000), chunks=(1000, ))
res = res.sum().compute()
assert res == 49995000
conf = ucp.get_config()
assert conf["SEG_SIZE"] == seg_size
async def test_cluster_info_sync():
async with LocalCluster(processes=False,
asynchronous=True,
scheduler_sync_interval="1ms") as cluster:
assert cluster._cluster_info["name"] == cluster.name
while "name" not in cluster.scheduler.get_metadata(
keys=["cluster-manager-info"]):
await asyncio.sleep(0.01)
info = await cluster.scheduler_comm.get_metadata(
keys=["cluster-manager-info"])
assert info["name"] == cluster.name
info = cluster.scheduler.get_metadata(keys=["cluster-manager-info"])
assert info["name"] == cluster.name
cluster._cluster_info["foo"] = "bar"
while "foo" not in cluster.scheduler.get_metadata(
keys=["cluster-manager-info"]):
await asyncio.sleep(0.01)
info = cluster.scheduler.get_metadata(keys=["cluster-manager-info"])
assert info["foo"] == "bar"
def test_adapt(loop):
with LocalCluster(scheduler_port=0, silence_logs=False, loop=loop,
diagnostics_port=0, processes=False, n_workers=0) as cluster:
cluster.adapt(minimum=0, maximum=2, interval='10ms')
assert cluster._adaptive.minimum == 0
assert cluster._adaptive.maximum == 2
ref = weakref.ref(cluster._adaptive)
cluster.adapt(minimum=1, maximum=2, interval='10ms')
assert cluster._adaptive.minimum == 1
gc.collect()
# the old Adaptive class sticks around, not sure why
# start = time()
# while ref():
# sleep(0.01)
# gc.collect()
# assert time() < start + 5
start = time()
while len(cluster.scheduler.workers) != 1:
sleep(0.01)
assert time() < start + 5
def _test_initialize_ucx_tcp():
initialize(enable_tcp_over_ucx=True)
with LocalCluster(
protocol="ucx",
dashboard_address=None,
n_workers=1,
threads_per_worker=1,
processes=True,
) as cluster:
with Client(cluster) as client:
res = da.from_array(numpy.arange(10000), chunks=(1000, ))
res = res.sum().compute()
assert res == 49995000
def check_ucx_options():
conf = ucp.get_config()
assert "TLS" in conf
assert "tcp" in conf["TLS"]
assert "sockcm" in conf["TLS"]
assert "cuda_copy" in conf["TLS"]
assert "sockcm" in conf["SOCKADDR_TLS_PRIORITY"]
return True
assert all(client.run(check_ucx_options).values())
def test_scale(loop):
""" Directly calling scale both up and down works as expected """
with LocalCluster(scheduler_port=0,
silence_logs=False,
loop=loop,
diagnostics_port=0,
processes=False,
n_workers=0) as cluster:
assert not cluster.scheduler.workers
cluster.scale(3)
start = time()
while len(cluster.scheduler.workers) != 3:
sleep(0.01)
assert time() < start + 5, len(cluster.scheduler.workers)
sleep(0.2) # let workers settle # TODO: remove need for this
cluster.scale(2)
start = time()
while len(cluster.scheduler.workers) != 2:
sleep(0.01)
assert time() < start + 5, len(cluster.scheduler.workers)
def test_local_tls(loop):
from distributed.utils_test import tls_only_security
security = tls_only_security()
with LocalCluster(scheduler_port=8786,
silence_logs=False,
security=security,
diagnostics_port=False,
ip='tls://0.0.0.0',
loop=loop) as c:
sync(loop,
assert_can_connect_from_everywhere_4,
c.scheduler.port,
connection_args=security.get_connection_args('client'),
protocol='tls',
timeout=3)
# If we connect to a TLS localculster without ssl information we should fail
sync(
loop,
assert_cannot_connect,
addr='tcp://127.0.0.1:%d' % c.scheduler.port,
connection_args=security.get_connection_args('client'),
exception_class=RuntimeError,
)
def _test_dataframe_merge(backend, protocol, n_workers):
with LocalCluster(
protocol=protocol,
dashboard_address=None,
n_workers=n_workers,
threads_per_worker=1,
processes=True,
) as cluster:
with Client(cluster):
nrows = n_workers * 10
# Let's make some dataframes that we can join on the "key" column
df1 = pd.DataFrame({"key": np.arange(nrows), "payload1": np.arange(nrows)})
key = np.arange(nrows)
np.random.shuffle(key)
df2 = pd.DataFrame(
{"key": key[nrows // 3 :], "payload2": np.arange(nrows)[nrows // 3 :]}
)
expected = df1.merge(df2).set_index("key")
if backend == "cudf":
df1 = cudf.DataFrame.from_pandas(df1)
df2 = cudf.DataFrame.from_pandas(df2)
ddf1 = dd.from_pandas(df1, npartitions=n_workers + 1)
ddf2 = dd.from_pandas(
df2, npartitions=n_workers - 1 if n_workers > 1 else 1
)
ddf3 = dataframe_merge(ddf1, ddf2, on="key").set_index("key")
got = ddf3.compute()
if backend == "cudf":
got = got.to_pandas()
got.index.names = ["key"] # TODO: this shouldn't be needed
pd.testing.assert_frame_equal(got, expected)
def test_ucx_localcluster(loop, processes):
if processes:
kwargs = {"env": {"UCX_MEMTYPE_CACHE": "n"}}
else:
kwargs = {}
ucx_addr = ucp.get_address()
with LocalCluster(
protocol="ucx",
interface="ib0",
dashboard_address=None,
n_workers=2,
threads_per_worker=1,
processes=processes,
loop=loop,
**kwargs
) as cluster:
with Client(cluster) as client:
x = client.submit(inc, 1)
x.result()
assert x.key in cluster.scheduler.tasks
if not processes:
assert any(w.data == {x.key: 2} for w in cluster.workers.values())
assert len(cluster.scheduler.workers) == 2
def test_defaults():
from distributed.worker import _ncores
with LocalCluster(scheduler_port=0,
silence_logs=False,
diagnostic_port=None) as c:
assert sum(w.ncores for w in c.workers) == _ncores
assert all(isinstance(w, Nanny) for w in c.workers)
assert all(w.ncores == 1 for w in c.workers)
with LocalCluster(nanny=False,
scheduler_port=0,
silence_logs=False,
diagnostic_port=None) as c:
assert sum(w.ncores for w in c.workers) == _ncores
assert all(isinstance(w, Worker) for w in c.workers)
assert len(c.workers) == 1
with LocalCluster(n_workers=2,
scheduler_port=0,
silence_logs=False,
diagnostic_port=None) as c:
assert sum(w.ncores for w in c.workers) == max(2, _ncores)
with LocalCluster(threads_per_worker=_ncores * 2,
scheduler_port=0,
silence_logs=False,
diagnostic_port=None) as c:
assert len(c.workers) == 1
with LocalCluster(n_workers=_ncores * 2,
scheduler_port=0,
silence_logs=False,
diagnostic_port=None) as c:
assert all(w.ncores == 1 for w in c.workers)
with LocalCluster(threads_per_worker=2,
n_workers=3,
scheduler_port=0,
silence_logs=False,
diagnostic_port=None) as c:
assert len(c.workers) == 3
assert all(w.ncores == 2 for w in c.workers)
def test_memory(loop):
with LocalCluster(scheduler_port=0, processes=False, silence_logs=False,
diagnostics_port=None, loop=loop) as cluster:
assert sum(w.memory_limit for w in cluster.workers) <= TOTAL_MEMORY
def test_only_local_access(loop):
with LocalCluster(scheduler_port=0,
silence_logs=False,
diagnostics_port=None,
loop=loop) as c:
sync(loop, assert_can_connect_locally_4, c.scheduler.port)
def test_protocol_ip(loop):
with LocalCluster(
host="tcp://127.0.0.2", loop=loop, n_workers=0, processes=False
) as cluster:
assert cluster.scheduler.address.startswith("tcp://127.0.0.2")
def test_death_timeout_raises(loop):
with pytest.raises(gen.TimeoutError):
with LocalCluster(scheduler_port=0, silence_logs=False,
death_timeout=1e-10, diagnostics_port=None,
loop=loop) as cluster:
pass
def test_protocol_tcp(loop):
with LocalCluster(
protocol="tcp", loop=loop, n_workers=0, processes=False
) as cluster:
assert cluster.scheduler.address.startswith("tcp://")
def test_protocol_inproc(loop):
with LocalCluster(protocol="inproc://", loop=loop, processes=False) as cluster:
assert cluster.scheduler.address.startswith("inproc://")
def test_only_local_access(loop):
with LocalCluster(
0, scheduler_port=0, silence_logs=False, dashboard_address=None, loop=loop
) as c:
sync(loop, assert_can_connect_locally_4, c.scheduler.port)
async def test_defaults(cleanup):
async with LocalCluster(
scheduler_port=0, silence_logs=False, dashboard_address=None, asynchronous=True
) as c:
assert sum(w.nthreads for w in c.workers.values()) == CPU_COUNT
assert all(isinstance(w, Nanny) for w in c.workers.values())
async def test_async_with():
async with LocalCluster(processes=False, asynchronous=True) as cluster:
w = cluster.workers
assert w
assert not w