add pre_configure hook that adds CUPTI lib dir

trz42 · May 25, 2024 · 2e78b71 · 2e78b71
1 parent e020246
commit 2e78b71
Showing 1 changed file with 24 additions and 0 deletions.
diff --git a/eb_hooks.py b/eb_hooks.py
@@ -369,6 +369,29 @@ def pre_configure_hook_openblas_optarch_generic(self, *args, **kwargs):
         raise EasyBuildError("OpenBLAS-specific hook triggered for non-OpenBLAS easyconfig?!")
 
 
+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':
+        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?!")
+
+
 def pre_configure_hook_libfabric_disable_psm3_x86_64_generic(self, *args, **kwargs):
     """Add --disable-psm3 to libfabric configure options when building with --optarch=GENERIC on x86_64."""
     if self.name == 'libfabric':
@@ -760,6 +783,7 @@ def inject_gpu_property(ec):
     'libfabric': pre_configure_hook_libfabric_disable_psm3_x86_64_generic,
     'MetaBAT': pre_configure_hook_metabat_filtered_zlib_dep,
     'OpenBLAS': pre_configure_hook_openblas_optarch_generic,
+    'PyTorch': pre_configure_hook_pytorch_add_cupti_libdir,
     'WRF': pre_configure_hook_wrf_aarch64,
     'at-spi2-core': pre_configure_hook_atspi2core_filter_ld_library_path,
 }