remove asset_graph argument from AssetGraphSubset.get_asset_subset (#31197)

Summary:
Goal is to move as much as possible of the asset-graph-dependant logic here to AssetGraphView and treat AssetGraphSubset as a relatively dumb data class.

Test Plan: BK

Author: gibsondan

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

@@ -238,9 +238,12 @@ def get_entity_subset_from_asset_graph_subset(
             self.asset_graph.has(key), f"Asset graph does not contain {key.to_user_string()}"
         )
 
-        serializable_subset = self._with_current_partitions_def(
-            asset_graph_subset.get_asset_subset(key, self.asset_graph)
-        )
+        serializable_subset = asset_graph_subset.get_asset_subset(key)
+
+        if not serializable_subset:
+            return self.get_empty_subset(key=key)
+
+        serializable_subset = self._with_current_partitions_def(serializable_subset)
 
         return EntitySubset(
             self, key=key, value=_ValidatedEntitySubsetValue(serializable_subset.value)
@@ -626,12 +629,8 @@ async def _compute_run_in_progress_asset_subset(self, key: AssetKey) -> EntitySu
     async def _compute_backfill_in_progress_asset_subset(
         self, key: AssetKey
     ) -> EntitySubset[AssetKey]:
-        value = (
-            self._queryer.get_active_backfill_in_progress_asset_graph_subset()
-            .get_asset_subset(asset_key=key, asset_graph=self.asset_graph)
-            .value
-        )
-        return EntitySubset(self, key=key, value=_ValidatedEntitySubsetValue(value))
+        asset_graph_subset = self._queryer.get_active_backfill_in_progress_asset_graph_subset()
+        return self.get_entity_subset_from_asset_graph_subset(asset_graph_subset, key)
 
     async def _compute_execution_failed_unpartitioned(self, key: AssetKey) -> bool:
         from dagster._core.event_api import AssetRecordsFilter

python_modules/dagster/dagster/_core/definitions/assets/graph/asset_graph_subset.py

@@ -67,35 +67,15 @@ def num_partitions_and_non_partitioned_assets(self) -> int:
     def is_empty(self) -> bool:
         return len(self.asset_keys) == 0
 
-    def _get_serializable_entity_subset(
-        self, asset_key: AssetKey
-    ) -> SerializableEntitySubset[AssetKey]:
+    def get_asset_subset(self, asset_key: AssetKey) -> Optional[SerializableEntitySubset[AssetKey]]:
         if asset_key in self.non_partitioned_asset_keys:
             return SerializableEntitySubset(key=asset_key, value=True)
         elif asset_key in self.partitions_subsets_by_asset_key:
             return SerializableEntitySubset(
                 key=asset_key, value=self.partitions_subsets_by_asset_key[asset_key]
             )
         else:
-            check.failed(f"Asset {asset_key} must be part of the AssetGraphSubset")
-
-    def get_asset_subset(
-        self, asset_key: AssetKey, asset_graph: BaseAssetGraph
-    ) -> SerializableEntitySubset[AssetKey]:
-        """Returns an AssetSubset representing the subset of a specific asset that this
-        AssetGraphSubset contains.
-        """
-        if (
-            asset_key in self.non_partitioned_asset_keys
-            or asset_key in self.partitions_subsets_by_asset_key
-        ):
-            return self._get_serializable_entity_subset(asset_key)
-        else:
-            partitions_def = asset_graph.get(asset_key).partitions_def
-            return SerializableEntitySubset(
-                key=asset_key,
-                value=partitions_def.empty_subset() if partitions_def else False,
-            )
+            return None
 
     def get_partitions_subset(
         self, asset_key: AssetKey, asset_graph: Optional[BaseAssetGraph] = None
@@ -127,7 +107,7 @@ def iterate_asset_subsets(self) -> Iterable[SerializableEntitySubset[AssetKey]]:
         AssetGraphSubset contains.
         """
         for asset_key in self.asset_keys:
-            yield self._get_serializable_entity_subset(asset_key)
+            yield check.not_none(self.get_asset_subset(asset_key))
 
     def __contains__(self, asset: Union[AssetKey, AssetKeyPartitionKey]) -> bool:  # pyright: ignore[reportIncompatibleMethodOverride]
         """If asset is an AssetKeyPartitionKey, check if the given AssetKeyPartitionKey is in the

python_modules/dagster/dagster/_core/definitions/auto_materialize_rule_impls.py

@@ -1056,14 +1056,19 @@ def evaluate_for_asset(self, context: "AutomationContext") -> "AutomationResult"
             AutomationResult,
         )
 
+        # this backfilling subset is aware of the current partitions definitions, and so will
+        # be valid
+        asset_subset = (
+            context.legacy_context.instance_queryer.get_active_backfill_target_asset_graph_subset().get_asset_subset(
+                context.legacy_context.asset_key
+            )
+            or context.asset_graph_view.get_empty_subset(
+                key=context.legacy_context.asset_key
+            ).convert_to_serializable_subset()
+        )
+
         backfilling_subset = ValidAssetSubset.coerce_from_subset(
-            # this backfilling subset is aware of the current partitions definitions, and so will
-            # be valid
-            (
-                context.legacy_context.instance_queryer.get_active_backfill_target_asset_graph_subset()
-            ).get_asset_subset(
-                context.legacy_context.asset_key, context.legacy_context.asset_graph
-            ),
+            asset_subset,
             context.legacy_context.partitions_def,
         )
 

python_modules/dagster/dagster/_core/execution/asset_backfill.py

@@ -1961,8 +1961,16 @@ def _get_cant_run_because_of_parent_reason(
     if parent_node.partitions_def != candidate_node.partitions_def:
         return f"parent {parent_node.key.to_user_string()} and {candidate_node.key.to_user_string()} have different partitions definitions so they cannot be materialized in the same run. {candidate_node.key.to_user_string()} can be materialized once {parent_node.key.to_user_string()} is materialized."
 
-    parent_target_subset = target_subset.get_asset_subset(parent_asset_key, asset_graph)
-    candidate_target_subset = target_subset.get_asset_subset(candidate_asset_key, asset_graph)
+    parent_target_subset = (
+        target_subset.get_asset_subset(parent_asset_key)
+        or asset_graph_view.get_empty_subset(key=parent_asset_key).convert_to_serializable_subset()
+    )
+    candidate_target_subset = (
+        target_subset.get_asset_subset(candidate_asset_key)
+        or asset_graph_view.get_empty_subset(
+            key=candidate_asset_key
+        ).convert_to_serializable_subset()
+    )
 
     num_parent_partitions_being_requested_this_tick = parent_being_requested_this_tick_subset.size