@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
sh 131 lines 4.27 KB
Raw
1 #!/usr/bin/env bash
2 # SPDX-License-Identifier: Apache-2.0
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 #
5 # Cache helper for ROCm base wheels
6 #
7 # This script manages caching of pre-built ROCm base wheels (torch, triton, etc.)
8 # to avoid rebuilding them when Dockerfile.rocm_base hasn't changed.
9 #
10 # Usage:
11 # cache-rocm-base-wheels.sh check - Check if cache exists, outputs "hit" or "miss"
12 # cache-rocm-base-wheels.sh upload - Upload wheels to cache
13 # cache-rocm-base-wheels.sh download - Download wheels from cache
14 # cache-rocm-base-wheels.sh key - Output the cache key
15 #
16 # Environment variables:
17 # S3_BUCKET - S3 bucket name (default: vllm-wheels)
18 #
19 # Note: ROCm version is determined by BASE_IMAGE in Dockerfile.rocm_base,
20 # so changes to ROCm version are captured by the Dockerfile hash.
21
22 set -euo pipefail
23
24 BUCKET="${S3_BUCKET:-vllm-wheels}"
25 DOCKERFILE="docker/Dockerfile.rocm_base"
26 CACHE_PREFIX="rocm/cache"
27
28 # Generate hash from Dockerfile content + build args
29 generate_cache_key() {
30 # Include Dockerfile content
31 if [[ ! -f "$DOCKERFILE" ]]; then
32 echo "ERROR: Dockerfile not found: $DOCKERFILE" >&2
33 exit 1
34 fi
35 local dockerfile_hash=$(sha256sum "$DOCKERFILE" | cut -c1-16)
36
37 echo "${dockerfile_hash}"
38 }
39
40 CACHE_KEY=$(generate_cache_key)
41 CACHE_PATH="s3://${BUCKET}/${CACHE_PREFIX}/${CACHE_KEY}/"
42
43 case "${1:-}" in
44 check)
45 echo "Checking cache for key: ${CACHE_KEY}" >&2
46 echo "Cache path: ${CACHE_PATH}" >&2
47
48 # Check if cache exists by listing objects
49 # We look for at least one .whl file
50 echo "Running: aws s3 ls ${CACHE_PATH}" >&2
51 S3_OUTPUT=$(aws s3 ls "${CACHE_PATH}" 2>&1) || true
52 echo "S3 ls output:" >&2
53 echo "$S3_OUTPUT" | head -5 >&2
54
55 if echo "$S3_OUTPUT" | grep -q "\.whl"; then
56 echo "hit"
57 else
58 echo "miss"
59 fi
60 ;;
61
62 upload)
63 echo "========================================"
64 echo "Uploading wheels to cache"
65 echo "========================================"
66 echo "Cache key: ${CACHE_KEY}"
67 echo "Cache path: ${CACHE_PATH}"
68 echo ""
69
70 if [[ ! -d "artifacts/rocm-base-wheels" ]]; then
71 echo "ERROR: artifacts/rocm-base-wheels directory not found" >&2
72 exit 1
73 fi
74
75 WHEEL_COUNT=$(find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
76 if [[ "$WHEEL_COUNT" -eq 0 ]]; then
77 echo "ERROR: No wheels found in artifacts/rocm-base-wheels/" >&2
78 exit 1
79 fi
80
81 echo "Uploading $WHEEL_COUNT wheels..."
82 aws s3 cp --recursive artifacts/rocm-base-wheels/ "${CACHE_PATH}"
83
84 echo ""
85 echo "Cache upload complete!"
86 echo "========================================"
87 ;;
88
89 download)
90 echo "========================================"
91 echo "Downloading wheels from cache"
92 echo "========================================"
93 echo "Cache key: ${CACHE_KEY}"
94 echo "Cache path: ${CACHE_PATH}"
95 echo ""
96 mkdir -p artifacts/rocm-base-wheels
97
98 # Use sync with include/exclude to only download .whl files
99 aws s3 sync "${CACHE_PATH}" artifacts/rocm-base-wheels/ \
100 --exclude "*" \
101 --include "*.whl"
102
103 echo ""
104 echo "Downloaded wheels:"
105 find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' -exec ls -lh {} \;
106 WHEEL_COUNT=$(find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
107 echo ""
108 echo "Total: $WHEEL_COUNT wheels"
109 echo "========================================"
110 ;;
111
112 key)
113 echo "${CACHE_KEY}"
114 ;;
115
116 path)
117 echo "${CACHE_PATH}"
118 ;;
119
120 *)
121 echo "Usage: $0 {check|upload|download|key|path}" >&2
122 echo "" >&2
123 echo "Commands:" >&2
124 echo " check - Check if cache exists, outputs 'hit' or 'miss'" >&2
125 echo " upload - Upload wheels from artifacts/rocm-base-wheels/ to cache" >&2
126 echo " download - Download wheels from cache to artifacts/rocm-base-wheels/" >&2
127 echo " key - Output the cache key" >&2
128 echo " path - Output the full S3 cache path" >&2
129 exit 1
130 ;;
131 esac