Script for setup of a UCX-enabled conda environment

ucx_setup.sh

#!/bin/bash
SOURCE_PATH=${HOME}/ucx-src
CONDA_ENV_NAME=ucx
RAPIDS_VERSION="0.9"
CUDA_VERSION="9.2"
# CUDA_VERSION="10.0"

if [ ! -d ${SOURCE_PATH} ]; then
    mkdir ${SOURCE_PATH}
fi

# Build cudf from source
cd ${SOURCE_PATH}
git clone https://github.com/rapidsai/cudf.git
cd ${SOURCE_PATH}/cudf
git checkout branch-${RAPIDS_VERSION}
git submodule update --init --remote --recursive
# This will by default create cudf_dev environment
# You can edit the YAML file to change that name
cp conda/environments/cudf_dev_cuda${CUDA_VERSION}.yml ${CONDA_ENV_NAME}_cuda${CUDA_VERSION}.yml
sed -i 's/name: cudf_env/name: '${CONDA_ENV_NAME}'/g' ${CONDA_ENV_NAME}_cuda${CUDA_VERSION}.yml
conda env create --name ${CONDA_ENV_NAME} --file ${CONDA_ENV_NAME}_cuda${CUDA_VERSION}.yml
source $(conda info --base)/bin/activate ${CONDA_ENV_NAME}
CUDACXX=/usr/local/cuda-${CUDA_VERSION}/bin/nvcc ./build.sh libcudf
CUDACXX=/usr/local/cuda-${CUDA_VERSION}/bin/nvcc ./build.sh cudf
CUDACXX=/usr/local/cuda-${CUDA_VERSION}/bin/nvcc ./build.sh dask_cudf

# Install required packages for ucx build
conda install -y python=3.7 libtool cmake automake autoconf cython bokeh pytest pkg-config

# Install additional packages
conda install -y -c conda-forge numba=0.45.1 arrow=0.14.1
conda install -y -c conda-forge boost
conda install -y nvidia::nccl=2 jakirkham/label/cupy::cupy

# Build and install ucx
cd ${SOURCE_PATH}
git clone https://github.com/openucx/ucx
cd ${SOURCE_PATH}/ucx
git remote add Akshay-Venkatesh https://github.com/Akshay-Venkatesh/ucx.git
git remote update Akshay-Venkatesh
git checkout ucx-cuda
git clean -xfd
export CUDA_HOME=/usr/local/cuda-${CUDA_VERSION}/
./autogen.sh
mkdir build
cd build
../configure --prefix=$CONDA_PREFIX --enable-debug --with-cuda=$CUDA_HOME --enable-mt --disable-cma CPPFLAGS="-I//usr/local/cuda-${CUDA_VERSION}/include"
make -j install

# Build and install ucx-py
cd ${SOURCE_PATH}
git clone https://github.com/rapidsai/ucx-py
cd ${SOURCE_PATH}/ucx-py
export UCX_PATH=$CONDA_PREFIX
make install

# Build and install rmm
cd ${SOURCE_PATH}
git clone https://github.com/rapidsai/rmm
cd ${SOURCE_PATH}/rmm
git submodule update --init --remote --recursive
git checkout branch-${RAPIDS_VERSION}
CUDACXX=/usr/local/cuda-${CUDA_VERSION}/bin/nvcc ./build.sh librmm rmm
cd python
pip install .

# Install dask-cuda
cd ${SOURCE_PATH}
git clone https://github.com/rapidsai/dask-cuda
cd ${SOURCE_PATH}/dask-cuda
git checkout branch-${RAPIDS_VERSION}
pip install .

# Install distributed with unmerged cudf serialization fixes
cd ${SOURCE_PATH}
git clone https://github.com/dask/distributed
cd ${SOURCE_PATH}/distributed
git remote add quasiben https://github.com/quasiben/distributed.git
git remote update
git checkout fix/update-cudf-serialization
pip install .