wenjin272 opened a new pull request, #1354: URL: https://github.com/apache/flink-agents/pull/1354
Closes #1353 ### Purpose of change On the JDK < 21 Java parallel engine, checkpoints can now capture actions waiting on external calls without draining the worker pool. Previously, a slow call blocked the pre-barrier drain, and an abort observed by that drain could fail the task even when checkpoint failures were tolerated. #### Runtime flow A worker records each pulled root ActionTask in keyed `inFlightActionTasks` state before executing it. The record remains until mailbox commit has persisted the Action's memory and processed its outputs. Checkpointing holds the execution lock for the complete synchronous backend capture; external calls that have handed off that lock can remain unfinished. Restore sorts the saved tasks by dispatch index, prepends them to the existing action queue, clears the in-flight map, and then resumes scheduling. #### Key decisions Use a separate state map so running workers are not re-dispatched from the live queue. Restore root tasks and regenerate child graphs through root replay. Keep snapshot capture and commit mutually exclusive instead of waiting for all external work to finish. Remove checkpoint-specific dispatch pausing and its abort-to-task-failure path. ### Behavioral Semantics #### Interaction decisions | Condition | Behavior | | --- | --- | | Low-JDK parallel worker is waiting in an async call | The unfinished root remains checkpointable; no pre-barrier drain. | | A later worker finished but cannot commit past the FIFO head | Its root remains in state and is replayed in dispatch order after recovery. | | A root is waiting on a child graph | Save the root; recovery regenerates the child graph. | | JDK 21 or Python execution | Existing scheduling remains; restore also recognizes the new in-flight map. | | Synchronous Action code holds the execution lock, or input admission is blocked | This change does not guarantee checkpoint progress through that blockage. | #### Behavioral contracts 1. Checkpoints can complete while async calls remain blocked and input admission is below its cap. 2. A tolerated checkpoint abort does not itself fail the Action or pause subsequent dispatch. 3. Recovery preserves committed memory and unfinished/queued inputs; actions committed before the snapshot are not re-executed by the new recovery mechanism. 4. Restored pulled roots precede unstarted tasks in their original per-key dispatch order. 5. Root/child recovery continues without independently dispatching the saved child graph twice. #### Failure behavior State write, snapshot, and Action execution errors still propagate through the existing Flink failure paths. Checkpoint storage failure/expiry follows Flink's configured checkpoint failure policy. Unfinished external calls can run again after recovery; this does not add exactly-once guarantees for external side effects. Worker exhaustion can still stall graph execution even when checkpoints succeed; the configuration documentation now states that distinction. ### Tests | Contract | Coverage | | --- | --- | | 1: CP progress during async waits | `CheckpointE2ETest.checkpointsProgressWhileIdleAndAcrossBusyKeys`; Python `checkpoint_test.py`; saturated-worker Heap/RocksDB test | | 2: Abort and continued progress | `checkpointFailureDoesNotStopActionsOrLaterCheckpoints` (async storage failure and real timeout); `checkpointAbortDoesNotFailOrPauseInFlightAction` | | 3: Recovery and commit boundary | `restoresCommittedMemoryAndPendingInputs`; `checkpointRestoresOnlyUncommittedActions`; existing memory-recovery tests | | 4: Ordered recovery | `restoresInFlightTasksInDispatchOrderBeforeQueuedTasks`; `ActionExecutionOperatorInFlightCheckpointTest` with pending and completed-but-uncommitted workers | | 5: Child graphs | `SubgraphParallelExecutionTest.restoresRootWhileChildModelCallIsInFlight` (chat and nested subagent); parallel failover regressions | Validation: - Java MiniCluster E2E: 4 passed each on JDK 11 / Flink 1.20.5, JDK 17 / Flink 2.3.0, and JDK 21 / Flink 2.3.0. - Python / Flink 2.3.0 checkpoint E2E: 1 passed each on JDK 17 and 21. - Targeted runtime regressions: JDK 17, 127 passed / 1 skipped; JDK 21, 74 passed / 14 skipped. - Spotless, Ruff, and whitespace checks. Not verified here: live Kafka/Fluss ActionStateStore outages, real network backpressure or input-admission saturation, and transfer of a real checkpoint between separate JDK installations. Input-limit parity fixes are deferred to a separate change. E2E recovery uses Flink state without an external ActionStateStore. <details> <summary>State invariants and validation setup</summary> - Root tasks are added under the execution lock after dequeue, then written again after the started notification so serialized backends preserve that flag. - The in-flight entry is removed at the end of `Work.commit()`. The outer synchronous `snapshotState(...)` holds the same non-reentrant lock; its inner callback avoids acquiring it twice. - Recovery scans keys that have only in-flight tasks as well as keys with queued work. Sorting does not depend on MapState iteration order; clearing the map prevents a later restore from merging tasks twice. - Focused state tests cover Heap and RocksDB, restored input-budget accounting, and existing migration behavior. Obsolete drain-specific tests are replaced by checkpoint/recovery assertions. - Distribution JARs were rebuilt through a targeted Maven reactor `clean install`. Python tests used a rebuilt/reinstalled local wheel containing those JARs. Java E2E used `-Dtest=CheckpointE2ETest` with the matching Flink profile. - The JDK 21 E2E Surefire profile also opens `java.util` and `java.nio` for Flink test serialization, while retaining the Continuation API export. </details> ### API No public API or configuration changes. Adds keyed `inFlightActionTasks` state using the existing ActionTask serializer; older snapshots have no entries in this new state. Checkpoints containing non-empty in-flight state must be restored with a version that understands it; downgrade to the old drain-only implementation is not supported by this change. ### Documentation - [ ] `doc-needed` - [ ] `doc-not-needed` - [x] `doc-included` ### Was this patch authored or co-authored using generative AI tooling? - [x] Yes - [ ] No Generated-by: Codex CLI 0.153.4 (GPT-6) -- 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]
