Skip to content

Commit 8538d3c

Browse files
allenwang28allenwang28
committed
pulls versions into its own file, modifies build_wheel.shs
1 parent 93cb0db commit 8538d3c

File tree

2 files changed

+133
-64
lines changed

2 files changed

+133
-64
lines changed

‎assets/versions.sh‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
# Version Configuration for Forge Wheel Building
8+
# This file contains all pinned versions and commits for dependencies
9+
10+
# PyTorch version
11+
PYTORCH_VERSION="2.9.0.dev20250905"
12+
13+
# vLLM branch
14+
VLLM_BRANCH="v0.10.0"
15+
16+
# Commit hashes
17+
MONARCH_COMMIT="d1c5ea4732704454efad82db678d4e66a4131bb2"
18+
TORCHTITAN_COMMIT="0cfbd0b3c2d827af629a107a77a9e47229c31663"
19+
TORCHSTORE_COMMIT="eed96eb55ce87d4a9880597dd7dfd0d291e9ac81"

‎scripts/build_wheels.sh‎

Lines changed: 114 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -14,42 +14,57 @@ YELLOW='\033[1;33m'
1414
BLUE='\033[0;34m'
1515
NC='\033[0m'
1616

17+
# Source version configuration
18+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
19+
VERSIONS_FILE="$SCRIPT_DIR/../assets/versions.sh"
20+
21+
if [ ! -f "$VERSIONS_FILE" ]; then
22+
echo -e "${RED}[ERROR]${NC} Versions file not found: $VERSIONS_FILE"
23+
exit 1
24+
fi
25+
26+
source "$VERSIONS_FILE"
27+
28+
# Validate required variables are set
29+
validate_versions() {
30+
local missing_vars=()
31+
32+
[ -z "${PYTORCH_VERSION:-}" ] && missing_vars+=("PYTORCH_VERSION")
33+
[ -z "${VLLM_BRANCH:-}" ] && missing_vars+=("VLLM_BRANCH")
34+
[ -z "${MONARCH_COMMIT:-}" ] && missing_vars+=("MONARCH_COMMIT")
35+
[ -z "${TORCHTITAN_COMMIT:-}" ] && missing_vars+=("TORCHTITAN_COMMIT")
36+
[ -z "${TORCHSTORE_COMMIT:-}" ] && missing_vars+=("TORCHSTORE_COMMIT")
37+
38+
if [ ${#missing_vars[@]} -gt 0 ]; then
39+
echo -e "${RED}[ERROR]${NC} Missing required variables in $VERSIONS_FILE:"
40+
for var in "${missing_vars[@]}"; do
41+
echo " - $var"
42+
done
43+
exit 1
44+
fi
45+
}
46+
47+
validate_versions
48+
1749
# Configuration
18-
PYTORCH_VERSION="2.9.0.dev20250905"
19-
VLLM_BRANCH="v0.10.0"
20-
MONARCH_COMMIT="6ca383aca99480aa1bf5853478d4d09fcb224035"
21-
TORCHTITAN_COMMIT="0cfbd0b3c2d827af629a107a77a9e47229c31663"
22-
TORCHSTORE_COMMIT="eed96eb55ce87d4a9880597dd7dfd0d291e9ac81"
2350
BUILD_DIR="$HOME/forge-build"
2451
WHEEL_DIR="$(pwd)/assets/wheels"
2552

2653
# Logging functions
2754
log_info() { echo -e "${GREEN}[INFO]${NC} $1"; }
2855
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
2956
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }
30-
log_step() { echo -e "${BLUE}[$1/$2]${NC} $3"; }
31-
32-
# Track current step for recovery
33-
CURRENT_STEP=0
34-
TOTAL_STEPS=8
57+
log_step() { echo -e "${BLUE}[STEP]${NC} $1"; }
3558

3659
# Function to handle step failures
3760
handle_failure() {
3861
local step_name="$1"
39-
local step_num="$2"
40-
local exit_code="$3"
41-
local retry_cmd="$4"
62+
local exit_code="$2"
4263

43-
log_error "Step $step_num failed: $step_name"
64+
log_error "Step failed: $step_name"
4465
log_error "Exit code: $exit_code"
4566
log_error "Working directory: $(pwd)"
4667
echo ""
47-
log_info "To retry this step manually:"
48-
echo " $retry_cmd"
49-
echo ""
50-
log_info "Or to resume from step $step_num:"
51-
echo " $0 --resume-from=$step_num"
52-
echo ""
5368
exit $exit_code
5469
}
5570

@@ -150,36 +165,67 @@ EOF
150165
}
151166

152167
# Parse command line arguments
153-
RESUME_FROM=1
168+
BUILD_TARGETS=()
169+
154170
while [[ $# -gt 0 ]]; do
155171
case $1 in
156-
--resume-from=*)
157-
RESUME_FROM="${1#*=}"
172+
vllm|monarch|torchtitan|torchstore)
173+
BUILD_TARGETS+=("$1")
158174
shift
159175
;;
176+
--help|-h)
177+
echo "Usage: $0 [TARGETS...]"
178+
echo ""
179+
echo "Build wheels for Forge dependencies."
180+
echo ""
181+
echo "Targets (default: all):"
182+
echo " vllm Build vLLM wheel"
183+
echo " monarch Build Monarch wheel"
184+
echo " torchtitan Build torchtitan wheel"
185+
echo " torchstore Build torchstore wheel"
186+
echo ""
187+
echo "Examples:"
188+
echo " $0 # Build all wheels"
189+
echo " $0 vllm # Build only vLLM"
190+
echo " $0 monarch torchtitan # Build Monarch and torchtitan"
191+
exit 0
192+
;;
160193
*)
161194
log_error "Unknown argument: $1"
195+
log_info "Use --help to see available options"
162196
exit 1
163197
;;
164198
esac
165199
done
166200

201+
# If no targets specified, build all
202+
if [ ${#BUILD_TARGETS[@]} -eq 0 ]; then
203+
BUILD_TARGETS=("vllm" "monarch" "torchtitan" "torchstore")
204+
log_info "No targets specified, building all wheels"
205+
else
206+
log_info "Building wheels: ${BUILD_TARGETS[*]}"
207+
fi
208+
209+
# Helper function to check if a target should be built
210+
should_build() {
211+
local target="$1"
212+
for t in "${BUILD_TARGETS[@]}"; do
213+
if [ "$t" == "$target" ]; then
214+
return 0
215+
fi
216+
done
217+
return 1
218+
}
219+
167220
# Step execution wrapper
168221
run_step() {
169-
local step_num="$1"
170-
local step_name="$2"
171-
local step_function="$3"
172-
173-
if [ "$step_num" -lt "$RESUME_FROM" ]; then
174-
log_info "Skipping step $step_num: $step_name (resuming from step $RESUME_FROM)"
175-
return 0
176-
fi
222+
local step_name="$1"
223+
local step_function="$2"
177224

178-
CURRENT_STEP=$step_num
179-
log_step "$step_num" "$TOTAL_STEPS" "$step_name"
225+
log_step "$step_name"
180226

181227
if ! $step_function; then
182-
handle_failure "$step_name" "$step_num" "$?" "Run step $step_num manually"
228+
handle_failure "$step_name" "$?"
183229
fi
184230
}
185231

@@ -196,6 +242,7 @@ step2_cuda_packages() {
196242

197243
# Step 3: Build vLLM wheel
198244
step3_vllm() {
245+
log_info "Building vLLM from branch: $VLLM_BRANCH (from $VERSIONS_FILE)"
199246
cd "$BUILD_DIR"
200247
if [ -d "vllm" ]; then
201248
log_warn "vLLM directory exists, removing..."
@@ -229,6 +276,7 @@ step4_rust_setup() {
229276

230277
# Step 5: Build Monarch wheel
231278
step5_monarch() {
279+
log_info "Building Monarch from commit: $MONARCH_COMMIT (from $VERSIONS_FILE)"
232280
cd "$BUILD_DIR"
233281
if [ -d "monarch" ]; then
234282
log_warn "Monarch directory exists, removing..."
@@ -245,6 +293,7 @@ step5_monarch() {
245293

246294
# Step 6: Build torchtitan wheel
247295
step6_torchtitan() {
296+
log_info "Building torchtitan from commit: $TORCHTITAN_COMMIT (from $VERSIONS_FILE)"
248297
cd "$BUILD_DIR"
249298
if [ -d "torchtitan" ]; then
250299
log_warn "torchtitan directory exists, removing..."
@@ -260,6 +309,7 @@ step6_torchtitan() {
260309

261310
# Step 7: Build torchstore wheel
262311
step7_torchstore() {
312+
log_info "Building torchstore from commit: $TORCHSTORE_COMMIT (from $VERSIONS_FILE)"
263313
cd "$BUILD_DIR"
264314
if [ -d "torchstore" ]; then
265315
log_warn "torchstore directory exists, removing..."
@@ -298,28 +348,35 @@ main() {
298348
echo "==================="
299349
echo ""
300350

301-
if [ "$RESUME_FROM" -gt 1 ]; then
302-
log_info "Resuming from step $RESUME_FROM..."
303-
# Source CUDA env if resuming
304-
if [ -f ~/.forge_cuda_env ]; then
305-
source ~/.forge_cuda_env
306-
fi
307-
# Source Rust env if resuming
308-
if [ -f ~/.cargo/env ]; then
309-
source ~/.cargo/env
310-
fi
311-
else
312-
validate_environment
313-
setup_build_dir
351+
validate_environment
352+
setup_build_dir
353+
354+
# PyTorch is needed for all builds
355+
run_step "Installing PyTorch nightly" step1_pytorch
356+
357+
# CUDA packages are needed for vLLM and Monarch
358+
if should_build "vllm" || should_build "monarch"; then
359+
run_step "Installing CUDA packages and setting environment" step2_cuda_packages
360+
fi
361+
362+
# Build requested wheels
363+
if should_build "vllm"; then
364+
run_step "Building vLLM wheel" step3_vllm
314365
fi
315366

316-
run_step 1 "Installing PyTorch nightly" step1_pytorch
317-
run_step 2 "Installing CUDA packages and setting environment" step2_cuda_packages
318-
run_step 3 "Building vLLM wheel" step3_vllm
319-
run_step 4 "Setting up Rust toolchain and additional packages" step4_rust_setup
320-
run_step 5 "Building Monarch wheel" step5_monarch
321-
run_step 6 "Building torchtitan wheel" step6_torchtitan
322-
run_step 7 "Building torchstore wheel" step7_torchstore
367+
# Rust setup is needed for Monarch
368+
if should_build "monarch"; then
369+
run_step "Setting up Rust toolchain and additional packages" step4_rust_setup
370+
run_step "Building Monarch wheel" step5_monarch
371+
fi
372+
373+
if should_build "torchtitan"; then
374+
run_step "Building torchtitan wheel" step6_torchtitan
375+
fi
376+
377+
if should_build "torchstore"; then
378+
run_step "Building torchstore wheel" step7_torchstore
379+
fi
323380

324381
verify_installation
325382

@@ -333,21 +390,14 @@ main() {
333390
log_info " conda activate forge"
334391
log_info " pip install torch==$PYTORCH_VERSION --index-url https://download.pytorch.org/whl/nightly/cu129"
335392
log_info " pip install $WHEEL_DIR/*.whl"
336-
log_info " source ~/.forge_cuda_env"
393+
if should_build "vllm" || should_build "monarch"; then
394+
log_info " source ~/.forge_cuda_env"
395+
fi
337396
log_info ""
338397
log_info "Build artifacts are in: $BUILD_DIR"
339398
log_info "You can remove them with: rm -rf $BUILD_DIR"
340399
}
341400

342-
# Set trap for cleanup on failure
343-
cleanup() {
344-
if [ $? -ne 0 ] && [ $CURRENT_STEP -gt 0 ]; then
345-
echo ""
346-
log_error "Setup failed at step $CURRENT_STEP"
347-
log_info "You can resume with: $0 --resume-from=$CURRENT_STEP"
348-
fi
349-
}
350-
trap cleanup EXIT
351401

352402
# Run main function
353403
main "$@"

0 commit comments

Comments
 (0)