def test_circular_resolve(memory_jobstore): from jobflow.core.reference import OutputReference # test catching circular resolve failure ref1 = OutputReference("12345") task_data = {"uuid": ref1.uuid, "index": 1, "output": ref1} memory_jobstore.update(task_data) with pytest.raises(RuntimeError): ref1.resolve(memory_jobstore)
def test_repr(): from jobflow import OutputReference ref = OutputReference("123") assert str(ref) == "OutputReference(123)" ref = OutputReference("123", (("a", "a"),)) assert str(ref) == "OutputReference(123, .a)" ref = OutputReference("123", (("a", "a"), ("i", 1))) assert str(ref) == "OutputReference(123, .a, [1])"
def test_store_inputs(memory_jobstore): from jobflow.core.job import OutputReference, store_inputs test_job = store_inputs(1) test_job.run(memory_jobstore) output = memory_jobstore.query_one({"uuid": test_job.uuid}, ["output"])["output"] assert output == 1 ref = OutputReference("abcd") test_job = store_inputs(ref) test_job.run(memory_jobstore) output = memory_jobstore.query_one({"uuid": test_job.uuid}, ["output"])["output"] assert OutputReference.from_dict(output) == ref
def test_schema(): from pydantic import BaseModel from jobflow import OutputReference class InnerSchema(BaseModel): n: float class MediumSchema(BaseModel): s: str nested: InnerSchema nested_opt: Optional[InnerSchema] nested_u: Union[InnerSchema, dict] nested_l: List[InnerSchema] nested_d: Dict[str, InnerSchema] class MySchema(BaseModel): number: int name: str nested: MediumSchema ref = OutputReference("123", output_schema=MySchema) assert ref.attributes == tuple() # check valid schema access works new_ref = ref.number assert new_ref.uuid == "123" assert new_ref.output_schema is None new_ref = ref["name"] assert new_ref.uuid == "123" assert new_ref.output_schema is None with pytest.raises(AttributeError): assert ref.a.uuid == "123" with pytest.raises(AttributeError): assert ref["a"].uuid == "123" with pytest.raises(AttributeError): assert ref[1].uuid == "123" # check valid nested schemas assert ref.nested.s.uuid == "123" with pytest.raises(AttributeError): assert ref.nested.m.uuid == "123" assert ref.nested.nested.n.uuid == "123" with pytest.raises(AttributeError): assert ref.nested.nested.m.uuid == "123" assert ref.nested.nested_opt.n.uuid == "123" with pytest.raises(AttributeError): assert ref.nested.nested_opt.m.uuid == "123" # Union, List and Dict are currently not recognized by their inner type # but check that there is no problem with them assert ref.nested.nested_u.n.uuid == "123" assert ref.nested.nested_l[0].n.uuid == "123" assert ref.nested.nested_d["a"].n.uuid == "123"
def test_get_set_attr(): from jobflow import OutputReference ref = OutputReference("123") # these should fail with pytest.raises(TypeError): ref["a"] = 1 with pytest.raises(TypeError): ref[1] = 1 with pytest.raises(TypeError): ref.a = 1 ref.uuid = 1 assert ref.uuid == 1
def test_access(): from jobflow import OutputReference # test empty ref = OutputReference("123") assert ref.attributes == tuple() # test bad init with pytest.raises(ValueError): OutputReference("123", (("x", 1),)) new_ref = ref.a assert new_ref.attributes == (("a", "a"),) assert new_ref.uuid == "123" assert isinstance(new_ref, OutputReference) new_ref = ref["a"] assert new_ref.attributes == (("i", "a"),) assert new_ref.uuid == "123" assert isinstance(new_ref, OutputReference) new_ref = ref[1] assert new_ref.attributes == (("i", 1),) assert new_ref.uuid == "123" assert isinstance(new_ref, OutputReference) # test filled ref = OutputReference("123", (("a", "b"),)) new_ref = ref.a assert new_ref.attributes == (("a", "b"), ("a", "a")) assert new_ref.uuid == "123" assert isinstance(new_ref, OutputReference) new_ref = ref["a"] assert new_ref.attributes == (("a", "b"), ("i", "a")) assert new_ref.uuid == "123" assert isinstance(new_ref, OutputReference) with pytest.raises(AttributeError): _ = ref.args with pytest.raises(AttributeError): _ = ref.__fake_variable
def test_not_iterable(): from jobflow.core.reference import OutputReference ref = OutputReference("12345") with pytest.raises(TypeError): next(ref) with pytest.raises(TypeError): for _ in ref: pass
def test_hash(): from jobflow import OutputReference assert hash(OutputReference("123")) == hash(OutputReference("123")) assert hash(OutputReference("123", (("i", 1), ("i", 2)))) == hash( OutputReference("123", (("i", 1), ("i", 2))) ) assert hash(OutputReference("123", (("a", "b"), ("i", 2)))) == hash( OutputReference("123", (("a", "b"), ("i", 2))) ) assert hash(OutputReference("123", (("a", "b"), ("i", "2")))) == hash( OutputReference("123", (("a", "b"), ("i", "2"))) )
def test_find_and_get_references(): from jobflow.core.reference import OutputReference, find_and_get_references ref1 = OutputReference("123") ref2 = OutputReference("1234", (("a", "a"),)) # test single reference assert find_and_get_references(ref1) == (ref1,) # test list and tuple of references assert find_and_get_references([ref1]) == (ref1,) assert set(find_and_get_references([ref1, ref2])) == {ref1, ref2} assert set(find_and_get_references((ref1, ref2))) == {ref1, ref2} # test dictionary dictionary values assert find_and_get_references({"a": ref1}) == (ref1,) assert set(find_and_get_references({"a": ref1, "b": ref2})) == {ref1, ref2} # test nested assert set(find_and_get_references({"a": [ref1, ref2]})) == {ref1, ref2} assert set(find_and_get_references([{"a": ref1}, {"b": ref2}])) == {ref1, ref2}
def test_reference_in_output(memory_jobstore): from jobflow.core.reference import OnMissing, OutputReference # test resolvable reference in job output ref1 = OutputReference("12345") ref2 = OutputReference("56789") task_data1 = {"uuid": ref1.uuid, "index": 1, "output": ref2} task_data2 = {"uuid": ref2.uuid, "index": 1, "output": "xyz"} memory_jobstore.update(task_data1) memory_jobstore.update(task_data2) assert "xyz" == ref1.resolve(memory_jobstore) # test missing reference in output ref1 = OutputReference("12345") ref2 = OutputReference("999") task_data = {"uuid": ref1.uuid, "index": 1, "output": ref2} memory_jobstore.update(task_data) assert ref1.resolve(memory_jobstore, on_missing=OnMissing.NONE) is None assert ref1.resolve(memory_jobstore, on_missing=OnMissing.PASS) == ref2 with pytest.raises(ValueError): ref1.resolve(memory_jobstore, on_missing=OnMissing.ERROR)
def test_job_input_references(): from jobflow.core.job import Job from jobflow.core.reference import OutputReference ref1 = OutputReference("12345") ref2 = OutputReference("54321") test_job = Job(add, function_args=(ref1, ), function_kwargs={"b": ref2}) assert set(test_job.input_references) == {ref1, ref2} assert set(test_job.input_uuids) == {"12345", "54321"} assert test_job.input_references_grouped == { "12345": (ref1, ), "54321": (ref2, ) } # test using reference attributes ref1 = OutputReference("12345", attributes=(("i", "name"), )) ref2 = OutputReference("12345", attributes=(("i", "value"), )) test_job = Job(add, function_args=(ref1, ), function_kwargs={"b": ref2}) assert set(test_job.input_references) == {ref1, ref2} assert set(test_job.input_uuids) == {"12345"} assert set(list(test_job.input_references_grouped)) == {"12345"} assert set(list( test_job.input_references_grouped["12345"])) == {ref1, ref2}
def test_set_uuid(): from jobflow import OutputReference ref = OutputReference("123") new_ref = ref.set_uuid("321") assert ref.uuid == "321" assert new_ref.uuid == "321" ref = OutputReference("123") new_ref = ref.set_uuid("321", inplace=False) assert ref.uuid == "123" assert new_ref.uuid == "321"
def test_as_dict(): from jobflow import OutputReference ref = OutputReference("123") d = ref.as_dict() assert d["@class"] == "OutputReference" assert d["@module"] == "jobflow.core.reference" assert d["uuid"] == "123" ref = OutputReference("123", (("a", "a"), ("i", "b"))) d = ref.as_dict() assert d["@class"] == "OutputReference" assert d["@module"] == "jobflow.core.reference" assert d["uuid"] == "123" assert d["attributes"] == (("a", "a"), ("i", "b"))
def test_job_resolve_args(memory_jobstore): from jobflow.core.job import Job from jobflow.core.reference import OutputReference # test basic run with no references test_job = Job(print, function_args=("I am a job", )) resolved_job = test_job.resolve_args(memory_jobstore) assert test_job == resolved_job ref = OutputReference("1234") memory_jobstore.update({"uuid": "1234", "index": 1, "output": 2}) # test run with input references test_job = Job(add, function_args=(1, ), function_kwargs={"b": ref}) resolved_job = test_job.resolve_args(memory_jobstore) assert test_job == resolved_job assert resolved_job.function_kwargs["b"] == 2 # test resolve with inplace=False test_job = Job(add, function_args=(1, ), function_kwargs={"b": ref}) resolved_job = test_job.resolve_args(memory_jobstore, inplace=False) assert test_job != resolved_job assert resolved_job.function_kwargs["b"] == 2 assert isinstance(test_job.function_kwargs["b"], OutputReference)
def test_resolve_references(memory_jobstore): from jobflow import OnMissing, OutputReference from jobflow.core.reference import resolve_references # resolve single using store ref = OutputReference("123") memory_jobstore.update({"uuid": "123", "index": 1, "output": "xyz"}) output = resolve_references([ref], memory_jobstore) assert len(output) == 1 assert output[ref] == "xyz" # resolve multiple using cache ref1 = OutputReference("123") ref2 = OutputReference("1234") memory_jobstore.update({"uuid": "1234", "index": 1, "output": 101}) output = resolve_references([ref1, ref2], memory_jobstore) assert len(output) == 2 assert output[ref1] == "xyz" assert output[ref2] == 101 # resolve group using cache ref1 = OutputReference("123", (("i", "a"),)) ref2 = OutputReference("123", (("i", "b"),)) ref3 = OutputReference("1234") memory_jobstore.update( {"uuid": "123", "index": 1, "output": {"a": "xyz", "b": "abc"}} ) output = resolve_references([ref1, ref2, ref3], memory_jobstore) assert len(output) == 3 assert output[ref1] == "xyz" assert output[ref2] == "abc" assert output[ref3] == 101 # test on missing ref1 = OutputReference("123") ref2 = OutputReference("12345") memory_jobstore.update({"uuid": "123", "index": 1, "output": "xyz"}) output = resolve_references( [ref1, ref2], memory_jobstore, on_missing=OnMissing.PASS ) assert len(output) == 2 assert output[ref1] == "xyz" assert output[ref2] == ref2 ref2 = OutputReference("12345") with pytest.raises(ValueError): resolve_references([ref1, ref2], memory_jobstore, on_missing=OnMissing.ERROR) # resolve using store and empty cache cache = {} output = resolve_references([ref], memory_jobstore, cache=cache) assert len(output) == 1 assert output[ref] == "xyz" # check cache supersedes store cache = {"123": {1: 101}} output = resolve_references([ref], memory_jobstore, cache=cache) assert len(output) == 1 assert output[ref] == 101 # test attributes ref = OutputReference("123", (("i", "a"), ("i", 1))) memory_jobstore.update({"uuid": "123", "index": 1, "output": {"a": [5, 6, 7]}}) output = resolve_references([ref], memory_jobstore) assert output[ref] == 6 ref = OutputReference("123", (("a", "__module__"),)) memory_jobstore.update({"uuid": "123", "index": 1, "output": OutputReference}) output = resolve_references([ref], memory_jobstore) assert output[ref] == "jobflow.core.reference"
def test_resolve(memory_jobstore): from jobflow import OnMissing, OutputReference ref = OutputReference("123") # fail if cache or store not provided with pytest.raises(TypeError): ref.resolve() # test on missing assert ref.resolve(memory_jobstore, on_missing=OnMissing.NONE) is None assert ref.resolve(memory_jobstore, on_missing=OnMissing.PASS) == ref with pytest.raises(ValueError): ref.resolve(memory_jobstore, on_missing=OnMissing.ERROR) # resolve using store memory_jobstore.update({"uuid": "123", "index": 1, "output": 101}) assert ref.resolve(store=memory_jobstore) == 101 # resolve using store and empty cache cache = {} assert ref.resolve(store=memory_jobstore, cache=cache) == 101 assert cache["123"][1] == 101 # check cache supersedes store cache = {"123": {1: "xyz"}} assert ref.resolve(store=memory_jobstore, cache=cache) == "xyz" assert cache["123"][1] == "xyz" # test indexing ref = OutputReference("123", (("i", "a"), ("i", 1))) memory_jobstore.update({"uuid": "123", "index": 1, "output": {"a": [5, 6, 7]}}) assert ref.resolve(memory_jobstore) == 6 # test attribute access ref = OutputReference("123", (("a", "__module__"),)) memory_jobstore.update({"uuid": "123", "index": 1, "output": OutputReference}) assert ref.resolve(memory_jobstore) == "jobflow.core.reference" # test missing attribute throws error ref = OutputReference("123", (("a", "b"),)) memory_jobstore.update({"uuid": "123", "index": 1, "output": [1234]}) with pytest.raises(AttributeError): ref.resolve(memory_jobstore) # test missing index throws error ref = OutputReference("123", (("i", "b"),)) with pytest.raises(TypeError): ref.resolve(memory_jobstore)
def test_find_and_resolve_references(memory_jobstore): from jobflow.core.reference import ( OnMissing, OutputReference, find_and_resolve_references, ) ref1 = OutputReference("123") ref2 = OutputReference("1234", (("i", "a"),)) memory_jobstore.update({"uuid": "123", "index": 1, "output": 101}) memory_jobstore.update({"uuid": "1234", "index": 1, "output": {"a": "xyz", "b": 5}}) # test no reference assert find_and_resolve_references(True, memory_jobstore) is True assert find_and_resolve_references("xyz", memory_jobstore) == "xyz" assert find_and_resolve_references([101], memory_jobstore) == [101] # test single reference assert find_and_resolve_references(ref1, memory_jobstore) == 101 # test list and tuple of references assert find_and_resolve_references([ref1], memory_jobstore) == [101] assert find_and_resolve_references([ref1, ref2], memory_jobstore) == [101, "xyz"] # test dictionary dictionary values output = find_and_resolve_references({"a": ref1}, memory_jobstore) assert output == {"a": 101} output = find_and_resolve_references({"a": ref1, "b": ref2}, memory_jobstore) assert output == { "a": 101, "b": "xyz", } # test nested output = find_and_resolve_references({"a": [ref1, ref2]}, memory_jobstore) assert output == {"a": [101, "xyz"]} output = find_and_resolve_references([{"a": ref1}, {"b": ref2}], memory_jobstore) assert output == [ {"a": 101}, {"b": "xyz"}, ] # test store, blank cache cache = {} output = find_and_resolve_references( {"a": [ref1, ref2]}, store=memory_jobstore, cache=cache ) assert output == {"a": [101, "xyz"]} assert cache["123"][1] == 101 # test cache overrides store output = find_and_resolve_references( {"a": [ref1, ref2]}, store=memory_jobstore, cache={"123": {1: 1}} ) assert output == {"a": [1, "xyz"]} # test on missing ref3 = OutputReference("12345", (("i", "a"),)) output = find_and_resolve_references( [ref1, ref3], memory_jobstore, on_missing=OnMissing.PASS ) assert output == [101, ref3] output = find_and_resolve_references( [ref1, ref3], memory_jobstore, on_missing=OnMissing.NONE ) assert output == [101, None] with pytest.raises(ValueError): find_and_resolve_references( [ref1, ref3], memory_jobstore, on_missing=OnMissing.ERROR )
def test_eq(): from jobflow import OutputReference assert OutputReference("123") == OutputReference("123") assert OutputReference("123") != OutputReference("1234") assert OutputReference("123", (("i", 1),)) == OutputReference("123", (("i", 1),)) assert OutputReference("123", (("i", "a"),)) == OutputReference( "123", (("i", "a"),) ) assert OutputReference("123", (("a", "a"),)) == OutputReference( "123", (("a", "a"),) ) assert OutputReference("123", (("a", "a"), ("i", "b"))) == OutputReference( "123", (("a", "a"), ("i", "b")) ) assert OutputReference("123", (("i", 1),)) != OutputReference("1234", (("i", 1),)) assert OutputReference("123", (("i", 1),)) != OutputReference("123", (("i", 2),)) assert OutputReference("123", (("i", 1),)) != OutputReference( "123", (("i", 2), ("i", 3), ("i", 4)) ) assert OutputReference("123", (("i", 1),)) != "OutputReference(123, [1])"
def test_job_config(memory_jobstore): from jobflow import ( CURRENT_JOB, Job, JobConfig, OnMissing, OutputReference, Response, ) def store_exposed(): return CURRENT_JOB.store is not None def reference_resolved(arg): return not isinstance(arg, OutputReference) def return_arg(arg): return arg # test expose store config = JobConfig(expose_store=False) test_job = Job(store_exposed, config=config) response = test_job.run(memory_jobstore) assert response.output is False config = JobConfig(expose_store=True) test_job = Job(store_exposed, config=config) response = test_job.run(memory_jobstore) assert response.output is True ref = OutputReference("1234") memory_jobstore.update({"uuid": "1234", "index": 0, "output": 5}) config = JobConfig(resolve_references=False) test_job = Job(reference_resolved, function_args=(ref, ), config=config) response = test_job.run(memory_jobstore) assert response.output is False config = JobConfig(resolve_references=True) test_job = Job(reference_resolved, function_args=(ref, ), config=config) response = test_job.run(memory_jobstore) assert response.output is True ref = OutputReference("xyz") config = JobConfig(on_missing_references=OnMissing.ERROR) test_job = Job(return_arg, function_args=(ref, ), config=config) with pytest.raises(ValueError): test_job.run(memory_jobstore) config = JobConfig(on_missing_references=OnMissing.NONE) test_job = Job(return_arg, function_args=(ref, ), config=config) response = test_job.run(memory_jobstore) assert response.output is None config = JobConfig(on_missing_references=OnMissing.PASS) test_job = Job(return_arg, function_args=(ref, ), config=config) response = test_job.run(memory_jobstore) assert isinstance(response.output, OutputReference) # test pass manager config def replace_job(): job = Job(add, function_args=(1, )) return Response(replace=job) def replace_list_job(): job1 = Job(add, function_args=(1, )) job2 = Job(add, function_args=(job1.output, )) return Response(replace=[job1, job2]) def replace_flow(): from jobflow import Flow job = Job(add, function_args=(1, )) flow = Flow([job], output=job.output) return Response(replace=flow) def addition_job(): job = Job(add, function_args=(1, )) return Response(addition=job) def detour_job(): job = Job(add, function_args=(1, )) return Response(detour=job) manager_config = {"abc": 1} manager_config2 = {"abc": 2} pass_config = JobConfig(manager_config=manager_config, pass_manager_config=True) nopass_config = JobConfig(manager_config=manager_config, pass_manager_config=False) response_config = JobConfig(manager_config=manager_config, response_manager_config=manager_config2) # test replace test_job = Job(replace_job, config=nopass_config) response = test_job.run(memory_jobstore) assert len(response.replace.jobs) == 1 assert response.replace.jobs[0].config.manager_config == {} test_job = Job(replace_job, config=pass_config) response = test_job.run(memory_jobstore) assert response.replace.jobs[0].config.manager_config == manager_config test_job = Job(replace_job, config=response_config) response = test_job.run(memory_jobstore) assert response.replace.jobs[0].config.manager_config == manager_config2 # test replace list of jobs test_job = Job(replace_list_job, config=nopass_config) response = test_job.run(memory_jobstore) for j in response.replace.jobs: assert j.config.manager_config == {} test_job = Job(replace_list_job, config=pass_config) response = test_job.run(memory_jobstore) for j in response.replace.jobs: assert j.config.manager_config == manager_config test_job = Job(replace_list_job, config=response_config) response = test_job.run(memory_jobstore) for j in response.replace.jobs: assert j.config.manager_config == manager_config2 # test replace with flow test_job = Job(replace_flow, config=nopass_config) response = test_job.run(memory_jobstore) for j in response.replace.jobs: assert j.config.manager_config == {} test_job = Job(replace_flow, config=pass_config) response = test_job.run(memory_jobstore) for j in response.replace.jobs: assert j.config.manager_config == manager_config test_job = Job(replace_flow, config=response_config) response = test_job.run(memory_jobstore) for j in response.replace.jobs: assert j.config.manager_config == manager_config2 # test addition test_job = Job(addition_job, config=nopass_config) response = test_job.run(memory_jobstore) assert len(response.addition.jobs) == 1 assert response.addition.jobs[0].config.manager_config == {} test_job = Job(addition_job, config=pass_config) response = test_job.run(memory_jobstore) assert len(response.addition.jobs) == 1 print(response.addition) assert response.addition.jobs[0].config.manager_config == manager_config assert response.addition.jobs[0].hosts == [response.addition.uuid] test_job = Job(addition_job, config=response_config) response = test_job.run(memory_jobstore) assert response.addition.jobs[0].config.manager_config == manager_config2 # test detour test_job = Job(detour_job, config=nopass_config) response = test_job.run(memory_jobstore) assert len(response.detour.jobs) == 1 assert response.detour.jobs[0].config.manager_config == {} test_job = Job(detour_job, config=pass_config) response = test_job.run(memory_jobstore) assert len(response.detour.jobs) == 1 assert response.detour.jobs[0].config.manager_config == manager_config assert response.detour.jobs[0].hosts == [response.detour.uuid] test_job = Job(detour_job, config=response_config) response = test_job.run(memory_jobstore) assert response.detour.jobs[0].config.manager_config == manager_config2