diff --git a/Custom_Hooks/eb_hooks.py b/Custom_Hooks/eb_hooks.py index 0e7663c6b23501d3d5082c0a38ac8aa2b3c0a4f7..3e3e3b931e9b32d3b02935e4c21f03c14e6a8066 100644 --- a/Custom_Hooks/eb_hooks.py +++ b/Custom_Hooks/eb_hooks.py @@ -404,8 +404,15 @@ def inject_ucx_tweaks(ec): if ec.name in 'UCX' and install_path().lower().startswith('/p/software'): key = "modluafooter" value = ''' -if not ( isloaded("UCX-settings") ) then - load("UCX-settings") +-- This weird construct is to prevent lmod from loading the default settings when reloading/swapping UCX +-- Unfortunately we can't do better (like preserving the UCX transport for CUDA, since the loaded +-- UCX-settings module is lost during the reload +if mode()=="load" then + if isloaded("mpi-settings/CUDA") then + try_load("UCX-settings/RC-CUDA") + else + try_load("UCX-settings") + end end ''' if key in ec_dict: