TaskFlow AirflowSkipException causes downstream step to fail when multiple_outputs is true

Apache Airflow version

2.5.1

What happened

Most of our code is based on TaskFlow API and we have many tasks that raise AirflowSkipException (or BranchPythonOperator) on purpose to skip the next downstream task (with trigger_rule = none_failed_min_one_success).
And these tasks are expecting a multiple output XCom result (local_file_path, file sizes, records count) from previous tasks and it's causing this error:
airflow.exceptions.XComNotFound: XComArg result from copy_from_data_lake_to_local_file at outbound_dag_AIR2070 with key="local_file_path" is not found!

What you think should happen instead

Considering trigger rule "none_failed_min_one_success", we expect that upstream task should be allowed to skip and downstream tasks will still run without raising any errors caused by not found XCom results.

How to reproduce

This is an aproximate example dag based on an existing one.

from os import path
import pendulum
from airflow import DAG
from airflow.decorators import task
from airflow.operators.python import BranchPythonOperator

DAG_ID = "testing_dag_AIR"
# PGP_OPERATION = None
PGP_OPERATION = "decrypt"
LOCAL_FILE_PATH = "/temp/example/example.csv"

with DAG(
    dag_id=DAG_ID,
    schedule='0 7-18 * * *',
    start_date=pendulum.datetime(2022, 12, 15, 7, 0, 0),
) as dag:
    @task(multiple_outputs=True, trigger_rule='none_failed_min_one_success')
    def copy_from_local_file_to_data_lake(local_file_path: str, dest_dir_path: str):
        destination_file_path = path.join(dest_dir_path, path.basename(local_file_path))

        return {
            "destination_file_path": destination_file_path,
            "file_size": 100
        }

    @task(multiple_outputs=True, trigger_rule='none_failed_min_one_success')
    def copy_from_data_lake_to_local_file(data_lake_file_path, local_dir_path):
        local_file_path = path.join(local_dir_path, path.basename(data_lake_file_path))

        return {
            "local_file_path": local_file_path,
            "file_size": 100
        }

    @task(multiple_outputs=True, task_id='get_pgp_file_info', trigger_rule='none_failed_min_one_success')
    def get_pgp_file_info(file_path, operation):
        import uuid
        import os

        src_file_name = os.path.basename(file_path)
        src_file_dir = os.path.dirname(file_path)
        run_id = str(uuid.uuid4())

        if operation == "decrypt":
            wait_pattern = f'*{src_file_name}'
        else:
            wait_pattern = f'*{src_file_name}.pgp'

        target_path = 'datalake/target'

        return {
            'src_file_path': file_path,
            'src_file_dir': src_file_dir,
            'target_path': target_path,
            'pattern': wait_pattern,
            'guid': run_id
        }

    @task(multiple_outputs=True, task_id='return_src_path', trigger_rule='none_failed_min_one_success')
    def return_src_path(src_file_path):
        return {
            'file_path': src_file_path,
            'file_size': 100
        }

    @task(multiple_outputs=True, task_id='choose_result', trigger_rule='none_failed_min_one_success')
    def choose_result(src_file_path, src_file_size, decrypt_file_path, decrypt_file_size):
        import os

        file_path = decrypt_file_path or src_file_path
        file_size = decrypt_file_size or src_file_size

        local_dir = os.path.dirname(file_path)

        return {
            'local_dir': local_dir,
            'file_path': file_path,
            'file_size': file_size,
            'file_name': os.path.basename(file_path)
        }

    def switch_branch_func(pgp_operation):
        if pgp_operation in ["decrypt", "encrypt"]:
            return 'get_pgp_file_info'
        else:
            return 'return_src_path'

    operation = PGP_OPERATION
    local_file_path = LOCAL_FILE_PATH

    check_need_to_decrypt = BranchPythonOperator(
        task_id='branch_task',
        python_callable=switch_branch_func,
        op_args=(operation,))

    pgp_file_info = get_pgp_file_info(local_file_path, operation)

    data_lake_file = copy_from_local_file_to_data_lake(pgp_file_info['src_file_path'], pgp_file_info['target_path'])

    decrypt_local_file = copy_from_data_lake_to_local_file(
        data_lake_file['destination_file_path'], pgp_file_info['src_file_dir'])

    src_result = return_src_path(local_file_path)

    result = choose_result(src_result['file_path'], src_result['file_size'],
                           decrypt_local_file['local_file_path'], decrypt_local_file['file_size'])

    check_need_to_decrypt >> [pgp_file_info, src_result]
    pgp_file_info >> decrypt_local_file
    [decrypt_local_file, src_result] >> result

Operating System

Windows 10

Versions of Apache Airflow Providers

No response

Deployment

Docker-Compose

Deployment details

docker-compose version: 3.7

Note: This also happens when it's deployed to one of our testing environments using official Airflow Helm Chart.

Anything else

This issue is similar to #24338, it was solved by #25661 but this case is related to multiple_outputs being set to True.

Are you willing to submit PR?

Yes I am willing to submit a PR!

Code of Conduct

I agree to follow this project's Code of Conduct

Thanks for opening your first issue here! Be sure to follow the issue template!

stuck with this too, hoping for this feature to be released soon

I also have this problem, please your help to solve it

Same here. We ran into this issue when trying to upgrade from v2.2.3 to v2.5.1

@renzo-sanchez-h please can you write a simple dag to reproduce this scenario. The dag above cannot run outside your environment. Have tweaked it but I can't reproduce the issue

@uranusjr

This fixes it:

diff --git a/airflow/models/xcom_arg.py b/airflow/models/xcom_arg.py
index 133fd4280b..1c7290f794 100644
--- a/airflow/models/xcom_arg.py
+++ b/airflow/models/xcom_arg.py
@@ -340,6 +340,8 @@ class PlainXComArg(XComArg):
             return result
         if self.key == XCOM_RETURN_KEY:
             return None
+        if isinstance(result, ArgNotSet):
+            return None
         raise XComNotFound(ti.dag_id, task_id, self.key)

But I'm wondering what the implications are. WDYT @uranusjr

This would change the behaviour of zip. ArgNotSet is fillvalue’s default value, and it needs to cause XComNotFound for the mechanism to work correctly. So either further logic is needed to distinguish between the cases, or one of these code paths need to change the sentinel value (not use ArgNotSet).

How does the XComNotFound exception propagate? It is likely better to catch this elsewhere in the stack and coerce to None.

This would change the behaviour of zip. ArgNotSet is fillvalue’s default value, and it needs to cause XComNotFound for the mechanism to work correctly. So either further logic is needed to distinguish between the cases, or one of these code paths need to change the sentinel value (not use ArgNotSet).

How does the XComNotFound exception propagate? It is likely better to catch this elsewhere in the stack and coerce to None.

This worked but I'm not sure how it will behave in the case of mapped tasks since I'm not querying with map_index:

from airflow.models.taskinstance import TaskInstance as TI
pti = session.query(TI).filter(
    TI.dag_id == ti.dag_id,
    TI.task_id == task_id,
    TI.run_id == ti.run_id,
).first()
if pti.state == State.SKIPPED:
    return None

Where does the query go? It may (or may not) make sense depending on the situation.

I think another problem is that technically the error is correct in that the XCom value is indeed missing. If we want to return None for SKIPPED, arguably the same should be done for other states as well.

Here's the full diff:

diff --git a/airflow/models/xcom_arg.py b/airflow/models/xcom_arg.py
index 133fd4280b..046bc4a6e2 100644
--- a/airflow/models/xcom_arg.py
+++ b/airflow/models/xcom_arg.py
@@ -33,6 +33,7 @@ from airflow.utils.context import Context
 from airflow.utils.edgemodifier import EdgeModifier
 from airflow.utils.mixins import ResolveMixin
 from airflow.utils.session import NEW_SESSION, provide_session
+from airflow.utils.state import State
 from airflow.utils.types import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:
@@ -340,6 +341,14 @@ class PlainXComArg(XComArg):
             return result
         if self.key == XCOM_RETURN_KEY:
             return None
+        from airflow.models.taskinstance import TaskInstance as TI
+        pti = session.query(TI).filter(
+            TI.dag_id == ti.dag_id,
+            TI.task_id == task_id,
+            TI.run_id == ti.run_id,
+        ).first()
+        if pti.state == State.SKIPPED:
+            return None
         raise XComNotFound(ti.dag_id, task_id, self.key)

Return None if an XComArg fails to resolve in a multiple_outputs Task…

Navigation Menu

Repository navigation

TaskFlow AirflowSkipException causes downstream step to fail when multiple_outputs is true #29199

Apache Airflow version

What happened

What you think should happen instead

How to reproduce

Operating System

Versions of Apache Airflow Providers

Deployment

Deployment details

Anything else

Are you willing to submit PR?

Code of Conduct

18 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

TaskFlow AirflowSkipException causes downstream step to fail when multiple_outputs is true #29199

Description

Apache Airflow version

What happened

What you think should happen instead

How to reproduce

Operating System

Versions of Apache Airflow Providers

Deployment

Deployment details

Anything else

Are you willing to submit PR?

Code of Conduct

Activity

boring-cyborg commented on Jan 27, 2023

CarrascoEr commented on Feb 2, 2023

emilymitacc commented on Feb 2, 2023

gmontero06 commented on Feb 2, 2023

ephraimbuddy commented on Feb 21, 2023

ephraimbuddy commented on Feb 21, 2023

uranusjr commented on Feb 22, 2023

ephraimbuddy commented on Feb 22, 2023

uranusjr commented on Feb 22, 2023

ephraimbuddy commented on Feb 22, 2023

18 remaining items

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions