[WIP][DO NOT MERGE] CANN: Improve docker image

cann/openeuler/Dockerfile

@@ -46,7 +46,8 @@ RUN yum install -y \
         pciutils \
         net-tools \
         sqlite-devel \
-        lapack-devel gcc-gfortran \
+        lapack-devel \
+        gcc-gfortran \
         python3-devel \
         python3-pip \
         wget \

cann/ubuntu/Dockerfile

@@ -11,21 +11,22 @@ ARG PLATFORM
 ARG CANN_CHIP
 ARG CANN_VERSION
 
-RUN apk add --no-cache \
+RUN sed -i 's/dl-cdn.alpinelinux.org/mirrors.tuna.tsinghua.edu.cn/g' /etc/apk/repositories && \
+    apk add --no-cache \
         bash \
         curl \
     && rm -rf /var/cache/apk/* \
     && rm -rf /tmp/*
 
 # Copy files
-COPY ../scripts/cann.sh /tmp/cann.sh
+COPY ./cann.sh /tmp/cann.sh
 
 # Download CANN
 RUN chmod +x /tmp/cann.sh && \
     bash /tmp/cann.sh --download
 
 # Phase 2: install CANN toolkit and kernels
-FROM ubuntu:${BASE_VERSION} as builder
+FROM ubuntu:${BASE_VERSION} as installer
 
 ARG PLATFORM
 ARG CANN_CHIP
@@ -34,12 +35,9 @@ ARG CANN_VERSION
 # Change the default shell
 SHELL [ "/bin/bash", "-c" ]
 
-# Set non-interactive mode for apt-get
-ENV DEBIAN_FRONTEND=noninteractive
-
 # Install dependencies
 RUN apt-get update \
-    && apt-get install --no-install-recommends -y \
+    && DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \
         apt-transport-https \
         ca-certificates \
         bash \
@@ -65,21 +63,48 @@ RUN apt-get update \
         python3 \
         python3-pip \
         python3-dev \
-        vim \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/* \
     && rm -rf /var/tmp/* \
     && rm -rf /tmp/*
 
+COPY --from=downloader /tmp /tmp
+
 # Install CANN
-RUN --mount=type=cache,from=downloader,target=/tmp,source=/tmp \
-    umask 0022 && \
-    chmod +x /tmp/cann.sh && \
+RUN chmod +x /tmp/cann.sh && \
     bash /tmp/cann.sh --install && \
     rm -rf /tmp/*
 
+# Phase 3:
+FROM ubuntu:${BASE_VERSION} as official
+
+# Change the default shell
+SHELL [ "/bin/bash", "-c" ]
+
+# Install dependencies
+RUN apt-get update \
+    && DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \
+        apt-transport-https \
+        ca-certificates \
+        bash \
+        openssl \
+        libblas3 \
+        python3 \
+        python3-pip \
+        python3-dev \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/* \
+    && rm -rf /var/tmp/* \
+    && rm -rf /tmp/*
+
 # Copy files
-COPY ../scripts/docker-entrypoint.sh /root/docker-entrypoint.sh
+COPY --from=installer /usr/local/Ascend /usr/local/Ascend
+COPY --from=installer /etc/Ascend /etc/Ascend
+COPY ./docker-entrypoint.sh /root/docker-entrypoint.sh
+
+# Install python packages
+RUN pip install --no-cache-dir attrs cython numpy decorator sympy cffi pyyaml pathlib2 psutil protobuf scipy requests absl-py \
+    -i https://repo.huaweicloud.com/repository/pypi/simple
 
 # Driver path
 ENV DRIVER_PATH=/usr/local/Ascend/driver

cann/ubuntu/cann.sh

@@ -0,0 +1,150 @@
+#!/bin/bash
+
+set -e
+
+get_architecture() {
+    # not case sensitive
+    shopt -s nocasematch
+
+    case "${PLATFORM}" in
+        "linux/x86_64"|"linux/amd64")
+            ARCH="x86_64"
+            ;;
+        "linux/aarch64"|"linux/arm64")
+            ARCH="aarch64"
+            ;;
+        *)
+            echo "Error: Unsupported architecture ${PLATFORM}."
+            exit 1
+            ;;
+    esac
+
+    echo "${ARCH}"
+}
+
+download_file() {
+    set +e
+
+    local max_retries=10
+    local retry_delay=10
+
+    local url="$1"
+    local path="$2"
+
+    for ((i=1; i<=max_retries; i++)); do
+        echo "Attempt $i of $max_retries..."
+
+        curl -fsSL "${url}" -o "${path}"
+
+        if [[ $? -eq 0 ]]; then
+            return 0
+        else
+            echo "Download failed with error code $?. Retrying in ${retry_delay} seconds..."
+            sleep ${retry_delay}
+        fi
+    done
+
+    echo "All attempts failed. Exiting."
+    return 1
+}
+
+download_cann() {
+    if [[ ${CANN_VERSION} =~ ^8.0.RC2.alpha ]]; then
+        local url_prefix="https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/Milan-ASL/Milan-ASL%20V100R001C18SPC805"
+    else
+        local url_prefix="https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/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}"
+
+    if [ ! -f "${TOOLKIT_PATH}" ]; then
+        echo "Downloading ${TOOLKIT_FILE} from ${toolkit_url}"
+        download_file "${toolkit_url}" "${TOOLKIT_PATH}"
+    fi
+
+    if [ ! -f "${KERNELS_PATH}" ]; then
+        echo "Downloading ${KERNELS_FILE} from ${kernels_url}"
+        download_file "${kernels_url}" "${KERNELS_PATH}"
+    fi
+
+    echo "CANN ${CANN_VERSION} download successful."
+}
+
+check_python() {
+    if ! command -v python &> /dev/null; then
+        if command -v python3 &> /dev/null; then
+            # Create symbolic link from python3 to python
+            ln -sf "$(command -v python3)" "$(dirname "$(command -v python3)")/python"
+            ln -sf "$(command -v pip3)" "$(dirname "$(command -v pip3)")/pip"
+            echo "Created symbolic link 'python' pointing to 'python3'."
+        else
+            echo "Python not installed."
+            exit 1
+        fi
+    fi
+}
+
+install_cann() {
+    # Check python
+    check_python
+
+    # Install dependencies
+    pip config set global.index-url https://repo.huaweicloud.com/repository/pypi/simple
+    pip install --no-cache-dir attrs cython numpy decorator sympy cffi pyyaml pathlib2 psutil protobuf scipy requests absl-py
+
+    # Download installers
+    if [ ! -f "${TOOLKIT_PATH}" ] || [ ! -f "${KERNELS_PATH}" ]; then
+        echo "[WARNING] Installers do not exist, re-download them."
+        download_cann
+    fi
+
+    # Install CANN Toolkit
+    echo "Installing ${TOOLKIT_FILE}"
+    chmod +x "${TOOLKIT_PATH}"
+    bash "${TOOLKIT_PATH}" --quiet --install --install-for-all --install-path="${CANN_HOME}"
+    rm -f "${TOOLKIT_PATH}"
+
+    # Set environment variables
+    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
+    else
+        cat >> /etc/profile <<'EOF'
+if [ -n "${DRIVER_PATH}" ]; then
+    export LD_LIBRARY_PATH=${DRIVER_PATH}/lib64/common/:${DRIVER_PATH}/lib64/driver/:${LD_LIBRARY_PATH}
+fi
+EOF
+        echo "source ${CANN_TOOLKIT_ENV_FILE}" >> /etc/profile
+        source ${CANN_TOOLKIT_ENV_FILE}
+    fi
+
+    # Install CANN Kernels
+    echo "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."
+}
+
+PLATFORM=${PLATFORM:=$(uname -s)/$(uname -m)}
+ARCH=$(get_architecture)
+CANN_HOME=${CANN_HOME:="/usr/local/Ascend"}
+CANN_CHIP=${CANN_CHIP:="910b"}
+CANN_VERSION=${CANN_VERSION:="8.0.RC1"}
+
+TOOLKIT_FILE="Ascend-cann-toolkit_${CANN_VERSION}_linux-${ARCH}.run"
+KERNELS_FILE="Ascend-cann-kernels-${CANN_CHIP}_${CANN_VERSION}_linux.run"
+TOOLKIT_PATH="/tmp/${TOOLKIT_FILE}"
+KERNELS_PATH="/tmp/${KERNELS_FILE}"
+
+if [ "$1" == "--download" ]; then
+    download_cann
+elif [ "$1" == "--install" ]; then
+    install_cann
+else
+    echo "Unexpected arguments, use '--download' or '--install' instead"
+    exit 1
+fi

cann/ubuntu/docker-entrypoint.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+set -e
+
+# Add the driver path to /etc/environment to make sure it can take effect
+# when connecting to the container
+printenv | grep DRIVER_PATH >> /etc/environment
+
+# Set environment variables for the Ascend toolkit
+source /usr/local/Ascend/ascend-toolkit/set_env.sh
+
+# Execute other commands
+exec "$@"