Repository: PrefectHQ/prefect

Issue:
Add `wait_for` for flows
## User story
As a Prefect 2 user, I want to use a flow as a subflow, with a dependency on an upstream flow or a task, even though they do not do not exchange data. I want to do this with the `wait_for` keyword argument, just as I can for tasks.

## Opened from the [Prefect Public Slack Community](https://prefect.io/slack)

**digangidomenico**: Hi all,
I am experimenting with tasks and subflows in prefect 2.0. While in tasks I can use the wait_for keyword to enforce task dependency, the same keyword does not work for subflows. I did not found much on the subflows docs, but it seems that in prefect 1.0 dependency between subflows was possible. Is such a feature available also in prefect 2 ?
**michael054**: This just hasn’t been added yet
**michael054**: However, since a subflow will block until it is complete, the `wait_for` kwarg is less useful
**michael054**: You might as well just wait for all the tasks to finish using `.wait()` before calling the subflow
**digangidomenico**: Thanks for the quick reply <@U01CEUST9B5>.
If I understand correctly, using wait_for in a task also injects a "logical dependency" not only a temporal ordering. For example, in the following

```
@task
def task_a():
    pass

@task
def task_b():
    pass

@flow
def my_flow():
    a = task_a()
    b = task_b(wait_for=[a])
```
if task_a  fails, task_b does not execute, right?
I am looking for a similar dependency between subflows or a subflow depending on a task
**michael054**: Ah yes, I see your point. You can pass the task futures to the subflow directly but when you do not want to add a data dependency this introduces a problem
**michael054**: You can add a `wait_for` argument to your subflow that takes a list of upstream tasks and just not use the data if you don’t need it. We’ll likely support this as a first-class feature in the future. I’ll open an issue for tracking.
**digangidomenico**: I think this will work. Many thanks for your help <@U01CEUST9B5>
**michael054**: <@ULVA73B9P> open “Add `wait_for` for flows”

Original thread can be found [here](https://prefect-community.slack.com/archives/CL09KU1K7/p1650646059038279?thread_ts=1650646059.038279&cid=CL09KU1K7).

## References
[Frequently Asked Questions - Enforcing ordering](https://docs.prefect.io/faq/?h=wait_for#how-do-i-enforce-ordering-between-tasks-that-dont-share-data)


Discussion:
+1 - would be great to have that this - this would unify the api a lot

Resolving patch:
diff --git a/src/prefect/engine.py b/src/prefect/engine.py
--- a/src/prefect/engine.py
+++ b/src/prefect/engine.py
@@ -99,7 +99,10 @@
 
 
 def enter_flow_run_engine_from_flow_call(
-    flow: Flow, parameters: Dict[str, Any], return_type: EngineReturnType
+    flow: Flow,
+    parameters: Dict[str, Any],
+    wait_for: Optional[Iterable[PrefectFuture]],
+    return_type: EngineReturnType,
 ) -> Union[State, Awaitable[State]]:
     """
     Sync entrypoint for flow calls.
@@ -127,10 +130,14 @@ def enter_flow_run_engine_from_flow_call(
     parent_flow_run_context = FlowRunContext.get()
     is_subflow_run = parent_flow_run_context is not None
 
+    if wait_for is not None and not is_subflow_run:
+        raise ValueError("Only flows run as subflows can wait for dependencies.")
+
     begin_run = partial(
         create_and_begin_subflow_run if is_subflow_run else create_then_begin_flow_run,
         flow=flow,
         parameters=parameters,
+        wait_for=wait_for,
         return_type=return_type,
         client=parent_flow_run_context.client if is_subflow_run else None,
     )
@@ -178,6 +185,7 @@ def enter_flow_run_engine_from_subprocess(flow_run_id: UUID) -> State:
 async def create_then_begin_flow_run(
     flow: Flow,
     parameters: Dict[str, Any],
+    wait_for: Optional[Iterable[PrefectFuture]],
     return_type: EngineReturnType,
     client: OrionClient,
 ) -> Any:
@@ -354,6 +362,7 @@ async def begin_flow_run(
             flow,
             flow_run=flow_run,
             parameters=parameters,
+            wait_for=None,
             client=client,
             partial_flow_run_context=flow_run_context,
             # Orchestration needs to be interruptible if it has a timeout
@@ -380,6 +389,7 @@ async def begin_flow_run(
 async def create_and_begin_subflow_run(
     flow: Flow,
     parameters: Dict[str, Any],
+    wait_for: Optional[Iterable[PrefectFuture]],
     return_type: EngineReturnType,
     client: OrionClient,
 ) -> Any:
@@ -400,6 +410,9 @@ async def create_and_begin_subflow_run(
     parent_logger.debug(f"Resolving inputs to {flow.name!r}")
     task_inputs = {k: await collect_task_run_inputs(v) for k, v in parameters.items()}
 
+    if wait_for:
+        task_inputs["wait_for"] = await collect_task_run_inputs(wait_for)
+
     rerunning = parent_flow_run_context.flow_run.run_count > 1
 
     # Generate a task in the parent flow run to represent the result of the subflow run
@@ -476,6 +489,7 @@ async def create_and_begin_subflow_run(
                 flow,
                 flow_run=flow_run,
                 parameters=parameters,
+                wait_for=wait_for,
                 # If the parent flow run has a timeout, then this one needs to be
                 # interruptible as well
                 interruptible=parent_flow_run_context.timeout_scope is not None,
@@ -512,6 +526,7 @@ async def orchestrate_flow_run(
     flow: Flow,
     flow_run: FlowRun,
     parameters: Dict[str, Any],
+    wait_for: Optional[Iterable[PrefectFuture]],
     interruptible: bool,
     client: OrionClient,
     partial_flow_run_context: PartialModel[FlowRunContext],
@@ -540,6 +555,21 @@ async def orchestrate_flow_run(
     )
     flow_run_context = None
 
+    try:
+        # Resolve futures in any non-data dependencies to ensure they are ready
+        if wait_for is not None:
+            await resolve_inputs(wait_for, return_data=False)
+    except UpstreamTaskError as upstream_exc:
+
+        return await propose_state(
+            client,
+            Pending(name="NotReady", message=str(upstream_exc)),
+            flow_run_id=flow_run.id,
+            # if orchestrating a run already in a pending state, force orchestration to
+            # update the state name
+            force=flow_run.state.is_pending(),
+        )
+
     state = await propose_state(client, Running(), flow_run_id=flow_run.id)
 
     while state.is_running():
diff --git a/src/prefect/flows.py b/src/prefect/flows.py
--- a/src/prefect/flows.py
+++ b/src/prefect/flows.py
@@ -39,6 +39,7 @@
     ParameterTypeError,
     UnspecifiedFlowError,
 )
+from prefect.futures import PrefectFuture
 from prefect.logging import get_logger
 from prefect.orion.schemas.core import raise_on_invalid_name
 from prefect.results import ResultSerializer, ResultStorage
@@ -148,7 +149,7 @@ def __init__(
         self.fn = fn
         self.isasync = is_async_fn(self.fn)
 
-        raise_for_reserved_arguments(self.fn, ["return_state"])
+        raise_for_reserved_arguments(self.fn, ["return_state", "wait_for"])
 
         # Version defaults to a hash of the function's file
         flow_file = inspect.getsourcefile(self.fn)
@@ -385,6 +386,7 @@ def __call__(
         self,
         *args: "P.args",
         return_state: bool = False,
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
         **kwargs: "P.kwargs",
     ):
         """
@@ -401,6 +403,7 @@ def __call__(
             *args: Arguments to run the flow with.
             return_state: Return a Prefect State containing the result of the
             flow run.
+            wait_for: Upstream task futures to wait for before starting the flow if called as a subflow
             **kwargs: Keyword arguments to run the flow with.
 
         Returns:
@@ -437,7 +440,10 @@ def __call__(
         return_type = "state" if return_state else "result"
 
         return enter_flow_run_engine_from_flow_call(
-            self, parameters, return_type=return_type
+            self,
+            parameters,
+            wait_for=wait_for,
+            return_type=return_type,
         )
 
     @overload
@@ -459,6 +465,7 @@ def _run(self: "Flow[P, T]", *args: P.args, **kwargs: P.kwargs) -> State[T]:
     def _run(
         self,
         *args: "P.args",
+        wait_for: Optional[Iterable[PrefectFuture]] = None,
         **kwargs: "P.kwargs",
     ):
         """
@@ -478,7 +485,10 @@ def _run(
         parameters = get_call_parameters(self.fn, args, kwargs)
 
         return enter_flow_run_engine_from_flow_call(
-            self, parameters, return_type="state"
+            self,
+            parameters,
+            wait_for=wait_for,
+            return_type="state",
         )
 
 


Test patch:
diff --git a/tests/test_engine.py b/tests/test_engine.py
--- a/tests/test_engine.py
+++ b/tests/test_engine.py
@@ -451,6 +451,7 @@ def foo():
                 flow=foo,
                 flow_run=flow_run,
                 parameters={},
+                wait_for=None,
                 client=orion_client,
                 interruptible=False,
                 partial_flow_run_context=partial_flow_run_context,
@@ -482,6 +483,7 @@ def foo():
                 flow=foo,
                 flow_run=flow_run,
                 parameters={},
+                wait_for=None,
                 client=orion_client,
                 interruptible=False,
                 partial_flow_run_context=partial_flow_run_context,
@@ -516,6 +518,7 @@ def flaky_function():
                 flow=flaky_function,
                 flow_run=flow_run,
                 parameters={},
+                wait_for=None,
                 client=orion_client,
                 interruptible=False,
                 partial_flow_run_context=partial_flow_run_context,
@@ -1156,6 +1159,7 @@ async def test_handles_bad_parameter_types(self, orion_client, parameterized_flo
         state = await create_then_begin_flow_run(
             flow=parameterized_flow,
             parameters={"dog": [1, 2], "cat": "not an int"},
+            wait_for=None,
             return_type="state",
             client=orion_client,
         )
@@ -1173,6 +1177,7 @@ async def test_handles_signature_mismatches(self, orion_client, parameterized_fl
         state = await create_then_begin_flow_run(
             flow=parameterized_flow,
             parameters={"puppy": "a string", "kitty": 42},
+            wait_for=None,
             return_type="state",
             client=orion_client,
         )
@@ -1198,6 +1203,7 @@ def flow_use_and_return_defaults(foo: str = "bar", bar: int = 1):
         state = await create_then_begin_flow_run(
             flow=flow_use_and_return_defaults,
             parameters={},
+            wait_for=None,
             return_type="state",
             client=orion_client,
         )
@@ -1218,6 +1224,7 @@ def raise_unspecified_exception(*args, **kwargs):
         state = await create_then_begin_flow_run(
             flow=parameterized_flow,
             parameters={"puppy": "a string", "kitty": 42},
+            wait_for=None,
             return_type="state",
             client=orion_client,
         )
@@ -1257,6 +1264,7 @@ async def test_handles_signature_mismatches(self, orion_client, parameterized_fl
         state = await create_then_begin_flow_run(
             flow=parameterized_flow,
             parameters={"puppy": "a string", "kitty": 42},
+            wait_for=None,
             return_type="state",
             client=orion_client,
         )
@@ -1283,6 +1291,7 @@ def raise_unspecified_exception(*args, **kwargs):
         state = await create_then_begin_flow_run(
             flow=parameterized_flow,
             parameters={"puppy": "a string", "kitty": 42},
+            wait_for=None,
             return_type="state",
             client=orion_client,
         )
@@ -1304,6 +1313,7 @@ async def test_handles_bad_parameter_types(
             state = await create_and_begin_subflow_run(
                 flow=parameterized_flow,
                 parameters={"dog": [1, 2], "cat": "not an int"},
+                wait_for=None,
                 return_type="state",
                 client=orion_client,
             )
@@ -1328,6 +1338,7 @@ async def test_handles_signature_mismatches(
             state = await create_and_begin_subflow_run(
                 flow=parameterized_flow,
                 parameters={"puppy": "a string", "kitty": 42},
+                wait_for=None,
                 return_type="state",
                 client=orion_client,
             )
@@ -1355,6 +1366,7 @@ def raise_unspecified_exception(*args, **kwargs):
         state = await create_then_begin_flow_run(
             flow=parameterized_flow,
             parameters={"puppy": "a string", "kitty": 42},
+            wait_for=None,
             return_type="state",
             client=orion_client,
         )
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
--- a/tests/test_tasks.py
+++ b/tests/test_tasks.py
@@ -1,6 +1,7 @@
 import datetime
 import inspect
 import warnings
+from asyncio import Event, sleep
 from typing import Any, Dict, List
 from unittest.mock import MagicMock
 from uuid import UUID
@@ -1819,6 +1820,157 @@ def unpacking_flow():
         )
 
 
+class TestSubflowWaitForTasks:
+    def test_downstream_does_not_run_if_upstream_fails(self):
+        @task
+        def fails():
+            raise ValueError("Fail task!")
+
+        @flow
+        def bar(y):
+            return y
+
+        @flow
+        def test_flow():
+            f = fails.submit()
+            b = bar._run(2, wait_for=[f])
+            return b
+
+        flow_state = test_flow._run()
+        subflow_state = flow_state.result(raise_on_failure=False)
+        assert subflow_state.is_pending()
+        assert subflow_state.name == "NotReady"
+
+    def test_downstream_runs_if_upstream_succeeds(self):
+        @flow
+        def foo(x):
+            return x
+
+        @flow
+        def bar(y):
+            return y
+
+        @flow
+        def test_flow():
+            f = foo(1)
+            b = bar(2, wait_for=[f])
+            return b
+
+        assert test_flow() == 2
+
+    async def test_backend_task_inputs_includes_wait_for_tasks(self, orion_client):
+        @task
+        def foo(x):
+            return x
+
+        @flow
+        def flow_foo(x):
+            return x
+
+        @flow
+        def test_flow():
+            a, b = foo.submit(1), foo.submit(2)
+            c = foo.submit(3)
+            d = flow_foo(c, wait_for=[a, b], return_state=True)
+            return (a, b, c, d)
+
+        a, b, c, d = test_flow()
+        d_subflow_run = await orion_client.read_flow_run(d.state_details.flow_run_id)
+        d_virtual_task_run = await orion_client.read_task_run(
+            d_subflow_run.parent_task_run_id
+        )
+
+        assert d_virtual_task_run.task_inputs["x"] == [
+            TaskRunResult(id=c.state_details.task_run_id)
+        ], "Data passing inputs are preserved"
+
+        assert set(d_virtual_task_run.task_inputs["wait_for"]) == {
+            TaskRunResult(id=a.state_details.task_run_id),
+            TaskRunResult(id=b.state_details.task_run_id),
+        }, "'wait_for' included as a key with upstreams"
+
+        assert set(d_virtual_task_run.task_inputs.keys()) == {
+            "x",
+            "wait_for",
+        }, "No extra keys around"
+
+    async def test_subflows_run_concurrently_with_tasks(self):
+        @task
+        async def waiter_task(event, delay):
+            await sleep(delay)
+            if event.is_set():
+                pass
+            else:
+                raise RuntimeError("The event hasn't been set!")
+
+        @flow
+        async def setter_flow(event):
+            event.set()
+            return 42
+
+        @flow
+        async def test_flow():
+            e = Event()
+            f = await waiter_task.submit(e, 1)
+            b = await setter_flow(e)
+            return b
+
+        assert (await test_flow()) == 42
+
+    async def test_subflows_waiting_for_tasks_can_deadlock(self):
+        @task
+        async def waiter_task(event, delay):
+            await sleep(delay)
+            if event.is_set():
+                pass
+            else:
+                raise RuntimeError("The event hasn't been set!")
+
+        @flow
+        async def setter_flow(event):
+            event.set()
+            return 42
+
+        @flow
+        async def test_flow():
+            e = Event()
+            f = await waiter_task.submit(e, 1)
+            b = await setter_flow(e, wait_for=[f])
+            return b
+
+        flow_state = await test_flow._run()
+        assert flow_state.is_failed()
+        assert "MissingResult: State data is missing" in flow_state.message
+
+    def test_using_wait_for_in_task_definition_raises_reserved(self):
+        with pytest.raises(
+            ReservedArgumentError, match="'wait_for' is a reserved argument name"
+        ):
+
+            @flow
+            def foo(wait_for):
+                pass
+
+    def test_downstream_runs_if_upstream_fails_with_allow_failure_annotation(self):
+        @task
+        def fails():
+            raise ValueError("Fail task!")
+
+        @flow
+        def bar(y):
+            return y
+
+        @flow
+        def test_flow():
+            f = fails.submit()
+            b = bar(2, wait_for=[allow_failure(f)], return_state=True)
+            return b
+
+        flow_state = test_flow(return_state=True)
+        subflow_state = flow_state.result(raise_on_failure=False)
+        assert subflow_state.result() == 2
+
+
 class TestTaskWaitFor:
     def test_downstream_does_not_run_if_upstream_fails(self):
         @task