Skip to content

Add async task store accessors for async tasks#68232

Open
guan404ming wants to merge 3 commits into
apache:mainfrom
guan404ming:async-task-store-accessors
Open

Add async task store accessors for async tasks#68232
guan404ming wants to merge 3 commits into
apache:mainfrom
guan404ming:async-task-store-accessors

Conversation

@guan404ming

@guan404ming guan404ming commented Jun 8, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Related Issue

closes #68214

Why

Async @task functions could not checkpoint to task state store because the existing get/set/delete/clear block the event loop.

How

  • Add aget/aset/adelete/aclear to TaskStateStoreAccessor, awaiting the supervisor socket round-trip (SUPERVISOR_COMMS.asend) instead of blocking the event loop.
  • Document the async accessors and fix the resumable-tasks note that said async tasks cannot crash-recover.
Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Claude Code with Claude Opus 4.8

  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

@guan404ming guan404ming force-pushed the async-task-store-accessors branch 2 times, most recently from d61e93b to 6b41fdd Compare June 8, 2026 18:19
@guan404ming guan404ming marked this pull request as ready for review June 8, 2026 23:10
@guan404ming guan404ming requested review from dabla and jason810496 June 8, 2026 23:10
dabla
dabla approved these changes Jun 9, 2026
View reviewed changes

@dabla dabla left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looking good to me

@dabla dabla mentioned this pull request Jun 9, 2026
Merged
@dabla

dabla commented Jun 10, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

It would be good if @ashb or @amoghrajesh could validate this PR.

@amoghrajesh

amoghrajesh commented Jun 10, 2026

Copy link
Copy Markdown
Contributor

I do not have b/w as of today, but I would wanna take a look at this one tomorrow.

@guan404ming guan404ming force-pushed the async-task-store-accessors branch 3 times, most recently from 7460b4d to a51d715 Compare June 11, 2026 03:29
@dabla dabla moved this from Backlog to In review in AIP-103: Task State Management Jun 13, 2026
jason810496
jason810496 reviewed Jun 15, 2026
View reviewed changes

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice! LGTM overall.

Comment thread airflow-core/docs/core-concepts/task-store.rst Outdated
Comment thread airflow-core/src/airflow/api_fastapi/execution_api/routes/task_store.py
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py
@guan404ming guan404ming force-pushed the async-task-store-accessors branch from 7fbc338 to 6251981 Compare June 15, 2026 04:47
jason810496
jason810496 approved these changes Jun 15, 2026
View reviewed changes

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks.

amoghrajesh
amoghrajesh requested changes Jun 15, 2026
View reviewed changes

@amoghrajesh amoghrajesh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The async boundary is at the task-supervisor socket. The supervisor is a separate OS process entirely and whatever it does (sync HTTP, sync DB) is irrelevant to the task event loop. The task just awaits the socket response via asyncio.to_thread, which parks in a thread pool without blocking the event loop.

The description says that "wiring the execution API routes to the async DB path via MetastoreStateStoreBackend.aset" is genuinely not what this PR does, and more importantly, it does not need to do that. The fix is complete at the socket layer.

You will have to edit that to avoid confusion and requesting changes to avoid accidental merge with missing adelete / aclear now.

Comment thread task-sdk/src/airflow/sdk/execution_time/context.py
@guan404ming

guan404ming commented Jun 15, 2026

Copy link
Copy Markdown
Member Author

Good catch, just added adelete and aclear so an async task can clean up without falling back to the blocking sync calls.
Also, I've updated the description which I forget to update after modification, thanks!

@guan404ming guan404ming requested a review from amoghrajesh June 15, 2026 06:40
@guan404ming guan404ming force-pushed the async-task-store-accessors branch from 71d2ffe to 47134bb Compare June 16, 2026 12:36
@guan404ming

guan404ming commented Jun 17, 2026

Copy link
Copy Markdown
Member Author

Hi @amoghrajesh gentle ping. PTAL, thanks!

@guan404ming guan404ming force-pushed the async-task-store-accessors branch from 47134bb to 6a8644b Compare June 22, 2026 06:43
@guan404ming

guan404ming commented Jun 23, 2026

Copy link
Copy Markdown
Member Author

Hi @amoghrajesh gentle ping. PTAL, thanks!

@amoghrajesh amoghrajesh mentioned this pull request Jun 24, 2026
Merged
ashb
ashb reviewed Jun 24, 2026
View reviewed changes
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py Outdated
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py Outdated
jason810496
jason810496 reviewed Jun 24, 2026
View reviewed changes

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would you mind to resolve the conflict when you have a moment? Thanks.

kaxil
kaxil reviewed Jun 24, 2026
View reviewed changes
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py Outdated
@guan404ming guan404ming force-pushed the async-task-store-accessors branch 3 times, most recently from 52901b0 to 2e3b501 Compare July 2, 2026 23:55
@guan404ming guan404ming force-pushed the async-task-store-accessors branch from 2e3b501 to 1864bb2 Compare July 4, 2026 08:31
@guan404ming

guan404ming commented Jul 5, 2026

Copy link
Copy Markdown
Member Author

Hi @amoghrajesh could you help take another look, thanks!

@amoghrajesh

amoghrajesh commented Jul 6, 2026

Copy link
Copy Markdown
Contributor

@guan404ming hello, I will take a look today

amoghrajesh
amoghrajesh reviewed Jul 6, 2026
View reviewed changes

@amoghrajesh amoghrajesh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks much better now, last set of comments, otherwise looks good.

Comment thread task-sdk/src/airflow/sdk/execution_time/context.py Outdated
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py
@guan404ming guan404ming force-pushed the async-task-store-accessors branch from 1864bb2 to 8c3f1dc Compare July 6, 2026 08:18
@guan404ming guan404ming requested a review from amoghrajesh July 6, 2026 08:19
amoghrajesh
amoghrajesh approved these changes Jul 6, 2026
View reviewed changes

@amoghrajesh amoghrajesh left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. @kaxil wanna take another look?

@amoghrajesh amoghrajesh requested a review from kaxil July 6, 2026 08:45
@guan404ming

guan404ming commented Jul 6, 2026

Copy link
Copy Markdown
Member Author

Thanks @amoghrajesh !

Lee-W
Lee-W approved these changes Jul 6, 2026
View reviewed changes
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py
Comment thread task-sdk/src/airflow/sdk/execution_time/context.py
@guan404ming guan404ming force-pushed the async-task-store-accessors branch from bd82974 to c9c382f Compare July 8, 2026 03:36
@guan404ming guan404ming force-pushed the async-task-store-accessors branch from c9c382f to 7e8440b Compare July 8, 2026 05:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ashb ashb ashb left review comments
@jason810496 jason810496 jason810496 approved these changes
@dabla dabla dabla approved these changes
@Lee-W Lee-W Lee-W approved these changes
@amoghrajesh amoghrajesh amoghrajesh approved these changes
@kaxil kaxil Awaiting requested review from kaxil kaxil is a code owner

Assignees

No one assigned

Labels

area:task-sdk kind:documentation

Projects

AIP-103: Task State Management
Status: In review

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Add async task store accessors (aget/aset) for use in async tasks

7 participants

@guan404ming @dabla @amoghrajesh @ashb @Lee-W @kaxil @jason810496