[AMORO-4198][optimizer] Support graceful shutdown for in-progress tasks#4197
Merged
Conversation
github-actions
Bot
added
module:ams-server
j1wonpark
changed the title
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #4197 +/- ##
============================================
+ Coverage 23.09% 30.24% +7.14%
- Complexity 2706 4390 +1684
============================================
Files 463 680 +217
Lines 42826 55337 +12511
Branches 6044 7102 +1058
============================================
+ Hits 9891 16735 +6844
- Misses 32076 37337 +5261
- Partials 859 1265 +406
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
On SIGTERM the optimizer flips its stopped flag and returns immediately, so in-flight task results are silently dropped (completeTask is gated by isStarted). On K8s this is compounded by `sh -c` swallowing SIGTERM, a 30s default grace period, and Hadoop's FileSystem cache cleanup racing JVM shutdown hooks. - Optimizer.stopOptimizing: join executors with a deadline, force interrupt only on timeout; keep toucher alive so AMS heartbeats continue while tasks drain. - OptimizerExecutor.completeTask: best-effort direct call after stop so the in-flight result still reaches AMS. - SparkOptimizer / StandaloneOptimizer: register on Hadoop ShutdownHookManager (priority above FS_CACHE / SparkContext) with an explicit per-hook timeout. - OptimizerConfig: new -st / --shutdown-timeout-ms (default 600s). - KubernetesOptimizerContainer: `sh -c 'exec <args>'` and an explicit terminationGracePeriodSeconds derived from -st + 30s buffer; user podTemplate values are respected. - optimizer.sh start-foreground: exec $CMDS so java gets PID 1. Signed-off-by: Jiwon Park <jpark92@outlook.kr>
j1wonpark
force-pushed
the
optimizer-graceful-shutdown
branch
from
May 31, 2026 00:50
3607baa to
2a6df8f
Compare
Contributor
Author
|
Gentle ping for review 🙏 @zhoujinsong @czy006 — this builds on the master-slave optimizer work you reviewed in #4174 / #3937, and you both merged |
czy006
reviewed
Jul 2, 2026
Contributor
There was a problem hiding this comment.
Thanks for your contributor cc @j1wonpark
- --shutdown-timeout-ms is added as a configurable parameter, but AMS does not append shutdown-timeout-ms from resource/container properties when generating optimizer startup arguments. Although K8s terminationGracePeriodSeconds parses this value from the startup arguments, the current AMS-managed K8s optimizer effectively can only use the default 600s + 30s buffer.
- Although the Flink optimizer inherits the shared stopOptimizing() / completeTask() changes, it does not register an equivalent drain hook for external termination. I suggest adding this part.
…rtup arg and K8s grace period The -st/--shutdown-timeout-ms option was parsed by the optimizer but AMS never emitted it into the startup args, so the configured shutdown timeout was unreachable for AMS-spawned optimizers and the K8s grace period was pinned at the default. - Emit '-st <value>' in buildOptimizerStartupArgsString when the optimizer group sets the shutdown-timeout-ms property (same pattern as -hb), which wires all containers sharing the args builder. - Derive terminationGracePeriodSeconds from the same resource property that drives the -st arg via PropertyUtil, instead of re-parsing the rendered CLI string, and remove parseShutdownTimeoutMs. Reading from one source guarantees the pod's SIGKILL deadline can never fall below the shutdown timeout the optimizer JVM actually honors. Signed-off-by: Jiwon Park <jpark92@outlook.kr>
…llation FlinkExecutor.close() previously interrupted the executor thread after a hardcoded 5s join, killing in-progress tasks on job cancellation and silently ignoring the -st/--shutdown-timeout-ms option on the Flink engine. - close() now stops the executor first, waits up to the configured shutdown timeout for the in-progress task to complete (so its result is reported to AMS), and only then force-interrupts. - The drain absorbs interrupts of the thread running close(): Flink's cancellation machinery (TaskCanceler/TaskInterrupter) interrupts the task thread to unblock I/O, which must not abort the drain early. - The drain budget is capped safely below Flink's cancellation watchdog (task.cancellation.timeout, default 180s), which would otherwise fail the whole TaskManager; operators wanting longer drains must raise both settings, as documented. - Document the shutdown-timeout-ms group property, the -st option for Flink/Spark external optimizers, and the interaction with task.cancellation.timeout. Signed-off-by: Jiwon Park <jpark92@outlook.kr>
… survive Fixes several defects in the graceful-shutdown paths that could still lose the result of a task that completed during the drain, or leave the process in a bad state: - Optimizer.startOptimizing published the executorThreads array before filling its elements; a concurrent shutdown hook could observe null elements and skip joining executor threads that were already running tasks. Build the array fully, then publish, then start the threads. - The force-interrupt pass in stopOptimizing interleaved interrupt with per-thread 1s joins: when the stopping thread was itself interrupted (e.g. the shutdown-hook watchdog cancelling a timed-out hook), join() threw instantly and executor threads after the first were never interrupted; the per-thread joins also overshot the deadline by up to N seconds. Interrupt all survivors first, then await them against one shared deadline, restoring the caller's interrupt status only after cleanup completes. - waitAShortTime restored the interrupt flag unconditionally, but no caller loop ever clears it: a stray interrupt while still running turned every subsequent sleep into an instant return, busy-spinning the poll/retry loops. Preserve the flag only once the operator is stopped. - completeTask checked isStarted() once at entry and fell back to a single best-effort call after stop, so a result finishing during shutdown was dropped on a stop/report race or one transient AMS error. It now retries through callAuthenticatedAmsWithDrain: a bounded window anchored at the moment shutdown is observed, reusing the existing transient-error and auth handling, aborted early by the shutdown force-interrupt. The mock AMS gains transient-failure injection to cover this. - The toucher stays alive during the drain for heartbeats, but scale-down unregisters the optimizer before the pod receives SIGTERM; the resulting auth error made the toucher re-register a ghost optimizer and rotate the executors' token so their final completeTask was rejected. Drain mode now skips re-registration once the token becomes invalid. - The graceful-shutdown test timed stop() with fixed sleeps against the poll cadence (~0.5s margins on CI); slow test tasks now signal a latch when execute() begins so the test stops the optimizer deterministically. Signed-off-by: Jiwon Park <jpark92@outlook.kr>
j1wonpark
force-pushed
the
optimizer-graceful-shutdown
branch
from
July 4, 2026 06:56
4551aff to
b4e5ae3
Compare
…link drain Flink cancels the FlinkToucher source before downstream operators, so heartbeats stop the moment the drain in FlinkExecutor.close() begins. AMS then expires the optimizer after optimizer.heart-beat-timeout (default 1 min) and, in the same sweep, unregisters it and resets its in-flight tasks — so any task finishing later than that lost its result even though the drain was still waiting, capping the effective drain window at the heartbeat timeout instead of the configured shutdown timeout. The drain loop now sends a best-effort touch with the existing token between join slices, at the configured heartbeat interval. Touching keeps the registration and the task assignment alive for exactly as long as the drain runs; it never re-registers (consistent with the drain-mode toucher), and once close() returns the touches stop and AMS cleans up through the normal expiration path. Signed-off-by: Jiwon Park <jpark92@outlook.kr>
Contributor
Author
|
Thanks for the review @czy006! Both points are addressed: 1. Fixed in 99367db — the optimizer group property 2. Added in 21d75de — I also fixed several races I found in the new shutdown paths (b4e5ae3). |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why are the changes needed?
Close #4198.
When an optimizer receives SIGTERM, in-progress tasks are silently dropped and
AMS re-schedules them — doubling work and potentially causing duplicate commits.
Brief change log
Optimizer.stopOptimizing(): join executor threads up to--shutdown-timeout-ms(default 10 min); keep toucher alive during drain so AMS heartbeats continue
OptimizerExecutor.completeTask(): best-effort direct call after shutdown soresults are not silently dropped
OptimizerToucher.stop(): interrupt runner thread to wake it from sleep immediatelyAbstractOptimizerOperator.waitAShortTime(): preserve interrupt flagOptimizerConfig: new-st/--shutdown-timeout-msoptionStandaloneOptimizer/SparkOptimizer: register graceful shutdown hook onHadoop's
ShutdownHookManageraboveFS_CACHEpriority with explicit timeoutKubernetesOptimizerContainer:execprefix in container command; deriveterminationGracePeriodSecondsfrom shutdown-timeout-ms + 30s bufferoptimizer.sh start-foreground:exec $CMDSso Java receives SIGTERM directlyHow was this patch tested?
Add some test cases that check the changes thoroughly including negative and positive cases if possible
Add screenshots for manual tests if appropriate
Run test locally before making a pull request
Documentation
usagestringReview updates
shutdown-timeout-msoptimizer group property to the-ststartup arg and the K8s grace period (single typed source; the CLI re-parser is removed), and documented it. (@czy006 review point 1)FlinkExecutor.close()waits up to the shutdown timeout, capped below Flink'stask.cancellation.timeout. (@czy006 review point 2)completeTaskretrying through a bounded drain window, and no toucher re-registration while draining.Known limitations (proposed follow-ups)
spark-submit --killdoes not extend the driver pod's termination grace period (K8s default 30s applies); operators can work around it viaspark-conf.spark.kubernetes.appKillPodDeletionGracePeriod.LocalOptimizerContainer.releaseResource()useskill -9, which skips JVM shutdown hooks; graceful drain only applies to externally sent SIGTERM.