Added high-speed interconnect support by enabling UCX

recipe/build-mpi.sh

@@ -21,8 +21,8 @@ if [[ "$target_platform" == osx-* ]]; then
     fi
 fi
 
-if [ $cuda_compiler_version == '10.2' ]; then
-    build_with_cuda="--with-cuda"
+if [[ $cuda_compiler_version == 10.2 ]]; then
+    build_with_cuda="--with-cuda --with-ucx=$PREFIX"
 else
     build_with_cuda=""
 fi
@@ -159,6 +159,11 @@ if [ ! -z "$build_with_cuda" ]; then
     echo "opal_warn_on_missing_libcuda = 0" >> $PREFIX/etc/openmpi-mca-params.conf
     echo "setting the mca opal_cuda_support to 0..."
     echo "opal_cuda_support = 0" >> $PREFIX/etc/openmpi-mca-params.conf
+
+    echo "setting the mca pml to ^ucx..."
+    echo "pml = ^ucx" >> $PREFIX/etc/openmpi-mca-params.conf
+    echo "setting the mca osc to ^ucx..."
+    echo "osc = ^ucx" >> $PREFIX/etc/openmpi-mca-params.conf
 
     POST_LINK=$PREFIX/bin/.openmpi-post-link.sh
     cp $RECIPE_DIR/post-link.sh $POST_LINK

recipe/meta.yaml

@@ -16,7 +16,7 @@ source:
     - ompi-8361.diff
 
 build:
-  number: 4
+  number: 5
   # remember to bump the CUDA version in build-mpi.sh too!
   skip: true  # [win or (linux64 and cuda_compiler_version != '10.2') or ((aarch64 or ppc64le) and cuda_compiler_version not in (undefined, "None"))]
 
@@ -27,7 +27,7 @@ outputs:
       run_exports:
         - {{ pin_subpackage('openmpi', min_pin='x.x.x', max_pin='x') }}
       ignore_run_exports_from:
-        - {{ compiler('cuda') }}
+        - {{ compiler('cuda') }}  # [linux64]
     requirements:
       build:
         - {{ compiler('c') }}
@@ -41,11 +41,14 @@ outputs:
         - perl 5.26.2
       host:
         - zlib
+        - ucx 1.9.0  # [linux64]
+        - ucx-proc =*=gpu  # [linux64]
       run:
         - zlib
         - mpi 1.0 openmpi
       run_constrained:
         - cudatoolkit >=9.2  # [linux64]
+        - ucx 1.9.0  # [linux64]
     test:
       script: run_test.sh
       files:

recipe/post-link.sh

@@ -2,7 +2,16 @@
 
 echo " " >> $PREFIX/.messages.txt
 echo "For Linux 64, Open MPI is built with CUDA awareness but this support is disabled by default." >> $PREFIX/.messages.txt
-echo "To enable it, please set the environmental variable OMPI_MCA_opal_cuda_support=true before" >> $PREFIX/.messages.txt
+echo "To enable it, please set the environment variable OMPI_MCA_opal_cuda_support=true before" >> $PREFIX/.messages.txt
 echo "launching your MPI processes. Equivalently, you can set the MCA parameter in the command line:" >> $PREFIX/.messages.txt
 echo "mpiexec --mca opal_cuda_support 1 ..." >> $PREFIX/.messages.txt
 echo " " >> $PREFIX/.messages.txt
+
+echo "In addition, the UCX support is also built but disabled by default." >> $PREFIX/.messages.txt
+echo "To enable it, first install UCX (conda install -c conda-forge ucx). Then, set the environment" >> $PREFIX/.messages.txt
+echo "variables OMPI_MCA_pml=\"ucx\" OMPI_MCA_osc=\"ucx\" before launching your MPI processes." >> $PREFIX/.messages.txt
+echo "Equivalently, you can set the MCA parameters in the command line:" >> $PREFIX/.messages.txt
+echo "mpiexec --mca pml ucx --mca osc ucx ..." >> $PREFIX/.messages.txt
+echo "Note that you might also need to set UCX_MEMTYPE_CACHE=n for CUDA awareness via UCX." >> $PREFIX/.messages.txt
+echo "Please consult UCX's documentation for detail." >> $PREFIX/.messages.txt
+echo " " >> $PREFIX/.messages.txt