Skip to content

Add async task store accessors for async tasks - #68232

Merged
guan404ming merged 3 commits into
apache:mainfrom
guan404ming:async-task-store-accessors
Jul 9, 2026
Merged

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

Conversation

@guan404ming

@guan404mingguan404ming 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
guan404mingforce-pushed the async-task-store-accessors branch 2 times, most recently from d61e93b to 6b41fddCompareJune 8, 2026 18:19
@guan404ming
guan404ming marked this pull request as ready for review June 8, 2026 23:10
dabla
dabla approved these changes Jun 9, 2026
View reviewed changes

@dabladabla 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

@dabladabla 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
guan404mingforce-pushed the async-task-store-accessors branch 3 times, most recently from 7460b4d to a51d715CompareJune 11, 2026 03:29
@dabladabla moved this from Backlog to In review in AIP-103: Task State ManagementJun 13, 2026
jason810496
jason810496 reviewed Jun 15, 2026
View reviewed changes

@jason810496jason810496 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 threadairflow-core/docs/core-concepts/task-store.rst Outdated
Comment threadairflow-core/src/airflow/api_fastapi/execution_api/routes/task_store.py
Comment threadtask-sdk/src/airflow/sdk/execution_time/context.py
@guan404ming
guan404mingforce-pushed the async-task-store-accessors branch from 7fbc338 to 6251981CompareJune 15, 2026 04:47
jason810496
jason810496 approved these changes Jun 15, 2026
View reviewed changes

@jason810496jason810496 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

@amoghrajeshamoghrajesh 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 threadtask-sdk/src/airflow/sdk/execution_time/context.py
@guan404ming

guan404ming commented Jun 15, 2026

Copy link
Copy Markdown
MemberAuthor

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
guan404mingforce-pushed the async-task-store-accessors branch from 71d2ffe to 47134bbCompareJune 16, 2026 12:36
@guan404ming

guan404ming commented Jun 17, 2026

Copy link
Copy Markdown
MemberAuthor

Hi @amoghrajesh gentle ping. PTAL, thanks!

@guan404ming
guan404mingforce-pushed the async-task-store-accessors branch from 47134bb to 6a8644bCompareJune 22, 2026 06:43
@guan404ming

guan404ming commented Jun 23, 2026

Copy link
Copy Markdown
MemberAuthor

Hi @amoghrajesh gentle ping. PTAL, thanks!

@guan404ming
guan404mingforce-pushed the async-task-store-accessors branch 3 times, most recently from 52901b0 to 2e3b501CompareJuly 2, 2026 23:55
@guan404ming
guan404mingforce-pushed the async-task-store-accessors branch from 2e3b501 to 1864bb2CompareJuly 4, 2026 08:31
@guan404ming

guan404ming commented Jul 5, 2026

Copy link
Copy Markdown
MemberAuthor

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

@amoghrajeshamoghrajesh 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 threadtask-sdk/src/airflow/sdk/execution_time/context.py Outdated
Comment threadtask-sdk/src/airflow/sdk/execution_time/context.py
@guan404ming
guan404mingforce-pushed the async-task-store-accessors branch from 1864bb2 to 8c3f1dcCompareJuly 6, 2026 08:18
amoghrajesh
amoghrajesh approved these changes Jul 6, 2026
View reviewed changes

@amoghrajeshamoghrajesh 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 kaxilJuly 6, 2026 08:45
@guan404ming

guan404ming commented Jul 6, 2026

Copy link
Copy Markdown
MemberAuthor

Thanks @amoghrajesh !

Lee-W
Lee-W approved these changes Jul 6, 2026
View reviewed changes
Comment threadtask-sdk/src/airflow/sdk/execution_time/context.py
Comment threadtask-sdk/src/airflow/sdk/execution_time/context.py
@guan404ming
guan404mingforce-pushed the async-task-store-accessors branch 2 times, most recently from c9c382f to 7e8440bCompareJuly 8, 2026 05:57
@guan404ming
Add async task store accessors for async tasks
cb84329
@guan404ming
Add adelete and aclear
847f36c
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
@guan404ming
Handle review comments
4d10476
Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>
@guan404ming
guan404mingforce-pushed the async-task-store-accessors branch from 7e8440b to 4d10476CompareJuly 9, 2026 04:38
@amoghrajesh

amoghrajesh commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

@guan404ming seems good enough for now, we can merge once CI is green.

@guan404ming

guan404ming commented Jul 9, 2026

Copy link
Copy Markdown
MemberAuthor

Sure, ci looks all green. I am going to merge this, thanks for all reviews!

@guan404ming
guan404ming merged commit e4ca33e into apache:mainJul 9, 2026
261 of 264 checks passed
@guan404ming
guan404ming deleted the async-task-store-accessors branch July 9, 2026 06:45
@dabla

dabla commented Jul 11, 2026

Copy link
Copy Markdown
Contributor

Sure, ci looks all green. I am going to merge this, thanks for all reviews!

Nice one, we finally got it merged :-)

@desusaitejadesusaiteja mentioned this pull request Aug 10, 2026
Open
1 task
This was referenced Aug 27, 2026
Open
Open
Open
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ashbashbashb left review comments
@jason810496jason810496jason810496 approved these changes
@dabladabladabla approved these changes
@Lee-WLee-WLee-W approved these changes
@amoghrajeshamoghrajeshamoghrajesh approved these changes
@kaxilkaxilAwaiting requested review from kaxilkaxil is a code owner

Assignees

No one assigned

Labels

area:task-sdkkind:documentation

Projects

AIP-103: Task State Management
Status: Done

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