This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,251 @@
|
||||
"""
|
||||
Package the PyInstaller --onedir ROCm build into two archives.
|
||||
|
||||
Takes the PyInstaller --onedir output directory and splits it into:
|
||||
1. voicebox-server-rocm.tar.gz — server core (exe + non-AMD deps)
|
||||
2. rocm-libs-{version}.tar.gz — AMD/ROCm runtime libraries only
|
||||
3. rocm-libs.json — version manifest for the ROCm libs
|
||||
|
||||
Mirrors scripts/package_cuda.py. The split lets the server core re-download on
|
||||
every app update while the much larger ROCm runtime stays cached until the
|
||||
toolkit version bumps.
|
||||
|
||||
Usage:
|
||||
python scripts/package_rocm.py backend/dist/voicebox-server-rocm/
|
||||
python scripts/package_rocm.py backend/dist/voicebox-server-rocm/ --output release-assets/
|
||||
python scripts/package_rocm.py backend/dist/voicebox-server-rocm/ --rocm-libs-version rocm7.2-v1
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import sys
|
||||
import tarfile
|
||||
from pathlib import Path
|
||||
|
||||
# DLL/.so name prefixes that identify AMD ROCm/HIP runtime libraries. They may
|
||||
# sit in torch/lib/ (torch's bundled HIP runtime) or inside the bundled ROCm SDK
|
||||
# packages. Matched case-insensitively against the file's base name.
|
||||
ROCM_DLL_PREFIXES = (
|
||||
"amdhip",
|
||||
"amd_comgr",
|
||||
"amdocl",
|
||||
"hiprtc",
|
||||
"hipblaslt",
|
||||
"hipblas",
|
||||
"hipfft",
|
||||
"hiprand",
|
||||
"hipsolver",
|
||||
"hipsparse",
|
||||
"hip",
|
||||
"rocblas",
|
||||
"rocfft",
|
||||
"rocrand",
|
||||
"rocsolver",
|
||||
"rocsparse",
|
||||
"rocprofiler",
|
||||
"roctracer",
|
||||
"roctx",
|
||||
"rocm_smi",
|
||||
"miopen",
|
||||
"rccl",
|
||||
"hsa-runtime",
|
||||
"hsa",
|
||||
)
|
||||
|
||||
# Directory markers for the bundled ROCm SDK runtime packages. Everything under
|
||||
# these trees except Python sources (the pure-python rocm_sdk glue) is part of
|
||||
# the runtime payload — this is where rocBLAS Tensile data and MIOpen kernel
|
||||
# databases live, which dominate the download size.
|
||||
ROCM_LIB_DIR_MARKERS = (
|
||||
"_rocm_sdk_core",
|
||||
"_rocm_sdk_libraries_custom",
|
||||
"rocm_sdk_core",
|
||||
"rocm_sdk_libraries_custom",
|
||||
)
|
||||
|
||||
# Heavy native/data extensions shipped by the ROCm runtime (HIP fat binaries,
|
||||
# rocBLAS Tensile data, MIOpen kernel DBs).
|
||||
ROCM_LIB_EXTS = (".dll", ".so", ".dat", ".db", ".kdb", ".hsaco", ".co", ".bc")
|
||||
|
||||
# Python sources stay in the server core so the rocm_sdk import glue remains
|
||||
# alongside the exe. (Both archives extract into backends/rocm/, so this only
|
||||
# affects which archive carries the file, not runtime resolution.)
|
||||
_PYTHON_EXTS = (".py", ".pyc", ".pyi")
|
||||
|
||||
|
||||
def is_rocm_file(rel_path: str) -> bool:
|
||||
"""Check if a relative path belongs to the AMD ROCm runtime libraries.
|
||||
|
||||
Identifies large ROCm/HIP runtime DLLs and the SDK runtime payload
|
||||
(kernel databases, Tensile data) regardless of where PyInstaller placed
|
||||
them, while keeping pure-python glue in the server core.
|
||||
"""
|
||||
rel_lower = rel_path.lower().replace("\\", "/")
|
||||
name = rel_lower.rsplit("/", 1)[-1]
|
||||
|
||||
# Never split out Python sources / stubs.
|
||||
if name.endswith(_PYTHON_EXTS):
|
||||
return False
|
||||
|
||||
# Native payload inside the bundled ROCm SDK package trees.
|
||||
if any(marker in rel_lower for marker in ROCM_LIB_DIR_MARKERS):
|
||||
if name.endswith(ROCM_LIB_EXTS):
|
||||
return True
|
||||
|
||||
# ROCm/HIP DLLs/shared objects anywhere (e.g. _internal/torch/lib/amdhip64.dll).
|
||||
if name.endswith((".dll", ".so")):
|
||||
name_no_ext = name.rsplit(".", 1)[0]
|
||||
for prefix in ROCM_DLL_PREFIXES:
|
||||
if name_no_ext.startswith(prefix):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def sha256_file(path: Path) -> str:
|
||||
"""Compute SHA-256 hex digest of a file."""
|
||||
h = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
while True:
|
||||
chunk = f.read(1024 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
h.update(chunk)
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def package(
|
||||
onedir_path: Path,
|
||||
output_dir: Path,
|
||||
rocm_libs_version: str,
|
||||
torch_compat: str,
|
||||
):
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Collect all files in the onedir output, split into core vs rocm.
|
||||
core_files = []
|
||||
rocm_files = []
|
||||
|
||||
for item in sorted(onedir_path.rglob("*")):
|
||||
if item.is_dir():
|
||||
continue
|
||||
rel = item.relative_to(onedir_path)
|
||||
rel_str = str(rel)
|
||||
if is_rocm_file(rel_str):
|
||||
rocm_files.append((rel_str, item))
|
||||
else:
|
||||
core_files.append((rel_str, item))
|
||||
|
||||
core_size = sum(f.stat().st_size for _, f in core_files)
|
||||
rocm_size = sum(f.stat().st_size for _, f in rocm_files)
|
||||
|
||||
print(f"Input directory: {onedir_path}")
|
||||
print(f"Core files: {len(core_files)} ({core_size / (1024**2):.1f} MB)")
|
||||
print(f"ROCm files: {len(rocm_files)} ({rocm_size / (1024**2):.1f} MB)")
|
||||
|
||||
if not rocm_files:
|
||||
print(
|
||||
f"ERROR: No ROCm files found in {onedir_path}. "
|
||||
"Refusing to create an empty ROCm libs archive.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
print(
|
||||
"Make sure you built with --rocm and the ROCm SDK packages are present. "
|
||||
"If the layout differs, adjust ROCM_DLL_PREFIXES / ROCM_LIB_DIR_MARKERS.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
# Create server core archive. Files are stored relative to the archive root
|
||||
# (no parent prefix) so extracting to backends/rocm/ lands at the right level.
|
||||
server_archive = output_dir / "voicebox-server-rocm.tar.gz"
|
||||
print(f"\nCreating server core archive: {server_archive.name}")
|
||||
with tarfile.open(server_archive, "w:gz") as tar:
|
||||
for rel_str, full_path in core_files:
|
||||
tar.add(full_path, arcname=rel_str)
|
||||
server_sha = sha256_file(server_archive)
|
||||
(output_dir / "voicebox-server-rocm.tar.gz.sha256").write_text(
|
||||
f"{server_sha} voicebox-server-rocm.tar.gz\n"
|
||||
)
|
||||
print(f" Size: {server_archive.stat().st_size / (1024**2):.1f} MB")
|
||||
print(f" SHA-256: {server_sha[:16]}...")
|
||||
|
||||
# Create ROCm libs archive.
|
||||
rocm_libs_archive = output_dir / f"rocm-libs-{rocm_libs_version}.tar.gz"
|
||||
print(f"\nCreating ROCm libs archive: {rocm_libs_archive.name}")
|
||||
with tarfile.open(rocm_libs_archive, "w:gz") as tar:
|
||||
for rel_str, full_path in rocm_files:
|
||||
tar.add(full_path, arcname=rel_str)
|
||||
rocm_sha = sha256_file(rocm_libs_archive)
|
||||
(output_dir / f"rocm-libs-{rocm_libs_version}.tar.gz.sha256").write_text(
|
||||
f"{rocm_sha} rocm-libs-{rocm_libs_version}.tar.gz\n"
|
||||
)
|
||||
print(f" Size: {rocm_libs_archive.stat().st_size / (1024**2):.1f} MB")
|
||||
print(f" SHA-256: {rocm_sha[:16]}...")
|
||||
|
||||
# Write rocm-libs.json manifest.
|
||||
manifest = {
|
||||
"version": rocm_libs_version,
|
||||
"torch_compat": torch_compat,
|
||||
"archive": rocm_libs_archive.name,
|
||||
"sha256": rocm_sha,
|
||||
}
|
||||
manifest_path = output_dir / "rocm-libs.json"
|
||||
manifest_path.write_text(json.dumps(manifest, indent=2) + "\n")
|
||||
print(f"\nManifest: {manifest_path.name}")
|
||||
print(json.dumps(manifest, indent=2))
|
||||
|
||||
# Summary
|
||||
total_input = core_size + rocm_size
|
||||
total_output = server_archive.stat().st_size + rocm_libs_archive.stat().st_size
|
||||
print(f"\nTotal input: {total_input / (1024**3):.2f} GB")
|
||||
print(f"Total output: {total_output / (1024**3):.2f} GB (compressed)")
|
||||
print(
|
||||
f"Server core: {server_archive.stat().st_size / (1024**2):.1f} MB (redownloaded on app update)"
|
||||
)
|
||||
print(
|
||||
f"ROCm libs: {rocm_libs_archive.stat().st_size / (1024**2):.1f} MB (cached until ROCm toolkit bump)"
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Package PyInstaller --onedir ROCm build into server + ROCm libs archives"
|
||||
)
|
||||
parser.add_argument(
|
||||
"input",
|
||||
type=Path,
|
||||
help="Path to PyInstaller --onedir output directory (e.g. backend/dist/voicebox-server-rocm/)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Output directory for archives (default: same as input parent)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rocm-libs-version",
|
||||
type=str,
|
||||
default="rocm7.2-v1",
|
||||
help="Version string for the ROCm libs archive (default: rocm7.2-v1)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--torch-compat",
|
||||
type=str,
|
||||
default=">=2.9.0,<2.10.0",
|
||||
help="Torch version compatibility range (default: >=2.9.0,<2.10.0)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.input.is_dir():
|
||||
print(f"Error: {args.input} is not a directory", file=sys.stderr)
|
||||
print("Expected a PyInstaller --onedir output directory.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
output_dir = args.output or args.input.parent
|
||||
package(args.input, output_dir, args.rocm_libs_version, args.torch_compat)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user