Repository: dagster-io/dagster
Issue: Re-execution from a step with downstream dynamic steps produces an error

https://dagster.slack.com/archives/C01U954MEER/p1640105707446800

If you try to re-execute the following subsets, you'll get an error:

```
+hello_dynamic
generate_subtasks*
generate_subtasks+
generate_subtasks, hello_dynamic
```

The error:

```
dagster.core.errors.DagsterExecutionStepNotFoundError: Can not build subset plan from unknown steps: hello_dynamic[subtask_1], hello_dynamic[subtask_2], hello_dynamic[subtask_3], hello_dynamic[subtask_4], hello_dynamic[subtask_5]
  File "/Users/owen/src/dagster/python_modules/dagster/dagster/grpc/impl.py", line 354, in get_external_execution_plan_snapshot
    create_execution_plan(
  File "/Users/owen/src/dagster/python_modules/dagster/dagster/core/execution/api.py", line 745, in create_execution_plan
    return ExecutionPlan.build(
  File "/Users/owen/src/dagster/python_modules/dagster/dagster/core/execution/plan/plan.py", line 973, in build
    return plan_builder.build()
  File "/Users/owen/src/dagster/python_modules/dagster/dagster/core/execution/plan/plan.py", line 227, in build
    plan = plan.build_subset_plan(
  File "/Users/owen/src/dagster/python_modules/dagster/dagster/core/execution/plan/plan.py", line 764, in build_subset_plan
    raise DagsterExecutionStepNotFoundError(
```

These subsets re-execute without issue:

```
hello_dynamic
hello_dynamic*
hello_dynamic[subtask_1]
```


Changed files:
python_modules/dagster/dagster/core/execution/backfill.py
python_modules/dagster/dagster_tests/execution_tests/dynamic_tests/test_dynamic_execution.py
python_modules/dagster/dagster/core/execution/plan/state.py
python_modules/dagster/dagster/core/execution/plan/resume_retry.py
python_modules/dagster/dagster/core/execution/api.py
python_modules/dagster-graphql/dagster_graphql/implementation/execution/run_lifecycle.py

diff --git a/python_modules/dagster-graphql/dagster_graphql/implementation/execution/run_lifecycle.py b/python_modules/dagster-graphql/dagster_graphql/implementation/execution/run_lifecycle.py
index 00270cde1e..d648446d2b 100644
--- a/python_modules/dagster-graphql/dagster_graphql/implementation/execution/run_lifecycle.py
+++ b/python_modules/dagster-graphql/dagster_graphql/implementation/execution/run_lifecycle.py
@@ -29,7 +29,6 @@ def compute_step_keys_to_execute(graphene_info, execution_params):
         if execution_params.execution_metadata.parent_run_id and execution_params.step_keys:
             known_state = KnownExecutionState.for_reexecution(
                 instance.all_logs(execution_params.execution_metadata.parent_run_id),
-                execution_params.step_keys,
             )
 
         return execution_params.step_keys, known_state
diff --git a/python_modules/dagster/dagster/core/execution/api.py b/python_modules/dagster/dagster/core/execution/api.py
index 3fc2a56fcf..e6ae0a3f06 100644
--- a/python_modules/dagster/dagster/core/execution/api.py
+++ b/python_modules/dagster/dagster/core/execution/api.py
@@ -1025,7 +1025,7 @@ def _resolve_reexecute_step_selection(
         run_config,
         mode,
         step_keys_to_execute=list(step_keys_to_execute),
-        known_state=KnownExecutionState.for_reexecution(parent_logs, step_keys_to_execute),
+        known_state=KnownExecutionState.for_reexecution(parent_logs),
         tags=parent_pipeline_run.tags,
     )
     return execution_plan
diff --git a/python_modules/dagster/dagster/core/execution/backfill.py b/python_modules/dagster/dagster/core/execution/backfill.py
index 18edca0f2a..93617021e9 100644
--- a/python_modules/dagster/dagster/core/execution/backfill.py
+++ b/python_modules/dagster/dagster/core/execution/backfill.py
@@ -249,7 +249,6 @@ def create_backfill_run(
         if last_run and last_run.status == PipelineRunStatus.SUCCESS:
             known_state = KnownExecutionState.for_reexecution(
                 instance.all_logs(parent_run_id),
-                step_keys_to_execute,
             )
         else:
             known_state = None
diff --git a/python_modules/dagster/dagster/core/execution/plan/resume_retry.py b/python_modules/dagster/dagster/core/execution/plan/resume_retry.py
index 123ab1387e..7aa831f70b 100644
--- a/python_modules/dagster/dagster/core/execution/plan/resume_retry.py
+++ b/python_modules/dagster/dagster/core/execution/plan/resume_retry.py
@@ -140,4 +140,4 @@ def get_retry_steps_from_parent_run(
         step_handle.to_key() for step_set in to_retry.values() for step_handle in step_set
     ]
 
-    return steps_to_retry, KnownExecutionState.for_reexecution(parent_run_logs, steps_to_retry)
+    return steps_to_retry, KnownExecutionState.for_reexecution(parent_run_logs)
diff --git a/python_modules/dagster/dagster/core/execution/plan/state.py b/python_modules/dagster/dagster/core/execution/plan/state.py
index 39565d2e5e..f306ca566a 100644
--- a/python_modules/dagster/dagster/core/execution/plan/state.py
+++ b/python_modules/dagster/dagster/core/execution/plan/state.py
@@ -108,16 +108,9 @@ def derive_from_logs(logs: List[EventLogEntry]) -> "KnownExecutionState":
         )
 
     @staticmethod
-    def for_reexecution(
-        parent_run_logs: List[EventLogEntry], step_keys_to_execute: List[str]
-    ) -> "KnownExecutionState":
+    def for_reexecution(parent_run_logs: List[EventLogEntry]) -> "KnownExecutionState":
         """
-        Copy over dynamic mappings from previous run, but drop any for steps that we intend to re-execute
+        Copy over dynamic mappings from previous run
         """
         parent_state = KnownExecutionState.derive_from_logs(parent_run_logs)
-        dynamic_mappings_to_use = {
-            step_key: parent_state.dynamic_mappings[step_key]
-            for step_key in parent_state.dynamic_mappings.keys()
-            if step_key not in step_keys_to_execute
-        }
-        return KnownExecutionState({}, dynamic_mappings_to_use)
+        return KnownExecutionState({}, parent_state.dynamic_mappings)
diff --git a/python_modules/dagster/dagster_tests/execution_tests/dynamic_tests/test_dynamic_execution.py b/python_modules/dagster/dagster_tests/execution_tests/dynamic_tests/test_dynamic_execution.py
index 4575a0ec1c..b48075a2a3 100644
--- a/python_modules/dagster/dagster_tests/execution_tests/dynamic_tests/test_dynamic_execution.py
+++ b/python_modules/dagster/dagster_tests/execution_tests/dynamic_tests/test_dynamic_execution.py
@@ -1,6 +1,7 @@
 import pytest
 
 from dagster import (
+    DynamicOut,
     DynamicOutput,
     DynamicOutputDefinition,
     Field,
@@ -9,11 +10,14 @@
     OutputDefinition,
     composite_solid,
     execute_pipeline,
+    job,
+    op,
     pipeline,
     reconstructable,
     solid,
 )
 from dagster.core.errors import DagsterExecutionStepNotFoundError
+from dagster.core.events import DagsterEventType
 from dagster.core.execution.api import create_execution_plan, reexecute_pipeline
 from dagster.core.execution.plan.state import KnownExecutionState
 from dagster.core.test_utils import default_mode_def_for_test, instance_for_test
@@ -283,20 +287,140 @@ def test_fan_out_in_out_in(run_config):
         assert empty_result.result_for_solid("sum_numbers").output_value() == 0
 
 
-def test_bad_step_selection():
+def define_real_dynamic_job():
+    @op(config_schema=list, out=DynamicOut(int))
+    def generate_subtasks(context):
+        for num in context.op_config:
+            yield DynamicOutput(num, mapping_key=str(num))
+
+    @op
+    def subtask(input_number: int):
+        print(f"Found {input_number} cereals")
+        return input_number
+
+    @job
+    def real_dynamic_job():
+        generate_subtasks().map(subtask)
+
+    return real_dynamic_job
+
+
+def test_select_dynamic_step_with_non_static_mapping():
+    with instance_for_test() as instance:
+        result_0 = execute_pipeline(
+            reconstructable(define_real_dynamic_job),
+            instance=instance,
+            run_config={"ops": {"generate_subtasks": {"config": [0, 2, 4]}}},
+        )
+        assert result_0.success
+
+        result_1 = reexecute_pipeline(
+            reconstructable(define_real_dynamic_job),
+            parent_run_id=result_0.run_id,
+            instance=instance,
+            step_selection=["generate_subtasks", "subtask[4]"],
+            run_config={"ops": {"generate_subtasks": {"config": [0, 2, 4]}}},
+        )
+        assert result_1.success
+        keys_1 = result_1.events_by_step_key.keys()
+        assert "generate_subtasks" in keys_1
+        assert "subtask[4]" in keys_1
+        assert "subtask[0]" not in keys_1
+        assert "subtask[2]" not in keys_1
+
+        result_2 = reexecute_pipeline(
+            reconstructable(define_real_dynamic_job),
+            parent_run_id=result_0.run_id,
+            instance=instance,
+            step_selection=["generate_subtasks", "subtask[4]"],
+            run_config={"ops": {"generate_subtasks": {"config": [1, 3, 5]}}},
+        )
+        # subtask[4] is skipped - bc generate_subtasks doesn't yield "4"
+        # subtask[1], subtask[3], subtask[5] have been generated but not selected to execute
+        assert result_2.success
+        keys_2 = result_2.events_by_step_key.keys()
+        assert "generate_subtasks" in keys_2
+        assert "subtask[4]" in keys_2
+        assert (
+            result_2.events_by_step_key["subtask[4]"][0].event_type == DagsterEventType.STEP_SKIPPED
+        )
+        assert "subtask[1]" not in keys_2
+        assert "subtask[3]" not in keys_2
+        assert "subtask[5]" not in keys_2
+
+        result_3 = reexecute_pipeline(
+            reconstructable(define_real_dynamic_job),
+            parent_run_id=result_0.run_id,
+            instance=instance,
+            step_selection=["generate_subtasks", "subtask[4]"],
+            run_config={"ops": {"generate_subtasks": {"config": [1, 2, 3, 4, 5]}}},
+        )
+        # subtask[4] is executed
+        # subtask[0], subtask[1], subtask[2], subtask[3], subtask[5] have been generated but not
+        # selected to execute
+        assert result_3.success
+        keys_3 = result_3.events_by_step_key.keys()
+        assert "generate_subtasks" in keys_3
+        assert "subtask[4]" in keys_3
+        assert DagsterEventType.STEP_SUCCESS in [
+            e.event_type for e in result_3.events_by_step_key["subtask[4]"]
+        ]
+        assert "subtask[0]" not in keys_3
+        assert "subtask[1]" not in keys_3
+        assert "subtask[2]" not in keys_3
+        assert "subtask[3]" not in keys_3
+        assert "subtask[5]" not in keys_3
+
+
+def test_select_dynamic_step_and_downstream():
     with instance_for_test() as instance:
         result_1 = execute_pipeline(dynamic_pipeline, instance=instance)
         assert result_1.success
 
-        # this exact error could be improved, but it should fail if you try to select
-        # both the dynamic outputting step key and something resolved by it in the previous run
-        with pytest.raises(DagsterExecutionStepNotFoundError):
-            reexecute_pipeline(
-                dynamic_pipeline,
-                parent_run_id=result_1.run_id,
-                instance=instance,
-                step_selection=["emit", "multiply_by_two[1]"],
-            )
+        result_2 = reexecute_pipeline(
+            dynamic_pipeline,
+            parent_run_id=result_1.run_id,
+            instance=instance,
+            step_selection=["emit*"],
+        )
+        assert result_2.success
+
+        keys_2 = result_2.events_by_step_key.keys()
+        assert "multiply_inputs[0]" in keys_2
+        assert "multiply_inputs[1]" in keys_2
+        assert "multiply_inputs[2]" in keys_2
+        assert "multiply_by_two[0]" in keys_2
+        assert "multiply_by_two[1]" in keys_2
+        assert "multiply_by_two[2]" in keys_2
+        assert result_2.result_for_solid("double_total").output_value() == 120
+
+        result_3 = reexecute_pipeline(
+            dynamic_pipeline,
+            parent_run_id=result_1.run_id,
+            instance=instance,
+            step_selection=["emit+"],
+        )
+        assert result_3.success
+
+        keys_3 = result_3.events_by_step_key.keys()
+        assert "multiply_inputs[0]" in keys_3
+        assert "multiply_inputs[1]" in keys_3
+        assert "multiply_inputs[2]" in keys_3
+        assert "multiply_by_two[0]" not in keys_3
+
+        result_4 = reexecute_pipeline(
+            dynamic_pipeline,
+            parent_run_id=result_1.run_id,
+            instance=instance,
+            step_selection=["emit_ten", "multiply_inputs[1]", "multiply_by_two[1]"],
+        )
+
+        keys_4 = result_4.events_by_step_key.keys()
+        assert "emit_ten" in keys_4
+        assert "multiply_inputs[1]" in keys_4
+        assert "multiply_by_two[1]" in keys_4
+        assert "emit" not in keys_4
+        assert "multiply_inputs[0]" not in keys_4
 
 
 @pytest.mark.parametrize(