@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
sh 140 lines 4.88 KB
Raw
1 #!/bin/bash
2 # SPDX-License-Identifier: Apache-2.0
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
5 # Script to install TorchCodec from source (required for ROCm compatibility)
6 # The PyPI wheel is built against upstream PyTorch and has ABI mismatches with
7 # ROCm's custom torch build, so we must compile from source.
8
9 set -e
10
11 TORCHCODEC_REPO="${TORCHCODEC_REPO:-https://github.com/pytorch/torchcodec.git}"
12 # Pin to a specific release for reproducibility; update as needed.
13 TORCHCODEC_BRANCH="${TORCHCODEC_BRANCH:-v0.10.0}"
14 # Cache directory for pre-built wheels to avoid redundant recompilation.
15 TORCHCODEC_WHEEL_CACHE="${TORCHCODEC_WHEEL_CACHE:-/root/.cache/torchcodec-wheels}"
16
17 echo "=== TorchCodec Installation Script ==="
18
19 # Check if torchcodec is already installed and working
20 if python3 -c "from torchcodec.decoders import VideoDecoder" 2>/dev/null; then
21 echo "TorchCodec is already installed and working. Skipping."
22 exit 0
23 fi
24
25 # Try to install from cached wheel first
26 ARCH_TAG="${PYTORCH_ROCM_ARCH:-all}"
27 # Normalize arch tag (replace ; with _) for use in filename
28 ARCH_TAG="${ARCH_TAG//;/_}"
29 CACHED_WHEEL="${TORCHCODEC_WHEEL_CACHE}/torchcodec-${TORCHCODEC_BRANCH}-${ARCH_TAG}.whl"
30
31 if [ -f "$CACHED_WHEEL" ]; then
32 echo "Found cached wheel: $CACHED_WHEEL"
33 pip install "$CACHED_WHEEL" && {
34 echo "Installed from cached wheel."
35 echo "=== TorchCodec installation complete ==="
36 exit 0
37 }
38 echo "Cached wheel installation failed, rebuilding from source..."
39 fi
40
41 echo "TorchCodec not found. Installing from source..."
42
43 # Install system dependencies (FFmpeg + pkg-config) if not already present.
44 # The Docker test image pre-installs these, so this is a fallback for other envs.
45 install_system_deps() {
46 if command -v apt-get &> /dev/null; then
47 echo "Installing system dependencies..."
48 apt-get update && apt-get install -y --no-install-recommends \
49 pkg-config \
50 ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
51 libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev
52 else
53 echo "Warning: apt-get did not work. Please install dependencies manually."
54 return 1
55 fi
56 }
57
58 # Check for pkg-config
59 if ! command -v pkg-config &> /dev/null; then
60 echo "pkg-config not found. Installing system dependencies..."
61 install_system_deps
62 fi
63
64 # Check for required FFmpeg libraries
65 echo "Checking for FFmpeg libraries..."
66 if ! pkg-config --exists libavcodec libavformat libavutil libswscale libavdevice libavfilter libswresample 2>/dev/null; then
67 echo "FFmpeg development libraries not found. Installing..."
68 install_system_deps
69 fi
70
71 # Install Python build dependencies
72 echo "Installing Python build dependencies..."
73 pip install pybind11 setuptools wheel
74
75 # Set pybind11 cmake path so CMake can find it
76 export pybind11_DIR=$(python3 -c "import pybind11; print(pybind11.get_cmake_dir())")
77 export CMAKE_PREFIX_PATH="${pybind11_DIR}:${CMAKE_PREFIX_PATH}"
78 echo "pybind11_DIR set to: $pybind11_DIR"
79
80 # Limit GPU architectures to only what this image targets.
81 # The default builds for all supported archs which is very slow.
82 if [ -n "$PYTORCH_ROCM_ARCH" ]; then
83 echo "Building for PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH"
84 fi
85
86 # Create temp directory for build
87 BUILD_DIR=$(mktemp -d -t torchcodec-XXXXXX)
88 echo "Building in temporary directory: $BUILD_DIR"
89
90 cleanup() {
91 echo "Cleaning up $BUILD_DIR"
92 rm -rf "$BUILD_DIR"
93 }
94 trap cleanup EXIT
95
96 # Clone and build
97 cd "$BUILD_DIR"
98 echo "Cloning TorchCodec from $TORCHCODEC_REPO (branch: $TORCHCODEC_BRANCH)..."
99 git clone --depth 1 --branch "$TORCHCODEC_BRANCH" "$TORCHCODEC_REPO" torchcodec
100
101 cd torchcodec
102
103 # Set build environment for ROCm compatibility
104 export TORCHCODEC_CMAKE_BUILD_DIR="${PWD}/build"
105 export TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR=1
106 export I_CONFIRM_THIS_IS_NOT_A_LICENSE_VIOLATION=1
107 # Use ninja for faster builds and parallelize compilation
108 export CMAKE_GENERATOR=Ninja
109 export MAX_JOBS="${MAX_JOBS:-$(nproc)}"
110 # Use ccache if available to speed up recompilation
111 if command -v ccache &> /dev/null; then
112 export CMAKE_C_COMPILER_LAUNCHER=ccache
113 export CMAKE_CXX_COMPILER_LAUNCHER=ccache
114 fi
115
116 echo "Building TorchCodec (MAX_JOBS=$MAX_JOBS)..."
117 pip wheel . --no-build-isolation --no-deps -w "$BUILD_DIR/dist"
118
119 # Install the built wheel
120 BUILT_WHEEL=$(find "$BUILD_DIR/dist" -maxdepth 1 -name 'torchcodec-*.whl' -print -quit)
121 if [ -z "$BUILT_WHEEL" ]; then
122 echo "Error: No wheel produced"
123 exit 1
124 fi
125
126 pip install "$BUILT_WHEEL"
127
128 # Cache the wheel for future runs
129 mkdir -p "$TORCHCODEC_WHEEL_CACHE"
130 cp "$BUILT_WHEEL" "$CACHED_WHEEL"
131 echo "Cached wheel to: $CACHED_WHEEL"
132
133 # Verify installation
134 echo "Verifying installation..."
135 if python3 -c "from torchcodec.decoders import VideoDecoder; print('TorchCodec installed successfully!')"; then
136 echo "=== TorchCodec installation complete ==="
137 else
138 echo "Error: TorchCodec installation failed verification"
139 exit 1
140 fi