fix: add asyncio.Lock to prevent concurrent CUDA downloads (#428)

* fix: add asyncio.Lock to prevent concurrent CUDA downloads

The startup auto-update task and the manual download endpoint can both
invoke download_cuda_binary() concurrently. Without mutual exclusion,
both coroutines write to the same temp file path, corrupting the
download. The progress-manager status check is a TOCTOU race because
the status is not set until after several synchronous checks complete.

Add a module-level asyncio.Lock acquired at the top of
download_cuda_binary() so only one download can proceed at a time.

* fix: fast-reject duplicate CUDA download when lock is held

Address CodeRabbit review feedback: check _download_lock.locked()
before awaiting the lock so concurrent callers return immediately
instead of queueing behind the first download. This prevents the
route handler from returning "started" to multiple callers when only
one download actually proceeds.
This commit is contained in:
Junghwan
2026-04-16 01:51:21 -07:00
committed by GitHub
parent c9d8142a78
commit be7c0cec12
+16
View File
@@ -11,6 +11,7 @@ Both archives are extracted into {data_dir}/backends/cuda/ which forms the
complete PyInstaller --onedir directory structure that torch expects.
"""
import asyncio
import hashlib
import json
import logging
@@ -34,6 +35,12 @@ PROGRESS_KEY = "cuda-backend"
# CUDA toolkit version or torch's CUDA dependency changes (e.g. cu126 -> cu128).
CUDA_LIBS_VERSION = "cu128-v1"
# Prevents concurrent download_cuda_binary() calls from racing on the same
# temp file. The auto-update background task and the manual HTTP endpoint
# can both invoke download_cuda_binary(); without this lock the progress-
# manager status check is a TOCTOU race.
_download_lock = asyncio.Lock()
def get_backends_dir() -> Path:
"""Directory where downloaded backend binaries are stored."""
@@ -241,6 +248,15 @@ async def download_cuda_binary(version: Optional[str] = None):
Args:
version: Version tag (e.g. "v0.3.0"). Defaults to current app version.
"""
if _download_lock.locked():
logger.info("CUDA download already in progress, skipping duplicate request")
return
async with _download_lock:
await _download_cuda_binary_locked(version)
async def _download_cuda_binary_locked(version: Optional[str] = None):
"""Inner implementation of download_cuda_binary, called under _download_lock."""
import httpx
if version is None: