Dev-iL opened a new pull request, #75381:
URL: https://github.com/apache/airflow/pull/75381
related: #67799
This PR converts the four Execution API task state store routes to `async
def` with `AsyncSessionDep`:
* `GET /store/ti/{task_instance_id}/{key}` (`get_task_state_store`)
* `PUT /store/ti/{task_instance_id}/{key}` (`set_task_state_store`)
* `DELETE /store/ti/{task_instance_id}/{key}` (`delete_task_state_store`)
* `DELETE /store/ti/{task_instance_id}` (`clear_task_state_store`)
Details:
* The routes call the configured backend's `aget` / `aset` / `adelete` /
`aclear` with the route's session. These methods are abstract on
`BaseStoreBackend`, so every backend already implements them. The state store
docs already tell custom backends with a synchronous client to offload their
async methods to a thread, and the `common.io` object-storage backend does so.
* The task instance lookup selects the five columns it uses and filters on
`working_set IS TRUE`. Before, `session.get(TI, ...)` loaded the full task
instance and its joined Dag run.
* `MetastoreBackend._aget_task_state_store` selects only `value` instead of
the full row. The column is `NOT NULL`, so a missing key still returns `None`.
* GET decodes the stored value and builds the response model in a worker
thread (`run_in_threadpool`). Without this, a large value blocks the event loop
longer than the synchronous route did (measurements below).
The HTTP contract does not change, so no Execution API version change is
needed.
<details><summary><h2>Benchmarks and Testing</h2></summary>
<p>
### Event-loop stall measurements
The harness is adapted from the review of #74540. Each request runs over
HTTP through `TestClient` from a separate thread. A 1 ms ticker on the server
loop records its longest sleep per request. Every case asserts the status and
the returned or stored value. 10 warm requests per case, cold request dropped,
2 rounds per version run interleaved in one container. GET/PUT values are
integer arrays. Each cell shows `median / max` stall in ms for round 1; round 2.
* **base**: the synchronous routes on `main`.
* **head**: async routes without any offload.
* **offload**: head plus offloading GET decoding/model building and PUT
`json.dumps`. **This PR ships head for PUT/DELETE/CLEAR and offload for GET.**
Python 3.14.8, SQLAlchemy 2.0.54, Breeze CI image. These are small local
samples of event-loop availability, not throughput.
**SQLite (aiosqlite)**
| Request | Value | base | head | offload |
|---|---:|---|---|---|
| GET | 0.4 KB | 2 / 2; 2 / 3 | 3 / 4; 3 / 4 | 4 / 5; 3 / 4 |
| GET | 58.9 KB | 14 / 17; 15 / 24 | 16 / 31; 14 / 31 | 16 / 28; 15 / 23 |
| GET | 688.9 KB | 106 / 144; 119 / 200 | 144 / 217; 144 / 158 | 123 / 145;
102 / 139 |
| PUT | 0.4 KB | 2 / 5; 3 / 6 | 3 / 6; 3 / 4 | 3 / 4; 3 / 5 |
| PUT | 58.9 KB | 7 / 17; 9 / 17 | 11 / 15; 12 / 25 | 12 / 14; 12 / 25 |
| PUT | 688.9 KB | 60 / 86; 81 / 159 | 89 / 109; 93 / 102 | 103 / 138; 100 /
119 |
| DELETE | - | 2 / 3; 4 / 5 | 3 / 4; 4 / 4 | 4 / 8; 4 / 4 |
| CLEAR, 1000 keys | - | 2 / 2; 2 / 3 | 3 / 4; 4 / 4 | 4 / 4; 4 / 4 |
| CLEAR, 5000 keys | - | 2 / 3; 2 / 4 | 3 / 4; 3 / 3 | 3 / 4; 3 / 4 |
**PostgreSQL (psycopg async)**
| Request | Value | base | head | offload |
|---|---:|---|---|---|
| GET | 0.4 KB | 3 / 5; 3 / 4 | 5 / 5; 4 / 5 | 4 / 5; 4 / 5 |
| GET | 58.9 KB | 19 / 27; 17 / 24 | 24 / 31; 21 / 27 | 19 / 27; 17 / 27 |
| GET | 688.9 KB | 143 / 153; 147 / 155 | 206 / 256; 199 / 232 | 152 / 171;
127 / 153 |
| PUT | 0.4 KB | 3 / 5; 3 / 4 | 4 / 4; 4 / 5 | 4 / 6; 4 / 6 |
| PUT | 58.9 KB | 10 / 18; 12 / 16 | 18 / 25; 15 / 19 | 19 / 24; 18 / 20 |
| PUT | 688.9 KB | 90 / 104; 86 / 95 | 133 / 161; 121 / 170 | 128 / 165; 125
/ 138 |
| DELETE | - | 3 / 4; 3 / 4 | 3 / 4; 4 / 6 | 4 / 4; 4 / 5 |
| CLEAR, 1000 keys | - | 3 / 4; 3 / 4 | 4 / 5; 4 / 5 | 4 / 5; 3 / 5 |
| CLEAR, 5000 keys | - | 3 / 5; 3 / 4 | 3 / 4; 4 / 6 | 4 / 5; 3 / 6 |
Findings:
* **GET**: without offload, the async route's median stall for a 689 KB
value is 21–44% above base, depending on backend and round. With decoding and
model building in a worker thread, GET is back at the base level on both
backends.
* **PUT is above base, and I found no remedy inside the route.** For a 689
KB value, the median stall is 15–78% above base, depending on backend and
round. For a 59 KB value, it is 2–9 ms longer. Offloading `json.dumps` does not
help (about 7 ms of the work). The loop-thread profiles show the same CPU on
the loop in all three versions: three pydantic validations of the `JsonValue`
body plus one cadwyn `model_dump`. The difference is where the loop gets a
break. My reading of the code, which I have not instrumented directly: cadwyn's
request migration re-runs FastAPI's `solve_dependencies`. With the synchronous
`SessionDep`, that resolves the generator dependency through
`contextmanager_in_threadpool`, which yields the loop in the middle of the
framework work. `AsyncSessionDep` gives no such yield. A probe with `await
asyncio.sleep(0)` at handler entry recovered only part of the gap (PostgreSQL,
PUT only):
| Value | base | head | head + yield |
|---:|---|---|---|
| 58.9 KB | 11 / 12; 10 / 15 | 18 / 26; 19 / 25 | 16 / 22; 16 / 26 |
| 688.9 KB | 79 / 91; 101 / 120 | 141 / 172; 145 / 183 | 124 / 145; 134 /
165 |
This cost belongs to framework request handling for large bodies and
likely applies to any async Execution API route with a large `JsonValue` body.
I left PUT without offload or yield, and I am raising it here for reviewers to
decide.
* **Small requests**: every route keeps the loop about 1 ms longer than
base, because async SQLAlchemy work now runs on the loop instead of in a worker
thread. Request duration is equal or lower.
I did not measure stalls on MySQL. The contract tests ran there.
## Tests
* The existing HTTP contract tests in
`versions/head/test_task_state_store.py` pass unchanged.
* `test_routes_select_columns_without_loading_entities` (all four routes)
asserts that no ORM entity is loaded during the request. It fails on `main`,
where `TaskInstance` and `DagRun` are loaded. Two negative controls confirm
that it detects loads through `AsyncSession`: restoring the full-row value
getter fails only the GET case, and restoring `session.get(TI, ...)` fails with
`TaskInstance` and `DagRun`.
* `test_get_builds_response_off_the_event_loop` asserts that the response
model is built without a running event loop. It fails when the offload is
removed from the async handler. It also passes on `main`, because the
synchronous handler already ran in a worker thread. It guards the offload.
Checks run:
* `pytest` on `test_task_state_store.py` and
`tests/unit/state/test_metastore.py`, in Breeze with Python 3.14.8: 107 passed
on each backend:
* SQLite (`sqlite` / `sqlite+aiosqlite`)
* PostgreSQL (`postgresql+psycopg` / `postgresql+psycopg_async`)
* MySQL (`mysql` / `mysql+aiomysql`)
* `prek run mypy-airflow-core`: no issues.
* `prek run --stage pre-commit` on the changed files: passed.
</p>
</details>
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR,
please
change below checkbox to `[X]` followed by the name of the tool, uncomment
the "Generated-by".
-->
- [x] Yes (please specify the tool below)
Generated-by: Claude Code (Opus 5.5) following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
---
* Read the **[Pull Request
Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#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](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
