@@ -162,21 +162,19 @@ def pre_module_hook(self, *args, **kwargs):
162162 self .log .info ("[pre-module hook] Set I_MPI_PMI_LIBRARY in impi module" )
163163 # Must be done this way, updating self.cfg['modextravars']
164164 # directly doesn't work due to templating.
165- en_templ = self .cfg .enable_templating
166- self .cfg .enable_templating = False
167- shlib_ext = get_shared_lib_ext ()
168- pmix_root = get_software_root ('PMIx' )
169- if pmix_root :
170- mpi_type = 'pmix_v3'
171- self .cfg ['modextravars' ].update ({
172- 'I_MPI_PMI_LIBRARY' : os .path .join (pmix_root , "lib" , "libpmi." + shlib_ext )
173- })
174- self .cfg ['modextravars' ].update ({'SLURM_MPI_TYPE' : mpi_type })
175- # Unfortunately UCX doesn't yet work for unknown reasons. Make sure it is off.
176- self .cfg ['modextravars' ].update ({'SLURM_PMIX_DIRECT_CONN_UCX' : 'false' })
177- else :
178- self .cfg ['modextravars' ].update ({'I_MPI_PMI_LIBRARY' : "/lap/slurm/lib/libpmi.so" })
179- self .cfg .enable_templating = en_templ
165+ with self .cfg .disable_templating ():
166+ shlib_ext = get_shared_lib_ext ()
167+ pmix_root = get_software_root ('PMIx' )
168+ if pmix_root :
169+ mpi_type = 'pmix_v3'
170+ self .cfg ['modextravars' ].update ({
171+ 'I_MPI_PMI_LIBRARY' : os .path .join (pmix_root , "lib" , "libpmi." + shlib_ext )
172+ })
173+ self .cfg ['modextravars' ].update ({'SLURM_MPI_TYPE' : mpi_type })
174+ # Unfortunately UCX doesn't yet work for unknown reasons. Make sure it is off.
175+ self .cfg ['modextravars' ].update ({'SLURM_PMIX_DIRECT_CONN_UCX' : 'false' })
176+ else :
177+ self .cfg ['modextravars' ].update ({'I_MPI_PMI_LIBRARY' : "/lap/slurm/lib/libpmi.so" })
180178
181179 if self .name == 'OpenBLAS' :
182180 self .log .info ("[pre-module hook] Set OMP_NUM_THREADS=1 in OpenBLAS module" )
@@ -197,18 +195,14 @@ def pre_module_hook(self, *args, **kwargs):
197195 self .log .info ("[pre-module hook] Set SLURM_MPI_TYPE=%s in OpenMPI module" % mpi_type )
198196 # Must be done this way, updating self.cfg['modextravars']
199197 # directly doesn't work due to templating.
200- en_templ = self .cfg .enable_templating
201- self .cfg .enable_templating = False
202- self .cfg ['modextravars' ].update ({'SLURM_MPI_TYPE' : mpi_type })
203- # Unfortunately UCX doesn't yet work for unknown reasons. Make sure it is off.
204- self .cfg ['modextravars' ].update ({'SLURM_PMIX_DIRECT_CONN_UCX' : 'false' })
205- self .cfg .enable_templating = en_templ
198+ with self .cfg .disable_templating ():
199+ self .cfg ['modextravars' ].update ({'SLURM_MPI_TYPE' : mpi_type })
200+ # Unfortunately UCX doesn't yet work for unknown reasons. Make sure it is off.
201+ self .cfg ['modextravars' ].update ({'SLURM_PMIX_DIRECT_CONN_UCX' : 'false' })
206202
207203 if self .name == 'PMIx' :
208204 # This is a, hopefully, temporary workaround for https://github.com/pmix/pmix/issues/1114
209205 if LooseVersion (self .version ) > LooseVersion ('2' ) and LooseVersion (self .version ) < LooseVersion ('3' ):
210206 self .log .info ("[pre-module hook] Set PMIX_MCA_gds=^ds21 in PMIx module" )
211- en_templ = self .cfg .enable_templating
212- self .cfg .enable_templating = False
213- self .cfg ['modextravars' ].update ({'PMIX_MCA_gds' : '^ds21' })
214- self .cfg .enable_templating = en_templ
207+ with self .cfg .disable_templating ():
208+ self .cfg ['modextravars' ].update ({'PMIX_MCA_gds' : '^ds21' })
0 commit comments