improve hook to add CUPTI/lib64 dir to LIBRARY_PATH

trz42 · May 25, 2024 · 8951b4f · 8951b4f
1 parent b053cbd
commit 8951b4f
Showing 1 changed file with 14 additions and 13 deletions.
diff --git a/eb_hooks.py b/eb_hooks.py
@@ -429,20 +429,21 @@ def pre_configure_hook_pytorch_add_cupti_libdir(self, *args, **kwargs):
     """
     Pre-configure hook for PyTorch: add directory $EESSI_SOFTWARE_PATH/software/CUDA/12.1.1/extras/CUPTI/lib64 to LIBRARY_PATH
     """
-    #if self.name == 'PyTorch' and 'CUDA' in self.versionsuffix:
     if self.name == 'PyTorch':
-        library_path = os.getenv('LIBRARY_PATH')
-        print_msg("library_path: '%s'", library_path)
-        #lib_dirs = library_path.split(':')
-        eessi_software_path = os.getenv('EESSI_SOFTWARE_PATH')
-        print_msg("eessi_software_path: '%s'", eessi_software_path)
-        #for lib_dir in lib_dirs:
-        #    if lib_dir.startswith(eessi_software_path + '/CUDA'):
-        #
-        cupti_lib_dir = os.path.join(eessi_software_path, 'software', 'CUDA', '12.1.1', 'extras', 'CUPTI', 'lib64')
-        print_msg("cupti_lib_dir: '%s'", cupti_lib_dir)
-        env.setvar('LIBRARY_PATH', ':'.join([library_path, cupti_lib_dir]))
-        print_msg("LIBRARY_PATH: '%s'", os.getenv('LIBRARY_PATH'))
+        if 'cudaver' in self.cfg.template_values and self.cfg.template_values['cudaver'] == '12.1.1':
+            _cudaver = self.cfg.template_values['cudaver']
+            print_msg("pre_configure_hook_pytorch_add_cupti_libdir: CUDA version: '%s'" % _cudaver)
+            _library_path = os.getenv('LIBRARY_PATH')
+            print_msg("pre_configure_hook_pytorch_add_cupti_libdir: library_path: '%s'", _library_path)
+            _eessi_software_path = os.getenv('EESSI_SOFTWARE_PATH')
+            print_msg("pre_configure_hook_pytorch_add_cupti_libdir: eessi_software_path: '%s'", _eessi_software_path)
+            _cupti_lib_dir = os.path.join(_eessi_software_path, 'software', 'CUDA', _cudaver, 'extras', 'CUPTI', 'lib64')
+            print_msg("pre_configure_hook_pytorch_add_cupti_libdir: cupti_lib_dir: '%s'", _cupti_lib_dir)
+            if _library_path:
+                env.setvar('LIBRARY_PATH', ':'.join([_library_path, _cupti_lib_dir]))
+            else:
+                env.setvar('LIBRARY_PATH', _cupti_lib_dir)
+            print_msg("pre_configure_hook_pytorch_add_cupti_libdir: LIBRARY_PATH: '%s'", os.getenv('LIBRARY_PATH'))
     else:
         raise EasyBuildError("PyTorch-specific hook triggered for non-PyTorch easyconfig?!")