feat(task-sdk): check invalid inlets or outlets before running tasks

Author: Lee-W

airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py

@@ -345,3 +345,9 @@ class TaskStatesResponse(BaseModel):
     """Response for task states with run_id, task and state."""
 
     task_states: dict[str, Any]
+
+
+class InvalidAssetsResponse(BaseModel):
+    """Response for invalid assets."""
+
+    invalid_assets: Annotated[list[AssetProfile], Field(default_factory=list)]

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

@@ -17,12 +17,15 @@
 
 from __future__ import annotations
 
+import contextlib
+import itertools
 import json
 from collections import defaultdict
 from collections.abc import Iterator
 from typing import TYPE_CHECKING, Annotated, Any
 from uuid import UUID
 
+import attrs
 import structlog
 from cadwyn import VersionedAPIRouter
 from fastapi import Body, HTTPException, Query, status
@@ -37,6 +40,7 @@
 from airflow.api_fastapi.common.db.common import SessionDep
 from airflow.api_fastapi.common.types import UtcDateTime
 from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
+    InvalidAssetsResponse,
     PrevSuccessfulDagRunResponse,
     TaskStatesResponse,
     TIDeferredStatePayload,
@@ -51,13 +55,16 @@
     TITerminalStatePayload,
 )
 from airflow.api_fastapi.execution_api.deps import JWTBearerTIPathDep
+from airflow.exceptions import TaskNotFound
+from airflow.models.asset import AssetActive
 from airflow.models.dagbag import DagBag
 from airflow.models.dagrun import DagRun as DR
 from airflow.models.taskinstance import TaskInstance as TI, _stop_remaining_tasks
 from airflow.models.taskreschedule import TaskReschedule
 from airflow.models.trigger import Trigger
 from airflow.models.xcom import XComModel
 from airflow.sdk.definitions._internal.expandinput import NotFullyPopulated
+from airflow.sdk.definitions.asset import Asset, AssetUniqueKey
 from airflow.sdk.definitions.taskgroup import MappedTaskGroup
 from airflow.utils import timezone
 from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState
@@ -840,5 +847,67 @@ def _get_group_tasks(dag_id: str, task_group_id: str, session: SessionDep, logic
     return group_tasks
 
 
+@ti_id_router.get(
+    "/{task_instance_id}/validate-inlets-and-outlets",
+    status_code=status.HTTP_200_OK,
+    responses={
+        status.HTTP_404_NOT_FOUND: {"description": "Task Instance not found"},
+    },
+)
+def validate_inlets_and_outlets(
+    task_instance_id: UUID,
+    session: SessionDep,
+    dag_bag: DagBagDep,
+) -> InvalidAssetsResponse:
+    """Validate whehter there're inactive assets in inlets and outlets of a given task instance."""
+    ti_id_str = str(task_instance_id)
+    bind_contextvars(ti_id=ti_id_str)
+
+    ti = session.scalar(select(TI).where(TI.id == ti_id_str))
+    if not ti or not ti.logical_date:
+        log.error("Task Instance not found")
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail={
+                "reason": "not_found",
+                "message": "Task Instance not found",
+            },
+        )
+
+    if not ti.task:
+        dag = dag_bag.get_dag(ti.dag_id)
+        if dag:
+            with contextlib.suppress(TaskNotFound):
+                ti.task = dag.get_task(ti.task_id)
+
+    inlets = [asset.asprofile() for asset in ti.task.inlets if isinstance(asset, Asset)]
+    outlets = [asset.asprofile() for asset in ti.task.outlets if isinstance(asset, Asset)]
+    if not (inlets or outlets):
+        return InvalidAssetsResponse(invalid_assets=[])
+
+    all_asset_unique_keys: set[AssetUniqueKey] = {
+        AssetUniqueKey.from_asset(inlet_or_outlet)  # type: ignore
+        for inlet_or_outlet in itertools.chain(inlets, outlets)
+    }
+    active_asset_unique_keys = {
+        AssetUniqueKey(name, uri)
+        for name, uri in session.execute(
+            select(AssetActive.name, AssetActive.uri).where(
+                tuple_(AssetActive.name, AssetActive.uri).in_(
+                    attrs.astuple(key) for key in all_asset_unique_keys
+                )
+            )
+        )
+    }
+    different = all_asset_unique_keys - active_asset_unique_keys
+
+    return InvalidAssetsResponse(
+        invalid_assets=[
+            asset_unique_key.to_asset().asprofile()  # type: ignore
+            for asset_unique_key in different
+        ]
+    )
+
+
 # This line should be at the end of the file to ensure all routes are registered
 router.include_router(ti_id_router)

airflow-core/src/airflow/models/dagrun.py

@@ -1894,6 +1894,7 @@ def schedule_tis(
                 and not ti.task.on_execute_callback
                 and not ti.task.on_success_callback
                 and not ti.task.outlets
+                and not ti.task.inlets
             ):
                 empty_ti_ids.append(ti.id)
             # check "start_trigger_args" to see whether the operator supports start execution from triggerer

task-sdk/src/airflow/sdk/api/client.py

@@ -40,6 +40,7 @@
     ConnectionResponse,
     DagRunStateResponse,
     DagRunType,
+    InvalidAssetsResponse,
     PrevSuccessfulDagRunResponse,
     TaskInstanceState,
     TaskStatesResponse,
@@ -272,6 +273,11 @@ def get_task_states(
         resp = self.client.get("task-instances/states", params=params)
         return TaskStatesResponse.model_validate_json(resp.read())
 
+    def validate_inlets_and_outlets(self, id: uuid.UUID) -> InvalidAssetsResponse:
+        """Validate whehter there're inactive assets in inlets and outlets of a given task instance."""
+        resp = self.client.get(f"task-instances/{id}/validate-inlets-and-outlets")
+        return InvalidAssetsResponse.model_validate_json(resp.read())
+
 
 class ConnectionOperations:
     __slots__ = ("client",)

task-sdk/src/airflow/sdk/api/datamodels/_generated.py

@@ -168,6 +168,14 @@ class IntermediateTIState(str, Enum):
     DEFERRED = "deferred"
 
 
+class InvalidAssetsResponse(BaseModel):
+    """
+    Response for invalid assets.
+    """
+
+    invalid_assets: Annotated[list[AssetProfile] | None, Field(title="Invalid Assets")] = None
+
+
 class PrevSuccessfulDagRunResponse(BaseModel):
     """
     Schema for response with previous successful DagRun information for Task Template Context.

task-sdk/src/airflow/sdk/execution_time/comms.py

@@ -61,6 +61,7 @@
     BundleInfo,
     ConnectionResponse,
     DagRunStateResponse,
+    InvalidAssetsResponse,
     PrevSuccessfulDagRunResponse,
     TaskInstance,
     TaskInstanceState,
@@ -206,6 +207,24 @@ def source_task_instance(self) -> AssetEventSourceTaskInstance | None:
         )
 
 
+class InvalidAssetsResult(InvalidAssetsResponse):
+    """Response of InvalidAssets requests."""
+
+    type: Literal["InvalidAssetsResult"] = "InvalidAssetsResult"
+
+    @classmethod
+    def from_invalid_assets_response(
+        cls, invalid_assets_response: InvalidAssetsResponse
+    ) -> InvalidAssetsResult:
+        """
+        Get InvalidAssetsResponse from InvalidAssetsResult.
+
+        InvalidAssetsResponse is autogenerated from the API schema, so we need to convert it to InvalidAssetsResult
+        for communication between the Supervisor and the task process.
+        """
+        return cls(**invalid_assets_response.model_dump(exclude_defaults=True), type="InvalidAssetsResult")
+
+
 class XComResult(XComResponse):
     """Response to ReadXCom request."""
 
@@ -354,6 +373,7 @@ class OKResponse(BaseModel):
         VariableResult,
         XComResult,
         XComCountResponse,
+        InvalidAssetsResult,
         OKResponse,
     ],
     Field(discriminator="type"),
@@ -557,6 +577,11 @@ class GetAssetEventByAssetAlias(BaseModel):
     type: Literal["GetAssetEventByAssetAlias"] = "GetAssetEventByAssetAlias"
 
 
+class ValidateInletsAndOutlets(BaseModel):
+    ti_id: UUID
+    type: Literal["ValidateInletsAndOutlets"] = "ValidateInletsAndOutlets"
+
+
 class GetPrevSuccessfulDagRun(BaseModel):
     ti_id: UUID
     type: Literal["GetPrevSuccessfulDagRun"] = "GetPrevSuccessfulDagRun"
@@ -623,6 +648,7 @@ class GetDRCount(BaseModel):
         SetXCom,
         SkipDownstreamTasks,
         SucceedTask,
+        ValidateInletsAndOutlets,
         TaskState,
         TriggerDagRun,
         DeleteVariable,

task-sdk/src/airflow/sdk/execution_time/supervisor.py

@@ -87,6 +87,7 @@
     GetXCom,
     GetXComCount,
     GetXComSequenceItem,
+    InvalidAssetsResult,
     PrevSuccessfulDagRunResult,
     PutVariable,
     RescheduleTask,
@@ -100,6 +101,7 @@
     TaskStatesResult,
     ToSupervisor,
     TriggerDagRun,
+    ValidateInletsAndOutlets,
     VariableResult,
     XComCountResponse,
     XComResult,
@@ -1154,6 +1156,10 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger):
             )
         elif isinstance(msg, DeleteVariable):
             resp = self.client.variables.delete(msg.key)
+        elif isinstance(msg, ValidateInletsAndOutlets):
+            invalid_assets_resp = self.client.task_instances.validate_inlets_and_outlets(msg.ti_id)
+            resp = InvalidAssetsResult.from_invalid_assets_response(invalid_assets_resp)
+            dump_opts = {"exclude_unset": True}
         else:
             log.error("Unhandled request", msg=msg)
             return

task-sdk/src/airflow/sdk/execution_time/task_runner.py

@@ -40,6 +40,7 @@
 
 from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersionLock
 from airflow.dag_processing.bundles.manager import DagBundlesManager
+from airflow.exceptions import AirflowInactiveAssetInInletOrOutletException
 from airflow.listeners.listener import get_listener_manager
 from airflow.sdk.api.datamodels._generated import (
     AssetProfile,
@@ -66,6 +67,7 @@
     GetTaskRescheduleStartDate,
     GetTaskStates,
     GetTICount,
+    InvalidAssetsResult,
     RescheduleTask,
     RetryTask,
     SetRenderedFields,
@@ -79,6 +81,7 @@
     ToSupervisor,
     ToTask,
     TriggerDagRun,
+    ValidateInletsAndOutlets,
 )
 from airflow.sdk.execution_time.context import (
     ConnectionAccessor,
@@ -778,6 +781,8 @@ def _prepare(ti: RuntimeTaskInstance, log: Logger, context: Context) -> ToSuperv
         # so that we do not call the API unnecessarily
         SUPERVISOR_COMMS.send_request(log=log, msg=SetRenderedFields(rendered_fields=rendered_fields))
 
+    _validate_task_inlets_and_outlets(ti=ti, log=log)
+
     try:
         # TODO: Call pre execute etc.
         get_listener_manager().hook.on_task_instance_running(
@@ -790,6 +795,17 @@ def _prepare(ti: RuntimeTaskInstance, log: Logger, context: Context) -> ToSuperv
     return None
 
 
+def _validate_task_inlets_and_outlets(*, ti: RuntimeTaskInstance, log: Logger) -> None:
+    SUPERVISOR_COMMS.send_request(msg=ValidateInletsAndOutlets(ti_id=ti.id), log=log)
+    invalid_assets_resp = SUPERVISOR_COMMS.get_message()
+    if TYPE_CHECKING:
+        assert isinstance(invalid_assets_resp, InvalidAssetsResult)
+    if invalid_assets := invalid_assets_resp.invalid_assets:
+        raise AirflowInactiveAssetInInletOrOutletException(
+            inactive_asset_keys=[AssetUniqueKey.from_asset(asset) for asset in invalid_assets]
+        )
+
+
 def _defer_task(
     defer: TaskDeferred, ti: RuntimeTaskInstance, log: Logger
 ) -> tuple[ToSupervisor, TaskInstanceState]: