Repository: PrefectHQ/prefect

Issue:
Support `resource_manager` tasks without a setup/cleanup step
The current `prefect.resource_manager` design requires both a `setup` and `cleanup` method defined. Sometimes you want to run a cleanup step without an explicit setup step (or the other way around, although that's less useful). It'd be good to modify the `resource_manager` to skip setup/cleanup task creation if the `setup`/`cleanup` methods are missing from the class.

See #3836 for a motivator for this. The `setup` method is unneeded, we only want to create a `cleanup` task that always runs after exiting that block.


Discussion:

Resolving patch:
diff --git a/src/prefect/core/flow.py b/src/prefect/core/flow.py
--- a/src/prefect/core/flow.py
+++ b/src/prefect/core/flow.py
@@ -401,10 +401,13 @@ def parameters(self) -> Set[Parameter]:
 
     @cache
     def _default_reference_tasks(self) -> Set[Task]:
-        from prefect.tasks.core.resource_manager import ResourceCleanupTask
+        from prefect.tasks.core.resource_manager import (
+            ResourceInitTask,
+            ResourceCleanupTask,
+        )
 
-        # Select all tasks that aren't ResourceCleanupTasks and have no
-        # downstream dependencies that aren't ResourceCleanupTasks
+        # Select all tasks that aren't a ResourceInitTask/ResourceCleanupTask
+        # and have no downstream dependencies that aren't ResourceCleanupTasks
         #
         # Note: this feels a bit gross, since it special cases a certain
         # subclass inside the flow runner. If this behavior expands to other
@@ -413,7 +416,7 @@ def _default_reference_tasks(self) -> Set[Task]:
         return {
             t
             for t in self.tasks
-            if not isinstance(t, ResourceCleanupTask)
+            if not isinstance(t, (ResourceInitTask, ResourceCleanupTask))
             and not any(
                 t
                 for t in self.downstream_tasks(t)
diff --git a/src/prefect/tasks/core/resource_manager.py b/src/prefect/tasks/core/resource_manager.py
--- a/src/prefect/tasks/core/resource_manager.py
+++ b/src/prefect/tasks/core/resource_manager.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, Union, Set, overload
+from typing import Any, Callable, Dict, Union, Set, Optional, overload
 
 import prefect
 from prefect import Task, Flow
@@ -28,7 +28,7 @@ def run(self, mgr: Any) -> Any:
 class ResourceCleanupTask(Task):
     """Cleanup a resource with its resource manager"""
 
-    def run(self, mgr: Any, resource: Any) -> None:
+    def run(self, mgr: Any, resource: Any = None) -> None:
         mgr.cleanup(resource)
 
 
@@ -67,7 +67,11 @@ class ResourceContext:
     """
 
     def __init__(
-        self, init_task: Task, setup_task: Task, cleanup_task: Task, flow: Flow
+        self,
+        init_task: Task,
+        setup_task: Optional[Task],
+        cleanup_task: Task,
+        flow: Flow,
     ):
         self.init_task = init_task
         self.setup_task = setup_task
@@ -88,7 +92,7 @@ def add_task(self, task: Task, flow: Flow) -> None:
             )
         self._tasks.add(task)
 
-    def __enter__(self) -> Task:
+    def __enter__(self) -> Optional[Task]:
         self.__prev_resource = prefect.context.get("resource")
         prefect.context.update(resource=self)
         return self.setup_task
@@ -106,13 +110,15 @@ def __exit__(self, *args: Any) -> None:
             # the resource cleanup should be set as a downstream task.
             upstream = self._flow.upstream_tasks(child)
             if (
-                not self._tasks.intersection(upstream)
+                self.setup_task is not None
+                and not self._tasks.intersection(upstream)
                 and self.setup_task not in upstream
             ):
                 child.set_upstream(self.setup_task, flow=self._flow)
             downstream = self._flow.downstream_tasks(child)
             if (
-                not self._tasks.intersection(downstream)
+                self.cleanup_task is not None
+                and not self._tasks.intersection(downstream)
                 and self.cleanup_task not in downstream
             ):
                 child.set_downstream(self.cleanup_task, flow=self._flow)
@@ -212,11 +218,18 @@ def __call__(self, *args: Any, flow: Flow = None, **kwargs: Any) -> ResourceCont
             *args, flow=flow, **kwargs
         )
 
-        setup_task = ResourceSetupTask(**self.setup_task_kwargs)(init_task, flow=flow)
-
-        cleanup_task = ResourceCleanupTask(**self.cleanup_task_kwargs)(
-            init_task, setup_task, flow=flow
-        )
+        if hasattr(self.resource_class, "setup"):
+            setup_task = ResourceSetupTask(**self.setup_task_kwargs)(
+                init_task, flow=flow
+            )
+            cleanup_task = ResourceCleanupTask(**self.cleanup_task_kwargs)(
+                init_task, setup_task, flow=flow
+            )
+        else:
+            setup_task = None
+            cleanup_task = ResourceCleanupTask(**self.cleanup_task_kwargs)(
+                init_task, flow=flow
+            )
 
         return ResourceContext(init_task, setup_task, cleanup_task, flow)
 
@@ -257,18 +270,19 @@ def resource_manager(
     """A decorator for creating a `ResourceManager` object.
 
     Used as a context manager, `ResourceManager` objects create tasks to setup
-    and cleanup temporary objects used within a block of tasks.  Examples might
-    include temporary Dask/Spark clusters, Docker containers, etc...
+    and/or cleanup temporary objects used within a block of tasks.  Examples
+    might include temporary Dask/Spark clusters, Docker containers, etc...
 
-    Through usage a ResourceManager object adds three tasks to the graph:
+    Through usage a ResourceManager object adds up to three tasks to the graph:
         - A `init` task, which returns an object that meets the `ResourceManager`
-          protocol. This protocol requires two methods:
+          protocol. This protocol contains two methods:
             * `setup(self) -> resource`: A method for creating the resource.
-                The return value from this will available to user tasks.
+                The return value from this will available to user tasks. If no
+                setup is required, the `setup` method may be left undefined.
             * `cleanup(self, resource) -> None`: A method for cleaning up the
-                resource.  This takes the return value from `setup` and
-                shouldn't return anything.
-        - A `setup` task, which calls the `setup` method on the `ResourceManager`
+                resource. This takes the return value from `setup` (or `None`
+                if no `setup` method) and shouldn't return anything.
+        - A `setup` task, which calls the optional `setup` method on the `ResourceManager`
         - A `cleanup` task, which calls the `cleanup` method on the `ResourceManager`.
 
     Args:
@@ -278,7 +292,7 @@ def resource_manager(
         - init_task_kwargs (dict, optional): keyword arguments that will be
             passed to the `Task` constructor for the `init` task.
         - setup_task_kwargs (dict, optional): keyword arguments that will be
-            passed to the `Task` constructor for the `setup` task.
+            passed to the `Task` constructor for the optional `setup` task.
         - cleanup_task_kwargs (dict, optional): keyword arguments that will be
             passed to the `Task` constructor for the `cleanup` task.
 


Test patch:
diff --git a/tests/tasks/core/test_resource_manager.py b/tests/tasks/core/test_resource_manager.py
--- a/tests/tasks/core/test_resource_manager.py
+++ b/tests/tasks/core/test_resource_manager.py
@@ -186,6 +186,56 @@ def test_resource_manager_generated_flow_structure(api):
     }
 
 
+def test_resource_manager_generated_flow_structure_no_setup():
+    @resource_manager
+    class MyResource:
+        def __init__(self, a):
+            self.a = a
+
+        def cleanup(self, val):
+            pass
+
+    with Flow("test") as flow:
+        a = inc(1)
+        context = MyResource(a)
+        with context as resource:
+            b = add(resource, a)
+            c = inc(b)
+            d = inc(2)
+            e = inc(d)
+            f = inc(3)
+        g = inc(f)
+
+    # task kwargs successfully forwarded to tasks
+    assert context.init_task.name == "MyResource"
+    assert context.setup_task is None
+    assert resource is None
+    assert context.cleanup_task.name == "MyResource.cleanup"
+    assert not context.cleanup_task.skip_on_upstream_skip
+
+    # Reference tasks setup properly
+    assert flow.reference_tasks() == {c, e, g}
+
+    # Check that:
+    # - Tasks with no downstream dependency in the resource context have
+    #   the cleanup task set as a downstream dependency
+    # - All other tasks only have explicit dependencies
+    assert flow.upstream_tasks(a) == set()
+    assert flow.upstream_tasks(context.init_task) == {a}
+    assert flow.upstream_tasks(b) == {a}
+    assert flow.upstream_tasks(c) == {b}
+    assert flow.upstream_tasks(d) == set()
+    assert flow.upstream_tasks(e) == {d}
+    assert flow.upstream_tasks(f) == set()
+    assert flow.upstream_tasks(g) == {f}
+    assert flow.upstream_tasks(context.cleanup_task) == {
+        context.init_task,
+        c,
+        e,
+        f,
+    }
+
+
 def test_resource_manager_execution_success():
     on_setup = MagicMock(return_value=100)
     on_cleanup = MagicMock()
@@ -206,6 +256,29 @@ def test_resource_manager_execution_success():
         assert r.is_successful()
 
 
+def test_resource_manager_execution_success_no_setup():
+    @resource_manager
+    class MyResource:
+        def __init__(self, on_cleanup):
+            self.on_cleanup = on_cleanup
+
+        def cleanup(self, val):
+            self.on_cleanup(val)
+
+    on_cleanup = MagicMock()
+
+    with Flow("test") as flow:
+        context = MyResource(on_cleanup)
+        with context:
+            inc(inc(1))
+
+    state = flow.run()
+    assert on_cleanup.call_args == ((None,), {})
+    assert state.is_successful()
+    for r in state.result.values():
+        assert r.is_successful()
+
+
 @pytest.mark.parametrize("kind", ["init", "setup", "cleanup"])
 def test_resource_manager_execution_with_failure_in_manager(kind):
     on_setup = MagicMock(return_value=100)