Added ability to run linux workflows on large runners

.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"