Download nightly wheels from main branch (#4257)

Use main branch only to download Triton wheels.
This allows building wheels for test branches without affecting normal
flow.

---------

Signed-off-by: Pavel Chekin <[email protected]>

Author: pbchekin

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

@@ -4,10 +4,13 @@ description: Install latest wheels from a specified workflow using a custom wild
 inputs:
   workflow:
     description: Name of the workflow to install wheels from
-    default: Triton wheels
+    default: nightly-wheels.yml
   repository:
     description: Repository name with owner
     default: intel/intel-xpu-backend-for-triton
+  branch:
+    description: Filter runs by branch (use empty string for all branches)
+    default: main
   wheels_pattern:
     # Example of specifying only some packages to install: '{intel_extension_for_pytorch-*,torch-*}'
     # Extended globbing is enabled for this pattern so for example to exclude intel_extension_for_pytorch use pattern
@@ -33,7 +36,11 @@ runs:
         GH_TOKEN: ${{ inputs.gh_token }}
       run: |
         set -e
-        run_id=$(gh run list -w "${{ inputs.workflow }}" -s success --repo ${{ inputs.repository }} --json databaseId --jq '.[0].databaseId')
+        if [[ -z "${{ inputs.branch }}" ]]; then
+          run_id=$(gh run list --workflow "${{ inputs.workflow }}" --status success --repo "${{ inputs.repository }}" --json databaseId --jq '.[0].databaseId')
+        else
+          run_id=$(gh run list --workflow "${{ inputs.workflow }}" --branch "${{ inputs.branch }}" --status success --repo "${{ inputs.repository }}" --json databaseId --jq '.[0].databaseId')
+        fi
         if [[ ! $run_id ]]; then
           exit 1
         fi

scripts/install-pytorch.sh

@@ -117,7 +117,7 @@ if [ "$BUILD_PYTORCH" = false ]; then
   fi
   echo "**** Download nightly builds. ****"
   PYTHON_VERSION=$(python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
-  RUN_ID=$(gh run list -w "Triton wheels" -R intel/intel-xpu-backend-for-triton --json databaseId,conclusion | jq -r '[.[] | select(.conclusion=="success")][0].databaseId')
+  RUN_ID=$(gh run list --workflow nightly-wheels.yml --branch main -R intel/intel-xpu-backend-for-triton --json databaseId,conclusion | jq -r '[.[] | select(.conclusion=="success")][0].databaseId')
   TEMP_DIR=$(mktemp -d)
   WHEEL_PATTERN="wheels-pytorch-py${PYTHON_VERSION}*"
   gh run download $RUN_ID \