Reduce cache usage in workflows (#3053)

Fixes #3052.

Author: pbchekin

.github/actions/install-wheels/action.yml

@@ -40,18 +40,7 @@ runs:
         echo "${{ inputs.workflow }}: $run_id"
         echo "run_id=$run_id" >> $GITHUB_OUTPUT
 
-    - name: Load wheels from cache
-      id: wheels-cache
-      uses: ./.github/actions/load
-      env:
-        # Increase this value to reset cache
-        CACHE_NUMBER: "1"
-      with:
-        path: ~/wheels
-        key: wheels-py${{ inputs.python_version }}-${{ steps.run_id.outputs.run_id }}-$CACHE_NUMBER
-
     - name: Download wheels from specified workflow run artifacts
-      if: ${{ steps.wheels-cache.outputs.status == 'miss' }}
       shell: bash
       env:
         GH_TOKEN: ${{ inputs.gh_token }}
@@ -70,10 +59,3 @@ runs:
         shopt -s extglob
         cd ~/wheels/*-py${{ inputs.python_version }}*
         ${{ inputs.install_cmd }} ${{ inputs.wheels_pattern }}
-
-    - name: Save wheels to cache
-      if: ${{ steps.wheels-cache.outputs.status == 'miss' }}
-      uses: ./.github/actions/save
-      with:
-        path: ${{ steps.wheels-cache.outputs.path }}
-        dest: ${{ steps.wheels-cache.outputs.dest }}

.github/actions/load/action.yml

@@ -2,7 +2,7 @@ name: load
 description: Load a directory from a cache based on a shared directory
 inputs:
   root:
-    description:
+    description: Directory for cache
     default: /cache
   path:
     description: A directory to load from a cache
@@ -12,8 +12,10 @@ inputs:
     required: true
   symlink:
     description: Create a symlink instead of copying from cache
-    type: bool
-    default: true
+    default: "true"
+  enabled:
+    description: Enable cache
+    default: "true"
 outputs:
   path:
     description: A directory to save to a cache
@@ -24,6 +26,7 @@ outputs:
   dest:
     description: Directory in cache
     value: ${{ steps.load.outputs.dest }}
+
 runs:
   using: "composite"
   steps:
@@ -37,7 +40,7 @@ runs:
           echo "Directory ${{ inputs.path }} exists and will not be restored from cache"
           exit 1
         fi
-        if [[ -d $ITEM_PATH ]]; then
+        if [[ ${{ inputs.enabled == 'true' }} && -d $ITEM_PATH ]]; then
           echo "Cache hit for ${{ inputs.key }}"
           echo "status=hit" >> $GITHUB_OUTPUT
           if [[ ${{ inputs.symlink }} == true ]]; then

.github/actions/save/action.yml

@@ -2,18 +2,22 @@ name: save
 description: Save a directory to a cache based on a shared directory
 inputs:
   root:
-    description:
+    description: Directory for cache
     default: /cache
   path:
     description: Directory to save to a cache
     required: true
   dest:
     description: Directory in cache
     required: true
+  enabled:
+    description: Enable cache
+    default: "true"
 runs:
   using: "composite"
   steps:
     - name: Save ${{ inputs.path }} to cache
+      if: ${{ inputs.enabled == 'true' }}
       shell: bash
       run: |
         TEMP_ITEM=$(mktemp -d -p ${{ inputs.root }})

.github/actions/setup-pytorch/action.yml

@@ -19,6 +19,9 @@ inputs:
   mode:
     description: Source or wheels
     default: source
+  cache:
+    description: Cache enabled or disabled
+    default: enabled
 runs:
   using: "composite"
   steps:
@@ -83,6 +86,7 @@ runs:
       with:
         path: pytorch
         key: pytorch-$PYTORCH_CACHE_KEY-$CACHE_NUMBER
+        enabled: ${{ inputs.cache == 'enabled' }}
 
     - name: Clone PyTorch repository
       if: ${{ steps.pytorch-cache.outputs.status == 'miss' }}
@@ -138,3 +142,4 @@ runs:
       with:
         path: ${{ steps.pytorch-cache.outputs.path }}
         dest: ${{ steps.pytorch-cache.outputs.dest }}
+        enabled: ${{ inputs.cache == 'enabled' }}

.github/workflows/conda-test-reusable.yml

@@ -78,15 +78,6 @@ jobs:
       - name: Checkout repository
         uses: actions/checkout@v4
 
-      - name: Load conda cache
-        id: conda-cache
-        uses: ./.github/actions/load
-        env:
-          CACHE_NUMBER: 7
-        with:
-          path: $HOME/miniforge3/envs/triton
-          key: conda-${{ inputs.env_manager }}-py${{ matrix.python }}-${{ hashFiles('scripts/triton.yml', 'python/pyproject.toml', 'python/setup.py') }}-${{ env.CACHE_NUMBER }}
-
       - name: Install Manager Environment
         shell: bash --noprofile --norc -eo pipefail {0}
         run: |
@@ -143,13 +134,6 @@ jobs:
         run: |
           ${{ env.TRITON_TEST_CMD }} --inductor --skip-pip-install
 
-      - name: Save conda cache
-        if: steps.conda-cache.outputs.status == 'miss'
-        uses: ./.github/actions/save
-        with:
-          path: ${{ steps.conda-cache.outputs.path }}
-          dest: ${{ steps.conda-cache.outputs.dest }}
-
       - name: Pass rate
         run: |
           pip install defusedxml

.github/workflows/inductor-tests-reusable.yml

@@ -49,14 +49,6 @@ jobs:
       - name: Checkout repository
         uses: actions/checkout@v4
 
-      - name: Load pip cache
-        id: pip-cache
-        uses: ./.github/actions/load
-        with:
-          path: $HOME/.cache/pip
-          # pip cache per commit id just to minimize network traffic
-          key: pip-$PYTHON_VERSION-$GITHUB_SHA
-
       - name: Install Python
         uses: actions/setup-python@v5
         with:
@@ -66,6 +58,7 @@ jobs:
         uses: ./.github/actions/setup-pytorch
         with:
           ref: ${{ inputs.pytorch_ref }}
+          cache: disabled
 
       - name: Setup Triton
         uses: ./.github/actions/setup-triton
@@ -126,10 +119,3 @@ jobs:
           name: logs-${{ env.PYTHON_VERSION }}
           path: pytorch/test/test-reports
           include-hidden-files: true
-
-      - name: Save pip cache
-        if: ${{ steps.pip-cache.outputs.status == 'miss' }}
-        uses: ./.github/actions/save
-        with:
-          path: ${{ steps.pip-cache.outputs.path }}
-          dest: ${{ steps.pip-cache.outputs.dest }}

.github/workflows/nightly-wheels.yml

@@ -70,34 +70,13 @@ jobs:
         run: |
           echo "TRITON_COMMIT_ID=$(git rev-parse HEAD)" >> $GITHUB_ENV
 
-      - name: Generate Triton cache key
-        id: triton-key
-        run: |
-          COMPOSITE_KEY=$(echo ${{ matrix.python }} $PYTORCH_VERSION $TRITON_COMMIT_ID | sha256sum - | cut -d\  -f1)
-          echo "key=triton-$COMPOSITE_KEY" >> $GITHUB_OUTPUT
-
-      - name: Load Triton wheels from a cache
-        id: triton-cache
-        uses: ./.github/actions/load
-        with:
-          path: python/dist
-          key: ${{ steps.triton-key.outputs.key }}
-
       - name: Build Triton wheels
-        if: ${{ steps.triton-cache.outputs.status == 'miss' }}
         uses: ./.github/actions/setup-triton
         with:
           command: >
             DEBUG=1
             python setup.py bdist_wheel && pip install dist/*.whl
 
-      - name: Save Triton wheels to a cache
-        if: ${{ steps.triton-cache.outputs.status == 'miss' }}
-        uses: ./.github/actions/save
-        with:
-          path: ${{ steps.triton-cache.outputs.path }}
-          dest: ${{ steps.triton-cache.outputs.dest }}
-
       - name: Install torchvision package
         uses: ./.github/actions/install-dependency
         with:

.github/workflows/triton-benchmarks.yml

@@ -82,14 +82,6 @@ jobs:
       - name: Checkout repository
         uses: actions/checkout@v4
 
-      - name: Load pip cache
-        id: pip-cache
-        uses: ./.github/actions/load
-        with:
-          path: $HOME/.cache/pip
-          # pip cache per commit id just to minimize network traffic
-          key: pip-$PYTHON_VERSION-$GITHUB_SHA
-
       - name: Install Python
         if: ${{ !(inputs.use_pyenv_python || false) }}
         uses: actions/setup-python@v5
@@ -273,13 +265,6 @@ jobs:
           cd benchmarks/micro_benchmarks
           python run_benchmarks.py --reports $REPORTS
 
-      - name: Save pip cache
-        if: ${{ steps.pip-cache.outputs.status == 'miss' }}
-        uses: ./.github/actions/save
-        with:
-          path: ${{ steps.pip-cache.outputs.path }}
-          dest: ${{ steps.pip-cache.outputs.dest }}
-
       - name: Upload benchmark reports
         if: ${{ steps.install.outcome == 'success' && !cancelled() }}
         uses: actions/upload-artifact@v4