refactor step context partition methods (#24293)

## Summary & Motivation

This is a non-behavior-changing PR that refactors some methods within
the internal execution context classes. This aim of this PR is to
prepare for #23824, by
centralizing the logic for getting the `PartitionsDefinition` for an
execution.

Working on this made me aware of a behavioral weirdness that I think we
should consider addressing in 1.9: `context.partition_keys` etc. will
return a non-`None` value in steps that target no partitioned assets, if
other steps in the same run do target partitioned assets.

## How I Tested These Changes

## Changelog

NOCHANGELOG

Author: sryza

python_modules/dagster/dagster/_core/execution/context/system.py

@@ -5,6 +5,7 @@
 """
 
 from abc import ABC, abstractmethod
+from functools import cached_property
 from typing import (
     TYPE_CHECKING,
     AbstractSet,
@@ -362,18 +363,6 @@ def scoped_resources_builder(self) -> ScopedResourcesBuilder:
     def log(self) -> DagsterLogManager:
         return self._log_manager
 
-    @property
-    def partitions_def(self) -> Optional[PartitionsDefinition]:
-        from dagster._core.definitions.job_definition import JobDefinition
-
-        job_def = self._execution_data.job_def
-        if not isinstance(job_def, JobDefinition):
-            check.failed(
-                "Can only call 'partitions_def', when using jobs, not legacy pipelines",
-            )
-        partitions_def = job_def.partitions_def
-        return partitions_def
-
     @property
     def has_partitions(self) -> bool:
         tags = self._plan_data.dagster_run.tags
@@ -386,68 +375,6 @@ def has_partitions(self) -> bool:
             )
         )
 
-    @property
-    def partition_key(self) -> str:
-        from dagster._core.definitions.multi_dimensional_partitions import (
-            MultiPartitionsDefinition,
-            get_multipartition_key_from_tags,
-        )
-
-        if not self.has_partitions:
-            raise DagsterInvariantViolationError(
-                "Cannot access partition_key for a non-partitioned run"
-            )
-
-        tags = self._plan_data.dagster_run.tags
-        if any([tag.startswith(MULTIDIMENSIONAL_PARTITION_PREFIX) for tag in tags.keys()]):
-            return get_multipartition_key_from_tags(tags)
-        elif PARTITION_NAME_TAG in tags:
-            return tags[PARTITION_NAME_TAG]
-        else:
-            range_start = tags[ASSET_PARTITION_RANGE_START_TAG]
-            range_end = tags[ASSET_PARTITION_RANGE_END_TAG]
-
-            if range_start != range_end:
-                raise DagsterInvariantViolationError(
-                    "Cannot access partition_key for a partitioned run with a range of partitions."
-                    " Call partition_key_range instead."
-                )
-            else:
-                if isinstance(self.partitions_def, MultiPartitionsDefinition):
-                    return self.partitions_def.get_partition_key_from_str(cast(str, range_start))
-                return cast(str, range_start)
-
-    @property
-    def partition_key_range(self) -> PartitionKeyRange:
-        from dagster._core.definitions.multi_dimensional_partitions import (
-            MultiPartitionsDefinition,
-            get_multipartition_key_from_tags,
-        )
-
-        if not self.has_partitions:
-            raise DagsterInvariantViolationError(
-                "Cannot access partition_key for a non-partitioned run"
-            )
-
-        tags = self._plan_data.dagster_run.tags
-        if any([tag.startswith(MULTIDIMENSIONAL_PARTITION_PREFIX) for tag in tags.keys()]):
-            multipartition_key = get_multipartition_key_from_tags(tags)
-            return PartitionKeyRange(multipartition_key, multipartition_key)
-        elif PARTITION_NAME_TAG in tags:
-            partition_key = tags[PARTITION_NAME_TAG]
-            return PartitionKeyRange(partition_key, partition_key)
-        else:
-            partition_key_range_start = tags[ASSET_PARTITION_RANGE_START_TAG]
-            if partition_key_range_start is not None:
-                if isinstance(self.partitions_def, MultiPartitionsDefinition):
-                    return PartitionKeyRange(
-                        self.partitions_def.get_partition_key_from_str(partition_key_range_start),
-                        self.partitions_def.get_partition_key_from_str(
-                            tags[ASSET_PARTITION_RANGE_END_TAG]
-                        ),
-                    )
-            return PartitionKeyRange(partition_key_range_start, tags[ASSET_PARTITION_RANGE_END_TAG])
-
     @property
     def has_partition_key(self) -> bool:
         return PARTITION_NAME_TAG in self._plan_data.dagster_run.tags
@@ -954,6 +881,100 @@ def get_output_asset_keys(self) -> AbstractSet[AssetKey]:
             output_keys.add(asset_key)
         return output_keys
 
+    @cached_property
+    def run_partitions_def(self) -> Optional[PartitionsDefinition]:
+        job_def_partitions_def = self.job_def.partitions_def
+        if job_def_partitions_def is not None:
+            return job_def_partitions_def
+
+        # In the case where a job targets assets with different PartitionsDefinitions,
+        # job_def.partitions_def will be None, but the assets targeted in this step might still be
+        # partitioned. All assets within a step are expected to either have the same partitions_def
+        # or no partitions_def. Get the partitions_def from one of the assets that has one.
+        return self.asset_partitions_def
+
+    @cached_property
+    def asset_partitions_def(self) -> Optional[PartitionsDefinition]:
+        """If the current step is executing a partitioned asset, returns the PartitionsDefinition
+        for that asset. If there are one or more partitioned assets executing in the step, they're
+        expected to all have the same PartitionsDefinition.
+        """
+        asset_layer = self.job_def.asset_layer
+        assets_def = asset_layer.assets_def_for_node(self.node_handle) if asset_layer else None
+        if assets_def is not None:
+            for asset_key in assets_def.keys:
+                partitions_def = self.job_def.asset_layer.get(asset_key).partitions_def
+                if partitions_def is not None:
+                    return partitions_def
+
+        return None
+
+    @property
+    def partition_key(self) -> str:
+        from dagster._core.definitions.multi_dimensional_partitions import (
+            MultiPartitionsDefinition,
+            get_multipartition_key_from_tags,
+        )
+
+        if not self.has_partitions:
+            raise DagsterInvariantViolationError(
+                "Cannot access partition_key for a non-partitioned run"
+            )
+
+        tags = self._plan_data.dagster_run.tags
+        if any([tag.startswith(MULTIDIMENSIONAL_PARTITION_PREFIX) for tag in tags.keys()]):
+            return get_multipartition_key_from_tags(tags)
+        elif PARTITION_NAME_TAG in tags:
+            return tags[PARTITION_NAME_TAG]
+        else:
+            range_start = tags[ASSET_PARTITION_RANGE_START_TAG]
+            range_end = tags[ASSET_PARTITION_RANGE_END_TAG]
+
+            if range_start != range_end:
+                raise DagsterInvariantViolationError(
+                    "Cannot access partition_key for a partitioned run with a range of partitions."
+                    " Call partition_key_range instead."
+                )
+            else:
+                if isinstance(self.run_partitions_def, MultiPartitionsDefinition):
+                    return self.run_partitions_def.get_partition_key_from_str(
+                        cast(str, range_start)
+                    )
+                return cast(str, range_start)
+
+    @property
+    def partition_key_range(self) -> PartitionKeyRange:
+        from dagster._core.definitions.multi_dimensional_partitions import (
+            MultiPartitionsDefinition,
+            get_multipartition_key_from_tags,
+        )
+
+        if not self.has_partitions:
+            raise DagsterInvariantViolationError(
+                "Cannot access partition_key for a non-partitioned run"
+            )
+
+        tags = self._plan_data.dagster_run.tags
+        if any([tag.startswith(MULTIDIMENSIONAL_PARTITION_PREFIX) for tag in tags.keys()]):
+            multipartition_key = get_multipartition_key_from_tags(tags)
+            return PartitionKeyRange(multipartition_key, multipartition_key)
+        elif PARTITION_NAME_TAG in tags:
+            partition_key = tags[PARTITION_NAME_TAG]
+            return PartitionKeyRange(partition_key, partition_key)
+        else:
+            partition_key_range_start = tags[ASSET_PARTITION_RANGE_START_TAG]
+            if partition_key_range_start is not None:
+                if isinstance(self.run_partitions_def, MultiPartitionsDefinition):
+                    return PartitionKeyRange(
+                        self.run_partitions_def.get_partition_key_from_str(
+                            partition_key_range_start
+                        ),
+                        self.run_partitions_def.get_partition_key_from_str(
+                            tags[ASSET_PARTITION_RANGE_END_TAG]
+                        ),
+                    )
+            return PartitionKeyRange(partition_key_range_start, tags[ASSET_PARTITION_RANGE_END_TAG])
+
     def has_asset_partitions_for_input(self, input_name: str) -> bool:
         asset_layer = self.job_def.asset_layer
         upstream_asset_key = asset_layer.asset_key_for_input(self.node_handle, input_name)
@@ -999,7 +1020,7 @@ def asset_partitions_subset_for_input(
             upstream_asset_partitions_def = asset_layer.get(upstream_asset_key).partitions_def
 
             if upstream_asset_partitions_def is not None:
-                partitions_def = assets_def.partitions_def if assets_def else None
+                partitions_def = self.asset_partitions_def if assets_def else None
                 partitions_subset = (
                     partitions_def.empty_subset().with_partition_key_range(
                         partitions_def,
@@ -1116,12 +1137,7 @@ def asset_partitions_time_window_for_output(self, output_name: str) -> TimeWindo
 
     @property
     def partition_time_window(self) -> TimeWindow:
-        asset_layer = self.job_def.asset_layer
-        partitions_def = self.job_def.partitions_def
-        if asset_layer:
-            assets_def = asset_layer.assets_def_for_node(self.node_handle)
-            if assets_def:
-                partitions_def = assets_def.partitions_def
+        partitions_def = self.run_partitions_def
 
         if partitions_def is None:
             raise DagsterInvariantViolationError("Partitions definition is not defined")

python_modules/dagster/dagster/_core/pipes/context.py

@@ -405,7 +405,7 @@ def build_external_execution_context_data(
         context.partition_time_window
         if context.has_partition_key
         and has_one_dimension_time_window_partitioning(
-            context.get_step_execution_context().partitions_def
+            context.get_step_execution_context().run_partitions_def
         )
         else None
     )