@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
sh 98 lines 3.65 KB
Raw
1 #!/bin/bash
2 set -e # Exit immediately if a command exits with a non-zero status.
3 # Script to build VLLM wheel for TPU with an optional version override.
4
5 SCRIPT_PATH_PARAM="$0"
6 TOOLS_DIR=$(cd "$(dirname "$SCRIPT_PATH_PARAM")" && pwd) # Absolute path to the script's directory
7 REPO_ROOT=$(cd "$TOOLS_DIR/../../" && pwd) # Absolute path to the repo root
8 VLLM_DIR="$REPO_ROOT/" # Path to the vllm sources
9
10 CHANGE_FILE_LIST=(
11 "vllm/entrypoints/cli/main.py"
12 "vllm/entrypoints/cli/run_batch.py"
13 "vllm/utils/__init__.py"
14 "vllm/platforms/__init__.py"
15 )
16
17 # Ensure we are not running from within the vllm directory if SCRIPT_PATH_PARAM is relative like "."
18 if [ "$TOOLS_DIR" = "$VLLM_DIR" ]; then
19 echo "Error: This script should not be run from the vllm directory directly if using relative paths."
20 echo "Place it in a subdirectory like 'tools/vllm-tpu' and run it from the repository root or via its full path."
21 exit 1
22 fi
23
24 # Optional version argument
25 if [ -n "$1" ]; then
26 USER_VERSION="$1"
27 export VLLM_VERSION_OVERRIDE="$USER_VERSION"
28 echo "User defined version: $USER_VERSION"
29 else
30 echo "No version override supplied. Using default version from source."
31 fi
32
33 PYPROJECT_FILE="$VLLM_DIR/pyproject.toml"
34
35 # Backup and update the project name.
36 if ! grep -q "name = \"vllm-tpu\"" "$PYPROJECT_FILE"; then
37 echo "Patching pyproject.toml project name to vllm-tpu..."
38 cp "$PYPROJECT_FILE" "${PYPROJECT_FILE}.bak"
39 sed -i '0,/^name = "vllm"/s//name = "vllm-tpu"/' "$PYPROJECT_FILE"
40
41 echo "Patching ${CHANGE_FILE_LIST[*]} vllm to vllm-tpu..."
42 # patching
43 # importlib.metadata.version('vllm') -> importlib.metadata.version('vllm-tpu')
44 # importlib.metadata.version("vllm") -> importlib.metadata.version("vllm-tpu")
45 # importlib.metadata.metadata('vllm') -> importlib.metadata.metadata('vllm-tpu')
46 # importlib.metadata.metadata("vllm") -> importlib.metadata.metadata("vllm-tpu")
47 # version('vllm') -> version('vllm-tpu')
48 # version("vllm") -> version("vllm-tpu")
49 sed -i \
50 -e "s/importlib.metadata.version(\(['\"]\)vllm\1)/importlib.metadata.version(\1vllm-tpu\1)/" \
51 -e "s/importlib.metadata.metadata(\(['\"]\)vllm\1)/importlib.metadata.metadata(\1vllm-tpu\1)/" \
52 -e "s/version(\(['\"]\)vllm\1)/version(\1vllm-tpu\1)/" \
53 "${CHANGE_FILE_LIST[@]}"
54 PATCHED=true
55 else
56 PATCHED=false
57 fi
58
59 # Navigate to the vllm directory
60 cd "$VLLM_DIR"
61
62 # Cleanup function to be called on exit or error
63 cleanup() {
64 echo "Cleaning up..."
65 if [ "$PATCHED" = true ]; then
66 echo "Restoring original pyproject.toml..."
67 cp "${PYPROJECT_FILE}.bak" "$PYPROJECT_FILE"
68 rm -f "${PYPROJECT_FILE}.bak"
69
70 echo "Restoring vllm code..."
71 sed -i \
72 -e "s/importlib.metadata.version(\(['\"]\)vllm-tpu\1)/importlib.metadata.version(\1vllm\1)/" \
73 -e "s/importlib.metadata.metadata(\(['\"]\)vllm-tpu\1)/importlib.metadata.metadata(\1vllm\1)/" \
74 -e "s/version(\(['\"]\)vllm-tpu\1)/version(\1vllm\1)/" \
75 "${CHANGE_FILE_LIST[@]}"
76 fi
77 }
78 trap cleanup EXIT HUP INT QUIT PIPE TERM # Register cleanup function to run on script exit and various signals
79
80 echo "Updating pyproject.toml completed. Proceeding with build..."
81
82 echo "Install dependencies for no-isolation build"
83 pip install -r "$VLLM_DIR/requirements/build/tpu.txt"
84
85 echo "Building wheel for TPU..."
86 rm -rf dist/
87 mkdir -p dist/
88
89 # User confirmed to use 'python -m build' directly
90 if ! VLLM_TARGET_DEVICE=tpu python -m build --no-isolation; then
91 echo "Error: Python build command failed. Check if 'python -m build' works and the 'build' module is installed."
92 exit 1
93 fi
94
95 trap - EXIT HUP INT QUIT PIPE TERM
96 cleanup
97
98 exit 0