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_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 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]