def test_set_uuid(): from jobflow import Job test_job = Job(add, function_args=(1, )) test_job.set_uuid("1234") assert test_job.uuid == "1234" assert test_job.output.uuid == "1234"
def test_pass_manager_config(): from jobflow import Flow, Job from jobflow.core.job import pass_manager_config manager_config = {"abc": 1} # test single job test_job1 = Job(add, function_args=(1, )) pass_manager_config(test_job1, manager_config) assert test_job1.config.manager_config == manager_config # test list job test_job1 = Job(add, function_args=(1, )) test_job2 = Job(add, function_args=(1, )) pass_manager_config([test_job1, test_job2], manager_config) assert test_job1.config.manager_config == manager_config assert test_job2.config.manager_config == manager_config # test flow test_job1 = Job(add, function_args=(1, )) test_job2 = Job(add, function_args=(1, )) flow = Flow([test_job1, test_job2]) pass_manager_config(flow, manager_config) assert test_job1.config.manager_config == manager_config assert test_job2.config.manager_config == manager_config # test bad input with pytest.raises(ValueError): pass_manager_config(["str"], manager_config)
def replace_list_flow(): job1 = Job(add, function_args=(1, )) job2 = Job(add, function_args=(job1.output, )) job3 = Job(add, function_args=(5, )) job4 = Job(add, function_args=(job3.output, )) flow1 = Flow([job1, job2], output={"1": job1.output, "2": job2.output}) flow2 = Flow([job3, job4], output={"3": job3.output, "4": job4.output}) return Response(replace=[flow1, flow2])
def test_graph(): from jobflow import Job # test no inputs test_job = Job(add, function_args=(1, ), function_kwargs={"b": 2}) graph = test_job.graph assert len(graph.nodes) == 1 assert len(graph.edges) == 0 # test arg inputs test_job1 = Job(add, function_args=(1, ), function_kwargs={"b": 2}) test_job = Job(add, function_args=(test_job1.output, ), function_kwargs={"b": 2}) graph = test_job.graph assert len(graph.nodes) == 2 assert len(graph.edges) == 1 # test arg input properties test_job1 = Job(add, function_args=(1, ), function_kwargs={"b": 2}) test_job = Job(add, function_args=(test_job1.output.value, ), function_kwargs={"b": 2}) graph = test_job.graph assert len(graph.nodes) == 2 assert len(graph.edges) == 1 assert graph.get_edge_data(test_job1.uuid, test_job.uuid)["properties"] == "value" # test arg and kwargs inputs test_job1 = Job(add, function_args=(1, ), function_kwargs={"b": 2}) test_job2 = Job(add, function_args=(1, ), function_kwargs={"b": 2}) test_job = Job(add, function_args=(test_job1.output, ), function_kwargs={"b": test_job2.output}) graph = test_job.graph assert len(graph.nodes) == 3 assert len(graph.edges) == 2 # test arg and kwargs input properties test_job1 = Job(add, function_args=(1, ), function_kwargs={"b": 2}) test_job = Job( add, function_args=(test_job1.output.value, ), function_kwargs={"b": test_job1.output.name}, ) graph = test_job.graph assert len(graph.nodes) == 2 assert len(graph.edges) == 1 assert set( graph.get_edge_data(test_job1.uuid, test_job.uuid)["properties"]) == { "value", "name", }
def test_hosts(memory_jobstore): from jobflow import Job test_job = Job(function=add, function_args=(1, )) assert test_job.hosts == [] test_job.add_hosts_uuids("12345") assert test_job.hosts == ["12345"] test_job.add_hosts_uuids(["67890"]) assert test_job.hosts == ["12345", "67890"] test_job.add_hosts_uuids(["09876"], prepend=True) assert test_job.hosts == ["09876", "12345", "67890"] test_job.run(memory_jobstore) result = memory_jobstore.query_one({"uuid": test_job.uuid}) assert result["hosts"] == ["09876", "12345", "67890"]
def test_append_name(): from jobflow import Job # test append test_job = Job(function=add) test_job.append_name(" test") assert test_job.name == "add test" # test prepend test_job = Job(function=add) test_job.append_name("test ", prepend=True) assert test_job.name == "test add"
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_serialization(): import json from monty.json import MontyDecoder, MontyEncoder from jobflow import Job test_job = Job(function=add, function_args=(1, ), function_kwargs={"b": 2}) uuid = test_job.uuid encoded_job = json.loads(MontyEncoder().encode(test_job)) decoded_job = MontyDecoder().process_decoded(encoded_job) assert decoded_job.uuid == uuid
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_update_config(): from jobflow import Job, JobConfig new_config = JobConfig( resolve_references=False, manager_config={"a": "b"}, pass_manager_config=False, ) new_config_dict = { "resolve_references": False, "manager_config": { "a": "b" }, "pass_manager_config": False, } # test no filter test_job = Job(add) test_job.update_config(new_config) assert test_job.config == new_config # test name filter test_job = Job(add) test_job.update_config(new_config, name_filter="add") assert test_job.config == new_config test_job = Job(add) test_job.update_config(new_config, name_filter="div") assert test_job.config != new_config # test function filter test_job = Job(add) test_job.update_config(new_config, function_filter=add) assert test_job.config == new_config test_job = Job(add) test_job.update_config(new_config, function_filter=list) assert test_job.config != new_config # test attributes test_job = Job(add) test_job.update_config(new_config, attributes="manager_config") assert test_job.config.manager_config == {"a": "b"} assert test_job.config.resolve_references test_job = Job(add) test_job.update_config(new_config, attributes=["manager_config"]) assert test_job.config.manager_config == {"a": "b"} assert test_job.config.resolve_references test_job = Job(add) test_job.update_config(new_config, attributes=["manager_config", "resolve_references"]) assert test_job.config.manager_config == {"a": "b"} assert not test_job.config.resolve_references assert test_job.config.pass_manager_config with pytest.raises(ValueError): test_job.update_config(new_config, attributes="abc_xyz") # test dictionary config updates test_job = Job(add) test_job.update_config(new_config_dict) assert test_job.config == new_config # test dict with attributes test_job = Job(add) test_job.update_config(new_config_dict, attributes="manager_config") assert test_job.config.manager_config == {"a": "b"} assert test_job.config.resolve_references test_job = Job(add) test_job.update_config(new_config_dict, attributes=["manager_config"]) assert test_job.config.manager_config == {"a": "b"} assert test_job.config.resolve_references test_job = Job(add) test_job.update_config(new_config_dict, attributes=["manager_config", "resolve_references"]) assert test_job.config.manager_config == {"a": "b"} assert not test_job.config.resolve_references assert test_job.config.pass_manager_config with pytest.raises(ValueError): test_job.update_config(new_config_dict, attributes="abc_xyz")
def test_replace_response(memory_jobstore): from jobflow import Flow, Job, Response 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(): job = Job(add, function_args=(1, )) flow = Flow([job], output=job.output) return Response(replace=flow) def replace_flow_multioutput(): job1 = Job(add, function_args=(1, )) job2 = Job(add, function_args=(job1.output, )) flow = Flow([job1, job2], output={"1": job1.output, "2": job2.output}) return Response(replace=flow) def replace_list_flow(): job1 = Job(add, function_args=(1, )) job2 = Job(add, function_args=(job1.output, )) job3 = Job(add, function_args=(5, )) job4 = Job(add, function_args=(job3.output, )) flow1 = Flow([job1, job2], output={"1": job1.output, "2": job2.output}) flow2 = Flow([job3, job4], output={"3": job3.output, "4": job4.output}) return Response(replace=[flow1, flow2]) # replace with job metadata = {"hi": "I exist"} test_job = Job(replace_job, metadata=metadata, output_schema="123") response = test_job.run(memory_jobstore) assert isinstance(response.replace, Flow) assert len(response.replace.jobs) == 1 assert response.replace.jobs[0].index == 2 assert response.replace.jobs[0].uuid == test_job.uuid assert response.replace.jobs[0].metadata == metadata assert response.replace.jobs[0].output_schema == "123" assert test_job.hosts == [] assert response.replace.hosts == [] assert response.replace.jobs[0].hosts == [response.replace.uuid] # replace with list of job test_job = Job(replace_list_job, metadata=metadata, output_schema="123") response = test_job.run(memory_jobstore) assert isinstance(response.replace, Flow) assert response.replace.jobs[-1].function == add assert len(response.replace.jobs) == 2 # currently output schema and metadata ignored in this case for j in response.replace.jobs: assert j.hosts == [response.replace.uuid] # replace with flow with outputs test_job = Job(replace_flow, metadata=metadata, output_schema="123") # wrap the job in a Flow to check hosts test_flow = Flow([test_job]) response = test_job.run(memory_jobstore) assert isinstance(response.replace, Flow) assert response.replace.jobs[-1].index == 2 assert response.replace.jobs[-1].uuid == test_job.uuid assert response.replace.jobs[-1].metadata == metadata assert response.replace.jobs[-1].output_schema == "123" assert response.replace.output is not None for j in response.replace.jobs: assert j.hosts == [response.replace.uuid, test_flow.uuid] # replace with flow with multi outputs test_job = Job(replace_flow_multioutput, metadata=metadata, output_schema="123") response = test_job.run(memory_jobstore) assert isinstance(response.replace, Flow) assert response.replace.jobs[-1].index == 2 assert response.replace.jobs[-1].uuid == test_job.uuid assert response.replace.jobs[-1].metadata == metadata assert response.replace.jobs[-1].output_schema == "123" assert response.replace.output is not None # replace with list of flow test_job = Job(replace_list_flow, metadata=metadata, output_schema="123") response = test_job.run(memory_jobstore) assert isinstance(response.replace, Flow) assert isinstance(response.replace.jobs[-1], Flow) assert len(response.replace.jobs) == 2 for f in response.replace.jobs: for j in f.jobs: assert j.hosts == [f.uuid, response.replace.uuid]
def add_schema_replace(a, b): new_job = Job(add, function_args=(a, b), output_schema=AddSchema) return Response(replace=new_job)
def replace_flow_multioutput(): job1 = Job(add, function_args=(1, )) job2 = Job(add, function_args=(job1.output, )) flow = Flow([job1, job2], output={"1": job1.output, "2": job2.output}) return Response(replace=flow)
def test_update_maker_kwargs(): from dataclasses import dataclass from jobflow import Job, Maker, job @dataclass class AddMaker(Maker): name: str = "add" b: int = 2 @job def make(self, a): return a + self.b @job def outer_job(maker1, maker2=None): pass # test no maker test_job = Job(add, function_args=(1, ), function_kwargs={"b": 2}) test_job.update_maker_kwargs({"b": 5}) assert test_job.function_kwargs["b"] == 2 add_maker = AddMaker(b=3) # test no filter test_job = add_maker.make(5) test_job.update_maker_kwargs({"b": 10}) assert test_job.maker.b == 10 # test bad kwarg test_job = add_maker.make(5) with pytest.raises(TypeError): test_job.update_maker_kwargs({"c": 10}) # test name filter test_job = add_maker.make(5) test_job.update_maker_kwargs({"b": 10}, name_filter="add") assert test_job.maker.b == 10 test_job = add_maker.make(5) test_job.update_maker_kwargs({"b": 10}, name_filter="div") assert test_job.maker.b == 3 # test class filter test_job = add_maker.make(5) test_job.update_maker_kwargs({"b": 10}, class_filter=AddMaker) assert test_job.maker.b == 10 test_job = add_maker.make(5) test_job.update_maker_kwargs({"b": 10}, class_filter=list) assert test_job.maker.b == 3 # test class filter with instance test_job = add_maker.make(5) test_job.update_maker_kwargs({"b": 10}, class_filter=add_maker) assert test_job.maker.b == 10 # test dict mod test_job = add_maker.make(5) test_job.update_maker_kwargs({"_inc": {"b": 10}}, dict_mod=True) assert test_job.maker.b == 13 # test makers in job args and kwargs maker1 = AddMaker(b=1) maker2 = AddMaker(b=2) test_job = outer_job(maker1, maker2=maker2) test_job.update_maker_kwargs({"_inc": {"b": 10}}, dict_mod=True) assert test_job.function_args[0].b == 11 assert test_job.function_kwargs["maker2"].b == 12 # test nested=False does not impact the jobs test_job.update_maker_kwargs({"_inc": { "b": 20 }}, dict_mod=True, nested=False) assert test_job.function_args[0].b == 11 assert test_job.function_kwargs["maker2"].b == 12
def test_update_kwargs(): from jobflow import Job # test no filter test_job = Job(add, function_args=(1, ), function_kwargs={"b": 2}) test_job.update_kwargs({"b": 5}) assert test_job.function_kwargs["b"] == 5 # test name filter test_job = Job(add, function_args=(1, ), function_kwargs={"b": 2}) test_job.update_kwargs({"b": 5}, name_filter="add") assert test_job.function_kwargs["b"] == 5 test_job = Job(add, function_args=(1, ), function_kwargs={"b": 2}) test_job.update_kwargs({"b": 5}, name_filter="div") assert test_job.function_kwargs["b"] == 2 # test function filter test_job = Job(add, function_args=(1, ), function_kwargs={"b": 2}) test_job.update_kwargs({"b": 5}, function_filter=add) assert test_job.function_kwargs["b"] == 5 test_job = Job(add, function_args=(1, ), function_kwargs={"b": 2}) test_job.update_kwargs({"b": 5}, function_filter=list) assert test_job.function_kwargs["b"] == 2 # test dict mod test_job = Job(add, function_args=(1, ), function_kwargs={"b": 2}) test_job.update_kwargs({"_inc": {"b": 5}}, dict_mod=True) assert test_job.function_kwargs["b"] == 7
def test_job_run(capsys, memory_jobstore, memory_data_jobstore): from jobflow.core.job import Job, Response # test basic run test_job = Job(print, function_args=("I am a job", )) response = test_job.run(memory_jobstore) assert capsys.readouterr().out == "I am a job\n" assert isinstance(response, Response) # test run with outputs test_job = Job(add, function_args=(1, ), function_kwargs={"b": 2}) response = test_job.run(memory_jobstore) assert isinstance(response, Response) assert response.output == 3 # test run with input references test_job = Job(add, function_args=(test_job.output, )) response = test_job.run(memory_jobstore) assert isinstance(response, Response) assert response.output == 8 def add_response(a, b): return Response(output=a + b, stop_children=True) test_job = Job(add_response, function_args=(1, 2)) response = test_job.run(memory_jobstore) assert isinstance(response, Response) assert response.output == 3 assert response.stop_children # test run with outputs and data store test_job = Job(add, function_args=(1, ), function_kwargs={"b": 2}, data=True) response = test_job.run(memory_data_jobstore) assert isinstance(response, Response) assert response.output == 3 # check output was not stored in the docs store result = memory_data_jobstore.query_one({"uuid": test_job.uuid}) assert isinstance(result["output"], dict) assert "blob_uuid" in result["output"] # check the output can be resolved result = memory_data_jobstore.query_one({"uuid": test_job.uuid}, load=True) assert result["output"] == 3 # test non MSONable output test_job = Job(bad_output) with pytest.raises(RuntimeError): test_job.run(memory_jobstore)
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
def replace_job(): job = Job(add, function_args=(1, )) return Response(replace=job)
def test_job_init(): from jobflow.core.job import Job # test basic init test_job = Job(function=add, function_args=("I am a job", )) assert test_job assert test_job.function == add assert test_job.name == "add" assert test_job.function_args == ("I am a job", ) assert test_job.function_kwargs == {} assert test_job.uuid is not None assert test_job.output.uuid == test_job.uuid # test init no args test_job = Job(function=add, function_args=tuple()) assert test_job assert test_job.function == add assert test_job.name == "add" assert test_job.function_args == tuple() assert test_job.function_kwargs == {} assert test_job.uuid is not None assert test_job.output.uuid == test_job.uuid # test job as another job as input with pytest.warns(UserWarning): Job(function=add, function_args=(test_job, )) # test init with kwargs test_job = Job(function=add, function_args=(1, ), function_kwargs={"b": 2}) assert test_job assert test_job.function == add assert test_job.name == "add" assert test_job.function_args == (1, ) assert test_job.function_kwargs == {"b": 2} assert test_job.uuid is not None assert test_job.uuid == test_job.output.value.uuid # test init with stores test_job = Job( function=add, function_args=(1, ), function_kwargs={"b": 2}, data="output", graphs="graph", ) assert test_job assert test_job.uuid == test_job.output.value.uuid assert test_job._kwargs == {"data": "output", "graphs": "graph"} # check giving True for multiple stores fails with pytest.raises(ValueError): Job(function=add, function_args=(1, ), data=True, graphs=True) # test changing job name (test needed due to setattr override test_job = Job(function=add, function_args=("I am a job", )) test_job.name = "abc" assert test_job.name == "abc" test_job.name += "xyz" assert test_job.name == "abcxyz"
def detour_job(): job = Job(add, function_args=(1, )) return Response(detour=job)
def addition_job(): job = Job(add, function_args=(1, )) return Response(addition=job)
def replace_flow(): from jobflow import Flow job = Job(add, function_args=(1, )) flow = Flow([job], output=job.output) return Response(replace=flow)
def replace_list_job(): job1 = Job(add, function_args=(1, )) job2 = Job(add, function_args=(job1.output, )) return Response(replace=[job1, job2])