Skip to content

Commit 5fed730

Browse files
Vamsi-kluVamsi-klu
committed
[dagster-graphql] Tolerate stale partition evaluation subsets
1 parent 4fb0017 commit 5fed730

2 files changed

Lines changed: 142 additions & 13 deletions

File tree

‎python_modules/dagster-graphql/dagster_graphql/schema/asset_condition_evaluations.py‎

Lines changed: 29 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
AutomationConditionEvaluation,
1212
get_expanded_label,
1313
)
14+
from dagster._core.errors import DagsterInvalidInvocationError
1415
from dagster._core.scheduler.instigation import AutoMaterializeAssetEvaluationRecord
1516

1617
from dagster_graphql.implementation.events import iterate_metadata_entries
@@ -49,10 +50,10 @@ class Meta:
4950

5051
def __init__(self, evaluation: AutomationConditionEvaluation):
5152
self._evaluation = evaluation
52-
if evaluation.true_subset.size > 0:
53+
if _get_historical_subset_size(evaluation.true_subset) > 0:
5354
status = AssetConditionEvaluationStatus.TRUE
5455
elif isinstance(evaluation.candidate_subset, SerializableEntitySubset) and (
55-
evaluation.candidate_subset.size > 0
56+
_get_historical_subset_size(evaluation.candidate_subset) > 0
5657
):
5758
status = AssetConditionEvaluationStatus.FALSE
5859
else:
@@ -110,11 +111,11 @@ def __init__(self, evaluation: AutomationConditionEvaluation):
110111
)
111112

112113
def resolve_numTrue(self, graphene_info: ResolveInfo) -> int | None:
113-
return self._evaluation.true_subset.size
114+
return _get_historical_subset_size(self._evaluation.true_subset)
114115

115116
def resolve_numCandidates(self, graphene_info: ResolveInfo) -> int | None:
116117
return (
117-
self._evaluation.candidate_subset.size
118+
_get_historical_subset_size(self._evaluation.candidate_subset)
118119
if isinstance(self._evaluation.candidate_subset, SerializableEntitySubset)
119120
else None
120121
)
@@ -143,12 +144,11 @@ def __init__(self, evaluation: AutomationConditionEvaluation, partition_key: str
143144
# where a sub-condition is not partitioned. In these cases, we can treat the expression
144145
# as SKIPPED
145146
status = AssetConditionEvaluationStatus.SKIPPED
146-
elif partition_key in evaluation.true_subset.subset_value:
147+
elif _historical_subset_contains_partition(evaluation.true_subset, partition_key):
147148
status = AssetConditionEvaluationStatus.TRUE
148-
elif (
149-
not isinstance(evaluation.candidate_subset, SerializableEntitySubset)
150-
or partition_key in evaluation.candidate_subset.subset_value
151-
):
149+
elif not isinstance(
150+
evaluation.candidate_subset, SerializableEntitySubset
151+
) or _historical_subset_contains_partition(evaluation.candidate_subset, partition_key):
152152
status = AssetConditionEvaluationStatus.FALSE
153153
else:
154154
status = AssetConditionEvaluationStatus.SKIPPED
@@ -171,7 +171,7 @@ def resolve_metadataEntries(
171171
(
172172
subset.metadata
173173
for subset in self._evaluation.subsets_with_metadata
174-
if self._partition_key in subset.subset.subset_value
174+
if _historical_subset_contains_partition(subset.subset, self._partition_key)
175175
),
176176
{},
177177
)
@@ -281,11 +281,11 @@ def __init__(self, evaluation: AutomationConditionEvaluation):
281281
)
282282

283283
def resolve_numTrue(self, graphene_info: ResolveInfo) -> int | None:
284-
return self._evaluation.true_subset.size
284+
return _get_historical_subset_size(self._evaluation.true_subset)
285285

286286
def resolve_numCandidates(self, graphene_info: ResolveInfo) -> int | None:
287287
return (
288-
self._evaluation.candidate_subset.size
288+
_get_historical_subset_size(self._evaluation.candidate_subset)
289289
if isinstance(self._evaluation.candidate_subset, SerializableEntitySubset)
290290
else None
291291
)
@@ -362,7 +362,7 @@ def __init__(
362362
)
363363

364364
def resolve_numRequested(self, graphene_info: ResolveInfo) -> int | None:
365-
return self._root_evaluation.true_subset.size
365+
return _get_historical_subset_size(self._root_evaluation.true_subset)
366366

367367

368368
class GrapheneAssetConditionEvaluationRecords(graphene.ObjectType):
@@ -388,3 +388,19 @@ def _flatten_evaluation(
388388
import itertools
389389

390390
return list(itertools.chain([e], *(_flatten_evaluation(ce) for ce in e.child_evaluations)))
391+
392+
393+
def _get_historical_subset_size(subset: SerializableEntitySubset) -> int:
394+
try:
395+
return subset.size
396+
except DagsterInvalidInvocationError:
397+
return 0
398+
399+
400+
def _historical_subset_contains_partition(
401+
subset: SerializableEntitySubset, partition_key: str
402+
) -> bool:
403+
try:
404+
return subset.is_partitioned and partition_key in subset.subset_value
405+
except DagsterInvalidInvocationError:
406+
return False

‎python_modules/dagster-graphql/dagster_graphql_tests/graphql/test_asset_condition_evaluations.py‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,16 @@
1717
HistoricalAllPartitionsSubsetSentinel,
1818
)
1919
from dagster._core.definitions.partitions.definition import (
20+
DynamicPartitionsDefinition,
2021
PartitionsDefinition,
2122
StaticPartitionsDefinition,
2223
)
24+
from dagster._core.definitions.partitions.partition_key_range import PartitionKeyRange
25+
from dagster._core.definitions.partitions.snap import PartitionsSnap
26+
from dagster._core.definitions.partitions.subset.key_ranges import KeyRangesPartitionsSubset
2327
from dagster._core.definitions.run_request import InstigatorType
2428
from dagster._core.definitions.sensor_definition import SensorType
29+
from dagster._core.errors import DagsterInvalidInvocationError
2530
from dagster._core.instance import DagsterInstance
2631
from dagster._core.remote_origin import RemoteInstigatorOrigin
2732
from dagster._core.scheduler.instigation import (
@@ -527,6 +532,34 @@ def _get_condition_evaluation(
527532
subsets_with_metadata=[],
528533
)
529534

535+
def _get_key_range_condition_evaluation(
536+
self,
537+
asset_key: AssetKey,
538+
description: str,
539+
partition_key_ranges: Sequence[PartitionKeyRange],
540+
) -> AutomationConditionEvaluation:
541+
partitions_snap = PartitionsSnap.from_def(DynamicPartitionsDefinition(name="dynamic"))
542+
subset = SerializableEntitySubset(
543+
key=asset_key,
544+
value=KeyRangesPartitionsSubset(
545+
partitions_snap=partitions_snap,
546+
key_ranges=partition_key_ranges,
547+
),
548+
)
549+
return AutomationConditionEvaluation(
550+
condition_snapshot=AutomationConditionNodeSnapshot(
551+
class_name="...",
552+
description=description,
553+
unique_id=str(random.randint(0, 100000000)),
554+
),
555+
true_subset=subset,
556+
candidate_subset=HistoricalAllPartitionsSubsetSentinel(),
557+
start_timestamp=123,
558+
end_timestamp=456,
559+
child_evaluations=[],
560+
subsets_with_metadata=[],
561+
)
562+
530563
def test_get_evaluations_with_partitions(self, graphql_context: WorkspaceRequestContext):
531564
asset_key = AssetKey("upstream_static_partitioned_asset")
532565
partitions_def = StaticPartitionsDefinition(["a", "b", "c", "d", "e", "f"])
@@ -822,6 +855,86 @@ def _get_node(id):
822855
"d",
823856
}
824857

858+
def test_get_evaluations_tolerates_deleted_historical_partition_size(
859+
self, graphql_context: WorkspaceRequestContext
860+
):
861+
asset_key = AssetKey("dynamic_partitioned_asset")
862+
graphql_context.instance.add_dynamic_partitions("dynamic", ["a"])
863+
evaluation = self._get_key_range_condition_evaluation(
864+
asset_key,
865+
"stale historical subset",
866+
[PartitionKeyRange("a", "a")],
867+
)
868+
869+
check.not_none(
870+
graphql_context.instance.schedule_storage
871+
).add_auto_materialize_asset_evaluations(
872+
evaluation_id=201,
873+
asset_evaluations=[
874+
AutomationConditionEvaluationWithRunIds(evaluation=evaluation, run_ids=frozenset())
875+
],
876+
)
877+
878+
with patch.object(
879+
KeyRangesPartitionsSubset,
880+
"__len__",
881+
side_effect=DagsterInvalidInvocationError("Partition key was deleted."),
882+
):
883+
results = execute_dagster_graphql(
884+
graphql_context,
885+
QUERY,
886+
variables={
887+
"assetKey": {"path": ["dynamic_partitioned_asset"]},
888+
"limit": 10,
889+
"cursor": None,
890+
},
891+
)
892+
893+
assert not results.errors
894+
records = results.data["assetConditionEvaluationRecordsOrError"]["records"]
895+
assert len(records) == 1
896+
assert records[0]["numRequested"] == 0
897+
assert records[0]["evaluationNodes"][0]["numTrue"] == 0
898+
899+
def test_get_partition_evaluation_tolerates_deleted_historical_partition_membership(
900+
self, graphql_context: WorkspaceRequestContext
901+
):
902+
asset_key = AssetKey("dynamic_partitioned_asset")
903+
graphql_context.instance.add_dynamic_partitions("dynamic", ["a"])
904+
evaluation = self._get_key_range_condition_evaluation(
905+
asset_key,
906+
"stale historical subset",
907+
[PartitionKeyRange("a", "a")],
908+
)
909+
910+
check.not_none(
911+
graphql_context.instance.schedule_storage
912+
).add_auto_materialize_asset_evaluations(
913+
evaluation_id=202,
914+
asset_evaluations=[
915+
AutomationConditionEvaluationWithRunIds(evaluation=evaluation, run_ids=frozenset())
916+
],
917+
)
918+
919+
with patch.object(
920+
KeyRangesPartitionsSubset,
921+
"__contains__",
922+
side_effect=DagsterInvalidInvocationError("Partition key was deleted."),
923+
):
924+
results = execute_dagster_graphql(
925+
graphql_context,
926+
LEGACY_QUERY_FOR_SPECIFIC_PARTITION,
927+
variables={
928+
"assetKey": {"path": ["dynamic_partitioned_asset"]},
929+
"partition": "a",
930+
"evaluationId": 202,
931+
},
932+
)
933+
934+
assert not results.errors
935+
evaluation = results.data["assetConditionEvaluationForPartition"]
936+
assert evaluation["evaluationNodes"][0]["status"] == "FALSE"
937+
825938
def test_since_metadata_field(self, graphql_context: WorkspaceRequestContext):
826939
"""Test that the sinceMetadata field is correctly populated for SinceCondition evaluations."""
827940
asset_key = AssetKey("test_asset")

0 commit comments

Comments
 (0)