Fix CloudwatchTaskHandler display error

[jason810496]

related: #49470

Why

As noted in #49470 (comment), PR #49470 unintentionally broke the display format of CloudWatchHandler.

What

This PR fixes the display issue for CloudWatchHandler and ensure it will display properly across different airflow-core versions.

[jason810496]

Just tested fixed CloudwatchTaskHandler with different airflow-core version setups and all of them display correctly.
I will keep on fixing the CI.

I setup the following airflow-core version matrix against fixed CloudwatchTaskHandler.

• 3.1.0 (unreleased)
  • On main branch then run breeze start-airflow --python 3.10 --backend postgres --integration localstack --mount-sources providers-and-tests
• 3.0.4 (unreleased)
  1. git checkout v3-0-test && rm dist/*
  2. breeze release-management prepare-airflow-distributions && breeze release-management prepare-task-sdk-distributions
  3. breeze start-airflow --python 3.10 --backend postgres --integration localstack --mount-sources providers-and-tests --use-airflow-version wheel
• 3.0.3
  • On main branch then run breeze start-airflow --python 3.10 --backend postgres --integration localstack --mount-sources providers-and-tests --use-airflow-version 3.0.3

After fix, the log should be format correctly on frontend, screenshot as below

• 3.1.0 (unreleased)
  • 
• 3.0.4 (unreleased)
  • 
• 3.0.3
  • 

cc @ashb @eladkal

[ashb]

This still goes and loads everything into memory for Cloudwatch doesn't it? As we call io.read() which does [x for x in self.stream()] -- I wonder if the read code in file_task_handler should check for a io.stream function and use that in preference over io.read?

That might be a bit complex, but either way I think we need a bit more thought here to not undo all the changes you had for paging/oom of logs

[jason810496]

This still goes and loads everything into memory for Cloudwatch doesn't it?

Yes, I'm still WIP on local for setup stream-based read with backward compatibility.
Since #49470 was merged, so after 3.0.3 ( get_base_airflow_version_tuple() >= (3, 0, 3) ), the airflow-core get accept log stream.

As we call io.read() which does [x for x in self.stream()] -- I wonder if the read code in file_task_handler should check for a io.stream function and use that in preference over io.read?

I had considered compact mechanism in #49470 by peeking the type of the _read_remote_logs.

https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/utils/log/file_task_handler.py#L611-L625

So we just need to make sure the provider with airflow-core after 3.0.3 version should return stream-based.
Somehow like:

# in `version_compat` module
# or maybe just use `get_base_airflow_version_tuple() >= (3, 0, 3)` directly in provider FileTaskHandler
STREAM_BASED_READ = get_base_airflow_version_tuple() >= (3, 0, 3)

# in provider FileTaskHandler
if STREAM_BASED_READ:
    return messages, log_streams # generator of StructuredLogMessage or str
else:
    return messages, log_list # list of str ( not memory efficient, but only after 3.0.3 airflow-core can support stream-based read )

[jason810496]

Just tested the latest provider with the following version matrix and also check the memory usage.

• Version matrix for Compatibility
  • 3.0.3
  • 3.0.4 -> start supporting stream-based read
  • main ( 5d302eb61f5878282c57105b12be0a8d1de936df version )
• Memory usage comparison
  • 3.0.3
  • 3.0.4 -> should not have peek memory usage when reading task instance log

Compatibility

3.0.3

Setup:

breeze start-airflow --python 3.10 --backend postgres --integration localstack --mount-sources providers-and-tests --use-airflow-version 3.0.3

Result:

3.0.4

Setup:

breeze start-airflow --python 3.10 --backend postgres --integration localstack --mount-sources providers-and-tests --use-airflow-version 3.0.4

Result:

3.1.x ( main branch )

Setup:

1. Switch to main
2. Run rm -rf dist/*
3. Run breeze release-management prepare-airflow-distributions
4. Run breeze release-management prepare-task-sdk-distributions
5. Switch back to branch of this PR
6. Run breeze start-airflow --use-airflow-version wheel --mount-sources providers-and-tests --integration localstack

Result:

Memory usage comparison

Setup:

1. Run start-airflow and trigger DagRun, then terminate all processes instead of API Server
2. Call get task instance log with json header
3. Call get task instance log with nd-json header

There is still memory peek for nd-json header.

No memory peek for nd-json header.

[jason810496]

The PR should be good to go, I just test it with both Airflow 3 and Airflow 2 with the following setup and both work well.

breeze start-airflow --backend postgres --integration localstack --db-reset --mount-sources providers-and-tests --use-airflow-version 3.1.3

breeze start-airflow --backend postgres --integration localstack --mount-sources providers-and-tests --use-airflow-version 2.11.0

[Lee-W]

looks like we have some compatibility issue 🤔