docker: add ucx configs to amd docker

devitocodes · Apr 19, 2023 · 5e8bab9 · 5e8bab9
1 parent 209f4c6
commit 5e8bab9
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/docker/Dockerfile.amd b/docker/Dockerfile.amd
@@ -113,10 +113,21 @@ ENV PATH=$OMPI_HOME/bin:$UCX_HOME/bin:$PATH \
     INCLUDE=$OMPI_HOME/include:$UCX_HOME/include:$INCLUDE \
     PKG_CONFIG_PATH=$OMPI_HOME/lib/pkgconfig:$UCX_HOME/lib/pkgconfig:$PKG_CONFIG_PATH
 
+
 # Adding environment variable for Running as ROOT and set MCA pml to UCX
-ENV OMPI_ALLOW_RUN_AS_ROOT=1  \
-    OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 \
-    OMPI_MCA_pml=ucx
+ENV OMPI_ALLOW_RUN_AS_ROOT=1
+ENV OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
+ENV OMPI_MCA_rmaps_base_oversubscribe=1
+ENV OMPI_MCA_btl_base_warn_component_unused=0
+ENV OMPI_MCA_hwloc_base_binding_policy=""
+ENV OMPI_MCA_pml=ucx
+# Add additional UCX environment variable
+ENV UCX_MEMTYPE_CACHE=no
+ENV UCX_NET_DEVICES=all
+ENV UCX_SHM_DEVICES=all
+ENV UCX_ACC_DEVICES=all
+ENV UCX_RNDV_THRESH=0
+ENV UCX_RNDV_SCHEME=get_zcopy
 
 ##############################################################
 # Additional Devito requirements

diff --git a/docker/Dockerfile.nvidia b/docker/Dockerfile.nvidia
@@ -62,6 +62,8 @@ ENV UCX_MEMTYPE_CACHE=no
 ENV UCX_NET_DEVICES=all
 ENV UCX_SHM_DEVICES=all
 ENV UCX_ACC_DEVICES=all
+ENV UCX_RNDV_THRESH=0
+ENV UCX_RNDV_SCHEME=get_zcopy
 ENV NCCL_UCX_RNDV_THRESH=0
 ENV NCCL_UCX_RNDV_SCHEME=get_zcopy
 ENV NCCL_PLUGIN_P2P=ucx