From 8951b4f2f721fc50e5a4eba084f8f83935902a44 Mon Sep 17 00:00:00 2001 From: Thomas Roeblitz Date: Sat, 25 May 2024 11:13:53 +0200 Subject: [PATCH] improve hook to add CUPTI/lib64 dir to LIBRARY_PATH --- eb_hooks.py | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/eb_hooks.py b/eb_hooks.py index 7b5c082ddc..f01e1ca55c 100644 --- 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?!")