From a91ae195ec0ddfded5840e77085581c6166bc564 Mon Sep 17 00:00:00 2001
From: Damian Alvarez <swmanage@hdfmll01.hdfml>
Date: Thu, 14 Apr 2022 15:52:11 +0200
Subject: [PATCH] To make the UCX module a bit smarter, and try to keep CUDA
settings if they were selected before
---
Custom_Hooks/eb_hooks.py | 11 +++++++++--
1 file changed, 9 insertions(+), 2 deletions(-)
diff --git a/Custom_Hooks/eb_hooks.py b/Custom_Hooks/eb_hooks.py
index 0e7663c6b..3e3e3b931 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:
--
GitLab