fix(framework): Prevent task executor shutdown deadlock - #7895
fix(framework): Prevent task executor shutdown deadlock#7895tanertopal wants to merge 17 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR improves shutdown reliability by ensuring background authenticated calls (heartbeats/log uploads) stop before a task is finalized (token revoked), and by making exit handling resilient to blocked handlers and re-entrant invocation.
Changes:
- Stop heartbeat sender (and log uploader where applicable) before pushing final task output.
- Start the force-exit watchdog before running exit handlers.
- Make
trigger_exit_handlersre-entrant-safe (avoid deadlocks) and add tests for both behaviors.
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| framework/py/flwr/supernode/runtime/run_clientapp.py | Stops heartbeats before pushing final task output to avoid post-finalization auth failures. |
| framework/py/flwr/supercore/task_process/model/run_model.py | Stops heartbeats earlier during shutdown to avoid token-revocation issues. |
| framework/py/flwr/supercore/task_process/connector/run_connector.py | Stops heartbeats before sending final task output. |
| framework/py/flwr/supercore/task_process/agent/run_agentapp.py | Stops log uploader/heartbeats before finalizing the task (token revocation). |
| framework/py/flwr/supercore/exit/exit_test.py | Adds test ensuring watchdog starts before exit handlers. |
| framework/py/flwr/supercore/exit/exit_handler_test.py | Adds re-entrancy test for exit handlers. |
| framework/py/flwr/supercore/exit/exit_handler.py | Avoids deadlock by invoking handlers outside the registry lock and clearing registry up-front. |
| framework/py/flwr/supercore/exit/exit.py | Reorders shutdown sequence so watchdog starts before running exit handlers. |
| framework/py/flwr/simulation/app.py | Stops authenticated background tasks before finishing the task and pushing final output. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c17c6ae100
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d64e823aa6
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d1fbc4e713
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6df0a49150
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ab51fefa0b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e273297b6a
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 29ab2994aa
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3593b630be
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7e1b1aa42b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 23 out of 23 changed files in this pull request and generated 1 comment.
Suppressed comments (3)
framework/py/flwr/common/logger.py:425
_log_uploaderuses an uninterruptibletime.sleep(LOG_UPLOAD_INTERVAL), sostop_log_uploader(..., timeout=...)cannot reliably stop the thread quickly (the join may time out while the uploader sleeps). To make shutdown truly bounded and responsive, replace the sleep with an interruptible wait (e.g., athreading.Event().wait(LOG_UPLOAD_INTERVAL)thatstop_log_uploadercan set, or block onlog_queue.get(timeout=...)so the sentinel wakes the thread immediately).
if exit_flag and not msgs:
break
time.sleep(LOG_UPLOAD_INTERVAL)
framework/py/flwr/supercore/heartbeat.py:148
- The gRPC heartbeat deadline is now
TASK_WORKER_CALL_TIMEOUT, but the heartbeat loop’s timing (inHeartbeatSender._run) is based on subtractingHEARTBEAT_CALL_TIMEOUTwhen computing the next interval. With a shorter RPC timeout, this mismatch changes the effective heartbeat cadence and can increase heartbeat frequency/load. Align the constants (useHEARTBEAT_CALL_TIMEOUTfor the RPC, or update the interval calculation to subtract the same timeout used for the RPC).
res = stub.SendTaskHeartbeat(req, timeout=TASK_WORKER_CALL_TIMEOUT)
framework/py/flwr/supercore/retry/retry_invoker.py:208
RetryInvokeris now mutated at runtime (viadisable_retries()) and is used from multiple threads in the new gRPC retry tests. Sincemax_tries/retries_disabledare read and written concurrently without synchronization, there is a risk of racey behavior across threads and runtimes. Consider guarding state changes/reads with a lock (or otherwise makingdisable_retries()thread-safe), especially aroundmax_trieschecks insideinvoke().
def disable_retries(self) -> None:
"""Disable further retries and interrupt an active retry wait."""
self.retries_disabled = True
self.max_tries = 1
if self.cancel_wait_function is not None:
self.cancel_wait_function()
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4a784e1a48
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
@codex review |
|
Codex Review: Didn't find any major issues. Breezy! Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
|
@codex review |
|
Codex Review: Didn't find any major issues. Can't wait for the next one! Reviewed commit: ℹ️ About Codex in GitHubCodex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback". |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 27 out of 27 changed files in this pull request and generated no new comments.
Suppressed comments (3)
framework/py/flwr/supercore/heartbeat.py:117
next_intervalcan become negative ifself._call_timeoutexceedsHEARTBEAT_DEFAULT_INTERVAL, which would cause the sender loop to run without sleeping (tight loop / high CPU). Clamp the computed base interval to a non-negative (or preferably a minimum heartbeat interval) before applying jitter/multipliers.
# Account for the RPC deadline when spacing heartbeat starts.
rd = random.uniform(*HEARTBEAT_RANDOM_RANGE)
next_interval: float = HEARTBEAT_DEFAULT_INTERVAL - self._call_timeout
next_interval *= HEARTBEAT_BASE_MULTIPLIER + rd
framework/py/flwr/supernode/runtime/run_clientapp.py:339
- Several exit paths compute timeouts via
max(0.0, deadline - time.monotonic()). Passingtimeout=0.0to gRPC effectively guarantees an immediateDEADLINE_EXCEEDED, which makes the final task output reporting fail deterministically once the budget is exhausted. Consider treatingtimeout <= 0as 'skip the RPC' (or enforce a small minimum positive timeout) to avoid guaranteed failures and noisy logs during shutdown.
if timeout is None:
stub.PushTaskOutput(request)
else:
stub.PushTaskOutput(request, timeout=timeout)
framework/py/flwr/supercore/heartbeat.py:98
stop()now raises only when the thread was never started (ident is None), but it no longer raises when the sender has already exited (it becomes effectively idempotent). Please update the docstring and/or the exception condition/message to reflect the intended semantics consistently (e.g., explicitly document idempotency, or restore the old behavior by checkingis_alive()if you want stop-on-stopped to error).
def stop(self, timeout: float | None = None) -> None:
"""Stop the heartbeat sender.
Parameters
----------
timeout : Optional[float] (default: None)
Maximum time in seconds to wait for the sender thread to stop.
"""
if self._thread.ident is None:
raise RuntimeError("Heartbeat sender is not running.")
self._stop_event.set()
self._thread.join(timeout=timeout)
|
Pull request was closed
Important
This original PR is now the umbrella and reference diff. Please review and merge the focused draft PRs below instead of merging this PR directly.
Root cause
PushTaskOutputcompletes a task and revokes its task token. Task executors previously sent this final RPC before stopping authenticated heartbeat and log-upload workers.A heartbeat racing with normal completion could therefore receive
UNAUTHENTICATEDand sendSIGINTwhile the main thread was already running exit callbacks. The nested exit attempted to acquire non-reentrant shutdown synchronization held by the interrupted thread. The force-exit watchdog also started only after callbacks returned, so this path could remain stuck indefinitely.The investigation also found related boundedness issues: active retry backoffs were not interrupted, heartbeat and log RPCs could outlive the shutdown window, sequential cleanup could consume the time needed for final task output, and runtime-environment cleanup could run before output reporting.
Focused PR stack
Merge these PRs in order. Each draft targets the preceding branch so its review diff contains only one focused change. After a predecessor merges, retarget the next PR to
main.flowchart TD P7914["#7914<br/>Exit foundation"] --> P7915["#7915<br/>Cancelable retry waits"] P7915 --> P7916["#7916<br/>gRPC retry shutdown"] P7916 --> P7917["#7917<br/>Bound heartbeat worker"] P7917 --> P7918["#7918<br/>Bound log uploader"] P7918 --> P7919["#7919<br/>Heartbeat-only executors"] P7919 --> P7920["#7920<br/>Logging executors"] P7920 --> P7921["#7921<br/>Runtime cleanup ordering"] P7921 --> P7922["#7922<br/>Fleet retry teardown"]Result
Together, the stack preserves the exact intended 27-file diff from this PR while separating process-exit infrastructure, retry cancellation, worker boundedness, executor ordering, cleanup priority, and Fleet teardown into independently reviewable changes.
Validation
git diff --checkmaincommit