diff --git a/eb_hooks.py b/eb_hooks.py index 67b031bc..33cb757c 100644 --- a/eb_hooks.py +++ b/eb_hooks.py @@ -17,6 +17,7 @@ from easybuild.tools import config from easybuild.tools.build_log import EasyBuildError, print_msg, print_warning from easybuild.tools.config import build_option, install_path, update_build_option +from easybuild.tools.config import ERROR from easybuild.tools.filetools import apply_regex_substitutions, copy_dir, copy_file, remove_file, symlink, which from easybuild.tools.modules import get_software_root, get_software_root_env_var_name from easybuild.tools.run import run_cmd @@ -49,6 +50,7 @@ EESSI_FORCE_ATTR = 'orig_force' EESSI_SUPPORTED_MODULE_ATTR = 'eessi_supported_module' EESSI_UNSUPPORTED_MODULE_ATTR = 'eessi_unsupported_module' +EESSI_SANITYCHECK_CUDA_ATTR = 'eessi_sanitycheck_initial_env_path' SYSTEM = EASYCONFIG_CONSTANTS['SYSTEM'][0] @@ -957,6 +959,64 @@ def post_module_hook_unsupported_module(self, *args, **kwargs): del self.initial_environ[unsup_mod.envvar] +def pre_sanitycheck_hook(self, *args, **kwargs): + """Main pre-sanitycheck hook: trigger custom functions.""" + pre_sanitycheck_hook_cuda(self, *args, **kwargs) + + +def post_sanitycheck_hook(self, *args, **kwargs): + """Main post-sanitycheck hook: trigger custom functions.""" + post_sanitycheck_hook_cuda(self, *args, **kwargs) + + +def pre_sanitycheck_hook_cuda(self, *args, **kwargs): + """ + If CUDA is a build-only dependency (demoted to build dep by inject_gpu_property), + temporarily load the module so CUDA tool cuobjdump path can be made + available during the sanity check step. + This is needed since EasyBuild 5.4.0 where build deps are no longer available + during the sanity check. + """ + cudaver = get_dependency_software_version("CUDA", ec=self.cfg, check_deps=False, check_builddeps=True) + if cudaver and EASYBUILD_VERSION >= '5.4.0': + # Get the CUDA dependency info from builddependencies + build_deps = self.cfg.get_ref('builddependencies') + for dep in build_deps: + if dep['name'] == 'CUDA': + # Load CUDA module + self.modules_tool.load([dep['full_mod_name']]) + cuobjdump_path = which('cuobjdump', on_error=ERROR) + cuobjdump_dir = os.path.dirname(cuobjdump_path) + self.cuobjdump_dir = cuobjdump_dir + self.modules_tool.unload([dep['full_mod_name']]) + # Store the original PATH for restoration in post hook + original_initial_environ_path = self.initial_environ['PATH'] + setattr(self, EESSI_SANITYCHECK_CUDA_ATTR, original_initial_environ_path) + # Modify stored initial environment (restored during the sanity check), not the current environment + self.initial_environ['PATH'] = original_initial_environ_path + os.pathsep + cuobjdump_dir + print_msg(f"Add location of cuobjdump ({cuobjdump_dir}) to initial environ PATH for CUDA sanity check") + break + + +def post_sanitycheck_hook_cuda(self, *args, **kwargs): + """ + Reverse the temporary CUDA dependency promotion from pre_sanitycheck_hook_cuda. + """ + if hasattr(self, EESSI_SANITYCHECK_CUDA_ATTR): + original_path = getattr(self, EESSI_SANITYCHECK_CUDA_ATTR) + + if not isinstance(original_path, str) or os.pathsep not in original_path: + raise EasyBuildError( + f"Expected {EESSI_SANITYCHECK_CUDA_ATTR} attribute to contain a PATH-like " + f"value, but got: {original_path}" + ) + + # Restore the original initial environ PATH + print_msg(f"Restoring stored initial environ PATH after CUDA sanity check") + self.initial_environ['PATH'] = original_path + delattr(self, EESSI_SANITYCHECK_CUDA_ATTR) + + def post_easyblock_hook_copy_easybuild_subdir(self, *args, **kwargs): """ Post easyblock hook that copies the easybuild subdirectory of every installed application