From b30e6df7214ff9fe553f3ac4b14359fd4967e42e Mon Sep 17 00:00:00 2001
From: andy-neuma <andy@neuralmagic.com>
Date: Thu, 2 May 2024 13:40:16 -0400
Subject: [PATCH] some minor patches

---
 .github/actions/nm-install-test-whl/action.yml    |  6 ++----
 .github/actions/nm-install-whl/action.yml         |  2 +-
 .github/workflows/build-test.yml                  |  1 +
 .github/workflows/nm-benchmark.yml                | 12 ++++++++++--
 .github/workflows/test.yml                        |  1 +
 neuralmagic/benchmarks/requirements-benchmark.txt |  1 -
 6 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/.github/actions/nm-install-test-whl/action.yml b/.github/actions/nm-install-test-whl/action.yml
index 7a34c1e31a8a1..193dad8f99820 100644
--- a/.github/actions/nm-install-test-whl/action.yml
+++ b/.github/actions/nm-install-test-whl/action.yml
@@ -44,14 +44,12 @@ runs:
         pip3 install coverage
         pip3 install pytest-cov
         pip3 install pytest-xdist
-        pip3 install --index-url http://${{ inputs.pypi }}:8080/ --trusted-host ${{ inputs.pypi }} nm-magic-wand-nightly
-        pip3 list
+        pip3 install -r requirements-dev.txt
         BASE=$(./.github/scripts/convert-version ${{ inputs.python }})
         WHL=$(find . -type f -iname "*${BASE}*.whl")
         WHL_BASENAME=$(basename ${WHL})
         echo "whl=${WHL_BASENAME}" >> "$GITHUB_OUTPUT"
-        pip3 install ${WHL}
-        pip3 install -r requirements-dev.txt
+        pip3 install ${WHL}[sparse]
         # report magic_wand version
         MAGIC_WAND=$(pip3 show nm-magic-wand-nightly | grep "Version" | cut -d' ' -f2)
         echo "magic_wand=${MAGIC_WAND}" >> "$GITHUB_OUTPUT"
diff --git a/.github/actions/nm-install-whl/action.yml b/.github/actions/nm-install-whl/action.yml
index 6d67522f09e7e..d67183a7239e8 100644
--- a/.github/actions/nm-install-whl/action.yml
+++ b/.github/actions/nm-install-whl/action.yml
@@ -19,7 +19,7 @@ runs:
         COMMIT=${{ github.sha }}
         VENV="${{ env.VENV_BASE }}-${COMMIT:0:7}"
         source $(pyenv root)/versions/${{ inputs.python }}/envs/${VENV}/bin/activate
-        pip3 install -r requirements-cuda.txt -r requirements-dev.txt
+        pip3 install -r requirements-dev.txt
         BASE=$(./.github/scripts/convert-version ${{ inputs.python }})
         WHL=$(find . -type f -iname "*${BASE}*.whl")
         WHL_BASENAME=$(basename ${WHL})
diff --git a/.github/workflows/build-test.yml b/.github/workflows/build-test.yml
index 2aa49091a90da..0b3ce56982081 100644
--- a/.github/workflows/build-test.yml
+++ b/.github/workflows/build-test.yml
@@ -172,6 +172,7 @@ jobs:
             timeout: ${{ inputs.benchmark_timeout }}
             gitref: ${{ github.ref }}
             python: ${{ inputs.python }}
+            whl: ${{ needs.BUILD.outputs.whl }}
             # Always push if it is a scheduled job
             push_benchmark_results_to_gh_pages: "${{ github.event_name == 'schedule' || inputs.push_benchmark_results_to_gh_pages }}"
         secrets: inherit
diff --git a/.github/workflows/nm-benchmark.yml b/.github/workflows/nm-benchmark.yml
index 91ca100646423..4733775621432 100644
--- a/.github/workflows/nm-benchmark.yml
+++ b/.github/workflows/nm-benchmark.yml
@@ -23,6 +23,10 @@ on:
         description: "python version, e.g. 3.10.12"
         type: string
         required: true
+      whl:
+        description: "whl to test (variable appears late binding so unusable outside 'download artifact')"
+        type: string
+        required: true
       push_benchmark_results_to_gh_pages:
         description: "When set to true, the workflow pushes all benchmarking results to gh-pages UI"
         type: string
@@ -51,6 +55,10 @@ on:
         description: "python version, e.g. 3.10.12"
         type: string
         required: true
+      whl:
+        description: "whl to test (variable appears late binding so unusable outside 'download artifact')"
+        type: string
+        required: true
       push_benchmark_results_to_gh_pages:
         description: "When set to true, the workflow pushes all benchmarking results to gh-pages UI"
         type: choice
@@ -85,8 +93,8 @@ jobs:
         uses: ./.github/actions/nm-set-env/
         with:
           hf_token: ${{ secrets.NM_HF_TOKEN }}
-          Gi_per_thread: ${{ inputs.Gi_per_thread }}
-          nvcc_threads: ${{ inputs.nvcc_threads }}
+          Gi_per_thread: 1
+          nvcc_threads: 0
 
       - name: set python
         id: set_python
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index e8e87bb4283b2..b081a63b7e9e1 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -82,6 +82,7 @@ jobs:
               with:
                 hf_token: ${{ secrets.NM_HF_TOKEN }}
                 Gi_per_thread: 1
+                nvcc_threads: 0
 
             - name: set python
               id: set_python
diff --git a/neuralmagic/benchmarks/requirements-benchmark.txt b/neuralmagic/benchmarks/requirements-benchmark.txt
index df1c80adcfc17..095bba70f1946 100644
--- a/neuralmagic/benchmarks/requirements-benchmark.txt
+++ b/neuralmagic/benchmarks/requirements-benchmark.txt
@@ -2,4 +2,3 @@
 requests
 aiohttp
 datasets
-nm-magic-wand-nightly