async def assert_balanced(inp, expected, c, s, *workers): steal = s.extensions["stealing"] steal._pc.stop() counter = itertools.count() tasks = list(concat(inp)) data_seq = itertools.count() futures = [] for w, ts in zip(workers, inp): for t in sorted(ts, reverse=True): if t: [dat] = await c.scatter([next(data_seq)], workers=w.address) ts = s.tasks[dat.key] # Ensure scheduler state stays consistent old_nbytes = ts.nbytes ts.nbytes = s.bandwidth * t for ws in ts.who_has: ws.nbytes += ts.nbytes - old_nbytes else: dat = 123 i = next(counter) f = c.submit( func, dat, key="%d-%d" % (int(t), i), workers=w.address, allow_other_workers=True, pure=False, priority=-i, ) futures.append(f) while len(s.rprocessing) < len(futures): await asyncio.sleep(0.001) for i in range(10): steal.balance() while steal.in_flight: await asyncio.sleep(0.001) result = [ sorted([int(key_split(k)) for k in s.processing[w.address]], reverse=True) for w in workers ] result2 = sorted(result, reverse=True) expected2 = sorted(expected, reverse=True) if config.get("pdb-on-err"): if result2 != expected2: import pdb pdb.set_trace() if result2 == expected2: return raise Exception("Expected: {}; got: {}".format(str(expected2), str(result2)))
def assert_balanced(inp, expected, c, s, *workers): steal = s.extensions['stealing'] steal._pc.stop() counter = itertools.count() tasks = list(concat(inp)) data_seq = itertools.count() futures = [] for w, ts in zip(workers, inp): for t in sorted(ts, reverse=True): if t: [dat] = yield c._scatter([next(data_seq)], workers=w.address) ts = s.tasks[dat.key] # Ensure scheduler state stays consistent old_nbytes = ts.nbytes ts.nbytes = BANDWIDTH * t for ws in ts.who_has: ws.nbytes += ts.nbytes - old_nbytes else: dat = 123 s.task_duration[str(int(t))] = 1 f = c.submit(func, dat, key='%d-%d' % (int(t), next(counter)), workers=w.address, allow_other_workers=True, pure=False) futures.append(f) while len(s.rprocessing) < len(futures): yield gen.sleep(0.001) for i in range(10): steal.balance() while steal.in_flight: yield gen.sleep(0.001) result = [ sorted([int(key_split(k)) for k in s.processing[w.address]], reverse=True) for w in workers ] result2 = sorted(result, reverse=True) expected2 = sorted(expected, reverse=True) if config.get('pdb-on-err'): if result2 != expected2: import pdb pdb.set_trace() if result2 == expected2: return raise Exception('Expected: {}; got: {}'.format(str(expected2), str(result2)))
def assert_balanced(inp, expected, c, s, *workers): steal = s.extensions['stealing'] steal._pc.stop() counter = itertools.count() tasks = list(concat(inp)) data_seq = itertools.count() futures = [] for w, ts in zip(workers, inp): for t in sorted(ts, reverse=True): if t: [dat] = yield c.scatter([next(data_seq)], workers=w.address) ts = s.tasks[dat.key] # Ensure scheduler state stays consistent old_nbytes = ts.nbytes ts.nbytes = BANDWIDTH * t for ws in ts.who_has: ws.nbytes += ts.nbytes - old_nbytes else: dat = 123 s.task_duration[str(int(t))] = 1 i = next(counter) f = c.submit(func, dat, key='%d-%d' % (int(t), i), workers=w.address, allow_other_workers=True, pure=False, priority=-i) futures.append(f) while len(s.rprocessing) < len(futures): yield gen.sleep(0.001) for i in range(10): steal.balance() while steal.in_flight: yield gen.sleep(0.001) result = [sorted([int(key_split(k)) for k in s.processing[w.address]], reverse=True) for w in workers] result2 = sorted(result, reverse=True) expected2 = sorted(expected, reverse=True) if config.get('pdb-on-err'): if result2 != expected2: import pdb pdb.set_trace() if result2 == expected2: return raise Exception('Expected: {}; got: {}'.format(str(expected2), str(result2)))
def assert_balanced(inp, expected, c, s, *workers): steal = s.extensions['stealing'] steal._pc.stop() counter = itertools.count() B = BANDWIDTH tasks = list(concat(inp)) data_seq = itertools.count() futures = [] for w, ts in zip(workers, inp): for t in ts: if t: [dat] = yield c._scatter([next(data_seq)], workers=w.address) s.nbytes[dat.key] = BANDWIDTH * t else: dat = 123 s.task_duration[str(int(t))] = 1 f = c.submit(func, dat, key='%d-%d' % (int(t), next(counter)), workers=w.address, allow_other_workers=True, pure=False) futures.append(f) while len(s.rprocessing) < len(futures): yield gen.sleep(0.001) s.extensions['stealing'].balance() result = [ sorted([int(key_split(k)) for k in s.processing[w.address]], reverse=True) for w in workers ] result2 = sorted(result, reverse=True) expected2 = sorted(expected, reverse=True) if config.get('pdb-on-err'): if result2 != expected2: import pdb pdb.set_trace() assert result2 == expected2