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
