install google.protobuf/protoc in the environment nsys-jax-combine ru…

…ns in
NVIDIA · Aug 21, 2024 · 399b428 · 399b428
1 parent 0161916
commit 399b428
Showing 1 changed file with 12 additions and 7 deletions.
diff --git a/.github/workflows/nsys-jax.yaml b/.github/workflows/nsys-jax.yaml
@@ -97,7 +97,7 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: '3.12'
-      - name: Add modern jax_nsys/ files to .zip inputs
+      - name: Add modern jax_nsys/ files to static .zip inputs
         run: |
           cd .github/container/jax_nsys
           for zip in ../../workflows/nsys-jax/test_data/pax_fsdp4_4proc_proc*.zip; do
@@ -108,20 +108,25 @@ jobs:
         shell: bash -x -e {0}
         run: |
           pip install virtualenv
-          .github/container/nsys-jax-combine \
+          virtualenv venv
+          ./venv/bin/pip install -r ./JAX-Toolbox/.github/container/requirements-nsys-jax.in
+          ./venv/bin/python ./JAX-Toolbox/.github/container/jax_nsys/install-protoc ./venv
+          ./venv/bin/python .github/container/nsys-jax-combine \
             --analysis summary \
             --analysis communication \
-            -o .github/workflows/nsys-jax/test_data/pax_fsdp4_4proc.zip \
+            -o pax_fsdp4_4proc.zip \
             .github/workflows/nsys-jax/test_data/pax_fsdp4_4proc_proc*.zip
-      - name: Mock up the structure of an extracted .zip file
-        run: unzip -d .github/container/jax_nsys/ .github/workflows/nsys-jax/test_data/pax_fsdp4_4proc.zip
+      - name: Extract the output .zip file
+        run: |
+          mkdir combined/
+          unzip -d combined/ pax_fsdp4_4proc.zip
       - name: Run the install script, but skip launching Jupyter Lab
         shell: bash -x -e {0}
-        run: NSYS_JAX_INSTALL_SKIP_LAUNCH=1 ./.github/container/jax_nsys/install.sh
+        run: NSYS_JAX_INSTALL_SKIP_LAUNCH=1 ./combined/install.sh
       - name: Test the Jupyter Lab installation and execute the notebook
         shell: bash -x -e {0}
         run: |
-          pushd .github/container/jax_nsys
+          pushd combined/
           ./nsys_jax_venv/bin/python -m jupyterlab --version
           # Run with ipython for the sake of getting a clear error message
           ./nsys_jax_venv/bin/ipython Analysis.ipynb