Add UCX-Py testing

buildlib/dockers/conda.Dockerfile

@@ -0,0 +1,106 @@
+ARG CUDA_VERSION=11.2
+ARG UBUNTU_VERSION=20.04
+FROM gpuci/miniconda-cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION}
+
+RUN apt-get update && \
+    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends tzdata && \
+    apt-get install -y \
+        automake \
+        dh-make \
+        g++ \
+        git \
+        libcap2 \
+        libnuma-dev \
+        libtool \
+        make \
+        udev \
+        wget \
+    && apt-get remove -y openjdk-11-* || apt-get autoremove -y \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+# MOFED
+ARG MOFED_VERSION=5.0-1.0.0.0
+ARG UBUNTU_VERSION
+ARG MOFED_OS=ubuntu${UBUNTU_VERSION}
+ENV MOFED_DIR MLNX_OFED_LINUX-${MOFED_VERSION}-${MOFED_OS}-x86_64
+ENV MOFED_SITE_PLACE MLNX_OFED-${MOFED_VERSION}
+ENV MOFED_IMAGE ${MOFED_DIR}.tgz
+RUN wget --no-verbose http://content.mellanox.com/ofed/${MOFED_SITE_PLACE}/${MOFED_IMAGE} && \
+    tar -xzf ${MOFED_IMAGE}
+RUN ${MOFED_DIR}/mlnxofedinstall --all -q \
+        --user-space-only \
+        --without-fw-update \
+        --skip-distro-check \
+        --without-ucx \
+        --without-hcoll \
+        --without-openmpi \
+        --without-sharp && \
+    rm -rf ${MOFED_DIR} && rm -rf *.tgz
+
+ENV CPATH /usr/local/cuda/include:${CPATH}
+ENV LD_LIBRARY_PATH /usr/local/cuda/lib64:/usr/local/cuda/compat:${LD_LIBRARY_PATH}
+ENV LIBRARY_PATH /usr/local/cuda/lib64:/usr/local/cuda/compat:${LIBRARY_PATH}
+ENV PATH /usr/local/cuda/compat:${PATH}
+
+# Required arguments
+ARG PYTHON_VER=3.8
+ARG BUILD_STACK_VER=9.3.0
+
+# Enables "source activate conda"
+SHELL ["/bin/bash", "-c"]
+
+# Add a condarc for channels and override settings
+RUN echo -e "\
+auto_update_conda: False \n\
+ssl_verify: False \n\
+channels: \n\
+  - gpuci \n\
+  - nvidia \n\
+  - conda-forge \n" > /opt/conda/.condarc \
+      && cat /opt/conda/.condarc
+
+# Update and add pkgs for gpuci builds
+RUN apt-get update -y --fix-missing \
+    && apt-get -qq install apt-utils -y --no-install-recommends \
+    && apt-get install -y \
+      chrpath \
+      file \
+      screen \
+      tzdata \
+      vim \
+      zlib1g-dev \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN apt-get update \
+    && apt-get -y upgrade \
+    && rm -rf /var/lib/apt/lists/*
+
+# Add core tools to base env
+RUN conda install -y gpuci-tools \
+    || conda install -y gpuci-tools
+
+# Create `ucx-py` conda env and make it default
+RUN gpuci_conda_retry create --no-default-packages --override-channels -n ucx-py \
+      -c nvidia \
+      -c conda-forge \
+      -c gpuci \
+      cudatoolkit=${CUDA_VERSION} \
+      git \
+      gpuci-tools \
+      libgcc-ng=${BUILD_STACK_VER} \
+      libstdcxx-ng=${BUILD_STACK_VER} \
+      python=${PYTHON_VER} \
+      'python_abi=*=*cp*' \
+      "setuptools<50" \
+    && sed -i 's/conda activate base/conda activate ucx-py/g' ~/.bashrc
+
+RUN source activate ucx-py \
+  && env \
+  && conda info \
+  && conda config --show-sources \
+  && conda list --show-channel-urls
+
+RUN chmod -R ugo+w /opt/conda \
+  && conda clean -tipy \
+  && chmod -R ugo+w /opt/conda

buildlib/pr/main.yml

@@ -8,6 +8,9 @@ resources:
     - container: fedora34
       image: rdmz-harbor.rdmz.labs.mlnx/ucx/fedora34:2
       options: -v /hpc/local:/hpc/local -v /auto/sw_tools:/auto/sw_tools
+    - container: conda
+      image: rdmz-harbor.rdmz.labs.mlnx/ucx/conda
+      options: -v /hpc/local:/hpc/local -v /auto/sw_tools:/auto/sw_tools
     - container: coverity_rh7
       image: rdmz-harbor.rdmz.labs.mlnx/ucx/coverity:mofed-5.1-2.3.8.0
       options: -v /hpc/local:/hpc/local -v /auto/sw_tools:/auto/sw_tools
@@ -173,3 +176,12 @@ stages:
         parameters:
             name: gpu
             demands: ucx_gpu -equals yes
+
+  - stage: python
+    dependsOn: [Static_check]
+    jobs:
+      - template: python/python-test.yml
+        parameters:
+          name: new
+          demands: ucx_docker -equals yes
+          container: conda

buildlib/pr/python/python-test.yml

@@ -0,0 +1,94 @@
+parameters:
+  name: python test
+  demands: []
+
+jobs:
+  - job: ${{ parameters.name }}
+
+    pool:
+      name: MLNX
+      demands: ${{ parameters.demands }}
+    container: conda
+
+    steps:
+      - checkout: self
+        fetchDepth: 100
+        clean: true
+        displayName: Checkout
+      - bash: |
+          set -x
+          source buildlib/az-helpers.sh
+          az_init_modules
+          try_load_cuda_env
+          set -eE
+          ./autogen.sh
+          ./contrib/configure-devel --prefix=$(Build.Repository.LocalPath)/install \
+            --enable-gtest=no --with-cuda=$have_cuda
+          make -j`nproc`
+          make install
+        displayName: Build UCX
+      - bash: |
+          set -x
+          source buildlib/az-helpers.sh
+          az_init_modules
+          try_load_cuda_env
+          set -eE
+
+          # Set path and build parallel level
+          export PATH=/opt/conda/bin:$PATH
+
+          ################################################################################
+          # SETUP - Install conda packages and check environment
+          ################################################################################
+          gpuci_logger "Check environment"
+          env
+
+          gpuci_logger "Activate conda env"
+          . /opt/conda/etc/profile.d/conda.sh
+          conda activate ucx-py
+          gpuci_conda_retry install \
+              "pytest" "pytest-asyncio" \
+              "dask" "distributed" \
+              "cupy=9" "numba" \
+              "cython" "libhwloc"
+
+          gpuci_logger "Clone and Build UCX-Py"
+          git clone https://github.com/rapidsai/ucx-py
+          cd ucx-py
+          python setup.py build_ext --inplace
+          python -m pip install -e .
+
+          gpuci_logger "Check versions"
+          python --version
+          conda info
+          conda config --show-sources
+          conda list --show-channel-urls
+          ucx_info -v
+
+          ################################################################################
+          # TEST - Run UCX-Py tests and benchmarks
+          ################################################################################
+          gpuci_logger "Check GPU usage"
+          nvidia-smi
+
+          test_tls="all"                       # All transports
+          test_tls+=" tcp,cuda_copy"           # Cuda without IPC
+          test_tls+=" tcp,cuda_copy,cuda_ipc"  # Cuda with IPC
+          if [ -z `get_rdma_interfaces` ]; then
+              test_tls+=" rc,cuda_copy"
+          fi
+
+          for tls in ${test_tls}; do
+              export UCX_TLS=$tls
+
+              # Test with TCP/Sockets
+              gpuci_logger "Tests (UCX_TLS=$UCX_TLS)"
+              pytest --cache-clear -vs ucp/_libs/tests
+              pytest --cache-clear -vs tests/
+
+              gpuci_logger "Benchmarks (UCX_TLS=$UCX_TLS)"
+              python benchmarks/send-recv.py -o cupy --server-dev 0 --client-dev 0 --reuse-alloc
+              python benchmarks/send-recv-core.py -o cupy --server-dev 0 --client-dev 0 --reuse-alloc
+          done
+
+        displayName: Run Python tests