ci: allow running multi-gpu tests through diffusers-bot using HF Jobs

.github/workflows/bot_pytest.yml

@@ -1,14 +1,20 @@
 name: "Bot: Pytest"
 
 # Reusable workflow called by diffusers_bot.yml for `/diffusers-bot pytest <args>`.
-# Runs pytest on a GPU runner and reports the result back on the PR.
 
 on:
   workflow_call:
     inputs:
       pytest_args:
         required: true
         type: string
+      pytest_flavor:
+        required: false
+        type: string
+        default: ''
+    secrets:
+      HF_JOBS_TOKEN:
+        required: false
 
 jobs:
   gate:
@@ -35,9 +41,15 @@ jobs:
           PR: ${{ github.event.issue.number }}
           COMMENTER: ${{ github.event.comment.user.login }}
           PYTEST_ARGS: ${{ inputs.pytest_args }}
+          PYTEST_FLAVOR: ${{ inputs.pytest_flavor }}
           RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
         run: |
-          BODY="⏳ Running \`pytest ${PYTEST_ARGS}\` on a GPU runner — [view logs](${RUN_URL}).
+          if [[ -n "$PYTEST_FLAVOR" ]]; then
+            LOCATION="HF Jobs (${PYTEST_FLAVOR})"
+          else
+            LOCATION="a GPU runner"
+          fi
+          BODY="⏳ Running \`pytest ${PYTEST_ARGS}\` on ${LOCATION} — [view logs](${RUN_URL}).
 
           Triggered by @${COMMENTER}."
           CID=$(gh api -X POST "repos/${REPO}/issues/${PR}/comments" -f body="$BODY" --jq '.id')
@@ -46,6 +58,7 @@ jobs:
   gpu:
     name: Run pytest on GPU
     needs: gate
+    if: ${{ inputs.pytest_flavor == '' }}
     # A newer command on the same PR supersedes an in-flight one. Scoped to this job
     # only so the superseded run's `report` still updates its comment.
     concurrency:
@@ -121,9 +134,53 @@ jobs:
           name: bot_gpu_test_reports
           path: reports
 
+  hf_job:
+    name: Run pytest on HF Jobs
+    needs: gate
+    if: ${{ inputs.pytest_flavor != '' }}
+    concurrency:
+      group: diffusers-bot-pytest-${{ github.event.issue.number }}
+      cancel-in-progress: true
+    runs-on: ubuntu-22.04
+    timeout-minutes: 150
+    permissions:
+      contents: read
+      pull-requests: read
+    outputs:
+      job_url: ${{ steps.run.outputs.job_url }}
+    steps:
+      - name: Resolve PR head
+        id: pr
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          REPO: ${{ github.repository }}
+          PR: ${{ github.event.issue.number }}
+        run: |
+          echo "sha=$(gh api "repos/${REPO}/pulls/${PR}" --jq '.head.sha')" >> "$GITHUB_OUTPUT"
+
+      - name: Checkout workflow code
+        uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6.1.0
+        with:
+          persist-credentials: false
+
+      - name: Install Hub client
+        run: python -m pip install 'huggingface_hub>=1.32,<2'
+
+      - name: Run HF Job
+        id: run
+        env:
+          HF_TOKEN: ${{ secrets.HF_JOBS_TOKEN }}
+          HF_JOBS_NAMESPACE: ${{ vars.DIFFUSERS_HF_JOBS_NAMESPACE || 'diffusers-internal-dev' }}
+          REPOSITORY: ${{ github.repository }}
+          PR_NUMBER: ${{ github.event.issue.number }}
+          PR_SHA: ${{ steps.pr.outputs.sha }}
+          PYTEST_ARGS: ${{ inputs.pytest_args }}
+          PYTEST_FLAVOR: ${{ inputs.pytest_flavor }}
+        run: python utils/bot_pytest_hf_job.py
+
   report:
     name: Report status
-    needs: [gate, gpu]
+    needs: [gate, gpu, hf_job]
     # Always run so the comment is updated on success, failure, or cancellation —
     # but only if `gate` actually posted a comment to update.
     if: ${{ always() && needs.gate.outputs.comment_id != '' }}
@@ -136,8 +193,10 @@ jobs:
           GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           REPO: ${{ github.repository }}
           CID: ${{ needs.gate.outputs.comment_id }}
-          RESULT: ${{ needs.gpu.result }}
+          RESULT: ${{ inputs.pytest_flavor == '' && needs.gpu.result || needs.hf_job.result }}
           PYTEST_ARGS: ${{ inputs.pytest_args }}
+          PYTEST_FLAVOR: ${{ inputs.pytest_flavor }}
+          JOB_URL: ${{ needs.hf_job.outputs.job_url }}
           RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
         run: |
           case "$RESULT" in
@@ -146,5 +205,13 @@ jobs:
             cancelled) EMOJI="⚠️"; MSG="was cancelled";;
             *)         EMOJI="⚠️"; MSG="did not run (${RESULT})";;
           esac
-          BODY="${EMOJI} \`pytest ${PYTEST_ARGS}\` ${MSG} on GPU — [view logs](${RUN_URL})."
+          if [[ -n "$PYTEST_FLAVOR" ]]; then
+            LOCATION="HF Jobs (${PYTEST_FLAVOR})"
+            if [[ -n "$JOB_URL" ]]; then
+              RUN_URL="$JOB_URL"
+            fi
+          else
+            LOCATION="GPU"
+          fi
+          BODY="${EMOJI} \`pytest ${PYTEST_ARGS}\` ${MSG} on ${LOCATION} — [view logs](${RUN_URL})."
           gh api -X PATCH "repos/${REPO}/issues/comments/${CID}" -f body="$BODY"

.github/workflows/diffusers_bot.yml

@@ -6,6 +6,7 @@ name: Diffusers Bot
 #   /diffusers-bot review          request a Serge AI review (also works from an inline review comment)
 #   /diffusers-bot pytest <args>   run `pytest <args>` on a GPU runner, e.g.
 #                                  `/diffusers-bot pytest tests/models/test_modeling_common.py -k "some_test"`
+#   /diffusers-bot pytest --flavor a100x4 <args>   run pytest on a multi-GPU HF Job
 #
 # Auth is checked once in the `authorize` job. Each subcommand lives in its
 # own reusable workflow file under .github/workflows/bot_*.yml.
@@ -43,6 +44,7 @@ jobs:
       is_authorized: ${{ steps.check.outputs.is_authorized }}
       subcommand: ${{ steps.check.outputs.subcommand }}
       pytest_args: ${{ steps.check.outputs.pytest_args }}
+      pytest_flavor: ${{ steps.check.outputs.pytest_flavor }}
     steps:
       - name: Authorize and parse command
         id: check
@@ -64,10 +66,18 @@ jobs:
 
             let subcommand = 'unknown';
             let pytestArgs = '';
+            let pytestFlavor = '';
 
-            if (firstLine.startsWith('/diffusers-bot pytest')) {
+            if (/^\/diffusers-bot pytest(?:\s|$)/.test(firstLine)) {
               subcommand = 'pytest';
-              pytestArgs = firstLine.replace('/diffusers-bot pytest', '').trim();
+              pytestArgs = firstLine.slice('/diffusers-bot pytest'.length).trim();
+              const flavor = pytestArgs.match(/^--flavor(?:=|\s+)(\S+)(?:\s+|$)/);
+              if (flavor) {
+                pytestFlavor = flavor[1];
+                pytestArgs = pytestArgs.slice(flavor[0].length).trim();
+              } else if (pytestArgs.startsWith('--flavor')) {
+                pytestFlavor = 'invalid';
+              }
             } else if (body.includes('/diffusers-bot style')) {
               subcommand = 'style';
             } else if (body.includes('/diffusers-bot review')) {
@@ -77,6 +87,7 @@ jobs:
             console.log(`Subcommand: ${subcommand}, pytestArgs: ${pytestArgs}`);
             core.setOutput('subcommand', subcommand);
             core.setOutput('pytest_args', pytestArgs);
+            core.setOutput('pytest_flavor', pytestFlavor);
 
   # ── Dispatch to sub-workflows ────────────────────────────────────────────────
   style:
@@ -117,3 +128,6 @@ jobs:
     uses: ./.github/workflows/bot_pytest.yml
     with:
       pytest_args: ${{ needs.authorize.outputs.pytest_args }}
+      pytest_flavor: ${{ needs.authorize.outputs.pytest_flavor }}
+    secrets:
+      HF_JOBS_TOKEN: ${{ secrets.DIFFUSERS_HF_JOBS_TOKEN }}

utils/bot_pytest_hf_job.py

@@ -0,0 +1,108 @@
+import os
+import signal
+import sys
+
+from huggingface_hub import HfApi
+
+
+COMMAND = r"""
+set -euo pipefail
+git init -q /tmp/diffusers
+cd /tmp/diffusers
+git remote add origin "https://github.com/${REPOSITORY}.git"
+git fetch -q --depth=2 origin "refs/pull/${PR_NUMBER}/head"
+git checkout -q --detach FETCH_HEAD
+test "$(git rev-parse HEAD)" = "$PR_SHA"
+nvidia-smi
+printf 'torch==2.10.0\ntorchvision==0.25.0\ntorchaudio==2.10.0\n' > "$UV_OVERRIDE"
+uv pip install -e ".[quality,training,test]"
+uv pip install peft@git+https://github.com/huggingface/peft.git
+uv pip uninstall accelerate && uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git
+uv pip uninstall transformers huggingface_hub && UV_PRERELEASE=allow uv pip install -U transformers@git+https://github.com/huggingface/transformers.git
+diffusers-cli env
+status=0
+python - <<'PY' || status=$?
+import os
+import shlex
+import subprocess
+
+args = shlex.split(os.environ["PYTEST_ARGS"])
+raise SystemExit(subprocess.call(["pytest", "--make-reports=tests_bot_gpu", *args]))
+PY
+if [[ "$status" -ne 0 ]]; then
+    cat reports/tests_bot_gpu_stats.txt || true
+    cat reports/tests_bot_gpu_failures_short.txt || true
+fi
+exit "$status"
+"""
+
+
+def main():
+    flavor = os.environ["PYTEST_FLAVOR"]
+    namespace = os.environ["HF_JOBS_NAMESPACE"]
+    token = os.environ["HF_TOKEN"]
+    if not namespace:
+        raise SystemExit("HF_JOBS_NAMESPACE must be set before running HF Jobs")
+    if not token:
+        raise SystemExit("Set the DIFFUSERS_HF_JOBS_TOKEN repository secret before running HF Jobs")
+
+    api = HfApi(token=token)
+    flavors = {
+        hardware.name
+        for hardware in api.list_jobs_hardware()
+        if hardware.accelerator is not None
+        and hardware.accelerator.type.lower() == "gpu"
+        and int(hardware.accelerator.quantity) > 1
+    }
+    if flavor not in flavors:
+        raise SystemExit(f"Unsupported multi-GPU flavor: {flavor}. Choose from: {', '.join(sorted(flavors))}")
+
+    job = api.run_job(
+        image="diffusers/diffusers-pytorch-cuda",
+        command=["bash", "-lc", COMMAND],
+        flavor=flavor,
+        namespace=namespace,
+        timeout="2h",
+        name=f"diffusers-pr-{os.environ['PR_NUMBER']}-pytest",
+        env={
+            "REPOSITORY": os.environ["REPOSITORY"],
+            "PR_NUMBER": os.environ["PR_NUMBER"],
+            "PR_SHA": os.environ["PR_SHA"],
+            "PYTEST_ARGS": os.environ["PYTEST_ARGS"],
+            "DIFFUSERS_IS_CI": "yes",
+            "OMP_NUM_THREADS": "8",
+            "MKL_NUM_THREADS": "8",
+            "HF_XET_HIGH_PERFORMANCE": "1",
+            "PYTEST_TIMEOUT": "600",
+            "UV_OVERRIDE": "/tmp/uv-overrides.txt",
+            "CUBLAS_WORKSPACE_CONFIG": ":16:8",
+        },
+    )
+    with open(os.environ["GITHUB_OUTPUT"], "a") as output:
+        output.write(f"job_url={job.url}\n")
+    print(f"HF Job: {job.url}", flush=True)
+
+    def cancel_job(signum, frame):
+        api.cancel_job(job_id=job.id, namespace=namespace)
+        raise SystemExit(128 + signum)
+
+    signal.signal(signal.SIGINT, cancel_job)
+    signal.signal(signal.SIGTERM, cancel_job)
+
+    try:
+        finished = api.wait_for_job(job_id=job.id, namespace=namespace, poll_interval=20)
+    except Exception:
+        api.cancel_job(job_id=job.id, namespace=namespace)
+        raise
+    try:
+        for line in api.fetch_job_logs(job_id=job.id, namespace=namespace, tail=200):
+            print(line, flush=True)
+    except Exception as error:
+        print(f"Could not fetch HF Job logs: {error}", flush=True)
+    print(f"HF Job status: {finished.status.stage}", flush=True)
+    if finished.status.stage != "COMPLETED":
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()