[CANN] Refactor cann.sh and workflows

.github/workflows/_docker-build.yml

@@ -84,7 +84,7 @@ jobs:
 
   docker:
     name: 'docker'
-    runs-on: ubuntu-22.04 # TODO: Use ubuntu-latest
+    runs-on: ubuntu-24.04-arm # TODO: Use ubuntu-latest
     needs:
       - prepare
     strategy:

.github/workflows/docker.yml

@@ -5,6 +5,7 @@ on:
   push:
     branches:
       - 'main'
+      - 'releases/**'
     paths:
       - '.github/actions/clean-up/action.yml'
       - '.github/actions/docker-login/action.yml'
@@ -20,6 +21,7 @@ on:
   pull_request:
     branches:
       - 'main'
+      - 'releases/**'
     types:
       - 'synchronize'
       - 'labeled'
@@ -72,7 +74,12 @@ jobs:
     name: python
     needs:
       - prepare
-    if: ${{ github.event_name != 'pull_request' || (!cancelled() && contains(needs.prepare.outputs.pr-labels, 'python')) }}
+    if: |
+      !cancelled() &&
+      (
+       (github.event_name != 'pull_request' && success()) ||
+       (github.event_name == 'pull_request' && contains(needs.prepare.outputs.pr-labels, 'python'))
+      )
     uses: ./.github/workflows/_docker-build.yml
     with:
       target: 'python'
@@ -88,7 +95,12 @@ jobs:
     needs:
       - prepare
       - python
-    if: ${{ github.event_name != 'pull_request' || (!cancelled() && contains(needs.prepare.outputs.pr-labels, 'cann')) }}
+    if: |
+      !cancelled() &&
+      (
+       (github.event_name != 'pull_request' && success()) ||
+       (github.event_name == 'pull_request' && contains(needs.prepare.outputs.pr-labels, 'cann'))
+      )
     uses: ./.github/workflows/_docker-build.yml
     with:
       target: 'cann'
@@ -100,7 +112,12 @@ jobs:
     needs:
       - prepare
       - cann
-    if: ${{ github.event_name != 'pull_request' || (!cancelled() && contains(needs.prepare.outputs.pr-labels, 'pytorch')) }}
+    if: |
+      !cancelled() &&
+      (
+       (github.event_name != 'pull_request' && success()) ||
+       (github.event_name == 'pull_request' && contains(needs.prepare.outputs.pr-labels, 'pytorch'))
+      )
     uses: ./.github/workflows/_docker-build.yml
     with:
       target: 'pytorch'
@@ -112,7 +129,12 @@ jobs:
     needs:
       - prepare
       - cann
-    if: ${{ github.event_name != 'pull_request' || (!cancelled() && contains(needs.prepare.outputs.pr-labels, 'mindspore')) }}
+    if: |
+      !cancelled() &&
+      (
+        (github.event_name != 'pull_request' && success()) ||
+        (github.event_name == 'pull_request' && contains(needs.prepare.outputs.pr-labels, 'mindspore'))
+      )
     uses: ./.github/workflows/_docker-build.yml
     with:
       target: 'mindspore'

cann/cann.sh

@@ -2,83 +2,86 @@
 
 set -e
 
-get_architecture() {
+# bash fonts colors
+cyan='\e[96m'
+yellow='\e[33m'
+red='\e[31m'
+none='\e[0m'
+
+_cyan() { echo -e "${cyan}$*${none}"; }
+_yellow() { echo -e "${yellow}$*${none}"; }
+_red() { echo -e "${red}$*${none}"; }
+
+_info() { _cyan "Info: $*"; }
+_warn() { _yellow "Warn: $*"; }
+_err() { _red "Error: $*" && exit 1; }
+
+_get_architecture() {
     # not case sensitive
     shopt -s nocasematch
 
     case "${PLATFORM}" in
         "linux/x86_64"|"linux/amd64")
-            ARCH="x86_64"
+            arch="x86_64"
             ;;
         "linux/aarch64"|"linux/arm64")
-            ARCH="aarch64"
+            arch="aarch64"
             ;;
         *)
-            echo "Error: Unsupported architecture ${PLATFORM}."
-            exit 1
+            _err "Unsupported architecture ${PLATFORM}."
             ;;
     esac
 
-    echo "${ARCH}"
+    echo "${arch}"
 }
 
-download_file() {
-    set +e
-
-    local max_retries=10
-    local retry_delay=10
-    local url="$1"
-    local path="$2"
-    echo "Downloading ${path} from ${url}"
-
-    for ((i=1; i<=max_retries; i++)); do
-        echo "Attempt $i of $max_retries..."
-        curl -fsSL -o "${path}" "${url}"
-        if [[ $? -eq 0 ]]; then
-            return 0
-        else
-            echo "Download failed with error code $?. Retrying in ${retry_delay} seconds..."
-            sleep ${retry_delay}
-        fi
-    done
+_retry() {
+    "$@" || (sleep 10 && "$@") || (sleep 20 && "$@") || (sleep 40 && "$@")
+}
 
-    echo "All attempts failed. Exiting."
-    exit 1
+_download_file() {
+    _info "Downloading file from $1"
+    _retry curl -fsSL -o "$2" "$1"
 }
 
 download_cann() {
+    declare -A version_dict
+    version_dict["8.0.RC2.alpha001"]="V100R001C18B800TP015"
+    version_dict["8.0.RC2.alpha002"]="V100R001C18SPC805"
+    version_dict["8.0.RC2.alpha003"]="V100R001C18SPC703"
+    version_dict["8.0.RC3.alpha002"]="V100R001C19SPC702"
+
     local url="https://ascend-repo.obs.cn-east-2.myhuaweicloud.com"
-    if [[ ${CANN_VERSION} == "8.0.RC2.alpha001" ]]; then
-        local url_prefix="${url}/Milan-ASL/Milan-ASL%20V100R001C18B800TP015"
-    elif [[ ${CANN_VERSION} == "8.0.RC2.alpha002" ]]; then
-        local url_prefix="${url}/Milan-ASL/Milan-ASL%20V100R001C18SPC805"
-    elif [[ ${CANN_VERSION} == "8.0.RC2.alpha003" ]]; then
-        local url_prefix="${url}/Milan-ASL/Milan-ASL%20V100R001C18SPC703"
-    elif [[ ${CANN_VERSION} == "8.0.RC3.alpha002" ]]; then
-        local url_prefix="${url}/Milan-ASL/Milan-ASL%20V100R001C19SPC702"
+    if [[ ${CANN_VERSION} == *alpha* ]]; then
+        local version=${version_dict[${CANN_VERSION}]}
+        if [[ ${version} ]]; then
+            local url_prefix="${url}/Milan-ASL/Milan-ASL%20${version}"
+        else
+            _err "The version is not defined for CANN ${CANN_VERSION}, please file an issue on GitHub."
+        fi
     else
         local url_prefix="${url}/CANN/CANN%20${CANN_VERSION}"
     fi
-    local url_suffix="response-content-type=application/octet-stream"
-    local toolkit_url="${url_prefix}/${TOOLKIT_FILE}?${url_suffix}"
-    local kernels_url="${url_prefix}/${KERNELS_FILE}?${url_suffix}"
 
+    # Download cann-toolkit
     if [ ! -f "${TOOLKIT_PATH}" ]; then
-        download_file "${toolkit_url}" "${TOOLKIT_PATH}"
+        local toolkit_url="${url_prefix}/${TOOLKIT_FILE}"
+        _download_file "${toolkit_url}" "${TOOLKIT_PATH}"
     fi
 
+    # Download cann-kernels
     if [ ! -f "${KERNELS_PATH}" ]; then
-        download_file "${kernels_url}" "${KERNELS_PATH}"
+        local kernels_url="${url_prefix}/${KERNELS_FILE}"
+        _download_file "${kernels_url}" "${KERNELS_PATH}"
     fi
 
-    echo "CANN ${CANN_VERSION} download successful."
+    _info "CANN ${CANN_VERSION} download successful."
 }
 
 set_env() {
     local cann_toolkit_env_file="${CANN_HOME}/ascend-toolkit/set_env.sh"
     if [ ! -f "${cann_toolkit_env_file}" ]; then
-        echo "CANN Toolkit ${CANN_VERSION} installation failed."
-        exit 1
+        _err "No such file: ${cann_toolkit_env_file}"
     else
         local driver_path_env="LD_LIBRARY_PATH=${CANN_HOME}/driver/lib64/common/:${CANN_HOME}/driver/lib64/driver/:\${LD_LIBRARY_PATH}" && \
         echo "export ${driver_path_env}" >> /etc/profile
@@ -92,7 +95,7 @@ set_env() {
 install_cann() {
     # Download installers
     if [ ! -f "${TOOLKIT_PATH}" ] || [ ! -f "${KERNELS_PATH}" ]; then
-        echo "[WARNING] Installers do not exist in the /tmp directory, re-download them."
+        _warn "Installers do not exist in the /tmp directory, re-download them."
         download_cann
     fi
 
@@ -103,7 +106,7 @@ install_cann() {
         psutil protobuf==3.20 scipy requests absl-py
 
     # Install CANN Toolkit
-    echo "Installing ${TOOLKIT_FILE}"
+    _info "Installing ${TOOLKIT_FILE}"
     chmod +x "${TOOLKIT_PATH}"
     bash "${TOOLKIT_PATH}" --quiet --install --install-for-all --install-path="${CANN_HOME}"
     rm -f "${TOOLKIT_PATH}"
@@ -112,22 +115,22 @@ install_cann() {
     set_env
 
     # Install CANN Kernels
-    echo "Installing ${KERNELS_FILE}"
+    _info "Installing ${KERNELS_FILE}"
     chmod +x "${KERNELS_PATH}"
     bash "${KERNELS_PATH}" --quiet --install --install-for-all --install-path="${CANN_HOME}"
     rm -f "${KERNELS_PATH}"
 
-    echo "CANN ${CANN_VERSION} installation successful."
+    _info "CANN ${CANN_VERSION} installation successful."
 }
 
 PLATFORM=${PLATFORM:=$(uname -s)/$(uname -m)}
-ARCH=$(get_architecture)
+ARCH=$(_get_architecture)
 CANN_HOME=${CANN_HOME:="/usr/local/Ascend"}
 CANN_CHIP=${CANN_CHIP:="910b"}
-CANN_VERSION=${CANN_VERSION:="8.0.RC1"}
+CANN_VERSION=${CANN_VERSION:="8.0.0"}
 
 # NOTE: kernels are arch-specific after 8.0.RC3.alpha002
-if [[ ${CANN_VERSION} == "8.0.RC3" ]]; then
+if [[ ${CANN_VERSION} == "8.0.RC3" || ${CANN_VERSION} == "8.0.0" ]]; then
   KERNELS_ARCH="linux-${ARCH}"
 else
   KERNELS_ARCH="linux"
@@ -146,6 +149,5 @@ elif [ "$1" == "--install" ]; then
 elif [ "$1" == "--set_env" ]; then
     set_env
 else
-    echo "Unexpected arguments, use '--download', '--install' or '--set_env' instead"
-    exit 1
+    _err "Unexpected arguments, use '--download', '--install' or '--set_env' instead"
 fi