def test_add_remove_worker(s): events = [] class MyPlugin(SchedulerPlugin): def add_worker(self, worker, scheduler): assert scheduler is s events.append(('add_worker', worker)) def remove_worker(self, worker, scheduler): assert scheduler is s events.append(('remove_worker', worker)) plugin = MyPlugin() s.add_plugin(plugin) assert events == [] a = Worker(s.address) b = Worker(s.address) yield a._start() yield b._start() yield a._close() yield b._close() assert events == [('add_worker', a.address), ('add_worker', b.address), ('remove_worker', a.address), ('remove_worker', b.address), ] events[:] = [] s.remove_plugin(plugin) a = Worker(s.address) yield a._start() yield a._close() assert events == []
def test_worker_name(): s = Scheduler(validate=True) s.start(0) w = Worker(s.ip, s.port, name='alice') yield w._start() assert s.workers[w.address].name == 'alice' assert s.aliases['alice'] == w.address with pytest.raises(ValueError): w2 = Worker(s.ip, s.port, name='alice') yield w2._start() yield w2._close() yield s.close() yield w._close()
def test_no_worker_to_memory_restrictions(c, s, a, b): x = delayed(slowinc)(1, delay=0.4) y = delayed(slowinc)(x, delay=0.4) z = delayed(slowinc)(y, delay=0.4) yy, zz = c.persist([y, z], workers={(x, y, z): 'alice'}) while not s.tasks: yield gen.sleep(0.01) w = Worker(s.ip, s.port, ncores=1, name='alice') w.put_key_in_memory(y.key, 3) yield w._start() while len(s.workers) < 3: yield gen.sleep(0.01) yield gen.sleep(0.3) assert s.get_task_status(keys={x.key, y.key, z.key}) == { x.key: 'released', y.key: 'memory', z.key: 'processing', } yield w._close()
def test_no_workers_to_memory(c, s): x = delayed(slowinc)(1, delay=0.4) y = delayed(slowinc)(x, delay=0.4) z = delayed(slowinc)(y, delay=0.4) yy, zz = c.persist([y, z]) while not s.tasks: yield gen.sleep(0.01) w = Worker(s.ip, s.port, ncores=1) w.put_key_in_memory(y.key, 3) yield w._start() start = time() while not s.workers: yield gen.sleep(0.01) assert s.get_task_status(keys={x.key, y.key, z.key}) == { x.key: 'released', y.key: 'memory', z.key: 'processing', } yield w._close()
def test_worker_arrives_with_processing_data(c, s, a, b): x = delayed(slowinc)(1, delay=0.4) y = delayed(slowinc)(x, delay=0.4) z = delayed(slowinc)(y, delay=0.4) yy, zz = c.persist([y, z]) while not any(w.processing for w in s.workers.values()): yield gen.sleep(0.01) w = Worker(s.ip, s.port, ncores=1) w.put_key_in_memory(y.key, 3) yield w._start() start = time() while len(s.workers) < 3: yield gen.sleep(0.01) assert s.get_task_status(keys={x.key, y.key, z.key}) == { x.key: 'released', y.key: 'memory', z.key: 'processing', } yield w._close()
def f(): s = Scheduler(loop=loop) done = s.start(0) a = Worker(s.ip, s.port, loop=loop, ncores=1) b = Worker(s.ip, s.port, loop=loop, ncores=1) yield [a._start(0), b._start(0)] progress = TextProgressBar([], scheduler=(s.ip, s.port), start=False, interval=0.01) yield progress.listen() assert progress.status == 'finished' check_bar_completed(capsys) yield [a._close(), b._close()] s.close() yield done
def test_services(s, a, b): c = Worker(s.ip, s.port, ncores=1, ip='127.0.0.1', services={'http': HTTPWorker}) yield c._start() assert isinstance(c.services['http'], HTTPServer) assert c.service_ports['http'] == c.services['http'].port assert s.worker_info[c.address]['services']['http'] == c.service_ports['http'] yield c._close()
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=2) 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(): yield a._close() with ignoring(): yield b._close() c.stop()
def f(): s = Scheduler(loop=loop) done = s.start(0) a = Worker(s.ip, s.port, loop=loop, ncores=1) b = Worker(s.ip, s.port, loop=loop, ncores=1) yield [a._start(0), b._start(0)] progress = TextProgressBar([], scheduler=(s.ip, s.port), start=False, interval=0.01) yield progress.listen() assert progress.status == 'finished' check_bar_completed(capsys) yield [a._close(), b._close()] s.close() yield done
def g(): c = Center('127.0.0.1', 8017) c.listen(c.port) a = Worker('127.0.0.2', 8018, c.ip, c.port, ncores=2) yield a._start() b = Worker('127.0.0.3', 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(): yield a._close() with ignoring(): yield b._close() c.stop()
def test_scatter_no_workers(s): with pytest.raises(gen.TimeoutError): yield gen.with_timeout(timedelta(seconds=0.1), s.scatter(data={'x': dumps(1)}, client='alice')) w = Worker(s.ip, s.port, ncores=3, ip='127.0.0.1') yield [s.scatter(data={'x': dumps(1)}, client='alice'), w._start()] assert w.data['x'] == 1 yield w._close()
def test_scatter_no_workers(c, s): with pytest.raises(gen.TimeoutError): yield gen.with_timeout(timedelta(seconds=0.1), s.scatter(data={'x': 1}, client='alice')) w = Worker(s.ip, s.port, ncores=3, ip='127.0.0.1') yield [c._scatter(data={'x': 1}), w._start()] assert w.data['x'] == 1 yield w._close()
def test_resources(c, s): assert not s.worker_resources assert not s.resources a = Worker(s.ip, s.port, loop=s.loop, resources={'GPU': 2}) b = Worker(s.ip, s.port, loop=s.loop, resources={'GPU': 1, 'DB': 1}) yield [a._start(), b._start()] assert s.resources == {'GPU': {a.address: 2, b.address: 1}, 'DB': {b.address: 1}} assert s.worker_resources == {a.address: {'GPU': 2}, b.address: {'GPU': 1, 'DB': 1}} yield b._close() assert s.resources == {'GPU': {a.address: 2}, 'DB': {}} assert s.worker_resources == {a.address: {'GPU': 2}} yield a._close()
def test_resources(c, s): assert not s.worker_resources assert not s.resources a = Worker(s.ip, s.port, loop=s.loop, resources={'GPU': 2}) b = Worker(s.ip, s.port, loop=s.loop, resources={'GPU': 1, 'DB': 1}) yield [a._start(), b._start()] assert s.resources == {'GPU': {a.address: 2, b.address: 1}, 'DB': {b.address: 1}} assert s.worker_resources == {a.address: {'GPU': 2}, b.address: {'GPU': 1, 'DB': 1}} yield b._close() assert s.resources == {'GPU': {a.address: 2}, 'DB': {}} assert s.worker_resources == {a.address: {'GPU': 2}} yield a._close()
def test_services(s, a, b): c = Worker(s.ip, s.port, ncores=1, ip='127.0.0.1', services={'http': HTTPWorker}) yield c._start() assert isinstance(c.services['http'], HTTPServer) assert c.service_ports['http'] == c.services['http'].port assert s.worker_services[c.address]['http'] == c.service_ports['http'] yield c._close()
def test_services_port(s, a, b): c = Worker(s.ip, s.port, ncores=1, ip='127.0.0.1', services={('http', 9898): HTTPWorker}) yield c._start() assert isinstance(c.services['http'], HTTPServer) assert (c.service_ports['http'] == c.services['http'].port == s.worker_info[c.address]['services']['http'] == 9898) c.services['http'].stop() yield c._close()
def test_services_port(s, a, b): c = Worker(s.ip, s.port, ncores=1, ip='127.0.0.1', services={('http', 9898): HTTPWorker}) yield c._start() assert isinstance(c.services['http'], HTTPServer) assert (c.service_ports['http'] == c.services['http'].port == s.worker_info[c.address]['services']['http'] == 9898) c.services['http'].stop() yield c._close()
def g(): c = Center("127.0.0.1", 8017) c.listen(c.port) a = Worker("127.0.0.2", 8018, c.ip, c.port, ncores=2) yield a._start() b = Worker("127.0.0.3", 8019, c.ip, c.port, ncores=1) yield b._start() start = time() try: while len(c.ncores) < 2: yield gen.sleep(0.01) if time() - start > 5: raise Exception("Cluster creation timeout") yield f(c, a, b) finally: logger.debug("Closing out test cluster") with ignoring(): yield a._close() with ignoring(): yield b._close() c.stop()
def test_worker_name(): s = Scheduler(validate=True) s.start(0) w = Worker(s.ip, s.port, name='alice') yield w._start() assert s.worker_info[w.address]['name'] == 'alice' assert s.aliases['alice'] == w.address with pytest.raises(ValueError): w = Worker(s.ip, s.port, name='alice') yield w._start() yield s.close() yield w._close()
def test_scatter_no_workers(c, s): with pytest.raises(gen.TimeoutError): yield s.scatter(data={"x": 1}, client="alice", timeout=0.1) start = time() with pytest.raises(gen.TimeoutError): yield c.scatter(123, timeout=0.1) assert time() < start + 1.5 w = Worker(s.ip, s.port, ncores=3) yield [c.scatter(data={"y": 2}, timeout=5), w._start()] assert w.data["y"] == 2 yield w._close()
def test_worker_name(): s = Scheduler() s.start(0) w = Worker(s.ip, s.port, name='alice') yield w._start() assert s.worker_info[w.address]['name'] == 'alice' assert s.aliases['alice'] == w.address with pytest.raises(ValueError): w = Worker(s.ip, s.port, name='alice') yield w._start() yield s.close() yield w._close()
def test_scatter_no_workers(c, s): with pytest.raises(gen.TimeoutError): yield s.scatter(data={'x': 1}, client='alice', timeout=0.1) start = time() with pytest.raises(gen.TimeoutError): yield c.scatter(123, timeout=0.1) assert time() < start + 1.5 w = Worker(s.ip, s.port, ncores=3) yield [c.scatter(data={'y': 2}, timeout=5), w._start()] assert w.data['y'] == 2 yield w._close()
def test_file_descriptors_dont_leak(s): psutil = pytest.importorskip('psutil') proc = psutil.Process() before = proc.num_fds() w = Worker(s.ip, s.port) yield w._start(0) yield w._close() during = proc.num_fds() start = time() while proc.num_fds() > before: yield gen.sleep(0.01) assert time() < start + 5
def test_file_descriptors_dont_leak(s): psutil = pytest.importorskip('psutil') proc = psutil.Process() before = proc.num_fds() w = Worker(s.ip, s.port) yield w._start(0) yield w._close() during = proc.num_fds() start = time() while proc.num_fds() > before: yield gen.sleep(0.01) assert time() < start + 5
def test_add_worker(s, a, b): w = Worker(s.ip, s.port, ncores=3) w.data['x-5'] = 6 w.data['y'] = 1 yield w._start(0) dsk = {('x-%d' % i): (inc, i) for i in range(10)} s.update_graph(tasks=valmap(dumps_task, dsk), keys=list(dsk), client='client', dependencies={k: set() for k in dsk}) s.add_worker(address=w.address, keys=list(w.data), ncores=w.ncores, services=s.services) s.validate_state() assert w.ip in s.host_info assert s.host_info[w.ip]['addresses'] == {a.address, b.address, w.address} yield w._close()
def test_add_worker(s, a, b): w = Worker(s.ip, s.port, ncores=3, ip='127.0.0.1') w.data['x-5'] = 6 w.data['y'] = 1 yield w._start(0) dsk = {('x-%d' % i).encode(): (inc, i) for i in range(10)} s.update_graph(tasks=valmap(dumps_task, dsk), keys=list(dsk), client='client', dependencies={k: set() for k in dsk}) s.add_worker(address=w.address, keys=list(w.data), ncores=w.ncores, services=s.services, coerce_address=False) s.validate_state() assert w.ip in s.host_info assert s.host_info[w.ip]['ports'] == set(map(str, [a.port, b.port, w.port])) yield w._close()
def test_add_worker(s, a, b): w = Worker(s.ip, s.port, ncores=3, ip='127.0.0.1') w.data['x-5'] = 6 w.data['y'] = 1 yield w._start(0) dsk = {('x-%d' % i): (inc, i) for i in range(10)} s.update_graph(tasks=valmap(dumps_task, dsk), keys=list(dsk), client='client', dependencies={k: set() for k in dsk}) s.add_worker(address=w.address, keys=list(w.data), ncores=w.ncores, services=s.services, coerce_address=False) s.validate_state() assert w.ip in s.host_info assert s.host_info[w.ip]['ports'] == set(map(str, [a.port, b.port, w.port])) yield w._close()
def test_resource_submit(c, s, a, b): x = c.submit(inc, 1, resources={'A': 3}) y = c.submit(inc, 2, resources={'B': 1}) z = c.submit(inc, 3, resources={'C': 2}) yield wait(x) assert x.key in a.data yield wait(y) assert y.key in b.data assert s.get_task_status(keys=[z.key]) == {z.key: 'no-worker'} d = Worker(s.ip, s.port, loop=s.loop, resources={'C': 10}) yield d._start() yield wait(z) assert z.key in d.data yield d._close()
def test_resource_submit(c, s, a, b): x = c.submit(inc, 1, resources={'A': 3}) y = c.submit(inc, 2, resources={'B': 1}) z = c.submit(inc, 3, resources={'C': 2}) yield wait(x) assert x.key in a.data yield wait(y) assert y.key in b.data assert s.get_task_status(keys=[z.key]) == {z.key: 'no-worker'} d = Worker(s.ip, s.port, loop=s.loop, resources={'C': 10}) yield d._start() yield wait(z) assert z.key in d.data yield d._close()
def test_resource_submit(c, s, a, b): x = c.submit(inc, 1, resources={'A': 3}) y = c.submit(inc, 2, resources={'B': 1}) z = c.submit(inc, 3, resources={'C': 2}) yield _wait(x) assert x.key in a.data yield _wait(y) assert y.key in b.data assert z.key in s.unrunnable d = Worker(s.ip, s.port, loop=s.loop, resources={'C': 10}) yield d._start() yield _wait(z) assert z.key in d.data yield d._close()
def test_new_worker_steals(c, s, a): yield wait(c.submit(slowinc, 1, delay=0.01)) futures = c.map(slowinc, range(100), delay=0.05) total = c.submit(sum, futures) while len(a.task_state) < 10: yield gen.sleep(0.01) b = Worker(s.ip, s.port, loop=s.loop, ncores=1, memory_limit=TOTAL_MEMORY) yield b._start() result = yield total assert result == sum(map(inc, range(100))) for w in [a, b]: assert all(isinstance(v, int) for v in w.data.values()) assert b.data yield b._close()
def test_resource_submit(c, s, a, b): x = c.submit(inc, 1, resources={'A': 3}) y = c.submit(inc, 2, resources={'B': 1}) z = c.submit(inc, 3, resources={'C': 2}) yield _wait(x) assert x.key in a.data yield _wait(y) assert y.key in b.data assert z.key in s.unrunnable d = Worker(s.ip, s.port, loop=s.loop, resources={'C': 10}) yield d._start() yield _wait(z) assert z.key in d.data yield d._close()
def test_new_worker_steals(c, s, a): yield _wait(c.submit(slowinc, 1, delay=0.01)._result()) futures = c.map(slowinc, range(100), delay=0.05) total = c.submit(sum, futures) while len(a.task_state) < 10: yield gen.sleep(0.01) b = Worker(s.ip, s.port, loop=s.loop, ncores=1) yield b._start() result = yield total._result() assert result == sum(map(inc, range(100))) for w in [a, b]: assert all(isinstance(v, int) for v in w.data.values()) assert b.data yield b._close()
def test_steal_resource_restrictions(c, s, a): future = c.submit(slowinc, 1, delay=0.10, workers=a.address) yield future._result() futures = c.map(slowinc, range(100), delay=0.2, resources={'A': 1}) while len(a.task_state) < 101: yield gen.sleep(0.01) assert len(a.task_state) == 101 b = Worker(s.ip, s.port, loop=s.loop, ncores=1, resources={'A': 4}) yield b._start() start = time() while not b.task_state or len(a.task_state) == 101: yield gen.sleep(0.01) assert time() < start + 3 assert len(b.task_state) > 0 assert len(a.task_state) < 101 yield b._close()
def test_steal_resource_restrictions(c, s, a): future = c.submit(slowinc, 1, delay=0.10, workers=a.address) yield future futures = c.map(slowinc, range(100), delay=0.2, resources={'A': 1}) while len(a.task_state) < 101: yield gen.sleep(0.01) assert len(a.task_state) == 101 b = Worker(s.ip, s.port, loop=s.loop, ncores=1, resources={'A': 4}) yield b._start() start = time() while not b.task_state or len(a.task_state) == 101: yield gen.sleep(0.01) assert time() < start + 3 assert len(b.task_state) > 0 assert len(a.task_state) < 101 yield b._close()
def test_errors_dont_block(): c = Center('127.0.0.1') c.listen(0) w = Worker(c.ip, c.port, ncores=1, ip='127.0.0.1') e = Executor((c.ip, c.port), start=False, loop=IOLoop.current()) yield w._start() yield e._start() L = [e.submit(inc, 1), e.submit(throws, 1), e.submit(inc, 2), e.submit(throws, 2)] start = time() while not (L[0].status == L[2].status == 'finished'): assert time() < start + 5 yield gen.sleep(0.01) result = yield e._gather([L[0], L[2]]) assert result == [2, 3] yield w._close() c.stop()
def test_no_worker_to_memory_restrictions(c, s, a, b): x = delayed(slowinc)(1, delay=0.4) y = delayed(slowinc)(x, delay=0.4) z = delayed(slowinc)(y, delay=0.4) yy, zz = c.persist([y, z], workers={(x, y, z): 'alice'}) while not s.task_state: yield gen.sleep(0.01) w = Worker(s.ip, s.port, ncores=1, name='alice') w.put_key_in_memory(y.key, 3) yield w._start() while len(s.workers) < 3: yield gen.sleep(0.01) yield gen.sleep(0.3) assert s.task_state[y.key] == 'memory' assert s.task_state[x.key] == 'released' assert s.task_state[z.key] == 'processing' yield w._close()
def test_errors_dont_block(): c = Center('127.0.0.1') c.listen(0) w = Worker(c.ip, c.port, ncores=1, ip='127.0.0.1') e = Executor((c.ip, c.port), start=False, loop=IOLoop.current()) yield w._start() yield e._start() L = [e.submit(inc, 1), e.submit(throws, 1), e.submit(inc, 2), e.submit(throws, 2)] start = time() while not (L[0].status == L[2].status == 'finished'): assert time() < start + 5 yield gen.sleep(0.01) result = yield e._gather([L[0], L[2]]) assert result == [2, 3] yield w._close() c.stop()
def test_scheduler_sees_memory_limits(s): w = Worker(s.ip, s.port, ncores=3, memory_limit=12345) yield w._start(0) assert s.workers[w.address].memory_limit == 12345 yield w._close()
def test_scheduler_sees_memory_limits(s): w = Worker(s.ip, s.port, ncores=3, memory_limit=12345) yield w._start(0) assert s.worker_info[w.address]['memory_limit'] == 12345 yield w._close()
def test_scheduler_sees_memory_limits(s): w = Worker(s.ip, s.port, ncores=3, ip='127.0.0.1', memory_limit=12345) yield w._start(0) assert s.worker_info[w.address]['memory_limit'] == 12345 yield w._close()