Skip to content

Python: fix: make concurrent FileCheckpointStorage saves not race on a shared temp path - #7757

Open
Manjunath Janardhan (manjunathshiva) wants to merge 20 commits into
microsoft:mainfrom
manjunathshiva:python-checkpoint-concurrent-save-7748
Open

Python: fix: make concurrent FileCheckpointStorage saves not race on a shared temp path#7757
Manjunath Janardhan (manjunathshiva) wants to merge 20 commits into
microsoft:mainfrom
manjunathshiva:python-checkpoint-concurrent-save-7748

Conversation

@manjunathshiva

@manjunathshiva Manjunath Janardhan (manjunathshiva) commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

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 then os.replaced that fixed path, so whichever save renamed first removed the temp file still being written by a competing save, which then failed — FileNotFoundError on POSIX (per the issue) and PermissionError [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>.tmp in the
    destination directory, created with os.open(O_CREAT | O_EXCL | O_WRONLY, 0o666) so the file
    keeps the process umask exactly as the previous open(..., "w") did. Same directory keeps
    os.replace atomic, 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.Future signal, awaited through a per-wait future fed by a
    done-callback rather than through asyncio.wrap_future, which chains cancellation into the future
    it 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 could
    fill the default pool, stall unrelated to_thread work including checkpoint loads, and deadlock
    once the write that had to finish first was queued behind those waiters. And because a
    concurrent.futures.Future is not bound to a loop, one chain orders writers enqueued from
    different event loops, which a per-path asyncio primitive cannot do.

    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 nothing was submitted. Releasing only from the coroutine would be
    a regression against the threading.Lock it 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 WorkflowCheckpoint generates a fresh
    UUID 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.replace is in flight.

    The drain absorbs re-delivered cancellations. A single await asyncio.shield(worker) is not
    enough: 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.run cancels pending tasks first and that drives a queued save through its cancellation
    path — 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.replace can transiently
    raise PermissionError even fully serialized, because a background indexer or AV scan briefly
    holds 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, where
    before 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 the
    same 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.Future so any loop can await
    it, 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_future was safe here on the strength of a
    check 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.py at 99% line coverage, poe syntax and
    poe typing clean across all five checkers, and ag-ui, declarative and foundry_hosting green as
    the 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.py rather than
    restating 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 -- fails
    59/100 then 85/100 with PermissionError on main (matching the 49-71 reported) and 0/100 here
    across three runs. The surviving checkpoint is also deterministic now: main leaves a different
    winner each run because the last os.replace to land is whichever the executor scheduled last,
    where this branch leaves the last save enqueued, every run.

Related Issue

Fixes #7748

Contribution Checklist

  • The code builds clean without any errors or warnings
  • All unit tests pass, and I have added new tests where possible
  • The PR follows the Contribution Guidelines
  • This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
  • This is not a breaking change.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 around os.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.

Comment thread python/packages/core/agent_framework/_workflows/_checkpoint.py Outdated
Comment thread python/packages/core/agent_framework/_workflows/_checkpoint.py Outdated
Comment thread python/packages/core/tests/workflow/test_checkpoint.py
@moonbox3

Copy link
Copy Markdown
Contributor

/review

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread python/packages/core/agent_framework/_workflows/_checkpoint.py Outdated
Comment thread python/packages/core/agent_framework/_workflows/_checkpoint.py Outdated
Comment thread python/packages/core/agent_framework/_workflows/_checkpoint.py Outdated
Comment thread python/packages/core/agent_framework/_workflows/_checkpoint.py Outdated
…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
@manjunathshiva

Copy link
Copy Markdown
Contributor Author

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.

@moonbox3

Copy link
Copy Markdown
Contributor

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 to_thread(lock.acquire) for waiting, as it still occupies executor threads and can deadlock when the actual write is queued behind those waiters.

We also need to address cancellation explicitly. shield() allows the caller to receive cancellation while the worker continues, and a per-path lock does not establish ordering between workers queued on different event loops. Please remove a cancelled write before submission, or drain its worker before propagating cancellation and releasing ownership.

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.
@manjunathshiva

Copy link
Copy Markdown
Contributor Author

Thanks — revised A is in, pushed as e4f30b885. Point by point:

  • Process-wide coordination keyed by canonical destination. A module-level registry keyed by the
    resolved destination path, guarded by a threading.Lock, so it spans coroutines, loops and
    FileCheckpointStorage instances.
  • Asynchronous waiters, and no to_thread(lock.acquire). Ownership is taken before the write is
    submitted; waiters suspend on asyncio.wrap_future over a concurrent.futures.Future and occupy
    no executor worker. Nothing blocks on a lock inside the pool, so the deadlock you described cannot
    form.
  • Cleanup after the final queued/running operation. Entries are reference-counted and the last
    release drops them.
  • Cancellation, explicitly. Cancelled before submission: the write is removed, never submitted,
    and ownership is still handed on so later saves are not stalled. Cancelled mid-write: the worker is
    drained before the cancellation propagates and before ownership is released.
  • Ordering across event loops. This is why the hand-off is a concurrent.futures.Future rather
    than anything loop-bound -- any loop can await it, so a single chain orders writers enqueued from
    different loops.
  • No separate executor per checkpoint path. None added; the default to_thread executor is still
    the only one, and now only the active write occupies it.
  • Deterministic tests, including distinct-ID registry cleanup. Ten new or rewritten tests, all
    event- and barrier-driven rather than timing-dependent, and every wait bounded so none can hang.
    Where a test needs several saves to be genuinely queued before it acts, it waits on the observable
    pending count rather than on a sleep. Verified across repeated runs under -n auto --dist worksteal, which is what poe test -P core already uses. Details on the three threads.

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
await asyncio.shield(worker) in the except CancelledError block does not drain, because the
pending cancellation makes the next await raise immediately -- the shield returns with the worker
still running. It now absorbs re-delivered cancellations until the worker is done, with a test that
cancels five times while a write is parked.

The second was a regression I introduced against the design being replaced. Taking ownership out of
the worker thread meant releasing it from the coroutine, and a coroutine whose loop is closed while
its task is pending never releases -- so the destination stayed owned for the life of the process and
later saves to it hung. A threading.Lock did not have that failure mode, because the worker thread
releases it regardless of the loop. Release now happens on the worker thread as the write's last act,
with the coroutine's finally as a backstop, and there is a test for it.

Separately, while checking the environment I found my local uv sync had been incomplete, which was
masking 111 skipped tests and about 150 pyright errors in files unrelated to this change. Re-synced;
core is 4341 passing with _checkpoint.py at 98% line coverage, and poe syntax and poe typing
are clean across all five checkers. Mentioning it because it means my earlier validation claims on
this PR were made against a thinner environment than I thought.

Two questions for you.

Draining changes what a caller observes: a cancelled save() now returns only after its own write
finishes, where before it returned promptly and the write continued behind it. That is the point of
the fix, but it is a real behaviour change for a caller that cancels and expects to proceed
immediately. The checklist on this PR says it is not a breaking change; I have left it as-is rather
than relabel your PR — tell me if you would rather it were marked.

And delete() deliberately does not join the queue, so a delete racing a save to the same path is
still unserialized. That was out of scope for #7748 and I did not want to widen the change without
asking, but say the word and it is a small addition.

One thing I added to the module comment rather than leaving implicit: the scope really is one
process. Replicas sharing a checkpoint directory over a mounted volume or network share get nothing
from this, since there is no shared state to coordinate through, and which write survives is
undefined -- the file not being seen half-written rests on os.replace being atomic on that
filesystem. Worth stating where the design is explained, because "process-wide" reads like a
guarantee if you are deploying more than one replica against the same volume.

Comment thread python/packages/core/agent_framework/_workflows/_checkpoint.py Outdated
Comment thread python/packages/core/agent_framework/_workflows/_checkpoint.py Outdated
…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.
@manjunathshiva

Copy link
Copy Markdown
Contributor Author

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
verified the latter rather than assumed it". That was wrong. asyncio cancels the wrapped future
from a done-callback that runs on a later loop iteration, and my check read the future immediately
after the CancelledError, so it reported cancelled=False when one iteration later it was
cancelled=True. A false negative presented as verification, which is worse than not having checked
at all — you found it by reading the code. The description is corrected and the mechanism is gone.

What the two findings had in common is worth stating, because it is the actual defect: asyncio
objects were standing in for the state of work that lives outside asyncio, and cancellation can mark
those done while the work continues. wrap_future chains cancellation into a shared signal; a
cancelled task reports done() while its thread is still inside os.replace. The rule now is that
the hand-off signal is the only source of truth and nothing asyncio-owned can mark it done.

Fixing them turned up three more failures of the same kind, all mine, all found by probing rather
than by review — flagging them because each is worse than what it replaced:

  • A hang. ensure_future(asyncio.to_thread(...)) makes the write a Task; shutdown cancels every
    task, possibly before it reaches the executor. Nothing released the destination and the drain
    waited on a signal that could never be resolved, so the save never completed and asyncio.run's
    shutdown hung. Now submitted with run_in_executor (a plain Future, invisible to
    asyncio.all_tasks()) plus a done-callback that releases in every terminal state.
  • A crash. Deferring a cancelled ticket's hand-off made release recursive; 1200 chained
    cancellations raised RecursionError on the worker thread and leaked the destination permanently.
    Releases are now drained in a loop.
  • Accumulation. The drain registered a done-callback per wait, so 200 cancellations left 200
    callbacks on one signal. Now one, pointing at the current waiter.

Also removed a submitted flag the rewrite made dead, and deleted a test of my own that asserted
nothing after I stripped its only assertion — it guarded an optimisation with no observable
behaviour.

Validation: poe check -P core green — lint, fmt, all five type checkers, 4349 tests, _checkpoint.py
at 97% line coverage. Deterministic across three runs under -n auto --dist worksteal. All seven
downstream consumers of checkpoint storage green: ag-ui 1142, declarative 989, foundry_hosting 286,
devui 200, orchestrations 195, azure-cosmos 59, hosting 23. Six regression tests, each failing
against the previous commit.

The two questions from the last round still stand: whether draining should be labelled a breaking
change, since a cancelled save() now returns only after its own write finishes; and whether
delete() should join the queue.

Comment on lines +351 to +352
with contextlib.suppress(RuntimeError):
loop.call_soon_threadsafe(_set)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. 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 is loop.run_in_executor now: a plain Future that asyncio.all_tasks() does not
    sweep, submitted synchronously, plus a done-callback that releases on success, failure and
    cancellation alike.
  2. RecursionError at 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.
  3. 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.
@manjunathshiva

Copy link
Copy Markdown
Contributor Author

Pushed 6783a30e5 and a86e94698, plus a merge with main to pick up #7948. The technical
answer to your queued-path question is in the inline thread; this is the part that belongs at PR
level.

This is the fourth consecutive round where you have found a real defect in code I had already
called validated.
Cancellation into the predecessor's future, the worker-signal ownership, the
wrap_future claim I reported as verified when my check read the future one loop iteration too
early, and now the queued path's dependence on a loop that had gone away. Every one of them you
found by reading the code, and every one of them was in a version I had convinced myself was done.
That is a pattern in how I was validating, not four unlucky bugs, so I have changed two things
rather than just fixing the fourth.

First, I stopped treating "I checked it" as evidence unless I can point at what failed. The
wrap_future claim is the cleanest example of the failure mode: a probe that reported what I
expected, run at the wrong moment, presented as verification. So for this round I re-ran all 25 new
or rewritten tests against main's _checkpoint.py instead of restating an earlier claim. All 25
fail there. I mention it because the description had been saying "ten" for several rounds — true
once, never updated, and exactly the kind of stale confidence that produced the wrap_future error.
That number is corrected now.

Second, I reviewed the cumulative diff rather than the increment. Reading +386 lines the way you
would see them, instead of the piece I had just changed, is what surfaced the last commit:
a86e94698 adds no source change, only tests for three error branches this PR introduced that had
none — the bounded os.replace retry giving up, the temp-file cleanup swallowing an OSError, and
the already-resolved fast path in the drain. The retry one is worth calling out: the suite was
reaching those lines, but only because the concurrency tests occasionally trip a real
PermissionError on Windows. That is environment-dependent and never happens on Linux CI, so the
coverage was an accident rather than a test. _checkpoint.py is at 99% now, and the three lines
still uncovered all predate this PR.

Fixing your finding also turned up three defects I had introduced myself — a shutdown hang, a
RecursionError at roughly 1200 chained cancellations with a permanent leak, and callback
accumulation in the drain. They are listed with their fixes in the inline reply. I would rather
show you the ones I caught than imply the fix arrived clean.

On #8214. It is approved and edits save() and _write_atomic — the same function this PR
rewrites — so there is a conflict coming, and I would rather flag what it must preserve than
discover it in the resolution. Two things: the save-time encode/decode validation, and
encoding="utf-8" on the write plus all three read sites.

That second one is a live bug, not a style fix. _write_atomic pairs ensure_ascii=False with a
stream opened at the platform default, so on Windows (cp1252 here) saving a checkpoint containing
any CJK text, emoji or many accented characters raises UnicodeEncodeError — I reproduced it. It
predates this PR and my rewrite of that line carried it forward unchanged.

I deliberately have not fixed it here. Fixing only the write side would be worse than leaving
it: files would go out as UTF-8 and still be read back at the locale default, turning a loud
write-time error into silent read-time corruption. Fixing both sides would duplicate an approved PR
and make the conflict worse. So #8214 should land it, and I will rebase on top preserving both of
its changes rather than resolving in favour of my version of the function. Flagging it because
"keep ours" is the natural resolution here and it would silently revert your approved fix.

End-to-end evidence, since CI has not run. I re-ran the issue's own repro -- 100 concurrent
save() calls for one checkpoint ID -- against both trees on Windows:

failures survivor
main 59/100, then 85/100 (PermissionError) w99, then w95 -- varies per run
this branch 0/100, three runs w99 every run

The failure rate matches the 49-71 the issue reported. The survivor column is the part I had not
tested before: on main the winner changes run to run because the last os.replace to land is
whichever the executor happened to schedule last, so a caller cannot know which checkpoint
survives. Here it is deterministically the last save enqueued. The description claims that ordering
property; this is the first time I have actually measured it rather than reasoned about it.
Distinct-ID saves: 0/100 failures, 100/100 readable back with correct content, no temp files left,
registry drained to empty.

On delete() being out of scope, you asked what that leaves open, and "unserialized" was too
vague an answer. Concretely, delete() is if file_path.exists(): file_path.unlink() on a worker
thread, outside the queue. Racing a save it can unlink the file a just-completed os.replace put
there, silently discarding a save that returned successfully; and two concurrent deletes can lose
the TOCTOU race so unlink() raises FileNotFoundError out of a method whose contract is to
return False. Neither corrupts a file -- os.replace stays atomic -- but the first loses a write
that reported success. Still my recommendation to keep it out of #7748's scope, now that it is
stated precisely enough for you to overrule me.

How long draining actually blocks, since the description asks whether it should be labelled a
breaking change and "a cancelled save now waits" is not something you can weigh without a number.
Measured on a local SSD, median of five:

checkpoint on-disk size save cancel blocks for
1 executor 662 B 3.0 ms 2.4 ms
100 executors 250 KB 4.3 ms 3.8 ms
500 executors 2.5 MB 16.9 ms 15.2 ms

And the case that matters more under contention -- cancelling saves that are still queued behind
an in-flight write -- costs 0.2 ms for five of them. They do not wait for the running write at
all; the deferred hand-off means they raise immediately, the in-flight write finishes undisturbed,
and the queue drains clean. So the new blocking applies only to the one caller whose own write is
actually in flight, bounded by that single write.

Both call sites in the engine (_runner_context.py:481, _functional.py:1216) await save()
inline, so this is what a cancelled workflow run now pays. Single-digit milliseconds on local
storage. It scales with write time, so a network share or mounted volume would cost proportionally
more -- worth knowing, but the bound is one write either way. My read is that this does not warrant
the breaking-change label, though it is your call and the behaviour change is real.

Two things I have deliberately left alone and want your call on:

  • The two loop-abandonment windows described in the inline thread and in the description. I
    chose to state them rather than build the loop-tracking and reaping needed to close them. Happy to
    add the machinery if you would rather have it.
  • _validate_file_path's path-traversal guard has no test_checkpoint.py:560, uncovered
    across the whole repo, and it predates this PR ([BREAKING] Python: Checkpoint refactor: encode/decode, checkpoint format, etc #3744). I have not touched it, since this PR is
    about write serialization and I would rather not widen a diff you are already four rounds into.
    Say the word and I will add it here, or open it separately.

One request: the Python workflow runs on this PR have been sitting at action_required since the
branch was pushed — Python - Tests, Code Quality, Test Coverage and Merge - Tests have never
executed, so the checks list makes the PR look untested. Everything above was validated locally on
Windows only. Could you approve the pending runs when you get a moment? I would rather you see CI
agree with me than take my word for it, especially on this PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

python Usage: [Issues, PRs], Target: Python

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Python: [Bug]: Concurrent FileCheckpointStorage saves fail due to shared temporary path

3 participants