def test_worker_killed(nworkers: int, sig: int) -> None:
"""Test what happens when 'funsies worker' gets killed."""
# std
import os
def kill_funsies_worker(*inp: bytes) -> bytes:
pid = os.getppid()
os.kill(pid, sig)
time.sleep(1.0)
return b"what"
def cap(inp: bytes) -> bytes:
return inp.upper()
with f.ManagedFun(nworkers=nworkers) as db:
wait_for_workers(db, nworkers)
s1 = f.reduce(kill_funsies_worker,
b"bla bla",
b"bla bla",
opt=f.options(timeout=5))
s1b = f.morph(cap, s1)
f.execute(s1b)
if nworkers == 1:
# no other workers to pick up the slack
with pytest.raises(TimeoutError):
f.wait_for(s1b, timeout=1)
else:
# everything is ok
f.wait_for(s1b, timeout=5)
assert f.take(s1b) == b"WHAT"
def test_job_killed(nworkers: int, sig: int) -> None:
"""Test what happens when 'funsies worker' is ok but its job gets killed."""
# std
import os
def kill_self(*inp: bytes) -> bytes:
pid = os.getpid()
os.kill(pid, sig)
time.sleep(2.0)
return b"what"
def cap(inp: bytes) -> bytes:
return inp.upper()
with f.ManagedFun(nworkers=nworkers) as db:
wait_for_workers(db, nworkers)
s1 = f.reduce(kill_self,
b"bla bla",
b"bla bla",
opt=f.options(timeout=3))
s1b = f.morph(cap, s1)
f.execute(s1b)
# error
f.wait_for(s1b, timeout=1)
err = f.take(s1b, strict=False)
assert isinstance(err, f.errors.Error)
assert err.kind == f.errors.ErrorKind.KilledBySignal
def optimize_conformer_xtb(structure: Artefact[bytes]) -> Artefact[bytes]:
"""Optimize a structure using xtb."""
# perform an xtb optimization
optim = f.shell(
"xtb input.xyz --opt vtight",
inp={"input.xyz": structure},
out=[".xtboptok", "xtbopt.xyz"],
)
# barrier for convergence (fails if .xtboptok is not found)
struct: Artefact[bytes] = f.reduce(lambda x, y: x, optim.out["xtbopt.xyz"],
optim.out[".xtboptok"])
return struct
def funsies_mergesort(art: Artefact[list[int]]) -> Artefact[list[int]]:
"""Mergesort a list of numbers with funsies."""
result = f.dynamic.sac(
# Recursive application of merge sort
# split -> generates two list or raises
# recurse(x) for each values of split
# merge(left, right)
split,
lambda element: funsies_mergesort(element),
lambda lr: f.reduce(merge, lr[0], lr[1]), # type:ignore
art,
out=Encoding.json,
)
return f.reduce(
# if the subdag fails, it's because split raised. In this case, we
# just forward the arguments
ignore_error,
result,
art,
strict=False,
out=Encoding.json,
)
def test_infer_errs() -> None:
"""Test inference applied to functions."""
with f.Fun(MockServer()):
a = f.put(b"bla bla")
b = f.put(3)
with pytest.raises(TypeError):
f.py(lambda x, y, z: (x, y), a, a, b)
# should NOT raise
f.py(
lambda x, y, z: (x, y),
a,
a,
b,
out=[types.Encoding.blob, types.Encoding.blob],
)
def i1o2(x: bytes) -> Tuple[bytes, bytes]:
return x, x
def i2o1(x: bytes, y: bytes) -> bytes:
return x
with pytest.raises(TypeError):
out = f.morph(i1o2, a) # type:ignore # noqa:F841
with pytest.raises(TypeError):
out = f.reduce(i1o2, a) # type:ignore # noqa:F841
with pytest.raises(TypeError):
out = f.reduce(lambda x, y: x, a, b) # type:ignore # noqa:F841
# If we pass out= then the inference is skipped
out = f.morph(i1o2, a,
out=types.Encoding.blob) # type:ignore # noqa:F841
out = f.reduce(i1o2, a,
out=types.Encoding.blob) # type:ignore # noqa:F841
def make_data_output(
structures: Sequence[Artefact[bytes]]) -> Artefact[list[Any]]:
"""Take xyz structure from xtb and parse them to a list of dicts."""
def to_dict(xyz: bytes) -> dict[str, Any]:
as_str = xyz.decode().strip()
energy = float(as_str.splitlines()[1].split()[1])
return {"structure": as_str, "energy": energy}
def sort_by_energy(*elements: dict[str, Any]) -> list[dict[str, Any]]:
out = [el for el in elements]
out = sorted(out, key=lambda x: x["energy"]) # type:ignore
return out
out = []
for s in structures:
out += [f.morph(to_dict, s, out=Encoding.json)] # elements to dicts
return f.reduce(sort_by_energy, *out) # transform to a sorted list
def test_parametric_store_recall() -> None:
"""Test storing and recalling parametrics."""
serv = MockServer()
with Fun(serv, options(distributed=False)):
a = put(3)
b = put(4)
s = reduce(lambda x, y: x + y, a, b)
s2 = morph(lambda x: 3 * x, s)
execute(s2)
assert take(s2) == 21
# parametrize
p.commit("math", inp=dict(a=a, b=b), out=dict(s=s, s2=s2))
with Fun(serv, options(distributed=False)):
out = p.recall("math", inp=dict(a=5, b=8))
execute(out["s2"])
assert take(out["s2"]) == 39
def test_raising_funsie() -> None:
"""Test funsie that raises an error.
This test is specifically designed to catch the bug fixed in fa9af6a4
where funsies that raised did not release their locks, leading to a race
condition.
"""
def raising_fun(*inp: str) -> bytes:
raise RuntimeError("this funsie raises.")
with f.ManagedFun(nworkers=2):
s0a = f.morph(lambda x: x, "bla blabla")
s0b = f.morph(lambda x: x, "blala")
s1 = f.reduce(raising_fun, "bla bla", s0a, s0b, strict=True)
f.execute(s1)
f.wait_for(s1, timeout=2)
with pytest.raises(UnwrapError):
_ = f.take(s1)
s2 = f.morph(lambda x: x, s1)
f.execute(s2)
f.wait_for(s2, timeout=0.5)
def test_timeout_deadlock() -> None:
"""Test funsies that time out.
Here we explicitly check if dependents are still enqueued or if the whole
thing deadlocks.
"""
def timeout_fun(*inp: str) -> bytes:
time.sleep(3.0)
return b"what"
def cap(inp: bytes) -> bytes:
return inp.capitalize()
with f.ManagedFun(nworkers=2):
# Test when python function times out
s1 = f.reduce(timeout_fun,
"bla bla",
"bla bla",
opt=f.options(timeout=1))
s1b = f.morph(cap, s1)
# Test when shell function times out
s2 = f.shell("sleep 20", "echo 'bla bla'", opt=f.options(timeout=1))
s2b = f.morph(cap, s2.stdouts[1])
f.execute(s1b, s2b)
# Check err for reduce
f.wait_for(s1b, timeout=1.5)
err = f.take(s1b, strict=False)
assert isinstance(err, f.errors.Error)
assert err.kind == f.errors.ErrorKind.JobTimedOut
assert err.source == s1.parent
# Check err for shell
f.wait_for(s2b, timeout=1.5)
err = f.take(s2b, strict=False)
assert isinstance(err, f.errors.Error)
assert err.kind == f.errors.ErrorKind.JobTimedOut
assert err.source == s2.hash
def test_cleanup() -> None:
"""Test truncation."""
# std
import os
def kill_self(*inp: bytes) -> bytes:
pid = os.getpid()
os.kill(pid, SIGKILL)
time.sleep(2.0)
return b"what"
with f.ManagedFun(nworkers=1) as db:
inp = "\n".join([f"{k}" for k in range(10)]).encode()
fun = f.reduce(kill_self, inp)
f.execute(fun)
time.sleep(0.5)
key1 = db.get(
f._constants.join(f._constants.OPERATIONS, fun.parent, "owner"))
f._context.cleanup_funsies(db)
key2 = db.get(
f._constants.join(f._constants.OPERATIONS, fun.parent, "owner"))
assert key1 is not None
assert key2 is None
def test_parametric_store_recall_optional() -> None:
"""Test storing a parametric with optional parameters."""
serv = MockServer()
with Fun(serv, options(distributed=False)):
a = put(3)
b = put("fun")
s = reduce(lambda x, y: x * y, a, b)
s2 = morph(lambda x: x.upper(), s)
# parametrize
p.commit("fun", inp=dict(a=a, b=b), out=dict(s=s2))
with Fun(serv, options(distributed=False)):
out = p.recall("fun", inp=dict(a=5))
execute(out["s"])
assert take(out["s"]) == "FUNFUNFUNFUNFUN"
# nested
out = p.recall("fun", inp=dict(b="lol"))
out = p.recall("fun", inp=dict(b=out["s"], a=2))
execute(out["s"])
assert take(out["s"]) == "LOLLOLLOLLOLLOLLOL"
def combine_inner(inp: Sequence[Artefact]) -> Artefact:
return funsies.reduce(sum_inputs, *inp)
def apply_inner(inp: Artefact) -> Artefact:
return funsies.reduce(sum_inputs, inp, 1)
def test_integration(reference: str, nworkers: int) -> None:
"""Test full integration."""
# make a temp file and copy reference database
dir = tempfile.mkdtemp()
if not make_reference:
shutil.copy(os.path.join(ref_dir, reference, "appendonly.aof"), dir)
shutil.copytree(os.path.join(ref_dir, reference, "data"),
os.path.join(dir, "data"))
shutil.copy(os.path.join(ref_dir, "redis.conf"), dir)
# data url
datadir = f"file://{os.path.join(dir, 'data')}"
# Dictionary for test data
test_data: dict[str, Any] = {}
def update_data(a: dict[int, int], b: list[int]) -> dict[int, int]:
for i in b:
a[i] = a.get(i, 0) + 1
return a
def sum_data(x: dict[int, int]) -> int:
return sum([int(k) * v for k, v in x.items()])
def make_secret(x: int) -> str:
return secrets.token_hex(x)
# Start funsie script
with ManagedFun(
nworkers=nworkers,
directory=dir,
data_url=datadir,
redis_args=["redis.conf"],
) as db:
integers = put([5, 4, 8, 9, 9, 10, 1, 3])
init_data = put({100: 9})
test_data["init_data"] = init_data
nbytes = put(4)
s1 = reduce(update_data, init_data, integers)
num = morph(sum_data, s1)
date = shell("date").stdout
test_data["date"] = date
rand = morph(make_secret, nbytes)
s4 = template(
"date:{{date}}\n" + "some random bytes:{{random}}\n" +
"a number: {{num}}\n" + "a string: {{string}}\n",
{
"date": date,
"random": rand,
"num": num,
"string": "wazza"
},
name="a template",
)
test_data["s4"] = s4
execute(s4)
wait_for(s4, 5)
# check that the db doesn't itself include data
for k in db.keys():
assert b"data" not in k
if make_reference:
folder = os.path.join(ref_dir, reference)
os.makedirs(folder, exist_ok=True)
for name, artefact in test_data.items():
with open(os.path.join(folder, name), "wb") as f:
execute(artefact)
wait_for(artefact, 10.0)
out = take(artefact)
data2 = _serdes.encode(artefact.kind, out)
assert isinstance(data2, bytes)
f.write(data2)
shutil.copy(
os.path.join(dir, "appendonly.aof"),
os.path.join(folder, "appendonly.aof"),
)
shutil.copytree(
os.path.join(dir, "data"),
os.path.join(folder, "data"),
)
else:
# Test against reference dbs
for name, artefact in test_data.items():
execute(artefact)
wait_for(artefact, 10.0)
with open(os.path.join(ref_dir, reference, name), "rb") as f:
data = f.read()
out = take(artefact)
data_ref = _serdes.encode(artefact.kind, out)
assert isinstance(data_ref, bytes)
assert data == data_ref
# delete tempdir
shutil.rmtree(dir)
def test_integration(reference: str, nworkers: int) -> None:
"""Test full integration."""
# make a temp file and copy reference database
dir = tempfile.mkdtemp()
if not make_reference:
shutil.copy(os.path.join(ref_dir, reference, "appendonly.aof"), dir)
shutil.copy(os.path.join(ref_dir, "redis.conf"), dir)
# Dictionary for test data
test_data = {}
# Start funsie script
with ManagedFun(nworkers=nworkers,
directory=dir,
redis_args=["redis.conf"]):
dat = put(b"bla bla")
step1 = morph(lambda x: x.decode().upper().encode(), dat)
step2 = shell(
"cat file1 file2; grep 'bla' file2 file1 > file3; date >> file3",
inp=dict(file1=step1, file2=dat),
out=["file2", "file3"],
)
echo = shell("sleep 1", "date")
merge = reduce(
join_bytes,
step2.out["file3"],
echo.stdouts[1],
name="merger",
)
def tolist(x: bytes, y: bytes) -> Dict[int, str]:
return {1: x.decode(), 8: y.decode()}
A = py(tolist, merge, echo.stdouts[1])
test_data["test1"] = A
def raises(inp: bytes) -> bytes:
raise RuntimeError("an error was raised")
def error_count(*inp: Result[bytes]) -> bytes:
out = utils.match_results(inp, lambda x: 0, lambda x: 1)
return str(sum(out)).encode()
err = morph(raises, dat)
count = reduce(error_count,
dat,
dat,
err,
dat,
err,
err,
echo.stdouts[0],
strict=False)
cat = utils.concat(merge,
dat,
err,
count,
echo.stdouts[1],
strict=False)
test_data["test2"] = cat
execute(step1)
wait_for(step1, timeout=10.0)
execute(step2)
wait_for(step2, timeout=10.0)
assert take(step1) == b"BLA BLA"
assert take(step2.stdout) == b"BLA BLAbla bla"
if make_reference:
folder = os.path.join(ref_dir, reference)
os.makedirs(folder, exist_ok=True)
for name, artefact in test_data.items():
with open(os.path.join(folder, name), "wb") as f:
execute(artefact)
wait_for(artefact, 10.0)
out = take(artefact)
data2 = _serdes.encode(artefact.kind, out)
assert isinstance(data2, bytes)
f.write(data2)
shutil.copy(
os.path.join(dir, "appendonly.aof"),
os.path.join(folder, "appendonly.aof"),
)
else:
# Test against reference dbs
for name, artefact in test_data.items():
execute(artefact)
wait_for(artefact, 10.0)
with open(os.path.join(ref_dir, reference, name), "rb") as f:
data = f.read()
out = take(artefact)
data_ref = _serdes.encode(artefact.kind, out)
assert isinstance(data_ref, bytes)
assert data == data_ref
shutil.rmtree(dir)
