Dev-iL commented on code in PR #73966:
URL: https://github.com/apache/airflow/pull/73966#discussion_r4157153992
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -1336,10 +1336,10 @@ def get_task_instance_states(
if map_index is not None:
query = query.where(TI.map_index == map_index)
- results = session.scalars(query).all()
+ results = (await session.scalars(query)).all()
Review Comment:
Thanks! I confirmed the fix: the new test fails with the query change
reverted (it lists the `task1` and `group1.add_one` entities) and passes
otherwise, and `TestGetTaskStates` + `TestGetCount` pass on SQLite, PostgreSQL
and MySQL.
One more step before this is ready, though. With 5,001 states the route
still blocks the event loop for longer than the old sync handler did (median
~30–40 ms vs ~10–14 ms). Profiling the loop showed it's no longer the query:
most of it is building the response row by row (`Row` attribute access plus a
one-item `dict.update` per row), plus ORM result processing in one piece.
Unpacking tuples and streaming the result in partitions brings it well below
the sync baseline, on SQLite and PostgreSQL:
```diff
- results = (await session.execute(query)).all()
+ def add_states(rows) -> None:
+ for run_id, task_id, ti_map_index, state in rows:
+ key = task_id if ti_map_index < 0 else
f"{task_id}_{ti_map_index}"
+ run_id_task_state_map[run_id][key] = state
- if task_group_id:
- group_tasks = await _get_group_tasks(
- dag_id, task_group_id, session, dag_bag, logical_dates,
run_ids, map_index
- )
-
- results = results + group_tasks if task_ids else group_tasks
+ # Without task_ids, task_group_id replaces the dag-wide match instead
of extending it.
+ if task_ids or not task_group_id:
+ streamed = await
session.stream(query.execution_options(yield_per=500))
+ async for partition in streamed.tuples().partitions():
+ add_states(partition)
- [
- run_id_task_state_map[task.run_id].update(
- {task.task_id: task.state}
- if task.map_index < 0
- else {f"{task.task_id}_{task.map_index}": task.state}
+ if task_group_id:
+ add_states(
+ await _get_group_tasks(dag_id, task_group_id, session, dag_bag,
logical_dates, run_ids, map_index)
)
- for task in results
- ]
```
It also skips the dag-wide query when `task_group_id` is given without
`task_ids`, since that result was discarded anyway. Your tests pass unchanged
on all three backends with this applied.
<details><summary>Measurements</summary>
Warm requests to `/execution/task-instances/states`, one run, 10 requests
per size with the first dropped, ticker on the app's event loop:
| Backend | States | | Sync (base) | This PR | With the change above |
|---|---:|---|---:|---:|---:|
| SQLite | 1,001 | median / max stall | 6.3–6.7 / 313–337 ms | 10–12 / 13–16
ms | 2.7 / 3.3–4.0 ms |
| SQLite | 5,001 | median / max stall | 8.7–10.4 / 305–359 ms | 27–46 /
46–55 ms | 3.3–6.2 / 7.2–8.2 ms |
| PostgreSQL 14 (psycopg) | 5,001 | median / max stall | 13.8 / 424 ms |
39.3 / 46.6 ms | 5.6 / 7.8 ms |
| PostgreSQL 14 (psycopg) | 5,001 | request time | 632 ms | 54 ms | 52 ms |
Small local samples, so treat them as event-loop availability rather than
throughput. I didn't measure MySQL timings or test PgBouncer; streaming opens a
server-side cursor inside the request's transaction.
</details>
--
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]
