Fixing nightly Github Action for AWS setup

.github/workflows/nightly.yml

@@ -15,108 +15,77 @@ permissions:
   id-token: write
 
 jobs:
+  ##############################################################################
+  # BUILD FVDB
+  ##############################################################################
+  start-build-runner:
+    name: Start CPU-only EC2 runner for build
+    runs-on: ubuntu-latest
+    if: github.repository == 'openvdb/fvdb-core'  # Prevent running in forks
+    outputs:
+      label: ${{ steps.start-build-runner.outputs.label }}
+      ec2-instance-id: ${{ steps.start-build-runner.outputs.ec2-instance-id }}
+    steps:
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: arn:aws:iam::420032683002:role/openvdb-fvdb-github-actions-role
+          aws-region: us-east-2
+      - name: Start EC2 runner
+        id: start-build-runner
+        uses: machulav/ec2-github-runner@v2.4.3
+        with:
+          mode: start
+          github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
+          ec2-image-id: ami-0e14a711dad782a70
+          ec2-instance-type: m6a.8xlarge
+          subnet-id: subnet-03f2320d6e6e0005b
+          security-group-id: sg-0cd08bd89d6212223
+
   fvdb-build:
-    if: ${{ !startsWith(github.event.pull_request.title, 'Draft:') }}
     name: fVDB Build
-    runs-on:
-      - self-hosted
+    needs: start-build-runner # required to start the main job when the runner is ready
+    runs-on: ${{ needs.start-build-runner.outputs.label }} # run the job on the newly created runner
     container:
-      image: aswf/ci-openvdb:2024
+      image: aswf/ci-openvdb:2024-clang17.2
       env:
         PYTHONPATH: ""
         CPM_SOURCE_CACHE: "/__w/cpm_cache"
+        CONDA_OVERRIDE_CUDA: "12.9"  # this is to build an environment on machines that lack a CUDA device and needs to be >= the CUDA version in the build_environment.yml file
       options: --rm
     defaults:
       run:
         shell: bash -el {0}
     steps:
-      ### Get the PR branch and apply changes to the target branch
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0 # Fetch all history for all branches
-      - name: Determine the target branch of the PR
-        run: |
-          echo "Fetching pull request metadata..."
-          # Install GitHub CLI and jq
-          dnf config-manager --add-repo https://cli.github.com/packages/rpm/gh-cli.repo
-          yum install -y gh
-          which jq || yum install -y jq
-          # Extract PR number from branch name if it follows pattern "pull-request/123"
-          BRANCH_NAME=${GITHUB_REF#refs/heads/}
-          echo "Branch name: $BRANCH_NAME"
-
-          # Setup GitHub CLI authentication
-          echo "${{ github.token }}" | gh auth login --with-token
-
-          # Get default branch from repository
-          echo "Getting default branch from repository..."
-          DEFAULT_BRANCH=$(gh repo view ${{ github.repository }} --json defaultBranchRef --jq '.defaultBranchRef.name' 2>/dev/null || echo 'develop')
-          echo "Repository default branch is: $DEFAULT_BRANCH"
-
-          if [[ $BRANCH_NAME =~ pull-request/([0-9]+) ]]; then
-            PR_NUMBER=${BASH_REMATCH[1]}
-            echo "Found PR number from branch name: $PR_NUMBER"
-
-            echo "Using gh CLI to get PR details..."
-            # Get PR info using GitHub CLI and extract the base branch
-            PR_INFO=$(gh pr view $PR_NUMBER --json baseRefName --repo ${{ github.repository }} 2>/dev/null || echo '{"baseRefName":"'$DEFAULT_BRANCH'"}')
-            TARGET_BRANCH=$(echo "$PR_INFO" | jq -r '.baseRefName')
-
-            echo "Target branch is $TARGET_BRANCH"
-
-            if [ -z "$TARGET_BRANCH" ] || [ "$TARGET_BRANCH" = "null" ]; then
-              echo "Could not determine target branch, using default branch"
-              TARGET_BRANCH="$DEFAULT_BRANCH"
-            fi
-          else
-            echo "Branch doesn't match the expected pattern, using default branch"
-            TARGET_BRANCH="$DEFAULT_BRANCH"
-          fi
+          ref: 'main'
 
-          echo "Final TARGET_BRANCH=$TARGET_BRANCH"
-          echo "TARGET_BRANCH=$TARGET_BRANCH" >> $GITHUB_ENV
-      - name: Set up Git user
-        run: |
-          git config --global user.name "github-actions[bot]"
-          git config --global user.email "github-actions[bot]@users.noreply.github.com"
-      - name: Merge pushed branch into target branch
-        run: |
-          if [ -z "$TARGET_BRANCH" ]; then
-            echo "No pull request metadata found. Skipping merge."
-            exit 0
-          fi
-          git config --global --add safe.directory "$(pwd)"
-          git fetch origin $TARGET_BRANCH
-          git checkout $TARGET_BRANCH
-          git merge $GITHUB_REF --no-ff --message "Merging $GITHUB_REF into $TARGET_BRANCH"
-        env:
-          GITHUB_REF: ${{ github.ref }}
-      #### End of git merge
 
       - name: Set up fvdb_build Conda env
-        uses: conda-incubator/setup-miniconda@v3
+        uses: mamba-org/setup-micromamba@v2
         with:
-          miniforge-version: latest
-          conda-remove-defaults: "true"
-          activate-environment: fvdb_build
-          environment-file: fvdb/env/build_environment.yml
+          post-cleanup: 'all'
+          environment-name: fvdb_build
+          environment-file: env/build_environment.yml
 
       - name: Build fvdb
         run: |
-          cd fvdb;
-          TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6+PTX" conda run --no-capture-output -n fvdb_build ./build.sh wheel verbose gtests benchmarks
+          micromamba activate fvdb_build
+          ./build.sh wheel verbose gtests benchmarks --cuda-arch-list '8.9+PTX'
 
       - name: Upload wheel
         uses: actions/upload-artifact@v4
         with:
             name: fvdb-test-package
-            path: fvdb/dist/*.whl
+            path: dist/*.whl
             retention-days: 2
 
       #NOTE: This tar step is here because directly uploading the build directory
       #      wasn't working due to losing executable permissions on the files.
       - name: Tar build directory
-        run: tar -cvf fvdb-gtest.tar fvdb/build/
+        run: tar -cvf fvdb-gtest.tar build/
 
       - name: Upload gtests
         uses: actions/upload-artifact@v4
@@ -125,22 +94,67 @@ jobs:
             path: fvdb-gtest.tar
             retention-days: 2
 
-      - name: Cleanup
-        if: always()
-        run: |
-          echo "Cleaning up /__w/_temp directory"
-          sudo rm -rf /__w/_temp/*
-          echo "Cleanup completed"
-
+  fvdb-build-stop-runner:
+    name: Stop CPU-only EC2 runner for build
+    needs:
+      - start-build-runner # required to get output from the start-build-runner job
+      - fvdb-build # required to wait when the main job is done
+    runs-on: ubuntu-latest
+    if: ${{ always() }} # required to stop the runner even if the error happened in the previous jobs
+    steps:
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: arn:aws:iam::420032683002:role/openvdb-fvdb-github-actions-role
+          aws-region: us-east-2
+      - name: Stop EC2 runner
+        uses: machulav/ec2-github-runner@v2.4.3
+        with:
+          mode: stop
+          github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
+          label: ${{ needs.start-build-runner.outputs.label }}
+          ec2-instance-id: ${{ needs.start-build-runner.outputs.ec2-instance-id }}
+
+
+  ##############################################################################
+  # START FVDB BENCHMARKS GPU RUNNER
+  ##############################################################################
+  start-benchmarks-gpu-runner:
+    name: Start EC2 GPU runner for benchmarks
+    needs: fvdb-build
+    runs-on: ubuntu-latest
+    outputs:
+      label: ${{ steps.start-benchmarks-gpu-runner.outputs.label }}
+      ec2-instance-id: ${{ steps.start-benchmarks-gpu-runner.outputs.ec2-instance-id }}
+    steps:
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: arn:aws:iam::420032683002:role/openvdb-fvdb-github-actions-role
+          aws-region: us-east-2
+      - name: Start EC2 GPU runner
+        id: start-benchmarks-gpu-runner
+        uses: machulav/ec2-github-runner@v2.4.3
+        with:
+          mode: start
+          github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
+          ec2-image-id: ami-0f5c0b65fde08ae43
+          ec2-instance-type: g6.xlarge # 4 CPU-core, L4 GPU
+          subnet-id: subnet-03f2320d6e6e0005b
+          security-group-id: sg-0cd08bd89d6212223
+
+  ##############################################################################
+  # RUN FVDB BENCHMARKS
+  ##############################################################################
   fvdb-benchmarks:
-    needs: [fvdb-build]
-    name: fVDB Continuous Benchmarking
-    runs-on:
-      - self-hosted
+    needs: start-benchmarks-gpu-runner # required to start the main job when the runner is ready
+    name: fVDB Benchmarks
+    runs-on: ${{ needs.start-benchmarks-gpu-runner.outputs.label }} # run the job on the newly created runner
     container:
-      image: aswf/ci-openvdb:2024
+      image: aswf/ci-openvdb:2024-clang17.2
       env:
         PYTHONPATH: ""
+        CPM_SOURCE_CACHE: "/__w/cpm_cache"
       options: --rm
     defaults:
       run:
@@ -150,64 +164,7 @@ jobs:
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0 # Fetch all history for all branches
-      - name: Determine the target branch of the PR
-        run: |
-          echo "Fetching pull request metadata..."
-          # Install GitHub CLI and jq
-          dnf config-manager --add-repo https://cli.github.com/packages/rpm/gh-cli.repo
-          yum install -y gh
-          which jq || yum install -y jq
-          # Extract PR number from branch name if it follows pattern "pull-request/123"
-          BRANCH_NAME=${GITHUB_REF#refs/heads/}
-          echo "Branch name: $BRANCH_NAME"
-
-          # Setup GitHub CLI authentication
-          echo "${{ github.token }}" | gh auth login --with-token
-
-          # Get default branch from repository
-          echo "Getting default branch from repository..."
-          DEFAULT_BRANCH=$(gh repo view ${{ github.repository }} --json defaultBranchRef --jq '.defaultBranchRef.name' 2>/dev/null || echo 'develop')
-          echo "Repository default branch is: $DEFAULT_BRANCH"
-
-          if [[ $BRANCH_NAME =~ pull-request/([0-9]+) ]]; then
-            PR_NUMBER=${BASH_REMATCH[1]}
-            echo "Found PR number from branch name: $PR_NUMBER"
-
-            echo "Using gh CLI to get PR details..."
-            # Get PR info using GitHub CLI and extract the base branch
-            PR_INFO=$(gh pr view $PR_NUMBER --json baseRefName --repo ${{ github.repository }} 2>/dev/null || echo '{"baseRefName":"'$DEFAULT_BRANCH'"}')
-            TARGET_BRANCH=$(echo "$PR_INFO" | jq -r '.baseRefName')
-
-            echo "Target branch is $TARGET_BRANCH"
-
-            if [ -z "$TARGET_BRANCH" ] || [ "$TARGET_BRANCH" = "null" ]; then
-              echo "Could not determine target branch, using default branch"
-              TARGET_BRANCH="$DEFAULT_BRANCH"
-            fi
-          else
-            echo "Branch doesn't match the expected pattern, using default branch"
-            TARGET_BRANCH="$DEFAULT_BRANCH"
-          fi
-
-          echo "Final TARGET_BRANCH=$TARGET_BRANCH"
-          echo "TARGET_BRANCH=$TARGET_BRANCH" >> $GITHUB_ENV
-      - name: Set up Git user
-        run: |
-          git config --global user.name "github-actions[bot]"
-          git config --global user.email "github-actions[bot]@users.noreply.github.com"
-      - name: Merge pushed branch into target branch
-        run: |
-          if [ -z "$TARGET_BRANCH" ]; then
-            echo "No pull request metadata found. Skipping merge."
-            exit 0
-          fi
-          git config --global --add safe.directory "$(pwd)"
-          git fetch origin $TARGET_BRANCH
-          git checkout $TARGET_BRANCH
-          git merge $GITHUB_REF --no-ff --message "Merging $GITHUB_REF into $TARGET_BRANCH"
-        env:
-          GITHUB_REF: ${{ github.ref }}
-      #### End of git merge
+          ref: 'main'
 
       - name: Set up fvdb_test Conda env
         uses: conda-incubator/setup-miniconda@v3
@@ -242,19 +199,43 @@ jobs:
         with:
             name: Python Benchmark with pytest-benchmark
             tool: 'pytest'
-            output-file-path: fvdb/tests/benchmark/output.json
+            output-file-path: tests/benchmark/output.json
             # Use personal access token instead of GITHUB_TOKEN due to https://github.community/t/github-action-not-triggering-gh-pages-upon-push/16096
             github-token: ${{ secrets.GITHUB_TOKEN }}
             auto-push: true
             # Show alert with commit comment on detecting possible performance regression
             alert-threshold: '200%'
             comment-on-alert: true
             fail-on-alert: true
-            alert-comment-cc-users: '@NVIDIA-Omniverse/fvdb-dev'
+            alert-comment-cc-users: '@openvdb/fvdb-dev'
 
       - name: Cleanup
         if: always()
         run: |
           echo "Cleaning up /__w/_temp directory"
           sudo rm -rf /__w/_temp/*
           echo "Cleanup completed"
+
+  ##############################################################################
+  # STOP FVDB BENCHMARKS GPU RUNNER
+  ##############################################################################
+  fvdb-benchmarks-stop-gpu-runner:
+    name: Stop GPU EC2 runner for benchmarks
+    needs:
+      - start-benchmarks-gpu-runner # required to get output from the start-benchmarks-gpu-runner job
+      - fvdb-benchmarks # required to wait when the main job is done
+    runs-on: ubuntu-latest
+    if: ${{ always() }} # required to stop the runner even if the error happened in the previous jobs
+    steps:
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: arn:aws:iam::420032683002:role/openvdb-fvdb-github-actions-role
+          aws-region: us-east-2
+      - name: Stop EC2 runner
+        uses: machulav/ec2-github-runner@v2.4.3
+        with:
+          mode: stop
+          github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
+          label: ${{ needs.start-benchmarks-gpu-runner.outputs.label }}
+          ec2-instance-id: ${{ needs.start-benchmarks-gpu-runner.outputs.ec2-instance-id }}

tests/wip/benchmark/benchmark/test_conv.py → tests/benchmark/test_conv.py

@@ -6,6 +6,7 @@
 import fvdb.nn as fvdbnn
 import pytest
 import torch
+from fvdb.convolution_plan import _CUTLASS_SUPPORTED_CHANNELS
 
 import fvdb
 
@@ -16,32 +17,37 @@
 PTS_CACHE = [torch.empty((10_000, 3), dtype=torch.float32).normal_() for _ in range(100)]
 
 
-@pytest.mark.parametrize("i_ch", [3, 8, 16, 32, 64, 128])
-@pytest.mark.parametrize("o_ch", [3, 8, 16, 32, 64, 128])
+@pytest.mark.parametrize("i_ch", [3, 8, 16, 32, 64, 128, 256, 512])
+@pytest.mark.parametrize("o_ch", [3, 8, 16, 32, 64, 128, 256, 512])
 @pytest.mark.parametrize("backend", ["default", "cutlass", "me", "halo", "igemm_mode0", "igemm_mode1", "igemm_mode2"])
 @pytest.mark.benchmark(
     group="sparse_conv3d",
     warmup=True,
     warmup_iterations=3,
 )
 def test_forward_conv3d(benchmark, i_ch, o_ch, backend):
+    if backend == "cutlass":
+        if (i_ch, o_ch) not in _CUTLASS_SUPPORTED_CHANNELS:
+            pytest.skip(f"Cutlass backend does not support channel pair {i_ch, o_ch}")
+
     device = torch.device("cuda", torch.cuda.current_device())
     pts = random.choice(PTS_CACHE).to(device=device) * 4
 
     coords = torch.floor(pts / 0.01).to(torch.int32)
     grid = fvdb.GridBatch.from_ijk(fvdb.JaggedTensor(coords), device=device)
 
     feature = torch.empty(grid.total_voxels, i_ch, dtype=torch.float32, device=device).random_()
+    feature_jt = fvdb.JaggedTensor([feature])
 
-    example_inputs = fvdbnn.VDBTensor(grid, fvdb.JaggedTensor([feature]))
-
+    # Create a convolution plan with the specified backend
+    plan = fvdb.ConvolutionPlan.from_grid_batch(
+        kernel_size=3, stride=1, source_grid=grid, channel_pairs=((i_ch, o_ch),), expert_config={"backend": backend}
+    )
     model = fvdbnn.SparseConv3d(in_channels=i_ch, out_channels=o_ch).to(device)
-    model.backend = backend
 
     model.eval()
 
     def run_model():
-        model(example_inputs)
+        return model(feature_jt, plan)
 
-    # benchmark(run_model)
     benchmark.pedantic(run_model, iterations=10, rounds=20)