Add pytorch_mode: source ot wheels

Author: pbchekin

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

@@ -16,6 +16,9 @@ inputs:
   ref:
     description: Branch, tag, commit id
     default: ""
+  mode:
+    description: Source or wheels
+    default: source
 runs:
   using: "composite"
   steps:
@@ -71,7 +74,7 @@ runs:
     - name: Generate PyTorch cache key
       shell: bash
       run: |
-        PYTORCH_CACHE_KEY=$(echo $PYTHON_VERSION $PYTORCH_COMMIT_ID ${{ hashFiles('scripts/patch-pytorch.sh') }} | sha256sum - | cut -d\  -f1)
+        PYTORCH_CACHE_KEY=$(echo $PYTHON_VERSION $PYTORCH_COMMIT_ID ${{ hashFiles('scripts/patch-pytorch.sh') }} ${{ inputs.mode }} | sha256sum - | cut -d\  -f1)
         echo "PYTORCH_CACHE_KEY=$PYTORCH_CACHE_KEY" | tee -a "$GITHUB_ENV"
 
     - name: Load PyTorch from a cache
@@ -94,7 +97,7 @@ runs:
         path: pytorch
 
     - name: Apply additional PR patches
-      if: ${{ steps.pytorch-cache.outputs.status == 'miss' && inputs.repository == 'pytorch/pytorch' }}
+      if: ${{ steps.pytorch-cache.outputs.status == 'miss' && inputs.repository == 'pytorch/pytorch' && inputs.mode == 'source' }}
       shell: bash
       run: |
         cd pytorch
@@ -108,7 +111,7 @@ runs:
         pip install 'numpy<2.0.0'
 
     - name: Build PyTorch
-      if: ${{ steps.pytorch-cache.outputs.status == 'miss' }}
+      if: ${{ steps.pytorch-cache.outputs.status == 'miss' && inputs.mode == 'source' }}
       shell: bash
       run: |
         source ${{ inputs.oneapi }}/setvars.sh
@@ -117,6 +120,13 @@ runs:
         pip install -r requirements.txt
         python setup.py bdist_wheel
 
+    - name: Download PyTorch wheels
+      if: ${{ steps.pytorch-cache.outputs.status == 'miss' && inputs.mode == 'wheels' }}
+      shell: bash
+      run: |
+        mkdir -p pytorch/dist
+        pip download torch --index-url https://download.pytorch.org/whl/nightly/xpu --dest pytorch/dist/
+
     - name: Install PyTorch
       shell: bash
       run: |

.github/workflows/build-test-gpu.yml

@@ -12,6 +12,10 @@ on:
         description: PyTorch ref, keep empty for default
         type: string
         default: ""
+      pytorch_mode:
+        description: PyTorch mode, source or wheels
+        type: string
+        default: "source"
       upload_test_reports:
         description: Upload test reports
         type: boolean
@@ -46,6 +50,7 @@ jobs:
       device: ${{ inputs.runner_label }}
       runner_label: ${{ inputs.runner_label }}
       pytorch_ref: ${{ inputs.pytorch_ref }}
+      pytorch_mode: ${{ inputs.pytorch_mode }}
       python_version: ${{ matrix.python }}
       upload_test_reports: ${{ inputs.upload_test_reports }}
       ignore_errors: ${{ inputs.ignore_errors }}

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

@@ -20,6 +20,10 @@ on:
         description: PyTorch ref, keep empty for default
         type: string
         default: ""
+      pytorch_mode:
+        description: PyTorch mode, source or wheels
+        type: string
+        default: "source"
       python_version:
         description: Python version
         type: string
@@ -96,6 +100,7 @@ jobs:
         with:
           repository: pytorch/pytorch
           ref: ${{ inputs.pytorch_ref }}
+          mode: ${{ inputs.pytorch_mode }}
 
       - name: Install pass_rate dependencies
         run: |

.github/workflows/build-test.yml

@@ -12,6 +12,10 @@ on:
         description: PyTorch ref, keep empty for default
         type: string
         default: ""
+      pytorch_mode:
+        description: PyTorch mode, source or wheels
+        type: string
+        default: "source"
       upload_test_reports:
         description: Upload test reports
         type: boolean
@@ -120,6 +124,7 @@ jobs:
       driver_version: ${{ matrix.driver }}
       runner_label: ${{ inputs.runner_label }}
       pytorch_ref: ${{ inputs.pytorch_ref }}
+      pytorch_mode: ${{ inputs.pytorch_mode }}
       python_version: ${{ matrix.python }}
       upload_test_reports: ${{ inputs.upload_test_reports || false }}
       ignore_errors: ${{ inputs.ignore_errors || false }}