Added ability to run linux workflows on large runners (#6273)

**Context:**

Currently the CI gets congested when large amounts of pull requests are
being updated simultaneously. This pull request gives PRs an escape
hatch and use large runners and use different queue to have CI jobs be
picked up.

**Description of the Change:**

This pull request adds two new features:
- Ability to add the `urgent` label to any pull request and switch it
over to large runners
- Automatic swap of rc branch to large runner
    - This assumes the rc branch is of the format `vX.Y.Z-rcN`

Large runners, albeit slightly more powerful than standard runners, can
be spawned at a much higher volume than standard runners ... this is
because we pay per minute for these runners vs being included on our
GitHub Plan.

If a PR needs CI run without waiting for a runner, **add the `urgent`
label to the pull request**.

Important Note:
- This only affect jobs that run on `pull_request` and use `ubuntu`
runners.
- This change is already in-place in lightning and catalyst.
    - PennyLaneAI/pennylane-lightning#774
    - PennyLaneAI/catalyst#846


**Benefits:**
Ability to leverage large runner to have quick time for a runner to pick
up a job.

**Possible Drawbacks:**
Though we dictate the pool size of large runners, it is possible to
still saturate it.

**Related GitHub Issues:**
None.
[sc-73711](https://app.shortcut.com/xanaduai/story/73711/update-pennylane-ci-to-use-large-runner-group)

---------

Co-authored-by: Mudit Pandey <[email protected]>

Author: rashidnhm

.github/workflows/determine-workflow-runner.yml

@@ -0,0 +1,65 @@
+name: Determine Workflow Runner group
+
+on:
+  workflow_call:
+    inputs:
+      default_runner:
+        description: The runner type that is used by the calling workflow by default
+        required: true
+        type: string
+      large_runner_group_name:
+        description: The name of the runner group that should be used for large jobs
+        required: false
+        type: string
+        default: 'pl-4-core-large-runner'
+      force_large_runner:
+        description: Whether to force the use of a large runner
+        required: false
+        type: boolean
+        default: false
+    outputs:
+      runner_group:
+        description: The runner all subsequent jobs within the calling workflow should run on
+        value: >-
+          ${{
+            jobs.determine_workflow_runner.outputs.runner_group ||
+            (
+              inputs.force_large_runner &&
+              inputs.large_runner_group_name ||
+              inputs.default_runner
+            )
+          }}
+
+jobs:
+  determine_workflow_runner:
+    runs-on: >-
+      ${{
+        (
+          github.event_name == 'pull_request'
+          && contains(github.event.pull_request.labels.*.name, 'urgent')
+        ) && inputs.large_runner_group_name || 'ubuntu-latest'
+      }}
+
+    outputs:
+      runner_group: ${{ steps.runner_group.outputs.runner_group }}
+
+    steps:
+      - name: Output Runner Group name
+        if: >-
+          ${{
+            github.event_name == 'pull_request'
+            && startsWith(inputs.default_runner, 'ubuntu')
+          }}
+        id: runner_group
+        env:
+          # We are not able to use \d to check numeric values as bash does not allow them (not POSIX compliant)
+          RC_BRANCH_FORMAT_REGEX: v[0-9]+\.[0-9]+\.[0-9]+-rc[0-9]?
+          REPO_FULL_NAME: PennyLaneAI/pennylane
+          LARGE_RUNNER_GROUP_NAME: ${{ inputs.large_runner_group_name }}
+        run: |
+          if [[ '${{ contains(github.event.pull_request.labels.*.name, 'urgent') }}' == 'true' || ('${{ github.event.pull_request.head.repo.full_name }}' == "$REPO_FULL_NAME" && '${{ github.event.pull_request.base.ref }}' =~ $RC_BRANCH_FORMAT_REGEX) ]]; then
+            echo "This job requires usage of the large runner group '$LARGE_RUNNER_GROUP_NAME'";
+            echo "runner_group=$LARGE_RUNNER_GROUP_NAME" >> $GITHUB_OUTPUT
+          else
+            echo "This job does not require usage of large runners ...";
+          fi

.github/workflows/docs.yml

@@ -24,11 +24,19 @@ on:
       - ready_for_review
 
 jobs:
+  determine_runner:
+    if: github.event.pull_request.draft == false
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-latest
+
   sphinx:
     if: github.event.pull_request.draft == false
     env:
       DEPS_BRANCH: bot/stable-deps-update
-    runs-on: ubuntu-latest
+    needs: [determine_runner]
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
     steps:
     - uses: actions/checkout@v4
     - uses: PennyLaneAI/sphinx-action@master

.github/workflows/format.yml

@@ -12,9 +12,17 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  determine_runner:
+    if: github.event.pull_request.draft == false
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-latest
+
   black-pylint:
     if: github.event.pull_request.draft == false
-    runs-on: ubuntu-latest
+    needs: [determine_runner]
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - name: Set up Python

.github/workflows/interface-unit-tests.yml

@@ -69,10 +69,28 @@ on:
         required: false
         type: boolean
         default: true
+      use_large_runner:
+        description: |
+          Indicate if the large runner should be used for the job.
+          If this is true, large runner is used for the build regardless of the context.
+        required: false
+        type: boolean
+        default: false
 
 jobs:
+  determine_runner:
+    if: github.event.pull_request.draft == false
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-latest
+      force_large_runner: ${{ inputs.use_large_runner }}
+
   setup-ci-load:
-    runs-on: ubuntu-latest
+    needs:
+      - determine_runner
+    
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - name: Setup Python Versions
@@ -162,7 +180,10 @@ jobs:
   # This job is the source of truth for the default versions of PyTorch, TensorFlow, and JAX.
   # Individual jobs can use these values or override at a per job level.
   default-dependency-versions:
-    runs-on: ubuntu-latest
+    needs:
+      - determine_runner
+    
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - name: Default JAX Version
@@ -196,6 +217,7 @@ jobs:
   torch-tests:
     needs:
       - setup-ci-load
+      - determine_runner
       - default-dependency-versions
     strategy:
       max-parallel: >-
@@ -212,6 +234,7 @@ jobs:
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'torch-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
+      job_runner_name: ${{ needs.determine_runner.outputs.runner_group }}
       job_name: ${{ inputs.job_name_prefix }}torch-tests (${{ matrix.python-version }})${{ inputs.job_name_suffix }}
       branch: ${{ inputs.branch }}
       coverage_artifact_name: core-interfaces-coverage-torch-${{ matrix.python-version }}
@@ -229,6 +252,7 @@ jobs:
   autograd-tests:
     needs:
       - setup-ci-load
+      - determine_runner
       - default-dependency-versions
     strategy:
       max-parallel: >-
@@ -245,6 +269,7 @@ jobs:
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'autograd-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
+      job_runner_name: ${{ needs.determine_runner.outputs.runner_group }}
       job_name: ${{ inputs.job_name_prefix }}autograd-tests (${{ matrix.python-version }})${{ inputs.job_name_suffix }}
       branch: ${{ inputs.branch }}
       coverage_artifact_name: core-interfaces-coverage-autograd-${{ matrix.python-version }}
@@ -259,6 +284,7 @@ jobs:
   tf-tests:
     needs:
       - setup-ci-load
+      - determine_runner
       - default-dependency-versions
     strategy:
       max-parallel: >-
@@ -276,6 +302,7 @@ jobs:
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'tf-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
+      job_runner_name: ${{ needs.determine_runner.outputs.runner_group }}
       job_name: ${{ inputs.job_name_prefix }}tf-tests (${{ matrix.group }}, ${{ matrix.python-version }})${{ inputs.job_name_suffix }}
       branch: ${{ inputs.branch }}
       coverage_artifact_name: core-interfaces-coverage-tf-${{ matrix.python-version }}-${{ matrix.group }}
@@ -295,6 +322,7 @@ jobs:
   jax-tests:
     needs:
       - setup-ci-load
+      - determine_runner
       - default-dependency-versions
     strategy:
       max-parallel: >-
@@ -312,6 +340,7 @@ jobs:
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'jax-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
+      job_runner_name: ${{ needs.determine_runner.outputs.runner_group }}
       job_name: ${{ inputs.job_name_prefix }}jax-tests (${{ matrix.group }}, ${{ matrix.python-version }})${{ inputs.job_name_suffix }}
       branch: ${{ inputs.branch }}
       coverage_artifact_name: core-interfaces-coverage-jax-${{ matrix.python-version }}-${{ matrix.group }}
@@ -331,6 +360,7 @@ jobs:
   core-tests:
     needs:
       - setup-ci-load
+      - determine_runner
       - default-dependency-versions
     strategy:
       max-parallel: >-
@@ -348,6 +378,7 @@ jobs:
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'core-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
+      job_runner_name: ${{ needs.determine_runner.outputs.runner_group }}
       job_name: ${{ inputs.job_name_prefix }}core-tests (${{ matrix.group }}, ${{ matrix.python-version }})${{ inputs.job_name_suffix }}
       branch: ${{ inputs.branch }}
       coverage_artifact_name: core-interfaces-coverage-core-${{ matrix.python-version }}-${{ matrix.group }}
@@ -365,6 +396,7 @@ jobs:
   all-interfaces-tests:
     needs:
       - setup-ci-load
+      - determine_runner
       - default-dependency-versions
     strategy:
       max-parallel: >-
@@ -381,6 +413,7 @@ jobs:
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'all-interfaces-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
+      job_runner_name: ${{ needs.determine_runner.outputs.runner_group }}
       job_name: ${{ inputs.job_name_prefix }}all-interfaces-tests (${{ matrix.python-version }})${{ inputs.job_name_suffix }}
       branch: ${{ inputs.branch }}
       coverage_artifact_name: all-interfaces-coverage
@@ -400,6 +433,7 @@ jobs:
   external-libraries-tests:
     needs:
       - setup-ci-load
+      - determine_runner
       - default-dependency-versions
     strategy:
       max-parallel: >-
@@ -416,6 +450,7 @@ jobs:
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'external-libraries-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
+      job_runner_name: ${{ needs.determine_runner.outputs.runner_group }}
       job_name: ${{ inputs.job_name_prefix }}external-libraries-tests (${{ matrix.python-version }})${{ inputs.job_name_suffix }}
       branch: ${{ inputs.branch }}
       coverage_artifact_name: external-libraries-tests-coverage
@@ -442,6 +477,7 @@ jobs:
   qcut-tests:
     needs:
       - setup-ci-load
+      - determine_runner
       - default-dependency-versions
     strategy:
       max-parallel: >-
@@ -458,6 +494,7 @@ jobs:
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'qcut-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
+      job_runner_name: ${{ needs.determine_runner.outputs.runner_group }}
       job_name: ${{ inputs.job_name_prefix }}qcut-tests (${{ matrix.python-version }})${{ inputs.job_name_suffix }}
       branch: ${{ inputs.branch }}
       coverage_artifact_name: qcut-coverage
@@ -478,6 +515,7 @@ jobs:
   qchem-tests:
     needs:
       - setup-ci-load
+      - determine_runner
       - default-dependency-versions
     strategy:
       max-parallel: >-
@@ -494,6 +532,7 @@ jobs:
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'qchem-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
+      job_runner_name: ${{ needs.determine_runner.outputs.runner_group }}
       job_name: ${{ inputs.job_name_prefix }}qchem-tests (${{ matrix.python-version }})${{ inputs.job_name_suffix }}
       branch: ${{ inputs.branch }}
       coverage_artifact_name: qchem-coverage
@@ -510,6 +549,7 @@ jobs:
   gradients-tests:
     needs:
       - setup-ci-load
+      - determine_runner
       - default-dependency-versions
     strategy:
       max-parallel: >-
@@ -529,6 +569,7 @@ jobs:
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'gradients-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
+      job_runner_name: ${{ needs.determine_runner.outputs.runner_group }}
       job_name: ${{ inputs.job_name_prefix }}gradients-tests (${{ matrix.config.suite }}, ${{ matrix.python-version }})${{ inputs.job_name_suffix }}
       branch: ${{ inputs.branch }}
       coverage_artifact_name: gradients-${{ matrix.config.suite }}-coverage
@@ -547,6 +588,7 @@ jobs:
   data-tests:
     needs:
       - setup-ci-load
+      - determine_runner
       - default-dependency-versions
     strategy:
       max-parallel: >-
@@ -563,6 +605,7 @@ jobs:
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'data-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
+      job_runner_name: ${{ needs.determine_runner.outputs.runner_group }}
       job_name: ${{ inputs.job_name_prefix }}data-tests (${{ matrix.python-version }})${{ inputs.job_name_suffix }}
       branch: ${{ inputs.branch }}
       coverage_artifact_name: data-coverage-${{ matrix.python-version }}
@@ -579,6 +622,7 @@ jobs:
   device-tests:
     needs:
       - setup-ci-load
+      - determine_runner
       - default-dependency-versions
     strategy:
       max-parallel: >-
@@ -602,6 +646,7 @@ jobs:
     if: ${{ !contains(fromJSON(needs.setup-ci-load.outputs.jobs-to-skip), 'device-tests') }}
     uses: ./.github/workflows/unit-test.yml
     with:
+      job_runner_name: ${{ needs.determine_runner.outputs.runner_group }}
       job_name: ${{ inputs.job_name_prefix }}device-tests (${{ matrix.config.device }}, ${{ matrix.config.shots }}, ${{ matrix.python-version }})${{ inputs.job_name_suffix }}
       branch: ${{ inputs.branch }}
       coverage_artifact_name: devices-coverage-${{ matrix.config.device }}-${{ matrix.config.shots }}

.github/workflows/unit-test.yml

@@ -22,6 +22,11 @@ on:
         description: Name of the artifact file that will contain the coverage file for codevoc
         required: true
         type: string
+      job_runner_name:
+        description: The name of the runner to use for the job
+        required: false
+        type: string
+        default: 'ubuntu-latest'
       checkout_fetch_depth:
         description: How many commits to checkout from HEAD of branch passed
         required: false
@@ -81,7 +86,7 @@ on:
 jobs:
   test:
     name: ${{ inputs.job_name }}
-    runs-on: ubuntu-latest
+    runs-on: ${{ inputs.job_runner_name }}
 
     steps:
       - name: Checkout

.github/workflows/upload.yml

@@ -4,25 +4,34 @@ on:
     types: [published]
 
 jobs:
+  determine_runner:
+    name: Change to Large Runner
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-latest
+      force_large_runner: true
+
   tests:
     uses: ./.github/workflows/interface-unit-tests.yml
     secrets:
       codecov_token: ${{ secrets.CODECOV_TOKEN }}
     with:
       branch: ${{ github.ref }}
+      use_large_runner: true
 
   upload:
-    runs-on: ubuntu-latest
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
     # Make sure that the PennyLane is not released unless the tests are passing.
     needs:
       - tests
+      - determine_runner
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 1
 
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: "3.10"