Swift for Tensorflow Ubuntu 18.04 idempotent install script

install_s4tf.sh

#! /usr/bin/env bash
#
# This is an install script to get a working Swift for Tensorflow
# environment with minimal fuss. Assumes Ubuntu 18.04 and bash.
#
# The script is idempotent so it can be safely run multiple times.
#
# The script uses an existing Anaconda/Miniconda installation
# or installs Miniconda to ~/miniconda3 if missing
#
# The script installs a bunch of system and conda packages so that
# all of the fastai_dev notebooks work (as much as possible).
#
# If Visual Studio Code and npm are present, a corresponding
# Swift extension is installed and configured to use sourcekit-lsp
#
# Notes:
# - This installs the latest nightly swift-tensorflow CUDA 10/CuDNN7 version
# - CUDA and CuDNN are installed as conda packages.
# - The conda Python is version 3.6 (3.7 might be compatible with swift-jupyter too)
#
set -Eeuo pipefail # fail early in case of errors

# Main script parameters:
CONDA_ENV=swiftai # conda environment name. Leave blank to use the default environment ("CONDA_ENV=")
BASE_DIR=~/swift
VSCODE_EXT=1 # 1 to install sourcekit-lsp vscode extension, any other value to disable

# See: https://github.com/tensorflow/swift/blob/master/notebooks/install_latest_swift.ipynb

apt_packages=(
    # General
    ubuntu-drivers-common libvorbis-dev libflac-dev libsndfile-dev cmake
    build-essential libgflags-dev libgoogle-glog-dev libgtest-dev google-mock
    zlib1g-dev libeigen3-dev libboost-all-dev libasound2-dev libogg-dev libtool
    libfftw3-dev libbz2-dev liblzma-dev libgoogle-glog0v5 gcc-6 gfortran-6
    g++-6 doxygen graphviz libsox-fmt-all parallel python3-pip jq curl
    exuberant-ctags vim-nox python-powerline

    # The latest Nvidia driver, marked as recommend by Ubuntu.
    nvidia-driver-430
    
    # General + Swift for Tensorflow
    git cmake ninja-build clang python uuid-dev libicu-dev icu-devtools
    libbsd-dev libedit-dev libxml2-dev libsqlite3-dev swig libpython-dev
    libncurses5-dev pkg-config libblocksruntime-dev libcurl4-openssl-dev
    systemtap-sdt-dev tzdata rsync

    # sourcekit-lsp
    cmake ninja-build clang systemtap-sdt-dev libbsd-dev linux-libc-dev

    # NB: libc6-dbg causes some warnings if installed!
)

set -x # Add some visibility to sudo commands

command -v apt-fast >/dev/null || sudo add-apt-repository --no-update -y ppa:apt-fast/stable
sudo add-apt-repository --no-update -y ppa:graphics-drivers/ppa
sudo apt-get update
command -v apt-fast >/dev/null || sudo apt-get -y install apt-fast

sudo apt-fast -y upgrade
sudo apt-fast -y install "${apt_packages[@]}"

ubuntu-drivers devices
sudo modprobe nvidia

# Because of the "set -e" flag, the whole script would stop executing at this point without the explicit "|| true"
# It's not super important if this fails. A reboot or something might fix it.
nvidia-smi || true

set +x # No need to echo everything now


TOOLCHAIN_PATH=$BASE_DIR/toolchain

# Download and extract the swift toolchain
url=https://storage.googleapis.com/swift-tensorflow-artifacts/nightlies/latest/swift-tensorflow-DEVELOPMENT-cuda10.0-cudnn7-ubuntu18.04.tar.gz
hash=$(curl -fs --head "$url" | grep '^x-goog-hash: md5' | cut -d' ' -f2)
cur_hash=$(cat "$BASE_DIR/.toolchain_hash" 2>/dev/null || true)
if [[ -d $TOOLCHAIN_PATH ]] && [[ $cur_hash = "$hash" ]]; then
    echo "swift-tensorflow not changed on server: Skipping update"
else
    rm -rf "$TOOLCHAIN_PATH"
    mkdir -p "$TOOLCHAIN_PATH"
    curl -f "$url" | tar xz -C "$TOOLCHAIN_PATH"
    echo "$hash" > "$BASE_DIR/.toolchain_hash"
fi &

# Install miniconda if there is no version of conda available:
if ! command -v conda >/dev/null; then
    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
    bash Miniconda3-latest-Linux-x86_64.sh -b

    # Ensure the installation comes into effect immediately
    . ~/miniconda3/etc/profile.d/conda.sh

    # And in subsequent shells
    cat >> ~/.bashrc <<EOF
[[ -r ~/miniconda3/etc/profile.d/conda.sh ]] && 
    . ~/miniconda3/etc/profile.d/conda.sh
[[ -z "$CONDA_ENV" ]] && conda activate
[[ -n "$CONDA_ENV" ]] && conda activate "$CONDA_ENV"
# PKG_CONFIG_PATH is required for making use of headers through conda packages, e.g. opencv4.
export PKG_CONFIG_PATH=\$CONDA_PREFIX/lib/pkgconfig
EOF
fi &

wait

set +u # conda uses some potentially uninitialized env vars
eval "$("$CONDA_EXE" shell.bash hook)"
[[ -z $CONDA_ENV ]] && conda activate
[[ -n $CONDA_ENV ]] && conda activate "$CONDA_ENV"
set -u

(
    # According to swift-jupyter, only Python 3.6 is supported. Python 3.7 probably works also.
    # However, if I don't specify a Python version, it will default to Python 2.7!!
    conda install -y -c conda-forge -c pytorch -c fastai fastai pytorch jupyter python=3.6 cuda100 cudnn opencv=4

    # Download and register the swift kernel
    [[ -d $BASE_DIR/swift-jupyter ]] ||
        git clone https://github.com/google/swift-jupyter.git "$BASE_DIR/swift-jupyter"

    cd "$BASE_DIR/swift-jupyter";
    git fetch # NB: currently not auto-updating

    python "$BASE_DIR/swift-jupyter/register.py" \
        --sys-prefix --swift-python-use-conda --use-conda-shared-libs \
        --swift-toolchain "$TOOLCHAIN_PATH"

    pip install jupyter_contrib_nbextensions
    jupyter contrib nbextension install --user
) &

(
    [[ -d $BASE_DIR/sourcekit-lsp ]] ||
        git clone https://github.com/apple/sourcekit-lsp "$BASE_DIR/sourcekit-lsp"

    cd "$BASE_DIR/sourcekit-lsp"
    git fetch # NB: currently not auto-updating

    swift package update
    swift build
    # As per the official instructions. However, this didn't work for me.
    #swift build -Xcxx -I"$BASE_PATH/toolchain/usr/lib/swift" -Xcxx -I"$BASE_PATH/toolchain/usr/lib/swift/Block"

    cp "$BASE_DIR/sourcekit-lsp/.build/x86_64-unknown-linux/debug/sourcekit-lsp" "$TOOLCHAIN_PATH/usr/bin"

    if [[ $VSCODE_EXT = 1 ]] && command -v code >/dev/null && command -v npm >/dev/null; then
        cd Editors/vscode
        npm run createDevPackage
        code --install-extension out/sourcekit-lsp-vscode-dev.vsix
        code --install-extension vknabel.vscode-swift-development-environment

        # Attempt to auto-configure the latter extension:
        code_config_path=~/.config/Code/User/settings.json
        if [[ -f $code_config_path ]]; then
            jq '. + {
                "sourcekit-lsp.toolchainPath": "'"$TOOLCHAIN_PATH"'"
                ,"swift.path.swift_driver_bin": "'"$TOOLCHAIN_PATH/usr/bin/swift"'"
                ,"sde.languageServerMode": "sourcekit-lsp"
                ,"sourcekit-lsp.serverPath": "'"$TOOLCHAIN_PATH/usr/bin/sourcekit-lsp"'"
            }' "$code_config_path" > "$code_config_path".new

            mv "$code_config_path" "$code_config_path".bak
            mv "$code_config_path".new "$code_config_path"
        fi
    fi
) &

(
    [[ -d $BASE_DIR/swift-corelibs-libdispatch ]] ||
        git clone https://github.com/apple/swift-corelibs-libdispatch "$BASE_DIR/swift-corelibs-libdispatch"

    cd "$BASE_DIR/swift-corelibs-libdispatch"
    git fetch # NB: currently not auto-updating

    rm -rf build/
    mkdir build/
    cd build/
    cmake -G Ninja -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ ..
    ninja
    ninja install
) &

wait

if ! grep -q '# swift toolchain path' ~/.profile; then
    cat >> ~/.profile <<EOF
PATH="$TOOLCHAIN_PATH/usr/bin:\$PATH" # swift toolchain path
EOF
fi