def test_chain_concat(test_env): test_env.start(1) with test_env.client.new_session() as s: t1 = tasks.concat((blob("a"), blob("b"))) t2 = tasks.concat((t1, blob("c"))) t3 = tasks.concat((t2, blob("d"))) t4 = tasks.concat((t3, blob("e"))) t5 = tasks.concat((t4, blob("f"))) t5.output.keep() s.submit() assert t5.output.fetch().get_bytes() == b"abcdef"
def run_small_gridcat(session): BLOB_SIZE = 5000 BLOB_COUNT = 10 rnd = test_rnd() def random_string(rnd, length): return "".join(rnd.choice(CHARS) for i in range(length)) cat = Program(("cat", Input("input1"), Input("input2")), stdout="output") md5sum = Program("md5sum", stdin="input", stdout="output") @remote() def take_first(ctx, data): return data.get_bytes().split()[0] consts = [blob(random_string(rnd, BLOB_SIZE)) for i in range(BLOB_COUNT)] ts = [] for i in range(BLOB_COUNT): for j in range(BLOB_COUNT): t1 = cat(input1=consts[i], input2=consts[j]) t2 = md5sum(input=t1) t3 = take_first(t2) ts.append(t3.output) result = md5sum(input=tasks.concat(ts)) result.output.keep() # session.pending_graph().write("/home/spirali/tmp/graph.dot") session.submit() result.output.fetch() == b"0a9612a2e855278d336a9e1a1589478f -\n"
def test_wait_unsubmitted_task(test_env): test_env.start(1) client = test_env.client s = client.new_session() with s: t1 = tasks.concat(()) with pytest.raises(RainException): t1.wait()
def test_concat2(test_env): """Merge empty list of blobs""" test_env.start(1) with test_env.client.new_session() as s: t1 = tasks.concat(()) t1.output.keep() s.submit() assert t1.output.fetch().get_bytes() == b""
def test_concat1(test_env): """Merge several short blobs""" test_env.start(1) with test_env.client.new_session() as s: t1 = tasks.concat([blob(x) for x in ("Hello ", "", "", "world", "!", "")]) t1.output.keep() s.submit() assert t1.output.fetch().get_bytes() == b"Hello world!"
def test_fetch_unsubmitted_task(test_env): test_env.start(1) client = test_env.client s = client.new_session() with s: t1 = tasks.concat(()) t1.keep_outputs() with pytest.raises(RainException): t1.fetch_outputs()
def test_update(test_env): test_env.start(1) client = test_env.client s = client.new_session() with s: t1 = tasks.concat((blob("a"), blob("b"))) s.submit() s.update((t1, )) t1.wait() s.update((t1, ))
def test_task_wait(test_env): test_env.start(1) client = test_env.client s = client.new_session() with s.bind_only(): t1 = tasks.concat((blob("a"), blob("b"))) assert t1.state is None s.submit() assert t1.state == rpc.common.TaskState.notAssigned t1.wait() assert t1.state == rpc.common.TaskState.finished
def test_concat3(test_env): """Merge large blobs""" test_env.start(1) a = b"a123" * 1000000 b = b"b43" * 2500000 c = b"c" * 1000 with test_env.client.new_session() as s: t1 = tasks.concat((blob(a), blob(c), blob(b), blob(c), blob(a))) t1.output.keep() s.submit() assert t1.output.fetch().get_bytes() == a + c + b + c + a
def test_wait_all(test_env): test_env.start(1) client = test_env.client s = client.new_session() with s: t1 = tasks.concat((blob("a"), blob("b"))) t2 = tasks.sleep(0.5, t1) s.submit() test_env.assert_duration(0.35, 0.65, lambda: s.wait_all()) assert t1.state == rpc.common.TaskState.finished assert t2.state == rpc.common.TaskState.finished test_env.assert_max_duration(0.1, lambda: t2.wait())
def test_dataobj_wait(test_env): test_env.start(1) client = test_env.client s = client.new_session() with s: t1 = tasks.concat((blob("a"), blob("b"))) o1 = t1.output assert t1.state is None s.submit() assert o1.state == rpc.common.DataObjectState.unfinished o1.wait() assert o1.state == rpc.common.DataObjectState.finished
def test_unkeep_unfinished(test_env): test_env.start(1) client = test_env.client s = client.new_session() with s: t1 = tasks.concat((blob("a"), blob("b"))) t1_output = t1.output t1_output.keep() t2 = tasks.sleep(0.3, t1) s.submit() assert t1_output.is_kept() is True t1_output.unkeep() assert t1_output.is_kept() is False t2.wait()
def test_execute_termination(test_env): test_env.start(1) import time with test_env.client.new_session() as s: tasks.execute("sleep 5") s.submit() time.sleep(0.5) with test_env.client.new_session() as s: t1 = tasks.concat((blob("a"), blob("b"))) t1.keep_outputs() s.submit() r = test_env.assert_max_duration(0.2, lambda: t1.output.fetch()) assert b"ab" == r.get_bytes()
def test_submit(test_env): test_env.no_final_check() test_env.start(1) client = test_env.client s = client.new_session() with s: t1 = tasks.concat((blob("a"), blob("b"))) t2 = tasks.sleep(1, t1) assert s.task_count == 2 assert s.dataobj_count == 4 # "a", "b", "ab", "ab" s.submit() assert s.task_count == 0 assert s.dataobj_count == 0 assert t1.state == rpc.common.TaskState.notAssigned assert t2.state == rpc.common.TaskState.notAssigned
def test_wait_some(test_env): test_env.start(1) client = test_env.client s = client.new_session() with s: t1 = tasks.concat(("a", "b")) t2 = tasks.sleep(0.4, t1) s.submit() finished = s.wait_some((t1, ), ()) assert t1.state == rpc.common.TaskState.finished assert t2.state == rpc.common.TaskState.notAssigned assert len(finished) == 2 assert len(finished[0]) == 1 assert len(finished[1]) == 0 assert finished[0][0].id == t1.id t2.wait()
def test_task_export(test_env): import os.path test1 = os.path.join(test_env.work_dir, "TEST1") test2 = os.path.join(test_env.work_dir, "TEST2") test_env.start(1) with test_env.client.new_session() as s: a = blob("Hello ") b = blob("World!") tasks.export(tasks.concat((a, b)), test1) tasks.export(tasks.execute("ls /", stdout="output"), test2) s.submit() s.wait_all() with open(test1) as f: assert f.read() == "Hello World!" with open(test2) as f: assert "bin" in f.read()
def build_fft_like(layers, datalen=1, session=None, submit_every=None): n = 1 << layers dos = [blob('B' * datalen) for i in range(n)] for j in range(layers): if submit_every is not None and j % submit_every == 0: for d in dos: d.keep() session.submit() print("Submit before layer ", j) ts = [ tasks.concat((dos[i], dos[(i + (1 << j)) % n])) for i in range(n) ] dos = [t.output for t in ts] for i, d in enumerate(dos): d.label = "L{}.{}".format(j, i) return dos
def test_big_diamond(test_env): @remote(outputs=("out1", "out2")) def splitter(ctx, data): data = data.get_bytes() left = data[0:len(data)] right = data[len(data):] return {"out1": left, "out2": right} @remote() def upper(ctx, data): return data.get_bytes().upper() LAYERS = 6 rnd = test_rnd() data = random_string(rnd, 100).lower().encode() test_env.start(4) with test_env.client.new_session() as s: layer = [blob(data)] for i in range(LAYERS): new_layer = [] for l in layer: task = splitter(l) new_layer.append(task.outputs["out1"]) new_layer.append(task.outputs["out2"]) layer = new_layer layer = [upper(t) for t in layer] for i in range(LAYERS): new_layer = [] for j in range(0, len(layer), 2): new_layer.append(tasks.concat((layer[j], layer[j+1]))) layer = new_layer # s.pending_graph().write("test.dot") assert len(layer) == 1 result = layer[0] result.output.keep() s.submit() result = result.output.fetch() assert result.get_bytes() == data.upper()
def build_random_layers(layers, n, inputs=2, datalen=1, session=None, rng=42, submit_every=None): if isinstance(rng, int): rng = random.Random(rng) dos = [blob('B' * datalen) for i in range(n)] for j in range(layers): if submit_every is not None and j % submit_every == 0: for d in dos: d.keep() session.submit() print("Submit before layer ", j) ts = [ tasks.concat(set(dos[inp] for inp in rng.sample(range(n), inputs))) for i in range(n) ] dos = [t.output for t in ts] for i, d in enumerate(dos): d.label = "L{}.{}".format(j, i) return dos