def test_compute_stream(s, a, b): stream = yield connect(a.ip, a.port) yield write(stream, {'op': 'compute-stream'}) msgs = [{ 'op': 'compute-task', 'function': dumps(inc), 'args': dumps((i, )), 'key': 'x-%d' % i } for i in range(10)] bstream = BatchedStream(stream, 0) for msg in msgs[:5]: yield write(stream, msg) for i in range(5): msg = yield read(bstream) assert msg['status'] == 'OK' assert msg['key'][0] == 'x' for msg in msgs[5:]: yield write(stream, msg) for i in range(5): msg = yield read(bstream) assert msg['status'] == 'OK' assert msg['key'][0] == 'x' yield write(stream, {'op': 'close'})
def test_server(s, a, b): stream = yield connect('127.0.0.1', s.port) yield write(stream, {'op': 'register-client', 'client': 'ident'}) yield write( stream, { 'op': 'update-graph', 'tasks': { 'x': (inc, 1), 'y': (inc, 'x') }, 'dependencies': { 'x': set(), 'y': {'x'} }, 'keys': ['y'], 'client': 'ident' }) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'y': break yield write(stream, {'op': 'close-stream'}) msg = yield read(stream) assert msg == {'op': 'stream-closed'} assert stream.closed() stream.close()
def test_filtered_communication(s, a, b): e = yield connect(ip=s.ip, port=s.port) f = yield connect(ip=s.ip, port=s.port) yield write(e, {'op': 'register-client', 'client': 'e'}) yield write(f, {'op': 'register-client', 'client': 'f'}) yield read(e) yield read(f) e = BatchedStream(e, 0) f = BatchedStream(f, 0) assert set(s.streams) == {'e', 'f'} yield write(e, {'op': 'update-graph', 'tasks': {'x': dumps_task((inc, 1)), 'y': dumps_task((inc, 'x'))}, 'dependencies': {'x': [], 'y': ['x']}, 'client': 'e', 'keys': ['y']}) yield write(f, {'op': 'update-graph', 'tasks': {'x': dumps_task((inc, 1)), 'z': dumps_task((add, 'x', 10))}, 'dependencies': {'x': [], 'z': ['x']}, 'client': 'f', 'keys': ['z']}) msg = yield read(e) assert msg['op'] == 'key-in-memory' assert msg['key'] == 'y' msg = yield read(f) assert msg['op'] == 'key-in-memory' assert msg['key'] == 'z'
def test_filtered_communication(s, a, b): c = yield connect(ip=s.ip, port=s.port) f = yield connect(ip=s.ip, port=s.port) yield write(c, {'op': 'register-client', 'client': 'c'}) yield write(f, {'op': 'register-client', 'client': 'f'}) yield read(c) yield read(f) assert set(s.streams) == {'c', 'f'} yield write(c, {'op': 'update-graph', 'tasks': {'x': dumps_task((inc, 1)), 'y': dumps_task((inc, 'x'))}, 'dependencies': {'x': [], 'y': ['x']}, 'client': 'c', 'keys': ['y']}) yield write(f, {'op': 'update-graph', 'tasks': {'x': dumps_task((inc, 1)), 'z': dumps_task((add, 'x', 10))}, 'dependencies': {'x': [], 'z': ['x']}, 'client': 'f', 'keys': ['z']}) msg, = yield read(c) assert msg['op'] == 'key-in-memory' assert msg['key'] == 'y' msg, = yield read(f) assert msg['op'] == 'key-in-memory' assert msg['key'] == 'z'
def f(): server = Server({'ping': pingpong}) server.listen(8887) stream = yield connect('127.0.0.1', 8887) yield write(stream, {'op': 'ping'}) response = yield read(stream) assert response == b'pong' yield write(stream, {'op': 'ping', 'close': True}) response = yield read(stream) assert response == b'pong' server.stop()
def f(): server = Server({"ping": pingpong}) server.listen(8887) stream = yield connect("127.0.0.1", 8887) yield write(stream, {"op": "ping"}) response = yield read(stream) assert response == b"pong" yield write(stream, {"op": "ping", "close": True}) response = yield read(stream) assert response == b"pong" server.stop()
def test_server(s, a, b): stream = yield connect('127.0.0.1', s.port) yield write(stream, {'op': 'start-control'}) yield write(stream, {'op': 'update-graph', 'dsk': {'x': (inc, 1), 'y': (inc, 'x')}, 'keys': ['y']}) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'y': break yield write(stream, {'op': 'close-stream'}) msg = yield read(stream) assert msg == {'op': 'stream-closed'} assert stream.closed()
def test_feed_setup_teardown(s, a, b): def setup(scheduler): return 1 def func(scheduler, state): assert state == 1 return b'OK' def teardown(scheduler, state): scheduler.flag = 'done' stream = yield connect(s.ip, s.port) yield write( stream, { 'op': 'feed', 'function': func, 'setup': setup, 'teardown': teardown, 'interval': 0.01 }) for i in range(5): response = yield read(stream) assert response == b'OK' stream.close() start = time() while not hasattr(s, 'flag'): yield gen.sleep(0.01) assert time() - start < 5
def test_feed_setup_teardown(s, a, b): def setup(scheduler): return 1 def func(scheduler, state): assert state == 1 return 'OK' def teardown(scheduler, state): scheduler.flag = 'done' stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'feed', 'function': dumps(func), 'setup': dumps(setup), 'teardown': dumps(teardown), 'interval': 0.01}) for i in range(5): response = yield read(stream) assert response == 'OK' stream.close() start = time() while not hasattr(s, 'flag'): yield gen.sleep(0.01) assert time() - start < 5
def test_monitor_resources(): pytest.importorskip('psutil') c = Center(ip='127.0.0.1') c.listen(0) n = Nanny(c.ip, c.port, ncores=2, ip='127.0.0.1') yield n._start() nn = rpc(ip=n.ip, port=n.port) assert n.process.is_alive() d = n.resource_collect() assert {'cpu_percent', 'memory_percent'}.issubset(d) assert isinstance(d['timestamp'], datetime) stream = yield connect(ip=n.ip, port=n.port) yield write(stream, {'op': 'monitor_resources', 'interval': 0.01}) for i in range(3): msg = yield read(stream) assert isinstance(msg, dict) assert {'cpu_percent', 'memory_percent'}.issubset(msg) stream.close() yield n._close() c.stop()
def test_filtered_communication(s, a, b): e = yield connect(ip=s.ip, port=s.port) f = yield connect(ip=s.ip, port=s.port) yield write(e, {'op': 'register-client', 'client': 'e'}) yield write(f, {'op': 'register-client', 'client': 'f'}) yield read(e) yield read(f) assert set(s.streams) == {'e', 'f'} yield write( e, { 'op': 'update-graph', 'tasks': { 'x': (inc, 1), 'y': (inc, 'x') }, 'dependencies': { 'x': set(), 'y': {'x'} }, 'client': 'e', 'keys': ['y'] }) yield write( f, { 'op': 'update-graph', 'tasks': { 'x': (inc, 1), 'z': (add, 'x', 10) }, 'dependencies': { 'x': set(), 'z': {'x'} }, 'client': 'f', 'keys': ['z'] }) msg = yield read(e) assert msg['op'] == 'key-in-memory' assert msg['key'] == 'y' msg = yield read(f) assert msg['op'] == 'key-in-memory' assert msg['key'] == 'z'
def handle_stream(self, stream, address): while True: try: msg = yield read(stream) self.count += 1 yield write(stream, msg) except StreamClosedError as e: return
def f(): server = Server({'ping': pingpong}) with pytest.raises(OSError): server.port server.listen(8887) assert server.port == 8887 stream = yield connect('127.0.0.1', 8887) yield write(stream, {'op': 'ping'}) response = yield read(stream) assert response == b'pong' yield write(stream, {'op': 'ping', 'close': True}) response = yield read(stream) assert response == b'pong' server.stop()
def test_server(s, a, b): stream = yield connect('127.0.0.1', s.port) yield write(stream, {'op': 'register-client', 'client': 'ident'}) yield write(stream, {'op': 'update-graph', 'tasks': {'x': (inc, 1), 'y': (inc, 'x')}, 'dependencies': {'x': set(), 'y': {'x'}}, 'keys': ['y'], 'client': 'ident'}) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'y': break yield write(stream, {'op': 'close-stream'}) msg = yield read(stream) assert msg == {'op': 'stream-closed'} assert stream.closed() stream.close()
def test_server(s, a, b): stream = yield connect('127.0.0.1', s.port) yield write(stream, {'op': 'register-client', 'client': 'ident'}) yield write(stream, {'op': 'update-graph', 'tasks': {'x': dumps_task((inc, 1)), 'y': dumps_task((inc, 'x'))}, 'dependencies': {'x': [], 'y': ['x']}, 'keys': ['y'], 'client': 'ident'}) while True: msg = yield readone(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'y': break yield write(stream, {'op': 'close-stream'}) msg = yield readone(stream) assert msg == {'op': 'stream-closed'} with pytest.raises(StreamClosedError): yield readone(stream) close(stream)
def test_feed(s, a, b): def func(scheduler): return scheduler.processing, scheduler.stacks stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'feed', 'function': func, 'interval': 0.01}) for i in range(5): response = yield read(stream) expected = s.processing, s.stacks stream.close()
def processing(): with log_errors(): from distributed.diagnostics.scheduler import processing stream = yield connect(ip=options["host"], port=options["tcp-port"]) yield write(stream, {"op": "feed", "function": dumps(processing), "interval": 0.200}) while True: try: msg = yield read(stream) except StreamClosedError: break else: messages["processing"] = msg
def test_server(s, a, b): stream = yield connect('127.0.0.1', s.port) yield write(stream, {'op': 'start-control'}) yield write( stream, { 'op': 'update-graph', 'dsk': { 'x': (inc, 1), 'y': (inc, 'x') }, 'keys': ['y'] }) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'y': break yield write(stream, {'op': 'close-stream'}) msg = yield read(stream) assert msg == {'op': 'stream-closed'} assert stream.closed()
def processing(): with log_errors(): from distributed.diagnostics.scheduler import processing stream = yield connect(ip=options['host'], port=options['tcp-port']) yield write(stream, {'op': 'feed', 'function': dumps(processing), 'interval': 0.200}) while True: try: msg = yield read(stream) except StreamClosedError: break else: messages['processing'] = msg
def test_feed(s, a, b): def func(scheduler): return dumps((scheduler.processing, scheduler.stacks)) stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'feed', 'function': dumps(func), 'interval': 0.01}) for i in range(5): response = yield read(stream) expected = s.processing, s.stacks assert cloudpickle.loads(response) == expected close(stream)
def test_feed(s, a, b): def func(scheduler): return dumps((scheduler.processing, scheduler.stacks)) stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'feed', 'function': dumps(func), 'interval': 0.01}) for i in range(5): response = yield read(stream) expected = s.processing, s.stacks assert cloudpickle.loads(response) == expected stream.close()
def test_compute_stream(s, a, b): stream = yield connect(a.ip, a.port) yield write(stream, {'op': 'compute-stream'}) msgs = [{'op': 'compute-task', 'function': dumps(inc), 'args': dumps((i,)), 'key': 'x-%d' % i} for i in range(10)] bstream = BatchedStream(stream, 0) for msg in msgs[:5]: yield write(stream, msg) for i in range(5): msg = yield read(bstream) assert msg['status'] == 'OK' assert msg['key'][0] == 'x' for msg in msgs[5:]: yield write(stream, msg) for i in range(5): msg = yield read(bstream) assert msg['status'] == 'OK' assert msg['key'][0] == 'x' yield write(stream, {'op': 'close'})
def test_feed_large_bytestring(s, a, b): np = pytest.importorskip('numpy') x = np.ones(10000000) def func(scheduler): y = x return True stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'feed', 'function': dumps(func), 'interval': 0.01}) for i in range(5): response = yield read(stream) assert response == True close(stream)
def f(): nn = rpc(ip=n.ip, port=n.port) yield n._start() assert n.process.is_alive() d = n.resource_collect() assert {'cpu_percent', 'memory_percent'}.issubset(d) assert isinstance(d['timestamp'], datetime) stream = yield connect(ip=n.ip, port=n.port) yield write(stream, {'op': 'monitor_resources', 'interval': 0.01}) for i in range(3): msg = yield read(stream) assert isinstance(msg, dict) assert {'cpu_percent', 'memory_percent'}.issubset(msg) stream.close() yield n._close() c.stop()
def test_monitor_resources(s): pytest.importorskip('psutil') n = Nanny(s.ip, s.port, ncores=2, ip='127.0.0.1', loop=s.loop) yield n._start() assert isalive(n.process) d = n.resource_collect() assert {'cpu_percent', 'memory_percent'}.issubset(d) assert 'timestamp' in d stream = yield connect(ip=n.ip, port=n.port) yield write(stream, {'op': 'monitor_resources', 'interval': 0.01}) for i in range(3): msg = yield read(stream) assert isinstance(msg, dict) assert {'cpu_percent', 'memory_percent'}.issubset(msg) close(stream) yield n._close() s.stop()
def handle_stream(self, stream, address): while True: msg = yield read(stream) self.count += 1 yield write(stream, msg)
def test_scheduler(s, a, b): stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'register-client', 'client': 'ident'}) stream = BatchedStream(stream, 10) msg = yield read(stream) assert msg['op'] == 'stream-start' # Test update graph yield write(stream, {'op': 'update-graph', 'tasks': valmap(dumps_task, {'x': (inc, 1), 'y': (inc, 'x'), 'z': (inc, 'y')}), 'dependencies': {'x': [], 'y': ['x'], 'z': ['y']}, 'keys': ['x', 'z'], 'client': 'ident'}) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'z': break assert a.data.get('x') == 2 or b.data.get('x') == 2 # Test erring tasks yield write(stream, {'op': 'update-graph', 'tasks': valmap(dumps_task, {'a': (div, 1, 0), 'b': (inc, 'a')}), 'dependencies': {'a': [], 'b': ['a']}, 'keys': ['a', 'b'], 'client': 'ident'}) while True: msg = yield read(stream) if msg['op'] == 'task-erred' and msg['key'] == 'b': break # Test missing data yield write(stream, {'op': 'missing-data', 'keys': ['z']}) s.ensure_idle_ready() while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'z': break # Test missing data without being informed for w in [a, b]: if 'z' in w.data: del w.data['z'] yield write(stream, {'op': 'update-graph', 'tasks': {'zz': dumps_task((inc, 'z'))}, 'dependencies': {'zz': ['z']}, 'keys': ['zz'], 'client': 'ident'}) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'zz': break write(stream, {'op': 'close'}) stream.close()
def test_scheduler(s, a, b): stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'register-client', 'client': 'ident'}) msg = yield read(stream) assert msg['op'] == 'stream-start' # Test update graph yield write( stream, { 'op': 'update-graph', 'tasks': { 'x': (inc, 1), 'y': (inc, 'x'), 'z': (inc, 'y') }, 'dependencies': { 'x': set(), 'y': {'x'}, 'z': {'y'} }, 'keys': ['x', 'z'], 'client': 'ident' }) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'z': break assert a.data.get('x') == 2 or b.data.get('x') == 2 # Test erring tasks yield write( stream, { 'op': 'update-graph', 'tasks': { 'a': (div, 1, 0), 'b': (inc, 'a') }, 'dependencies': { 'a': set(), 'b': {'a'} }, 'keys': ['a', 'b'], 'client': 'ident' }) while True: msg = yield read(stream) if msg['op'] == 'task-erred' and msg['key'] == 'b': break # Test missing data yield write(stream, {'op': 'missing-data', 'missing': ['z']}) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'z': break # Test missing data without being informed for w in [a, b]: if 'z' in w.data: del w.data['z'] yield write( stream, { 'op': 'update-graph', 'tasks': { 'zz': (inc, 'z') }, 'dependencies': { 'zz': {'z'} }, 'keys': ['zz'], 'client': 'ident' }) while True: msg = yield read(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'zz': break write(stream, {'op': 'close'}) stream.close()
def disconnect(ip, port): stream = yield connect(ip, port) yield write(stream, {'op': 'terminate', 'close': True}) response = yield read(stream) stream.close()