[MNT] Comments on numba caching and release CI fix (#2916)

* release CI jank

* comment

* swap

Author: MatthewMiddlehurst

.github/actions/numba_cache/action.yml

@@ -24,19 +24,23 @@ runs:
       run: echo "NUMBA_CACHE_DIR=${{ github.workspace }}/.numba_cache" >> $GITHUB_ENV
       shell: bash
 
-    # Sets the cpu name for numba to use
-    - name: Set numba cpu env
+    # Sets the CPU name for numba to use
+    # This is required for reusing the cache across different runners of the same
+    # architecture and OS
+    - name: Set numba CPU env
       run: echo "NUMBA_CPU_NAME=generic" >> $GITHUB_ENV
       shell: bash
 
-    # Sets the cpu features for numba to use
-    - name: Set numba cpu features env
+    # Sets the CPU features for numba to use
+    # Found through a bit of trial and error to run on different CPUs found in GitHub
+    # runners
+    - name: Set numba CPU features env
       run: echo "NUMBA_CPU_FEATURES=+64bit +adx +aes +avx +avx2 -avx512bf16 -avx512bitalg -avx512bw -avx512cd -avx512dq -avx512er -avx512f -avx512ifma -avx512pf -avx512vbmi -avx512vbmi2 -avx512vl -avx512vnni -avx512vpopcntdq +bmi +bmi2 -cldemote -clflushopt -clwb -clzero +cmov +cx16 +cx8 -enqcmd +f16c +fma -fma4 +fsgsbase +fxsr -gfni +invpcid -lwp +lzcnt +mmx +movbe -movdir64b -movdiri -mwaitx +pclmul -pconfig -pku +popcnt -prefetchwt1 +prfchw -ptwrite -rdpid +rdrnd +rdseed +rtm +sahf -sgx -sha -shstk +sse +sse2 +sse3 +sse4.1 +sse4.2 -sse4a +ssse3 -tbm -vaes -vpclmulqdq -waitpkg -wbnoinvd -xop +xsave -xsavec +xsaveopt -xsaves" >> $GITHUB_ENV
       shell: bash
 
-    # Set the CICD_RUNNING env so that numba knows it is running in a CI environment
-    - name: Set CICD_RUNNING env
-      run: echo "CICD_RUNNING=1" >> $GITHUB_ENV
+    # Set the CACHING_CICD_RUNNING env so we know to use custom numba caching
+    - name: Set CACHING_CICD_RUNNING env
+      run: echo "CACHING_CICD_RUNNING=1" >> $GITHUB_ENV
       shell: bash
 
     # Get current date for cache restore

.github/workflows/periodic_tests.yml

@@ -178,7 +178,7 @@ jobs:
         if: runner.os == 'Linux'
         uses: pierotofy/set-swap-space@v1.0
         with:
-          swap-size-gb: 4
+          swap-size-gb: 6
 
       - name: Use numba cache to set env variables but not restore cache
         uses: ./.github/actions/numba_cache

.github/workflows/pr_pytest.yml

@@ -76,7 +76,7 @@ jobs:
         if: runner.os == 'Linux'
         uses: pierotofy/set-swap-space@v1.0
         with:
-          swap-size-gb: 4
+          swap-size-gb: 6
 
       - if: ${{ github.event_name != 'pull_request' || !contains(github.event.pull_request.labels.*.name, 'no numba cache') }}
         name: Restore numba cache

.github/workflows/release.yml

@@ -69,7 +69,7 @@ jobs:
         if: runner.os == 'Linux'
         uses: pierotofy/set-swap-space@v1.0
         with:
-          swap-size-gb: 4
+          swap-size-gb: 6
 
       - uses: actions/download-artifact@v4
         with:
@@ -83,6 +83,14 @@ jobs:
         name: Unix wheel filename
         run: echo "WHEELNAME=$(ls ./dist/aeon-*-none-any.whl)" >> $GITHUB_ENV
 
+      - name: Use numba cache to set env variables but not restore cache
+        uses: ./.github/actions/numba_cache
+        with:
+            cache_name: "release"
+            runner_os: ${{ runner.os }}
+            python_version: ${{ matrix.python-version }}
+            restore_cache: "false"
+
       - if: runner.os == 'Windows'
         name: Windows install
         uses: nick-fields/retry@v3

aeon/testing/_cicd_numba_caching.py

@@ -6,7 +6,7 @@
 
 import numba.core.caching
 
-if os.environ.get("CICD_RUNNING") == "1":  # pragma: no cover
+if os.environ.get("CACHING_CICD_RUNNING") == "1":  # pragma: no cover
 
     def get_invalid_numba_files():
         """Get the files that have been changed since the last commit.