chore: add a new workflow for the tests (#565)

* update

* update

* add pr tests

* finished pr-test

* update

* update

* update

* update

* bring back test ampere

* update pr-test

* style: make linter happy

* change to on pr

.github/workflows/test-blackwell.yaml → .github/workflows/pr-test.yaml

-name: Blackwell Tests
+name: PR Tests
 on:
-  workflow_dispatch:
-    inputs:
-      test_target:
-        description: 'What to test: "pr" or "branch"'
-        required: true
-        type: choice
-        options:
-          - pr
-          - branch
-      pr_number:
-        description: 'Pull Request Number (only if test_target == "pr")'
-        required: false
-      branch_name:
-        description: 'Branch name (only if test_target == "branch")'
-        default: 'main'
-        required: false
-#  push:
-#    branches: [ main ]
-#    paths:
-#      - "nunchaku/**"
-#      - "src/**"
-#      - "tests/**"
-#      - "examples/**"
-#  pull_request:
-#    types: [ opened, synchronize, reopened, edited ]
-#    paths:
-#      - "nunchaku/**"
-#      - "src/**"
-#      - "tests/**"
-#      - "examples/**"
-#  issue_comment:
-#    types: [ created ]
-
+  pull_request:
+    types: [opened, synchronize, reopened]
+    paths:
+      - "nunchaku/**"
+      - "src/**"
+      - "tests/**"
+      - "examples/**"
+      - "scripts/**"
+  repository_dispatch:
+    types: [pr-test]
 concurrency:
   group: ${{ github.repository }}-${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
   cancel-in-progress: true
 jobs:
-  check-comment:
-    if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'issue_comment' && github.event.issue.pull_request && !github.event.pull_request.draft) }}
-    runs-on: [self-hosted, blackwell]
-    outputs:
-      should_run: ${{ steps.check.outputs.should_run }}
-    steps:
-      - id: check
-        run: |
-          body="${{ github.event.comment.body }}"
-          body_lower=$(echo "$body" | tr '[:upper:]' '[:lower:]')
-          if [[ "$body_lower" == "run tests" || "$body_lower" == "run test" ]]; then
-            echo "should_run=true" >> $GITHUB_OUTPUT
-          else
-            echo "should_run=false" >> $GITHUB_OUTPUT
-          fi
   run-tests:
-    runs-on: [self-hosted, blackwell]
-    needs: [check-comment]
-    if: ${{ github.event_name != 'issue_comment' || needs.check-comment.outputs.should_run == 'true' }}
+    strategy:
+      matrix:
+        gpu: ['3090', '5090']
+    runs-on:
+      - self-hosted
+      - ${{ matrix.gpu }}
+    # Only run for non-draft PRs, or for repository_dispatch events
+    if: |
+      (github.event_name == 'pull_request' && !github.event.pull_request.draft) ||
+      (github.event_name == 'repository_dispatch')
     steps:
       - name: Determine ref
         id: set-ref
         run: |
-          if [[ "${{ github.event.inputs.test_target }}" == "pr" ]]; then
-            echo "ref=refs/pull/${{ github.event.inputs.pr_number }}/merge" >> $GITHUB_OUTPUT
+          if [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            echo "ref=${{ github.event.pull_request.head.ref }}" >> $GITHUB_OUTPUT
+          elif [[ "${{ github.event_name }}" == "repository_dispatch" ]]; then
+            # Use github.event.client_payload.ref if provided, else default to main
+            if [[ -n "${{ github.event.client_payload.ref }}" ]]; then
+              echo "ref=${{ github.event.client_payload.ref }}" >> $GITHUB_OUTPUT
             else
-            echo "ref=refs/heads/${{ github.event.inputs.branch_name }}" >> $GITHUB_OUTPUT
+              echo "ref=main" >> $GITHUB_OUTPUT
+            fi
           fi
       - name: Checkout
         uses: actions/checkout@v4
@@ -76,7 +50,7 @@ jobs:
         run: |
           which python
           echo "Setting up Python with Conda"
-          conda create -n test_env python=3.11 -y
+          conda create -n test_env python=3.12 -y
       - name: Install dependencies
         run: |
           source $(conda info --base)/etc/profile.d/conda.sh
@@ -102,11 +76,11 @@ jobs:
           pwd
           cd ..
           pip install comfy-cli
-          yes | comfy --here install --nvidia --skip-torch-or-directml --version 0.3.43
+          yes | comfy --here install --nvidia --skip-torch-or-directml --version 0.3.44
           cd ComfyUI
           rm -r models
-          mkdir -p ${{ secrets.COMFYUI_MODELS_ROOT_BLACKWELL }}
-          ln -s ${{ secrets.COMFYUI_MODELS_ROOT_BLACKWELL }} models
+          mkdir -p ${{ secrets.COMFYUI_MODELS_ROOT_AMPERE }}
+          ln -s ${{ secrets.COMFYUI_MODELS_ROOT_AMPERE }} models
           cd custom_nodes
           git clone -b dev https://github.com/mit-han-lab/ComfyUI-nunchaku.git
           cd ..
@@ -125,33 +99,15 @@ jobs:
           cd ../ComfyUI
           python nunchaku_tests/scripts/nunchaku-flux1-dev.py
           pytest -v nunchaku_tests/
-      - name: Nunchaku FLUX memory tests
-        run: |
-          pwd
-          source $(conda info --base)/etc/profile.d/conda.sh
-          conda activate test_env || { echo "Failed to activate conda env"; exit 1; }
-          which python
-          NUNCHAKU_TEST_CACHE_ROOT=${{ secrets.NUNCHAKU_TEST_CACHE_ROOT_BLACKWELL }} HF_TOKEN=${{ secrets.HF_TOKEN }} pytest -v tests/flux/test_flux_memory.py
-      - name: Nunchaku FLUX example tests
-        run: |
-          source $(conda info --base)/etc/profile.d/conda.sh
-          conda activate test_env || { echo "Failed to activate conda env"; exit 1; }
-          which python
-          NUNCHAKU_TEST_CACHE_ROOT=${{ secrets.NUNCHAKU_TEST_CACHE_ROOT_BLACKWELL }} HF_TOKEN=${{ secrets.HF_TOKEN }} pytest -v tests/flux/test_flux_examples.py
-      - name: Nunchaku FLUX other tests
-        run: |
-          source $(conda info --base)/etc/profile.d/conda.sh
-          conda activate test_env || { echo "Failed to activate conda env"; exit 1; }
-          which python
-          NUNCHAKU_TEST_CACHE_ROOT=${{ secrets.NUNCHAKU_TEST_CACHE_ROOT_BLACKWELL }} HF_TOKEN=${{ secrets.HF_TOKEN }} pytest -v tests/flux --ignore=tests/flux/test_flux_memory.py --ignore=tests/flux/test_flux_examples.py
-      - name: Nunchaku SANA tests
+      - name: Run nunchaku tests
         run: |
           source $(conda info --base)/etc/profile.d/conda.sh
           conda activate test_env || { echo "Failed to activate conda env"; exit 1; }
           which python
-          NUNCHAKU_TEST_CACHE_ROOT=${{ secrets.NUNCHAKU_TEST_CACHE_ROOT_BLACKWELL }} HF_TOKEN=${{ secrets.HF_TOKEN }} pytest -v tests/sana
+          NUNCHAKU_TEST_CACHE_ROOT=${{ secrets.NUNCHAKU_TEST_CACHE_ROOT_AMPERE }} HF_TOKEN=${{ secrets.HF_TOKEN }} pytest -v tests/flux/test_flux_examples.py
+          NUNCHAKU_TEST_CACHE_ROOT=${{ secrets.NUNCHAKU_TEST_CACHE_ROOT_AMPERE }} HF_TOKEN=${{ secrets.HF_TOKEN }} python scripts/run_all_tests.py
       - name: clean up
-        if: always() && (github.event_name != 'issue_comment' || needs.check-comment.outputs.should_run == 'true')
+        if: always()
         run: |
           cd ..
-          rm -rf ComfyUI
+          rm -rf ComfyUI ComfyUI-nunchaku

scripts/run_all_tests.py

+import subprocess
+from pathlib import Path
+
+
+def run_all_tests():
+    test_dir = Path(__file__).parent.parent / "tests"
+    test_files = []
+    for file_path in test_dir.rglob("test_*.py"):
+        # Ignore tests/flux/test_flux_examples.py
+        try:
+            rel_path = file_path.relative_to(test_dir)
+        except ValueError:
+            continue
+        if str(rel_path) == "flux/test_flux_examples.py":
+            continue
+        test_files.append(str(file_path))
+    print("Running tests:")
+    for test_file in sorted(test_files):
+        print(f"  {test_file}")
+
+    failed_tests = []
+    for test_file in test_files:
+        print(f"Running {test_file} ...")
+        result = subprocess.run(["pytest", test_file])
+        if result.returncode != 0:
+            print(f"Test failed: {test_file}")
+            failed_tests.append(test_file)
+        else:
+            print(f"Test passed: {test_file}")
+
+    if failed_tests:
+        print("Some tests failed.")
+        for test_file in failed_tests:
+            print(f"  {test_file}")
+        exit(1)
+    else:
+        print("All tests passed.")
+
+
+if __name__ == "__main__":
+    run_all_tests()