Repository: PrefectHQ/prefect

Issue:
Support Task.map for autonomous tasks
When using a task outside of a flow, AKA an "autonomous task," you should be able to call `Task.map` to map task invocations over an iterable.


Discussion:

Resolving patch:
diff --git a/src/prefect/engine.py b/src/prefect/engine.py
--- a/src/prefect/engine.py
+++ b/src/prefect/engine.py
@@ -1386,6 +1386,7 @@ def enter_task_run_engine(
 
         return submit_autonomous_task_run_to_engine(
             task=task,
+            task_run=None,
             parameters=parameters,
             task_runner=task_runner,
             wait_for=wait_for,
@@ -1421,12 +1422,13 @@ def enter_task_run_engine(
 
 async def begin_task_map(
     task: Task,
-    flow_run_context: FlowRunContext,
+    flow_run_context: Optional[FlowRunContext],
     parameters: Dict[str, Any],
     wait_for: Optional[Iterable[PrefectFuture]],
     return_type: EngineReturnType,
     task_runner: Optional[BaseTaskRunner],
-) -> List[Union[PrefectFuture, Awaitable[PrefectFuture]]]:
+    autonomous: bool = False,
+) -> List[Union[PrefectFuture, Awaitable[PrefectFuture], TaskRun]]:
     """Async entrypoint for task mapping"""
     # We need to resolve some futures to map over their data, collect the upstream
     # links beforehand to retain relationship tracking.
@@ -1493,18 +1495,29 @@ async def begin_task_map(
         # Collapse any previously exploded kwargs
         call_parameters = collapse_variadic_parameters(task.fn, call_parameters)
 
-        task_runs.append(
-            partial(
-                get_task_call_return_value,
-                task=task,
-                flow_run_context=flow_run_context,
-                parameters=call_parameters,
-                wait_for=wait_for,
-                return_type=return_type,
-                task_runner=task_runner,
-                extra_task_inputs=task_inputs,
+        if autonomous:
+            task_runs.append(
+                await create_autonomous_task_run(
+                    task=task,
+                    parameters=call_parameters,
+                )
             )
-        )
+        else:
+            task_runs.append(
+                partial(
+                    get_task_call_return_value,
+                    task=task,
+                    flow_run_context=flow_run_context,
+                    parameters=call_parameters,
+                    wait_for=wait_for,
+                    return_type=return_type,
+                    task_runner=task_runner,
+                    extra_task_inputs=task_inputs,
+                )
+            )
+
+    if autonomous:
+        return task_runs
 
     # Maintain the order of the task runs when using the sequential task runner
     runner = task_runner if task_runner else flow_run_context.task_runner
diff --git a/src/prefect/tasks.py b/src/prefect/tasks.py
--- a/src/prefect/tasks.py
+++ b/src/prefect/tasks.py
@@ -808,16 +808,16 @@ def submit(
             PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING.value()
             and not FlowRunContext.get()
         ):
-            create_autonomous_task_run = create_call(
+            create_autonomous_task_run_call = create_call(
                 create_autonomous_task_run, task=self, parameters=parameters
             )
             if self.isasync:
                 return from_async.wait_for_call_in_loop_thread(
-                    create_autonomous_task_run
+                    create_autonomous_task_run_call
                 )
             else:
                 return from_sync.wait_for_call_in_loop_thread(
-                    create_autonomous_task_run
+                    create_autonomous_task_run_call
                 )
 
         return enter_task_run_engine(
@@ -984,7 +984,7 @@ def map(
             [[11, 21], [12, 22], [13, 23]]
         """
 
-        from prefect.engine import enter_task_run_engine
+        from prefect.engine import begin_task_map, enter_task_run_engine
 
         # Convert the call args/kwargs to a parameter dict; do not apply defaults
         # since they should not be mapped over
@@ -997,6 +997,25 @@ def map(
                 "`task.map()` is not currently supported by `flow.visualize()`"
             )
 
+        if (
+            PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING.value()
+            and not FlowRunContext.get()
+        ):
+            map_call = create_call(
+                begin_task_map,
+                task=self,
+                parameters=parameters,
+                flow_run_context=None,
+                wait_for=wait_for,
+                return_type=return_type,
+                task_runner=None,
+                autonomous=True,
+            )
+            if self.isasync:
+                return from_async.wait_for_call_in_loop_thread(map_call)
+            else:
+                return from_sync.wait_for_call_in_loop_thread(map_call)
+
         return enter_task_run_engine(
             self,
             parameters=parameters,


Test patch:
diff --git a/tests/test_autonomous_tasks.py b/tests/test_autonomous_tasks.py
--- a/tests/test_autonomous_tasks.py
+++ b/tests/test_autonomous_tasks.py
@@ -1,12 +1,12 @@
 import asyncio
 from pathlib import Path
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Iterable, Tuple
 from unittest import mock
 
 import pytest
 
 import prefect.results
-from prefect import Task, task
+from prefect import Task, task, unmapped
 from prefect.blocks.core import Block
 from prefect.client.orchestration import PrefectClient, get_client
 from prefect.client.schemas import TaskRun
@@ -27,7 +27,7 @@
 
 
 @sync_compatible
-async def result_factory_from_task(task):
+async def result_factory_from_task(task) -> ResultFactory:
     return await ResultFactory.from_autonomous_task(task)
 
 
@@ -63,7 +63,7 @@ async def clear_cached_filesystems():
 
 
 @pytest.fixture
-def foo_task():
+def foo_task() -> Task:
     @task
     def foo(x: int) -> int:
         print(x)
@@ -73,7 +73,7 @@ def foo(x: int) -> int:
 
 
 @pytest.fixture
-def async_foo_task():
+def async_foo_task() -> Task:
     @task
     async def async_foo(x: int) -> int:
         print(x)
@@ -183,11 +183,6 @@ async def test_async_task_submission_creates_a_scheduled_task_run(
     assert parameters == dict(x=42)
 
 
-async def test_task_submission_via_map_raises_error(async_foo_task_with_result_storage):
-    with pytest.raises(RuntimeError, match="Tasks cannot be run outside of a flow"):
-        async_foo_task_with_result_storage.map([42])
-
-
 async def test_scheduled_tasks_are_enqueued_server_side(
     foo_task_with_result_storage: Task,
 ):
@@ -320,3 +315,86 @@ async def test_call_with_return_state(self, async_foo_task):
         assert state.is_completed()
 
         assert await state.result() == 42
+
+
+class TestMap:
+    async def test_map(self, async_foo_task):
+        task_runs = await async_foo_task.map([1, 2, 3])
+
+        assert len(task_runs) == 3
+
+        result_factory = await result_factory_from_task(async_foo_task)
+
+        for i, task_run in enumerate(task_runs):
+            assert task_run.state.is_scheduled()
+            assert await result_factory.read_parameters(
+                task_run.state.state_details.task_parameters_id
+            ) == {"x": i + 1}
+
+    async def test_map_with_implicitly_unmapped_kwargs(self):
+        @task
+        def bar(x: int, unmappable: int) -> Tuple[int, int]:
+            return (x, unmappable)
+
+        task_runs = bar.map([1, 2, 3], unmappable=42)
+
+        assert len(task_runs) == 3
+
+        result_factory = await result_factory_from_task(bar)
+
+        for i, task_run in enumerate(task_runs):
+            assert task_run.state.is_scheduled()
+            assert await result_factory.read_parameters(
+                task_run.state.state_details.task_parameters_id
+            ) == {"x": i + 1, "unmappable": 42}
+
+    async def test_async_map_with_implicitly_unmapped_kwargs(self):
+        @task
+        async def bar(x: int, unmappable: int) -> Tuple[int, int]:
+            return (x, unmappable)
+
+        task_runs = await bar.map([1, 2, 3], unmappable=42)
+
+        assert len(task_runs) == 3
+
+        result_factory = await result_factory_from_task(bar)
+
+        for i, task_run in enumerate(task_runs):
+            assert task_run.state.is_scheduled()
+            assert await result_factory.read_parameters(
+                task_run.state.state_details.task_parameters_id
+            ) == {"x": i + 1, "unmappable": 42}
+
+    async def test_map_with_explicit_unmapped_kwargs(self):
+        @task
+        def bar(x: int, mappable: Iterable) -> Tuple[int, Iterable]:
+            return (x, mappable)
+
+        task_runs = bar.map([1, 2, 3], mappable=unmapped(["some", "iterable"]))
+
+        assert len(task_runs) == 3
+
+        result_factory = await result_factory_from_task(bar)
+
+        for i, task_run in enumerate(task_runs):
+            assert task_run.state.is_scheduled()
+            assert await result_factory.read_parameters(
+                task_run.state.state_details.task_parameters_id
+            ) == {"x": i + 1, "mappable": ["some", "iterable"]}
+
+    async def test_async_map_with_explicit_unmapped_kwargs(self):
+        @task
+        async def bar(x: int, mappable: Iterable) -> Tuple[int, Iterable]:
+            return (x, mappable)
+
+        task_runs = await bar.map([1, 2, 3], mappable=unmapped(["some", "iterable"]))
+
+        assert len(task_runs) == 3
+
+        result_factory = await result_factory_from_task(bar)
+
+        for i, task_run in enumerate(task_runs):
+            assert task_run.state.is_scheduled()
+            assert await result_factory.read_parameters(
+                task_run.state.state_details.task_parameters_id
+            ) == {"x": i + 1, "mappable": ["some", "iterable"]}