Enable Cache in ICON4Py checks (#493)

Author: kotsaloscv

checks/apps/icon4py/icon4py_check.py

@@ -3,6 +3,8 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
+import os
+import re
 import reframe as rfm
 import reframe.utility.sanity as sn
 
@@ -28,16 +30,28 @@ class ICON4PyBenchmarks(rfm.RunOnlyRegressionTest):
         'HUGETLB_MORECORE': 'no',
         'GT4PY_UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE': '1',
         'PYTHONOPTIMIZE': '2',
-        # GT4Py cache does not work properly for dace backend yet
-        # 'GT4PY_BUILD_CACHE_LIFETIME': 'persistent',
-        # 'GT4PY_BUILD_CACHE_DIR': '...',
+        'GT4PY_BUILD_CACHE_LIFETIME': 'persistent',
     }
     executable = './_run.sh'
     executable_opts = ['2>&1']
 
     @run_before('run')
     def prepare_env(self):
         gpu_arch = self.current_partition.select_devices('gpu')[0].arch
+
+        cache_folder = (
+            f"{os.environ.get('SCRATCH')}/"
+            f".cache/reframe_bencher_icon4py/"
+        )
+        sub_folder = (
+            f"{self.current_system.name}="
+            f"{gpu_arch}="
+            f"{self.current_environ}"
+        )
+        sub_folder = re.sub(r'[^a-zA-Z0-9=]', '', sub_folder)
+        cache_folder = os.path.join(cache_folder, sub_folder)
+        self.env_vars['GT4PY_BUILD_CACHE_DIR'] = cache_folder
+
         if 'gfx' in gpu_arch:  # AMD GPU
             self.env_vars['CUPY_INSTALL_USE_HIP'] = '1'
             self.env_vars['HCC_AMDGPU_TARGET'] = gpu_arch
@@ -54,12 +68,12 @@ def validate_test(self):
             (r'^\s*model/atmosphere/diffusion/tests/'
              r'diffusion/integration_tests/'
              r'test_benchmark_diffusion\.py'
-             r'::test_diffusion_benchmark\s*PASSED'
+             r'::test_diffusion_benchmark[\s\S]*?PASSED'
              ), self.stdout)
         dycore_granule = sn.assert_found(
             (r'^\s*model/atmosphere/dycore/tests/'
              r'dycore/integration_tests/test_benchmark_solve_nonhydro\.py'
-             r'::test_benchmark_solve_nonhydro\[True-False\]\s*PASSED'
+             r'::test_benchmark_solve_nonhydro\[True-False\][\s\S]*?PASSED'
              ), self.stdout)
 
         return diffusion_granule and dycore_granule

checks/apps/icon4py/src/_install.sh

@@ -7,7 +7,7 @@ unset PYTHONPATH
 
 git clone https://github.com/C2SM/icon4py.git
 cd icon4py
-git checkout 5485bcacb1dbc7688b1e7d276d4e2e28362c5444  # Commit: Update to GT4Py v1.1.0 (#933)
+git checkout 15b7406d9385a189abaaa71b14851d1491b231f1  # Commit: Update to GT4Py v1.1.2 (#969)
 
 # Install uv locally
 curl -LsSf https://astral.sh/uv/install.sh | UV_UNMANAGED_INSTALL="$PWD/bin" sh
@@ -25,6 +25,54 @@ mpi4py_ver=$(uv pip show mpi4py | awk '/Version:/ {print $2}')
 uv pip uninstall mpi4py && uv pip install --no-binary mpi4py "mpi4py==$mpi4py_ver"
 uv pip install git+https://github.com/cupy/cupy.git
 
+# Patch Gt4Py to avoid cache issues
+uv pip uninstall gt4py
+git clone --branch v1.1.2 https://github.com/GridTools/gt4py.git
+python3 -c '
+import sys
+from pathlib import Path
+
+file_path = Path("gt4py/src/gt4py/next/otf/stages.py")
+lines = file_path.read_text().splitlines()
+
+new_lines = []
+iterator = iter(lines)
+
+found = False
+for line in iterator:
+    # 1. Detect the start of the block we want to change
+    if "program_hash = utils.content_hash(" in line:
+        found = True
+        # Insert the NEW pre-calculation line
+        # We steal the indentation from the current line to be safe
+        indent = line[:line.find("program_hash")]
+        new_lines.append(f"{indent}offset_provider_arrays = {{key: value.ndarray if hasattr(value, \"ndarray\") else value for key, value in offset_provider.items()}}")
+        
+        # Add the modified content_hash call
+        new_lines.append(f"{indent}program_hash = utils.content_hash(")
+        new_lines.append(f"{indent}    (")
+        new_lines.append(f"{indent}        program.fingerprint(),")
+        new_lines.append(f"{indent}        sorted(offset_provider_arrays.items(), key=lambda el: el[0]),")
+        
+        # Skip the OLD lines from the iterator until we hit "column_axis"
+        # We blindly consume lines until we find the one we keep
+        while True:
+            skipped_line = next(iterator)
+            if "column_axis," in skipped_line:
+                new_lines.append(skipped_line) # Add column_axis line back
+                break
+    else:
+        new_lines.append(line)
+
+if found:
+    file_path.write_text("\n".join(new_lines) + "\n")
+    print("Patch applied.")
+else:
+    print("Target line not found.")
+    sys.exit(1)
+'
+uv pip install -e ./gt4py
+
 ################################################################################
 # NVHPC runtime auto-discovery for serialbox (libnvhpcatm.so)
 ################################################################################

checks/apps/icon4py/src/_run.sh

@@ -7,7 +7,7 @@ unset PYTHONPATH
 cd icon4py
 source .venv/bin/activate
 
-pytest -v \
+pytest -sv \
     -m continuous_benchmarking \
     --benchmark-only \
     --benchmark-warmup=on \
@@ -16,8 +16,8 @@ pytest -v \
     --backend=dace_gpu \
     --grid=icon_benchmark_regional \
     --benchmark-time-unit=ms \
-    model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py::test_diffusion_benchmark \
-    model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py::test_benchmark_solve_nonhydro[True-False]
+    "model/atmosphere/diffusion/tests/diffusion/integration_tests/test_benchmark_diffusion.py::test_diffusion_benchmark" \
+    "model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py::test_benchmark_solve_nonhydro[True-False]"
 echo
 
 # Cleanup