On-Board Lightning Linux workflows (that are run on PRs) to large runners based on label

.github/workflows/build_and_cache_Kokkos_linux.yml

@@ -9,6 +9,14 @@ on:
       os:
         required: true
         type: string
+      runs_on:
+        description: |
+          The runner that should run the jobs. If left blank, the value from inputs.os is used.
+          This is useful if you want the jobs to run in a specific runner group, while not using that group name as part 
+          of the cache key.
+        required: false
+        type: string
+        default: ''
     outputs:
       exec_model:
         description: "The execution model for Kokkos."
@@ -43,7 +51,7 @@ jobs:
         kokkos_version: ${{ fromJson(needs.linux-set-builder-matrix.outputs.kokkos_version) }}
     timeout-minutes: 30
     name: Kokkos core (${{ matrix.exec_model }})
-    runs-on: ${{ inputs.os }}
+    runs-on: ${{ inputs.runs_on || inputs.os }}
 
     steps:
       - name: Cache installation directories

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

@@ -0,0 +1,40 @@
+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
+    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.default_runner }}
+
+env:
+  LARGE_RUNNER_GROUP_NAME: pl-4-core-large-runner
+
+jobs:
+  determine_workflow_runner:
+    runs-on: >-
+      ${{
+        (
+          github.event_name == 'pull_request'
+          && contains(github.event.pull_request.labels.*.name, 'urgent')
+        ) && 'pl-4-core-large-runner' || 'ubuntu-latest'
+      }}
+
+    outputs:
+      runner_group: ${{ steps.runner_group.outputs.runner_group }}
+
+    steps:
+      - name: Output Runner Group name
+        if: >-
+          ${{
+            github.event_name == 'pull_request'
+            && contains(github.event.pull_request.labels.*.name, 'urgent')
+            && startsWith(inputs.default_runner, 'ubuntu')
+          }}
+        id: runner_group
+        run: echo "runner_group=$LARGE_RUNNER_GROUP_NAME" >> $GITHUB_OUTPUT

.github/workflows/tests_lgpu_cpp.yml

@@ -42,8 +42,6 @@ jobs:
     strategy:
       max-parallel: 1
       matrix:
-        os: [ubuntu-22.04]
-        pl_backend: ["lightning_gpu"]
         cuda_version: ["12"]
 
     steps:
@@ -62,13 +60,12 @@ jobs:
     needs: [builddeps]
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_gpu"]
         cuda_version: ["12"]
 
     name: C++ tests (Lightning-GPU)
     runs-on:
-      - ${{ matrix.os }}
+      - ubuntu-22.04
       - self-hosted
       - gpu
 

.github/workflows/tests_lgpu_python.yml

@@ -42,8 +42,6 @@ jobs:
     strategy:
       max-parallel: 1
       matrix:
-        os: [ubuntu-22.04]
-        pl_backend: ["lightning_gpu"]
         cuda_version: ["12"]
 
     steps:
@@ -61,14 +59,13 @@ jobs:
     needs: [builddeps]
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_gpu"]
         default_backend: ["lightning_qubit"]
         cuda_version: ["12"]
 
     name: Python tests with LGPU
     runs-on:
-      - ${{ matrix.os }}
+      - ubuntu-22.04
       - self-hosted
       - gpu
 

.github/workflows/tests_linux_cpp.yml

@@ -33,15 +33,21 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  determine_runner:
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-22.04
+
   cpptests:
     if: ${{ !contains(fromJSON('["schedule", "workflow_dispatch"]'), github.event_name) }}
+    needs: [determine_runner]
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
     timeout-minutes: 60
     name: C++ tests
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - uses: actions/setup-python@v5
@@ -123,13 +129,13 @@ jobs:
 
   cpptestswithOpenBLAS:
     if: ${{ !contains(fromJSON('["schedule", "workflow_dispatch"]'), github.event_name) }}
+    needs: [determine_runner]
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
     timeout-minutes: 60
     name: C++ tests (OpenBLAS)
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - uses: actions/setup-python@v5
@@ -180,13 +186,15 @@ jobs:
 
   build_and_cache_Kokkos:
     name: "Build and cache Kokkos"
+    needs: [determine_runner]
     uses: ./.github/workflows/build_and_cache_Kokkos_linux.yml
     with:
+      runs_on: ${{ needs.determine_runner.outputs.runner_group }}
       os: ubuntu-22.04
 
   cpptestswithKokkos:
     if: ${{ !contains(fromJSON('["schedule", "workflow_dispatch"]'), github.event_name) }}
-    needs: [build_and_cache_Kokkos]
+    needs: [build_and_cache_Kokkos, determine_runner]
     strategy:
       matrix:
         os: [ubuntu-22.04]
@@ -195,7 +203,7 @@ jobs:
         kokkos_version: ${{ fromJson(needs.build_and_cache_Kokkos.outputs.kokkos_version) }}
     timeout-minutes: 60
     name: C++ tests (Kokkos)
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - uses: actions/setup-python@v5
@@ -281,15 +289,15 @@ jobs:
   cpptestsWithMultipleBackends:
   # Device-specific tests are performed for both. Device-agnostic tests default to LightningQubit.
     if: ${{ !contains(fromJSON('["schedule", "workflow_dispatch"]'), github.event_name) }}
-    needs: [build_and_cache_Kokkos]
+    needs: [build_and_cache_Kokkos, determine_runner]
     strategy:
       matrix:
         os: [ubuntu-22.04]
         exec_model: ${{ fromJson(needs.build_and_cache_Kokkos.outputs.exec_model) }}
         kokkos_version: ${{ fromJson(needs.build_and_cache_Kokkos.outputs.kokkos_version) }}
     timeout-minutes: 60
     name: C++ tests (multiple backends)
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - uses: actions/setup-python@v5

.github/workflows/tests_linux_python.yml

@@ -31,15 +31,21 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  determine_runner:
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-22.04
+
 
   pythontests:
+    needs: [determine_runner]
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
     timeout-minutes: 75
     name: Python tests
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - uses: actions/setup-python@v4
@@ -140,13 +146,13 @@ jobs:
 
 
   pythontestswithOpenBLAS:
+    needs: [determine_runner]
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
     timeout-minutes: 75
     name: Python tests with OpenBLAS
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - name: Checkout PennyLane-Lightning
@@ -242,12 +248,14 @@ jobs:
 
   build_and_cache_Kokkos:
     name: "Build and cache Kokkos"
+    needs: [determine_runner]
     uses: ./.github/workflows/build_and_cache_Kokkos_linux.yml
     with:
       os: ubuntu-22.04
+      runs_on: ${{ needs.determine_runner.outputs.runner_group }}
 
   pythontestswithKokkos:
-    needs: [build_and_cache_Kokkos]
+    needs: [build_and_cache_Kokkos, determine_runner]
     strategy:
       matrix:
         os: [ubuntu-22.04]
@@ -259,7 +267,7 @@ jobs:
             exec_model: OPENMP
     timeout-minutes: 75
     name: Python tests with Kokkos
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - name: Checkout PennyLane-Lightning

.github/workflows/tests_without_binary.yml

@@ -25,13 +25,19 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  determine_runner:
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-22.04
+
   pythontests:
+    needs: [determine_runner]
     timeout-minutes: 30
     name: Python tests
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
     strategy:
       matrix:
-        os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit", "lightning_kokkos", "lightning_gpu"]
 
     steps:

.github/workflows/wheel_linux_x86_64.yml

@@ -32,17 +32,22 @@ jobs:
     with:
       event_name: ${{ github.event_name }}
 
+  determine_runner:
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-latest
+
   build_dependencies:
-    needs: [set_wheel_build_matrix]
+    needs: [set_wheel_build_matrix, determine_runner]
     strategy:
       matrix:
-        os: [ubuntu-latest]
         exec_model: ${{ fromJson(needs.set_wheel_build_matrix.outputs.exec_model) }}
         kokkos_version: ${{ fromJson(needs.set_wheel_build_matrix.outputs.kokkos_version) }}
         container_img: ["quay.io/pypa/manylinux2014_x86_64"]
     timeout-minutes: 30
     name: Kokkos core (${{ matrix.exec_model }})
-    runs-on: ${{ matrix.os }}
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
     container: ${{ matrix.container_img }}
 
     steps:
@@ -90,20 +95,19 @@ jobs:
           cd -
 
   linux-wheels-x86-64:
-    needs: [set_wheel_build_matrix, build_dependencies]
+    needs: [set_wheel_build_matrix, build_dependencies, determine_runner]
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-latest]
         arch: [x86_64]
         pl_backend: ["lightning_kokkos", "lightning_qubit"]
         cibw_build: ${{ fromJson(needs.set_wheel_build_matrix.outputs.python_version) }}
         exec_model: ${{ fromJson(needs.set_wheel_build_matrix.outputs.exec_model) }}
         kokkos_version: ${{ fromJson(needs.set_wheel_build_matrix.outputs.kokkos_version) }}
         container_img: ["quay.io/pypa/manylinux2014_x86_64"]
     timeout-minutes: 30
-    name: ${{ matrix.os }}::${{ matrix.arch }} - ${{ matrix.pl_backend }} (Python ${{ fromJson('{ "cp39-*":"3.9","cp310-*":"3.10","cp311-*":"3.11","cp312-*":"3.12" }')[matrix.cibw_build] }})
-    runs-on: ${{ matrix.os }}
+    name: Linux::${{ matrix.arch }} - ${{ matrix.pl_backend }} (Python ${{ fromJson('{ "cp39-*":"3.9","cp310-*":"3.10","cp311-*":"3.11","cp312-*":"3.12" }')[matrix.cibw_build] }})
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
     container: ${{ matrix.container_img }}
 
     steps:

.github/workflows/wheel_linux_x86_64_cuda.yml

@@ -32,20 +32,26 @@ jobs:
     with:
       event_name: ${{ github.event_name }}
 
+  determine_runner:
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-latest
+
+
   linux-wheels-x86-64:
-    needs: [set_wheel_build_matrix]
+    needs: [set_wheel_build_matrix, determine_runner]
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-latest]
         arch: [x86_64]
         pl_backend: ["lightning_gpu"]
         cuda_version: ["12"]
         cibw_build: ${{ fromJson(needs.set_wheel_build_matrix.outputs.python_version) }}
         container_img: ["quay.io/pypa/manylinux2014_x86_64"]
     timeout-minutes: 30
-    name: ${{ matrix.os }}::${{ matrix.arch }} - ${{ matrix.pl_backend }} CUDA ${{ matrix.cuda_version }} (Python ${{ fromJson('{ "cp39-*":"3.9","cp310-*":"3.10","cp311-*":"3.11","cp312-*":"3.12" }')[matrix.cibw_build] }})
-    runs-on: ${{ matrix.os }}
+    name: Linux::${{ matrix.arch }} - ${{ matrix.pl_backend }} CUDA ${{ matrix.cuda_version }} (Python ${{ fromJson('{ "cp39-*":"3.9","cp310-*":"3.10","cp311-*":"3.11","cp312-*":"3.12" }')[matrix.cibw_build] }})
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
     container: ${{ matrix.container_img }}
 
     steps:

.github/workflows/wheel_noarch.yml

@@ -19,18 +19,25 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  determine_runner:
+    name: Determine runner type to use
+    uses: ./.github/workflows/determine-workflow-runner.yml
+    with:
+      default_runner: ubuntu-latest
+
+
   build-pure-python-wheel:
     if: |
       github.event_name != 'pull_request' ||
       contains(github.event.pull_request.labels.*.name, 'ci:build_wheels') ||
       github.event_name == 'workflow_dispatch'
+    needs: [determine_runner]
     strategy:
       matrix:
-        os: [ubuntu-latest]
         pl_backend: ["lightning_gpu", "lightning_kokkos", "lightning_qubit"]
     timeout-minutes: 30
-    name: ${{ matrix.os }} - Pure Python wheels - ${{ matrix.pl_backend }} (Python 3.9)
-    runs-on: ${{ matrix.os }}
+    name: Linux - Pure Python wheels - ${{ matrix.pl_backend }} (Python 3.9)
+    runs-on: ${{ needs.determine_runner.outputs.runner_group }}
 
     steps:
       - name: Checkout PennyLane-Lightning