diff --git a/src/jobflow/core/job.py b/src/jobflow/core/job.py index d00501b2..18c11337 100644 --- a/src/jobflow/core/job.py +++ b/src/jobflow/core/job.py @@ -28,6 +28,23 @@ logger = logging.getLogger(__name__) +def _replace_job_or_flow_with_output(value): + """Replace Jobs and Flows in nested containers with their outputs.""" + from jobflow.core.flow import Flow + + 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 { + key: _replace_job_or_flow_with_output(item) for key, item in value.items() + } + return value + + @dataclass class JobConfig(MSONable): """ @@ -213,6 +230,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 ) @@ -604,7 +625,7 @@ def run(self, store: jobflow.JobStore, job_dir: Path = None) -> Response: from datetime import datetime from jobflow import CURRENT_JOB - from jobflow.core.flow import get_flow + from jobflow.core.flow import flow_build_context, get_flow from jobflow.core.schemas import JobStoreDocument index_str = f", {self.index}" if self.index != 1 else "" @@ -626,12 +647,17 @@ def run(self, store: jobflow.JobStore, job_dir: Path = None) -> Response: if bound is not None and not isinstance(bound, types.ModuleType): function = types.MethodType(function, bound) - response = function(*self.function_args, **self.function_kwargs) + dynamic_children = [] + with flow_build_context(dynamic_children): + response = function(*self.function_args, **self.function_kwargs) response = Response.from_job_returns( response, self.output_schema, job_dir=job_dir ) if response.replace is not None: + response.replace = _expand_dynamic_dependencies( + response.replace, dynamic_children + ) response.replace = prepare_replace(response.replace, self) if response.addition is not None: @@ -1481,6 +1507,51 @@ def prepare_replace( return replace +def _expand_dynamic_dependencies(replace, dynamic_children): + """Add dynamically created dependencies required by a replacement.""" + from jobflow.core.flow import Flow + + candidates = [child for child in dynamic_children if child.host is None] + if not candidates: + return replace + + if isinstance(replace, dict): + explicit_jobs = list(replace.values()) + output = {key: child.output for key, child in replace.items()} + elif isinstance(replace, (list, tuple)): + explicit_jobs = list(replace) + output = type(replace)(child.output for child in replace) + else: + explicit_jobs = [replace] + output = replace.output + + def all_uuids(child): + if isinstance(child, Flow): + return set(child.all_uuids) + return {child.uuid} + + explicit_uuids = set().union(*(all_uuids(child) for child in explicit_jobs)) + required_uuids = set(explicit_uuids) + + while True: + selected = [ + child for child in candidates if all_uuids(child) & required_uuids + ] + expanded_uuids = set().union( + required_uuids, *(set(child.graph.nodes) for child in selected) + ) + if expanded_uuids == required_uuids: + break + required_uuids = expanded_uuids + + selected = [child for child in candidates if all_uuids(child) & required_uuids] + selected_uuids = set().union(*(all_uuids(child) for child in selected)) + if selected_uuids.issubset(explicit_uuids): + return replace + + return Flow(selected, output=output) + + def pass_manager_config( jobs: Job | jobflow.Flow | list[Job | jobflow.Flow], manager_config: dict[str, Any], diff --git a/tests/core/test_job.py b/tests/core/test_job.py index f9416221..c87f853b 100644 --- a/tests/core/test_job.py +++ b/tests/core/test_job.py @@ -230,6 +230,38 @@ def replace_list_flow(): # currently output schema and metadata ignored for all but the last `store_inputs` +def test_replace_response_collects_dynamic_dependencies(): + """Test replacement jobs include dynamically created dependencies.""" + from jobflow import Response, job, run_locally + + @job + def add_job(a, b): + return a + b + + @job + def multiply_job(value, factor): + return value * factor + + @job + def make_dynamic_workflow(): + add_job(100, 200) + first = add_job(1, 2) + final = multiply_job(first, 3) + return Response(replace=final) + + dynamic_job = make_dynamic_workflow() + responses = run_locally(dynamic_job, ensure_success=True) + + outputs = [ + response.output + for index_to_response in responses.values() + for response in index_to_response.values() + if isinstance(response.output, int) + ] + assert sorted(outputs) == [3, 9] + assert len(responses) == 3 + + def test_job_config(memory_jobstore): from jobflow import ( CURRENT_JOB,