[dj] Wire job entity evaluation and run request pipeline for Declarative Automation (#22184)

Updates the core run launching machinery for the asset daemon to kick
off runs of the relevant jobs when interacting w/ automation results
associated with an asset job (in contrast to the previous behavior of
always just launching runs against the `__ASSET_JOB`)

---------

Co-authored-by: Catherine Noll <catherine@dagsterlabs.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Internal-RevId: ebbc2c79ef0a05a340802990bdffaf3076a31180

Author: 3 people

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

@@ -331,12 +331,8 @@ def get_subset_from_serializable_subset(
         self, serializable_subset: SerializableEntitySubset[T_EntityKey]
     ) -> EntitySubset[T_EntityKey] | None:
         key = serializable_subset.key
-        # serialized subsets come from persisted cursors, which may reference key types this
-        # view does not yet support
-        if (
-            isinstance(key, (AssetKey, AssetCheckKey))
-            and self.asset_graph.has(key)
-            and serializable_subset.is_compatible_with_partitions_def(self._get_partitions_def(key))
+        if self.asset_graph.has(key) and serializable_subset.is_compatible_with_partitions_def(
+            self._get_partitions_def(key)
         ):
             return EntitySubset(
                 self, key=key, value=_ValidatedEntitySubsetValue(serializable_subset.value)

python_modules/dagster/dagster/_core/definitions/automation_condition_sensor_definition.py

@@ -1,14 +1,20 @@
 from collections.abc import Mapping, Sequence
 from functools import partial
-from typing import Any, cast
+from typing import TYPE_CHECKING, AbstractSet, Any, Union, cast  # noqa: UP035
 
 import dagster._check as check
-from dagster._annotations import beta_param, public
+from dagster._annotations import (
+    beta_param,
+    hidden_param,
+    only_allow_hidden_params_in_kwargs,
+    public,
+)
+from dagster._core.definitions.asset_key import AssetJobKey
 from dagster._core.definitions.asset_selection import AssetSelection, CoercibleToAssetSelection
 from dagster._core.definitions.declarative_automation.automation_condition import (
     AutomationCondition,
 )
-from dagster._core.definitions.metadata import RawMetadataMapping
+from dagster._core.definitions.metadata import JsonMetadataValue, RawMetadataMapping
 from dagster._core.definitions.run_request import SensorResult
 from dagster._core.definitions.sensor_definition import (
     DefaultSensorStatus,
@@ -22,11 +28,46 @@
 from dagster._utils import IHasInternalInit
 from dagster._utils.tags import normalize_tags
 
+if TYPE_CHECKING:
+    from dagster._core.remote_representation.external_data import SensorMetadataSnap
+
 MAX_ENTITIES = 500
 EMIT_BACKFILLS_METADATA_KEY = "dagster/emit_backfills"
+ASSET_JOB_KEYS_METADATA_KEY = "dagster/asset_job_keys"
 DEFAULT_AUTOMATION_CONDITION_SENSOR_NAME = "default_automation_condition_sensor"
 
 
+def asset_job_keys_from_sensor_metadata(
+    metadata: Union[Mapping[str, Any], "SensorMetadataSnap", None],
+) -> AbstractSet[AssetJobKey]:
+    """Parses the asset-job keys a sensor is responsible for evaluating out of its metadata.
+
+    Accepts two shapes because get_default_automation_condition_sensor_target iterates a
+    mixed list of sensors: a SensorDefinition's metadata (a normalized
+    Mapping[str, MetadataValue]) or a RemoteSensor's SensorMetadataSnap. Missing metadata
+    or a missing key means the sensor claims no job keys; a present key with unexpected
+    contents is an error.
+    """
+    # local import: external_data imports from this module
+    from dagster._core.remote_representation.external_data import SensorMetadataSnap
+
+    mapping: Mapping[str, Any] | None = (
+        metadata.standard_metadata if isinstance(metadata, SensorMetadataSnap) else metadata
+    )
+    if not mapping or ASSET_JOB_KEYS_METADATA_KEY not in mapping:
+        return set()
+    # the sorted list of job names written in __init__ normalizes to a JsonMetadataValue,
+    # on the definition side and (after the serdes round-trip) on the snapshot side alike
+    job_names = check.inst(
+        mapping[ASSET_JOB_KEYS_METADATA_KEY],
+        JsonMetadataValue,
+        f"unexpected value for sensor metadata key {ASSET_JOB_KEYS_METADATA_KEY}",
+    ).value
+    return {
+        AssetJobKey(job_name=check.str_param(job_name, "job_name")) for job_name in job_names or []
+    }
+
+
 def _evaluate(sensor_def: "AutomationConditionSensorDefinition", context: SensorEvaluationContext):
     from dagster._core.definitions.automation_tick_evaluation_context import (
         AutomationTickEvaluationContext,
@@ -51,6 +92,7 @@ def _evaluate(sensor_def: "AutomationConditionSensorDefinition", context: Sensor
         observe_run_tags={},
         auto_observe_asset_keys=set(),
         asset_selection=sensor_def.asset_selection,
+        asset_job_keys=sensor_def.asset_job_keys,
         emit_backfills=sensor_def.emit_backfills,
         default_condition=sensor_def.default_condition,
         logger=context.log,
@@ -81,6 +123,13 @@ def not_supported(context) -> None:
 @public
 @beta_param(param="use_user_code_server")
 @beta_param(param="default_condition")
+@hidden_param(
+    param="asset_job_keys",
+    breaking_version="",
+    additional_warn_text="This parameter is internal: it is populated automatically on the "
+    "default automation condition sensor and will be replaced by a job-selection API.",
+    emit_runtime_warning=False,
+)
 class AutomationConditionSensorDefinition(SensorDefinition, IHasInternalInit):
     """Targets a set of assets and repeatedly evaluates all the AutomationConditions on all of
     those assets to determine which to request runs for.
@@ -159,7 +208,9 @@ def __init__(
         emit_backfills: bool = True,
         use_user_code_server: bool = False,
         default_condition: AutomationCondition | None = None,
+        **kwargs: Any,
     ):
+        only_allow_hidden_params_in_kwargs(AutomationConditionSensorDefinition, kwargs)
         self._use_user_code_server = use_user_code_server
         check.bool_param(emit_backfills, "allow_backfills")
 
@@ -180,6 +231,20 @@ def __init__(
         if emit_backfills:
             metadata = {**(metadata or {}), EMIT_BACKFILLS_METADATA_KEY: True}
 
+        # asset-job keys are carried in metadata so they cross the snapshot boundary to
+        # the daemon (the same channel emit_backfills uses)
+        asset_job_keys = kwargs.get("asset_job_keys")
+        if asset_job_keys:
+            metadata = {
+                **(metadata or {}),
+                ASSET_JOB_KEYS_METADATA_KEY: sorted(
+                    key.job_name
+                    for key in check.set_param(
+                        asset_job_keys, "asset_job_keys", of_type=AssetJobKey
+                    )
+                ),
+            }
+
         super().__init__(
             name=check_valid_name(name),
             job_name=None,
@@ -211,6 +276,11 @@ def emit_backfills(self) -> bool:
     def default_condition(self) -> AutomationCondition | None:
         return self._default_condition
 
+    @property
+    def asset_job_keys(self) -> AbstractSet[AssetJobKey]:
+        """The asset-job entity keys this sensor is responsible for evaluating."""
+        return asset_job_keys_from_sensor_metadata(self.metadata)
+
     @property
     def sensor_type(self) -> SensorType:
         return SensorType.AUTOMATION if self._use_user_code_server else SensorType.AUTO_MATERIALIZE
@@ -229,6 +299,7 @@ def dagster_internal_init(  # type: ignore
         emit_backfills: bool,
         use_user_code_server: bool,
         default_condition: AutomationCondition | None,
+        **kwargs: Any,
     ) -> "AutomationConditionSensorDefinition":
         return AutomationConditionSensorDefinition(
             name=name,
@@ -242,6 +313,7 @@ def dagster_internal_init(  # type: ignore
             emit_backfills=emit_backfills,
             use_user_code_server=use_user_code_server,
             default_condition=default_condition,
+            **kwargs,
         )
 
     def with_attributes(

python_modules/dagster/dagster/_core/definitions/automation_tick_evaluation_context.py

@@ -10,7 +10,12 @@
 from dagster import PartitionKeyRange
 from dagster._core.asset_graph_view.entity_subset import EntitySubset
 from dagster._core.definitions.asset_daemon_cursor import AssetDaemonCursor
-from dagster._core.definitions.asset_key import AssetCheckKey, AssetOrCheckKey, EntityKey
+from dagster._core.definitions.asset_key import (
+    AssetCheckKey,
+    AssetJobKey,
+    AssetOrCheckKey,
+    EntityKey,
+)
 from dagster._core.definitions.asset_selection import AssetSelection
 from dagster._core.definitions.assets.graph.asset_graph_subset import AssetGraphSubset
 from dagster._core.definitions.assets.graph.base_asset_graph import BaseAssetGraph
@@ -57,14 +62,23 @@ def __init__(
         emit_backfills: bool,
         default_condition: AutomationCondition | None = None,
         evaluation_time: datetime.datetime | None = None,
+        asset_job_keys: AbstractSet[AssetJobKey] | None = None,
     ):
-        resolved_entity_keys = {
+        resolved_entity_keys: set[EntityKey] = {
             entity_key
             for entity_key in (
                 asset_selection.resolve(asset_graph) | asset_selection.resolve_checks(asset_graph)
             )
             if default_condition or asset_graph.get(entity_key).automation_condition is not None
         }
+        # asset selections cannot express job keys, so conditioned jobs are added
+        # directly from the graph
+        resolved_entity_keys |= {
+            key
+            for key in (asset_job_keys or set())
+            # default_condition does not extend to jobs
+            if asset_graph.has(key) and asset_graph.get(key).automation_condition is not None
+        }
         self._total_keys = len(resolved_entity_keys)
         self._evaluation_id = evaluation_id
         self._evaluator = AutomationConditionEvaluator(
@@ -323,6 +337,44 @@ def _flood_fill_asset_subsets(k: AssetOrCheckKey):
     return backfill_request, list(entity_subsets_by_key.values())
 
 
+def _build_job_run_requests(
+    asset_job_subsets: Sequence[EntitySubset[AssetJobKey]],
+    asset_graph: BaseAssetGraph,
+    run_tags: Mapping[str, str] | None,
+) -> Sequence[RunRequest]:
+    """Builds RunRequests for job entities whose automation conditions fired this tick.
+
+    Unlike asset run requests, which list the specific assets and checks to execute, a
+    job run request launches the whole job by name: job_name is set and there is no
+    asset selection.
+
+    An unpartitioned job launches as a single run. A partitioned job launches as one
+    run per requested partition, each tagged with its partition key.
+    """
+    run_requests = []
+    for subset in asset_job_subsets:
+        if subset.is_empty:
+            continue
+        partitions_def = asset_graph.get(subset.key).partitions_def
+        if partitions_def is None:
+            run_requests.append(
+                RunRequest(job_name=subset.key.job_name, tags=dict(run_tags) if run_tags else {})
+            )
+        else:
+            run_requests.extend(
+                RunRequest(
+                    job_name=subset.key.job_name,
+                    partition_key=partition_key,
+                    tags={
+                        **(run_tags or {}),
+                        **partitions_def.get_tags_for_partition_key(partition_key),
+                    },
+                )
+                for partition_key in sorted(subset.expensively_compute_partition_keys())
+            )
+    return sorted(run_requests, key=lambda rr: (rr.job_name or "", rr.partition_key or ""))
+
+
 def build_run_requests(
     entity_subsets: Sequence[EntitySubset],
     asset_graph: BaseAssetGraph,
@@ -332,22 +384,28 @@ def build_run_requests(
     """For a single asset in a given tick, the asset will only be part of a run or a backfill, not both.
     If the asset is targetd by a backfill, there will only be one backfill that targets the asset.
     """
+    # job entity subsets become whole-job run requests; everything downstream of this
+    # split (backfills, partition mapping) deals only in assets and checks
+    asset_job_subsets = [es for es in entity_subsets if isinstance(es.key, AssetJobKey)]
+    asset_and_check_subsets = [es for es in entity_subsets if not isinstance(es.key, AssetJobKey)]
+    job_run_requests = _build_job_run_requests(asset_job_subsets, asset_graph, run_tags)
+
     if emit_backfills:
-        backfill_run_request, entity_subsets = _build_backfill_request(
-            entity_subsets, asset_graph, run_tags
+        backfill_run_request, asset_and_check_subsets = _build_backfill_request(
+            asset_and_check_subsets, asset_graph, run_tags
         )
     else:
         backfill_run_request = None
 
     run_requests = _build_run_requests_from_partitions_def_mapping(
-        _get_mapping_from_entity_subsets(entity_subsets, asset_graph),
+        _get_mapping_from_entity_subsets(asset_and_check_subsets, asset_graph),
         asset_graph,
         run_tags,
     )
     if backfill_run_request:
         run_requests = [backfill_run_request, *run_requests]
 
-    return run_requests
+    return [*job_run_requests, *run_requests]
 
 
 def _any_check_uses_automation_condition(

python_modules/dagster/dagster/_core/definitions/job_definition.py

@@ -243,6 +243,15 @@ def __init__(
                 "_automation_condition",
                 "AutomationCondition can only be provided for asset jobs.",
             )
+            # every partitioned job synthesizes a PartitionedConfig internally, so only
+            # reject when the user actually supplied config for it to resolve
+            if self._original_config_argument is not None and self.partitioned_config is not None:
+                raise DagsterInvalidDefinitionError(
+                    f"Job '{self.name}' has both an automation_condition and partitioned run"
+                    " config. Declarative automation submits partitioned-job runs without"
+                    " resolving per-partition config, so this combination is not currently"
+                    " supported."
+                )
 
     def dagster_internal_init(
         *,