Repository: PrefectHQ/prefect
Pull request: Adds `deferred` kwarg to `Task.map` to executed mapped tasks on a task server
<!-- 
Thanks for opening a pull request to Prefect! We've got a few requests to help us review contributions:

- Make sure that your title neatly summarizes the proposed changes.
- Provide a short overview of the change and the value it adds.
- Share an example to help us understand the change in user experience.
- Confirm that you've done common tasks so we can give a timely review.
- Review our contribution guidelines: https://docs.prefect.io/latest/contributing/overview/

Happy engineering!
-->

<!-- Include an overview here -->
Adds a `deferred` kwarg to signify that mapped tasks should be executed on a task server. `deferred` can be used inside or outside of flows. Calling `Task.map` outside of flow without `deferred=True` will raise an error.

Closes #13722

### Example
<!-- 
Share an example of the change in action.

A code blurb is best. Changes to features should include an example that is executable by a new user.
If changing documentation, a link to a preview of the page is great.
 -->

### Checklist
<!-- These boxes may be checked after opening the pull request. -->

- [ ] This pull request references any related issue by including "closes `<link to issue>`"
 	- If no issue exists and your change is not a small fix, please [create an issue](https://github.com/PrefectHQ/prefect/issues/new/choose) first.
- [x] If this pull request adds new functionality, it includes unit tests that cover the changes
- [x] This pull request includes a label categorizing the change e.g. `maintenance`, `fix`, `feature`, `enhancement`, `docs`.
  <!-- If you do not have permission to add a label, a maintainer will add one for you -->

For documentation changes:

- [ ] This pull request includes redirect settings in `mint.json` for files that are removed or renamed.

For new functions or classes in the Python SDK:

- [ ] This pull request includes helpful docstrings.
- [ ] If a new Python file was added, this pull request contains a stub page in the Python SDK docs and an entry in `docs/mint.json` navigation.


Problem statement:
Add a `deferred` parameter to `Task.map` to control whether to defer or submit tasks
Currently, `Task.map` always submits tasks within a flow or defers tasks outside of a flow. But now that we support using `delay` and `apply_async` within flows correctly, a flow might want to defer tasks with `map` and may not be using `PrefectTaskRunner`, so we should allow controlling map behavior with a `deferred` parameter that defaults to `False`.

Hints from issue discussion:

Commit message: Fixes failing tests

Patch to review:
diff --git a/src/prefect/tasks.py b/src/prefect/tasks.py
index ec2dcd497d0b..96d0a153ba69 100644
--- a/src/prefect/tasks.py
+++ b/src/prefect/tasks.py
@@ -849,7 +849,11 @@ def submit(
         flow_run_context = FlowRunContext.get()
 
         if not flow_run_context:
-            raise ValueError("Task.submit() must be called within a flow")
+            raise RuntimeError(
+                "Unable to determine task runner to use for submission. If you are"
+                " submitting a task outside of a flow, please use `.delay`"
+                " to submit the task run for deferred execution."
+            )
 
         task_viz_tracker = get_task_viz_tracker()
         if task_viz_tracker:
@@ -897,6 +901,7 @@ def map(
         *args: Any,
         return_state: bool = False,
         wait_for: Optional[Iterable[PrefectFuture]] = None,
+        deferred: bool = False,
         **kwargs: Any,
     ):
         """
@@ -1010,6 +1015,7 @@ def map(
             [[11, 21], [12, 22], [13, 23]]
         """
 
+        from prefect.task_runners import TaskRunner
         from prefect.utilities.visualization import (
             VisualizationUnsupportedError,
             get_task_viz_tracker,
@@ -1026,22 +1032,21 @@ def map(
                 "`task.map()` is not currently supported by `flow.visualize()`"
             )
 
-        if not flow_run_context:
-            # TODO: Should we split out background task mapping into a separate method
-            # like we do for the `submit`/`apply_async` split?
+        if deferred:
             parameters_list = expand_mapping_parameters(self.fn, parameters)
-            # TODO: Make this non-blocking once we can return a list of futures
-            # instead of a list of task runs
-            return [
-                run_coro_as_sync(self.create_run(parameters=parameters, deferred=True))
+            futures = [
+                self.apply_async(kwargs=parameters, wait_for=wait_for)
                 for parameters in parameters_list
             ]
-
-        from prefect.task_runners import TaskRunner
-
-        task_runner = flow_run_context.task_runner
-        assert isinstance(task_runner, TaskRunner)
-        futures = task_runner.map(self, parameters, wait_for)
+        elif task_runner := getattr(flow_run_context, "task_runner", None):
+            assert isinstance(task_runner, TaskRunner)
+            futures = task_runner.map(self, parameters, wait_for)
+        else:
+            raise RuntimeError(
+                "Unable to determine task runner to use for mapping. If you are"
+                " mapping a task outside of a flow, please provide `deferred=True`"
+                " to submit the mapped task runs for deferred execution."
+            )
         if return_state:
             states = []
             for future in futures:
diff --git a/tests/test_background_tasks.py b/tests/test_background_tasks.py
index 40a9e34c3015..9dd3a70047d8 100644
--- a/tests/test_background_tasks.py
+++ b/tests/test_background_tasks.py
@@ -333,7 +333,7 @@ async def test_call_with_return_state(self, async_foo_task):
 
 class TestMap:
     async def test_map(self, async_foo_task):
-        task_runs = async_foo_task.map([1, 2, 3])
+        task_runs = async_foo_task.map([1, 2, 3], deferred=True)
 
         assert len(task_runs) == 3
 
@@ -350,7 +350,7 @@ async def test_map_with_implicitly_unmapped_kwargs(self):
         def bar(x: int, unmappable: int) -> Tuple[int, int]:
             return (x, unmappable)
 
-        task_runs = bar.map([1, 2, 3], unmappable=42)
+        task_runs = bar.map([1, 2, 3], unmappable=42, deferred=True)
 
         assert len(task_runs) == 3
 
@@ -367,7 +367,7 @@ async def test_async_map_with_implicitly_unmapped_kwargs(self):
         async def bar(x: int, unmappable: int) -> Tuple[int, int]:
             return (x, unmappable)
 
-        task_runs = bar.map([1, 2, 3], unmappable=42)
+        task_runs = bar.map([1, 2, 3], unmappable=42, deferred=True)
 
         assert len(task_runs) == 3
 
@@ -384,7 +384,9 @@ async def test_map_with_explicit_unmapped_kwargs(self):
         def bar(x: int, mappable: Iterable) -> Tuple[int, Iterable]:
             return (x, mappable)
 
-        task_runs = bar.map([1, 2, 3], mappable=unmapped(["some", "iterable"]))
+        task_runs = bar.map(
+            [1, 2, 3], mappable=unmapped(["some", "iterable"]), deferred=True
+        )
 
         assert len(task_runs) == 3
 
@@ -404,7 +406,9 @@ async def test_async_map_with_explicit_unmapped_kwargs(self):
         async def bar(x: int, mappable: Iterable) -> Tuple[int, Iterable]:
             return (x, mappable)
 
-        task_runs = bar.map([1, 2, 3], mappable=unmapped(["some", "iterable"]))
+        task_runs = bar.map(
+            [1, 2, 3], mappable=unmapped(["some", "iterable"]), deferred=True
+        )
 
         assert len(task_runs) == 3
 
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
index 12eaea31ad28..08bfa4470514 100644
--- a/tests/test_tasks.py
+++ b/tests/test_tasks.py
@@ -18,7 +18,7 @@
 from prefect.client.orchestration import PrefectClient
 from prefect.client.schemas.filters import LogFilter, LogFilterFlowRunId
 from prefect.client.schemas.objects import StateType, TaskRunResult
-from prefect.context import TaskRunContext, get_run_context
+from prefect.context import FlowRunContext, TaskRunContext, get_run_context
 from prefect.exceptions import (
     MappingLengthMismatch,
     MappingMissingIterable,
@@ -76,6 +76,11 @@ def test_flow():
     return test_flow
 
 
+async def get_background_task_run_parameters(task, parameters_id):
+    factory = await ResultFactory.from_autonomous_task(task)
+    return await factory.read_parameters(parameters_id)
+
+
 class TestTaskName:
     def test_name_from_function(self):
         @task
@@ -203,19 +208,6 @@ async def bar():
 
         assert await bar() == 1
 
-    # Will not be supported in new engine
-    @pytest.mark.skip(reason="Fails with new engine, passed on old engine")
-    def test_async_task_called_inside_sync_flow(self):
-        @task
-        async def foo(x):
-            return x
-
-        @flow
-        def bar():
-            return foo(1)
-
-        assert bar() == 1
-
     def test_task_call_with_debug_mode(self):
         @task
         def foo(x):
@@ -371,21 +363,6 @@ async def bar():
         assert isinstance(task_state, State)
         assert await task_state.result() == 1
 
-    # Will not be supported in new engine
-    @pytest.mark.skip(reason="Fails with new engine, passed on old engine")
-    def test_async_task_run_inside_sync_flow(self):
-        @task
-        async def foo(x):
-            return x
-
-        @flow
-        def bar():
-            return foo(1, return_state=True)
-
-        task_state = bar()
-        assert isinstance(task_state, State)
-        assert task_state.result() == 1
-
     def test_task_failure_does_not_affect_flow(self):
         @task
         def foo():
@@ -425,6 +402,14 @@ def my_flow():
 
 
 class TestTaskSubmit:
+    def test_raises_outside_of_flow(self):
+        @task
+        def foo(x):
+            return x
+
+        with pytest.raises(RuntimeError):
+            foo.submit(1)
+
     async def test_sync_task_submitted_inside_sync_flow(self):
         @task
         def foo(x):
@@ -3165,6 +3150,67 @@ def my_flow():
         task_states = my_flow()
         assert [await state.result() for state in task_states] == [2, 3, 4]
 
+    def test_map_raises_outside_of_flow_when_not_deferred(self):
+        @task
+        def test_task(x):
+            print(x)
+
+        with pytest.raises(RuntimeError):
+            test_task.map([1, 2, 3])
+
+    async def test_deferred_map_outside_flow(self):
+        @task
+        def test_task(x):
+            print(x)
+
+        mock_task_run_id = uuid4()
+        mock_future = PrefectDistributedFuture(task_run_id=mock_task_run_id)
+        mapped_args = [1, 2, 3]
+
+        futures = test_task.map(x=mapped_args, wait_for=[mock_future], deferred=True)
+        assert all(isinstance(future, PrefectDistributedFuture) for future in futures)
+        for future, parameter_value in zip(futures, mapped_args):
+            assert await get_background_task_run_parameters(
+                test_task, future.state.state_details.task_parameters_id
+            ) == {
+                "parameters": {"x": parameter_value},
+                "wait_for": [mock_future],
+                "context": ANY,
+            }
+
+    async def test_deferred_map_inside_flow(self):
+        @task
+        def test_task(x):
+            print(x)
+
+        @flow
+        async def test_flow():
+            mock_task_run_id = uuid4()
+            mock_future = PrefectDistributedFuture(task_run_id=mock_task_run_id)
+            mapped_args = [1, 2, 3]
+
+            flow_run_context = FlowRunContext.get()
+            flow_run_id = flow_run_context.flow_run.id
+            futures = test_task.map(
+                x=mapped_args, wait_for=[mock_future], deferred=True
+            )
+            for future, parameter_value in zip(futures, mapped_args):
+                saved_data = await get_background_task_run_parameters(
+                    test_task, future.state.state_details.task_parameters_id
+                )
+                assert saved_data == {
+                    "parameters": {"x": parameter_value},
+                    "wait_for": [mock_future],
+                    "context": ANY,
+                }
+                # Context should contain the current flow run ID
+                assert (
+                    saved_data["context"]["flow_run_context"]["flow_run"]["id"]
+                    == flow_run_id
+                )
+
+        await test_flow()
+
 
 class TestTaskConstructorValidation:
     async def test_task_cannot_configure_too_many_custom_retry_delays(self):
@@ -4119,10 +4165,6 @@ def commit(txn):
 
 
 class TestApplyAsync:
-    async def get_background_task_run_parameters(self, task, parameters_id):
-        factory = await ResultFactory.from_autonomous_task(task)
-        return await factory.read_parameters(parameters_id)
-
     @pytest.mark.parametrize(
         "args, kwargs",
         [
@@ -4139,7 +4181,7 @@ def multiply(x, y):
 
         future = multiply.apply_async(args, kwargs)
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             multiply, future.state.state_details.task_parameters_id
         ) == {"parameters": {"x": 42, "y": 42}, "context": ANY}
 
@@ -4170,7 +4212,7 @@ def add(x, y=42):
 
         future = add.apply_async((42,))
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             add, future.state.state_details.task_parameters_id
         ) == {"parameters": {"x": 42, "y": 42}, "context": ANY}
 
@@ -4181,7 +4223,7 @@ def add(x, y=42):
 
         future = add.apply_async((42,), {"y": 100})
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             add, future.state.state_details.task_parameters_id
         ) == {"parameters": {"x": 42, "y": 100}, "context": ANY}
 
@@ -4192,7 +4234,7 @@ def add_em_up(*args):
 
         future = add_em_up.apply_async((42, 42))
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             add_em_up, future.state.state_details.task_parameters_id
         ) == {"parameters": {"args": (42, 42)}, "context": ANY}
 
@@ -4203,7 +4245,7 @@ def add_em_up(**kwargs):
 
         future = add_em_up.apply_async(kwargs={"x": 42, "y": 42})
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             add_em_up, future.state.state_details.task_parameters_id
         ) == {"parameters": {"kwargs": {"x": 42, "y": 42}}, "context": ANY}
 
@@ -4214,7 +4256,7 @@ def add_em_up(*args, **kwargs):
 
         future = add_em_up.apply_async((42,), {"y": 42})
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             add_em_up, future.state.state_details.task_parameters_id
         ) == {"parameters": {"args": (42,), "kwargs": {"y": 42}}, "context": ANY}
 
@@ -4228,7 +4270,7 @@ def multiply(x, y):
 
         future = multiply.apply_async((42, 42), wait_for=[wait_for_future])
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             multiply, future.state.state_details.task_parameters_id
         ) == {
             "parameters": {"x": 42, "y": 42},
@@ -4246,7 +4288,7 @@ def the_answer():
 
         future = the_answer.apply_async(wait_for=[wait_for_future])
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             the_answer, future.state.state_details.task_parameters_id
         ) == {"wait_for": [wait_for_future], "context": ANY}
 
@@ -4268,10 +4310,6 @@ def add(x, y):
 
 
 class TestDelay:
-    async def get_background_task_run_parameters(self, task, parameters_id):
-        factory = await ResultFactory.from_autonomous_task(task)
-        return await factory.read_parameters(parameters_id)
-
     @pytest.mark.parametrize(
         "args, kwargs",
         [
@@ -4288,7 +4326,7 @@ def multiply(x, y):
 
         future = multiply.delay(*args, **kwargs)
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             multiply, future.state.state_details.task_parameters_id
         ) == {"parameters": {"x": 42, "y": 42}, "context": ANY}
 
@@ -4319,7 +4357,7 @@ def add(x, y=42):
 
         future = add.delay(42)
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             add, future.state.state_details.task_parameters_id
         ) == {"parameters": {"x": 42, "y": 42}, "context": ANY}
 
@@ -4330,7 +4368,7 @@ def add(x, y=42):
 
         future = add.delay(42, y=100)
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             add, future.state.state_details.task_parameters_id
         ) == {"parameters": {"x": 42, "y": 100}, "context": ANY}
 
@@ -4341,7 +4379,7 @@ def add_em_up(*args):
 
         future = add_em_up.delay(42, 42)
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             add_em_up, future.state.state_details.task_parameters_id
         ) == {"parameters": {"args": (42, 42)}, "context": ANY}
 
@@ -4352,7 +4390,7 @@ def add_em_up(**kwargs):
 
         future = add_em_up.delay(x=42, y=42)
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             add_em_up, future.state.state_details.task_parameters_id
         ) == {"parameters": {"kwargs": {"x": 42, "y": 42}}, "context": ANY}
 
@@ -4363,6 +4401,6 @@ def add_em_up(*args, **kwargs):
 
         future = add_em_up.delay(42, y=42)
 
-        assert await self.get_background_task_run_parameters(
+        assert await get_background_task_run_parameters(
             add_em_up, future.state.state_details.task_parameters_id
         ) == {"parameters": {"args": (42,), "kwargs": {"y": 42}}, "context": ANY}


Code review comments:
diff_hunk: @@ -1026,22 +1032,21 @@ def map(
                 "`task.map()` is not currently supported by `flow.visualize()`"
             )
 
-        if not flow_run_context:
-            # TODO: Should we split out background task mapping into a separate method
-            # like we do for the `submit`/`apply_async` split?
+        if deferred:
             parameters_list = expand_mapping_parameters(self.fn, parameters)
-            # TODO: Make this non-blocking once we can return a list of futures
-            # instead of a list of task runs
-            return [
-                run_coro_as_sync(self.create_run(parameters=parameters, deferred=True))
+            futures = [
+                self.apply_async(kwargs=parameters, wait_for=wait_for)
                 for parameters in parameters_list
             ]
-
-        from prefect.task_runners import TaskRunner
-
-        task_runner = flow_run_context.task_runner
-        assert isinstance(task_runner, TaskRunner)
-        futures = task_runner.map(self, parameters, wait_for)
+        elif task_runner := getattr(flow_run_context, "task_runner", None):
+            assert isinstance(task_runner, TaskRunner)
+            futures = task_runner.map(self, parameters, wait_for)
+        else:
+            raise RuntimeError(
+                "Unable to determine task runner to use for mapping. If you are"
path: src/prefect/tasks.py
text: @user1:
nit
```suggestion
                "Unable to determine task runner to use for mapped task runs. If you are"
```