Allow specifying mii branch for nv-a6000 workflow (#4936)

deepspeedai · Jan 11, 2024 · 4738985 · 4738985
1 parent ee7db48
commit 4738985
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/.github/workflows/nv-a6000.yml b/.github/workflows/nv-a6000.yml
@@ -8,6 +8,12 @@ on:
       - 'tests/unit/inference/v2/**'
       - '.github/workflows/nv-a6000.yml'
   workflow_dispatch:
+    inputs:
+      mii_branch:
+        description: 'DeepSpeed-MII Branch'
+        required: false
+        default: 'main'
+        type: string
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
@@ -54,7 +60,9 @@ jobs:
           python -m pytest --color=yes --durations=0 --verbose -rF -m 'inference_v2_ops' unit/ --torch_ver="2.0" --cuda_ver="12"
       - name: MII unit tests
         run: |
-          git clone --depth=1 https://github.com/microsoft/DeepSpeed-MII.git
+          BRANCH="${{ github.event.inputs.mii_branch }}"
+          echo "Cloning DeepSpeed-MII branch: $BRANCH"
+          git clone -b $BRANCH --depth=1 https://github.com/microsoft/DeepSpeed-MII.git
           cd DeepSpeed-MII
           pip install .[dev]
           cd tests