Skip to content

Commit eaa6d2a

Browse files
sirutBuasaisirutBuasai
committed
add pytorch runtime
Signed-off-by: sirutBuasai <sirutbuasai27@outlook.com>
1 parent d928d7a commit eaa6d2a

4 files changed

Lines changed: 230 additions & 0 deletions

File tree

‎scripts/ci/build/pytorch_runtime/lib/fetch_wheels.sh‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
#!/usr/bin/env bash
2+
# Download pre-built wheels from S3 into the build context.
3+
#
4+
# Usage:
5+
# bash fetch_wheels.sh --dest-dir <dir> --bucket <bucket> --cuda-version <ver> \
6+
# --packages "flash-attn:2.8.3,transformer-engine-torch:2.12.0"
7+
#
8+
# Exit code: 0 if all wheels found, 1 if any cache miss.
9+
# Cache key: s3://<bucket>/wheels/<cuda>/<pkg_underscore>/<wheel_filename>
10+
11+
set -euo pipefail
12+
13+
DEST_DIR=""
14+
BUCKET="dlc-cicd-wheels"
15+
CUDA=""
16+
PACKAGES=""
17+
18+
while [[ $# -gt 0 ]]; do
19+
case "$1" in
20+
--dest-dir) DEST_DIR="$2"; shift 2 ;;
21+
--bucket) BUCKET="$2"; shift 2 ;;
22+
--cuda-version) CUDA="$2"; shift 2 ;;
23+
--packages) PACKAGES="$2"; shift 2 ;;
24+
*) echo "Unknown argument: $1" >&2; exit 1 ;;
25+
esac
26+
done
27+
28+
[[ -n "$DEST_DIR" ]] || { echo "ERROR: --dest-dir is required" >&2; exit 1; }
29+
[[ -n "$CUDA" ]] || { echo "ERROR: --cuda-version is required" >&2; exit 1; }
30+
31+
if [[ -z "$BUCKET" ]]; then
32+
echo "No wheel cache bucket configured — skipping"
33+
exit 1
34+
fi
35+
36+
mkdir -p "${DEST_DIR}"
37+
38+
ALL_HIT=true
39+
IFS=',' read -ra SPECS <<< "$PACKAGES"
40+
for spec in "${SPECS[@]}"; do
41+
[[ -z "$spec" ]] && continue
42+
PKG="${spec%%:*}"
43+
VER="${spec##*:}"
44+
PKG_UNDER="${PKG//-/_}"
45+
PREFIX="wheels/${CUDA}/${PKG_UNDER}/"
46+
47+
echo "Looking for ${PKG}==${VER} (${CUDA}) in s3://${BUCKET}/${PREFIX} ..."
48+
if aws s3 cp "s3://${BUCKET}/${PREFIX}" "${DEST_DIR}/" --recursive --exclude "*" --include "*.whl" 2>/dev/null; then
49+
echo "Cache hit: ${PKG}==${VER}"
50+
else
51+
echo "Cache miss: ${PKG}==${VER}"
52+
ALL_HIT=false
53+
fi
54+
done
55+
56+
if [[ "$ALL_HIT" == "true" ]]; then
57+
exit 0
58+
else
59+
exit 1
60+
fi

‎scripts/ci/build/pytorch_runtime/lib/upload_wheels.sh‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
#!/usr/bin/env bash
2+
# Extract built wheels from Docker wheel-export stage and upload to S3.
3+
#
4+
# Usage:
5+
# bash upload_wheels.sh --bucket <bucket> --cuda-version <ver> --image-uri <uri> \
6+
# --dockerfile <path> --packages "flash-attn:2.8.3,transformer-engine-torch:2.12.0"
7+
8+
set -euo pipefail
9+
10+
BUCKET="dlc-cicd-wheels"
11+
CUDA=""
12+
IMAGE=""
13+
DOCKERFILE=""
14+
PACKAGES=""
15+
16+
while [[ $# -gt 0 ]]; do
17+
case "$1" in
18+
--bucket) BUCKET="$2"; shift 2 ;;
19+
--cuda-version) CUDA="$2"; shift 2 ;;
20+
--image-uri) IMAGE="$2"; shift 2 ;;
21+
--dockerfile) DOCKERFILE="$2"; shift 2 ;;
22+
--packages) PACKAGES="$2"; shift 2 ;;
23+
*) echo "Unknown argument: $1" >&2; exit 1 ;;
24+
esac
25+
done
26+
27+
[[ -n "$BUCKET" ]] || { echo "ERROR: --bucket is required" >&2; exit 1; }
28+
[[ -n "$CUDA" ]] || { echo "ERROR: --cuda-version is required" >&2; exit 1; }
29+
[[ -n "$IMAGE" ]] || { echo "ERROR: --image-uri is required" >&2; exit 1; }
30+
[[ -n "$DOCKERFILE" ]] || { echo "ERROR: --dockerfile is required" >&2; exit 1; }
31+
32+
EXPORT_DIR=$(mktemp -d)
33+
34+
docker buildx build --progress=plain --target wheel-export --output "type=local,dest=${EXPORT_DIR}" \
35+
-f "${DOCKERFILE}" . 2>/dev/null || {
36+
echo "wheel-export stage not available — extracting from runtime image"
37+
CID=$(docker create "${IMAGE}" /bin/true)
38+
docker cp "${CID}:/tmp/built_wheels/" "${EXPORT_DIR}/wheels/" 2>/dev/null || true
39+
docker rm "${CID}" &>/dev/null || true
40+
}
41+
42+
IFS=',' read -ra SPECS <<< "$PACKAGES"
43+
for spec in "${SPECS[@]}"; do
44+
[[ -z "$spec" ]] && continue
45+
PKG="${spec%%:*}"
46+
PKG_UNDER="${PKG//-/_}"
47+
48+
WHL=$(find "${EXPORT_DIR}" -name "${PKG_UNDER}*.whl" 2>/dev/null | head -1)
49+
if [[ -z "${WHL}" ]]; then
50+
echo "No wheel found for ${PKG}"
51+
continue
52+
fi
53+
54+
FNAME=$(basename "${WHL}")
55+
S3_KEY="wheels/${CUDA}/${PKG_UNDER}/${FNAME}"
56+
57+
if aws s3 ls "s3://${BUCKET}/${S3_KEY}" &>/dev/null; then
58+
echo "Already cached: ${S3_KEY}"
59+
continue
60+
fi
61+
62+
echo "Uploading ${FNAME} -> s3://${BUCKET}/${S3_KEY}"
63+
aws s3 cp "${WHL}" "s3://${BUCKET}/${S3_KEY}" || echo "Upload failed (non-fatal)"
64+
done
65+
66+
rm -rf "${EXPORT_DIR}"

‎scripts/ci/build/pytorch_runtime/post_build.sh‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
#!/usr/bin/env bash
2+
# Post-build hook for PyTorch.
3+
# Uploads built wheels to S3 cache for future builds.
4+
#
5+
# Usage:
6+
# bash .github/scripts/build/pytorch/post_build.sh --config-file <path>
7+
#
8+
# Inputs:
9+
# --config-file - config file path
10+
# WHEEL_CACHE_HIT - "true" to skip upload (env var from pre_build)
11+
# WHEELS_BUCKET - S3 bucket (env var)
12+
# CI_IMAGE_URI - built image URI (env var)
13+
#
14+
# Outputs: none (uploads to S3)
15+
16+
set -euo pipefail
17+
18+
if [[ "${WHEEL_CACHE_HIT:-}" == "true" ]]; then
19+
echo "Wheel cache hit — skipping upload"
20+
exit 0
21+
fi
22+
23+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
24+
CONFIG_FILE=""
25+
26+
while [[ $# -gt 0 ]]; do
27+
case "$1" in
28+
--config-file) CONFIG_FILE="$2"; shift 2 ;;
29+
*) echo "Unknown argument: $1" >&2; exit 1 ;;
30+
esac
31+
done
32+
33+
[[ -n "$CONFIG_FILE" ]] || { echo "ERROR: --config-file is required" >&2; exit 1; }
34+
[[ -f "$CONFIG_FILE" ]] || { echo "ERROR: Config file not found: $CONFIG_FILE" >&2; exit 1; }
35+
36+
CUDA_VERSION=$(yq '.build.cuda_version' "$CONFIG_FILE")
37+
FLASH_ATTN_VERSION=$(yq '.build.flash_attn_version // ""' "$CONFIG_FILE")
38+
TRANSFORMER_ENGINE_VERSION=$(yq '.build.transformer_engine_version // ""' "$CONFIG_FILE")
39+
DOCKERFILE=$(yq '.build.dockerfile' "$CONFIG_FILE")
40+
41+
PACKAGES=()
42+
[[ -n "$FLASH_ATTN_VERSION" ]] && PACKAGES+=("flash-attn:${FLASH_ATTN_VERSION}")
43+
[[ -n "$TRANSFORMER_ENGINE_VERSION" ]] && PACKAGES+=("transformer-engine-torch:${TRANSFORMER_ENGINE_VERSION}")
44+
45+
PACKAGES_STR=$(IFS=','; echo "${PACKAGES[*]}")
46+
echo "Uploading wheels to cache: ${PACKAGES_STR:-none}"
47+
bash "$SCRIPT_DIR/lib/upload_wheels.sh" --bucket "${WHEELS_BUCKET:-dlc-cicd-wheels}" \
48+
--cuda-version "$CUDA_VERSION" --image-uri "${CI_IMAGE_URI}" \
49+
--dockerfile "$DOCKERFILE" --packages "$PACKAGES_STR" || true

‎scripts/ci/build/pytorch_runtime/pre_build.sh‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
#!/usr/bin/env bash
2+
# Pre-build hook for PyTorch.
3+
# Fetches cached wheels from S3 to avoid recompilation.
4+
#
5+
# Usage:
6+
# bash .github/scripts/build/pytorch/pre_build.sh --config-file <path>
7+
#
8+
# Inputs:
9+
# --config-file - config file path
10+
# WHEELS_BUCKET - S3 bucket for wheel cache (env var, from workflow vars)
11+
#
12+
# Outputs (env vars written to $GITHUB_ENV):
13+
# WHEEL_CACHE_HIT - "true" if all wheels found, "false" otherwise
14+
#
15+
# Side effects:
16+
# Places wheels in docker/pytorch/{version}/wheels/ if cache hit
17+
18+
set -euo pipefail
19+
20+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
21+
CONFIG_FILE=""
22+
23+
while [[ $# -gt 0 ]]; do
24+
case "$1" in
25+
--config-file) CONFIG_FILE="$2"; shift 2 ;;
26+
*) echo "Unknown argument: $1" >&2; exit 1 ;;
27+
esac
28+
done
29+
30+
[[ -n "$CONFIG_FILE" ]] || { echo "ERROR: --config-file is required" >&2; exit 1; }
31+
[[ -f "$CONFIG_FILE" ]] || { echo "ERROR: Config file not found: $CONFIG_FILE" >&2; exit 1; }
32+
33+
CUDA_VERSION=$(yq '.build.cuda_version' "$CONFIG_FILE")
34+
FLASH_ATTN_VERSION=$(yq '.build.flash_attn_version // ""' "$CONFIG_FILE")
35+
TRANSFORMER_ENGINE_VERSION=$(yq '.build.transformer_engine_version // ""' "$CONFIG_FILE")
36+
FRAMEWORK_VERSION=$(yq '.metadata.framework_version' "$CONFIG_FILE")
37+
PYTORCH_SHORT=$(echo "$FRAMEWORK_VERSION" | cut -d. -f1,2)
38+
39+
DEST="docker/pytorch/${PYTORCH_SHORT}/wheels"
40+
mkdir -p "$DEST"
41+
42+
PACKAGES=()
43+
[[ -n "$FLASH_ATTN_VERSION" ]] && PACKAGES+=("flash-attn:${FLASH_ATTN_VERSION}")
44+
[[ -n "$TRANSFORMER_ENGINE_VERSION" ]] && PACKAGES+=("transformer-engine-torch:${TRANSFORMER_ENGINE_VERSION}")
45+
46+
PACKAGES_STR=$(IFS=','; echo "${PACKAGES[*]}")
47+
echo "Fetching cached wheels: ${PACKAGES_STR:-none}"
48+
if bash "$SCRIPT_DIR/lib/fetch_wheels.sh" --dest-dir "$DEST" --bucket "${WHEELS_BUCKET:-dlc-cicd-wheels}" \
49+
--cuda-version "$CUDA_VERSION" --packages "$PACKAGES_STR"; then
50+
echo "WHEEL_CACHE_HIT=true" >> "${GITHUB_ENV:-/dev/null}"
51+
echo "Wheel cache hit"
52+
else
53+
echo "WHEEL_CACHE_HIT=false" >> "${GITHUB_ENV:-/dev/null}"
54+
echo "Wheel cache miss — will build from source"
55+
fi

0 commit comments

Comments
 (0)