def test_use_config_file(loop):
with new_config_file(tls_only_config()):
with popen(['dask-scheduler', '--no-bokeh', '--host', 'tls://']) as s:
with popen(['dask-worker', '--no-bokeh', 'tls://127.0.0.1:8786']) as w:
with Client('tls://127.0.0.1:8786', loop=loop,
security=tls_security()) as c:
wait_for_cores(c)
def test_pid_file(loop):
with tmpfile() as s:
with popen(['dask-scheduler', '--pid-file', s]) as sched:
while not os.path.exists(s):
sleep(0.01)
text = False
while not text:
sleep(0.01)
with open(s) as f:
text = f.read()
pid = int(text)
assert sched.pid == pid
with tmpfile() as w:
with popen(['dask-worker', '127.0.0.1:8786', '--pid-file', w]) as worker:
while not os.path.exists(w):
sleep(0.01)
text = False
while not text:
sleep(0.01)
with open(w) as f:
text = f.read()
pid = int(text)
assert worker.pid == pid
def test_pid_file(loop):
def check_pidfile(proc, pidfile):
while not os.path.exists(pidfile):
sleep(0.01)
text = False
while not text:
sleep(0.01)
with open(pidfile) as f:
text = f.read()
pid = int(text)
if sys.platform.startswith('win'):
# On Windows, `dask-XXX` invokes the dask-XXX.exe
# shim, but the PID is written out by the child Python process
assert pid
else:
assert proc.pid == pid
with tmpfile() as s:
with popen(['dask-scheduler', '--pid-file', s, '--no-bokeh']) as sched:
check_pidfile(sched, s)
with tmpfile() as w:
with popen(['dask-worker', '127.0.0.1:8786', '--pid-file', w,
'--no-bokeh']) as worker:
check_pidfile(worker, w)
def test_scheduler_file(loop, nanny):
with tmpfile() as fn:
with popen(['dask-scheduler', '--no-bokeh', '--scheduler-file', fn]) as sched:
with popen(['dask-worker', '--scheduler-file', fn, nanny, '--no-bokeh']):
with Client(scheduler_file=fn, loop=loop) as c:
start = time()
while not c.scheduler_info()['workers']:
sleep(0.1)
assert time() < start + 10
def test_multiple_workers(loop):
with popen(['dask-scheduler', '--no-bokeh']) as s:
with popen(['dask-worker', 'localhost:8786']) as a:
with popen(['dask-worker', 'localhost:8786']) as b:
with Client('127.0.0.1:%d' % Scheduler.default_port, loop=loop) as c:
start = time()
while len(c.ncores()) < 2:
sleep(0.1)
assert time() < start + 10
def test_resources(loop):
with popen(['dask-scheduler', '--no-bokeh']) as sched:
with popen(['dask-worker', 'tcp://127.0.0.1:8786', '--no-bokeh',
'--resources', 'A=1 B=2,C=3']) as worker:
with Client('127.0.0.1:8786', loop=loop) as c:
while not c.scheduler_info()['workers']:
sleep(0.1)
info = c.scheduler_info()
worker = list(info['workers'].values())[0]
assert worker['resources'] == {'A': 1, 'B': 2, 'C': 3}
def test_no_reconnect(nanny, loop):
with popen(['dask-scheduler', '--no-bokeh']) as sched:
wait_for_port('127.0.0.1:8786')
with popen(['dask-worker', '127.0.0.1:8786', '--no-reconnect', nanny,
'--no-bokeh']) as worker:
sleep(2)
terminate_process(sched)
start = time()
while isalive(worker):
sleep(0.1)
assert time() < start + 10
def test_memory_limit(loop):
with popen(['dask-scheduler', '--no-bokeh']) as sched:
with popen(['dask-worker', '127.0.0.1:8786', '--memory-limit', '2e3MB',
'--no-bokeh']) as worker:
with Client('127.0.0.1:8786', loop=loop) as c:
while not c.ncores():
sleep(0.1)
info = c.scheduler_info()
d = first(info['workers'].values())
assert isinstance(d['memory_limit'], int)
assert d['memory_limit'] == 2e9
def test_no_reconnect(nanny, loop):
with popen(['dask-scheduler', '--no-bokeh']) as sched:
wait_for_port(('127.0.0.1', 8786))
with popen(['dask-worker', 'tcp://127.0.0.1:8786', '--no-reconnect', nanny,
'--no-bokeh']) as worker:
sleep(2)
terminate_process(sched)
start = time()
while worker.poll() is None:
sleep(0.1)
assert time() < start + 10
def test_local_directory(loop, nanny):
with tmpfile() as fn:
with popen(['dask-scheduler', '--no-bokeh']) as sched:
with popen(['dask-worker', '127.0.0.1:8786', nanny,
'--no-bokeh', '--local-directory', fn]) as worker:
with Client('127.0.0.1:8786', loop=loop, timeout=10) as c:
start = time()
while not c.scheduler_info()['workers']:
sleep(0.1)
assert time() < start + 8
info = c.scheduler_info()
worker = list(info['workers'].values())[0]
assert worker['local_directory'].startswith(fn)
def test_nanny_worker_ports(loop):
with popen(['dask-scheduler', '--port', '8989']) as sched:
with popen(['dask-worker', '127.0.0.1:8989', '--host', '127.0.0.1',
'--worker-port', '8788', '--nanny-port', '8789']) as worker:
with Client('127.0.0.1:8989', loop=loop) as c:
start = time()
while True:
d = sync(c.loop, c.scheduler.identity)
if d['workers']:
break
else:
assert time() - start < 5
sleep(0.1)
assert d['workers']['127.0.0.1:8788']['services']['nanny'] == 8789
def test_nanny_worker_ports(loop):
with popen(['dask-scheduler', '--port', '9359', '--no-bokeh']) as sched:
with popen(['dask-worker', '127.0.0.1:9359', '--host', '127.0.0.1',
'--worker-port', '9684', '--nanny-port', '5273',
'--no-bokeh']) as worker:
with Client('127.0.0.1:9359', loop=loop) as c:
start = time()
while True:
d = sync(c.loop, c.scheduler.identity)
if d['workers']:
break
else:
assert time() - start < 5
sleep(0.1)
assert d['workers']['tcp://127.0.0.1:9684']['services']['nanny'] == 5273
def test_bokeh_whitelist(loop):
pytest.importorskip('bokeh')
with pytest.raises(Exception):
requests.get('http://localhost:8787/status/').ok
with popen(['dask-scheduler', '--bokeh-whitelist', '127.0.0.2:8787',
'--bokeh-whitelist', '127.0.0.3:8787']) as proc:
with Client('127.0.0.1:%d' % Scheduler.default_port, loop=loop) as c:
pass
while True:
line = proc.stderr.readline()
if b'Bokeh UI' in line:
break
start = time()
while True:
try:
for name in ['127.0.0.2', '127.0.0.3']:
response = requests.get('http://%s:8787/status/' % name)
assert response.ok
break
except Exception as f:
print(f)
sleep(0.1)
assert time() < start + 20
def test_bokeh(loop):
pytest.importorskip('bokeh')
with popen(['dask-scheduler']) as proc:
with Client('127.0.0.1:%d' % Scheduler.default_port, loop=loop) as c:
pass
while True:
line = proc.stderr.readline()
if b'Bokeh UI' in line:
break
names = ['localhost', '127.0.0.1', get_ip()]
if 'linux' in sys.platform:
names.append(socket.gethostname())
start = time()
while True:
try:
for name in names:
response = requests.get('http://%s:8787/status/' % name)
assert response.ok
break
except Exception as f:
print(f)
sleep(0.1)
assert time() < start + 10
with pytest.raises(Exception):
requests.get('http://127.0.0.1:8787/status/')
def test_nprocs_expands_name(loop):
with popen(['dask-scheduler', '--no-bokeh']) as sched:
with popen(['dask-worker', '127.0.0.1:8786', '--nprocs', '2',
'--name', 'foo']) as worker:
with popen(['dask-worker', '127.0.0.1:8786', '--nprocs', '2']) as worker:
with Client('tcp://127.0.0.1:8786', loop=loop) as c:
start = time()
while len(c.scheduler_info()['workers']) < 4:
sleep(0.2)
assert time() < start + 10
info = c.scheduler_info()
names = [d['name'] for d in info['workers'].values()]
foos = [n for n in names if n.startswith('foo')]
assert len(foos) == 2
assert len(set(names)) == 4
def test_bokeh_non_standard_ports(loop):
pytest.importorskip('bokeh')
with popen(['dask-scheduler',
'--port', '3448',
'--http-port', '4824',
'--bokeh-port', '4832']) as proc:
with Client('127.0.0.1:3448', loop=loop) as c:
pass
while True:
line = proc.stderr.readline()
if b'Bokeh UI' in line:
break
start = time()
while True:
try:
response = requests.get('http://localhost:4832/status/')
assert response.ok
break
except:
sleep(0.1)
assert time() < start + 20
with pytest.raises(Exception):
requests.get('http://127.0.0.1:4832/status/')
def test_no_bokeh(loop):
pytest.importorskip('bokeh')
with popen(['dask-scheduler', '--no-bokeh']) as proc:
with Executor('127.0.0.1:%d' % Scheduler.default_port, loop=loop) as e:
for i in range(3):
line = proc.stderr.readline()
assert b'bokeh' not in line.lower()
def test_no_bokeh(loop):
pytest.importorskip('bokeh')
with popen(['dask-scheduler', '--no-bokeh']) as proc:
with Client('127.0.0.1:%d' % Scheduler.default_port, loop=loop) as c:
for i in range(3):
line = proc.stderr.readline()
assert b'bokeh' not in line.lower()
with pytest.raises(Exception):
requests.get('http://127.0.0.1:8787/status/')
def test_respect_host_listen_address(loop, nanny, host):
with popen(['dask-scheduler', '--no-bokeh']) as sched:
with popen(['dask-worker', '127.0.0.1:8786',
nanny, '--no-bokeh',
'--host', host]) as worker:
with Client('127.0.0.1:8786') as client:
while not client.ncores():
sleep(0.1)
info = client.scheduler_info()
# roundtrip works
assert client.submit(lambda x: x + 1, 10).result() == 11
def func(dask_worker):
return dask_worker.listener.listen_address
listen_addresses = client.run(func)
assert all(host in v for v in listen_addresses.values())
def test_contact_listen_address(loop, nanny, listen_address):
with popen(['dask-scheduler', '--no-bokeh']) as sched:
with popen(['dask-worker', '127.0.0.1:8786',
nanny, '--no-bokeh',
'--contact-address', 'tcp://127.0.0.2:39837',
'--listen-address', listen_address]) as worker:
with Client('127.0.0.1:8786') as client:
while not client.ncores():
sleep(0.1)
info = client.scheduler_info()
assert 'tcp://127.0.0.2:39837' in info['workers']
# roundtrip works
assert client.submit(lambda x: x + 1, 10).result() == 11
def func(dask_worker):
return dask_worker.listener.listen_address
assert client.run(func) == {'tcp://127.0.0.2:39837': listen_address}
def test_no_scheduler(loop):
with tmpfile() as fn:
with popen(['mpirun', '--np', '2', 'dask-mpi', '--scheduler-file', fn],
stdin=subprocess.DEVNULL):
with Client(scheduler_file=fn) as c:
start = time()
while len(c.scheduler_info()['workers']) != 1:
assert time() < start + 10
sleep(0.2)
assert c.submit(lambda x: x + 1, 10).result() == 11
with popen(['mpirun', '--np', '1', 'dask-mpi',
'--scheduler-file', fn, '--no-scheduler']):
start = time()
while len(c.scheduler_info()['workers']) != 2:
assert time() < start + 10
sleep(0.2)
def test_defaults(loop):
with popen(['dask-scheduler', '--no-bokeh']) as proc:
with Client('127.0.0.1:%d' % Scheduler.default_port, loop=loop) as c:
response = requests.get('http://127.0.0.1:9786/info.json')
assert response.ok
assert response.json()['status'] == 'running'
with pytest.raises(Exception):
requests.get('http://127.0.0.1:9786/info.json')
with pytest.raises(Exception):
requests.get('http://127.0.0.1:8787/status/')
def test_basic(loop, nanny):
with tmpfile() as fn:
with popen(['mpirun', '--np', '4', 'dask-mpi', '--scheduler-file', fn, nanny],
stdin=subprocess.DEVNULL):
with Client(scheduler_file=fn) as c:
start = time()
while len(c.scheduler_info()['workers']) != 3:
assert time() < start + 10
sleep(0.2)
assert c.submit(lambda x: x + 1, 10, workers=1).result() == 11
def test_defaults(loop):
with popen(["dask-scheduler", "--no-dashboard"]) as proc:
async def f():
# Default behaviour is to listen on all addresses
await assert_can_connect_from_everywhere_4_6(8786, timeout=5.0)
with Client("127.0.0.1:%d" % Scheduler.default_port, loop=loop) as c:
c.sync(f)
response = requests.get("http://127.0.0.1:8787/status/")
assert response.status_code == 404
with pytest.raises(Exception):
response = requests.get("http://127.0.0.1:9786/info.json")
def test_dashboard_non_standard_ports(loop):
pytest.importorskip("bokeh")
try:
import jupyter_server_proxy # noqa: F401
proxy_exists = True
except ImportError:
proxy_exists = False
with popen(["dask-scheduler", "--port", "3449"]):
with popen([
"dask-worker", "tcp://127.0.0.1:3449", "--dashboard-address",
":4833"
]) as proc:
with Client("127.0.0.1:3449", loop=loop) as c:
pass
start = time()
while True:
try:
response = requests.get("http://127.0.0.1:4833/status")
assert response.ok
redirect_resp = requests.get("http://127.0.0.1:4833/main")
redirect_resp.ok
# TEST PROXYING WORKS
if proxy_exists:
url = "http://127.0.0.1:8787/proxy/4833/127.0.0.1/status"
response = requests.get(url)
assert response.ok
break
except Exception:
sleep(0.5)
assert time() < start + 20
with pytest.raises(Exception):
requests.get("http://localhost:4833/status/")
async def test_nanny_worker_ports(c, s):
with popen([
"dask-worker",
s.address,
"--host",
"127.0.0.1",
"--worker-port",
"9684",
"--nanny-port",
"5273",
"--no-dashboard",
]):
await c.wait_for_workers(1)
d = await c.scheduler.identity()
assert d["workers"]["tcp://127.0.0.1:9684"][
"nanny"] == "tcp://127.0.0.1:5273"
def test_defaults(loop):
with popen(["dask-scheduler", "--no-bokeh"]) as proc:
@gen.coroutine
def f():
# Default behaviour is to listen on all addresses
yield [assert_can_connect_from_everywhere_4_6(8786,
5.0)] # main port
with Client("127.0.0.1:%d" % Scheduler.default_port, loop=loop) as c:
c.sync(f)
with pytest.raises(Exception):
requests.get("http://127.0.0.1:8787/status/")
with pytest.raises(Exception):
response = requests.get("http://127.0.0.1:9786/info.json")
def test_hostport(loop):
with popen(['dask-scheduler', '--no-bokeh', '--host', '127.0.0.1:8978',
'--http-port', '8979']):
@gen.coroutine
def f():
yield [
# The scheduler's main port can't be contacted from the outside
assert_can_connect_locally_4(8978, 2.0),
# ... but its HTTP port can
assert_can_connect_from_everywhere_4_6(8979, 2.0),
]
loop.run_sync(f)
with Client('127.0.0.1:8978', loop=loop) as c:
assert len(c.ncores()) == 0
def test_preload_file(loop):
def check_scheduler():
import scheduler_info
return scheduler_info.get_scheduler_address()
tmpdir = tempfile.mkdtemp()
try:
path = os.path.join(tmpdir, "scheduler_info.py")
with open(path, "w") as f:
f.write(PRELOAD_TEXT)
with tmpfile() as fn:
with popen(["dask-scheduler", "--scheduler-file", fn, "--preload", path]):
with Client(scheduler_file=fn, loop=loop) as c:
assert c.run_on_scheduler(check_scheduler) == c.scheduler.address
finally:
shutil.rmtree(tmpdir)
async def test_no_reconnect(c, s, nanny):
with popen([
"dask-worker",
s.address,
"--no-reconnect",
nanny,
"--no-dashboard",
]) as worker:
# roundtrip works
assert await c.submit(lambda x: x + 1, 10) == 11
(comm, ) = s.stream_comms.values()
comm.abort()
# worker terminates as soon as the connection is aborted
await to_thread(worker.communicate, timeout=5)
assert worker.returncode == 0
def test_dashboard(loop, mpirun):
with tmpfile(extension="json") as fn:
cmd = mpirun + [
"-np",
"2",
"dask-mpi",
"--scheduler-file",
fn,
"--dashboard-address",
":59583",
]
with popen(cmd, stdin=FNULL):
check_port_okay(59583)
with pytest.raises(Exception):
requests.get("http://localhost:59583/status/")
def test_bokeh_external_port(loop):
pytest.importorskip('bokeh')
with popen(['dask-scheduler', '--bokeh-external-port', '4833']) as proc:
with Client('127.0.0.1:8786', loop=loop) as c:
d = c.scheduler_info()
start = time()
while True:
try:
response = requests.get('http://localhost:4833/workers/')
assert response.ok
break
except:
sleep(0.1)
assert time() < start + 20
with pytest.raises(Exception):
requests.get('http://localhost:4833/workers/')
async def test_preload_config(c, s):
# Ensure dask-worker pulls the preload from the Dask config if
# not specified via a command line option
preload_text = """
def dask_setup(worker):
worker.foo = 'setup'
"""
env = os.environ.copy()
env["DASK_DISTRIBUTED__WORKER__PRELOAD"] = preload_text
with popen(
[
"dask-worker",
s.address,
],
env=env,
):
await c.wait_for_workers(1)
[foo] = (await c.run(lambda dask_worker: dask_worker.foo)).values()
assert foo == "setup"
def test_bokeh(loop):
with tmpfile() as fn:
with popen(['mpirun', '--np', '2', 'dask-mpi', '--scheduler-file', fn,
'--bokeh-port', '59583', '--bokeh-worker-port', '59584'],
stdin=subprocess.DEVNULL):
for port in [59853, 59584]:
start = time()
while True:
try:
response = requests.get('http://localhost:%d/status/' % port)
assert response.ok
break
except Exception:
sleep(0.1)
assert time() < start + 20
with pytest.raises(Exception):
requests.get('http://localhost:59583/status/')
def test_defaults(loop):
with popen(['dask-scheduler', '--no-bokeh']) as proc:
@gen.coroutine
def f():
# Default behaviour is to listen on all addresses
yield [
assert_can_connect_from_everywhere_4_6(8786, 5.0), # main port
]
loop.run_sync(f)
with Client('127.0.0.1:%d' % Scheduler.default_port, loop=loop) as c:
pass
with pytest.raises(Exception):
requests.get('http://127.0.0.1:8787/status/')
with pytest.raises(Exception):
response = requests.get('http://127.0.0.1:9786/info.json')
async def test_nanny_worker_port_range_too_many_workers_raises(s):
with popen(
[
"dask-worker",
s.address,
"--nworkers",
"3",
"--host",
"127.0.0.1",
"--worker-port",
"9684:9685",
"--nanny-port",
"9686:9687",
"--no-dashboard",
],
flush_output=False,
) as worker:
assert any(b"Could not start" in worker.stdout.readline()
for _ in range(100))
def test_non_default_ports(loop, nanny, mpirun):
with tmpfile(extension="json") as fn:
cmd = mpirun + [
"-np",
"2",
"dask-mpi",
"--scheduler-file",
fn,
nanny,
"--scheduler-port",
"56723",
]
with popen(cmd):
with Client(scheduler_file=fn) as c:
sched_info = c.scheduler_info()
_, sched_port = get_address_host_port(sched_info["address"])
assert sched_port == 56723
def test_bokeh_non_standard_ports(loop):
pytest.importorskip('bokeh')
with popen(['dask-scheduler', '--port', '3448', '--bokeh-port',
'4832']) as proc:
with Client('127.0.0.1:3448', loop=loop) as c:
pass
start = time()
while True:
try:
response = requests.get('http://localhost:4832/status/')
assert response.ok
break
except Exception:
sleep(0.1)
assert time() < start + 20
with pytest.raises(Exception):
requests.get('http://localhost:4832/status/')
def test_bokeh(loop):
with tmpfile() as fn:
with popen(['mpirun', '--np', '2', 'dask-mpi', '--scheduler-file', fn,
'--bokeh-port', '59583', '--bokeh-worker-port', '59584'],
stdin=subprocess.DEVNULL):
for port in [59853, 59584]:
start = time()
while True:
try:
response = requests.get('http://localhost:%d/status/' % port)
assert response.ok
break
except Exception:
sleep(0.1)
assert time() < start + 20
with pytest.raises(Exception):
requests.get('http://localhost:59583/status/')
async def test_nanny_worker_port_range(c, s):
with popen([
"dask-worker",
s.address,
"--nworkers",
"3",
"--host",
"127.0.0.1",
"--worker-port",
"10000:11000",
"--nanny-port",
"11000:12000",
"--no-dashboard",
]):
await c.wait_for_workers(3)
worker_ports = await c.run(lambda dask_worker: dask_worker.port)
assert all(10000 <= p <= 11000 for p in worker_ports.values())
nanny_ports = await c.run(lambda dask_worker: dask_worker.port,
nanny=True)
assert all(11000 <= p <= 12000 for p in nanny_ports.values())
def test_preload_module(loop):
def check_scheduler():
import scheduler_info
return scheduler_info.get_scheduler_address()
tmpdir = tempfile.mkdtemp()
try:
path = os.path.join(tmpdir, 'scheduler_info.py')
with open(path, 'w') as f:
f.write(PRELOAD_TEXT)
with tmpfile() as fn:
with popen(['dask-scheduler', '--scheduler-file', fn,
'--preload', 'scheduler_info'],
env=dict(os.environ, PYTHONPATH=tmpdir)):
with Client(scheduler_file=fn, loop=loop) as c:
assert c.run_on_scheduler(check_scheduler) == \
c.scheduler.address
finally:
shutil.rmtree(tmpdir)
def test_dashboard_whitelist(loop):
pytest.importorskip("bokeh")
with pytest.raises(Exception):
requests.get("http://localhost:8787/status/").ok
with popen(["dask-scheduler"]) as proc:
with Client("127.0.0.1:%d" % Scheduler.default_port, loop=loop) as c:
pass
start = time()
while True:
try:
for name in ["127.0.0.2", "127.0.0.3"]:
response = requests.get("http://%s:8787/status/" % name)
assert response.ok
break
except Exception as f:
print(f)
sleep(0.1)
assert time() < start + 20
def test_dashboard_non_standard_ports(loop):
pytest.importorskip("bokeh")
with popen(
["dask-scheduler", "--port", "3448", "--dashboard-address",
":4832"]) as proc:
with Client("127.0.0.1:3448", loop=loop) as c:
pass
start = time()
while True:
try:
response = requests.get("http://localhost:4832/status/")
assert response.ok
break
except Exception:
sleep(0.1)
assert time() < start + 20
with pytest.raises(Exception):
requests.get("http://localhost:4832/status/")
async def test_preload_config(cleanup):
# Ensure dask-worker pulls the preload from the Dask config if
# not specified via a command line option
preload_text = """
def dask_setup(worker):
worker.foo = 'setup'
"""
async with Scheduler(port=0) as s:
async with Client(s.address, asynchronous=True) as c:
env = {"DASK_DISTRIBUTED__WORKER__PRELOAD": preload_text}
with popen(
[
"dask-worker",
s.address,
],
env=env,
):
await c.wait_for_workers(1)
[foo] = (await
c.run(lambda dask_worker: dask_worker.foo)).values()
assert foo == "setup"
def test_bokeh_whitelist(loop):
pytest.importorskip('bokeh')
with pytest.raises(Exception):
requests.get('http://localhost:8787/status/').ok
with popen(['dask-scheduler', '--bokeh-whitelist', '127.0.0.2:8787',
'--bokeh-whitelist', '127.0.0.3:8787']) as proc:
with Client('127.0.0.1:%d' % Scheduler.default_port, loop=loop) as c:
pass
start = time()
while True:
try:
for name in ['127.0.0.2', '127.0.0.3']:
response = requests.get('http://%s:8787/status/' % name)
assert response.ok
break
except Exception as f:
print(f)
sleep(0.1)
assert time() < start + 20
def test_preload_command(loop):
def check_passthrough():
import passthrough_info
return passthrough_info.get_passthrough()
tmpdir = tempfile.mkdtemp()
try:
path = os.path.join(tmpdir, 'passthrough_info.py')
with open(path, 'w') as f:
f.write(PRELOAD_COMMAND_TEXT)
with tmpfile() as fn:
print(fn)
with popen(['dask-scheduler', '--scheduler-file', fn,
'--preload', path, "--passthrough", "foobar"]):
with Client(scheduler_file=fn, loop=loop) as c:
assert c.run_on_scheduler(check_passthrough) == \
"foobar"
finally:
shutil.rmtree(tmpdir)
async def test_reconnect(c, s, nanny):
with popen([
"dask-worker",
s.address,
"--reconnect",
nanny,
"--no-dashboard",
]) as worker:
# roundtrip works
await c.submit(lambda x: x + 1, 10) == 11
(comm, ) = s.stream_comms.values()
comm.abort()
# roundtrip still works, which means the worker reconnected
assert await c.submit(lambda x: x + 1, 11) == 12
# closing the scheduler cleanly does terminate the worker
await s.close()
await to_thread(worker.communicate, timeout=5)
assert worker.returncode == 0
def test_non_default_ports(loop, nanny, mpirun):
with tmpfile(extension="json") as fn:
cmd = mpirun + [
"-np",
"2",
"dask-mpi",
"--scheduler-file",
fn,
nanny,
"--scheduler-port",
"56723",
"--worker-port",
"58464",
"--nanny-port",
"50164",
]
with popen(cmd):
with Client(scheduler_file=fn) as c:
start = time()
while len(c.scheduler_info()["workers"]) != 1:
assert time() < start + 10
sleep(0.2)
sched_info = c.scheduler_info()
sched_host, sched_port = get_address_host_port(
sched_info["address"])
assert sched_port == 56723
for worker_addr, worker_info in sched_info["workers"].items():
worker_host, worker_port = get_address_host_port(
worker_addr)
assert worker_port == 58464
if nanny == "--nanny":
_, nanny_port = get_address_host_port(
worker_info["nanny"])
assert nanny_port == 50164
assert c.submit(lambda x: x + 1, 10).result() == 11
async def test_worker_class(cleanup, tmp_path, nanny):
# Create module with custom worker class
WORKER_CLASS_TEXT = """
from distributed.worker import Worker
class MyWorker(Worker):
pass
"""
tmpdir = str(tmp_path)
tmpfile = str(tmp_path / "myworker.py")
with open(tmpfile, "w") as f:
f.write(WORKER_CLASS_TEXT)
# Put module on PYTHONPATH
env = os.environ.copy()
if "PYTHONPATH" in env:
env["PYTHONPATH"] = tmpdir + ":" + env["PYTHONPATH"]
else:
env["PYTHONPATH"] = tmpdir
async with Scheduler(port=0) as s:
async with Client(s.address, asynchronous=True) as c:
with popen(
[
"dask-worker",
s.address,
nanny,
"--worker-class",
"myworker.MyWorker",
],
env=env,
) as worker:
await c.wait_for_workers(1)
def worker_type(dask_worker):
return type(dask_worker).__name__
worker_types = await c.run(worker_type)
assert all(name == "MyWorker"
for name in worker_types.values())
async def test_contact_listen_address(c, s, nanny, listen_address):
with popen([
"dask-worker",
s.address,
nanny,
"--no-dashboard",
"--contact-address",
"tcp://127.0.0.2:39837",
"--listen-address",
listen_address,
]):
await c.wait_for_workers(1)
info = c.scheduler_info()
assert info["workers"].keys() == {"tcp://127.0.0.2:39837"}
# roundtrip works
assert await c.submit(lambda x: x + 1, 10) == 11
def func(dask_worker):
return dask_worker.listener.listen_address
assert await c.run(func) == {"tcp://127.0.0.2:39837": listen_address}
def test_bokeh(loop):
pytest.importorskip('bokeh')
with popen(['dask-scheduler']) as proc:
with Executor('127.0.0.1:%d' % Scheduler.default_port, loop=loop) as e:
pass
while True:
line = proc.stderr.readline()
if b'Bokeh UI' in line:
break
start = time()
while True:
try:
for name in [socket.gethostname(), 'localhost', '127.0.0.1', get_ip()]:
response = requests.get('http://%s:8787/status/' % name)
assert response.ok
break
except Exception as f:
print(f)
sleep(0.1)
assert time() < start + 10
def test_preload_command_default(loop):
def check_passthrough():
import passthrough_info
return passthrough_info.get_passthrough()
tmpdir = tempfile.mkdtemp()
try:
path = os.path.join(tmpdir, "passthrough_info.py")
with open(path, "w") as f:
f.write(PRELOAD_COMMAND_TEXT)
with tmpfile() as fn2:
print(fn2)
with popen(
["dask-scheduler", "--scheduler-file", fn2, "--preload", path],
stdout=sys.stdout,
stderr=sys.stderr,
):
with Client(scheduler_file=fn2, loop=loop) as c:
assert c.run_on_scheduler(check_passthrough) == "default"
finally:
shutil.rmtree(tmpdir)
def test_bokeh(loop):
pytest.importorskip('bokeh')
with popen(['dask-scheduler']) as proc:
with Client('127.0.0.1:%d' % Scheduler.default_port, loop=loop) as c:
pass
names = ['localhost', '127.0.0.1', get_ip()]
if 'linux' in sys.platform:
names.append(socket.gethostname())
start = time()
while True:
try:
for name in names:
response = requests.get('http://%s:8787/status/' % name)
assert response.ok
break
except Exception as f:
print(f)
sleep(0.1)
assert time() < start + 10
with pytest.raises(Exception):
requests.get('http://127.0.0.1:8787/status/')
def test_defaults(loop):
with popen(['dask-scheduler', '--no-bokeh']) as proc:
with Executor('127.0.0.1:%d' % Scheduler.default_port, loop=loop) as e:
response = requests.get('http://127.0.0.1:9786/info.json')
assert response.ok
assert response.json()['status'] == 'running'
def test_nprocs_negative(loop):
with popen(["dask-scheduler", "--no-dashboard"]) as sched:
with popen(["dask-worker", "127.0.0.1:8786", "--nprocs=-1"]) as worker:
with Client("tcp://127.0.0.1:8786", loop=loop) as c:
c.wait_for_workers(cpu_count(), timeout="10 seconds")
def test_nprocs_requires_nanny(loop):
with popen(['dask-scheduler', '--no-bokeh']) as sched:
with popen(['dask-worker', '127.0.0.1:8786', '--nprocs=2',
'--no-nanny']) as worker:
assert any(b'Failed to launch worker' in worker.stderr.readline()
for i in range(15))
def test_no_nanny(loop):
with popen(['dask-scheduler', '--no-bokeh']) as sched:
with popen(['dask-worker', '127.0.0.1:8786', '--no-nanny',
'--no-bokeh']) as worker:
assert any(b'Registered' in worker.stderr.readline()
for i in range(15))
def test_hostport(loop):
with popen(['dask-scheduler', '--no-bokeh', '--host', '127.0.0.1:8978']):
with Client('127.0.0.1:8978', loop=loop) as c:
assert len(c.ncores()) == 0
def test_no_dashboard(loop):
pytest.importorskip("bokeh")
with popen(["dask-scheduler", "--no-dashboard"]) as proc:
with Client("127.0.0.1:%d" % Scheduler.default_port, loop=loop) as c:
response = requests.get("http://127.0.0.1:8787/status/")
assert response.status_code == 404
def test_hostport(loop):
with popen(['dask-scheduler', '--no-bokeh', '--host', '127.0.0.1:8978']):
with Executor('127.0.0.1:8978', loop=loop) as e:
assert len(e.ncores()) == 0