test/scripts/package_vendor_gpu_image.sh

133 lines
3.8 KiB
Bash

#!/usr/bin/env bash
set -euo pipefail
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
VENDOR=""
BASE_IMAGE=""
REGISTRY="unis"
IMAGE_NAME="qwen3-asr"
VERSION="$(date +"%Y%m%d_%H%M%S")"
OUTPUT_DIR="${ROOT_DIR}/build-file"
EXPORT_TAR="true"
NO_CACHE="false"
PYTHON_BIN=""
usage() {
cat <<EOF
Usage:
./scripts/package_vendor_gpu_image.sh --vendor metax|iluvatar|mthreads --base-image IMAGE [options]
Options:
--vendor NAME Vendor accelerator: metax, iluvatar, or mthreads
--base-image IMAGE Official vendor vLLM image already loaded/pullable by Docker
-v, --version VER Output image version tag; default: timestamp
-r, --registry REG Image registry/namespace; default: ${REGISTRY}
-o, --output DIR Export directory; default: ${OUTPUT_DIR}
--python-bin BIN Python inside base image; default: /opt/conda/bin/python for metax, python3 for iluvatar
--no-export Build image only, do not docker save tar.gz
--no-cache Disable Docker build cache
-h, --help Show help
Example:
./scripts/package_vendor_gpu_image.sh \\
--vendor metax \\
--base-image metax/vllm:3.7.0.38-c500-py310 \\
-v n260-3.7.0.38
./scripts/package_vendor_gpu_image.sh \\
--vendor iluvatar \\
--base-image registry.iluvatar.com.cn:10443/customer/sz/vllm0.17.0-4.4.0-x86:v5 \\
-v vllm0.17.0-4.4.0-v5
./scripts/package_vendor_gpu_image.sh \\
--vendor mthreads \\
--base-image registry.mthreads.com/presale/devtech/vllm_musa:s4000_4.3.5_d0519 \\
-v s4000_4.3.5_d0519
EOF
}
die() { echo "[ERROR] $1" >&2; exit 1; }
info() { echo "[INFO] $1" >&2; }
compressor() { command -v pigz >/dev/null 2>&1 && echo "pigz -f" || echo "gzip -f"; }
while [[ $# -gt 0 ]]; do
case "$1" in
--vendor) VENDOR="$2"; shift 2 ;;
--base-image) BASE_IMAGE="$2"; shift 2 ;;
-v|--version) VERSION="$2"; shift 2 ;;
-r|--registry) REGISTRY="$2"; shift 2 ;;
-o|--output) OUTPUT_DIR="$2"; shift 2 ;;
--python-bin) PYTHON_BIN="$2"; shift 2 ;;
--no-export) EXPORT_TAR="false"; shift ;;
--no-cache) NO_CACHE="true"; shift ;;
-h|--help) usage; exit 0 ;;
*) die "Unknown argument: $1" ;;
esac
done
[[ -n "$VENDOR" ]] || die "--vendor is required"
[[ -n "$BASE_IMAGE" ]] || die "--base-image is required"
case "$VENDOR" in
metax)
dockerfile="Dockerfile.metax"
tag="${REGISTRY}/${IMAGE_NAME}:metax-${VERSION}"
target="metax"
build_arg_name="METAX_BASE_IMAGE"
default_python_bin="/opt/conda/bin/python"
;;
iluvatar|ix|tianshu)
dockerfile="Dockerfile.iluvatar"
tag="${REGISTRY}/${IMAGE_NAME}:iluvatar-${VERSION}"
target="iluvatar"
build_arg_name="ILUVATAR_BASE_IMAGE"
default_python_bin="python3"
;;
mthreads|musa|moorethreads)
dockerfile="Dockerfile.mthreads"
tag="${REGISTRY}/${IMAGE_NAME}:mthreads-${VERSION}"
target="mthreads"
build_arg_name="MTHREADS_BASE_IMAGE"
default_python_bin="python3"
;;
*)
die "Unsupported vendor: ${VENDOR}"
;;
esac
PYTHON_BIN="${PYTHON_BIN:-$default_python_bin}"
docker image inspect "$BASE_IMAGE" >/dev/null 2>&1 || {
die "Base image not found locally: ${BASE_IMAGE}. Run docker load first or docker pull it."
}
args=(
build
-f "$dockerfile"
-t "$tag"
--build-arg "${build_arg_name}=${BASE_IMAGE}"
--build-arg "PYTHON_BIN=${PYTHON_BIN}"
)
[[ "$NO_CACHE" == "true" ]] && args+=(--no-cache)
info "Building fused image: ${tag}"
info "Vendor base image: ${BASE_IMAGE}"
info "Base image Python: ${PYTHON_BIN}"
(
cd "$ROOT_DIR"
docker "${args[@]}" .
)
if [[ "$EXPORT_TAR" == "true" ]]; then
mkdir -p "$OUTPUT_DIR"
tar_path="${OUTPUT_DIR}/${IMAGE_NAME}-${target}-${VERSION}-amd64.tar"
info "Exporting image archive: ${tar_path}"
docker save -o "$tar_path" "$tag"
info "Compressing archive"
$(compressor) "$tar_path"
info "Done: ${tar_path}.gz"
else
info "Built image: ${tag}"
fi