Python: fix: make concurrent FileCheckpointStorage saves not race on a shared temp path - #7757
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 addresses a concurrency race in FileCheckpointStorage.save() when multiple saves target the same checkpoint ID, and adds a regression test to ensure concurrent saves don’t fail.
Changes:
- Added a concurrency regression test covering concurrent saves with the same checkpoint ID.
- Updated
FileCheckpointStorage.save()to use unique temp files and added per-checkpoint-ID serialization plus a retry loop aroundos.replace()for Windows. - Implemented best-effort cleanup for temp files when failures occur.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| python/packages/core/tests/workflow/test_checkpoint.py | Adds a regression test that exercises concurrent save() calls for the same checkpoint ID. |
| python/packages/core/agent_framework/_workflows/_checkpoint.py | Makes save() more robust under concurrency by changing temp-file strategy, adding per-ID locking, and retrying replace on Windows. |
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
a32ff36 to
fd0a15a
Compare
|
/review |
There was a problem hiding this comment.
MAF Automated Review — Iteration 1
Result: Findings reported
Scope: full PR (1 commit(s)): fd0a15ab45a6
Model: gpt-5.6-sol
Overview
The change removes the shared temporary-path race by giving each save a unique same-directory file, publishing it atomically, and serializing same-ID saves at the coroutine level. The regression test verifies concurrent saves complete and leave a parseable checkpoint. Residual compatibility and lifecycle issues remain around active-lock eviction, maximum-length checkpoint IDs, event-loop retention, and changed file permissions.
Reviewed the supplied pull-request change set across correctness, security/reliability, architecture, and failure behavior.
4 verified findings remained after source verification (4 medium) across 1 file. Details are attached to the affected lines below.
Affected areas: python/packages/core/agent_framework/_workflows/_checkpoint.py
…a shared temp path FileCheckpointStorage.save() wrote to and then renamed a fixed "<checkpoint-id>.json.tmp" path, so concurrent saves of the same checkpoint ID raced over the shared temp file. Whichever save renamed first removed the temp file still being written by another save, which then failed with FileNotFoundError / PermissionError in os.replace. Create a unique temp file per save in the destination directory (so os.replace remains atomic), serialize same-ID writes with a per-ID lock, and retry the atomic move briefly to absorb the transient Windows background-handle PermissionError that surfaces even for fully serialized replaces. Fixes microsoft#7748
fd0a15a to
24eb113
Compare
|
Evan Mattson (@moonbox3) Gentle ping — no rush at all, just checking in. I laid out options A/B/C/D above and leaned toward (A) as the middle ground, but I'm equally happy to implement (B) if you'd prefer the stronger isolation. Whenever you have a moment to weigh in, I'm ready to proceed. |
|
Manjunath Janardhan (@manjunathshiva), Thanks for laying these out. Let’s proceed with a revised A: process-wide coordination keyed by canonical destination, with asynchronous waiters and cleanup after the final queued/running operation. Please don’t use We also need to address cancellation explicitly. Lastly, please cover those cases with deterministic tests, including distinct-ID registry cleanup. I’d avoid a separate executor per checkpoint path. |
…rite Implements the revised A shape chosen in review: process-wide coordination keyed by canonical destination, asynchronous waiters, and registry entries released after the final queued or running operation. Three review findings shared one cause -- ownership was established inside the executor thread rather than before submission. Acquiring the destination lock inside the worker meant every same-path save occupied an `asyncio.to_thread` worker while merely waiting, so a burst could fill the default pool and stall unrelated work including checkpoint loads, and could deadlock once the write that had to finish first was queued behind those waiters. Ownership now comes from a queue keyed by the canonical path and is taken before the write is submitted, so only the active write holds a worker. Waiters suspend on `asyncio.wrap_future` over a `concurrent.futures.Future`, which is not bound to a loop, so one chain also orders writers enqueued from different event loops -- something a per-path lock could not do. The registry was a dict that never removed an entry. Because `WorkflowCheckpoint` generates a fresh UUID by default, every save retained one entry for the life of the process. Entries are now reference-counted by queued-or-running operations and dropped by the last release. Cancellation is handled explicitly, both ways the review allowed. A save cancelled while still waiting for the destination has nothing in the executor, so it is simply removed -- and it still hands ownership on, or every later save for that path would wait forever. A save cancelled mid-write drains its worker before propagating, because releasing first would let the next save start while this replace is still in flight. That drain absorbs re-delivered cancellations: once a task has a cancellation pending its next await raises immediately, so a single shielded await returns with the worker still running and the write can still land after ownership is released. Ownership is released as the write's last act on the worker thread, with the coroutine's `finally` as a backstop for the case where no write was submitted. Releasing only from the coroutine regressed on the design it replaced: a `threading.Lock` was released by the worker thread, which outlives the loop that submitted it, whereas a coroutine that never resumes -- its loop closed with the task still pending -- would leave the destination owned for the life of the process and hang every later save to it. The two releasers are idempotent, so exactly one of them accounts for each ticket. Tests are deterministic and cover each case: registry release for repeated, distinct-ID and concurrent saves; queued saves not occupying executor threads, gated by a two-worker executor with three saves in flight; serialization across two real event loops; cancellation before submission writing nothing and not stalling the follower; repeated cancellation still draining; a write failing while draining being reported rather than lost; a failed save releasing its destination and leaving the path usable; recovery of a destination whose loop was closed mid-write; and the defensive release path.
|
Thanks — revised A is in, pushed as
Two things worth surfacing rather than leaving in the diff. The drain was wrong in my first attempt, and the existing test caught it. A single The second was a regression I introduced against the design being replaced. Taking ownership out of Separately, while checking the environment I found my local Two questions for you. Draining changes what a caller observes: a cancelled And One thing I added to the module comment rather than leaving implicit: the scope really is one |
…tate Follow-up review on microsoft#7757 found two ways a later save could still overtake an earlier one. Both had the same cause: asyncio objects were standing in for the state of work that lives outside asyncio, and cancellation can mark those done while the work continues. `asyncio.wrap_future` chains cancellation into the future it wraps. Cancelling the middle of three queued saves therefore cancelled the hand-off signal the first save still had to resolve, and the unconditional release then let the third reach `os.replace` beside the first -- after which the first write could land last and overwrite the newer checkpoint. Waiting now goes through `_wait_for_signal`, which feeds a fresh per-wait future from a done-callback and leaves the shared signal untouched no matter what happens to the waiter. A save cancelled while queued no longer resolves its own signal either; it defers the hand-off until its predecessor has actually finished, so the queue stays ordered while still guaranteeing nothing waits behind it forever. Draining on the task wrapping `asyncio.to_thread` was the second: a cancelled task reports `done()` while its function is still inside `os.replace`, so loop shutdown released ownership with the write in flight. The drain now waits on the signal the worker thread resolves, which cancellation cannot mark done. Fixing those exposed two more failures of the same kind, both found by probing rather than by review. `ensure_future(asyncio.to_thread(...))` makes the write a Task, and shutdown cancels every task -- possibly before it reaches the executor. Nothing then released the destination, and the drain waited for a signal that could never be resolved, so the save never completed at all. The write is now submitted with `run_in_executor`, which submits synchronously and returns a plain Future that `asyncio.all_tasks()` does not include, plus a done-callback that releases on success, failure and cancellation alike. Deferring a cancelled ticket's hand-off onto its predecessor made release recursive: resolving one signal ran the next ticket's callback synchronously, so a run of cancelled queued saves longer than the recursion limit raised `RecursionError` on the worker thread partway through and left the destination owned for good. Releases are now collected per thread and drained in a loop, so the depth is flat however long the run. Two smaller points from the same review. A write that fails while draining is recorded on the ticket rather than read back from the worker future, because a cancelled task hides its own exception and the cancelled caller never sees it. And the drain registers one done-callback pointing at whichever waiter is current, rather than one per wait, so a caller cancelling in a loop no longer grows that list in step with it. Tests: cancelling the middle of three queued saves cannot let the third overtake; a cancelled worker task does not release while the thread writes; shutdown before the write starts neither hangs nor strands the destination, asserted structurally on the write not being a task so the regression fails cleanly instead of hanging CI; a deferred-release run of 2000 links does not recurse; submission against a dead executor releases; and resolving a signal after its loop closed does not raise.
|
Both findings are fixed and pushed. Before the detail, a correction to this PR's description. It said the hand-off "relies on ... a cancelled waiter not poisoning its predecessor's future — I What the two findings had in common is worth stating, because it is the actual defect: asyncio Fixing them turned up three more failures of the same kind, all mine, all found by probing rather
Also removed a Validation: The two questions from the last round still stand: whether draining should be labelled a breaking |
| with contextlib.suppress(RuntimeError): | ||
| loop.call_soon_threadsafe(_set) |
There was a problem hiding this comment.
Following up on the earlier loop-shutdown ownership report: could the queued path also release without depending on its event loop? This commit makes the worker signal authoritative after submission, but _wait_for_signal() suppresses RuntimeError when a predecessor finishes after the waiter's loop has closed. The queued ticket then never reaches _release_write_after(), so its completion stays unresolved and every later save for that destination hangs. Could the predecessor callback own an abandonment or release path that does not require the closed loop to resume the coroutine?
There was a problem hiding this comment.
Yes — and you were right that the suppression was the whole problem. Fixed in 6783a30e5.
_wait_for_signal now takes an on_abandoned hook, and the queued path passes it
_release_write(ticket). The RuntimeError is still caught, but instead of being swallowed it is
the signal to release on the thread that resolved the predecessor, rather than on a loop that
will never run again:
try:
loop.call_soon_threadsafe(_set)
except RuntimeError:
if on_abandoned is not None:
on_abandoned()One property makes this correct rather than merely unblocking: on_abandoned can only fire from
inside the predecessor's own done-callback, so at that moment the predecessor has genuinely
finished. Releasing there preserves the same ordering invariant _release_write_after exists to
protect — the successor cannot start its os.replace beside a write that is still running.
test_file_checkpoint_storage_abandoned_loop_while_queued_releases_the_destination pins it: a save
queued behind a held write, its loop closed underneath it, and a third save on another thread that
must still complete.
Two windows this does not close, stated rather than papered over. A loop that closes after
call_soon_threadsafe succeeded but before the callback runs, and a loop abandoned without being
closed at all. Both leave a queued ticket unreleased. Both need a loop closed with tasks still
pending, which asyncio already reports as an error. Closing them means tracking each ticket's owning
loop and reaping, which is real machinery in this path for a failure only reachable that way — I
chose to state them. Say if you would rather have the machinery and I will add it.
A graceful shutdown is unaffected, and that is the case callers actually depend on: asyncio.run
cancels pending tasks before closing, which drives a queued save through its cancellation path and
defers the hand-off onto the predecessor. test_file_checkpoint_storage_graceful_shutdown_releases_a_queued_save
pins that separately, so the guarantee is not resting on the argument above.
Three defects I introduced while fixing this, all found before pushing and all now covered by
tests, since the pattern on this PR has been that my first shape is not the last one:
- Shutdown hang. My first attempt submitted the write with
ensure_future(asyncio.to_thread(...)), which makes it a Task — and loop shutdown cancels every
task, so the write could be cancelled before it ever reached the executor, leaving nobody to
release. It isloop.run_in_executornow: a plain Future thatasyncio.all_tasks()does not
sweep, submitted synchronously, plus a done-callback that releases on success, failure and
cancellation alike. RecursionErrorat roughly 1200 chained cancellations, with a permanent leak when it hit —
each deferred release ran the next one synchronously, one stack frame per link. Flattened with a
per-thread trampoline (_ReleaseTrampoline) that collects and drains in a loop.- Callback accumulation in the drain: 200 cancellations left 200 callbacks on one signal. Now
one callback pointing at whichever waiter is current.
Follow-up review on microsoft#7757. A submitted write is released by its worker thread, which outlives the loop that submitted it, but a ticket still waiting for its predecessor has no worker to fall back on. If that waiter's loop closes, `call_soon_threadsafe` raises `RuntimeError`, the suppression swallowed it, and the coroutine never resumed -- so nothing resolved the ticket's hand-off signal and every later save for that destination waited forever. Reproduced: the entry stays at `pending=1` and a later save from a fresh loop times out. The predecessor's callback now owns that case. `_wait_for_signal` takes an `on_abandoned` hook, invoked when the waiter's loop has closed, and the queued wait passes its own release -- performed on the thread that resolved the predecessor rather than on a loop that will never run again. The drain keeps suppressing, and the two guards now differ deliberately: by the time anything drains, the write has been submitted, so the worker thread and the submitted future's callback both release without needing that loop. The comment says so, since the asymmetry is otherwise unexplained. Two windows this cannot close, both stated in the comment rather than implied: a loop that closes after `call_soon_threadsafe` succeeded but before the callback runs, and one abandoned without being closed at all. Reaching either needs a loop closed with tasks still pending, which asyncio already reports as an error. Graceful shutdown is unaffected -- `asyncio.run` cancels pending tasks first, which drives a queued save through its cancellation path and defers the hand-off onto its predecessor. That path now has its own test, since it is the guarantee callers actually depend on.
Three branches added by this PR had no test. Coverage of `_checkpoint.py` was reporting 97% with the gaps sitting exactly on the error handling: * the bounded `os.replace` retry giving up after five attempts. The suite only ever reached the surrounding lines by accident -- the concurrency tests trip a real `PermissionError` on Windows now and then, which is environment-dependent and never happens on Linux CI. The new test pins both halves that matter: the retry is bounded, and giving up still releases the destination rather than wedging every later save to it. * the temp-file cleanup swallowing an `OSError`. It runs in a `finally` while a write exception is propagating, so letting one out would substitute a misleading "cannot remove temp file" for the real disk failure. Reverting the handler makes the new test fail with exactly that substitution. * the fast path in `_await_signal_through_cancellation` for a signal that has already resolved. Exercised directly, since arranging the race through `save()` is not deterministic. No source change. Every line this PR adds is now covered; the three that remain uncovered in the file all predate it.
…current-save-7748
|
Pushed This is the fourth consecutive round where you have found a real defect in code I had already First, I stopped treating "I checked it" as evidence unless I can point at what failed. The Second, I reviewed the cumulative diff rather than the increment. Reading +386 lines the way you Fixing your finding also turned up three defects I had introduced myself — a shutdown hang, a On #8214. It is approved and edits That second one is a live bug, not a style fix. I deliberately have not fixed it here. Fixing only the write side would be worse than leaving End-to-end evidence, since CI has not run. I re-ran the issue's own repro -- 100 concurrent
The failure rate matches the 49-71 the issue reported. The survivor column is the part I had not On How long draining actually blocks, since the description asks whether it should be labelled a
And the case that matters more under contention -- cancelling saves that are still queued behind Both call sites in the engine ( Two things I have deliberately left alone and want your call on:
One request: the Python workflow runs on this PR have been sitting at |
Motivation & Context
Concurrent calls to
FileCheckpointStorage.save()for the same checkpoint ID raced over a shared temporary file (<checkpoint-id>.json.tmp). Each save wrote to and thenos.replaced that fixed path, so whichever save renamed first removed the temp file still being written by a competing save, which then failed —FileNotFoundErroron POSIX (per the issue) andPermissionError [WinError 5]on Windows (where the concurrent-replace also trips the destination-lock check). In the issue's 100-concurrent-save repro, between 49 and 71 saves failed.Description & Review Guide
Three cooperating parts in
_checkpoint.py, the last two reshaped by review.Unique temp file per save. Each save writes to its own
.maf-ckpt-<uuid>.tmpin thedestination directory, created with
os.open(O_CREAT | O_EXCL | O_WRONLY, 0o666)so the filekeeps the process umask exactly as the previous
open(..., "w")did. Same directory keepsos.replaceatomic, and per-save temp names remove the shared-path collision that Python: [Bug]: Concurrent FileCheckpointStorage saves fail due to shared temporary path #7748 reported.Ownership taken before the write is submitted. A process-wide queue keyed by the canonical
destination path hands out ownership; only then is the write submitted to the executor. Waiters
suspend on a
concurrent.futures.Futuresignal, awaited through a per-wait future fed by adone-callback rather than through
asyncio.wrap_future, which chains cancellation into the futureit wraps.
Two properties follow from that choice. A queued same-path save occupies no executor worker, where
previously each one held a worker purely to block on
threading.Lock.acquire()— a burst couldfill the default pool, stall unrelated
to_threadwork including checkpoint loads, and deadlockonce the write that had to finish first was queued behind those waiters. And because a
concurrent.futures.Futureis not bound to a loop, one chain orders writers enqueued fromdifferent event loops, which a per-path
asyncioprimitive cannot do.Ownership is released as the write's last act on the worker thread, with the coroutine's
finallyas a backstop for the case where nothing was submitted. Releasing only from the coroutine would be
a regression against the
threading.Lockit replaces: a lock is released by the worker thread,which outlives the loop that submitted it, whereas a coroutine whose loop is closed while its task
is pending never releases -- leaving the destination owned for the life of the process.
Registry entries are reference-counted by queued-or-running operations and dropped by the last
release. The previous dict never removed an entry, and since
WorkflowCheckpointgenerates a freshUUID by default, ordinary saves each retained one for the life of the process.
Cancellation, handled explicitly in both directions. Cancelled while still waiting for the
destination: the write is never submitted, so there is nothing left to land later — and ownership is
still handed on, or every later save for that path would wait forever. Cancelled once the write is
running: the worker is drained before the cancellation propagates and before ownership is released,
since releasing first would let the next save begin while this
os.replaceis in flight.The drain absorbs re-delivered cancellations. A single
await asyncio.shield(worker)is notenough: once a task has a cancellation pending its next await raises immediately, so the shield
returns with the worker still running.
Ownership survives the loop that took it. A submitted write is released by its worker thread,
which outlives the loop that submitted it. A ticket still waiting for its predecessor has no
worker to fall back on, so the predecessor's callback releases it directly when the waiter's loop
has closed — on the thread that resolved the predecessor, rather than on a loop that will never
run again. Without that, a queued save whose loop went away left its hand-off signal unresolved
and every later save for that destination waited forever.
Two windows this cannot close: a loop that closes after the wake-up was queued but before it runs,
and one abandoned without being closed at all. Both need a loop closed with tasks still pending,
which asyncio already reports as an error. A graceful shutdown is unaffected, because
asyncio.runcancels pending tasks first and that drives a queued save through its cancellationpath — pinned by its own test, since that is the path callers actually depend on.
Bounded replace retry. Retained from the original fix. On Windows
os.replacecan transientlyraise
PermissionErroreven fully serialized, because a background indexer or AV scan brieflyholds a handle to the destination (19 failures in 200 purely sequential replaces locally). Five
attempts with 1→2→4→8→16 ms backoff absorb it.
What are the major changes? Destination ownership moved out of the worker thread and in front of
submission, reference-counted registry entries, and an explicit cancellation contract. No public
API change.
What is the impact of these changes? Concurrent saves of one checkpoint ID no longer fail, and
last-writer-wins now follows the order callers observed rather than the order the executor happened
to schedule. Queued saves no longer consume executor capacity, so a save burst cannot starve
checkpoint loads.
One observable behaviour change: a cancelled
save()returns after its own write finishes, wherebefore it returned promptly and the write continued behind it. That is the fix — a cancelled write
landing later could overwrite a newer checkpoint — but it does mean a caller that cancels no longer
proceeds immediately.
Deliberately out of scope:
delete()does not join the queue, so a delete racing a save to thesame path remains unserialized. Python: [Bug]: Concurrent FileCheckpointStorage saves fail due to shared temporary path #7748 is about concurrent saves; happy to add it if wanted.
What do you want reviewers to focus on? The cancellation contract, since it is the part with a
caller-visible consequence, and whether draining should be labelled a breaking change.
Also the two loop-abandonment windows described above. I chose to state them rather than engineer
around them: closing them means tracking each ticket's loop and reaping, which is real complexity
in this path for a failure only reachable through a loop closed with tasks still pending. Say if
you would rather have the machinery.
Second, the ownership hand-off. The signal is a
concurrent.futures.Futureso any loop can awaitit, and nothing asyncio-owned may mark it done — cancelling a waiter must not reach it, and the
submitted write must not be a Task, or loop shutdown can cancel it before it reaches the executor.
An earlier revision of this description claimed
wrap_futurewas safe here on the strength of acheck that read the wrapped future one loop iteration too early; it is not, and the mechanism has
been replaced.
Validation: core 5147 passing with
_checkpoint.pyat 99% line coverage,poe syntaxandpoe typingclean across all five checkers, and ag-ui, declarative and foundry_hosting green asthe downstream users of checkpoint storage. The three lines still uncovered in the file all
predate this PR.
This PR adds or rewrites 25 tests. I re-ran all 25 against
main's_checkpoint.pyrather thanrestating an earlier claim, and all 25 fail there. An earlier revision of this description said
"ten", which had been true several rounds ago and was never updated.
End to end, the issue's own repro -- 100 concurrent
save()calls for one checkpoint ID -- fails59/100 then 85/100 with
PermissionErroronmain(matching the 49-71 reported) and 0/100 hereacross three runs. The surviving checkpoint is also deterministic now:
mainleaves a differentwinner each run because the last
os.replaceto land is whichever the executor scheduled last,where this branch leaves the last save enqueued, every run.
Related Issue
Fixes #7748
Contribution Checklist