Track subflow results along with task results (#18338)

Author: jakekaplan

src/integrations/prefect-dask/prefect_dask/task_runners.py

@@ -92,7 +92,7 @@ def count_to(highest_number):
 import distributed.deploy.cluster
 from typing_extensions import ParamSpec
 
-from prefect.client.schemas.objects import State, TaskRunInput
+from prefect.client.schemas.objects import RunInput, State
 from prefect.futures import PrefectFuture, PrefectFutureList, PrefectWrappedFuture
 from prefect.logging.loggers import get_logger
 from prefect.task_runners import TaskRunner
@@ -405,7 +405,7 @@ def submit(
         task: "Task[P, Coroutine[Any, Any, R]]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectDaskFuture[R]] | None = None,
-        dependencies: dict[str, Set[TaskRunInput]] | None = None,
+        dependencies: dict[str, Set[RunInput]] | None = None,
     ) -> PrefectDaskFuture[R]: ...
 
     @overload
@@ -414,15 +414,15 @@ def submit(
         task: "Task[Any, R]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectDaskFuture[R]] | None = None,
-        dependencies: dict[str, Set[TaskRunInput]] | None = None,
+        dependencies: dict[str, Set[RunInput]] | None = None,
     ) -> PrefectDaskFuture[R]: ...
 
     def submit(
         self,
         task: "Union[Task[P, R], Task[P, Coroutine[Any, Any, R]]]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectDaskFuture[R]] | None = None,
-        dependencies: dict[str, Set[TaskRunInput]] | None = None,
+        dependencies: dict[str, Set[RunInput]] | None = None,
     ) -> PrefectDaskFuture[R]:
         if not self._started:
             raise RuntimeError(

src/integrations/prefect-ray/prefect_ray/task_runners.py

@@ -87,7 +87,7 @@ def count_to(highest_number):
 
 from typing_extensions import ParamSpec, Self
 
-from prefect.client.schemas.objects import TaskRunInput
+from prefect.client.schemas.objects import RunInput
 from prefect.context import serialize_context
 from prefect.futures import PrefectFuture, PrefectFutureList, PrefectWrappedFuture
 from prefect.logging.loggers import get_logger
@@ -227,7 +227,7 @@ def submit(
         task: "Task[P, Coroutine[Any, Any, R]]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ) -> PrefectRayFuture[R]: ...
 
     @overload
@@ -236,15 +236,15 @@ def submit(
         task: "Task[P, R]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ) -> PrefectRayFuture[R]: ...
 
     def submit(
         self,
         task: Task[P, R],
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ):
         if not self._started:
             raise RuntimeError(
@@ -342,7 +342,7 @@ def _run_prefect_task(
         context: dict[str, Any],
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
-        dependencies: dict[str, set[TaskRunInput]] | None = None,
+        dependencies: dict[str, set[RunInput]] | None = None,
     ) -> Any:
         """Resolves Ray futures before calling the actual Prefect task function.
 

src/prefect/client/orchestration/__init__.py

@@ -101,10 +101,11 @@
     WorkQueueFilterName,
 )
 from prefect.client.schemas.objects import (
-    Constant,
+    TaskRunResult,
+    FlowRunResult,
     Parameter,
+    Constant,
     TaskRunPolicy,
-    TaskRunResult,
     WorkQueue,
     WorkQueueStatusDetail,
 )
@@ -778,13 +779,7 @@ async def create_task_run(
         task_inputs: Optional[
             dict[
                 str,
-                list[
-                    Union[
-                        TaskRunResult,
-                        Parameter,
-                        Constant,
-                    ]
-                ],
+                list[Union[TaskRunResult, FlowRunResult, Parameter, Constant]],
             ]
         ] = None,
     ) -> TaskRun:
@@ -1440,6 +1435,7 @@ def create_task_run(
                 list[
                     Union[
                         TaskRunResult,
+                        FlowRunResult,
                         Parameter,
                         Constant,
                     ]

src/prefect/client/schemas/__init__.py

@@ -11,11 +11,12 @@
         BlockType,
         FlowRun,
         FlowRunPolicy,
+        FlowRunResult,
         State,
         StateDetails,
         StateType,
         TaskRun,
-        TaskRunInput,
+        RunInput,
         TaskRunPolicy,
         TaskRunResult,
         Workspace,
@@ -36,6 +37,7 @@
     "DEFAULT_BLOCK_SCHEMA_VERSION": (__package__, ".objects"),
     "FlowRun": (__package__, ".objects"),
     "FlowRunPolicy": (__package__, ".objects"),
+    "FlowRunResult": (__package__, ".objects"),
     "OrchestrationResult": (__package__, ".responses"),
     "SetStateStatus": (__package__, ".responses"),
     "State": (__package__, ".objects"),
@@ -60,6 +62,7 @@
     "DEFAULT_BLOCK_SCHEMA_VERSION",
     "FlowRun",
     "FlowRunPolicy",
+    "FlowRunResult",
     "OrchestrationResult",
     "SetStateStatus",
     "State",
@@ -70,7 +73,7 @@
     "StateRejectDetails",
     "StateType",
     "TaskRun",
-    "TaskRunInput",
+    "RunInput",
     "TaskRunPolicy",
     "TaskRunResult",
     "Workspace",

src/prefect/client/schemas/actions.py

@@ -443,6 +443,7 @@ class TaskRunCreate(ActionBaseModel):
         list[
             Union[
                 objects.TaskRunResult,
+                objects.FlowRunResult,
                 objects.Parameter,
                 objects.Constant,
             ]

src/prefect/client/schemas/objects.py

@@ -81,6 +81,11 @@
 DEFAULT_AGENT_WORK_POOL_NAME: Literal["default-agent-pool"] = "default-agent-pool"
 
 
+class RunType(AutoEnum):
+    FLOW_RUN = "flow_run"
+    TASK_RUN = "task_run"
+
+
 class StateType(AutoEnum):
     """Enumeration of state types."""
 
@@ -164,7 +169,6 @@ class ConcurrencyLimitConfig(PrefectBaseModel):
 class StateDetails(PrefectBaseModel):
     flow_run_id: Optional[UUID] = None
     task_run_id: Optional[UUID] = None
-    # for task runs that represent subflows, the subflow's run ID
     child_flow_run_id: Optional[UUID] = None
     scheduled_time: Optional[DateTime] = None
     cache_key: Optional[str] = None
@@ -182,6 +186,16 @@ class StateDetails(PrefectBaseModel):
     # Captures the trace_id and span_id of the span where this state was created
     traceparent: Optional[str] = None
 
+    def to_run_result(
+        self, run_type: RunType
+    ) -> Optional[Union[FlowRunResult, TaskRunResult]]:
+        if run_type == run_type.FLOW_RUN and self.flow_run_id:
+            return FlowRunResult(id=self.flow_run_id)
+        elif run_type == run_type.TASK_RUN and self.task_run_id:
+            return TaskRunResult(id=self.task_run_id)
+        else:
+            return None
+
 
 def data_discriminator(x: Any) -> str:
     if isinstance(x, dict) and "storage_key" in x:
@@ -734,7 +748,7 @@ def validate_jitter_factor(cls, v: Optional[float]) -> Optional[float]:
         return validate_not_negative(v)
 
 
-class TaskRunInput(PrefectBaseModel):
+class RunInput(PrefectBaseModel):
     """
     Base class for classes that represent inputs to task runs, which
     could include, constants, parameters, or other task runs.
@@ -747,21 +761,26 @@ class TaskRunInput(PrefectBaseModel):
         input_type: str
 
 
-class TaskRunResult(TaskRunInput):
+class TaskRunResult(RunInput):
     """Represents a task run result input to another task run."""
 
     input_type: Literal["task_run"] = "task_run"
     id: UUID
 
 
-class Parameter(TaskRunInput):
+class FlowRunResult(RunInput):
+    input_type: Literal["flow_run"] = "flow_run"
+    id: UUID
+
+
+class Parameter(RunInput):
     """Represents a parameter input to a task run."""
 
     input_type: Literal["parameter"] = "parameter"
     name: str
 
 
-class Constant(TaskRunInput):
+class Constant(RunInput):
     """Represents constant input value to a task run."""
 
     input_type: Literal["constant"] = "constant"
@@ -811,7 +830,9 @@ class TaskRun(TimeSeriesBaseModel, ObjectBaseModel):
     state_id: Optional[UUID] = Field(
         default=None, description="The id of the current task run state."
     )
-    task_inputs: dict[str, list[Union[TaskRunResult, Parameter, Constant]]] = Field(
+    task_inputs: dict[
+        str, list[Union[TaskRunResult, FlowRunResult, Parameter, Constant]]
+    ] = Field(
         default_factory=dict,
         description=(
             "Tracks the source of inputs to a task run. Used for internal bookkeeping. "

src/prefect/context.py

@@ -34,6 +34,7 @@
 from prefect.assets import Asset
 from prefect.client.orchestration import PrefectClient, SyncPrefectClient, get_client
 from prefect.client.schemas import FlowRun, TaskRun
+from prefect.client.schemas.objects import RunType
 from prefect.events.worker import EventsWorker
 from prefect.exceptions import MissingContextError
 from prefect.results import (
@@ -363,7 +364,7 @@ class EngineContext(RunContext):
         flow: The flow instance associated with the run
         flow_run: The API metadata for the flow run
         task_runner: The task runner instance being used for the flow run
-        task_run_results: A mapping of result ids to task run states for this flow run
+        run_results: A mapping of result ids to run states for this flow run
         log_prints: Whether to log print statements from the flow run
         parameters: The parameters passed to the flow run
         detached: Flag indicating if context has been serialized and sent to remote infrastructure
@@ -394,9 +395,10 @@ class EngineContext(RunContext):
     # Counter for flow pauses
     observed_flow_pauses: dict[str, int] = Field(default_factory=dict)
 
-    # Tracking for result from task runs in this flow run for dependency tracking
-    # Holds the ID of the object returned by the task run and task run state
-    task_run_results: dict[int, State] = Field(default_factory=dict)
+    # Tracking for result from task runs and sub flows in this flow run for
+    # dependency tracking. Holds the ID of the object returned by
+    # the run and state
+    run_results: dict[int, tuple[State, RunType]] = Field(default_factory=dict)
 
     # Tracking information needed to track asset linage between
     # tasks and materialization

src/prefect/flow_engine.py

@@ -105,7 +105,7 @@
 from prefect.utilities.collections import visit_collection
 from prefect.utilities.engine import (
     capture_sigterm,
-    link_state_to_result,
+    link_state_to_flow_run_result,
     propose_state,
     propose_state_sync,
     resolve_to_final_result,
@@ -338,6 +338,7 @@ def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]":
             self._return_value, State
         ):
             _result = self._return_value
+            link_state_to_flow_run_result(self.state, _result)
 
             if asyncio.iscoroutine(_result):
                 # getting the value for a BaseResult may return an awaitable
@@ -373,6 +374,7 @@ def handle_success(self, result: R) -> R:
         self.set_state(terminal_state)
         self._return_value = resolved_result
 
+        link_state_to_flow_run_result(terminal_state, resolved_result)
         self._telemetry.end_span_on_success()
 
         return result
@@ -903,6 +905,7 @@ async def result(self, raise_on_failure: bool = True) -> "Union[R, State, None]"
             self._return_value, State
         ):
             _result = self._return_value
+            link_state_to_flow_run_result(self.state, _result)
 
             if asyncio.iscoroutine(_result):
                 # getting the value for a BaseResult may return an awaitable
@@ -1426,7 +1429,7 @@ def run_generator_flow_sync(
                     while True:
                         gen_result = next(gen)
                         # link the current state to the result for dependency tracking
-                        link_state_to_result(engine.state, gen_result)
+                        link_state_to_flow_run_result(engine.state, gen_result)
                         yield gen_result
                 except StopIteration as exc:
                     engine.handle_success(exc.value)
@@ -1468,7 +1471,7 @@ async def run_generator_flow_async(
                         # can't use anext in Python < 3.10
                         gen_result = await gen.__anext__()
                         # link the current state to the result for dependency tracking
-                        link_state_to_result(engine.state, gen_result)
+                        link_state_to_flow_run_result(engine.state, gen_result)
                         yield gen_result
                 except (StopAsyncIteration, GeneratorExit) as exc:
                     await engine.handle_success(None)

src/prefect/runner/submit.py

@@ -17,7 +17,13 @@
     FlowRunFilterParentFlowRunId,
     TaskRunFilter,
 )
-from prefect.client.schemas.objects import Constant, FlowRun, Parameter, TaskRunResult
+from prefect.client.schemas.objects import (
+    Constant,
+    FlowRun,
+    FlowRunResult,
+    Parameter,
+    TaskRunResult,
+)
 from prefect.context import FlowRunContext
 from prefect.flows import Flow
 from prefect.logging import get_logger
@@ -66,9 +72,9 @@ async def _submit_flow_to_runner(
 
         parent_flow_run_context = FlowRunContext.get()
 
-        task_inputs: dict[str, list[TaskRunResult | Parameter | Constant]] = {
-            k: list(await collect_task_run_inputs(v)) for k, v in parameters.items()
-        }
+        task_inputs: dict[
+            str, list[Union[TaskRunResult, FlowRunResult, Parameter, Constant]]
+        ] = {k: list(await collect_task_run_inputs(v)) for k, v in parameters.items()}
         parameters = await resolve_inputs(parameters)
         dummy_task = Task(name=flow.name, fn=flow.fn, version=flow.version)
         parent_task_run = await client.create_task_run(

src/prefect/server/database/orm_models.py

@@ -643,7 +643,10 @@ def __table_args__(cls) -> Iterable[sa.Index]:
 
 
 _TaskInput = Union[
-    schemas.core.TaskRunResult, schemas.core.Parameter, schemas.core.Constant
+    schemas.core.TaskRunResult,
+    schemas.core.FlowRunResult,
+    schemas.core.Parameter,
+    schemas.core.Constant,
 ]
 _TaskInputs = dict[str, list[_TaskInput]]
 

src/prefect/server/schemas/actions.py

@@ -501,6 +501,7 @@ class TaskRunCreate(ActionBaseModel):
         List[
             Union[
                 schemas.core.TaskRunResult,
+                schemas.core.FlowRunResult,
                 schemas.core.Parameter,
                 schemas.core.Constant,
             ]