mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-15 12:50:42 -07:00
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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user