asset partitions numFailed (#12570)

gql for counting failed partitions.

Update: to make it so that numMaterialized and numFailed can't total to
greater than numPartitions, we now filter out failed partitions from
materialized.

Author: johannkm

python_modules/dagster-graphql/dagster_graphql/implementation/fetch_assets.py

@@ -45,6 +45,7 @@
 from dagster._core.instance import DynamicPartitionsStore
 from dagster._core.storage.partition_status_cache import (
     CACHEABLE_PARTITION_TYPES,
+    filter_incomplete_materialized_runs_to_failed,
     get_and_update_asset_status_cache_value,
     get_materialized_multipartitions,
     get_validated_partition_keys,
@@ -321,19 +322,18 @@ def get_unique_asset_id(
     )
 
 
-def get_materialized_partitions_subset(
+def get_materialized_and_failed_partition_subsets(
     instance: DagsterInstance,
     asset_key: AssetKey,
     dynamic_partitions_loader: DynamicPartitionsStore,
     partitions_def: Optional[PartitionsDefinition] = None,
-) -> Optional[PartitionsSubset]:
+) -> Tuple[Optional[PartitionsSubset], Optional[PartitionsSubset]]:
     """
-    Returns the materialization status for each partition key. The materialization status
-    is a boolean indicating whether the partition has been materialized: True if materialized,
-    False if not.
+    Returns a tuple of two PartitionSubset objects: the first is the materialized partitions,
+    the second is the failed partitions.
     """
     if not partitions_def:
-        return None
+        return None, None
 
     if instance.can_cache_asset_status_data() and isinstance(
         partitions_def, CACHEABLE_PARTITION_TYPES
@@ -348,8 +348,13 @@ def get_materialized_partitions_subset(
             if updated_cache_value
             else partitions_def.empty_subset()
         )
+        failed_subset = (
+            updated_cache_value.deserialize_failed_partition_subsets(partitions_def)
+            if updated_cache_value
+            else partitions_def.empty_subset()
+        )
 
-        return materialized_subset
+        return materialized_subset, failed_subset
 
     else:
         # If the partition status can't be cached, fetch partition status from storage
@@ -370,12 +375,36 @@ def get_materialized_partitions_subset(
             dynamic_partitions_loader, partitions_def, set(materialized_keys)
         )
 
-        return (
+        materialized_subset = (
             partitions_def.empty_subset().with_partition_keys(validated_keys)
             if validated_keys
             else partitions_def.empty_subset()
         )
 
+        # failed
+        incomplete_materialization_runs = instance.event_log_storage.get_latest_asset_partition_materialization_attempts_without_materializations(
+            asset_key
+        )
+
+        if not incomplete_materialization_runs:
+            failed_subset = partitions_def.empty_subset()
+        else:
+            validated_keys = get_validated_partition_keys(
+                dynamic_partitions_loader,
+                partitions_def,
+                filter_incomplete_materialized_runs_to_failed(
+                    instance, incomplete_materialization_runs
+                ),
+            )
+
+            failed_subset = (
+                partitions_def.empty_subset().with_partition_keys(validated_keys)
+                if validated_keys
+                else partitions_def.empty_subset()
+            )
+
+        return materialized_subset, failed_subset
+
 
 def build_materialized_partitions(
     dynamic_partitions_store: DynamicPartitionsStore,

python_modules/dagster-graphql/dagster_graphql/schema/asset_graph.py

@@ -46,7 +46,7 @@
 from ..implementation.fetch_assets import (
     build_materialized_partitions,
     get_freshness_info,
-    get_materialized_partitions_subset,
+    get_materialized_and_failed_partition_subsets,
 )
 from ..implementation.loader import (
     BatchMaterializationLoader,
@@ -748,14 +748,14 @@ def resolve_materializedPartitions(
         if not self._dynamic_partitions_loader:
             check.failed("dynamic_partitions_loader must be provided to get partition keys")
 
-        materialized_partition_subset = get_materialized_partitions_subset(
+        materialized_partition_subset = get_materialized_and_failed_partition_subsets(
             graphene_info.context.instance,
             asset_key,
             self._dynamic_partitions_loader,
             self._external_asset_node.partitions_def_data.get_partitions_definition()
             if self._external_asset_node.partitions_def_data
             else None,
-        )
+        )[0]
 
         return build_materialized_partitions(
             self._dynamic_partitions_loader,
@@ -770,7 +770,10 @@ def resolve_partitionStats(self, graphene_info) -> Optional[GraphenePartitionSta
             if not self._dynamic_partitions_loader:
                 check.failed("dynamic_partitions_loader must be provided to get partition keys")
 
-            materialized_partition_subset = get_materialized_partitions_subset(
+            (
+                materialized_partition_subset,
+                failed_partition_subset,
+            ) = get_materialized_and_failed_partition_subsets(
                 graphene_info.context.instance,
                 asset_key,
                 self._dynamic_partitions_loader,
@@ -779,14 +782,24 @@ def resolve_partitionStats(self, graphene_info) -> Optional[GraphenePartitionSta
                 else None,
             )
 
-            if materialized_partition_subset is None:
+            if materialized_partition_subset is None or failed_partition_subset is None:
                 check.failed("Expected partitions subset for a partitioned asset")
 
+            num_materialized = len(materialized_partition_subset)
+            num_materialized_and_not_failed = num_materialized - len(
+                [
+                    k
+                    for k in failed_partition_subset.get_partition_keys()
+                    if k in materialized_partition_subset
+                ]
+            )
+
             return GraphenePartitionStats(
-                numMaterialized=len(materialized_partition_subset),
+                numMaterialized=num_materialized_and_not_failed,
                 numPartitions=partitions_def_data.get_partitions_definition().get_num_partitions(
                     dynamic_partitions_store=self._dynamic_partitions_loader
                 ),
+                numFailed=len(failed_partition_subset),
             )
         else:
             return None

python_modules/dagster-graphql/dagster_graphql/schema/pipelines/pipeline.py

@@ -148,6 +148,7 @@ class Meta:
 class GraphenePartitionStats(graphene.ObjectType):
     numMaterialized = graphene.NonNull(graphene.Int)
     numPartitions = graphene.NonNull(graphene.Int)
+    numFailed = graphene.NonNull(graphene.Int)
 
     class Meta:
         name = "PartitionStats"

python_modules/dagster-graphql/dagster_graphql_tests/graphql/repo.py

@@ -1651,6 +1651,20 @@ def yield_partition_materialization():
 )
 
 
+@asset(partitions_def=StaticPartitionsDefinition(["a", "b", "c", "d"]))
+def fail_partition_materialization(context):
+    if context.run.tags.get("fail") == "true":
+        raise Exception("fail_partition_materialization")
+    yield Output(5)
+
+
+fail_partition_materialization_job = build_assets_job(
+    "fail_partition_materialization_job",
+    assets=[fail_partition_materialization],
+    executor_def=in_process_executor,
+)
+
+
 @asset
 def asset_yields_observation():
     yield AssetObservation(asset_key=AssetKey("asset_yields_observation"), metadata={"text": "FOO"})
@@ -1926,6 +1940,7 @@ def define_pipelines():
         dynamic_partitioned_assets_job,
         time_partitioned_assets_job,
         partition_materialization_job,
+        fail_partition_materialization_job,
         observation_job,
         failure_assets_job,
         asset_group_job,