move nsys-jax CI under main pipeline so it can inherit the publish se…

…tting

.github/workflows/nsys-jax.yaml → .github/workflows/_nsys_jax.yaml

@@ -1,18 +1,17 @@
-name: nsys-jax pure-Python CI
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
-  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+name: ~nsys-jax pure-Python CI
 
 on:
-  pull_request:
-    types:
-      - opened
-      - reopened
-      - ready_for_review
-      - synchronize
-    paths-ignore:
-      - '**.md'
+  workflow_call:
+    inputs:
+      RENDERED_NOTEBOOK_GIST_ID:
+        type: string
+        description: GitHub Gist ID under nvjax user where outputs will be copied if PUBLISH_RENDERED_NOTEBOOK is true
+        required: true
+      PUBLISH_RENDERED_NOTEBOOK:
+        type: boolean
+        description: Whether to copy newly rendered notebooks to RENDERED_NOTEBOOK_GIST_ID
+        default: false
+        required: true
 
 env:
   NSYS_JAX_PYTHON_FILES: JAX-Toolbox/.github/container/nsys-jax JAX-Toolbox/.github/container/jax_nsys
@@ -72,6 +71,9 @@ jobs:
           ./venv/bin/mypy ${NSYS_JAX_PYTHON_FILES}
 
   notebook:
+    env:
+      # Name/bash regex for shields.io endpoint JSON files
+      PUBLISH_NOTEBOOK_FILES: '(.*\.ipynb)'
     runs-on: ubuntu-24.04
     steps:
       - name: Check out the repository under ${GITHUB_WORKSPACE}
@@ -105,7 +107,7 @@ jobs:
           workdir=$(mktemp -d)
           ./nsys_jax_venv/bin/jupyter nbconvert --execute --to notebook --output-dir=$workdir Analysis.ipynb
           echo "WORKDIR=$workdir" >> $GITHUB_OUTPUT
-      - name: Upload rendered notebook to gist
+      - name: Upload rendered notebook to Gist
         id: upload
         uses: actions/github-script@v7
         with:
@@ -137,6 +139,49 @@ jobs:
             });
             console.log(gist)
             return gist.data.id;
+      - name: Copy rendered notebook to Gist with well-known ID
+        uses: actions/github-script@v7
+        with:
+          github-token: ${{ secrets.NVJAX_GIST_TOKEN }}
+          script: |
+            const srcId = "${{ steps.upload.outputs.result }}";
+            const dstId = "${{ inputs.PUBLISH_RENDERED_NOTEBOOK && vars.RENDERED_NOTEBOOK_GIST_ID || vars.MOCK_RENDERED_NOTEBOOK_GIST_ID }}";
+            const { PUBLISH_NOTEBOOK_FILES } = process.env;
+            const 
+
+            // Fetch existing files from destination gist
+            const { data: dstData } = await github.rest.gists.get({
+              gist_id: dstId
+            });
+
+            // Mark existing files in destination gist for deletion
+            let filesToUpdate = {};
+            for (const filename of Object.keys(dstData.files)) {
+              filesToUpdate[filename] = null;
+            }
+
+            // Fetch files from source gist
+            const { data: srcData } = await github.rest.gists.get({
+              gist_id: srcId
+            });
+
+            // Add or update files based on the pattern
+            const pattern = new RegExp(`${PUBLISH_NOTEBOOK_FILES}`);
+            for (const [filename, fileObj] of Object.entries(srcData.files)) {
+              if (filename.match(pattern)) {
+                filesToUpdate[filename] = {
+                  content: fileObj.content
+                };
+              }
+            }
+
+            // Update files in destination gist
+            await github.rest.gists.update({
+              gist_id: dstId,
+              files: filesToUpdate
+            });
+            console.log("Files copied successfully.");
+            console.log(Object.keys(filesToUpdate));
 
   ruff:
     runs-on: ubuntu-24.04

.github/workflows/ci.yaml

@@ -113,6 +113,16 @@ jobs:
             echo "Error: If BUMP_MANIFEST=false, MERGE_BUMPED_MANIFEST cannot be true" >&2
             exit 1
           fi
+
+  nsys-jax:
+    needs: metadata
+    uses: ./.github/workflows/_nsys_jax.yaml
+    with:
+      PUBLISH_RENDERED_NOTEBOOK: ${{ needs.metadata.outputs.PUBLISH == 'true' }}
+      # TODO: these could/should be saved in the repository settings instead
+      RENDERED_NOTEBOOK_GIST_ID: e2cd3520201caab6b67385ed36fad3c1
+      MOCK_RENDERED_NOTEBOOK_GIST_ID: 16698d9e9e52320243165d61b5bb3975
+    secrets: inherit
 
   bump-manifest:
     needs: metadata