generated from Labyricorn/labyricorn-project-template
Initial commit (forked from jamiepine/voicebox)
This commit is contained in:
@@ -0,0 +1,467 @@
|
||||
"""
|
||||
ROCm backend download, assembly, and verification.
|
||||
|
||||
Downloads two archives from GitHub Releases:
|
||||
1. Server core (voicebox-server-rocm.tar.gz) — the exe + non-AMD deps,
|
||||
versioned with the app.
|
||||
2. ROCm libs (rocm-libs-{version}.tar.gz) — AMD runtime libraries,
|
||||
versioned independently (only redownloaded on ROCm toolkit bump).
|
||||
|
||||
Both archives are extracted into {data_dir}/backends/rocm/ which forms the
|
||||
complete PyInstaller --onedir directory structure that torch expects.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tarfile
|
||||
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 = "rocm-backend"
|
||||
|
||||
# The current expected ROCm libs version. Bump this when we change the
|
||||
# ROCm toolkit version or torch's ROCm dependency changes (e.g. rocm7.2 -> rocm7.4).
|
||||
ROCM_LIBS_VERSION = "rocm7.2-v1"
|
||||
|
||||
# Prevents concurrent download_rocm_binary() calls from racing on the same
|
||||
# temp file. The auto-update background task and the manual HTTP endpoint
|
||||
# can both invoke download_rocm_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."""
|
||||
d = get_data_dir() / "backends"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return d
|
||||
|
||||
|
||||
def get_rocm_dir() -> Path:
|
||||
"""Directory where the ROCm backend (onedir) is extracted."""
|
||||
d = get_backends_dir() / "rocm"
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return d
|
||||
|
||||
|
||||
def get_rocm_exe_name() -> str:
|
||||
"""Platform-specific ROCm executable filename."""
|
||||
if sys.platform == "win32":
|
||||
return "voicebox-server-rocm.exe"
|
||||
return "voicebox-server-rocm"
|
||||
|
||||
|
||||
def get_rocm_binary_path() -> Optional[Path]:
|
||||
"""Return path to the ROCm executable if it exists inside the onedir."""
|
||||
p = get_rocm_dir() / get_rocm_exe_name()
|
||||
if p.exists():
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
def get_rocm_libs_manifest_path() -> Path:
|
||||
"""Path to the rocm-libs.json manifest inside the ROCm dir."""
|
||||
return get_rocm_dir() / "rocm-libs.json"
|
||||
|
||||
|
||||
def get_installed_rocm_libs_version() -> Optional[str]:
|
||||
"""Read the installed ROCm libs version from rocm-libs.json, or None."""
|
||||
manifest_path = get_rocm_libs_manifest_path()
|
||||
if not manifest_path.exists():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(manifest_path.read_text())
|
||||
return data.get("version")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not read rocm-libs.json: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def is_rocm_active() -> bool:
|
||||
"""Check if the current process is the ROCm binary.
|
||||
|
||||
The ROCm binary sets this env var on startup (see server.py).
|
||||
"""
|
||||
return os.environ.get("VOICEBOX_BACKEND_VARIANT") == "rocm"
|
||||
|
||||
|
||||
def get_rocm_status() -> dict:
|
||||
"""Get current ROCm backend status for the API."""
|
||||
progress_manager = get_progress_manager()
|
||||
rocm_path = get_rocm_binary_path()
|
||||
progress = progress_manager.get_progress(PROGRESS_KEY)
|
||||
rocm_libs_version = get_installed_rocm_libs_version()
|
||||
|
||||
return {
|
||||
"available": rocm_path is not None,
|
||||
"active": is_rocm_active(),
|
||||
"binary_path": str(rocm_path) if rocm_path else None,
|
||||
"rocm_libs_version": rocm_libs_version,
|
||||
"downloading": progress is not None and progress.get("status") == "downloading",
|
||||
"download_progress": progress,
|
||||
}
|
||||
|
||||
|
||||
def _needs_server_download(version: Optional[str] = None) -> bool:
|
||||
"""Check if the server core archive needs to be (re)downloaded."""
|
||||
rocm_path = get_rocm_binary_path()
|
||||
if not rocm_path:
|
||||
return True
|
||||
# Check if the binary version matches the expected app version
|
||||
installed = get_rocm_binary_version()
|
||||
expected = version or __version__
|
||||
if expected.startswith("v"):
|
||||
expected = expected[1:]
|
||||
return installed != expected
|
||||
|
||||
|
||||
def _needs_rocm_libs_download() -> bool:
|
||||
"""Check if the ROCm libs archive needs to be (re)downloaded."""
|
||||
installed = get_installed_rocm_libs_version()
|
||||
if installed is None:
|
||||
return True
|
||||
return installed != ROCM_LIBS_VERSION
|
||||
|
||||
|
||||
async def _download_and_extract_archive(
|
||||
client,
|
||||
url: str,
|
||||
sha256_url: Optional[str],
|
||||
dest_dir: Path,
|
||||
label: str,
|
||||
progress_offset: int,
|
||||
total_size: int,
|
||||
):
|
||||
"""Download a .tar.gz archive and extract it into dest_dir.
|
||||
|
||||
Args:
|
||||
client: httpx.AsyncClient
|
||||
url: URL of the .tar.gz archive
|
||||
sha256_url: URL of the .sha256 checksum file (optional)
|
||||
dest_dir: Directory to extract into
|
||||
label: Human-readable label for progress updates
|
||||
progress_offset: Byte offset for progress reporting (when downloading
|
||||
multiple archives sequentially)
|
||||
total_size: Total bytes across all downloads (for progress bar)
|
||||
"""
|
||||
progress = get_progress_manager()
|
||||
temp_path = dest_dir / f".download-{label.replace(' ', '-')}.tmp"
|
||||
|
||||
# Clean up leftover partial download
|
||||
if temp_path.exists():
|
||||
temp_path.unlink()
|
||||
|
||||
# Fetch expected checksum (fail-fast: never extract an unverified archive)
|
||||
expected_sha = None
|
||||
if sha256_url:
|
||||
try:
|
||||
sha_resp = await client.get(sha256_url)
|
||||
sha_resp.raise_for_status()
|
||||
expected_sha = sha_resp.text.strip().split()[0]
|
||||
logger.info(f"{label}: expected SHA-256: {expected_sha[:16]}...")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"{label}: failed to fetch checksum from {sha256_url}") from e
|
||||
|
||||
# Stream download, verify, and extract — always clean up temp file
|
||||
downloaded = 0
|
||||
try:
|
||||
async with client.stream("GET", url) as response:
|
||||
response.raise_for_status()
|
||||
with open(temp_path, "wb") as f:
|
||||
async for chunk in response.aiter_bytes(chunk_size=1024 * 1024):
|
||||
f.write(chunk)
|
||||
downloaded += len(chunk)
|
||||
progress.update_progress(
|
||||
PROGRESS_KEY,
|
||||
current=progress_offset + downloaded,
|
||||
total=total_size,
|
||||
filename=f"Downloading {label}",
|
||||
status="downloading",
|
||||
)
|
||||
|
||||
# Verify integrity
|
||||
if expected_sha:
|
||||
progress.update_progress(
|
||||
PROGRESS_KEY,
|
||||
current=progress_offset + downloaded,
|
||||
total=total_size,
|
||||
filename=f"Verifying {label}...",
|
||||
status="downloading",
|
||||
)
|
||||
sha256 = hashlib.sha256()
|
||||
with open(temp_path, "rb") as f:
|
||||
while True:
|
||||
data = f.read(1024 * 1024)
|
||||
if not data:
|
||||
break
|
||||
sha256.update(data)
|
||||
actual = sha256.hexdigest()
|
||||
if actual != expected_sha:
|
||||
raise ValueError(
|
||||
f"{label} integrity check failed: expected {expected_sha[:16]}..., got {actual[:16]}..."
|
||||
)
|
||||
logger.info(f"{label}: integrity verified")
|
||||
|
||||
# Extract (use data filter for path traversal protection on Python 3.12+)
|
||||
progress.update_progress(
|
||||
PROGRESS_KEY,
|
||||
current=progress_offset + downloaded,
|
||||
total=total_size,
|
||||
filename=f"Extracting {label}...",
|
||||
status="downloading",
|
||||
)
|
||||
with tarfile.open(temp_path, "r:gz") as tar:
|
||||
tar.extractall(path=dest_dir, filter="data")
|
||||
|
||||
logger.info(f"{label}: extracted to {dest_dir}")
|
||||
finally:
|
||||
if temp_path.exists():
|
||||
temp_path.unlink()
|
||||
return downloaded
|
||||
|
||||
|
||||
async def download_rocm_binary(version: Optional[str] = None):
|
||||
"""Download the ROCm backend (server core + ROCm libs if needed).
|
||||
|
||||
Downloads both archives from GitHub Releases, extracts them into
|
||||
{data_dir}/backends/rocm/, and writes the rocm-libs.json manifest.
|
||||
|
||||
Only downloads what's needed:
|
||||
- Server core: always redownloaded (versioned with app)
|
||||
- ROCm libs: only if missing or version mismatch
|
||||
|
||||
Args:
|
||||
version: Version tag (e.g. "v0.3.0"). Defaults to current app version.
|
||||
"""
|
||||
if _download_lock.locked():
|
||||
logger.info("ROCm download already in progress, skipping duplicate request")
|
||||
return
|
||||
async with _download_lock:
|
||||
await _download_rocm_binary_locked(version)
|
||||
|
||||
|
||||
async def _download_rocm_binary_locked(version: Optional[str] = None):
|
||||
"""Inner implementation of download_rocm_binary, called under _download_lock."""
|
||||
import httpx
|
||||
|
||||
if version is None:
|
||||
version = f"v{__version__}"
|
||||
|
||||
progress = get_progress_manager()
|
||||
rocm_dir = get_rocm_dir()
|
||||
|
||||
need_server = _needs_server_download(version)
|
||||
need_libs = _needs_rocm_libs_download()
|
||||
|
||||
if not need_server and not need_libs:
|
||||
logger.info("ROCm backend is up to date, nothing to download")
|
||||
return
|
||||
|
||||
logger.info(
|
||||
f"Starting ROCm backend download for {version} "
|
||||
f"(server={'yes' if need_server else 'cached'}, "
|
||||
f"libs={'yes' if need_libs else 'cached'})"
|
||||
)
|
||||
progress.update_progress(
|
||||
PROGRESS_KEY,
|
||||
current=0,
|
||||
total=0,
|
||||
filename="Preparing download...",
|
||||
status="downloading",
|
||||
)
|
||||
|
||||
# Server core and libs archive are both published under the app-version
|
||||
# release tag; the libs content version is encoded in the filename only.
|
||||
server_base_url = f"{GITHUB_RELEASES_URL}/{version}"
|
||||
libs_base_url = server_base_url
|
||||
server_archive = "voicebox-server-rocm.tar.gz"
|
||||
libs_archive = f"rocm-libs-{ROCM_LIBS_VERSION}.tar.gz"
|
||||
|
||||
# Always stage when any download is needed, then atomically rename over
|
||||
# rocm_dir on success. This prevents a failed mid-extraction from leaving
|
||||
# rocm_dir in a partially-installed state that still passes the
|
||||
# get_rocm_binary_path() existence check. Existing files are pre-copied
|
||||
# into staging so partial updates (e.g. libs-only or server-only) preserve
|
||||
# whatever isn't being re-downloaded.
|
||||
use_staging = need_server or need_libs
|
||||
staging_dir = get_backends_dir() / "rocm-staging"
|
||||
|
||||
if use_staging:
|
||||
if staging_dir.exists():
|
||||
shutil.rmtree(staging_dir)
|
||||
staging_dir.mkdir(parents=True, exist_ok=True)
|
||||
# Preserve existing files (server or libs) that don't need re-downloading.
|
||||
# Extracted archives will overwrite only what we actually download.
|
||||
if rocm_dir.exists():
|
||||
shutil.copytree(rocm_dir, staging_dir, dirs_exist_ok=True)
|
||||
extract_dir = staging_dir
|
||||
else:
|
||||
extract_dir = rocm_dir
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(follow_redirects=True, timeout=30.0) as client:
|
||||
# Estimate total download size
|
||||
total_size = 0
|
||||
if need_server:
|
||||
try:
|
||||
head = await client.head(f"{server_base_url}/{server_archive}")
|
||||
total_size += int(head.headers.get("content-length", 0))
|
||||
except Exception:
|
||||
pass
|
||||
if need_libs:
|
||||
try:
|
||||
head = await client.head(f"{libs_base_url}/{libs_archive}")
|
||||
total_size += int(head.headers.get("content-length", 0))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger.info(f"Total download size: {total_size / 1024 / 1024:.1f} MB")
|
||||
|
||||
offset = 0
|
||||
|
||||
# Download server core
|
||||
if need_server:
|
||||
server_downloaded = await _download_and_extract_archive(
|
||||
client,
|
||||
url=f"{server_base_url}/{server_archive}",
|
||||
sha256_url=f"{server_base_url}/{server_archive}.sha256",
|
||||
dest_dir=extract_dir,
|
||||
label="ROCm server",
|
||||
progress_offset=offset,
|
||||
total_size=total_size,
|
||||
)
|
||||
offset += server_downloaded
|
||||
|
||||
# Make executable on Unix
|
||||
exe_path = extract_dir / get_rocm_exe_name()
|
||||
if sys.platform != "win32" and exe_path.exists():
|
||||
exe_path.chmod(0o755)
|
||||
|
||||
# Download ROCm libs
|
||||
if need_libs:
|
||||
await _download_and_extract_archive(
|
||||
client,
|
||||
url=f"{libs_base_url}/{libs_archive}",
|
||||
sha256_url=f"{libs_base_url}/{libs_archive}.sha256",
|
||||
dest_dir=extract_dir,
|
||||
label="ROCm libraries",
|
||||
progress_offset=offset,
|
||||
total_size=total_size,
|
||||
)
|
||||
|
||||
# Write local rocm-libs.json manifest
|
||||
manifest = {"version": ROCM_LIBS_VERSION}
|
||||
(extract_dir / "rocm-libs.json").write_text(json.dumps(manifest, indent=2) + "\n")
|
||||
|
||||
# Atomic swap: replace rocm_dir with the fully-extracted staging dir
|
||||
if use_staging:
|
||||
backup_dir = get_backends_dir() / "rocm-backup"
|
||||
if backup_dir.exists():
|
||||
shutil.rmtree(backup_dir)
|
||||
if rocm_dir.exists():
|
||||
rocm_dir.rename(backup_dir)
|
||||
try:
|
||||
staging_dir.rename(rocm_dir)
|
||||
except Exception:
|
||||
if backup_dir.exists() and not rocm_dir.exists():
|
||||
backup_dir.rename(rocm_dir)
|
||||
raise
|
||||
else:
|
||||
if backup_dir.exists():
|
||||
shutil.rmtree(backup_dir)
|
||||
|
||||
logger.info(f"ROCm backend ready at {rocm_dir}")
|
||||
progress.mark_complete(PROGRESS_KEY)
|
||||
|
||||
except Exception as e:
|
||||
if use_staging and staging_dir.exists():
|
||||
shutil.rmtree(staging_dir)
|
||||
logger.error(f"ROCm backend download failed: {e}")
|
||||
progress.mark_error(PROGRESS_KEY, str(e))
|
||||
raise
|
||||
|
||||
|
||||
def get_rocm_binary_version() -> Optional[str]:
|
||||
"""Get the version of the installed ROCm binary, or None if not installed."""
|
||||
import subprocess
|
||||
|
||||
rocm_path = get_rocm_binary_path()
|
||||
if not rocm_path:
|
||||
return None
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[str(rocm_path), "--version"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
cwd=str(rocm_path.parent), # Run from the onedir directory
|
||||
)
|
||||
# Output format: "voicebox-server 0.3.0"
|
||||
for line in result.stdout.strip().splitlines():
|
||||
if "voicebox-server" in line:
|
||||
return line.split()[-1]
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not get ROCm binary version: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def check_and_update_rocm_binary():
|
||||
"""Check if the ROCm binary is outdated and auto-download if so.
|
||||
|
||||
Called on server startup. Checks both server version and ROCm libs
|
||||
version. Downloads only what's needed.
|
||||
"""
|
||||
rocm_path = get_rocm_binary_path()
|
||||
if not rocm_path:
|
||||
return # No ROCm binary installed, nothing to update
|
||||
|
||||
if is_rocm_active():
|
||||
logger.info("ROCm backend is active; skipping auto-update to avoid replacing the running backend")
|
||||
return
|
||||
|
||||
need_server = _needs_server_download()
|
||||
need_libs = _needs_rocm_libs_download()
|
||||
|
||||
if not need_server and not need_libs:
|
||||
logger.info(f"ROCm binary is up to date (server=v{__version__}, libs={get_installed_rocm_libs_version()})")
|
||||
return
|
||||
|
||||
reasons = []
|
||||
if need_server:
|
||||
rocm_version = get_rocm_binary_version()
|
||||
reasons.append(f"server v{rocm_version} != v{__version__}")
|
||||
if need_libs:
|
||||
installed_libs = get_installed_rocm_libs_version()
|
||||
reasons.append(f"libs {installed_libs} != {ROCM_LIBS_VERSION}")
|
||||
|
||||
logger.info(f"ROCm backend needs update ({', '.join(reasons)}). Auto-downloading...")
|
||||
|
||||
try:
|
||||
await download_rocm_binary()
|
||||
except Exception as e:
|
||||
logger.error(f"Auto-update of ROCm binary failed: {e}")
|
||||
|
||||
|
||||
async def delete_rocm_binary() -> bool:
|
||||
"""Delete the downloaded ROCm backend directory. Returns True if deleted."""
|
||||
import shutil
|
||||
|
||||
rocm_dir = get_rocm_dir()
|
||||
if rocm_dir.exists() and any(rocm_dir.iterdir()):
|
||||
shutil.rmtree(rocm_dir)
|
||||
logger.info(f"Deleted ROCm backend directory: {rocm_dir}")
|
||||
return True
|
||||
return False
|
||||
Reference in New Issue
Block a user