Skip to content

{2023.06}[foss/2023a] TensorFlow v2.15.1 w/ CUDA 12.1.1 + eb_hooks.py #35

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
easyconfigs:
- TensorFlow-2.15.1-foss-2023a-CUDA-12.1.1.eb
51 changes: 50 additions & 1 deletion eb_hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -256,6 +256,53 @@ def post_prepare_hook(self, *args, **kwargs):
post_prepare_hook_ignore_zen4_gcccore1220_error(self, *args, **kwargs)


def parse_hook_tensorflow_CUDA(ec, eprefix):
"""
Fix the Python and environment used while building and running tests for TensorFlow with CUDA
"""
if ec.name == 'TensorFlow' and ec.version == '2.15.1' :
# Check if CUDA is in dependencies
has_cuda = any(
(isinstance(dep, (list, tuple)) and dep[0] == 'CUDA') or
(isinstance(dep, dict) and dep.get('name') == 'CUDA')
for dep in ec.get('dependencies', [])
)

if has_cuda:
ec['preconfigopts'] = ec.get('preconfigopts', '') + (
'export TF_NEED_CUDA=1 && '
'export CUDA_TOOLKIT_PATH=$EBROOTCUDA && '
'export TF_CUDA_INCLUDE_PATH=$EBROOTCUDA/include && '
'export CUDNN_INSTALL_PATH=$EBROOTCUDNN && '
'export GCC_HOST_COMPILER_PATH=$EBROOTGCC/bin/gcc && '
'sed -i \'s|--define=PREFIX=/usr|--define=PREFIX=\\$EESSI_EPREFIX|g\' .bazelrc && '
)

current_opts = ec.get('buildopts', [])
if isinstance(current_opts, str):
current_opts = current_opts.split()

ec['buildopts'] = current_opts + [
'--linkopt=-Wl,--disable-new-dtags --host_linkopt=-Wl,--disable-new-dtags --action_env=GCC_HOST_COMPILER_PATH=$EBROOTGCC/bin/gcc --host_action_env=GCC_HOST_COMPILER_PATH=$EBROOTGCC/bin/gcc --linkopt=-Wl,-rpath,$EBROOTCUDA/lib:$EBROOTCUDNN/lib:$EBROOTNCCL/lib --host_linkopt=-Wl,-rpath,$EBROOTCUDA/lib:$EBROOTCUDNN/lib:$EBROOTNCCL/lib',
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This doesn't look correct to me, why aren't you telling it to use our rpath wrappers, e.g., $(which gcc)?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is happening within Bazel, so i tried to use the options available within the easyblock

]

ec['pretestopts'] = (
"""interppath=$(find "$EESSI_EPREFIX/lib64" -name 'ld-*' | grep -E 'so\\.1|so\\.2' | head -n1) && """
"""pybin=$(find "%(builddir)s/%(name)s/bazel-root/" -type f -path "*/external/python_%(arch)s-unknown-linux-gnu/bin/python%(pyshortver)s" | head -n1) && """
"""patchelf --set-interpreter "$interppath" "$pybin" && """
"""export LD_LIBRARY_PATH="$EBROOTCUDA/lib:$EBROOTCUDNN/lib:$EBROOTNCCL/lib:$LD_LIBRARY_PATH" && """
)

ec['postinstallcmds'] = [
'mkdir -p %(installdir)s/bin',
'ln -s $EBROOTCUDA/bin/cuobjdump %(installdir)s/bin/cuobjdump',
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not a fan of this, why is it necessary?

Copy link
Contributor Author

@TopRichard TopRichard Aug 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It was necessary for EasyBuild CUDA check, not related to TensorFlow or EESSI

ERROR EasyBuild encountered an error (at easybuild/tools/build_log.py:166 in caller_info): Installation of TensorFlow-2.15.1-foss-2023a-CUDA-12.1.1.eb failed: 'Failed to execute CUDA sanity check: cuobjdump not found\nCUDA module must be loaded for sanity check (or cuobjdump available in PATH)' (at easybuild/main.py:205 in build_and_install_software)

]

print_msg("TensorFlow-CUDA related changes have been applied")
else:
raise EasyBuildError("TensorFlow-CUDA specific hook triggered for non-TensorFlow-CUDA easyconfig?!")


def parse_hook_casacore_disable_vectorize(ec, eprefix):
"""
Disable 'vectorize' toolchain option for casacore 3.5.0 on aarch64/neoverse_v1
Expand All @@ -275,7 +322,7 @@ def parse_hook_casacore_disable_vectorize(ec, eprefix):
if 'toolchainopts' not in ec or ec['toolchainopts'] is None:
ec['toolchainopts'] = {}
ec['toolchainopts']['vectorize'] = False
print_msg("Changed toochainopts for %s: %s", ec.name, ec['toolchainopts'])
print_msg("Changed toolchainopts for %s: %s", ec.name, ec['toolchainopts'])
else:
print_msg("Not changing option vectorize for %s on non-neoverse_v1", ec.name)
else:
Expand Down Expand Up @@ -1298,6 +1345,7 @@ def post_module_hook(self, *args, **kwargs):


PARSE_HOOKS = {
'TensorFlow': parse_hook_tensorflow_CUDA,
'casacore': parse_hook_casacore_disable_vectorize,
'CGAL': parse_hook_cgal_toolchainopts_precise,
'fontconfig': parse_hook_fontconfig_add_fonts,
Expand Down Expand Up @@ -1400,3 +1448,4 @@ def set_maximum(parallel, max_value):
CPU_TARGET_A64FX: (set_maximum, 8),
},
}

Loading