Fix an issue of mpi4py (NVIDIA#475)

Author: Shixiaowei02

docker/Dockerfile.multi

@@ -29,6 +29,10 @@ ENV LD_LIBRARY_PATH=/usr/local/tensorrt/lib:${LD_LIBRARY_PATH}
 COPY docker/common/install_polygraphy.sh install_polygraphy.sh
 RUN bash ./install_polygraphy.sh && rm install_polygraphy.sh
 
+# Install mpi4py
+COPY docker/common/install_mpi4py.sh install_mpi4py.sh
+RUN bash ./install_mpi4py.sh && rm install_mpi4py.sh
+
 # Install PyTorch
 ARG TORCH_INSTALL_TYPE="skip"
 COPY docker/common/install_pytorch.sh install_pytorch.sh

docker/common/install_base.sh

@@ -18,7 +18,6 @@ init_ubuntu() {
         apt-get remove --purge -y tensorrt*
     fi
     pip uninstall -y tensorrt
-    pip install mpi4py
 }
 
 install_gcc_centos() {

docker/common/install_mpi4py.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+set -ex
+
+MPI4PY_VERSION="3.1.5"
+RELEASE_URL="https://github.com/mpi4py/mpi4py/archive/refs/tags/${MPI4PY_VERSION}.tar.gz"
+curl -L ${RELEASE_URL} | tar -zx -C /tmp
+# Bypassing compatibility issues with higher versions (>= 69) of setuptools.
+sed -i 's/>= 40\.9\.0/>= 40.9.0, < 69/g' /tmp/mpi4py-${MPI4PY_VERSION}/pyproject.toml
+pip install /tmp/mpi4py-${MPI4PY_VERSION}
+rm -rf /tmp/mpi4py*