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_program_name(fake_session):
p1 = Program(["/bin/ls"], name="program1")
p2 = Program(["/bin/ls"])
with fake_session:
t1 = p1()
assert t1.name == "program1"
t2 = p1(name="xxx")
assert t2.name == "xxx"
t3 = p2()
assert t3.name is None
t4 = p2(name="yyy")
assert t4.name == "yyy"
def test_program_sleep_1(test_env):
"""Sleep followed by wait"""
test_env.start(1)
program = Program("sleep 1")
with test_env.client.new_session() as s:
t1 = program()
s.submit()
test_env.assert_duration(0.99, 1.1, lambda: t1.wait())
def test_program_stdout_only(test_env):
"""Capturing stdout"""
test_env.start(1)
program = Program("ls /", stdout="output")
with test_env.client.new_session() as s:
t1 = program()
t1.output.keep()
s.submit()
assert b"etc\n" in t1.output.fetch().get_bytes()
def test_program_fail(test_env):
"""Setting input file for program"""
test_env.start(1)
args = ("ls", "/non-existing-dir")
program = Program(args, stdout="output")
with test_env.client.new_session() as s:
t1 = program()
s.submit()
pytest.raises(RainException, lambda: t1.wait())
def test_program_input_file(test_env):
"""Setting input file for program"""
test_env.start(1)
program = Program(("/bin/grep", "ab", Input("in1")), stdout="output")
with test_env.client.new_session() as s:
t1 = program(in1=blob("abc\nNOTHING\nabab"))
t1.output.keep()
s.submit()
assert t1.output.fetch().get_bytes() == b"abc\nabab\n"
def test_program_invalid_filename(test_env):
"""Setting input file for program"""
test_env.start(1)
args = ("/bin/non-existing-program", )
program = Program(args, stdout="output")
with test_env.client.new_session() as s:
t1 = program()
t1.output.keep()
s.submit()
pytest.raises(RainException, lambda: t1.wait())
def test_program_create_file(test_env):
"""Capturing file"""
test_env.start(1)
args = ("/bin/bash", "-c", "echo ABC > output.txt")
program = Program(args,
output_files=[Output("my_output", path="output.txt")])
with test_env.client.new_session() as s:
t1 = program()
t1.outputs["my_output"].keep()
s.submit()
assert t1.outputs["my_output"].fetch().get_bytes() == b"ABC\n"
def test_execute_shell(test_env):
test_env.start(1)
p1 = Program(("echo", "$HOME"), stdout=True)
p2 = Program(("echo", "$HOME"), stdout=True, shell=True)
with test_env.client.new_session() as s:
t1 = tasks.execute(("echo", "$HOME"), stdout=True)
t1.output.keep()
t2 = tasks.execute(("echo", "$HOME"), stdout=True, shell=True)
t2.output.keep()
t3 = p1()
t3.output.keep()
t4 = p2()
t4.output.keep()
s.submit()
assert b"$HOME\n" == t1.output.fetch().get_bytes()
assert (os.getenv("HOME") +
"\n").encode() == t2.output.fetch().get_bytes()
assert b"$HOME\n" == t3.output.fetch().get_bytes()
assert (os.getenv("HOME") +
"\n").encode() == t4.output.fetch().get_bytes()
def test_fetch_from_failed_session_immediate(test_env):
"""Setting input file for program"""
test_env.start(1)
args = ("/bin/non-existing-program", )
program = Program(args, stdout="output")
with test_env.client.new_session() as s:
t1 = program()
t1.output.keep()
s.submit()
with pytest.raises(TaskException):
t1.output.fetch()
with pytest.raises(TaskException):
t1.output.fetch()
def test_early_wait_all_failed_(test_env):
test_env.start(1)
client = test_env.client
s = client.new_session()
with s:
t0 = tasks.sleep(0.4, blob("test"))
args = ("/bin/non-existing-program")
program = Program(args, stdout="output", stdin="input")
t1 = program(input=t0)
t1_output = t1.output
t1_output.keep()
s.submit()
with pytest.raises(TaskException):
s.wait_all()
def test_late_wait_all_failed(test_env):
test_env.start(1)
client = test_env.client
s = client.new_session()
with s:
args = ("/bin/non-existing-program", )
program = Program(args, stdout="output")
t1 = program()
t1_output = t1.output
t1_output.keep()
s.submit()
time.sleep(0.3)
with pytest.raises(TaskException):
s.wait_all()
def test_update_from_failed_session(test_env):
"""Setting input file for program"""
test_env.start(1)
args = ("/bin/non-existing-program", )
program = Program(args, stdout="output")
with test_env.client.new_session() as s:
t1 = program()
t1.output.keep()
s.submit()
time.sleep(0.6)
with pytest.raises(RainException):
t1.update()
with pytest.raises(RainException):
t1.output.update()
def test_program_outputs(test_env):
"Specify program content type on spec and instantiation."
obj = ["1", 2.0, {'a': 42}]
program1 = Program(["cat", Input("i")], stdout="o")
program2 = Program(["cat", Input("i", content_type='pickle')],
stdout=Output(content_type='pickle'))
test_env.start(1)
with test_env.client.new_session() as s:
# Dynamic content-type, forgotten by cat
t1a = program1(i=pickled(obj))
t1a.output.keep()
# Static content-type by instantiation
t1b = program1(i=pickled(obj), output=Output(content_type='pickle'))
t1b.output.keep()
# No content type
t1c = program1(i=blob(pickle.dumps(obj)))
t1c.output.keep()
# Static content-type by Program spec
t2 = program2(i=pickled(obj))
t2.output.keep()
s.submit()
assert t1a.output.content_type is None
with pytest.raises(RainException):
assert t1a.output.fetch().load() == obj
assert t1a.output.fetch().get_bytes() == pickle.dumps(obj)
# TODO(gavento): Needs OutputSpec and Output merging
# assert t1b.output.fetch().load() == obj
assert t1c.output.content_type is None
with pytest.raises(RainException):
t1c.output.fetch().load()
assert t1a.output.fetch().get_bytes() == pickle.dumps(obj)
assert t2.output.fetch().load() == obj
def test_py_pass_through(test_env):
@remote(outputs=("out1", "out2"))
def test(ctx, data1, data2):
return {"out1": data1, "out2": data2}
test_env.start(1)
cat = Program("/bin/cat input1", stdout="output", input_paths=[Input("input1")])
with test_env.client.new_session() as s:
data = b"ABC" * 10000
t0 = cat(input1=blob(data))
t1 = test(t0, blob("Hello!"))
t1.outputs["out1"].keep()
t1.outputs["out2"].keep()
s.submit()
assert data == t1.outputs["out1"].fetch().get_bytes()
assert b"Hello!" == t1.outputs["out2"].fetch().get_bytes()
def test_unkeep_failed(test_env):
test_env.start(1)
client = test_env.client
s = client.new_session()
with s:
args = ("/bin/non-existing-program", )
program = Program(args, stdout="output")
t1 = program()
t1_output = t1.output
t1_output.keep()
s.submit()
time.sleep(0.6)
with pytest.raises(RainException):
t1_output.unkeep()
with pytest.raises(RainException):
t1_output.unkeep()