From 399b4289bd736a016f14f3a139f993a5a94d8360 Mon Sep 17 00:00:00 2001
From: Olli Lupton <olupton@nvidia.com>
Date: Wed, 21 Aug 2024 11:02:56 +0000
Subject: [PATCH] install google.protobuf/protoc in the environment
 nsys-jax-combine runs in

---
 .github/workflows/nsys-jax.yaml | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/nsys-jax.yaml b/.github/workflows/nsys-jax.yaml
index 271f805cd..731bedea7 100644
--- 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