fix(mm): close TOCTOU race in _restore_incomplete_installs

[lstein]

Summary

This PR fixes an intermittent unit test failure in test_model_install.py and, following review, several further races in the same code.

• _restore_incomplete_installs snapshotted active sources under the lock, then iterated tmpdirs against the stale snapshot. A foreground import_model call landing during iteration could be missed, the same source enqueued twice, and the second download would fail with FileNotFoundError when trying to rename the .downloading file the first download had already moved.
• Fix: move the membership check inside the lock immediately before _install_jobs.append(job) so check-and-append is atomic.
• Review round 2 (JPPhoto):
  • The duplicate-tmpdir dedup ran before the active-source check, so restore could delete an active job's tmpdir when two markers existed for one source. The dedup now runs inside the locked section after the active check, and no tmpdir is deleted while its source has an active job.
  • The lock was one-sided: import_model checked/registered without holding _lock. import_model now performs an atomic check-and-register.
• Review round 3 (JPPhoto):
  • Holding the installer lock across import_model's download enqueue inverted the lock order against download-queue callbacks (queue lock → installer lock) and could deadlock a new import against an in-flight callback. import_model now reserves the source in _pending_sources under the lock, runs the import helpers unlocked, and registers the job under the lock; concurrent imports of the same source wait on a condition variable. The lock is a plain Lock again, with the lock-order rule documented at its declaration.
  • An owner job completing between restore's marker read and its locked active check could look inactive while its marker was still on disk, letting restore enqueue a duplicate for a directory the owner was about to clean up. Restore now snapshots actively-owned sources before scanning: a source owned when the scan begins stays owned for the whole scan, and the per-marker locked check still covers imports that start mid-scan.

Fixes #9141.

Test plan

• pytest tests/app/services/model_install/test_model_install.py::test_simple_download — ran 10× consecutively, all pass (was flaky ~30–50% on main)
• Regression test test_restore_skips_source_queued_during_restore: pauses restore in the race window, queues a foreground job for the same source, asserts no duplicate is enqueued — fails on main, passes with this fix
• Regression test test_restore_preserves_active_jobs_tmpdir: two markers for one source, active job owns the later-visited dir — asserts the active dir survives restore
• Regression test test_restore_preserves_tmpdir_of_job_in_download_cache: same, with the active job tracked only in _download_cache with a different tmpdir
• Regression test test_concurrent_import_and_restore_register_single_job: a real import_model() racing restore registers exactly one job and completes end-to-end
• Regression test test_import_during_paused_download_callback_does_not_deadlock: a second import_model() requests a download job ID while a download callback is paused inside the queue lock — asserts both threads finish and both installs complete
• Regression test test_restore_skips_marker_of_job_completing_mid_scan: the owner transitions to COMPLETED after restore reads its marker — asserts no duplicate job or download is queued
• All review-round regression tests verified to fail against the pre-fix code
• Timing-sensitive tests (deadlock, mid_scan, concurrent_import_and_restore) pass 10/10 consecutive runs
• pytest tests/app/services/model_install/ — 39/39 pass

🤖 Generated with Claude Code

[JPPhoto]

This looks good when inspecting visually, but I'd also like to see a test added for this in tests/app/services/model_install/test_model_install.py, especially since it's a race condition. It may be a little challenging to get the test right, though:

1. Create a tmp install dir with an active marker for a remote source.
2. Monkeypatch _resume_remote_download() so restore pauses at a controlled point.
3. Insert an active job for the same source into _install_jobs while restore is between marker parsing and the new locked duplicate check.
4. Let restore continue.
5. Assert restore did not append a duplicate job and did not enqueue/resume a second download.

That maps directly to the PR’s fix around invokeai/app/services/model_install/model_install_default.py:249.

The only slightly fiddly part is choosing the pause point. Because the fixed code does the duplicate check under _lock immediately before append, the test probably needs to monkeypatch either _guess_source() or _read_install_marker() to signal “restore has observed the marker”, then add the foreground job before restore reaches the locked block. That is manageable with two threading.Events, similar to the existing test_import_waits_for_startup_restore.

The test should not need real network downloads if _resume_remote_download() is monkeypatched to fail the test if called for the skipped duplicate.

[lstein]

Thanks for the detailed test sketch @JPPhoto — added in b953bba, following it closely:

1. The test writes a tmpinstall_* dir with a version-1 marker (status downloading) for a URL source before starting the installer.
2. _guess_source is monkeypatched to signal "restore has observed the marker" and then block on a threading.Event, which pauses restore exactly between marker parsing and the locked duplicate check (two events, same pattern as test_import_waits_for_startup_restore).
3. While restore is paused, the test appends an active job for the same source to _install_jobs under _lock, as a concurrent import_model would.
4. Restore is released and the test waits for it to complete.
5. It asserts the foreground job is the only job for that source and that _resume_remote_download (monkeypatched to record calls) was never invoked — no network needed.

To confirm it actually covers the race, I ran the test against the pre-fix model_install_default.py from main: it fails there (duplicate job appended and resume attempted) and passes with the fix. Full tests/app/services/model_install/ suite passes (34/34).

[JPPhoto]

• invokeai/app/services/model_install/model_install_default.py:222: _restore_incomplete_installs() now marks a source as seen before it checks whether that source is already active. If there are two tmpinstall_* dirs with markers for the same source and an active job is already using the second dir, restore can see the stale first dir, add the source to seen_sources, skip it as active at lines 249-255, then treat the active second dir as a duplicate and delete it at lines 222-224. That can remove the active job's partial download directory while the job is still tracked in _install_jobs or _download_cache, causing later resume/download/install work to fail or restart incorrectly. To expose this issue, add a test that creates two install marker dirs for the same source, registers an active job whose _install_tmpdir is the second dir, forces restore to visit the first dir before the active dir, then asserts the active dir still exists and no duplicate restore job was queued.

• invokeai/app/services/model_install/model_install_default.py:249: The new lock is one-sided and does not make registration atomic with import_model(). The production import path checks _install_jobs at line 477, updates _download_cache at line 1264, and appends at line 495 without acquiring _lock. An import that passes _wait_for_restore_complete() immediately before start() clears the initially-set event can therefore race restoration: both paths can observe no active job and enqueue the same source. The new test masks this because it manually acquires _lock before appending at invokeai/tests/app/services/model_install/test_model_install.py:442, unlike the production path. To expose this issue, add a test that pauses a real import_model() call after its duplicate check, starts restoration for the same source, lets restore perform its locked check, and then allows both paths to register their jobs; assert that only one job and one download are created.

• invokeai/tests/app/services/model_install/test_model_install.py:381: The new regression test covers the single-marker duplicate-source case, but it does not cover the tmpdir ownership/deletion case above because it appends a foreground job using the same tmpdir as the marker at lines 436-443. It also bypasses the real remote import shape where _enqueue_remote_download() stores the job in _download_cache before import_model() appends it to _install_jobs (invokeai/app/services/model_install/model_install_default.py:1264 and :495). To expose this issue, add coverage where the active job is represented by _download_cache and has a different _install_tmpdir from an older duplicate marker for the same source.

[lstein]

Thanks @JPPhoto — all three findings were real, and they're addressed in ebaec12.

1. Restore deleting an active job's tmpdir. Confirmed exactly as you described: the seen_sources dedup ran before the active-source check, so a stale first dir marked the source seen and the active second dir was then rmtree'd as a "duplicate". The dedup now lives inside the locked section, after the active check, and when a source has an active job restore never deletes any of its tmpdirs (stale duplicates are collected on a later startup when the source is idle). The actual _safe_rmtree call happens after the lock is released. Covered by test_restore_preserves_active_jobs_tmpdir, which builds two marker dirs for one source, registers an active job owning the second, and forces restore to visit the stale dir first (sorted Path.glob).

2. One-sided lock. Also confirmed. _lock is now a threading.RLock (the import helpers call _next_id(), which acquires it), and import_model() holds it across its duplicate check, job creation, and registration — so import's check-and-register and restore's check-and-append are mutually atomic. One tradeoff worth flagging: the HF metadata fetch in _import_from_hf now runs under the lock, which briefly stalls download progress callbacks during an import's metadata lookup. Given imports are infrequent and human-triggered I think that's acceptable, but happy to hoist the fetch out with a re-check if you'd prefer. Covered by test_concurrent_import_and_restore_register_single_job, which drives a real import_model() (no manual locking this time): the import pauses inside _import_from_url after its duplicate check, start() runs in a second thread, and the test asserts exactly one job is registered, _resume_remote_download is never called, and the import then completes end-to-end. The test buffers _put_in_queue until after the assertions so the import job can't reach a terminal state (and delete its marker) before restore's active check — without that the test itself was flaky.

3. _download_cache coverage. Added test_restore_preserves_tmpdir_of_job_in_download_cache: the active job is tracked only in _download_cache with a different _install_tmpdir than an older duplicate marker for the same source.

Verification: all three new tests fail against the pre-fix code (checked by reverting model_install_default.py and re-running) and pass with the fixes; the concurrent test passed 15/15 consecutive runs; full tests/app/services/model_install/ suite passes 37/37; ruff clean.

[JPPhoto]

Two issues:

• invokeai/app/services/model_install/model_install_default.py:493: import_model() now holds the installer lock while remote import reaches _multifile_download(), which acquires the download queue lock through _next_id(). Download callbacks take these locks in the opposite order: invokeai/app/services/download/download_default.py:691 holds the download queue lock while invoking _download_progress_callback(), which acquires the installer lock at invokeai/app/services/model_install/model_install_default.py:1400. If a second remote import begins while an earlier download callback is running, the callback can wait for the installer lock while the import waits for the download queue lock, permanently deadlocking both threads. This can freeze active downloads and leave the new import request blocked indefinitely. To expose this issue, add a test that pauses an existing multifile callback while it holds the download queue lock, starts another remote import_model() call until it holds the installer lock and requests a download job ID, then releases the callback and asserts both threads finish within a timeout.

• invokeai/app/services/model_install/model_install_default.py:253: The new per-marker active check excludes terminal jobs, but job terminal transitions, marker deletion, and temporary-directory cleanup are not synchronized with this lock. Restore can read an active job's marker, the installer thread can then mark that job completed at line 1487, and restore can subsequently see no nonterminal owner and append a second job for the same marker at line 264. The original up-front active_sources snapshot removed by this PR would continue treating a job that was active when restoration began as owned. The original job then deletes the marker and temporary directory while the duplicate restore job resumes from that directory, causing a duplicate installation or filesystem failure. To expose this issue, add a test that starts restore with an active job owning the marker directory, pauses restore after reading the marker, transitions the owner to COMPLETED while leaving its marker and directory present, resumes restore, and asserts that no second job or download is queued for that source.

[lstein]

Thanks @JPPhoto — both findings were real, and both are fixed in f60cbf7 along with the two regression tests you outlined.

1. Lock-order deadlock in import_model()

Confirmed: _mfd_started/_mfd_progress hold the download queue's lock while invoking the installer callbacks (which acquire the installer lock), and the round-2 version of import_model() held the installer lock while reaching the queue's _next_id() — a classic ABBA inversion.

Fix: import_model() no longer holds the installer lock across the import helpers. Instead it reserves the source in a new _pending_sources set under the lock (after the duplicate check), runs the helpers — metadata fetch, tmpdir selection, download enqueue — with the lock released, then registers the job and clears the reservation under the lock. Restore's per-marker check now also treats a pending source as active, so check-and-register remains atomic with respect to restore. Two concurrent import_model() calls for the same source are handled with a condition variable: the second waits for the first to register (or fail) and then re-runs the duplicate check, so it returns the first call's job instead of enqueuing twice.

Since nothing acquires the lock reentrantly anymore, the RLock from round 2 is reverted to a plain Lock, and the lock-order rule (never call into the download queue while holding the installer lock) is documented at the lock's declaration. This also resolves the tradeoff I flagged last round — the HF metadata fetch no longer runs under the lock, so download progress callbacks are never stalled behind it.

Test: test_import_during_paused_download_callback_does_not_deadlock follows your outline — it pauses the first import's on_start callback on the queue thread (queue lock held, installer lock not yet acquired), drives a second real import_model() to the point of requesting a download job ID from the queue (observed via a _multifile_download wrapper), releases the callback, and asserts both threads finish within a timeout and both installs complete. Two teardown details: the test uses a private DownloadQueueService rather than the shared fixture, and a daemon import thread — on the pre-fix code the deadlock is unbreakable, and with the shared fixture its teardown joins the wedged worker threads and hangs the entire test run (pytest-timeout's thread method can't interrupt a blocked lock acquire); with the private queue the test instead fails cleanly in ~15s.

2. Owner completing mid-scan re-enqueued by restore

Confirmed: the per-marker active check alone can't see a job that was active when restore read the marker but transitioned to COMPLETED before the locked check, since terminal transitions and marker/tmpdir cleanup aren't synchronized with the installer lock.

Fix: _restore_incomplete_installs() now snapshots the set of actively-owned sources (non-terminal _install_jobs, non-terminal _download_cache entries, and _pending_sources) under the lock before scanning, restoring the original code's semantics: a source owned when the scan begins stays owned for the whole scan, even if the owner reaches a terminal state mid-scan. The per-marker locked check is kept in addition, to cover imports that start during the scan.

Test: test_restore_skips_marker_of_job_completing_mid_scan follows your outline — restore is paused after reading the owner's marker (via the _guess_source wrapper), the owner transitions to COMPLETED with its marker and directory left on disk, restore resumes, and the test asserts no second job or download is queued for the source and the directory is untouched.

Verification

• Both new tests fail against the previous commit (deadlock test: import_model deadlocked against a download callback after the 15s join timeout; mid-scan test: duplicate job appended and _resume_remote_download called) and pass with the fixes.
• Full tests/app/services/model_install/ suite: 39/39.
• The three timing-sensitive tests (deadlock, mid_scan, concurrent_import_and_restore) pass 10/10 consecutive runs.

[JPPhoto]

The reservation closes the deadlock, but it introduces an unsolved handoff case: restore treats a pending import as an owner and permanently skips its marker, even though that import may fail before registering any job:

• invokeai/app/services/model_install/model_install_default.py:222: _restore_incomplete_installs() permanently copies _pending_sources into owned_sources, but a pending import is only a reservation and may fail before registering a job. If import_model() is pending when restoration starts and its helper raises at lines 531-548, the reservation is removed, but restore still skips every marker for that source at lines 273-282. Startup then finishes with the incomplete install neither restored nor owned by an active job, deferring recovery until another restart or manual retry. Test: create an incomplete marker, pause import_model() after it reserves the source, start restoration so the reservation is included in owned_sources, make the import helper raise before registering a job, and assert restoration rechecks the failed reservation and queues the marker instead of completing with no owner.