Output from traditional Task can not be used in task decorator without a *reduce* task

[MaiHoangViet1809]

Apache Airflow version

2.4.1

What happened

When build a flow between traditional task and taskflow API, I saw this error log:

[2022-10-11, 11:23:58 +07] {taskinstance.py:1851} ERROR - Task failed with exception
Traceback (most recent call last):
  File "/home/prod_account/env3.10.5/lib/python3.10/site-packages/airflow/utils/session.py", line 72, in wrapper
    return func(*args, **kwargs)
  File "/home/prod_account/env3.10.5/lib/python3.10/site-packages/airflow/models/taskinstance.py", line 2378, in xcom_push
    XCom.set(
  File "/home/prod_account/env3.10.5/lib/python3.10/site-packages/airflow/utils/session.py", line 72, in wrapper
    return func(*args, **kwargs)
  File "/home/prod_account/env3.10.5/lib/python3.10/site-packages/airflow/models/xcom.py", line 206, in set
    value = cls.serialize_value(
  File "/home/prod_account/env3.10.5/lib/python3.10/site-packages/airflow/models/xcom.py", line 597, in serialize_value
    return json.dumps(value).encode('UTF-8')
  File "/home/prod_account/apps/python3_10/lib/python3.10/json/__init__.py", line 231, in dumps
    return _default_encoder.encode(obj)
  File "/home/prod_account/apps/python3_10/lib/python3.10/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/home/prod_account/apps/python3_10/lib/python3.10/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
  File "/home/prod_account/apps/python3_10/lib/python3.10/json/encoder.py", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type _LazyXComAccess is not JSON serializable

What you think should happen instead

It a task expand call follow a S3keySensor, then I try to wrap it with another reduce (return a list or do aggregation), it work again.
I think the result push through xcom should be reduced by default to be used in next task if that next task is not a map/expand task ?

How to reproduce

this is the sample code included both working and bug flow:

from airflow import DAG
from datetime import datetime, timedelta
from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor
from airflow.sensors.base import PokeReturnValue

from airflow.decorators import task
from airflow.operators.python import get_current_context
from typing import Dict, Union

PARAMS = {
    "INPUT_FOLDER": "/input",
    "DATA_FOLDER": "/data",
}

DEBUG = True


class S3KeySensor_test(S3KeySensor):
    def poke(self, context) -> Union[bool, PokeReturnValue]:
        check_result = True  # all(self._check_key(key) for key in self.bucket_key)
        if check_result:
            return PokeReturnValue(is_done=True, xcom_value=",".join(self.bucket_key))
        else:
            return False


with DAG(
    "AA.REPORT_BUG", start_date=datetime(2022, 1, 1), catchup=False, schedule=None, params=PARAMS
) as dag:

    @task
    def generate_dependency():
        ctx = get_current_context()
        params = ctx["params"]
        logical_date = ctx["logical_date"]
        return [
            {
                "name": "test1",
                "bucket_key": f"s3:/{ params['INPUT_FOLDER'] }/{ logical_date.strftime('%Y/%m/%d') }",
            },
            {
                "name": "test2",
                "bucket_key": f"s3:/{ params['DATA_FOLDER'] }/{ logical_date.strftime('%Y/%m/%d') }",
            },
        ]

    @task(multiple_outputs=True)
    def map_dependency(data) -> Dict[str, str]:
        return {"bucket_key": data["bucket_key"]}

    @task
    def reduce_task(data):
        return list(data or [])

    @task
    def non_reduce_task(arg_1):
        return arg_1

    # bug flow
    list_dependency = generate_dependency.override(task_id="list_dependency")()
    dependencies = map_dependency.override(task_id="map_dependency").expand(data=list_dependency)

    list_task_sensor = S3KeySensor_test.partial(
        task_id="check_s3",
        wildcard_match=True,
        aws_conn_id="aws_default",
        timeout=300,
        poke_interval=120,
        mode="reschedule",
    ).expand_kwargs(dependencies)

    bug_task = non_reduce_task.override(task_id="bug_task", retries=3, retry_delay=timedelta(minutes=15))(
        arg_1=list_task_sensor.output
    )

    # work flow
    list_dependency_2 = generate_dependency.override(task_id="list_dependency_2")()
    dependencies_2 = map_dependency.override(task_id="map_dependency_2").expand(data=list_dependency_2)

    list_task_sensor_2 = S3KeySensor_test.partial(
        task_id="check_s3_2",
        wildcard_match=True,
        aws_conn_id="aws_default",
        timeout=300,
        poke_interval=120,
        mode="reschedule",
    ).expand_kwargs(dependencies_2)

    normal_task = non_reduce_task.override(
        task_id="normal_task", retries=3, retry_delay=timedelta(minutes=15)
    )(arg_1=reduce_task(list_task_sensor_2.output))

Operating System

REHL 7

Versions of Apache Airflow Providers

apache-airflow-providers-amazon==4.1.0

Deployment

Other

Deployment details

No response

Anything else

No response

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

[boring-cyborg[bot]]

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

[potiuk]

I believe you should make sure that your task returns a serializable value. Maybe @uranusjr have more to talk about it but IMHO returning JSON serializable value is a responsibility of the person who wants the operator/task to work in dynamic task context.

Converting it into discussion if needed.

[MaiHoangViet1809]

the output return from S3KeySensor with class PokeReturnValue already wrap xcom there, it should be in json style right ?
if I convert the code to return something like this, it just make thing complicate in my opinion:

if check_result:
    return PokeReturnValue(is_done=True, xcom_value={"return_value": ",".join(self.bucket_key) } )

I think PokeReturnValue should change to adapt to dynamic task context or S3KeySensor should allow us to return json directly.

[uranusjr]

The example is far too long as a sample, but I guess the relevant part is basically

non_reduce_task(arg_1=reduce_task(list_task_sensor_2.output))

where non_reduce_task is a taskflow @task, list_task_sensor_2 is a classic operator instance. I think it’s reasonable to expect this to work without an intermediate reduce_task and it should be doable by having some special handling for the lazy XCom proxy.

[MaiHoangViet1809]

Correct, that was handled by PokeReturnValue class of sensorbase class, as I dig down to source code, this class was partial document on git only, not exist in official website, PokeReturnValue should return json to adapt task flow api, or else, the airflow will not generate any error, silently make the task in bug flow skipped when we don't use a "reduce task" wrap for output of s3keySensor return via PokeReturnValue.

[uranusjr]

By the way, the easiest workaround would be to cast the return value explicitly:

@task
def non_reduce_task(arg_1):
    return list(arg_1)

This would allow the value to be serialised as expected without an intermediate task.

[uranusjr]

I filed two PRs #27250 and #27251 as two approaches to resolve the issue (by either documenting the list() solution, or trying to call list() auto-magically).