Skip to content

Race condition in ModelInstallService._restore_incomplete_installs causes duplicate download enqueue #9141

Description

@lstein

Summary

tests/app/services/model_install/test_model_install.py::test_simple_download is flaky on main (sometimes 1 failure, sometimes more). The root cause is a TOCTOU race between the background "restore incomplete installs" thread and the foreground import_model path.

Reproduction

pytest tests/app/services/model_install/test_model_install.py::test_simple_download

Repro rate is ~30–50% on a fast machine. Not consistent.

Diagnostic

Captured failing log (key lines):

INFO --> Queueing model install: https://www.test.foo/download/test_embedding.safetensors (1 file)
INFO --> Queueing model install: https://www.test.foo/download/test_embedding.safetensors (1 file)   ← duplicate
INFO --> Finished restoring incomplete installs
INFO --> File download started: https://www.test.foo/download/test_embedding.safetensors
INFO --> File download started: https://www.test.foo/download/test_embedding.safetensors            ← duplicate
INFO --> Download complete
ERROR --> FileNotFoundError: '.../test_embedding.safetensors.downloading' -> '.../test_embedding.safetensors'

The same source is enqueued twice. The second download then fails because the first one already renamed .downloading → final filename.

Root cause

invokeai/app/services/model_install/model_install_default.py:

def _restore_incomplete_installs(self) -> None:
    path = self._app_config.models_path
    seen_sources: set[str] = set()
    with self._lock:                                                                       # acquired …
        active_sources = {str(j.source) for j in self._install_jobs if not j.in_terminal_state}
        active_sources.update(str(j.source) for j in self._download_cache.values() if not j.in_terminal_state)
                                                                                           # … released here
    for tmpdir in path.glob(f"{TMPDIR_PREFIX}*"):                                          # race window opens
        ...
        if source_str in active_sources:                                                   # stale snapshot
            continue
        ...
        self._install_jobs.append(job)
        self._resume_remote_download(job)                                                  # duplicate queue

The lock only protects the active_sources snapshot. The for-loop processes the on-disk tmpdirs against that stale snapshot. Race window:

  1. Service init kicks off _restore_incomplete_installs_async in a background thread.
  2. Restore thread acquires lock, snapshots active_sources = {}, releases.
  3. Test thread calls import_model(URL) → creates tmpdir + marker → appends job → queues download.
  4. Restore thread enters the for-loop with the stale empty snapshot, sees the tmpdir created in step 3, sees source_str not in active_sources, enqueues it a second time.

If step 3 happens before step 2, the test passes. If after, it fails.

Suggested fix

Re-check membership inside the lock right before appending, instead of trusting a snapshot:

for tmpdir in path.glob(f"{TMPDIR_PREFIX}*"):
    ...
    with self._lock:
        if any(str(j.source) == source_str for j in self._install_jobs if not j.in_terminal_state):
            continue
        if any(str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state):
            continue
        self._install_jobs.append(job)
    if job.paused:
        continue
    ...

This makes the check-and-append atomic. Holding the lock across the entire for-loop is also correct but blocks import_model calls for the duration of the restore.

Environment

  • Branch: main
  • Python 3.12.12

Activity

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

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions