diff --git a/pyproject.toml b/pyproject.toml index 30fdcd9c..f5ad9e63 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,7 +79,7 @@ jobflow = ["py.typed"] [tool.versioningit.vcs] method = "git" -default-tag = "0.0.1" +default-tag = "0.3.1" [tool.flake8] max-line-length = 88 diff --git a/src/jobflow/core/flow.py b/src/jobflow/core/flow.py index 17fb2fd7..c6c1e472 100644 --- a/src/jobflow/core/flow.py +++ b/src/jobflow/core/flow.py @@ -13,7 +13,12 @@ import jobflow from jobflow.core.reference import find_and_get_references -from jobflow.utils import ValueEnum, contains_flow_or_job, suid +from jobflow.utils import ( + ValueEnum, + contains_flow_or_job, + replace_job_or_flow_with_output, + suid, +) if TYPE_CHECKING: from collections.abc import Iterator, Sequence @@ -965,15 +970,14 @@ def __init__(self, fn, *args, **kwargs): if isinstance(output, (jobflow.Job, jobflow.Flow)): warnings.warn( - f"@flow decorated function '{name}' contains a Flow or" - f"Job as an output. Usually the output should be the output of" - f"a Job or another Flow (e.g. job.output). Replacing the" - f"output of the @flow with the output of the Flow/Job." - f"If this message is unexpected then double check the outputs" - f"of your @flow decorated function.", + f"@flow decorated function '{name}' contains a Flow or Job as an " + "output. Usually the output should be the output of a Job or another " + "Flow (e.g. job.output). Replacing the output of the @flow with the " + "output of the Flow/Job. If this message is unexpected then double " + "check the outputs of your @flow decorated function.", stacklevel=2, ) - output = output.output + output = replace_job_or_flow_with_output(output) super().__init__(name=name, jobs=children_list, output=output) diff --git a/src/jobflow/core/job.py b/src/jobflow/core/job.py index d00501b2..47e8d6dd 100644 --- a/src/jobflow/core/job.py +++ b/src/jobflow/core/job.py @@ -13,6 +13,7 @@ from jobflow.core.flow import _current_flow_context from jobflow.core.reference import OnMissing, OutputReference +from jobflow.utils.find import replace_job_or_flow_with_output from jobflow.utils.uid import suid if typing.TYPE_CHECKING: @@ -213,6 +214,10 @@ def get_job(*args, **kwargs) -> Job: f = met args = args[1:] + if _current_flow_context.get() is not None: + args = replace_job_or_flow_with_output(args) + kwargs = replace_job_or_flow_with_output(kwargs) + return Job( function=f, function_args=args, function_kwargs=kwargs, **job_kwargs ) diff --git a/src/jobflow/utils/__init__.py b/src/jobflow/utils/__init__.py index ed6e1648..906c8121 100644 --- a/src/jobflow/utils/__init__.py +++ b/src/jobflow/utils/__init__.py @@ -5,6 +5,7 @@ contains_flow_or_job, find_key, find_key_value, + replace_job_or_flow_with_output, update_in_dictionary, ) from jobflow.utils.log import initialize_logger diff --git a/src/jobflow/utils/find.py b/src/jobflow/utils/find.py index 5dfc98e9..1c5c0240 100644 --- a/src/jobflow/utils/find.py +++ b/src/jobflow/utils/find.py @@ -11,6 +11,24 @@ from monty.json import MSONable +def replace_job_or_flow_with_output(value): + """Replace Jobs and Flows in a value with their outputs.""" + from jobflow import Flow, Job + + if isinstance(value, (Job, Flow)): + return value.output + if isinstance(value, list): + return [replace_job_or_flow_with_output(item) for item in value] + if isinstance(value, tuple): + return tuple(replace_job_or_flow_with_output(item) for item in value) + if isinstance(value, dict): + return { + replace_job_or_flow_with_output(key): replace_job_or_flow_with_output(item) + for key, item in value.items() + } + return value + + def find_key( d: dict[Hashable, Any] | list[Any], key: Hashable | type[MSONable], diff --git a/tests/core/test_flow_decorator.py b/tests/core/test_flow_decorator.py index aa6b755d..155f98c8 100644 --- a/tests/core/test_flow_decorator.py +++ b/tests/core/test_flow_decorator.py @@ -159,6 +159,67 @@ def my_flow(a, b): assert result[flow1.output.uuid][1].output == 7 +def test_flow_resolves_minimal_job_input_example(): + """Test the minimal Job-as-input example from issue #884.""" + from jobflow import flow, job + from jobflow.managers.local import run_locally + + @job + def mult(a, b): + return a * b + + @flow + def workflow(a, b): + sum_result = add(a, b) + return mult(a, sum_result) + + flow1 = workflow(1, 2) + result = run_locally(flow1, ensure_success=True) + + assert result[flow1.output.uuid][1].output == 3 + + +def test_flow_resolves_job_inputs_to_outputs(): + """Test top-level container Jobs resolve to outputs.""" + from jobflow import flow, job + + @job + def consume(values): + return values + + @flow + def my_flow(): + source = add(1, 2) + list_job = consume([source]) + tuple_job = consume((source,)) + dict_job = consume({"source": source}) + return [ + list_job.output, + tuple_job.output, + dict_job.output, + ] + + flow1 = my_flow() + source = flow1.jobs[0] + + assert flow1.jobs[1].function_args == ([source.output],) + assert flow1.jobs[2].function_args == ((source.output,),) + assert flow1.jobs[3].function_args == ({"source": source.output},) + + +def test_flow_normalizes_top_level_job_outputs(): + """Test Jobs at the top level of a decorated flow output are normalized.""" + from jobflow import flow + + @flow + def my_flow(): + return {"job": add(1, 2)} + + flow1 = my_flow() + + assert flow1.output == {"job": flow1.jobs[0].output} + + def test_flow_returns_list(): """Test that a flow that returns a list of OutputReferences can be created and run.""" diff --git a/tests/core/test_job.py b/tests/core/test_job.py index f9416221..9b43a61d 100644 --- a/tests/core/test_job.py +++ b/tests/core/test_job.py @@ -83,7 +83,6 @@ def test_job_run(capsys, memory_jobstore, memory_data_jobstore): 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) @@ -1155,7 +1154,7 @@ def make(self): test_job = maker.make() test_job.update_metadata( {"v": 5}, - callback_filter=lambda job: (job.maker is not None and job.maker.value == 42), + callback_filter=lambda job: job.maker is not None and job.maker.value == 42, ) assert test_job.metadata["v"] == 5