[pac] Update status resolution logic

Author: OwenKephart

python_modules/dagster/dagster/_core/asset_graph_view/asset_graph_view.py

@@ -1,3 +1,4 @@
+import asyncio
 import functools
 from collections.abc import Awaitable, Iterable
 from datetime import datetime, timedelta
@@ -162,10 +163,7 @@ def get_inner_queryer_for_back_compat(self) -> "CachingInstanceQueryer":
         return self._queryer
 
     def _get_partitions_def(self, key: T_EntityKey) -> Optional["PartitionsDefinition"]:
-        if isinstance(key, AssetKey):
-            return self.asset_graph.get(key).partitions_def
-        else:
-            return None
+        return self.asset_graph.get(key).partitions_def
 
     @cached_method
     @use_partition_loading_context
@@ -374,6 +372,16 @@ def get_asset_subset_from_asset_partitions(
         )
         return EntitySubset(self, key=key, value=_ValidatedEntitySubsetValue(value))
 
+    @use_partition_loading_context
+    def get_subset_from_partition_keys(
+        self,
+        key: T_EntityKey,
+        partitions_def: "PartitionsDefinition",
+        partition_keys: AbstractSet[str],
+    ) -> EntitySubset[T_EntityKey]:
+        value = partitions_def.subset_with_partition_keys(partition_keys)
+        return EntitySubset(self, key=key, value=_ValidatedEntitySubsetValue(value))
+
     @use_partition_loading_context
     def compute_parent_subset_and_required_but_nonexistent_subset(
         self, parent_key, subset: EntitySubset[T_EntityKey]
@@ -548,11 +556,19 @@ def compute_latest_time_window_subset(
             check.failed(f"Unsupported partitions_def: {partitions_def}")
 
     async def compute_subset_with_status(
-        self, key: AssetCheckKey, status: Optional["AssetCheckExecutionResolvedStatus"]
-    ):
+        self,
+        key: AssetCheckKey,
+        status: Optional["AssetCheckExecutionResolvedStatus"],
+        from_subset: EntitySubset,
+    ) -> EntitySubset[AssetCheckKey]:
+        """Returns the subset of an asset check that matches a given status."""
         from dagster._core.storage.event_log.base import AssetCheckSummaryRecord
 
-        """Returns the subset of an asset check that matches a given status."""
+        # Handle partitioned asset checks with partition-level granularity
+        if self._get_partitions_def(key):
+            return await self._get_partitioned_check_subset_with_status(key, status, from_subset)
+
+        # Handle non-partitioned asset checks with existing logic
         summary = await AssetCheckSummaryRecord.gen(self, key)
         latest_record = summary.last_check_execution_record if summary else None
         resolved_status = (
@@ -586,31 +602,61 @@ async def compute_subset_with_freshness_state(
             return self.get_empty_subset(key=key)
 
     async def _compute_run_in_progress_check_subset(
-        self, key: AssetCheckKey
+        self, key: AssetCheckKey, from_subset: EntitySubset
     ) -> EntitySubset[AssetCheckKey]:
         from dagster._core.storage.asset_check_execution_record import (
             AssetCheckExecutionResolvedStatus,
         )
 
         return await self.compute_subset_with_status(
-            key, AssetCheckExecutionResolvedStatus.IN_PROGRESS
+            key, AssetCheckExecutionResolvedStatus.IN_PROGRESS, from_subset
         )
 
     async def _compute_execution_failed_check_subset(
-        self, key: AssetCheckKey
+        self, key: AssetCheckKey, from_subset: EntitySubset
     ) -> EntitySubset[AssetCheckKey]:
         from dagster._core.storage.asset_check_execution_record import (
             AssetCheckExecutionResolvedStatus,
         )
 
         return await self.compute_subset_with_status(
-            key, AssetCheckExecutionResolvedStatus.EXECUTION_FAILED
+            key, AssetCheckExecutionResolvedStatus.EXECUTION_FAILED, from_subset
         )
 
     async def _compute_missing_check_subset(
-        self, key: AssetCheckKey
+        self, key: AssetCheckKey, from_subset: EntitySubset
+    ) -> EntitySubset[AssetCheckKey]:
+        return await self.compute_subset_with_status(key, None, from_subset)
+
+    async def _get_asset_check_partition_status(
+        self, key: AssetCheckKey, partition: str
+    ) -> Optional["AssetCheckExecutionResolvedStatus"]:
+        # NOTE: we should add a LoadingContext-native version of this
+        record = self.instance.event_log_storage.get_latest_asset_check_execution_by_key(
+            [key], partition
+        ).get(key)
+        if record:
+            targets_latest = await record.targets_latest_materialization(self)
+            return await record.resolve_status(self) if targets_latest else None
+        else:
+            return None
+
+    async def _get_partitioned_check_subset_with_status(
+        self,
+        key: AssetCheckKey,
+        status: Optional["AssetCheckExecutionResolvedStatus"],
+        from_subset: EntitySubset,
     ) -> EntitySubset[AssetCheckKey]:
-        return await self.compute_subset_with_status(key, None)
+        check_node = self.asset_graph.get(key)
+        if not check_node or not check_node.partitions_def:
+            check.failed(f"Asset check {key} not found or not partitioned.")
+
+        partitions = list(from_subset.expensively_compute_partition_keys())
+        statuses = await asyncio.gather(
+            *(self._get_asset_check_partition_status(key, p) for p in partitions)
+        )
+        matching_partitions = {p for p, s in zip(partitions, statuses) if s == status}
+        return from_subset.compute_intersection_with_partition_keys(matching_partitions)
 
     async def _compute_run_in_progress_asset_subset(self, key: AssetKey) -> EntitySubset[AssetKey]:
         from dagster._core.storage.partition_status_cache import AssetStatusCacheValue
@@ -735,15 +781,21 @@ async def _compute_missing_asset_subset(
             )
 
     @cached_method
-    async def compute_run_in_progress_subset(self, *, key: EntityKey) -> EntitySubset:
+    async def compute_run_in_progress_subset(
+        self, *, key: EntityKey, from_subset: EntitySubset
+    ) -> EntitySubset:
         return await _dispatch(
             key=key,
-            check_method=self._compute_run_in_progress_check_subset,
+            check_method=functools.partial(
+                self._compute_run_in_progress_check_subset, from_subset=from_subset
+            ),
             asset_method=self._compute_run_in_progress_asset_subset,
         )
 
     @cached_method
-    async def compute_backfill_in_progress_subset(self, *, key: EntityKey) -> EntitySubset:
+    async def compute_backfill_in_progress_subset(
+        self, *, key: EntityKey, from_subset: EntitySubset
+    ) -> EntitySubset:
         async def get_empty_subset(key: EntityKey) -> EntitySubset:
             return self.get_empty_subset(key=key)
 
@@ -755,10 +807,14 @@ async def get_empty_subset(key: EntityKey) -> EntitySubset:
         )
 
     @cached_method
-    async def compute_execution_failed_subset(self, *, key: EntityKey) -> EntitySubset:
+    async def compute_execution_failed_subset(
+        self, *, key: EntityKey, from_subset: EntitySubset
+    ) -> EntitySubset:
         return await _dispatch(
             key=key,
-            check_method=self._compute_execution_failed_check_subset,
+            check_method=functools.partial(
+                self._compute_execution_failed_check_subset, from_subset=from_subset
+            ),
             asset_method=self._compute_execution_failed_asset_subset,
         )
 
@@ -768,7 +824,9 @@ async def compute_missing_subset(
     ) -> EntitySubset:
         return await _dispatch(
             key=key,
-            check_method=self._compute_missing_check_subset,
+            check_method=functools.partial(
+                self._compute_missing_check_subset, from_subset=from_subset
+            ),
             asset_method=functools.partial(
                 self._compute_missing_asset_subset, from_subset=from_subset
             ),

python_modules/dagster/dagster/_core/asset_graph_view/entity_subset.py

@@ -69,10 +69,13 @@ def convert_to_serializable_subset(self) -> SerializableEntitySubset[T_EntityKey
         return SerializableEntitySubset(key=self._key, value=self._value)
 
     def expensively_compute_partition_keys(self) -> AbstractSet[str]:
-        return {
-            check.not_none(akpk.partition_key, "No None partition keys")
-            for akpk in self.expensively_compute_asset_partitions()
-        }
+        internal_value = self.get_internal_value()
+        if isinstance(internal_value, PartitionsSubset):
+            return set(internal_value.get_partition_keys())
+        elif internal_value:
+            check.failed("Subset is not partitioned")
+        else:
+            return set()
 
     def expensively_compute_asset_partitions(self) -> AbstractSet[AssetKeyPartitionKey]:
         if not isinstance(self.key, AssetKey):
@@ -106,11 +109,13 @@ def compute_intersection(self, other: Self) -> Self:
         return self._oper(other, operator.and_)
 
     def compute_intersection_with_partition_keys(
-        self: "EntitySubset[AssetKey]", partition_keys: AbstractSet[str]
-    ) -> "EntitySubset[AssetKey]":
-        key = check.inst(self.key, AssetKey)
-        partition_subset = self._asset_graph_view.get_asset_subset_from_asset_partitions(
-            self.key, {AssetKeyPartitionKey(key, pk) for pk in partition_keys}
+        self: "EntitySubset[T_EntityKey]", partition_keys: AbstractSet[str]
+    ) -> "EntitySubset[T_EntityKey]":
+        if self.partitions_def is None:
+            return self._asset_graph_view.get_empty_subset(key=self.key)
+
+        partition_subset = self._asset_graph_view.get_subset_from_partition_keys(
+            self.key, self.partitions_def, partition_keys
         )
         return self.compute_intersection(partition_subset)
 

python_modules/dagster/dagster/_core/definitions/declarative_automation/operands/operands.py

@@ -102,7 +102,9 @@ def name(self) -> str:
         return "run_in_progress"
 
     async def compute_subset(self, context: AutomationContext) -> EntitySubset:  # pyright: ignore[reportIncompatibleMethodOverride]
-        return await context.asset_graph_view.compute_run_in_progress_subset(key=context.key)
+        return await context.asset_graph_view.compute_run_in_progress_subset(
+            key=context.key, from_subset=context.candidate_subset
+        )
 
 
 @whitelist_for_serdes
@@ -113,7 +115,9 @@ def name(self) -> str:
         return "backfill_in_progress"
 
     async def compute_subset(self, context: AutomationContext) -> EntitySubset:  # pyright: ignore[reportIncompatibleMethodOverride]
-        return await context.asset_graph_view.compute_backfill_in_progress_subset(key=context.key)
+        return await context.asset_graph_view.compute_backfill_in_progress_subset(
+            key=context.key, from_subset=context.candidate_subset
+        )
 
 
 @whitelist_for_serdes(storage_name="FailedAutomationCondition")
@@ -124,7 +128,9 @@ def name(self) -> str:
         return "execution_failed"
 
     async def compute_subset(self, context: AutomationContext) -> EntitySubset:  # pyright: ignore[reportIncompatibleMethodOverride]
-        return await context.asset_graph_view.compute_execution_failed_subset(key=context.key)
+        return await context.asset_graph_view.compute_execution_failed_subset(
+            key=context.key, from_subset=context.candidate_subset
+        )
 
 
 @whitelist_for_serdes
@@ -322,5 +328,5 @@ async def compute_subset(  # pyright: ignore[reportIncompatibleMethodOverride]
             else AssetCheckExecutionResolvedStatus.FAILED
         )
         return await context.asset_graph_view.compute_subset_with_status(
-            key=context.key, status=target_status
+            key=context.key, status=target_status, from_subset=context.candidate_subset
         )

python_modules/dagster/dagster/_core/storage/asset_check_execution_record.py

@@ -173,17 +173,28 @@ async def resolve_status(
             check.failed(f"Unexpected status {self.status}")
 
     async def targets_latest_materialization(self, loading_context: LoadingContext) -> bool:
-        from dagster._core.storage.event_log.base import AssetRecord
+        from dagster._core.storage.event_log.base import AssetRecord, AssetRecordsFilter
 
         resolved_status = await self.resolve_status(loading_context)
         if resolved_status == AssetCheckExecutionResolvedStatus.IN_PROGRESS:
             # all in-progress checks execute against the latest version
             return True
 
-        asset_record = await AssetRecord.gen(loading_context, self.key.asset_key)
-        latest_materialization = (
-            asset_record.asset_entry.last_materialization_record if asset_record else None
-        )
+        if self.partition is None:
+            asset_record = await AssetRecord.gen(loading_context, self.key.asset_key)
+            latest_materialization = (
+                asset_record.asset_entry.last_materialization_record if asset_record else None
+            )
+        else:
+            records = loading_context.instance.fetch_materializations(
+                AssetRecordsFilter(
+                    asset_key=self.key.asset_key,
+                    asset_partitions=[self.partition],
+                ),
+                limit=1,
+            )
+            latest_materialization = records.records[0] if records.records else None
+
         if not latest_materialization:
             # no previous materialization, so it's executing against the lastest version
             return True