dabla commented on PR #62922: URL: https://github.com/apache/airflow/pull/62922#issuecomment-6337721568
@kaxil a heads-up on a change since your last review: [e5abeccb1e](https://github.com/apache/airflow/pull/62922/commits/e5abeccb1e0112a70c6e00b18585c499033e9a87) changes how the indexed tasks' checkpoints are written. It doesn't touch the points of earlier rounds; this comment explains why and what it trades off. ### Why Every indexed task wrote its checkpoint under its own task state store key (`_iterable_<index>`), and a retry read them back one by one. That is one request per index each way, each through the supervisor to the API server, on top of the XCom push. The task process sends its requests one at a time, so on large inputs the checkpoints alone took half of the channel. A real run over 16,194 items made 16,194 checkpoint writes on its first attempt and 16,194 reads before its retry could start running anything. ### What changed - **Writes per round:** an indexed task's outcome is kept in memory and written with its round: `task_concurrency` outcomes as one `_iterable_batch_<n>` key. Rounds follow completion order, so a slow indexed task holds back no outcome but its own. - **Reads once:** a retry reads the batches once when it starts, in order until the first missing one, the latest entry of an index winning. Replaying an index reads nothing more. - **The rest, then the marker:** what's left when the event loop has closed is written with the sync store call, then the completion marker, so the marker still comes after everything it covers. - **Smaller classes:** the logic sits in small classes with one job each (`CheckpointBatch`, `PendingCheckpoints`, `EarlierCheckpoints`, `CheckpointBatches`) behind `Checkpoints.read` / `record`. The lifecycle no longer writes to the store itself. An `IndexedTaskOutcome` replaces the `(task, result, raised)` tuple, and its `fail_fast()` replaces `_fail_fast_for`. ### Trade-offs, all on purpose - **Reported before durable.** An indexed task is reported, success callback included, when it finishes, not when its round is written. After a hard kill (SIGKILL, OOM), the indexed tasks of the last round run again on the retry and their success callbacks fire again. That is what happens to an operator's own work after a kill during `execute`. Waiting for the round would hold every callback back by up to a round. - **A failed write is not an indexed task's outcome.** It's kept and written with the next round; a final write that fails is logged, and those indexed tasks run again on a retry. Before, a failed checkpoint write failed its indexed task. - **A retry that can't read the store fails before running anything.** Running an index whose checkpoint may say "done" would repeat work. - **The per-index keys are gone** (`IndexedTaskInstance.aget_state` / `aset_state`), with no fallback: nothing with the old format was released. ### Alternatives considered - **Overwrite one key per attempt** with every outcome so far: each write rewrites everything before it, so the bytes grow quadratically with the input, and the value gets as big as the whole result. - **A background writer thread with a queue:** it sends the same requests over the same single channel, so it gains nothing. It would also bring back what the earlier rounds removed: a sync SDK call from another thread while the loop waits, a writer to stop and drain on kill and timeout, and the marker-after-last-write order turning into a queue drain. ### Measured | | one key per index | per round | |---|---|---| | store writes, first attempt (N items, 16 workers) | N | ⌈N / 16⌉ | | store reads, retry | N + 1 | ⌈N / 16⌉ + 2 | | wall time, first attempt / retry (same 2 ms per supervisor request) | 1× | **1.56× / 1.8× faster** at 1,000, 5,000 and 16,194 items | | real run, 16,194 items, first attempt | 2 h 14 | 27 min | ### Docs `mapped-tasks-vs-iterable-tasks.rst` says how the checkpoints are written. It now also advises object storage behind the task state store for large results, as for XComs, with the `StateStoreObjectStorageBackend` of `common.io`. That advice holds with or without this change: - a database alone fills the table as the XCom table would without an XCom backend; - workers warn above `[state_store] max_value_storage_bytes` (64 KB by default); - MySQL caps a row at 16 MiB. It also notes that this backend doesn't apply `expires_at`, so the bucket's lifecycle rules decide how long the checkpoints stay. ### Tests - **New:** `TestCheckpointBatches`. It covers: - one write per round, each outcome once; - a round written once under the lock; - failed writes kept; - a `BaseException` during a write; - the latest entry winning and reading stopping at a gap; - the rest written before the marker; - a retry through `execute` that runs only the failed item. - **Rewritten for the rules above:** six tests (the read failure, the success callback, failed skip / up-for-retry writes, the write count, the kill window). The kill now lands while the input is read and rendered, since the checkpoint read no longer awaits. - **End to end:** `airflow dags test` with a retry. Only the failed index ran again, each success callback fired once, and the downstream task read the replayed results. -- 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]
