Repository: PrefectHQ/prefect

Issue:
Question: Retry on mapped task
## Description
*A clear description of the bug*

Just experimenting and trying to figure out if I am doing something incorrectly or if it isn't possible. In the case where I have mapped task that has a retry attached to it. If any instance of that mapped task fails, it causes all downstream tasks to fail.


## Expected Behavior
*What did you expect to happen instead?*

When an instance of the mapped task fails, it will retry until all instances are successful, or, if `max_retries` is met and still not successful, to fail the rest of the `Flow`.


## Reproduction
*A minimal example that exhibits the behavior.*

*Code:*
```python
from datetime import timedelta
from random import random

from prefect import Flow, context, task


@task(max_retries=3, retry_delay=timedelta(seconds=2))
def add_ten(data: float):
    if random() < 0.33:
        raise ValueError("Random error")
    return data + 10


@task
def sum_data(d):
    log = context.get("logger")
    log.info(f"Sum of data {sum(d) / len(d)}")


with Flow("Mapped Tasks Should Retry on Fail") as flow:
    d = add_ten.map(data=[1, 2, 3])
    sum_data(d)

flow.run()
```

*Output:*
```
[2019-08-30 22:54:31,086] INFO - prefect.FlowRunner | Beginning Flow run for 'Mapped Tasks Should Retry on Fail'
[2019-08-30 22:54:31,086] INFO - prefect.FlowRunner | Starting flow run.
[2019-08-30 22:54:31,091] INFO - prefect.TaskRunner | Task '1': Starting task run...
[2019-08-30 22:54:31,091] INFO - prefect.TaskRunner | Task '1': finished task run for task with final state: 'Success'
[2019-08-30 22:54:31,092] INFO - prefect.TaskRunner | Task '3': Starting task run...
[2019-08-30 22:54:31,092] INFO - prefect.TaskRunner | Task '3': finished task run for task with final state: 'Success'
[2019-08-30 22:54:31,092] INFO - prefect.TaskRunner | Task '2': Starting task run...
[2019-08-30 22:54:31,093] INFO - prefect.TaskRunner | Task '2': finished task run for task with final state: 'Success'
[2019-08-30 22:54:31,093] INFO - prefect.TaskRunner | Task 'List': Starting task run...
[2019-08-30 22:54:31,093] INFO - prefect.TaskRunner | Task 'List': finished task run for task with final state: 'Success'
[2019-08-30 22:54:31,094] INFO - prefect.TaskRunner | Task 'add_ten': Starting task run...
[2019-08-30 22:54:31,095] INFO - prefect.TaskRunner | Task 'add_ten[0]': Starting task run...
[2019-08-30 22:54:31,095] INFO - prefect.TaskRunner | Task 'add_ten[0]': finished task run for task with final state: 'Success'
[2019-08-30 22:54:31,095] INFO - prefect.TaskRunner | Task 'add_ten[2]': Starting task run...
[2019-08-30 22:54:31,096] INFO - prefect.TaskRunner | Task 'add_ten[2]': finished task run for task with final state: 'Success'
[2019-08-30 22:54:31,096] INFO - prefect.TaskRunner | Task 'add_ten[1]': Starting task run...
[2019-08-30 22:54:31,096] INFO - prefect.TaskRunner | Unexpected error: ValueError('Random error')
[2019-08-30 22:54:31,096] INFO - prefect.TaskRunner | Task 'add_ten[1]': finished task run for task with final state: 'Retrying'
[2019-08-30 22:54:31,097] INFO - prefect.TaskRunner | Task 'add_ten': finished task run for task with final state: 'Mapped'
[2019-08-30 22:54:31,097] INFO - prefect.TaskRunner | Task 'sum_data': Starting task run...
[2019-08-30 22:54:31,097] INFO - prefect.TaskRunner | Task 'sum_data': finished task run for task with final state: 'TriggerFailed'
[2019-08-30 22:54:31,098] INFO - prefect.FlowRunner | Flow run FAILED: some reference tasks failed.
[2019-08-30 22:54:31,098] INFO - prefect.Flow | Waiting for next available Task run at 2019-08-30T22:54:33.096859+00:00
```

To expand on what I believe the ideal behavior would be: In the above output we see that the `add_ten` task instance for index: `1` failed due to `Random Error` and caused the downstream tasks to also fail. But, the definition for the `add_ten` task has three retries allocated. While the `add_ten` task instances for indices: `0` and `2` succeeded on the first try, the `add_ten` task should be retried prior to continuing to the downstream `sum_data` task.

## Environment
*Any additional information about your environment*

OS: macOS 10.14.6
Env Manager: conda 4.6.14
Python: 3.7.4

prefect installed with pip
prefect: 0.6.1



Discussion:

Resolving patch:
diff --git a/src/prefect/engine/task_runner.py b/src/prefect/engine/task_runner.py
--- a/src/prefect/engine/task_runner.py
+++ b/src/prefect/engine/task_runner.py
@@ -351,7 +351,17 @@ def check_upstream_finished(
         Raises:
             - ENDRUN: if upstream tasks are not finished.
         """
-        if not all(s.is_finished() for s in upstream_states.values()):
+        all_states = set()  # type: Set[State]
+        for edge, upstream_state in upstream_states.items():
+            # if the upstream state is Mapped, and this task is also mapped,
+            # we want each individual child to determine if it should
+            # proceed or not based on its upstream parent in the mapping
+            if isinstance(upstream_state, Mapped) and not edge.mapped:
+                all_states.update(upstream_state.map_states)
+            else:
+                all_states.add(upstream_state)
+
+        if not all(s.is_finished() for s in all_states):
             self.logger.debug(
                 "Task '{name}': not all upstream states are finished; ending run.".format(
                     name=prefect.context.get("task_full_name", self.task.name)


Test patch:
diff --git a/tests/core/test_task_map.py b/tests/core/test_task_map.py
--- a/tests/core/test_task_map.py
+++ b/tests/core/test_task_map.py
@@ -594,6 +594,67 @@ def take_sum(x):
     assert isinstance(state.result[s], prefect.engine.state.TriggerFailed)
 
 
+@pytest.mark.parametrize(
+    "executor", ["local", "sync", "mproc", "mthread"], indirect=True
+)
+def test_reduce_task_properly_applies_trigger_across_all_mapped_states(executor):
+    """
+    The `take_sum` task reduces over the `div` task, which is going to return one Success
+    and one Retrying. The `take_sum` should fail its upstream finished check.
+    """
+
+    @prefect.task
+    def ll():
+        return [0, 1]
+
+    @prefect.task(max_retries=5, retry_delay=datetime.timedelta(hours=1))
+    def div(x):
+        return 1 / x
+
+    @prefect.task
+    def take_sum(x):
+        return sum(x)
+
+    with Flow(name="test") as f:
+        d = div.map(ll)
+        t = div.map(d)
+        s = take_sum(d)
+
+    state = FlowRunner(flow=f).run(executor=executor, return_tasks=[s, t])
+    assert state.is_running()
+    assert state.result[s].is_pending()
+    assert state.result[t].is_mapped()
+    assert state.result[t].map_states[0].is_pending()
+    assert all(s.is_successful() for s in state.result[t].map_states[1:])
+
+
+@pytest.mark.parametrize(
+    "executor", ["local", "sync", "mproc", "mthread"], indirect=True
+)
+def test_reduce_task_properly_applies_trigger_across_all_mapped_states_for_deep_pipelines(
+    executor
+):
+    @prefect.task
+    def ll():
+        return [0, 1]
+
+    @prefect.task(max_retries=5, retry_delay=datetime.timedelta(hours=1))
+    def div(x):
+        return 1 / x
+
+    @prefect.task
+    def take_sum(x):
+        return sum(x)
+
+    with Flow(name="test") as f:
+        d = div.map(ll)
+        s = take_sum(d)
+
+    state = FlowRunner(flow=f).run(executor=executor, return_tasks=[s])
+    assert state.is_running()
+    assert state.result[s].is_pending()
+
+
 @pytest.mark.parametrize(
     "executor", ["local", "sync", "mproc", "mthread"], indirect=True
 )
diff --git a/tests/engine/test_task_runner.py b/tests/engine/test_task_runner.py
--- a/tests/engine/test_task_runner.py
+++ b/tests/engine/test_task_runner.py
@@ -461,6 +461,26 @@ def test_raises_with_one_unfinished(self):
                 state=state, upstream_states={1: Success(), 2: Running()}
             )
 
+    def test_raises_if_mapped_upstream_retrying(self):
+        state = Pending()
+        task = Task()
+        with pytest.raises(ENDRUN) as exc:
+            edge = Edge(1, 2, mapped=False)
+            new_state = TaskRunner(task).check_upstream_finished(
+                state=state,
+                upstream_states={edge: Mapped(map_states=[Retrying(), Success()])},
+            )
+
+    def test_doesnt_raise_if_mapped_upstream_complete(self):
+        state = Pending()
+        task = Task()
+        edge = Edge(1, 2, mapped=False)
+        new_state = TaskRunner(task).check_upstream_finished(
+            state=state,
+            upstream_states={edge: Mapped(map_states=[Failed(), Success()])},
+        )
+        assert new_state is state
+
 
 class TestCheckUpstreamSkipped:
     def test_empty(self):