[wunianze666-netizen]

Summary

Protect application-initiated model offloads against concurrent model construction. Parameter assignment during Accelerate's process-global init_empty_weights patch can otherwise strand offloaded weights on meta.

Acquire the shared model-load read lock before the cache lock at architecture-switch, session-end and node API offloads. Both public node offload methods now own that lock; Qwen no longer acquires an outer lock. The public API docstrings, cache lock contracts and generated invocation-context documentation describe the ownership rule. Independent cache offloads retain shared-reader concurrency.

Related Issues / Discussions

Closes #10009.

QA Instructions

• Final regression fixture against the original make_room_in_vram production method: one intended failure, with weights on meta; the same fixture passes with the fix.
• Focused Qwen, node API, session offload, lock and documentation-export suites: 72 passed.
• Separate eight-file related run: 80 passed, 3 CUDA-dependent skips. The focused and related suites overlap.
• pnpm run generate-docs-data exited 0. Ruff 0.11.2 check and format check, git diff --check, and normal commit hooks passed. The unchanged uv lock hook had no files to check.
• The pytest runs used a task-local offline runner invoking pytest -n 0, with isolated application/temp roots and BLAS threads limited to one. Exact invocations are recorded in the task's command logs.

The focused run used these pytest arguments after the task-local offline runner:

-n 0 -o cache_dir=<isolated-cache> --basetemp=<isolated-temp>
tests/app/services/shared/test_invocation_context_make_room_in_vram.py
tests/app/services/shared/test_invocation_context_offload_from_vram.py
tests/app/services/session_processor/test_architecture_switch.py
tests/app/services/session_processor/test_release_vram_after_session.py
tests/app/invocations/test_qwen_image_text_encoder.py
tests/backend/model_manager/load/model_cache/test_model_load_rw_lock.py
tests/test_docs_json_export.py --durations=10

The tests exercise real Torch parameter assignment, Accelerate construction, ModelCache and the shared locks with small float32 CPU modules; VRAM residency/capabilities and allocator I/O are simulated. Windows / Python 3.12.14 / Torch 2.13 CPU / Diffusers 0.42.1 source overlay. Physical GPU, full model, family E2E, full locked environment and unpaced offload duration were not measured.

On the previous head, Windows ARM64 job 129337633587 failed before pytest because Chocolatey's FFmpeg 9.0.2 feed returned HTTP 504. That attribution applies only to its actual checkout; checks for the new head must be evaluated separately.

Review

The public lock ownership inconsistency and review nits were addressed; relevant checks were rerun. Lock waits remain non-cancellable, following the existing model-load lock contract. AI assistance was used for implementation and test authoring.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

---

^{Open in [code]smith · Need help on this PR? Tag @codesmith-bot with what you need.}

[Pfannkuchensack]

Thanks for picking up #10009. The fix closes the race at all three sites and I found no blocker. One thing I'd like changed before merge, the rest are nits.

Reviewed at 35d0b87.

Should fix

The two offload methods on the node API now follow opposite lock rules, and neither docstring says so.

• context.models.offload_from_vram now takes the read lock itself (invocation_context.py:722).
• context.models.make_room_in_vram does not (invocation_context.py:743). Its only caller holds the lock from outside (qwen_image_text_encoder.py:547).

Both are public API for custom nodes, so two failures are possible:

1. A node copies the Qwen pattern and wraps offload_from_vram in MODEL_LOAD_LOCK.read_lock(). If a construction queues between the outer and the inner acquire, the thread deadlocks with itself, because the lock is non-reentrant and write-preferring. Every VRAM move and every construction in the process then stalls. Before this PR that pattern was safe.
2. A node calls make_room_in_vram without the lock, as its docstring suggests is fine. That path reaches _move_model_to_ram and load_state_dict(assign=True) too, so it is the same race as #10009.

Suggested fix: take the read lock inside ModelsInterface.make_room_in_vram as well, drop the outer with in the Qwen encoder (and adjust its test), and add "the caller must not hold MODEL_LOAD_LOCK" to both docstrings, as model_load_base.py:46 does for its own contract. The docstrings are published in docs/src/generated/invocation-context.json, so that file needs regenerating.

If you'd rather keep this PR to the three sites, the minimum is the two docstrings plus the regenerated JSON.

Nits

• Cache docstrings. offload_model_from_vram, offload_models_from_vram_except, offload_unlocked_models_from_vram and make_room_in_vram in model_cache.py don't say that the caller must hold the read lock. That gap predates this PR and is why three callers missed it; one sentence per method would close it.
• Offloads under the read lock are not paced. Each of the three sites holds the read lock across the whole offload, including the gc.collect() and empty_cache() in _offload_entries_from_vram.
  • With the default keep_ram_copy_of_weights: true this is short, since tensors are only re-pointed.
  • With false every resident model is copied back, and a queued construction plus every VRAM move on other GPUs wait for that long. I have not measured it.
  • The first lock() pass and the calibration offload already behave this way, so it is not a new kind of stall.
  • The comment on VRAM_MOVE_PASS_BYTES ("bounds any queued construction's wait to roughly one pass") does not hold for these paths. Worth a word in that comment.
• Test run time. Each of the three race tests takes about 0.8 s here, nearly all of it the real gc.collect() calls in the offload path. tests/AGENTS.md asks for gc.freeze() around repeated collections (see test_model_cache_ram_budget.py), or patch gc.collect in the fixture since it is not under test.
• Tests that pass without the fix.
  • test_offloads_on_independent_caches_can_assign_weights_concurrently uses two SimpleNamespace stand-ins and no construction, so the weight assertions cannot fail. What it checks is that readers overlap, which belongs in test_model_load_rw_lock.py if it isn't covered there.
  • test_failed_offload_releases_load_lock_and_preserves_exception and the added assert_load_lock_idle() pairs test the lock's own finally, which a with at the call site cannot break.
• Fixture wiring. offload_race = model_offload.offload_race in each test module is a third way of sharing fixtures. The repo registers shared ones in tests/conftest.py (e.g. lost_races).
• PR description.
  • The ARM64 paragraph ("still being observed") is stale.
  • The "equivalent repository commands" are not the ones that ran.
  • The model-cache scope doesn't match: all production changes are under invokeai/app/.

What I checked

• Every path into _move_model_to_ram from the three sites is inside the read lock (partial and full unload).
• No in-repo caller already holds MODEL_LOAD_LOCK or a cache lock: the seven PiD nodes, _on_before_run_session and the worker loop. The order is global lock before cache lock everywhere, also with several GPUs.
• Exceptions, return values and the empty_cache(force=True) after a successful session-end offload are unchanged; the latter now runs outside the lock.
• Tests on Windows, CPU only:
  • with the fix: 23 passed;
  • with the three production files reverted to base: 3 failed with weights on meta, 20 passed;
  • each of the three files passes on its own;
  • 150 repeated runs under parallel load: no flake.
• Ruff 0.11.2 check and format are clean.
• CI is green except windows-arm64, which fails before pytest on a Chocolatey 504.

Not checked: GPU runs with real models, the hold time with keep_ram_copy_of_weights: false, and custom node packs outside this repo.

[wunianze666-netizen]

Thank you for the detailed review. Addressed in 08998bb.

make_room_in_vram now owns the read lock before the cache lock, and the Qwen caller no longer acquires a nested read lock. Both public node API docstrings state that callers must not already hold the model-load or cache lock; the generated invocation-context JSON was regenerated. Cache method contracts and the unpaced-offload comment now state the relevant lock/pacing boundaries.

The shared race fixture is registered through tests/conftest.py, avoids unrelated gc.collect() calls, and redundant lock-only assertions were removed; reader overlap remains covered in the lock suite. With the final regression fixture, the original implementation failed by leaving weights on meta, and the fixed focused suite passed 72 tests. A separate eight-file related run passed 80 tests with 3 CUDA-dependent skips; the suites overlap and should not be added together. Documentation generation and Ruff 0.11.2 check/format passed.

These runs used Windows, Python 3.12, CPU Torch 2.13 and a Diffusers 0.42.1 overlay; they do not establish physical-GPU behavior, full-model behavior or unpaced offload duration. The previous ARM64 failure was an external Chocolatey FFmpeg 504 before pytest, and the new head's CI is pending.