def f():
s = Scheduler()
s.listen(0)
x = Worker(s.ip, s.port, ip='127.0.0.1')
y = Worker(s.ip, s.port, ip='127.0.0.1')
z = Worker(s.ip, s.port, ip='127.0.0.1')
x.data['a'] = 1
y.data['a'] = 2
yield [x._start(), y._start(), z._start()]
zz = rpc(ip=z.ip, port=z.port)
yield zz.compute(function=dumps(inc),
args=dumps(('a',)),
who_has={'a': [x.address]},
key='b')
assert z.data['b'] == 2
if 'a' in z.data:
del z.data['a']
yield zz.compute(function=dumps(inc),
args=dumps(('a',)),
who_has={'a': [y.address]},
key='c')
assert z.data['c'] == 3
yield [x._close(), y._close(), z._close()]
zz.close_streams()
def test_worker_with_port_zero():
s = Scheduler()
s.start(8007)
w = Worker(s.address)
yield w._start()
assert isinstance(w.port, int)
assert w.port > 1024
yield w._close()
def test_scheduler_file():
with tmpfile() as fn:
s = Scheduler(scheduler_file=fn)
s.start(8009)
w = Worker(scheduler_file=fn)
yield w._start()
assert set(s.workers) == {w.address}
yield w._close()
s.stop()
def f():
c = Center(ip="127.0.0.1")
c.listen(0)
x = Worker(c.ip, c.port, ip="127.0.0.1")
y = Worker(c.ip, c.port, ip="127.0.0.1")
z = Worker(c.ip, c.port, ip="127.0.0.1")
x.data["a"] = 1
y.data["a"] = 2
yield [x._start(), y._start(), z._start()]
zz = rpc(ip=z.ip, port=z.port)
yield zz.compute(function=inc, args=("a",), needed=["a"], who_has={"a": {x.address}}, key="b")
assert z.data["b"] == 2
yield zz.compute(function=inc, args=("a",), needed=["a"], who_has={"a": {y.address}}, key="c")
assert z.data["c"] == 3
yield [x._close(), y._close(), z._close()]
zz.close_streams()
def g():
c = Center('127.0.0.1', 8017)
c.listen(c.port)
a = Worker('127.0.0.1', 8018, c.ip, c.port, ncores=1)
yield a._start()
b = Worker('127.0.0.1', 8019, c.ip, c.port, ncores=1)
yield b._start()
while len(c.ncores) < 2:
yield gen.sleep(0.01)
try:
yield f(c, a, b)
finally:
with ignoring(Exception):
yield a._close()
with ignoring(Exception):
yield b._close()
c.stop()
def test_service_hosts_match_worker(s):
pytest.importorskip('bokeh')
from distributed.bokeh.worker import BokehWorker
services = {('bokeh', 0): BokehWorker}
for host in ['tcp://0.0.0.0', 'tcp://127.0.0.2']:
w = Worker(s.address, services=services)
yield w._start(host)
sock = first(w.services['bokeh'].server._http._sockets.values())
assert sock.getsockname()[0] == host.split('://')[1]
yield w._close()
def f():
c = Center(ip='127.0.0.1')
c.listen(0)
x = Worker(c.ip, c.port, ip='127.0.0.1')
y = Worker(c.ip, c.port, ip='127.0.0.1')
z = Worker(c.ip, c.port, ip='127.0.0.1')
x.data['a'] = 1
y.data['a'] = 2
yield [x._start(), y._start(), z._start()]
zz = rpc(ip=z.ip, port=z.port)
yield zz.compute(function=inc, args=('a',),
who_has={'a': {x.address}}, key='b')
assert z.data['b'] == 2
yield zz.compute(function=inc, args=('a',),
who_has={'a': {y.address}}, key='c')
assert z.data['c'] == 3
yield [x._close(), y._close(), z._close()]
zz.close_streams()
def test_Executor(c, s):
with ThreadPoolExecutor(2) as e:
w = Worker(s.ip, s.port, executor=e)
assert w.executor is e
yield w._start()
future = c.submit(inc, 1)
result = yield future
assert result == 2
assert e._threads # had to do some work
yield w._close()
def test_Executor(c, s):
with ThreadPoolExecutor(2) as e:
w = Worker(s.ip, s.port, executor=e)
assert w.executor is e
yield w._start()
future = c.submit(inc, 1)
result = yield future
assert result == 2
assert e._threads # had to do some work
yield w._close()
def test_avoid_memory_monitor_if_zero_limit(c, s):
worker = Worker(s.address, loop=s.loop, memory_limit=0,
memory_monitor_interval=10)
yield worker._start()
assert type(worker.data) is dict
assert 'memory' not in worker.periodic_callbacks
future = c.submit(inc, 1)
assert (yield future) == 2
yield gen.sleep(worker.memory_monitor_interval / 1000)
yield c.submit(inc, 2) # worker doesn't pause
yield worker._close()
def test_service_hosts_match_worker(s):
pytest.importorskip("bokeh")
from distributed.bokeh.worker import BokehWorker
services = {("bokeh", ":0"): BokehWorker}
w = Worker(s.address, services={("bokeh", ":0"): BokehWorker})
yield w._start("tcp://0.0.0.0")
sock = first(w.services["bokeh"].server._http._sockets.values())
assert sock.getsockname()[0] in ("::", "0.0.0.0")
yield w._close()
w = Worker(s.address, services={("bokeh", ":0"): BokehWorker})
yield w._start("tcp://127.0.0.1")
sock = first(w.services["bokeh"].server._http._sockets.values())
assert sock.getsockname()[0] in ("::", "0.0.0.0")
yield w._close()
w = Worker(s.address, services={("bokeh", 0): BokehWorker})
yield w._start("tcp://127.0.0.1")
sock = first(w.services["bokeh"].server._http._sockets.values())
assert sock.getsockname()[0] == "127.0.0.1"
yield w._close()
def test_avoid_memory_monitor_if_zero_limit(c, s):
worker = Worker(s.address, loop=s.loop, memory_limit=0,
memory_monitor_interval=10)
yield worker._start()
assert type(worker.data) is dict
assert 'memory' not in worker.periodic_callbacks
future = c.submit(inc, 1)
assert (yield future) == 2
yield gen.sleep(worker.memory_monitor_interval / 1000)
yield c.submit(inc, 2) # worker doesn't pause
yield worker._close()
def test_worker_fds(s):
psutil = pytest.importorskip('psutil')
start = psutil.Process().num_fds()
worker = Worker(s.address, loop=s.loop)
yield worker._start()
middle = psutil.Process().num_fds()
assert middle > start
yield worker._close()
start = time()
while psutil.Process().num_fds() > start:
yield gen.sleep(0.01)
assert time() < start + 0.5
def test_worker_fds(s):
psutil = pytest.importorskip('psutil')
yield gen.sleep(0.05)
start = psutil.Process().num_fds()
worker = Worker(s.address, loop=s.loop)
yield worker._start()
yield gen.sleep(0.1)
middle = psutil.Process().num_fds()
start = time()
while middle > start:
yield gen.sleep(0.01)
assert time() < start + 1
yield worker._close()
start = time()
while psutil.Process().num_fds() > start:
yield gen.sleep(0.01)
assert time() < start + 0.5
def test_spill_to_disk(e, s):
np = pytest.importorskip('numpy')
w = Worker(s.ip, s.port, loop=s.loop, memory_limit=1000)
yield w._start()
x = e.submit(np.random.randint, 0, 255, size=500, dtype='u1', key='x')
yield _wait(x)
y = e.submit(np.random.randint, 0, 255, size=500, dtype='u1', key='y')
yield _wait(y)
assert set(w.data) == {x.key, y.key}
assert set(w.data.fast) == {x.key, y.key}
z = e.submit(np.random.randint, 0, 255, size=500, dtype='u1', key='z')
yield _wait(z)
assert set(w.data) == {x.key, y.key, z.key}
assert set(w.data.fast) == {y.key, z.key}
assert set(w.data.slow) == {x.key}
yield x._result()
assert set(w.data.fast) == {x.key, z.key}
assert set(w.data.slow) == {y.key}
yield w._close()
def test_spill_to_disk(e, s):
np = pytest.importorskip('numpy')
w = Worker(s.address, loop=s.loop, memory_limit=1000)
yield w._start()
x = e.submit(np.random.randint, 0, 255, size=500, dtype='u1', key='x')
yield _wait(x)
y = e.submit(np.random.randint, 0, 255, size=500, dtype='u1', key='y')
yield _wait(y)
assert set(w.data) == {x.key, y.key}
assert set(w.data.fast) == {x.key, y.key}
z = e.submit(np.random.randint, 0, 255, size=500, dtype='u1', key='z')
yield _wait(z)
assert set(w.data) == {x.key, y.key, z.key}
assert set(w.data.fast) == {y.key, z.key}
assert set(w.data.slow) == {x.key} or set(w.data.slow) == {x.key, y.key}
yield x._result()
assert set(w.data.fast) == {x.key, z.key}
assert set(w.data.slow) == {y.key} or set(w.data.slow) == {x.key, y.key}
yield w._close()
def test_spill_to_disk(c, s):
np = pytest.importorskip('numpy')
w = Worker(s.address, loop=s.loop, memory_limit=1200 / 0.6,
memory_pause_fraction=None, memory_spill_fraction=None)
yield w._start()
x = c.submit(np.random.randint, 0, 255, size=500, dtype='u1', key='x')
yield wait(x)
y = c.submit(np.random.randint, 0, 255, size=500, dtype='u1', key='y')
yield wait(y)
assert set(w.data) == {x.key, y.key}
assert set(w.data.fast) == {x.key, y.key}
z = c.submit(np.random.randint, 0, 255, size=500, dtype='u1', key='z')
yield wait(z)
assert set(w.data) == {x.key, y.key, z.key}
assert set(w.data.fast) == {y.key, z.key}
assert set(w.data.slow) == {x.key} or set(w.data.slow) == {x.key, y.key}
yield x
assert set(w.data.fast) == {x.key, z.key}
assert set(w.data.slow) == {y.key} or set(w.data.slow) == {x.key, y.key}
yield w._close()
def test_spill_to_disk(c, s):
np = pytest.importorskip('numpy')
w = Worker(s.address, loop=s.loop, memory_limit=1200 / 0.6,
memory_pause_fraction=None, memory_spill_fraction=None)
yield w._start()
x = c.submit(np.random.randint, 0, 255, size=500, dtype='u1', key='x')
yield wait(x)
y = c.submit(np.random.randint, 0, 255, size=500, dtype='u1', key='y')
yield wait(y)
assert set(w.data) == {x.key, y.key}
assert set(w.data.fast) == {x.key, y.key}
z = c.submit(np.random.randint, 0, 255, size=500, dtype='u1', key='z')
yield wait(z)
assert set(w.data) == {x.key, y.key, z.key}
assert set(w.data.fast) == {y.key, z.key}
assert set(w.data.slow) == {x.key} or set(w.data.slow) == {x.key, y.key}
yield x
assert set(w.data.fast) == {x.key, z.key}
assert set(w.data.slow) == {y.key} or set(w.data.slow) == {x.key, y.key}
yield w._close()
def test_scheduler_address_config(c, s):
with dask.config.set({'scheduler-address': s.address}):
worker = Worker(loop=s.loop)
yield worker._start()
assert worker.scheduler.address == s.address
yield worker._close()
def test_register_worker_callbacks(c, s, a, b):
#preload function to run
def mystartup(dask_worker):
dask_worker.init_variable = 1
def mystartup2():
import os
os.environ['MY_ENV_VALUE'] = 'WORKER_ENV_VALUE'
return "Env set."
#Check that preload function has been run
def test_import(dask_worker):
return hasattr(dask_worker, 'init_variable')
# and dask_worker.init_variable == 1
def test_startup2():
import os
return os.getenv('MY_ENV_VALUE', None) == 'WORKER_ENV_VALUE'
# Nothing has been run yet
assert len(s.worker_setups) == 0
result = yield c.run(test_import)
assert list(result.values()) == [False] * 2
result = yield c.run(test_startup2)
assert list(result.values()) == [False] * 2
# Start a worker and check that startup is not run
worker = Worker(s.address, loop=s.loop)
yield worker._start()
result = yield c.run(test_import, workers=[worker.address])
assert list(result.values()) == [False]
yield worker._close()
# Add a preload function
response = yield c.register_worker_callbacks(setup=mystartup)
assert len(response) == 2
assert len(s.worker_setups) == 1
# Check it has been ran on existing worker
result = yield c.run(test_import)
assert list(result.values()) == [True] * 2
# Start a worker and check it is ran on it
worker = Worker(s.address, loop=s.loop)
yield worker._start()
result = yield c.run(test_import, workers=[worker.address])
assert list(result.values()) == [True]
yield worker._close()
# Register another preload function
response = yield c.register_worker_callbacks(setup=mystartup2)
assert len(response) == 2
assert len(s.worker_setups) == 2
# Check it has been run
result = yield c.run(test_startup2)
assert list(result.values()) == [True] * 2
# Start a worker and check it is ran on it
worker = Worker(s.address, loop=s.loop)
yield worker._start()
result = yield c.run(test_import, workers=[worker.address])
assert list(result.values()) == [True]
result = yield c.run(test_startup2, workers=[worker.address])
assert list(result.values()) == [True]
yield worker._close()
# Final exception test
with pytest.raises(ZeroDivisionError):
yield c.register_worker_callbacks(setup=lambda: 1 / 0)
def test_scheduler_address_config(c, s):
with dask.config.set({'scheduler-address': s.address}):
worker = Worker(loop=s.loop)
yield worker._start()
assert worker.scheduler.address == s.address
yield worker._close()
def test_register_worker_callbacks(c, s, a, b):
#preload function to run
def mystartup(dask_worker):
dask_worker.init_variable = 1
def mystartup2():
import os
os.environ['MY_ENV_VALUE'] = 'WORKER_ENV_VALUE'
return "Env set."
#Check that preload function has been run
def test_import(dask_worker):
return hasattr(dask_worker, 'init_variable')
# and dask_worker.init_variable == 1
def test_startup2():
import os
return os.getenv('MY_ENV_VALUE', None) == 'WORKER_ENV_VALUE'
# Nothing has been run yet
assert len(s.worker_setups) == 0
result = yield c.run(test_import)
assert list(result.values()) == [False] * 2
result = yield c.run(test_startup2)
assert list(result.values()) == [False] * 2
# Start a worker and check that startup is not run
worker = Worker(s.address, loop=s.loop)
yield worker._start()
result = yield c.run(test_import, workers=[worker.address])
assert list(result.values()) == [False]
yield worker._close()
# Add a preload function
response = yield c.register_worker_callbacks(setup=mystartup)
assert len(response) == 2
assert len(s.worker_setups) == 1
# Check it has been ran on existing worker
result = yield c.run(test_import)
assert list(result.values()) == [True] * 2
# Start a worker and check it is ran on it
worker = Worker(s.address, loop=s.loop)
yield worker._start()
result = yield c.run(test_import, workers=[worker.address])
assert list(result.values()) == [True]
yield worker._close()
# Register another preload function
response = yield c.register_worker_callbacks(setup=mystartup2)
assert len(response) == 2
assert len(s.worker_setups) == 2
# Check it has been run
result = yield c.run(test_startup2)
assert list(result.values()) == [True] * 2
# Start a worker and check it is ran on it
worker = Worker(s.address, loop=s.loop)
yield worker._start()
result = yield c.run(test_import, workers=[worker.address])
assert list(result.values()) == [True]
result = yield c.run(test_startup2, workers=[worker.address])
assert list(result.values()) == [True]
yield worker._close()
# Final exception test
with pytest.raises(ZeroDivisionError):
yield c.register_worker_callbacks(setup=lambda: 1 / 0)
