mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-10-03 09:05:17 -07:00
feat: CUDA backend swap via binary download and restart
Add the ability to download a CUDA-enabled backend binary (~2.4 GB) and swap it in via a backend-only restart, solving the #1 user pain point (19 open 'GPU not detected' issues caused by GitHub's 2 GB asset limit). Backend: - cuda_download.py: download from R2 (primary) or GitHub split-parts (fallback), SHA-256 verification, atomic writes, progress via SSE - 4 new endpoints: GET/POST/DELETE /backend/cuda-*, GET cuda-progress - server.py: --version flag, auto-detect variant from binary name - build_binary.py: --cuda flag for CUDA PyInstaller builds - split_binary.py: split large binaries into <2GB GitHub Release assets - CI workflow for building CUDA binary Tauri: - restart_server command (stop -> wait -> start) - start_server prefers CUDA binary from {data_dir}/backends/ if present - Version mismatch check: runs --version before launching CUDA binary Frontend: - GpuAcceleration component: download, progress, restart, switch, delete - API client + types for CUDA status and management - Platform lifecycle: restartServer() on Tauri/Web - Aggressive 1s health polling during restart for fast reconnection
This commit is contained in:
James Pine
@@ -0,0 +1,200 @@
|
||||
"""
|
||||
CUDA backend binary download, assembly, and verification.
|
||||
|
||||
Downloads split parts of the CUDA-enabled voicebox-server binary from
|
||||
GitHub Releases, reassembles them, verifies integrity via SHA-256,
|
||||
and places the binary in the app's data directory for use on next
|
||||
backend restart.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from .config import get_data_dir
|
||||
from .utils.progress import get_progress_manager
|
||||
from . import __version__
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
GITHUB_RELEASES_URL = "https://github.com/jamiepine/voicebox/releases/download"
|
||||
|
||||
PROGRESS_KEY = "cuda-backend"
|
||||
|
||||
|
||||
def get_backends_dir() -> Path:
|
||||
"""Directory where downloaded backend binaries are stored."""
|
||||
d = get_data_dir() / "backends"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return d
|
||||
|
||||
|
||||
def get_cuda_binary_name() -> str:
|
||||
"""Platform-specific CUDA binary filename."""
|
||||
if sys.platform == "win32":
|
||||
return "voicebox-server-cuda.exe"
|
||||
return "voicebox-server-cuda"
|
||||
|
||||
|
||||
def get_cuda_binary_path() -> Optional[Path]:
|
||||
"""Return path to CUDA binary if it exists."""
|
||||
p = get_backends_dir() / get_cuda_binary_name()
|
||||
if p.exists():
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
def is_cuda_active() -> bool:
|
||||
"""Check if the current process is the CUDA binary.
|
||||
|
||||
The CUDA binary sets this env var on startup (see server.py).
|
||||
"""
|
||||
return os.environ.get("VOICEBOX_BACKEND_VARIANT") == "cuda"
|
||||
|
||||
|
||||
def get_cuda_status() -> dict:
|
||||
"""Get current CUDA backend status for the API."""
|
||||
progress_manager = get_progress_manager()
|
||||
cuda_path = get_cuda_binary_path()
|
||||
progress = progress_manager.get_progress(PROGRESS_KEY)
|
||||
|
||||
return {
|
||||
"available": cuda_path is not None,
|
||||
"active": is_cuda_active(),
|
||||
"binary_path": str(cuda_path) if cuda_path else None,
|
||||
"downloading": progress is not None and progress.get("status") == "downloading",
|
||||
"download_progress": progress,
|
||||
}
|
||||
|
||||
|
||||
async def download_cuda_binary(version: Optional[str] = None):
|
||||
"""Download the CUDA backend binary from GitHub Releases.
|
||||
|
||||
Downloads split parts listed in a manifest file, concatenates them,
|
||||
and verifies the SHA-256 checksum for integrity. Atomic write
|
||||
(temp file -> rename).
|
||||
|
||||
Args:
|
||||
version: Version tag (e.g. "v0.2.0"). Defaults to current app version.
|
||||
"""
|
||||
import httpx
|
||||
|
||||
if version is None:
|
||||
version = f"v{__version__}"
|
||||
|
||||
progress = get_progress_manager()
|
||||
binary_name = get_cuda_binary_name()
|
||||
dest_dir = get_backends_dir()
|
||||
final_path = dest_dir / binary_name
|
||||
temp_path = dest_dir / f"{binary_name}.download"
|
||||
|
||||
# Clean up any leftover partial download
|
||||
if temp_path.exists():
|
||||
temp_path.unlink()
|
||||
|
||||
logger.info(f"Starting CUDA backend download for {version}")
|
||||
progress.update_progress(
|
||||
PROGRESS_KEY, current=0, total=0,
|
||||
filename="Fetching manifest...", status="downloading",
|
||||
)
|
||||
|
||||
base_url = f"{GITHUB_RELEASES_URL}/{version}"
|
||||
stem = Path(binary_name).stem # voicebox-server-cuda
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(follow_redirects=True, timeout=30.0) as client:
|
||||
# Fetch the manifest (list of split part filenames)
|
||||
manifest_url = f"{base_url}/{stem}.manifest"
|
||||
manifest_resp = await client.get(manifest_url)
|
||||
manifest_resp.raise_for_status()
|
||||
parts = [p.strip() for p in manifest_resp.text.strip().splitlines() if p.strip()]
|
||||
|
||||
if not parts:
|
||||
raise ValueError("Empty manifest — no split parts found")
|
||||
|
||||
logger.info(f"Found {len(parts)} split parts to download")
|
||||
|
||||
# Fetch expected checksum (optional — for integrity verification)
|
||||
expected_sha = None
|
||||
try:
|
||||
sha_url = f"{base_url}/{stem}.sha256"
|
||||
sha_resp = await client.get(sha_url)
|
||||
if sha_resp.status_code == 200:
|
||||
# Format: "sha256hex filename\n"
|
||||
expected_sha = sha_resp.text.strip().split()[0]
|
||||
logger.info(f"Expected SHA-256: {expected_sha[:16]}...")
|
||||
except Exception:
|
||||
logger.warning("Could not fetch checksum file — skipping verification")
|
||||
|
||||
# Download and concatenate parts
|
||||
total_downloaded = 0
|
||||
with open(temp_path, "wb") as f:
|
||||
for i, part_name in enumerate(parts):
|
||||
part_url = f"{base_url}/{part_name}"
|
||||
logger.info(f"Downloading part {i + 1}/{len(parts)}: {part_name}")
|
||||
|
||||
async with client.stream("GET", part_url) as response:
|
||||
response.raise_for_status()
|
||||
async for chunk in response.aiter_bytes(chunk_size=1024 * 1024):
|
||||
f.write(chunk)
|
||||
total_downloaded += len(chunk)
|
||||
progress.update_progress(
|
||||
PROGRESS_KEY, current=total_downloaded, total=0,
|
||||
filename=f"Part {i + 1}/{len(parts)}",
|
||||
status="downloading",
|
||||
)
|
||||
|
||||
# Verify integrity if checksum was available
|
||||
if expected_sha:
|
||||
progress.update_progress(
|
||||
PROGRESS_KEY, current=total_downloaded, total=total_downloaded,
|
||||
filename="Verifying integrity...", status="downloading",
|
||||
)
|
||||
sha256 = hashlib.sha256()
|
||||
with open(temp_path, "rb") as f:
|
||||
while True:
|
||||
chunk = f.read(1024 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
sha256.update(chunk)
|
||||
|
||||
actual = sha256.hexdigest()
|
||||
if actual != expected_sha:
|
||||
raise ValueError(
|
||||
f"Integrity check failed: expected {expected_sha[:16]}..., "
|
||||
f"got {actual[:16]}..."
|
||||
)
|
||||
logger.info(f"Integrity verified: {actual[:16]}...")
|
||||
|
||||
# Atomic move into place
|
||||
if final_path.exists():
|
||||
final_path.unlink()
|
||||
temp_path.rename(final_path)
|
||||
|
||||
# Make executable on Unix
|
||||
if sys.platform != "win32":
|
||||
final_path.chmod(0o755)
|
||||
|
||||
logger.info(f"CUDA backend downloaded to {final_path}")
|
||||
progress.mark_complete(PROGRESS_KEY)
|
||||
|
||||
except Exception as e:
|
||||
# Clean up on failure
|
||||
if temp_path.exists():
|
||||
temp_path.unlink()
|
||||
logger.error(f"CUDA backend download failed: {e}")
|
||||
progress.mark_error(PROGRESS_KEY, str(e))
|
||||
raise
|
||||
|
||||
|
||||
async def delete_cuda_binary() -> bool:
|
||||
"""Delete the downloaded CUDA binary. Returns True if deleted."""
|
||||
path = get_cuda_binary_path()
|
||||
if path and path.exists():
|
||||
path.unlink()
|
||||
logger.info(f"Deleted CUDA binary: {path}")
|
||||
return True
|
||||
return False
|
||||
Reference in New Issue
Block a user