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' close(stream) start = time() while not hasattr(s, 'flag'): yield gen.sleep(0.01) assert time() - start < 5
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) 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 assert cloudpickle.loads(response) == expected close(stream)
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 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_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 test_scheduler(s, a, b): stream = yield connect(s.ip, s.port) yield write(stream, {'op': 'register-client', 'client': 'ident'}) msg = yield readone(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 readone(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 readone(stream) if msg['op'] == 'task-erred' and msg['key'] == 'b': break # Test missing data yield write(stream, {'op': 'missing-data', 'keys': ['z']}) s.ensure_occupied() while True: msg = yield readone(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 readone(stream) if msg['op'] == 'key-in-memory' and msg['key'] == 'zz': break write(stream, {'op': 'close'}) close(stream)