Make AssetGraphSubset and AssetBackfillData serializable

dagster-io · Nov 13, 2023 · 79652f9 · 79652f9
1 parent 71ba29e
commit 79652f9
Show file tree

Hide file tree

Showing 13 changed files with 335 additions and 182 deletions.
diff --git a/python_modules/dagster-graphql/dagster_graphql/implementation/execution/backfill.py b/python_modules/dagster-graphql/dagster_graphql/implementation/execution/backfill.py
@@ -95,7 +95,7 @@ def get_asset_backfill_preview(
 
     asset_partitions = []
 
-    for asset_key in asset_backfill_data.get_targeted_asset_keys_topological_order():
+    for asset_key in asset_backfill_data.get_targeted_asset_keys_topological_order(asset_graph):
         if asset_graph.get_partitions_def(asset_key):
             partitions_subset = asset_backfill_data.target_subset.partitions_subsets_by_asset_key[
                 asset_key

diff --git a/python_modules/dagster-graphql/dagster_graphql_tests/graphql/test_asset_backfill.py b/python_modules/dagster-graphql/dagster_graphql_tests/graphql/test_asset_backfill.py
@@ -284,7 +284,9 @@ def test_launch_asset_backfill_all_partitions_asset_selection():
                 target_subset.get_partitions_subset(AssetKey("asset2")).get_partition_keys()
                 == all_partition_keys
             )
-            assert not target_subset.get_partitions_subset(AssetKey("asset1")).get_partition_keys()
+            assert not target_subset.get_partitions_subset(
+                AssetKey("asset1"), asset_graph=repo.asset_graph
+            ).get_partition_keys()
 
 
 def test_launch_asset_backfill_partitions_by_asset():
@@ -536,9 +538,8 @@ def test_launch_asset_backfill_with_upstream_anchor_asset():
                 launch_backfill_result, instance, repo
             )
             target_subset = asset_backfill_data.target_subset
-            asset_graph = target_subset.asset_graph
+            asset_graph = repo.asset_graph
             assert target_subset == AssetGraphSubset(
-                target_subset.asset_graph,
                 partitions_subsets_by_asset_key={
                     AssetKey("hourly"): asset_graph.get_partitions_def(
                         AssetKey("hourly")
@@ -605,9 +606,8 @@ def test_launch_asset_backfill_with_two_anchor_assets():
                 launch_backfill_result, instance, repo
             )
             target_subset = asset_backfill_data.target_subset
-            asset_graph = target_subset.asset_graph
+            asset_graph = repo.asset_graph
             assert target_subset == AssetGraphSubset(
-                target_subset.asset_graph,
                 partitions_subsets_by_asset_key={
                     AssetKey("hourly1"): asset_graph.get_partitions_def(
                         AssetKey("hourly1")
@@ -663,9 +663,8 @@ def test_launch_asset_backfill_with_upstream_anchor_asset_and_non_partitioned_as
                 launch_backfill_result, instance, repo
             )
             target_subset = asset_backfill_data.target_subset
-            asset_graph = target_subset.asset_graph
+            asset_graph = repo.asset_graph
             assert target_subset == AssetGraphSubset(
-                target_subset.asset_graph,
                 non_partitioned_asset_keys={AssetKey("non_partitioned")},
                 partitions_subsets_by_asset_key={
                     AssetKey("hourly"): (

diff --git a/python_modules/dagster-graphql/dagster_graphql_tests/graphql/test_partition_backfill.py b/python_modules/dagster-graphql/dagster_graphql_tests/graphql/test_partition_backfill.py
@@ -234,6 +234,7 @@ def _execute_asset_backfill_iteration_no_side_effects(
     updated_backfill = backfill.with_asset_backfill_data(
         cast(AssetBackfillIterationResult, result).backfill_data,
         dynamic_partitions_store=graphql_context.instance,
+        asset_graph=asset_graph,
     )
     graphql_context.instance.update_backfill(updated_backfill)
 

diff --git a/python_modules/dagster/dagster/_core/definitions/asset_graph.py b/python_modules/dagster/dagster/_core/definitions/asset_graph.py
@@ -574,20 +574,19 @@ def bfs_filter_subsets(
 
         queued_subsets_by_asset_key: Dict[AssetKey, Optional[PartitionsSubset]] = {
             initial_asset_key: (
-                initial_subset.get_partitions_subset(initial_asset_key)
+                initial_subset.get_partitions_subset(initial_asset_key, self)
                 if self.get_partitions_def(initial_asset_key)
                 else None
             ),
         }
-        result = AssetGraphSubset(self)
+        result = AssetGraphSubset()
 
         while len(queue) > 0:
             asset_key = queue.popleft()
             partitions_subset = queued_subsets_by_asset_key.get(asset_key)
 
             if condition_fn(asset_key, partitions_subset):
                 result |= AssetGraphSubset(
-                    self,
                     non_partitioned_asset_keys={asset_key} if partitions_subset is None else set(),
                     partitions_subsets_by_asset_key=(
                         {asset_key: partitions_subset} if partitions_subset is not None else {}