def test_get_local_alloc():
async def f():
ref = tsk.alloc(8)
np_v = np.frombuffer(ref.get_local(), dtype=np.int)
np_v[0] = 10
np_v2 = np.frombuffer(await ref.get(), dtype=np.int)
assert np_v2[0] == 10
tsk.zmq_run(f=f)
def test_delete_shm():
initial_dev_shm = os.listdir("/dev/shm")
cfg = tsk.Cfg(initializer=setup_worker)
r = tsk.zmq_run(cfg=cfg, f=f)
assert r == 99
final_dev_shm = os.listdir("/dev/shm")
new_files = set(final_dev_shm) - set(initial_dev_shm)
assert len(new_files) == 0
def main():
cfg = tsk.Cfg(n_workers=2,
log_level=logging.WARN,
initializer=setup_worker)
async def f():
nrows = int(1e7)
nnz = nrows * 10
n_repeats = 1
mat = random_test_matrix(nrows, nnz)
vec = np.random.rand(nrows) - 0.5
t = tsk.Timer()
for i in range(n_repeats):
correct = mat.dot(vec)
t.report("simple dot")
gang = await tsk.ctx().wait_for_workers(cfg.n_workers)
t.report("wait for workers")
t.report("launch profiler")
tsk_vec = TskArray(vals=vec)
t.report("shmem v")
tsk_mat = distribute(mat, gang)
t.report("distribute mat")
result = await tsk_mat.dot(tsk_vec)
t.report("first dot")
async with tsk.Profiler(gang):
t.restart()
for i in range(n_repeats):
result = await tsk_mat.dot(tsk_vec)
t.report("parallel dot")
print(np.sum(correct))
print(np.sum(result))
assert np.sum(result) == np.sum(correct)
tsk.zmq_run(cfg=cfg, f=f)
def run_tsk_parallel():
async def f():
gang = await tsk.ctx().wait_for_workers(2)
fnc_dref = tsk.put(long_fnc)
async with tsk.Profiler(gang):
start = time.time()
await wait_all([
tsk.task(fnc_dref, to=gang[i % len(gang)])
for i in range(n_jobs)
])
print("inside: ", time.time() - start)
def setup_worker(name):
tsk.cfg.stdout_logging(name, "taskloaf.profile")
cfg = tsk.Cfg(initializer=setup_worker)
return tsk.zmq_run(cfg=cfg, f=f)
import sys
import taskloaf as tsk
if __name__ == "__main__":
cfg = tsk.Cfg()
cfg.n_workers = int(sys.argv[1])
cfg.base_port = int(sys.argv[2])
if len(sys.argv) > 3:
connect_port = int(sys.argv[3])
cfg.connect_to = (cfg.hostname, connect_port)
tsk.zmq_run(cfg=cfg)
def test_get_local_put():
async def f():
ref = tsk.put(10)
assert ref.get_local() == 10
tsk.zmq_run(f=f)
import taskloaf as tsk
import logging
async def submit():
gang = await tsk.ctx().wait_for_workers(2)
X = 3.1
n = 10
pr = tsk.task(lambda: tsk.task(lambda: X))
for i in range(n):
pr = pr.then(lambda x: x + 1)
async def asum(x):
return sum(x)
pr = tsk.when_all([pr, tsk.task(lambda: X, to=gang[1])]).then(asum)
print("answer is", await pr)
if __name__ == "__main__":
cfg = tsk.Cfg(log_level=logging.WARN)
tsk.zmq_run(cfg=cfg, f=submit)
start = time.time()
while len(tsk.ctx().messenger.endpts.keys()) < 2:
await asyncio.sleep(0)
# print("HI")
gang = list(tsk.ctx().messenger.endpts.keys())
end = time.time()
print(end - start)
print(end - start)
print(end - start)
print(end - start)
print(gang)
print(gang)
print(gang)
print(gang)
print(gang)
print(gang)
print("LAUNCH")
await tsk.task(lambda: print("hI"), to=gang[0])
print("DONE")
return 123
if __name__ == "__main__":
cfg = tsk.Cfg()
cfg.n_workers = 1
cfg.connect_to = "tcp://127.0.0.1:5400"
r = tsk.zmq_run(cfg=cfg, f=f)
print(r)
assert r == 123
import taskloaf as tsk
N = 4
async def f():
gang = await tsk.ctx().wait_for_workers(N)
await tsk.task(lambda: print("hI"), to=gang[N - 2])
return 123
if __name__ == "__main__":
for i in range(2):
r = tsk.zmq_run(cfg=tsk.Cfg(n_workers=N), f=f)
print(r)
assert r == 123
