Repository: dagster-io/dagster
Pull request: explicitly include asset check selection on backfill runs
## Summary & Motivation

Fixes https://github.com/dagster-io/dagster/issues/25328.

Currently, when asset backfills launch, all the checks for the selected assets are executed. However, these checks aren't explicitly included on the `asset_check_selection` property of the `DagsterRun`.  A consequence of this is that, when the run is re-executed, the asset checks aren't included in the subset job, which causes the error described in the issue linked above.

This problem points to some deeper mismatched assumptions that would be good to fix as well.  I.e. it seems like we should either:
- be stricter about the selections included on runs in general
- be looser about how we construct subsetted jobs when doing re-execution

## How I Tested These Changes

Ran through the reproduction steps listed on https://github.com/dagster-io/dagster/issues/25328 and verified that I could successfully re-execute the asset check for the run inside the backfill.

Also added a unit test to verify that the check selection is included inside the run request created by the backfill.

Ideally this would include a larger unit test that does something closer to the manual reproduction described above. I explored a couple routes for this, but alas was not able to find one that worked before my PTO.

## Changelog

Fixed a bug that caused `DagsterExecutionStepNotFoundError` errors when trying to execute an asset check step of a run launched by a backfill.


Problem statement:
Error when re-executing asset check from backfill run
### Reproduction
Use this code location:

```python
from dagster import BackfillPolicy, DailyPartitionsDefinition, Definitions, asset, asset_check

pdef = DailyPartitionsDefinition("2024-10-01")

@asset(partitions_def=pdef, backfill_policy=BackfillPolicy.single_run())
def adummy() -> None:
    pass

@asset_check(asset=adummy)
def acheck() -> None:
    pass

defs = Definitions(assets=[adummy], asset_checks=[acheck])
```

- Launch a backfill of two or more partitions
- Wait for it to complete
- Click into the backfill run
- Select the step corresponding to the asset check
- Click the re-execute button
- Observe this error:

```python
dagster._core.errors.DagsterExecutionStepNotFoundError: Can not build subset plan from unknown step: adummy_acheck

  File "/Users/sryza/dagster/python_modules/dagster/dagster/_grpc/server.py", line 578, in ExecutionPlanSnapshot
    execution_plan_snapshot_or_error = get_external_execution_plan_snapshot(
  File "/Users/sryza/dagster/python_modules/dagster/dagster/_grpc/impl.py", line 509, in get_external_execution_plan_snapshot
    create_execution_plan(
  File "/Users/sryza/dagster/python_modules/dagster/dagster/_core/execution/api.py", line 716, in create_execution_plan
    return ExecutionPlan.build(
  File "/Users/sryza/dagster/python_modules/dagster/dagster/_core/execution/plan/plan.py", line 955, in build
    ).build()
  File "/Users/sryza/dagster/python_modules/dagster/dagster/_core/execution/plan/plan.py", line 224, in build
    plan = plan.build_subset_plan(
  File "/Users/sryza/dagster/python_modules/dagster/dagster/_core/execution/plan/plan.py", line 851, in build_subset_plan
    raise DagsterExecutionStepNotFoundError(
```

Interestingly, this does not occur in the case that the original run is a non-backfill run for a single partition.

### Investigation

I did some investigation into what's going on under the covers and discovered that:
- The `JobSubsetSelector` passed to `get_remote_job_or_raise` has an empty `asset_check_selection`.
- This value comes from the frontend, which gets it from the `DagsterRun` object that's being re-executed.
- That `DagsterRun` object has no `asset_check_selection` on it probably because the asset check wasn't explicitly selected when the run was launched.
- In the non-backfill case, the frontend explicitly provides the asset check selection

### Possible solution

In the backfill run creation logic, include asset checks for selected assets in the selection.

Hints from issue discussion:

Commit message: add efficient override for get_check_keys_for_assets

Patch to review:
diff --git a/python_modules/dagster/dagster/_core/definitions/automation_tick_evaluation_context.py b/python_modules/dagster/dagster/_core/definitions/automation_tick_evaluation_context.py
index 2dbe9169e5a0e..c6bdc353e0006 100644
--- a/python_modules/dagster/dagster/_core/definitions/automation_tick_evaluation_context.py
+++ b/python_modules/dagster/dagster/_core/definitions/automation_tick_evaluation_context.py
@@ -403,17 +403,25 @@ def build_run_requests_with_backfill_policies(
         partition_keys,
         backfill_policy,
     ), asset_keys in assets_to_reconcile_by_partitions_def_partition_keys_backfill_policy.items():
+        asset_check_keys = asset_graph.get_check_keys_for_assets(asset_keys)
         if partitions_def is None and partition_keys is not None:
             check.failed("Partition key provided for unpartitioned asset")
         if partitions_def is not None and partition_keys is None:
             check.failed("Partition key missing for partitioned asset")
         if partitions_def is None and partition_keys is None:
             # non partitioned assets will be backfilled in a single run
-            run_requests.append(RunRequest(asset_selection=list(asset_keys), tags={}))
+            run_requests.append(
+                RunRequest(
+                    asset_selection=list(asset_keys),
+                    asset_check_keys=list(asset_check_keys),
+                    tags={},
+                )
+            )
         else:
             run_requests.extend(
                 _build_run_requests_with_backfill_policy(
                     list(asset_keys),
+                    list(asset_check_keys),
                     check.not_none(backfill_policy),
                     check.not_none(partition_keys),
                     check.not_none(partitions_def),
@@ -426,6 +434,7 @@ def build_run_requests_with_backfill_policies(
 
 def _build_run_requests_with_backfill_policy(
     asset_keys: Sequence[AssetKey],
+    asset_check_keys: Sequence[AssetCheckKey],
     backfill_policy: BackfillPolicy,
     partition_keys: FrozenSet[str],
     partitions_def: PartitionsDefinition,
@@ -444,6 +453,7 @@ def _build_run_requests_with_backfill_policy(
             run_requests.append(
                 _build_run_request_for_partition_key_range(
                     asset_keys=list(asset_keys),
+                    asset_check_keys=list(asset_check_keys),
                     partition_range_start=partition_key_range.start,
                     partition_range_end=partition_key_range.end,
                     run_tags=tags,
@@ -453,6 +463,7 @@ def _build_run_requests_with_backfill_policy(
             run_requests.extend(
                 _build_run_requests_for_partition_key_range(
                     asset_keys=list(asset_keys),
+                    asset_check_keys=list(asset_check_keys),
                     partitions_def=partitions_def,
                     partition_key_range=partition_key_range,
                     max_partitions_per_run=check.int_param(
@@ -467,6 +478,7 @@ def _build_run_requests_with_backfill_policy(
 
 def _build_run_requests_for_partition_key_range(
     asset_keys: Sequence[AssetKey],
+    asset_check_keys: Sequence[AssetCheckKey],
     partitions_def: PartitionsDefinition,
     partition_key_range: PartitionKeyRange,
     max_partitions_per_run: int,
@@ -492,7 +504,11 @@ def _build_run_requests_for_partition_key_range(
         partition_chunk_end_key = partition_keys[partition_chunk_end_index]
         run_requests.append(
             _build_run_request_for_partition_key_range(
-                asset_keys, partition_chunk_start_key, partition_chunk_end_key, run_tags
+                asset_keys,
+                asset_check_keys,
+                partition_chunk_start_key,
+                partition_chunk_end_key,
+                run_tags,
             )
         )
         partition_chunk_start_index = partition_chunk_end_index + 1
@@ -501,6 +517,7 @@ def _build_run_requests_for_partition_key_range(
 
 def _build_run_request_for_partition_key_range(
     asset_keys: Sequence[AssetKey],
+    asset_check_keys: Sequence[AssetCheckKey],
     partition_range_start: str,
     partition_range_end: str,
     run_tags: Dict[str, str],
@@ -512,4 +529,9 @@ def _build_run_request_for_partition_key_range(
         ASSET_PARTITION_RANGE_END_TAG: partition_range_end,
     }
     partition_key = partition_range_start if partition_range_start == partition_range_end else None
-    return RunRequest(asset_selection=asset_keys, partition_key=partition_key, tags=tags)
+    return RunRequest(
+        asset_selection=asset_keys,
+        partition_key=partition_key,
+        tags=tags,
+        asset_check_keys=asset_check_keys,
+    )
diff --git a/python_modules/dagster/dagster/_core/definitions/base_asset_graph.py b/python_modules/dagster/dagster/_core/definitions/base_asset_graph.py
index 5d3b6e05a47ed..500eacca0cc4c 100644
--- a/python_modules/dagster/dagster/_core/definitions/base_asset_graph.py
+++ b/python_modules/dagster/dagster/_core/definitions/base_asset_graph.py
@@ -355,6 +355,15 @@ def root_executable_asset_keys(self) -> AbstractSet[AssetKey]:
     @abstractmethod
     def asset_check_keys(self) -> AbstractSet[AssetCheckKey]: ...
 
+    def get_check_keys_for_assets(
+        self, asset_keys: AbstractSet[AssetKey]
+    ) -> AbstractSet[AssetCheckKey]:
+        return {
+            asset_check_key
+            for asset_check_key in self.asset_check_keys
+            if asset_check_key.asset_key in asset_keys
+        }
+
     @cached_property
     def all_partitions_defs(self) -> Sequence[PartitionsDefinition]:
         return sorted(
diff --git a/python_modules/dagster/dagster/_core/definitions/remote_asset_graph.py b/python_modules/dagster/dagster/_core/definitions/remote_asset_graph.py
index ca4d7f5687335..49c8cce4e5573 100644
--- a/python_modules/dagster/dagster/_core/definitions/remote_asset_graph.py
+++ b/python_modules/dagster/dagster/_core/definitions/remote_asset_graph.py
@@ -445,6 +445,16 @@ def _asset_check_nodes_by_asset_key(self) -> Mapping[AssetKey, Sequence[RemoteAs
     def get_checks_for_asset(self, asset_key: AssetKey) -> Sequence[RemoteAssetCheckNode]:
         return self._asset_check_nodes_by_asset_key.get(asset_key, [])
 
+    def get_check_keys_for_assets(
+        self, asset_keys: AbstractSet[AssetKey]
+    ) -> AbstractSet[AssetCheckKey]:
+        return set().union(
+            *(
+                {check.asset_check.key for check in self.get_checks_for_asset(asset_key)}
+                for asset_key in asset_keys
+            )
+        )
+
     @cached_property
     def asset_check_keys(self) -> AbstractSet[AssetCheckKey]:
         return set(self.remote_asset_check_nodes_by_key.keys())
diff --git a/python_modules/dagster/dagster_tests/core_tests/execution_tests/test_asset_backfill.py b/python_modules/dagster/dagster_tests/core_tests/execution_tests/test_asset_backfill.py
index 79dba96254df8..7ad4fb7bf76d1 100644
--- a/python_modules/dagster/dagster_tests/core_tests/execution_tests/test_asset_backfill.py
+++ b/python_modules/dagster/dagster_tests/core_tests/execution_tests/test_asset_backfill.py
@@ -16,10 +16,12 @@
 
 import pytest
 from dagster import (
+    AssetCheckResult,
     AssetIn,
     AssetKey,
     AssetOut,
     AssetsDefinition,
+    BackfillPolicy,
     DagsterInstance,
     DagsterRunStatus,
     DailyPartitionsDefinition,
@@ -33,6 +35,7 @@
     TimeWindowPartitionMapping,
     WeeklyPartitionsDefinition,
     asset,
+    asset_check,
     materialize,
     multi_asset,
 )
@@ -1795,3 +1798,46 @@ def foo_child():
         "fake_id", asset_backfill_data, asset_graph, instance, assets_by_repo_name
     )
     assert asset_backfill_data.requested_subset == asset_backfill_data.target_subset
+
+
+def test_asset_backfill_with_asset_check():
+    instance = DagsterInstance.ephemeral()
+    partitions_def = DailyPartitionsDefinition("2023-10-01")
+
+    @asset(partitions_def=partitions_def, backfill_policy=BackfillPolicy.single_run())
+    def foo():
+        pass
+
+    @asset_check(asset=foo)
+    def foo_check():
+        return AssetCheckResult(passed=True)
+
+    assets_by_repo_name = {"repo": [foo, foo_check]}
+    asset_graph = get_asset_graph(assets_by_repo_name)
+
+    target_partitions_subset = partitions_def.empty_subset().with_partition_key_range(
+        partitions_def, PartitionKeyRange("2023-11-01", "2023-11-03")
+    )
+    asset_backfill_data = AssetBackfillData.from_asset_graph_subset(
+        asset_graph_subset=AssetGraphSubset(
+            partitions_subsets_by_asset_key={foo.key: target_partitions_subset}
+        ),
+        dynamic_partitions_store=MagicMock(),
+        backfill_start_timestamp=create_datetime(2023, 12, 5, 0, 0, 0).timestamp(),
+    )
+    assert set(asset_backfill_data.target_subset.iterate_asset_partitions()) == {
+        AssetKeyPartitionKey(foo.key, "2023-11-01"),
+        AssetKeyPartitionKey(foo.key, "2023-11-02"),
+        AssetKeyPartitionKey(foo.key, "2023-11-03"),
+    }
+
+    result = execute_asset_backfill_iteration_consume_generator(
+        backfill_id="fake_id",
+        asset_backfill_data=asset_backfill_data,
+        asset_graph=asset_graph,
+        instance=instance,
+    )
+    assert len(result.run_requests) == 1
+    run_request = result.run_requests[0]
+    assert run_request.asset_selection == [foo.key]
+    assert run_request.asset_check_keys == [foo_check.check_key]


Code review comments:
diff_hunk: @@ -445,6 +445,16 @@ def _asset_check_nodes_by_asset_key(self) -> Mapping[AssetKey, Sequence[RemoteAs
     def get_checks_for_asset(self, asset_key: AssetKey) -> Sequence[RemoteAssetCheckNode]:
         return self._asset_check_nodes_by_asset_key.get(asset_key, [])
 
+    def get_check_keys_for_assets(
+        self, asset_keys: AbstractSet[AssetKey]
+    ) -> AbstractSet[AssetCheckKey]:
+        return set().union(
+            *(
+                {check.asset_check.key for check in self.get_checks_for_asset(asset_key)}
+                for asset_key in asset_keys
+            )
+        )
path: python_modules/dagster/dagster/_core/definitions/remote_asset_graph.py
text: @user1:
`set().union(*(...))`  will raise a `TypeError` if `asset_keys` is empty. Consider this safer alternative:

```python
set().union(*(
    {check.asset_check.key for check in self.get_checks_for_asset(asset_key)}
    for asset_key in asset_keys
) if asset_keys else set())
```

*Spotted by [Graphite Reviewer](https://app.graphite.dev/graphite-reviewer/?org=dagster-io&ref=ai-review-comment)*<i class='graphite__hidden'><br /><br />Is this helpful? React 👍 or 👎 to let us know.</i>