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)