def test_get():
one_serialized = pickle.dumps(1)
async def f(w):
ref = put(w, 1)
ref2 = alloc(w, len(one_serialized))
buf = await ref2.get()
buf[:] = one_serialized
assert (await ref.get() == 1)
assert (await ref2.get() == one_serialized)
async def g(w):
assert (await ref.get() == 1)
assert (ref.key() in w.object_cache)
assert (await ref2.get() == one_serialized)
assert (ref2.key() in w.object_cache)
v1 = asyncio.ensure_future(ref._remote_get())
await asyncio.sleep(0)
assert (isinstance(w.object_cache[ref.key()], asyncio.Future))
assert (await ref.get() == 1)
assert (isinstance(w.object_cache[ref.key()], int))
assert (await v1 == 1)
submit_ref_work(w, 0, taskloaf.worker.shutdown)
submit_ref_work(w, 1, g)
await wait_to_die()
cluster(2, f)
def test_auto_unwrap():
async def f(w):
y = await task(w, lambda w: 10,
to=1).then(lambda w, x: task(w, lambda w: 2 * x), to=0)
assert (y == 20)
cluster(2, f)
def test_when_all():
async def f(w):
y = await when_all([task(w, lambda w: 10),
task(w, lambda w: 5,
to=1)]).then(lambda w, x: sum(x), to=1)
assert (y == 15)
cluster(2, f)
def test_task_exception():
async def f(w):
def g(w):
raise Implosion()
with pytest.raises(Implosion):
await task(w, g)
with pytest.raises(Implosion):
await task(w, g, to=1)
cluster(2, f)
def test_submit_ref_work():
async def f(w):
ref = put(w, 1)
assert (ref._id == 0)
async def g(w):
assert (ref.ref._id == 0)
submit_ref_work(w, 0, taskloaf.worker.shutdown)
submit_ref_work(w, 1, g)
await wait_to_die()
cluster(2, f)
def test_task_await_elsewhere():
async def f(w):
def g(w):
return 71
pr = task(w, g, to=1)
async def h(w):
a = await pr
return a + 1
pr2 = task(w, h, to=2)
assert (await pr2 == 72)
cluster(3, f)
def test_multiuse_msgs():
"""
Here, ref is deserialized multiple times on different workers.
But, ref is only serialized (and the child count incremented)
once on the main worker.
So, the ref count will be 2 while there will be 3 live references.
"""
async def f(w):
async def fnc():
v = await fnc.ref.get()
assert (v == 1)
del fnc.ref
fnc.ref = put(w, 1)
w.finished = False
async def g(w):
fnc = await g.ref_fnc.get()
await fnc()
del fnc, g.ref_fnc
def h(w):
w.finished = True
w.submit_work(0, h)
w.object_cache.clear()
import gc
gc.collect()
g.ref_fnc = put(w, fnc)
submit_ref_work(w, 1, g)
submit_ref_work(w, 2, g)
w.submit_work(1, taskloaf.worker.shutdown)
w.submit_work(2, taskloaf.worker.shutdown)
del fnc.ref, g.ref_fnc
while not w.finished:
await asyncio.sleep(0)
for i in range(200):
if not w.allocator.empty():
await asyncio.sleep(0.01)
continue
assert (w.allocator.empty())
cluster(3, f)
def test_task():
async def f(w):
w.abc = 0
def here(w):
w.abc = 1
await task(w, here)
assert (w.abc == 1)
def here_args(w, x):
w.abc = x
await task(w, here_args, 17)
assert (w.abc == 17)
def there(w):
assert (w.addr == 1)
return 13
assert (await task(w, there, to=1) == 13)
thirteen = taskloaf.serialize.dumps(w, 13)
def there2(w):
assert (w.addr == 1)
return thirteen
assert (await task(w, there2, to=1) == thirteen)
there_ref = put(w, there)
assert (await task(w, there_ref, to=1) == 13)
def f_with_args(w, a):
return a
f_ref = put(w, f_with_args)
assert (await task(w, f_ref, 14, to=1) == 14)
a_ref = put(w, 15)
assert (await task(w, f_ref, a_ref, to=1) == 15)
cluster(2, f)
def test_remote_work():
async def f(w):
if w.addr != 0:
return
async def g(w):
assert (w.addr == 1)
def h(w):
assert (w.addr == 0)
taskloaf.worker.shutdown(w)
w.submit_work(0, h)
taskloaf.worker.shutdown(w)
w.submit_work(1, g)
while True:
await asyncio.sleep(0)
cluster(2, f, runner=mpiexisting)
def test_cluster_death_cleansup():
def check(n, onoff):
for i in range(n):
path = '/dev/shm/taskloaf_' + str(i) + '_0'
assert (os.path.exists(path) == onoff)
async def f(w):
ptr = w.allocator.malloc(1)
check(1, True)
raise FunnyException()
check(2, False)
raises = False
try:
cluster(2, f, runner=mpiexisting)
except FunnyException as e:
raises = True
assert (rank() == 1 or raises)
import gc
gc.collect()
check(2, False)
def test_remote_double_get():
async def f(w):
ref = put(w, 1)
async def g(w):
assert (await ref.get() == 1)
v1 = ref._remote_get()
assert (await v1 == 1)
def h(w):
if not hasattr(w, 'x'):
w.x = 0
else:
taskloaf.worker.shutdown(w)
submit_ref_work(w, 0, h)
submit_ref_work(w, 1, g)
submit_ref_work(w, 1, g)
await wait_to_die()
cluster(2, f)
def test_exception():
async def f(w):
def count_except(w):
if not hasattr(w, 'testcounter'):
w.testcounter = 0
w.testcounter += 1
if w.testcounter % 2:
raise Exception("failed task!")
if w.testcounter > 10:
#TODO: Look up condition variables in python and asyncio. Could
# be handy for doing automatic wakeup of threads. asyncio.Future
# works, but is it the best option?
def success(w):
w.wait_fut.set_result(None)
w.submit_work(0, success)
for i in range(12):
w.submit_work(1, count_except)
w.wait_fut = asyncio.Future()
await w.wait_fut
cluster(2, f)
def test_then():
async def f(w):
y = await task(w, lambda w: 10, to=1).then(lambda w, x: 2 * x, to=0)
assert (y == 20)
cluster(2, f)
def test_cluster_output():
async def f(w):
return 1
assert (cluster(1, f) == 1)
