101 lines
3.3 KiB
Bash
101 lines
3.3 KiB
Bash
#!/bin/bash
|
|
|
|
set -euo pipefail
|
|
|
|
# arguments:
|
|
# $1: SOC_ARG (ascend910b, ascend910_93, ascend950)
|
|
|
|
SOC_ARG="${1:-}"
|
|
|
|
log() {
|
|
echo "[install_batch_invariant] $*"
|
|
}
|
|
|
|
# validate arguments
|
|
if [[ -z "${SOC_ARG}" ]]; then
|
|
log "ERROR: SOC_ARG is required as first argument"
|
|
exit 1
|
|
fi
|
|
|
|
log "Starting batch_invariant installation..."
|
|
log "SOC_ARG=${SOC_ARG}"
|
|
|
|
# determine device type from SOC_ARG
|
|
case "${SOC_ARG}" in
|
|
ascend910b)
|
|
BATCH_INVARIANT_DEVICE="910b"
|
|
;;
|
|
ascend910_93)
|
|
BATCH_INVARIANT_DEVICE="A3"
|
|
;;
|
|
*)
|
|
log "Warning: batch_invariant not available for SOC_ARG=${SOC_ARG}; skipping"
|
|
exit 0
|
|
;;
|
|
esac
|
|
|
|
# detect system architecture
|
|
ARCH_INFO=$(uname -m)
|
|
case "${ARCH_INFO}" in
|
|
aarch64)
|
|
ARCH_SUFFIX="aarch64"
|
|
;;
|
|
x86_64)
|
|
ARCH_SUFFIX="x86_64"
|
|
;;
|
|
*)
|
|
log "Warning: unknown architecture ${ARCH_INFO}; cannot determine batch_invariant package"
|
|
exit 0
|
|
;;
|
|
esac
|
|
|
|
# download and install run package
|
|
BATCH_INVARIANT_RUN_URL="https://vllm-ascend.obs.cn-north-4.myhuaweicloud.com/vllm-ascend/cann-ops-batch_invariant-${BATCH_INVARIANT_DEVICE}-1.0.0-linux.${ARCH_SUFFIX}.run"
|
|
BATCH_INVARIANT_RUN_FILE="cann-ops-batch_invariant-${BATCH_INVARIANT_DEVICE}-1.0.0-linux.${ARCH_SUFFIX}.run"
|
|
|
|
log "Downloading batch_invariant run package..."
|
|
unset ASCEND_CUSTOM_OPP_PATH
|
|
if curl --max-time 60 -sS -k -O "${BATCH_INVARIANT_RUN_URL}" && [[ -f "${BATCH_INVARIANT_RUN_FILE}" ]]; then
|
|
chmod +x "${BATCH_INVARIANT_RUN_FILE}"
|
|
log "Running installer: ${BATCH_INVARIANT_RUN_FILE}"
|
|
if "./${BATCH_INVARIANT_RUN_FILE}"; then
|
|
log "batch_invariant run package installed successfully"
|
|
else
|
|
log "Failed to install batch_invariant run package"
|
|
fi
|
|
else
|
|
log "Failed to download batch_invariant run package: ${BATCH_INVARIANT_RUN_URL}"
|
|
fi
|
|
# clean up downloaded run file (always clean, regardless of success/failure)
|
|
rm -f "${BATCH_INVARIANT_RUN_FILE}"
|
|
|
|
# download and install whl package
|
|
BATCH_INVARIANT_WHL_URL="https://vllm-ascend.obs.cn-north-4.myhuaweicloud.com/vllm-ascend/batch_invariant-torch_ops_extension-1.0.0.zip"
|
|
BATCH_INVARIANT_WHL_FILE="batch_invariant-torch_ops_extension-1.0.0.zip"
|
|
|
|
log "Downloading batch_invariant whl package..."
|
|
if curl --max-time 3 -sS -k -O "${BATCH_INVARIANT_WHL_URL}" >/dev/null 2>&1 && [[ -f "${BATCH_INVARIANT_WHL_FILE}" ]]; then
|
|
if unzip -o "${BATCH_INVARIANT_WHL_FILE}" >/dev/null 2>&1; then
|
|
if [[ -d "torch_ops_extension/batch_invariant_ops" ]]; then
|
|
cd torch_ops_extension/batch_invariant_ops
|
|
log "Building and installing batch_invariant whl package..."
|
|
if bash build_and_install.sh; then
|
|
log "batch_invariant whl package installed successfully"
|
|
else
|
|
log "Failed to build and install batch_invariant whl package"
|
|
fi
|
|
cd -
|
|
else
|
|
log "batch_invariant_ops directory not found in zip"
|
|
fi
|
|
else
|
|
log "Failed to unzip batch_invariant whl package"
|
|
fi
|
|
else
|
|
log "Failed to download batch_invariant whl package: ${BATCH_INVARIANT_WHL_URL}"
|
|
fi
|
|
# clean up downloaded files (always clean, regardless of success/failure)
|
|
rm -rf "${BATCH_INVARIANT_WHL_FILE}" torch_ops_extension
|
|
|
|
log "batch_invariant_ops build completed"
|