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

easystacks/software.eessi.io/2023.06/accel/nvidia/eessi-2023.06-eb-5.1.1-2023a-CUDA.yml

@@ -0,0 +1,2 @@
+easyconfigs:
+  - TensorFlow-2.15.1-foss-2023a-CUDA-12.1.1.eb

eb_hooks.py

@@ -256,6 +256,49 @@ 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'] = (
+                '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 && '
+            )
+
+            ec['buildopts'] = [
+                '--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',
+            ]
+
+            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',
+            ]
+
+            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
@@ -275,7 +318,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:
@@ -1298,6 +1341,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,
@@ -1400,3 +1444,4 @@ def set_maximum(parallel, max_value):
         CPU_TARGET_A64FX: (set_maximum, 8),
     },
 }
+