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:

Resolving patch:
diff --git a/src/prefect/engine/flow_runner.py b/src/prefect/engine/flow_runner.py
--- a/src/prefect/engine/flow_runner.py
+++ b/src/prefect/engine/flow_runner.py
@@ -427,8 +427,8 @@ def extra_context(task: Task, task_index: int = None) -> dict:
                 ):
                     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 @@ def extra_context(task: Task, task_index: int = None) -> dict:
                         (
                             prefect.tasks.core.resource_manager.ResourceSetupTask,
                             prefect.tasks.core.resource_manager.ResourceCleanupTask,
+                            prefect.tasks.secrets.SecretBase,
                         ),
                     )
                     and task_state is not None


Test patch:
diff --git a/tests/engine/test_flow_runner.py b/tests/engine/test_flow_runner.py
--- 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()