ci(*:skip) Fix cache cleanup (#3521)

.github/workflows/cache-cleanup.yml

@@ -0,0 +1,63 @@
+name: Cache cleanup
+
+on:
+  workflow_dispatch
+
+jobs:
+  e2e-cleanup:
+    runs-on: ubuntu-22.04
+    timeout-minutes: 10
+    # Using approach described here:
+    # https://docs.github.com/en/actions/using-jobs/using-a-matrix-for-your-jobs
+    strategy:
+      matrix:
+        include:
+          - directory: bare
+          - directory: bare-https
+          - directory: bare-client-auth
+          - directory: jax
+          - directory: pytorch
+          - directory: tensorflow
+          - directory: tabnet
+          - directory: opacus
+          - directory: pytorch-lightning
+          - directory: scikit-learn
+          - directory: fastai
+          - directory: pandas
+
+    name: Framework / ${{ matrix.directory }}
+
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.8
+
+      - name: Cleanup caches by directories
+        # Only keep caches that match the latest keys for each directory
+        run: |
+          gh extension install actions/gh-actions-cache
+          
+          REPO=${{ github.repository }}
+          LATEST_KEY=pythonloc-${{ matrix.directory }}-${{ env.pythonLocation }}-${{ hashFiles(format('./e2e/{0}/pyproject.toml', matrix.directory)) }}
+
+          echo "Fetching list of cache keys"
+          cacheKeys=$(gh actions-cache list -R $REPO | grep "${{ matrix.directory }}" | cut -f 1 )
+
+          ## Setting this to not fail the workflow while deleting cache keys. 
+          set +e
+          echo "Deleting caches..."
+          for cacheKey in $cacheKeys
+          do
+            if [ "$cacheKey" != "$LATEST_KEY" ]; then
+              echo -e "\tOld key found -> $cacheKey"
+              gh actions-cache delete "$cacheKey" --confirm
+            fi
+          done
+
+          echo "Done"
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/e2e.yml

@@ -129,13 +129,11 @@ jobs:
       # Using approach described here for Python location caching:
       # https://blog.allenai.org/python-caching-in-github-actions-e9452698e98d
       - name: Cache Python location
-        id: cache-python
-        uses: actions/cache@v4
+        id: cache-restore-python
+        uses: actions/cache/restore@v4
         with:
           path: ${{ env.pythonLocation }}
-          key: pythonloc-${{ runner.os }}-${{ matrix.directory }}-${{ env.pythonLocation }}-${{ hashFiles('**/pyproject.toml') }}
-          restore-keys: |
-            pythonloc-${{ runner.os }}-${{ matrix.directory }}-${{ env.pythonLocation }}
+          key: pythonloc-${{ runner.os }}-${{ matrix.directory }}-${{ env.pythonLocation }}-${{ hashFiles(format('./e2e/{0}/pyproject.toml', matrix.directory)) }}
       - name: Install dependencies
         run: python -m pip install --upgrade .
       - name: Install Flower wheel from artifact store
@@ -166,6 +164,13 @@ jobs:
       - name: Run driver test with client authentication
         if: ${{ matrix.directory == 'bare-client-auth' }}
         run: ./../test_driver.sh bare client-auth
+      - name: Cache save Python location
+        id: cache-save-python
+        uses: actions/cache/save@v4
+        if: ${{ github.ref_name == 'main' &&  !steps.cache-restore-python.outputs.cache-hit }}
+        with:
+          path: ${{ env.pythonLocation }}
+          key: pythonloc-${{ runner.os }}-${{ matrix.directory }}-${{ env.pythonLocation }}-${{ hashFiles(format('./e2e/{0}/pyproject.toml', matrix.directory)) }}
 
   strategies:
     runs-on: ubuntu-22.04