Repository: PrefectHQ/prefect

Issue:
Recompute secret tasks when retrying flow from failed
## Opened from the [Prefect Public Slack Community](https://join.slack.com/t/prefect-public/shared_invite/enQtNzE5OTU3OTQwNzc1LTQ5M2FkZmQzZjI0ODg1ZTBmOTc0ZjVjYWFjMWExZDAyYzBmYjVmMTE1NTQ1Y2IxZTllOTc4MmI3NzYxMDlhYWU)

**lk.novotny**: Hello Prefect community. I have an issue with retrying failed flows runs with checkpointing where some tasks have secret outputs. Since their output value is not persisted I would expect them to be re-computed when restarting the flow run, but they are just always `None` . Reproducible example in thread. Thanks in advance for any help :slightly_smiling_face:

**lk.novotny**: My flow run may sometimes fail and I want to be able to restart it. Because of that I'm using the `flow.result` to store checkpoints. However, some tasks contain secret values that I don't want to store anywhere - so I've subclassed the `SecretBase` which indeed doesn't store the result. However, when I restart the flow from the failed state, the value of the secret is `None` (I suppose because there is no `Result` object for it), however I would expect it to be re-computed. Below is a simple example I was able to reproduce my issue with, here is the error I get when I re-run the flow from failed state. Is there something I'm doing wrong?


```
Unexpected error: TypeError("'NoneType' object is not subscriptable",)
Traceback (most recent call last):
  File ".../home/user/bss/dice/Tournesol/prefect_flows/flows//venv/lib/python3.6/site-packages/prefect/engine/runner.py", line 48, in inner
    new_state = method(self, state, *args, **kwargs)
  File ".../venv/lib/python3.6/site-packages/prefect/engine/task_runner.py", line 860, in get_task_run_state
    logger=self.logger,
  File ".../venv/lib/python3.6/site-packages/prefect/utilities/executors.py", line 298, in run_task_with_timeout
    return task.run(*args, **kwargs)  # type: ignore
  File ".../secret_result.py", line 35, in use_secret
    <http://prefect.context.logger.info|prefect.context.logger.info>(secret['path'])
TypeError: 'NoneType' object is not subscriptable
```


```
import os
import random
from os import environ

import prefect
from prefect import task, Flow, Parameter
from prefect.engine.results import LocalResult
from prefect.tasks.secrets import SecretBase

result_directory = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "..", "results")
)

result = LocalResult(
    dir=result_directory,
    location="{flow_name}/"
    "{scheduled_start_time:%d-%m_%H-%M-%S}/"
    "{task_full_name}-{task_run_id}.prefect_result",
)


class CustomSecret(SecretBase):
    def run(self):
        return {'path': environ.get('PATH'), 'home': environ.get('HOME')}


@task
def fragile_task(error_probability: float) -&gt; None:
    if random.random() &lt;= error_probability:
        raise ValueError('I\'m fragile, just retry me')


@task
def use_secret(secret: CustomSecret) -&gt; None:
    <http://prefect.context.logger.info|prefect.context.logger.info>(secret['path'])
    <http://prefect.context.logger.info|prefect.context.logger.info>(secret['home'])


with Flow('test', result=result) as flow:
    error_prob = Parameter('error_probability', default=0.7)
    secret = CustomSecret()

    a = fragile_task(error_prob)
    b = use_secret(secret)

    flow.set_dependencies(b, upstream_tasks=[a])

if __name__ == '__main__':
    flow.register(project_name='default') 
```



**josh**: Hi <@U017U6GA1NK> thanks for providing a reproducible example! This looks to be a common occurrence with _any_ task that inherits from the SecretBase class (even the PrefectSecret task). I’m going to open this as an issue to investigate how we can make sure these secret tasks are recomputed on retry. In the meantime you can always access secrets directly through the `prefect.client.Secret` object inside of your task or instead of a custom secret you could have a task that grabs the secret value and has `checkpoint=False` to not persist the secret data when passing it to the next task.

**josh**: <@ULVA73B9P> open “Recompute secret tasks when retrying flow from failed”

Original thread can be found [here](https://prefect-community.slack.com/archives/CL09KU1K7/p1609949303328800?thread_ts=1609949303.328800&cid=CL09KU1K7).



Discussion:

Patch:
diff --git a/changes/pr3975.yaml b/changes/pr3975.yaml
deleted file mode 100644
index b4b643e41b..0000000000
--- a/changes/pr3975.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-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/changes/pr3977.yaml b/changes/pr3977.yaml
new file mode 100644
index 0000000000..2b409e56be
--- /dev/null
+++ b/changes/pr3977.yaml
@@ -0,0 +1,2 @@
+enhancement:
+  - "Rerun secret tasks on flow-run restart - [#3977](https://github.com/PrefectHQ/prefect/pull/3977)"
diff --git a/src/prefect/engine/flow_runner.py b/src/prefect/engine/flow_runner.py
index 2c1ff6817e..5150b15fb2 100644
--- a/src/prefect/engine/flow_runner.py
+++ b/src/prefect/engine/flow_runner.py
@@ -427,8 +427,8 @@ class FlowRunner(Runner):
                 ):
                     task_states[task] = task_state = Success(result=task.value)
 
-                # Always restart completed resource setup/cleanup tasks unless
-                # they were explicitly cached.
+                # Always restart completed resource setup/cleanup tasks and
+                # secret tasks unless they were explicitly cached.
                 # TODO: we only need to rerun these tasks if any pending
                 # downstream tasks depend on them.
                 if (
@@ -437,6 +437,7 @@ class FlowRunner(Runner):
                         (
                             prefect.tasks.core.resource_manager.ResourceSetupTask,
                             prefect.tasks.core.resource_manager.ResourceCleanupTask,
+                            prefect.tasks.secrets.SecretBase,
                         ),
                     )
                     and task_state is not None
diff --git a/src/prefect/utilities/tasks.py b/src/prefect/utilities/tasks.py
index 19d4aad7cb..2915aaf251 100644
--- a/src/prefect/utilities/tasks.py
+++ b/src/prefect/utilities/tasks.py
@@ -97,20 +97,16 @@ 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 and subflow. Constant arguments
-    #   are added later as needed.
+    # - Add all non-constant arguments to the flow. Constant arguments are
+    #   added later as needed.
     def preprocess(a: Any) -> "prefect.Task":
-        # 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
-
+        a2 = as_task(a, flow=flow2)
+        is_mapped = not isinstance(a, prefect.utilities.edges.unmapped)
+        is_constant = isinstance(a2, Constant)
         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/engine/test_flow_runner.py b/tests/engine/test_flow_runner.py
index 9459b16ba2..1ffdfa5687 100644
--- a/tests/engine/test_flow_runner.py
+++ b/tests/engine/test_flow_runner.py
@@ -281,6 +281,23 @@ def test_secrets_dynamically_pull_from_context():
     assert flow_state.is_successful()
 
 
+def test_secrets_are_rerun_on_restart():
+    @prefect.task
+    def identity(x):
+        return x
+
+    with Flow("test") as flow:
+        secret = PrefectSecret("key")
+        val = identity(secret)
+
+    with prefect.context(secrets={"key": "val"}):
+        state = FlowRunner(flow=flow).run(
+            task_states={secret: Success()}, return_tasks=[val]
+        )
+    assert state.is_successful()
+    assert state.result[val].result == "val"
+
+
 def test_flow_runner_doesnt_return_by_default():
     flow = Flow(name="test")
     task1 = SuccessTask()
diff --git a/tests/utilities/test_tasks.py b/tests/utilities/test_tasks.py
index d8a8700e23..62ebe61107 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, resource_manager
+from prefect import Flow, Task, case, Parameter
 from prefect.engine.flow_runner import FlowRunner
 from prefect.engine.state import Paused, Resume
 from prefect.utilities import tasks, edges
@@ -302,50 +302,24 @@ class TestApplyMap:
         assert res == sol
 
     def test_apply_map_inside_case_statement_works(self):
-        def func(x, a):
-            return add(x, 1), add(x, a)
+        def func(x):
+            return add(x, 1), add(x, 2)
 
         with Flow("test") as flow:
             branch = Parameter("branch")
             with case(branch, True):
-                a = inc(1)
-                b, c = apply_map(func, range(4), edges.unmapped(a))
-                d = add.map(b, c)
+                a, b = apply_map(func, range(4))
+                c = add.map(a, b)
 
         state = flow.run(branch=True)
-        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]
+        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]
 
         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
@@ -406,7 +380,7 @@ class TestAsTask:
         assert res.result[val].result == obj
 
     def test_as_task_toggles_constants(self):
-        with Flow("test"):
+        with Flow("test") as f:
             t = tasks.as_task(4)
 
         assert isinstance(t, Task)
@@ -435,7 +409,7 @@ class TestAsTask:
         ],
     )
     def test_nested_collections_of_mixed_constants_are_not_constants(self, val):
-        with Flow("test"):
+        with Flow("test") as f:
             task = tasks.as_task(val)
         assert not isinstance(task, Constant)