Cambricon PyTorch Quick Install (Ubuntu and MLU370 only)

install.sh

#!/usr/bin/env bash

############################### Usage ###############################
# Make sure this script is in the same directory as MLU370 SDK
############################## Example ##############################
# $ ls
# install.sh  MLU370_v1.20
# $ ./install.sh
#####################################################################

info() {
    echo -e "\e[1;32m$1\e[0m"
}

error() {
    echo -e "\e[1;31mError: $1\e[0m"
    exit 1
}

set -e

# Check OS
if [ "$(lsb_release -is)" != "Ubuntu" ]; then
    error "The current system is not Ubuntu"
fi

# Define matching files
ubuntu_version=$(lsb_release -rs)
wildcard_paths=(
    "MLU370_v*.*/driver_v*/Ubuntu/${ubuntu_version}/cambricon-mlu-driver-ubuntu${ubuntu_version}-dkms_*_amd64.deb"
    "MLU370_v*.*/cntoolkit_v*/Ubuntu/${ubuntu_version}/cntoolkit_*_amd64.deb"
    "MLU370_v*.*/cnnl_v*/Ubuntu/${ubuntu_version}/cnnl_*_amd64.deb"
    "MLU370_v*.*/cncl_v*/Ubuntu/${ubuntu_version}/cncl_*_amd64.deb"
    "MLU370_v*.*/cncv_v*/Ubuntu/${ubuntu_version}/cncv_*_amd64.deb"
    "MLU370_v*.*/dali_v*/cambricon_dali-*.whl"
    "MLU370_v*.*/cnnl_extra_v*/Ubuntu/${ubuntu_version}/cnnlextra_*_amd64.deb"
    "MLU370_v*.*/mlu_ops_v*/Ubuntu/${ubuntu_version}/mluops_*_amd64.deb"
)

# Check files
matched_files=()
for pattern in "${wildcard_paths[@]}"; do
    matches=($pattern)
    if [ ${#matches[@]} -eq 0 ] || [ ! -f "${matches[0]}" ]; then
        error "No matching file found: $pattern"
    fi
    for match in "${matches[@]}"; do
        absolute_path=$(realpath "$match")
        matched_files+=("$absolute_path")
    done
done

# Confirm
read -p "All checks have passed. Do you want to proceed with the installation? [Y/n]" choice
choice=${choice:-y}
if [[ "$choice" != "y" && "$choice" != "Y" ]]; then
    exit 0
fi

# Install
sudo apt update
sudo apt install -y python3 python3-pip
for file in "${matched_files[@]}"; do
    info "Installing $file"
    if [[ $file == *cntoolkit_*_amd64.deb ]]; then
        sudo apt install -y "$file"
        sudo apt update
        sudo apt install -y cntoolkit-cloud
        sudo rm -f /etc/apt/sources.list.d/cntoolkit.list        
    elif [[ $file == *.deb ]]; then
        sudo apt install -y "$file"
    elif [[ $file == *.whl ]]; then
        pip install "$file"
    fi
done

# Pytorch
# TODO: currently only install torch==2.4.0
whl_file=$(find MLU370_v*.*/pytorch_v*_torch2.4/wheel/torch_mlu-*.whl -type f 2>/dev/null | head -n 1)
current_python_version=$(python3 -c 'import sys; print(".".join(map(str, sys.version_info[:2])))')
requird_version=$(echo "$whl_file" | grep -oP 'cp\d+' | head -n 1 | sed 's/cp//')
requird_version_major=${requird_version:0:1}
requird_version_minor=${requird_version:1}
requird_version="$requird_version_major.$requird_version_minor"
if [[ "$PYTHON_VERSION" != "$REQUIRED_VERSION" ]]; then
    error "Current Python version: $PYTHON_VERSION, need $requird_version."
fi
info "Installing torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0"
pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cpu
info "Installing $whl_file"
pip install $whl_file

# Set environment variables
lines_to_add=(
    "export NEUWARE_HOME=/usr/local/neuware"
    "export PATH=\$NEUWARE_HOME/bin:\$PATH"
    "export LD_LIBRARY_PATH=\$NEUWARE_HOME/lib64:\$LD_LIBRARY_PATH"
)
bashrc="$HOME/.bashrc"
for line in "${lines_to_add[@]}"; do
    if ! grep -Fxq "$line" $bashrc; then
        echo "$line" >> $bashrc
        info "Added '$line' to $bashrc"
    fi
    eval "$line"
done

# Final check
python3 - <<EOF
import torch
import torch_mlu

GREEN = "\033[1;32m"
RED = "\033[1;31m"
RESET = "\033[0m"

try:
    a = torch.randn(2,3).mlu()
    a.abs()
    print(f"{GREEN}✅ torch_mlu is running successfully. Installation is complete!{RESET}")
except Exception as e:
    print(f"{RED}❌ Installation failed. Error message: {e}{RESET}")
    exit(1)
EOF