From 45046752faa327d5990e00b3006c1f6f188c3924 Mon Sep 17 00:00:00 2001 From: Thomas Roeblitz Date: Fri, 22 Nov 2024 09:12:10 +0100 Subject: [PATCH] add pre_configure hook for PyTorch to add CUPTI lib64 dir to LIBRARY_PATH --- eb_hooks.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/eb_hooks.py b/eb_hooks.py index 03642656ea..63e79ebf51 100644 --- a/eb_hooks.py +++ b/eb_hooks.py @@ -604,6 +604,32 @@ def pre_configure_hook_LAMMPS_zen4(self, *args, **kwargs): raise EasyBuildError("LAMMPS-specific hook triggered for non-LAMMPS 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' and self.version == '2.1.2': + 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: + print_msg("PyTorch/2.1.2-specific pre_configure hook triggered for non-CUDA or non-CUDA/12.1.1 easyconfig triggered; NOT adding CUPTI lib64 dir to LIBRARY_PATH") + else: + raise EasyBuildError("PyTorch/2.1.2-specific pre_configure hook triggered for non-PyTorch/2.1.2 easyconfig?!") + + def pre_test_hook(self, *args, **kwargs): """Main pre-test hook: trigger custom functions based on software name.""" if self.name in PRE_TEST_HOOKS: @@ -995,6 +1021,7 @@ def inject_gpu_property(ec): 'OpenBLAS': pre_configure_hook_openblas_optarch_generic, 'WRF': pre_configure_hook_wrf_aarch64, 'LAMMPS': pre_configure_hook_LAMMPS_zen4, + 'PyTorch': pre_configure_hook_pytorch_add_cupti_libdir, 'Score-P': pre_configure_hook_score_p, }