From ee13e1d658d1d9109e3b8fdcce477a065614a01d Mon Sep 17 00:00:00 2001
From: Samuel Burnham <45365069+samuelburnham@users.noreply.github.com>
Date: Wed, 13 Dec 2023 13:43:22 -0500
Subject: [PATCH] Add custom `cargo` command on `workflow_dispatch`

---
 .github/workflows/cargo-command.yml | 56 +++++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)
 create mode 100644 .github/workflows/cargo-command.yml

diff --git a/.github/workflows/cargo-command.yml b/.github/workflows/cargo-command.yml
new file mode 100644
index 00000000..11ea9ba3
--- /dev/null
+++ b/.github/workflows/cargo-command.yml
@@ -0,0 +1,56 @@
+# Commit a temporary `cargo-command.env` on the desired local branch before running workflow
+name: Run custom `cargo` command
+
+on:
+  workflow_dispatch:
+    inputs:
+      command:
+        description: 'Cargo command to run'
+        required: true
+        type: string
+      runner:
+        description: 'GitHub Actions runner'
+        required: true
+        default: 'ubuntu-latest'
+        type: choice
+        options:
+          - ubuntu-latest
+          - buildjet-32vcpu-ubuntu-2204
+          - gpu-bench
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+  gpu-script:
+    name: Run custom `ci-script.sh` on GPU
+    if: github.event.inputs.runner == 'gpu-bench'
+    runs-on: [self-hosted, "${{ github.event.inputs.runner }}", "NVIDIA L4"]
+    steps:
+      - uses: actions/checkout@v4
+      # Install dependencies
+      - uses: actions-rs/toolchain@v1
+      - uses: Swatinem/rust-cache@v2
+      - name: Set up GPU if applicable
+        run: |
+          nvidia-smi
+          nvcc --version
+          CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | sed 's/\.//g')
+          echo "EC_GPU_CUDA_NVCC_ARGS=--fatbin --gpu-architecture=sm_$CUDA_ARCH --generate-code=arch=compute_$CUDA_ARCH,code=sm_$CUDA_ARCH" | tee -a $GITHUB_ENV
+          echo "CUDA_ARCH=$CUDA_ARCH" | tee -a $GITHUB_ENV
+          echo "EC_GPU_FRAMEWORK=cuda" | tee -a $GITHUB_ENV
+          env | grep -E "LURK|EC_GPU|CUDA"
+      - uses: cardinalby/export-env-action@v2
+        with:
+          envFile:
+            'cargo-command.env'
+      # Checks if command starts with `cargo`
+      - name: Run command
+        run: |
+          CARGO_CMD=$(echo ${{ inputs.command }} | awk '{ print $1 }')
+          if [ "$CARGO_CMD" != "cargo" ]; 
+          then
+            echo 'Invalid `cargo` command'
+            exit 1
+          else
+            ${{ inputs.command }}
+          fi
\ No newline at end of file