126 lines
4.4 KiB
Bash
Executable File
126 lines
4.4 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
# Build TorchCodec against the ROCm PyTorch installation. The PyPI wheel is
|
|
# built against upstream PyTorch and is not ABI-compatible with ROCm builds.
|
|
|
|
set -euo pipefail
|
|
|
|
TORCHCODEC_REPO="${TORCHCODEC_REPO:-https://github.com/pytorch/torchcodec.git}"
|
|
# Keep this pinned for ROCm/PyTorch ABI compatibility. When rebasing the ROCm
|
|
# base image or vLLM, revalidate this version against the active torch build.
|
|
TORCHCODEC_BRANCH="${TORCHCODEC_BRANCH:-v0.10.0}"
|
|
TORCHCODEC_WHEEL_CACHE="${TORCHCODEC_WHEEL_CACHE:-/root/.cache/torchcodec-wheels}"
|
|
|
|
echo "=== TorchCodec ROCm installation ==="
|
|
|
|
verify_torchcodec() {
|
|
python3 -c "from torchcodec.decoders import VideoDecoder"
|
|
}
|
|
|
|
if verify_torchcodec >/dev/null 2>&1; then
|
|
echo "TorchCodec is already installed and working."
|
|
exit 0
|
|
fi
|
|
|
|
install_system_deps() {
|
|
if ! command -v apt-get >/dev/null 2>&1; then
|
|
echo "TorchCodec build dependencies are missing and apt-get is unavailable."
|
|
return 1
|
|
fi
|
|
apt-get update
|
|
apt-get install -y --no-install-recommends \
|
|
pkg-config \
|
|
ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
|
|
libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev
|
|
}
|
|
|
|
if ! command -v pkg-config >/dev/null 2>&1 || \
|
|
! pkg-config --exists libavcodec libavformat libavutil libswscale \
|
|
libavdevice libavfilter libswresample 2>/dev/null; then
|
|
echo "Installing missing FFmpeg development dependencies..."
|
|
install_system_deps
|
|
fi
|
|
|
|
if ! python3 -c "import packaging, pybind11" >/dev/null 2>&1; then
|
|
python3 -m pip install packaging pybind11 setuptools wheel
|
|
fi
|
|
|
|
PYBIND11_DIR=$(python3 -c "import pybind11; print(pybind11.get_cmake_dir())")
|
|
export pybind11_DIR="$PYBIND11_DIR"
|
|
export CMAKE_PREFIX_PATH="${PYBIND11_DIR}:${CMAKE_PREFIX_PATH:-}"
|
|
|
|
safe_component() {
|
|
printf '%s' "$1" | sed 's/[^A-Za-z0-9_.-]/_/g'
|
|
}
|
|
|
|
PYTHON_SOABI=$(python3 -c \
|
|
"import sysconfig; print(sysconfig.get_config_var('SOABI') or 'unknown')")
|
|
TORCH_VERSION=$(python3 -c "import torch; print(torch.__version__)")
|
|
ROCM_VERSION=$(python3 -c "import torch; print(torch.version.hip or 'none')")
|
|
FFMPEG_VERSION=$(pkg-config --modversion libavcodec 2>/dev/null || echo unknown)
|
|
ARCH_TAG="${PYTORCH_ROCM_ARCH:-all}"
|
|
|
|
CACHE_DIR="${TORCHCODEC_WHEEL_CACHE}/$(safe_component "$TORCHCODEC_BRANCH")/$(safe_component "$PYTHON_SOABI")/$(safe_component "$TORCH_VERSION")/$(safe_component "$ROCM_VERSION")/$(safe_component "$FFMPEG_VERSION")/$(safe_component "$ARCH_TAG")"
|
|
CACHED_WHEEL=""
|
|
if [ -d "$CACHE_DIR" ]; then
|
|
CACHED_WHEEL=$(find "$CACHE_DIR" -maxdepth 1 -type f -name 'torchcodec-*.whl' -print -quit)
|
|
fi
|
|
|
|
validate_wheel_filename() {
|
|
WHEEL_PATH="$1" python3 - <<'PY'
|
|
from os import environ
|
|
from pathlib import Path
|
|
|
|
from packaging.utils import parse_wheel_filename
|
|
|
|
wheel = Path(environ["WHEEL_PATH"])
|
|
parse_wheel_filename(wheel.name)
|
|
PY
|
|
}
|
|
|
|
if [ -n "$CACHED_WHEEL" ] && validate_wheel_filename "$CACHED_WHEEL"; then
|
|
echo "Trying cached wheel: $CACHED_WHEEL"
|
|
if python3 -m pip install --no-deps --force-reinstall "$CACHED_WHEEL" && \
|
|
verify_torchcodec; then
|
|
echo "Installed TorchCodec from cache."
|
|
exit 0
|
|
fi
|
|
echo "Cached wheel is incompatible; rebuilding TorchCodec."
|
|
rm -f "$CACHED_WHEEL"
|
|
fi
|
|
|
|
BUILD_DIR=$(mktemp -d -t torchcodec-XXXXXX)
|
|
trap 'rm -rf "$BUILD_DIR"' EXIT
|
|
|
|
cd "$BUILD_DIR"
|
|
echo "Cloning TorchCodec $TORCHCODEC_BRANCH..."
|
|
git clone --depth 1 --branch "$TORCHCODEC_BRANCH" "$TORCHCODEC_REPO" torchcodec
|
|
cd torchcodec
|
|
|
|
export TORCHCODEC_CMAKE_BUILD_DIR="${PWD}/build"
|
|
export TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR=1
|
|
export I_CONFIRM_THIS_IS_NOT_A_LICENSE_VIOLATION=1
|
|
export CMAKE_GENERATOR=Ninja
|
|
export MAX_JOBS="${MAX_JOBS:-$(nproc)}"
|
|
|
|
if command -v ccache >/dev/null 2>&1; then
|
|
export CMAKE_C_COMPILER_LAUNCHER=ccache
|
|
export CMAKE_CXX_COMPILER_LAUNCHER=ccache
|
|
fi
|
|
|
|
echo "Building TorchCodec (MAX_JOBS=$MAX_JOBS)..."
|
|
python3 -m pip wheel . --no-build-isolation --no-deps -w "$BUILD_DIR/dist"
|
|
BUILT_WHEEL=$(find "$BUILD_DIR/dist" -maxdepth 1 -type f -name 'torchcodec-*.whl' -print -quit)
|
|
test -n "$BUILT_WHEEL"
|
|
validate_wheel_filename "$BUILT_WHEEL"
|
|
|
|
python3 -m pip install --no-deps --force-reinstall "$BUILT_WHEEL"
|
|
verify_torchcodec
|
|
|
|
mkdir -p "$CACHE_DIR"
|
|
cp "$BUILT_WHEEL" "$CACHE_DIR/$(basename "$BUILT_WHEEL")"
|
|
echo "Cached TorchCodec wheel in $CACHE_DIR"
|
|
echo "=== TorchCodec ROCm installation complete ==="
|