diff --git a/python_modules/dagster/dagster/_core/asset_graph_view/asset_graph_view.py b/python_modules/dagster/dagster/_core/asset_graph_view/asset_graph_view.py index 895677cac429b..1c8e0cf57088c 100644 --- a/python_modules/dagster/dagster/_core/asset_graph_view/asset_graph_view.py +++ b/python_modules/dagster/dagster/_core/asset_graph_view/asset_graph_view.py @@ -1,4 +1,3 @@ -import asyncio import functools from collections.abc import Awaitable, Iterable from datetime import datetime, timedelta @@ -22,6 +21,7 @@ from dagster._core.definitions.freshness import FreshnessState from dagster._core.definitions.partitions.context import ( PartitionLoadingContext, + partition_loading_context, use_partition_loading_context, ) from dagster._core.definitions.partitions.definition import ( @@ -564,9 +564,12 @@ async def compute_subset_with_status( """Returns the subset of an asset check that matches a given status.""" from dagster._core.storage.event_log.base import AssetCheckSummaryRecord - # Handle partitioned asset checks with partition-level granularity + # Handle partitioned asset checks if self._get_partitions_def(key): - return await self._get_partitioned_check_subset_with_status(key, status, from_subset) + with partition_loading_context(new_ctx=self._partition_loading_context): + 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) @@ -628,35 +631,38 @@ async def _compute_missing_check_subset( ) -> 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 - + @use_partition_loading_context async def _get_partitioned_check_subset_with_status( self, key: AssetCheckKey, status: Optional["AssetCheckExecutionResolvedStatus"], from_subset: EntitySubset, ) -> EntitySubset[AssetCheckKey]: + from dagster._core.storage.asset_check_state import AssetCheckState + 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) + cache_value = ( + await AssetCheckState.gen(self, (key, check_node.partitions_def)) + or AssetCheckState.empty() ) - matching_partitions = {p for p, s in zip(partitions, statuses) if s == status} - return from_subset.compute_intersection_with_partition_keys(matching_partitions) + + if status is None: + known_statuses = self.get_empty_subset(key=key) + for serializable_subset in cache_value.subsets.values(): + subset = self.get_subset_from_serializable_subset(serializable_subset) + if subset: + known_statuses = known_statuses.compute_union(subset) + return from_subset.compute_difference(known_statuses) or self.get_empty_subset(key=key) + else: + serializable_subset = cache_value.subsets.get(status) + if serializable_subset is None: + return self.get_empty_subset(key=key) + return self.get_subset_from_serializable_subset( + serializable_subset + ) or self.get_empty_subset(key=key) async def _compute_run_in_progress_asset_subset(self, key: AssetKey) -> EntitySubset[AssetKey]: from dagster._core.storage.partition_status_cache import AssetStatusCacheValue diff --git a/python_modules/dagster/dagster/_core/storage/asset_check_state.py b/python_modules/dagster/dagster/_core/storage/asset_check_state.py new file mode 100644 index 0000000000000..2f6f42540c4aa --- /dev/null +++ b/python_modules/dagster/dagster/_core/storage/asset_check_state.py @@ -0,0 +1,263 @@ +from collections import defaultdict +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Optional, TypeAlias + +from dagster_shared.record import record +from dagster_shared.serdes import whitelist_for_serdes + +from dagster._core.asset_graph_view.serializable_entity_subset import SerializableEntitySubset +from dagster._core.definitions.asset_key import AssetCheckKey +from dagster._core.definitions.partitions.definition.partitions_definition import ( + PartitionsDefinition, +) +from dagster._core.loader import LoadableBy, LoadingContext +from dagster._core.storage.asset_check_execution_record import ( + AssetCheckExecutionRecordStatus, + AssetCheckExecutionResolvedStatus, + AssetCheckPartitionInfo, +) +from dagster._core.storage.dagster_run import FINISHED_STATUSES, DagsterRunStatus, RunsFilter + +if TYPE_CHECKING: + from dagster._core.instance import DagsterInstance + +StatusSubsets: TypeAlias = Mapping[ + AssetCheckExecutionResolvedStatus, SerializableEntitySubset[AssetCheckKey] +] +InProgressRuns: TypeAlias = Mapping[str, SerializableEntitySubset[AssetCheckKey]] + + +@whitelist_for_serdes +@record +class AssetCheckState(LoadableBy[tuple[AssetCheckKey, Optional[PartitionsDefinition]]]): + latest_storage_id: int + subsets: StatusSubsets + in_progress_runs: InProgressRuns + + def compatible_with(self, partitions_def: Optional[PartitionsDefinition]) -> bool: + subset = next(iter(self.subsets.values()), None) + if subset is None: + return True + return subset.is_compatible_with_partitions_def(partitions_def) + + @classmethod + def empty(cls) -> "AssetCheckState": + return cls(latest_storage_id=0, subsets={}, in_progress_runs={}) + + @classmethod + def _blocking_batch_load( + cls, + keys: Iterable[tuple[AssetCheckKey, Optional[PartitionsDefinition]]], + context: LoadingContext, + ) -> Iterable[Optional["AssetCheckState"]]: + keys = list(keys) + mapping = context.instance.event_log_storage.get_asset_check_state(keys) + return [mapping.get(check_key) for check_key, _ in keys] + + def with_updates( + self, + key: AssetCheckKey, + partitions_def: Optional[PartitionsDefinition], + partition_records: Sequence[AssetCheckPartitionInfo], + run_statuses: Mapping[str, DagsterRunStatus], + ) -> "AssetCheckState": + latest_storage_id = max( + ( + max(r.latest_check_event_storage_id, r.latest_materialization_storage_id or 0) + for r in partition_records + ), + default=self.latest_storage_id, + ) + + subsets = { + status: self.subsets.get(status, SerializableEntitySubset.empty(key, partitions_def)) + for status in AssetCheckExecutionResolvedStatus + } + in_progress_runs = dict(self.in_progress_runs) + + # update all subsets based on the new partition records + subsets, in_progress_runs = _process_partition_records( + key, partitions_def, subsets, in_progress_runs, partition_records + ) + # then check the run statuses and resolve any previously in-progress runs that have completed + subsets, in_progress_runs = _process_run_statuses( + key, partitions_def, subsets, in_progress_runs, run_statuses + ) + return AssetCheckState( + latest_storage_id=latest_storage_id, + subsets=subsets, + in_progress_runs=in_progress_runs, + ) + + +def bulk_update_asset_check_state( + instance: "DagsterInstance", + keys: Sequence[tuple[AssetCheckKey, Optional[PartitionsDefinition]]], + initial_states: Mapping[AssetCheckKey, "AssetCheckState"], +) -> Mapping[AssetCheckKey, "AssetCheckState"]: + check_keys = [key for key, _ in keys] + partitions_defs_by_key = {key: partitions_def for key, partitions_def in keys} + + # we prefer to do a single fetch for all keys, so we use the minimum storage id of the initial states + storage_id = min((state.latest_storage_id for state in initial_states.values()), default=0) + infos = instance.event_log_storage.get_asset_check_partition_info( + check_keys, after_storage_id=storage_id + ) + + # find the set of run ids we need to fetch to resolve the in-progress runs, and + # group the partition infos by check key + run_ids_to_fetch: set[str] = set().union( + *(state.in_progress_runs.keys() for state in initial_states.values()) + ) + infos_by_key: dict[AssetCheckKey, list[AssetCheckPartitionInfo]] = defaultdict(list) + for info in infos: + infos_by_key[info.check_key].append(info) + if info.latest_execution_status == AssetCheckExecutionRecordStatus.PLANNED: + run_ids_to_fetch.add(info.latest_planned_run_id) + + # do a bulk fetch for runs across all states + finished_runs = ( + instance.get_runs( + filters=RunsFilter(run_ids=list(run_ids_to_fetch), statuses=FINISHED_STATUSES) + ) + if len(run_ids_to_fetch) > 0 + else [] + ) + finished_runs_status_by_id = {run.run_id: run.status for run in finished_runs} + return { + key: initial_states[key].with_updates( + key, partitions_defs_by_key[key], infos_by_key[key], finished_runs_status_by_id + ) + for key in check_keys + } + + +def _valid_partition_key( + partition_key: Optional[str], partitions_def: Optional[PartitionsDefinition] +) -> bool: + if partitions_def is None: + return partition_key is None + else: + return partition_key is not None and partitions_def.has_partition_key(partition_key) + + +def _process_partition_records( + key: AssetCheckKey, + partitions_def: Optional[PartitionsDefinition], + subsets: StatusSubsets, + in_progress_runs: InProgressRuns, + partition_infos: Sequence[AssetCheckPartitionInfo], +) -> tuple[StatusSubsets, InProgressRuns]: + """Returns a set of updated subsets based on new partition records and the latest materialization storage ids.""" + new_subsets = dict(subsets) + new_in_progress_runs = dict(in_progress_runs) + + for partition_record in partition_infos: + pk = partition_record.partition_key + if not _valid_partition_key(pk, partitions_def): + continue + + partition_subset = SerializableEntitySubset.from_coercible_value(key, pk, partitions_def) + + if partition_record.latest_execution_status == AssetCheckExecutionRecordStatus.PLANNED: + # Add to IN_PROGRESS and track run + new_subsets[AssetCheckExecutionResolvedStatus.IN_PROGRESS] = new_subsets[ + AssetCheckExecutionResolvedStatus.IN_PROGRESS + ].compute_union(partition_subset) + run_id = partition_record.latest_planned_run_id + new_in_progress_runs[run_id] = new_in_progress_runs.get( + run_id, SerializableEntitySubset.empty(key, partitions_def) + ).compute_union(partition_subset) + + elif partition_record.latest_execution_status in ( + AssetCheckExecutionRecordStatus.SUCCEEDED, + AssetCheckExecutionRecordStatus.FAILED, + ): + if partition_record.is_current: + # Check is current, set appropriate status + status = ( + AssetCheckExecutionResolvedStatus.SUCCEEDED + if partition_record.latest_execution_status + == AssetCheckExecutionRecordStatus.SUCCEEDED + else AssetCheckExecutionResolvedStatus.FAILED + ) + new_subsets[status] = new_subsets[status].compute_union(partition_subset) + else: + # new materialization, clear the check status + for status in new_subsets: + new_subsets[status] = new_subsets[status].compute_difference(partition_subset) + + return new_subsets, new_in_progress_runs + + +def _process_run_statuses( + key: AssetCheckKey, + partitions_def: Optional[PartitionsDefinition], + subsets: StatusSubsets, + in_progress_runs: InProgressRuns, + run_statuses: Mapping[str, DagsterRunStatus], +) -> tuple[StatusSubsets, InProgressRuns]: + """Resolve in-progress runs that have completed. + + This checks if any runs tracked in in_progress_runs have finished, + and moves their partitions to SKIPPED or EXECUTION_FAILED. + """ + if not in_progress_runs: + return subsets, in_progress_runs + + delta_skipped, delta_execution_failed, resolved_run_ids = _resolve_in_progress_subsets( + key, partitions_def, in_progress_runs, run_statuses + ) + + new_in_progress_runs = { + run_id: subset + for run_id, subset in in_progress_runs.items() + if run_id not in resolved_run_ids + } + + new_subsets = dict(subsets) + new_subsets[AssetCheckExecutionResolvedStatus.IN_PROGRESS] = ( + new_subsets[AssetCheckExecutionResolvedStatus.IN_PROGRESS] + .compute_difference(delta_skipped) + .compute_difference(delta_execution_failed) + ) + new_subsets[AssetCheckExecutionResolvedStatus.SKIPPED] = new_subsets[ + AssetCheckExecutionResolvedStatus.SKIPPED + ].compute_union(delta_skipped) + new_subsets[AssetCheckExecutionResolvedStatus.EXECUTION_FAILED] = new_subsets[ + AssetCheckExecutionResolvedStatus.EXECUTION_FAILED + ].compute_union(delta_execution_failed) + + return new_subsets, new_in_progress_runs + + +def _resolve_in_progress_subsets( + key: AssetCheckKey, + partitions_def: Optional[PartitionsDefinition], + in_progress_runs: InProgressRuns, + run_statuses: Mapping[str, DagsterRunStatus], +) -> tuple[ + SerializableEntitySubset[AssetCheckKey], + SerializableEntitySubset[AssetCheckKey], + set[str], +]: + """Resolve in-progress runs that have completed. + + Returns: + Tuple of (delta_skipped, delta_execution_failed, resolved_run_ids) + """ + empty_subset = SerializableEntitySubset.empty(key, partitions_def) + delta_skipped = empty_subset + delta_execution_failed = empty_subset + resolved_run_ids: set[str] = set() + + for run_id, run_status in run_statuses.items(): + if run_status in FINISHED_STATUSES and run_id in in_progress_runs: + resolved_run_ids.add(run_id) + run_subset = in_progress_runs[run_id] + if run_status == DagsterRunStatus.FAILURE: + delta_execution_failed = delta_execution_failed.compute_union(run_subset) + else: + delta_skipped = delta_skipped.compute_union(run_subset) + + return delta_skipped, delta_execution_failed, resolved_run_ids diff --git a/python_modules/dagster/dagster/_core/storage/event_log/base.py b/python_modules/dagster/dagster/_core/storage/event_log/base.py index d3e22ec7b3357..c9ffb890eed29 100644 --- a/python_modules/dagster/dagster/_core/storage/event_log/base.py +++ b/python_modules/dagster/dagster/_core/storage/event_log/base.py @@ -50,6 +50,7 @@ if TYPE_CHECKING: from dagster._core.events.log import EventLogEntry + from dagster._core.storage.asset_check_state import AssetCheckState from dagster._core.storage.partition_status_cache import AssetStatusCacheValue @@ -659,6 +660,28 @@ def get_asset_check_partition_info( """Get asset check partition records with execution status and planned run info.""" pass + def get_checkpointed_asset_check_state( + self, keys: Sequence[AssetCheckKey] + ) -> Mapping[AssetCheckKey, "AssetCheckState"]: + """Get the current stored asset check state for a list of asset checks and their + associated partitions definitions. This method is not guaranteed to return a + state object that is up to date with the latest events. + """ + from dagster._core.storage.asset_check_state import AssetCheckState + + return {key: AssetCheckState.empty() for key in keys} + + def get_asset_check_state( + self, keys: Sequence[tuple[AssetCheckKey, Optional[PartitionsDefinition]]] + ) -> Mapping[AssetCheckKey, "AssetCheckState"]: + from dagster._core.storage.asset_check_state import bulk_update_asset_check_state + + return bulk_update_asset_check_state( + self._instance, + keys, + initial_states=self.get_checkpointed_asset_check_state([key for key, _ in keys]), + ) + @abstractmethod def fetch_materializations( self, diff --git a/python_modules/dagster/dagster/_core/storage/legacy_storage.py b/python_modules/dagster/dagster/_core/storage/legacy_storage.py index a0ed8875bb0ba..028f673b6cd4d 100644 --- a/python_modules/dagster/dagster/_core/storage/legacy_storage.py +++ b/python_modules/dagster/dagster/_core/storage/legacy_storage.py @@ -5,7 +5,6 @@ from dagster import _check as check from dagster._config.config_schema import UserConfigSchema -from dagster._core.definitions.asset_checks.asset_check_spec import AssetCheckKey from dagster._core.definitions.asset_key import EntityKey from dagster._core.definitions.declarative_automation.serialized_objects import ( AutomationConditionEvaluationWithRunIds, @@ -62,6 +61,7 @@ ) from dagster._core.snap.execution_plan_snapshot import ExecutionPlanSnapshot from dagster._core.snap.job_snapshot import JobSnap + from dagster._core.storage.asset_check_state import AssetCheckState from dagster._core.storage.dagster_run import ( DagsterRun, DagsterRunStatsSnapshot, @@ -776,6 +776,11 @@ def get_asset_check_partition_info( keys=keys, after_storage_id=after_storage_id, partition_keys=partition_keys ) + def get_checkpointed_asset_check_state( + self, keys: Sequence["AssetCheckKey"] + ) -> Mapping["AssetCheckKey", "AssetCheckState"]: + return self._storage.event_log_storage.get_checkpointed_asset_check_state(keys) + class LegacyScheduleStorage(ScheduleStorage, ConfigurableClass): def __init__(self, storage: DagsterStorage, inst_data: Optional[ConfigurableClassData] = None):