Repository: PrefectHQ/prefect

Issue:
Cannot use apply_map in a ResourceManager context
## Description
<!-- A clear description of the bug -->
Prefect version 0.14.2 introduced a bugfix #3920 which disallows the use of apply_map
within a ResourceManager context. We heavily rely on both ResourceManager and apply_map, e.g. we want to 
scale a presto/trino cluster using the setup/cleanup of the ResourceManager, and add tasks to the flow using apply_map.


## Expected Behavior
<!-- What did you expect to happen instead? -->

The following unittest must passed with latest prefect versions.
```
import pytest
from prefect import Flow, apply_map, task
from prefect.tasks.core.resource_manager import resource_manager


@resource_manager
class DummyResource:

    def cleanup(self, resource):
        pass


@task
def adder(x):
    return x + 1


def func(x):
    return adder(x)


@pytest.mark.parametrize("inputs, expected", [
    ([1], [2]),
    ([2, 3], [3, 4])
])
def test_apply_map_within_resource_manager(inputs, expected):
    with Flow(name='test') as flow:
        with DummyResource():
            y = task(fn=lambda x: x)(inputs)
            result = apply_map(func, y)
    assert flow.run().result[result].result == expected

```

## Reproduction
<!-- A minimal example that exhibits the behavior. -->

Test passes with prefect version 0.14.1. Test fails with exception `ValueError: Multiple flows cannot be used with the same resource block` with prefect version 0.14.2 and 0.14.3.

## Environment
<!-- Any additional information about your environment

Optionally run `prefect diagnostics` from the command line and paste the information here. -->
Unittest works with
`pip install prefect==0.14.1`

Flow doesn't work with
`pip install prefect==0.14.2` or
`pip install prefect==0.14.3`


Comments:

Patch:
diff --git a/changes/pr3975.yaml b/changes/pr3975.yaml
new file mode 100644
index 0000000000..b4b643e41b
--- /dev/null
+++ b/changes/pr3975.yaml
@@ -0,0 +1,2 @@
+fix:
+  - "Fix regression in `apply_map` which prevented use in `case`/`resource_manager` blocks - [#3975](https://github.com/PrefectHQ/prefect/pull/3975)"
diff --git a/src/prefect/utilities/tasks.py b/src/prefect/utilities/tasks.py
index 2915aaf251..19d4aad7cb 100644
--- a/src/prefect/utilities/tasks.py
+++ b/src/prefect/utilities/tasks.py
@@ -97,16 +97,20 @@ def apply_map(func: Callable, *args: Any, flow: "Flow" = None, **kwargs: Any) ->
     # Preprocess inputs to `apply_map`:
     # - Extract information about each argument (is unmapped, is constant, ...)
     # - Convert all arguments to instances of `Task`
-    # - Add all non-constant arguments to the flow. Constant arguments are
-    #   added later as needed.
+    # - Add all non-constant arguments to the flow and subflow. Constant arguments
+    #   are added later as needed.
     def preprocess(a: Any) -> "prefect.Task":
-        a2 = as_task(a, flow=flow2)
-        is_mapped = not isinstance(a, prefect.utilities.edges.unmapped)
-        is_constant = isinstance(a2, Constant)
+        # Clear external case/resource when adding tasks to flow2
+        with prefect.context(case=None, resource=None):
+            a2 = as_task(a, flow=flow2)
+            is_mapped = not isinstance(a, prefect.utilities.edges.unmapped)
+            is_constant = isinstance(a2, Constant)
+            if not is_constant:
+                flow2.add_task(a2)  # type: ignore
+
         arg_info[a2] = (is_mapped, is_constant)
         if not is_constant:
             flow.add_task(a2)  # type: ignore
-            flow2.add_task(a2)  # type: ignore
         if is_mapped and is_constant:
             id_to_const[id(a2.value)] = a2  # type: ignore
         return a2


Test patch:
diff --git a/tests/utilities/test_tasks.py b/tests/utilities/test_tasks.py
index 62ebe61107..d8a8700e23 100644
--- a/tests/utilities/test_tasks.py
+++ b/tests/utilities/test_tasks.py
@@ -1,6 +1,6 @@
 import pytest
 
-from prefect import Flow, Task, case, Parameter
+from prefect import Flow, Task, case, Parameter, resource_manager
 from prefect.engine.flow_runner import FlowRunner
 from prefect.engine.state import Paused, Resume
 from prefect.utilities import tasks, edges
@@ -302,24 +302,50 @@ class TestApplyMap:
         assert res == sol
 
     def test_apply_map_inside_case_statement_works(self):
-        def func(x):
-            return add(x, 1), add(x, 2)
+        def func(x, a):
+            return add(x, 1), add(x, a)
 
         with Flow("test") as flow:
             branch = Parameter("branch")
             with case(branch, True):
-                a, b = apply_map(func, range(4))
-                c = add.map(a, b)
+                a = inc(1)
+                b, c = apply_map(func, range(4), edges.unmapped(a))
+                d = add.map(b, c)
 
         state = flow.run(branch=True)
-        assert state.result[a].result == [1, 2, 3, 4]
-        assert state.result[b].result == [2, 3, 4, 5]
-        assert state.result[c].result == [3, 5, 7, 9]
+        assert state.result[a].result == 2
+        assert state.result[b].result == [1, 2, 3, 4]
+        assert state.result[c].result == [2, 3, 4, 5]
+        assert state.result[d].result == [3, 5, 7, 9]
 
         state = flow.run(branch=False)
         assert state.result[a].is_skipped()
         assert state.result[b].is_skipped()
         assert state.result[c].is_skipped()
+        assert state.result[d].is_skipped()
+
+    def test_apply_map_inside_resource_manager_works(self):
+        @resource_manager
+        class MyResource:
+            def setup(self):
+                return 1
+
+            def cleanup(self, _):
+                pass
+
+        def func(x, a):
+            return add(x, a), add(x, 2)
+
+        with Flow("test") as flow:
+            with MyResource() as a:
+                b, c = apply_map(func, range(4), edges.unmapped(a))
+                d = add.map(b, c)
+
+        state = flow.run()
+        assert state.result[a].result == 1
+        assert state.result[b].result == [1, 2, 3, 4]
+        assert state.result[c].result == [2, 3, 4, 5]
+        assert state.result[d].result == [3, 5, 7, 9]
 
     def test_apply_map_inputs_added_to_subflow_before_calling_func(self):
         """We need to ensure all args to `appy_map` are added to the temporary
@@ -380,7 +406,7 @@ class TestAsTask:
         assert res.result[val].result == obj
 
     def test_as_task_toggles_constants(self):
-        with Flow("test") as f:
+        with Flow("test"):
             t = tasks.as_task(4)
 
         assert isinstance(t, Task)
@@ -409,7 +435,7 @@ class TestAsTask:
         ],
     )
     def test_nested_collections_of_mixed_constants_are_not_constants(self, val):
-        with Flow("test") as f:
+        with Flow("test"):
             task = tasks.as_task(val)
         assert not isinstance(task, Constant)