Serialize imports behind startup restoration - #9433
Conversation
github-actions
Bot
added
python
lstein
left a comment
lstein
left a comment
There was a problem hiding this comment.
Reviewed at 842b01d against #9142, since the two are framed as alternatives. Up front: I think the ordering invariant in this PR is the better design, and I'd rather land it than my own PR. But I don't think it closes #9141 as it stands, so I'm requesting changes rather than approving.
Everything below was verified against the code and, where noted, reproduced.
Context worth stating: main already has the barrier
#9239 (merged, db0b08b) already added _restore_completed_event and made import_model() / wait_for_installs() wait on it. So this PR isn't introducing serialization — it's closing the remaining ways to slip past a barrier that already exists (event starts unset instead of set; startup failures propagate; wait_for_installs honors its timeout at the barrier). Worth saying explicitly in the PR body, because the "69 vs 411 lines" comparison reads differently once you account for the barrier already being in main.
Why I prefer this design over #9142
With a genuine barrier, restore's active_sources snapshot at L215-217 is a snapshot of a quiesced _install_jobs. That single fact makes most of #9142 unnecessary — the _pending_sources deferral, _source_import_generations, the deferred-marker recheck and DEFERRED_RESTORE_TIMEOUT all exist only because restore and import can overlap there. Removing the concurrency beats coordinating it, and restoration is a one-shot startup operation where nothing is gained by letting imports through. I'd rather maintain this.
Blocking: #9141 is still reachable, via import-vs-import
The barrier serializes restore against import. It doesn't serialize import against import, and import_model (L498-520) still takes no lock: similar_jobs is computed at L501, but self._install_jobs.append(install_job) happens at L519 — after the metadata fetch and after submit_multifile_download. The dedup window spans a full network round trip.
That matters because of _find_reusable_tmpdir (L191-208). _enqueue_remote_download writes the marker with status WAITING at L1313 before submitting at L1317, so a second import arriving at L1225 after the first reached L1313 is handed the first import's tmpdir, and both download into it.
Reproduced on this branch (both threads run well after restoration completed, so the barrier is not what's being tested):
# tests/app/services/model_install/test_probe.py
import threading, time
from pathlib import Path
import pytest
from pydantic_core import Url
from invokeai.app.services.config import InvokeAIAppConfig
from invokeai.app.services.model_install import ModelInstallServiceBase
from invokeai.app.services.model_install.model_install_common import (
InstallStatus, ModelInstallJob, URLModelSource,
)
from invokeai.app.services.model_install.model_install_default import TMPDIR_PREFIX
from invokeai.app.services.model_records import ModelRecordChanges
from tests.backend.model_manager.model_manager_fixtures import * # noqa F403
@pytest.mark.timeout(timeout=60, method="thread")
def test_concurrent_imports_of_same_source(
mm2_installer: ModelInstallServiceBase, mm2_app_config: InvokeAIAppConfig
) -> None:
source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors"))
mm2_installer._restore_completed_event.wait(timeout=10)
# A prior interrupted install leaves a marker, so _find_reusable_tmpdir() hands out the SAME dir.
tmpdir = Path(mm2_app_config.models_path) / f"{TMPDIR_PREFIX}reusable"
tmpdir.mkdir()
stub = ModelInstallJob(id=99998, source=source, config_in=ModelRecordChanges(), local_path=tmpdir)
stub._install_tmpdir = tmpdir
mm2_installer._write_install_marker(stub, status=InstallStatus.DOWNLOADING)
real = mm2_installer._import_from_url
def _slow(src, config=None):
time.sleep(1.0) # widen the existing window; the interleaving itself is legal
return real(src, config)
mm2_installer._import_from_url = _slow
threads = [threading.Thread(target=lambda: mm2_installer.import_model(source)) for _ in range(2)]
for t in threads:
t.start()
for t in threads:
t.join(timeout=25)
jobs = mm2_installer.get_job_by_source(source)
assert len(jobs) == 1, f"{len(jobs)} jobs for one source, tmpdirs={[j._install_tmpdir for j in jobs]}"Results:
| branch | jobs for one source | outcome |
|---|---|---|
main @ ae5694d |
2, same tmpdir | both ERROR — FileNotFoundError: .../test_embedding.safetensors |
| this PR @ 842b01d | 2, same tmpdir | identical failure |
| #9142 @ 74622e1 | 1 | COMPLETED |
That's the #9141 symptom verbatim, with the barrier fully in place.
Production path, no test harness: heuristic_import() is called from invokeai/app/invocations/flux_redux.py:156 and ip_adapter.py:220 on session-processor worker threads, and multi-GPU parallel sessions (#9263) makes two of those genuinely concurrent — two graphs needing the same not-yet-installed SigLIP / CLIP-Vision encoder. Both search_by_attrs return empty, both pass similar_jobs == [], both land in the same tmpdir. The HTTP routes are async def so they can't race each other; the node path is the exposure. Note the codebase already recognized this exact hazard for the sibling path and fixed it with _download_cache_lock (L757-797, comment: "parallel (multi-GPU) sessions ... don't race to download into the same cache directory"). import_model never got the equivalent.
Suggested fix: the reservation half of #9142 — a _pending_sources: set[str] plus a Condition on _lock; reserve the source under the lock, run the import helpers unlocked, register under the lock, and have a concurrent import of the same source wait and then re-run the duplicate check. That's ~40 lines and nothing else from #9142 is needed once this barrier exists. The unlocked-helpers part is not optional: download_default.py:691 dispatches _execute_cb while holding the queue lock, and those callbacks take the installer _lock (L1414/1429/1442/1451/1463), so holding _lock across submit_multifile_download inverts the order.
Non-blocking, cheap
-
L374
except Exceptionshould beexcept BaseException(or afinally). The event is cleared at L355 and only set by the restore thread spawned at L373; aBaseExceptionescaping in between leaves the barrier permanently closed with_running=Trueand_startup_error=None, so_wait_for_restore_completepasses both guards and blocks forever at L305. In fairness I could not reach this in production — insidecatch_sigint()Ctrl-C hard-kills viaSIG_DFL+raise_signal, and otherwiserun_app.py:144tears the process down — so this is robustness, not a live bug. Still a one-word fix, andmaindidn't have the hole (the event was set in__init__and only cleared as the last statement ofstart()). -
_startup_erroris read before the wait and never re-read. L299 snapshots it, L301 releases, L305 blocks. A startup failure occurring while a caller is parked at L305 sets_startup_errorat L375 and sets the event at L376 — so the parked waiter is woken by the very event signalling failure, returnsTrue, and proceeds as if startup succeeded. Re-check under_lockafter L305 returns. Needs a stop/start cycle to hit, so low severity, but it defeats the contract this PR adds. -
L294 skips the new checks on the lock-timeout branch.
return Falsethere bypasses L297-303 entirely, so a contended never-started service surfacesTimeoutErrorfromwait_for_installsinstead of the intendedRuntimeError("...is not running"). Practically unreachable today, but it makes the new contract conditional on lock contention. -
L297's guard is
and, so it only ever fires for a never-started service.stop()(L379-389) sets_running = Falsebut never clears_restore_completed_event, so after a normal stop the conjunction is False and imports sail straight through into a stopped service. Pre-existing behavior, not a regression — but the error message claims a check that isn't performed. Relatedly,stop()never clears_stop_eventand never joins the restore thread, so on a stop→start cycle the old restore thread'sfinallyat L286 can open the new barrier. Not reachable in the app (nothing restarts the service), but it's a sharp edge under the new "the event means restoration finished" semantics. -
test_import_waits_for_startup_restorehas no@pytest.mark.timeout, unlike its two new siblings, and its teardowninstaller.stop()→_install_thread.join()is unbounded.
What I attacked and couldn't break
For the record, so this isn't re-litigated:
- No lock-order inversion, and none introduced. I enumerated every installer
_lockholder and every download-queue method reachable from it. The callbacks that run under_lockreach onlyDownloadQueueService.cancel_job/pause_job, neither of which takes the queue lock. Every_enqueue_remote_downloadcall site runs with_locknot held. Single-direction graph. _wait_for_restore_completereleases_lockat L301 before waiting at L305. That one line is what makes this design deadlock-safe against the restore thread's L215 acquisition — worth a comment so nobody "simplifies" it into awithblock.- No self-deadlock from
start()holding the non-reentrant_lockacross its whole body. Traced all six callees;_next_id()(L1152, which takes_lock) is unreachable from any of them. - Every production caller of
import_model/heuristic_importis provably post-start(), so nothing gets a spurious "not running":Invoker.__init__startsmodel_managerbeforesession_processor(fixed insertion order ininvocation_services.py),sync_configured_external_starter_modelsruns afterInvoker.__init__returns, and all install routes areasync defserved after lifespan startup. wait_for_installstimeout accounting is correct —start = time.time()at L548 is taken before the barrier and counted once;max(0.0, ...)at L294/L304 rules out negative timeouts (which would mean "block forever"). There are no production callers ofwait_for_installsanyway.- All three new tests genuinely fail against pre-change
main— I ran them at ae5694d:3 failed.test_import_waits_for_startup_restorestill fails with the white-box first assertion deleted (DID NOT RAISE TimeoutError), so it's behavioral, not just a flag check. The monkeypatched_observed_wait_for_restorecloses over the real bound method, so it doesn't dodge the code under test. download_and_cache_model,resume_job,restart_*don't bypass the barrier meaningfully — the first writes to a different tree and is already per-source locked;resume_jobreturns unlessjob.paused, and restorecontinues past paused jobs without resuming them.tests/app/services/model_install/passes 35/35 locally, CI green.
Summary
Land this design. Add the import_model reservation so #9141 is actually closed, and I'll close #9142 in favor of it. Items 1-3 above are one-liners; 4 and 5 are your call.
|
Thanks for the thorough review. I confirmed the blocking import-vs-import race. I added a deterministic regression test that holds the first import after its duplicate check while starting a second import for the same source. Before the fix, it registered two jobs; after the fix, both callers receive the same job. I implemented the narrow reservation portion you suggested:
Additional changes:
I agree the PR description should explicitly state that #9239 introduced the original barrier. This PR closes the remaining early-pass paths and now adds same-source import serialization; the earlier raw line-count comparison with #9142 omitted that context. On the other non-blocking points:
Validation:
|
lstein
left a comment
lstein
left a comment
There was a problem hiding this comment.
Re-reviewed at 33812c958a. The blocking finding is fixed — approving. Thanks for taking the reservation approach; the result is tighter than what I had in #9142, and I'll close that one in favour of this.
I pushed two small hardenings to the branch as 99c69bee1a (see below) — revert them if you disagree, the approval stands either way.
Verification
I confirmed the fix two ways, neither of them your tests.
1. My round-1 probe (2 threads, seeded reusable tmpdir, time.sleep(1.0) inside _import_from_url) now passes.
2. A randomized stress probe — 8 threads across 2 sources, 0-0.6 s jitter, a DOWNLOADING marker seeded per source so _find_reusable_tmpdir() hands out a shared directory. It asserts the actual #9141 symptom rather than a job count: no two jobs share an _install_tmpdir, and no job ends up errored.
| branch | result |
|---|---|
33812c958a |
8/8 pass |
842b01d6f4 (previous head) |
5/5 fail — FileNotFoundError: .../tmpinstall_reuse0/test_embedding.safetensors.downloading -> .../test_embedding.safetensors |
That's #9141 verbatim on the old head and gone on this one.
Your new tests also genuinely fail against 842b01d6f4: test_concurrent_imports_of_same_source_return_one_job → assert 2 == 1, test_base_exception_during_startup_releases_import_waiters → event unset. (test_failed_import_releases_source_reservation passes on the old head too — it's a guard for the new mechanism rather than a regression test, which is fine, just noting it isn't load-bearing.) CI green, 17/17.
I also accept both of your rebuttals:
- Pre-wait
_startup_errorread — you're right and I was wrong to raise it.start()holds_lockacross 355-381, and_wait_for_restore_completeacquires_lockat 297 before snapshotting at 303, so no caller can take the snapshot mid-startup. There's no interleaving that readsNoneand is then woken by a sync-path failure. - Lock-timeout branch — agreed, determining the state requires the lock;
import_modelalways takes the blocking-acquire path anyway.
What I pushed
Both are one-liners in the invariant this PR now owns, each with a regression test that fails against 33812c958a.
(a) Check similar_jobs before new_jobs. A waiter can be released into a state where the source has both a job that was registered while it waited and has since gone terminal, and a live one — and new_jobs[0] returns the dead one. Reproduced:
T1 reserves S, parked in its helper. T2 enters, known_job_ids = {}, parks at :509.
T1 appends job 88001, notify_all. job 88001 is/goes ERROR.
T3 enters: similar_jobs empty (88001 terminal) -> reserves S, appends live job 88002, notify_all.
T2 finally wins the lock: new_jobs = [88001, 88002] -> returns new_jobs[0].
→ T2 got job 88001 (status=InstallStatus.ERROR) instead of the live job 88002
To be precise about the scope: this only bites when a live job co-exists. It does not change the case you documented in the comment — where the owner's job is the only new one and is already terminal, the waiter still receives it under either ordering, which I agree is the right call. Worth knowing that case is more reachable than it looks: submit_multifile_download runs at :1342 but the append is at :542, so a fast download failure (DNS, connection refused) can run _download_error_callback → _set_error in between, and _put_in_queue → cancel_job (:420-421) does the same after stop().
(b) prune_jobs() filters and reassigns under the condition. It was an unsynchronized read-modify-write on the attribute the dedup now treats as authoritative, while import_model appends under _install_condition at :542:
prune :669 unfinished_jobs = [...] # snapshot, job not yet appended
import :542 self._install_jobs.append(job) # appends to the OLD list object
prune :670 self._install_jobs = unfinished_jobs
The consequence is exactly the bug this PR closes: the source ends up in neither _install_jobs nor _pending_sources, its marker still says WAITING, so the next import passes both checks and _find_reusable_tmpdir() hands it the same tmpdir. Reachable — DELETE /api/v2/models/install is async def on the event loop, while heuristic_import runs on session-processor threads (flux_redux.py:156, ip_adapter.py:220).
Two things worth stating plainly, because I got the first one wrong on my first attempt:
- It still rebinds, deliberately. My first cut mutated in place (
self._install_jobs[:] = [...]), which is worse: a CPython list iterator is index-based, so shrinking the list under an unlocked reader (get_job_by_id,get_job_by_source,list_jobs) silently skips the tail. I reproducedValueError: No job with id 66003 knownfor a live job that way. The rebind keeps existing iterators on a stable list object, which is the property the original code had for free. - It costs something:
prune_jobsnow waits on_lock, and_download_error_callback(:1487) and_download_cancelled_callback(:1508) hold that lock across_safe_rmtree, which doesgc.collect()+rmtree()and up to 1.5 s of sleeps on win32. So cancelling a partly-downloaded multi-GB model and then hitting "clear finished" can now stall the event loop for the duration of that rmtree. It's the same lockinstall_modelalready blocks the loop on at :297, so this widens an existing exposure rather than creating a new one — but if you'd rather not take that trade for a window this narrow, say so and I'll drop (b).
I also could not trigger the lost update by blocking inside the comprehension, because a list comprehension picks up concurrent appends — the real window is only the gap between the two statements. The test asserts the invariant directly instead: while a prune is mid-filter, _lock must be held and a concurrent import_model must not get a job registered.
40 passed locally (38 + 2), 3× under random ordering, ruff clean.
Observations, no action needed on this PR
-
The reservation wait is unbounded and can wedge a multi-GPU worker.
_install_condition.wait()(:509) has no timeout, and the reservation is held across_remote_files_from_source→HuggingFaceMetadataFetch—grep -rn timeout invokeai/backend/model_manager/metadata/fetch/returns nothing, so a stalled connection blocks forever. Twosession_processor_Nworkers hittingflux_redux.py:156for the same SigLIP source is the designed-for case: worker 0 wins the reservation and hangs, worker 1 parks at :509. A worker'scancel_eventcan't interrupt aCondition.wait(), so that GPU's queue is stuck until restart, andstop()(:383-393) touches neither_pending_sourcesnornotify_all(). Workers aredaemon=True, so this doesn't hang process exit — it's a runtime wedge, not a shutdown one. The underlying no-timeout HTTP is the real bug and predates you; a boundedwait()plus anotify_all()instop()would contain it. -
A single bad marker aborts restoration of every other one, and the barrier still reports success. In
_restore_incomplete_installs, the per-markertryends at :249, butModelRecordChanges(**...)(:251),InstallStatus(status)(:262) and_put_in_queue(:270) are outside it. Anything raised there escapes theforloop;_runcatches at :287, logs, and thefinallysets the event without setting_startup_error. So_wait_for_restore_complete()returnsTrueand imports proceed believing restoration completed, when nothing was restored. Widening thetryto cover :251-270 with acontinuewould make the barrier's guarantee actually hold. (Trigger is cross-version markers —_read_install_markeronly rejects a version mismatch, so a marker from a build with a newInstallStatusmember is accepted at :180 and explodes at :262 on rollback.) -
The
andat :301 still doesn't fire in the case that actually happens. You're right that stop/restart is out of scope, but plain shutdown isn't:Invoker.stop()iteratesvars(self.services)in attribute order, andmodel_manageris assigned atinvocation_services.py:109versussession_processorat:117— soinstall.stop()runs before the session workers are stopped. A worker still insideheuristic_importthen finds_runningFalse but the event set, sails past the guard, and gets a silently-cancelled job (local) or a download submitted to an already-stopping queue. Not a regression —mainhas no_runningcheck at all — but the check you added misses its most likely real caller. Clearing_restore_completed_eventinstop()would make the guard mean what its message says. -
test_concurrent_imports_of_same_source_return_one_jobhas a real flake path.second_helper_enteredcan never be set — T2 blocks at :509, which is before the monkeypatched_import_from_url— sosecond_helper_entered.wait(timeout=1)is a fixed 1 s sleep whose result is discarded. If T2 is descheduled past that second on a loaded runner,release_first_importfires before T2 reaches :506; T1's job then completes (the mock session serves instantly), T2'sknown_job_idsalready contains it,similar_jobsis empty because it's terminal, and T2 creates a second job →assert len(jobs) == 1fails. Waiting on observable state instead would make it deterministic, andassert not second_helper_entered.wait(...)would state the intent. -
QA instructions still say "Expected result: 35 tests pass" — the directory has 40 now.
What I attacked and couldn't break
- Lock order. Helpers run with
_lockreleased (:522-534), so the established queue-lock → installer-lock direction is preserved; the only installer→queue edges (cancel_jobat :1457/:1480) callDownloadQueueService.cancel_job, which takes no queue lock. No cycle. - Re-entrancy on the non-reentrant
_lock._next_id()(:1178) is reachable only from the helpers, which run outside the condition._put_in_queue(:419) takes no lock, so_download_complete_callback→_put_in_queueunder_lockis safe, and_install_queueis unbounded soput()can't block. - Reservation leaks / double-remove / lost wakeups. Both exits (:535-539, :541-544) remove and
notify_allunder the lock; the check-and-add at :508-520 never releases the lock, so reservation is exclusive andremovecan'tKeyError. - Three-thread barging. :508 is a
while, so a waiter re-tests after every wakeup — the only fallout is (a) above. - Key/predicate drift.
source_key = str(source)andjob.source == sourceboth reduce tostr()viaStringLikeSource.__eq__, and__str__omitsaccess_token, so_import_from_hf's in-place token mutation at :1205-1206 is invisible to dedup. - Cross-source tmpdir sharing.
_find_reusable_tmpdirfilters onmarker["source"] != source_str(:203), so a per-source reservation is sufficient. except BaseExceptioninstart(). It does close the hole:_restore_incomplete_installs_async()is the last statement in the try, nothing runs between the try and thewithexit, and aThread.start()failure propagates into the handler, which sets the event._run'sfinally(:289-290) also runs forBaseException.- Barrier deadlock.
_lockis released in thefinallyat :305 beforeEvent.wait()at :309, so no waiter holds it while blocked; the restore thread'swith self._lockat :219 and_next_id()at :253 don't overlap. wait_for_installstimeout accounting.start = time.time()at :573 precedes the barrier wait, bothmax(0.0, ...)clamps hold, no double-counting. (_install_queue.join()at :583 is still unbounded, but that's pre-existing and outside this diff.)download_and_cache_modelbypassing the barrier. It does bypass it, but it never writes_install_jobs/_download_cacheand writes only undermodels/.download_cache, which can't matchmodels_path.glob("tmpinstall_*"). Benign.
Two hardenings to the import reservation. import_model() checked new_jobs before similar_jobs, so a waiter released into a state where the source has both a job registered while it waited that has since gone terminal and a live one would return the dead job, reporting a failure for a source that is actively installing. Check for a live job first; the documented case where the owner's job is the only new one and is already terminal is unchanged. prune_jobs() filtered and reassigned _install_jobs with no lock, so an import_model() registration landing between the two was dropped - leaving a live install invisible to the duplicate check, its marker still WAITING, and the next import of that source reusing its tmpdir. Do both under the install condition. Rebind rather than mutating in place so that readers already iterating the old list are not silently truncated. Both tests fail against the previous commit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
99c69be to
a7e2600
Compare
|
Thanks for the re-review, independent stress testing, and the two hardenings. I agree with both changes in Preferring an active I also agree that the concurrent-import test has a scheduling-dependent path. I've replaced the discarded one-second wait with observable synchronization. The unbounded reservation wait, malformed-marker isolation, and shutdown behavior are useful findings, but are broader pre-existing lifecycle concerns. I would handle those separately rather than expand this fix further. I added some tests to verify everything (the current test count is now 43):
|
…9448) The four timeout=5 marks added in #9433 use method="thread", so the 5-second budget covers fixture setup and teardown as well as the test body. Tearing down the mm2_download_queue fixture alone takes up to ~1s (five worker threads polling the queue at a 1-second interval), and on a heavily loaded CI runner the total easily exceeds 5s: the py3.11 windows-cpu job on #9447 timed out in test_base_exception_during_startup_releases_import_waiters during download-queue teardown, and an unrelated branch hit the same timeout in test_import_fails_after_startup_failure on linux-cpu the same day. Neither dump showed a deadlock - the workers were in their normal poll loop. Bump these marks to 30s, matching the other timeout marks in this file. The timeouts exist to catch hangs, not to enforce speed, so the larger budget loses nothing. Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2 participants
Summary
Fixes a TOCTOU race between startup restoration and foreground model imports.
PR #9239 introduced the startup restoration barrier by making
import_model()andwait_for_installs()wait for_restore_completed_event. However, the event initially started set, leaving paths that could bypass the barrier. This PR closes those paths, hardens failure and timeout handling, and serializes concurrent imports of the same source._restore_incomplete_installs()previously scanned temporary install directories using an active-source snapshot that could become stale. A concurrentimport_model()could register the same source during that scan, causing duplicate downloads and a laterFileNotFoundErrorwhen both jobs attempted to move the same.downloadingfile.This change establishes startup restoration as a barrier:
import_model()waits for restoration to finish before examining or registering jobs.wait_for_installs()applies its timeout while waiting at the barrier.This removes restore/import concurrency instead of coordinating ownership while both operations run.
Difference From #9142
This is an alternative to #9142 and accepting one should close the other.
Unlike #9142, this PR builds on the barrier introduced by #9239 and removes restore/import overlap. It only adds per-source coordination for concurrent foreground imports; #9142 coordinates restoration and imports concurrently through a broader ownership and deferred-restoration protocol.
PR #9142 allows foreground imports and startup restoration to execute concurrently. It makes that concurrency safe using:
This PR changes the ordering invariant: foreground imports cannot begin until startup restoration has finished. Because restore and import no longer overlap, it does not need a second source-ownership protocol or deferred-marker state machine.
The tradeoff is that an import requested during startup waits for restoration. Restoration already gates reliable knowledge of active installs, so this keeps synchronization at the startup boundary and leaves the normal import path unchanged after startup.
Scope comparison against current
main:Related Issues / Discussions
Fixes #9141.
Alternative implementation to #9142.
QA Instructions
Run:
pytest tests/app/services/model_install/Expected result: all tests pass.
Regression coverage verifies:
wait_for_installs()honors its timeout during restoration.Merge Plan
Checklist
What's Newcopy (if doing a release after this PR)