Add feature to build rosetta with custom patches provided thru the web UI

.github/workflows/_build_rosetta.yaml

@@ -17,6 +17,11 @@ on:
         description: "Build date in YYYY-MM-DD format"
         required: false
         default: 'NOT SPECIFIED'
+      CUSTOM_PATCHLISTS:
+        type: string
+        description: 'Example: "plist1=p1,p2;plist2=," updates rosetta/plist1 to have lines "p1\np2" and rosetta/plist2 to be empty.'
+        required: false
+        default: ""
       PLATFORMS:
         type: string
         description: 'JSON list of platforms. Ex: ["amd64"]'
@@ -59,6 +64,21 @@ jobs:
 
       - name: Check out the repository under ${GITHUB_WORKSPACE}
         uses: actions/checkout@v3
+
+      - name: Apply any custom patches
+        shell: bash -x -e {0}
+        run: |
+          pushd rosetta
+          IFS=";"
+          CUSTOM_PATCHLISTS='${{ inputs.CUSTOM_PATCHLISTS }}'
+          for plist_rule in $CUSTOM_PATCHLISTS; do
+            path=$(cut -d= -f1 <<<"$plist_rule")
+            patches=$(cut -d= -f2 <<<"$plist_rule")
+            [[ -n "$(dirname $path)" ]] && mkdir -p $(dirname $path)
+            echo $patches | tr ',' '\n' | tee $path
+          done
+          unset IFS
+          popd
 
       - name: Login to GitHub Container Registry
         uses: docker/login-action@v2

.github/workflows/nightly-rosetta-pax-build.yaml

@@ -10,9 +10,19 @@ on:
     inputs:
       BASE_IMAGE:
         type: string
-        description: 'PAX image built by NVIDIA/JAX-Toolbox'
+        description: 'Upstream PAX image built by NVIDIA/JAX-Toolbox'
         default: 'ghcr.io/nvidia/upstream-pax:latest'
         required: true
+      PAXML_PATCHES:
+        type: string
+        description: "[Paxml] patches (Ex1: 'p1,p2' converted to 'p1\np2'). (Ex2: ',' == empty patchlist). Default: use upstream"
+        default: ''
+        required: false
+      PRAXIS_PATCHES:
+        type: string
+        description: "[Praxis] patches (Ex1: 'p1,p2' converted to 'p1\np2'). (Ex2: ',' == empty patchlist). Default: use upstream"
+        default: ''
+        required: false
       PUBLISH:
         type: boolean
         description: Publish dated images and update the 'latest' tag?
@@ -22,6 +32,9 @@ on:
 env:
   BASE_LIBRARY: pax
   DOCKER_REGISTRY: ghcr.io/nvidia
+  # Relative to rosetta/
+  CUSTOM_PAXML_PATCHLIST: patchlist-paxml.txt
+  CUSTOM_PRAXIS_PATCHLIST: patchlist-praxis.txt
 
 permissions:
   contents: read  # to fetch code
@@ -37,6 +50,7 @@ jobs:
       BUILD_DATE: ${{ steps.meta-vars.outputs.BUILD_DATE }}
       BASE_LIBRARY: ${{ steps.meta-vars.outputs.BASE_LIBRARY }}
       BASE_IMAGE: ${{ steps.meta-vars.outputs.BASE_IMAGE }}
+      CUSTOM_PATCHLISTS: ${{ steps.meta-vars.outputs.CUSTOM_PATCHLISTS }}
       PUBLISH: ${{ steps.meta-vars.outputs.PUBLISH }}
     steps:
       - name: Set build metadata
@@ -52,6 +66,12 @@ jobs:
           echo "BUILD_DATE=${BUILD_DATE}" >> $GITHUB_OUTPUT
           echo "BASE_LIBRARY=${{ env.BASE_LIBRARY }}" >> $GITHUB_OUTPUT
           echo "BASE_IMAGE=${BASE_IMAGE}" >> $GITHUB_OUTPUT
+          CUSTOM_PATCHLISTS=()
+          [[ -n "${{ inputs.PAXML_PATCHES }}" ]] && CUSTOM_PATCHLISTS+=(${{ env.CUSTOM_PAXML_PATCHLIST }}=${{ inputs.PAXML_PATCHES }})
+          [[ -n "${{ inputs.PRAXIS_PATCHES }}" ]] && CUSTOM_PATCHLISTS+=(${{ env.CUSTOM_PRAXIS_PATCHLIST }}=${{ inputs.PRAXIS_PATCHES }})
+          IFS=";"
+          echo "CUSTOM_PATCHLISTS=${CUSTOM_PATCHLISTS[*]}" >> $GITHUB_OUTPUT
+          unset IFS
           echo "PUBLISH=${{ inputs.PUBLISH }}" >> $GITHUB_OUTPUT
 
   build:
@@ -62,6 +82,7 @@ jobs:
       BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
       BASE_LIBRARY: ${{ needs.metadata.outputs.BASE_LIBRARY }}
       BASE_IMAGE: ${{ needs.metadata.outputs.BASE_IMAGE }}
+      CUSTOM_PATCHLISTS: ${{ needs.metadata.outputs.CUSTOM_PATCHLISTS }}
     secrets: inherit
 
   publish-build:

.github/workflows/nightly-rosetta-t5x-build-test.yaml

@@ -10,9 +10,19 @@ on:
     inputs:
       BASE_IMAGE:
         type: string
-        description: 'T5x image built by NVIDIA/JAX-Toolbox'
+        description: 'Upstream T5x image built by NVIDIA/JAX-Toolbox'
         default: 'ghcr.io/nvidia/upstream-t5x:latest'
         required: true
+      T5X_PATCHES:
+        type: string
+        description: "[T5x] patches (Ex1: 'p1,p2' converted to 'p1\np2'). (Ex2: ',' == empty patchlist). Default: use upstream"
+        default: ''
+        required: false
+      FLAX_PATCHES:
+        type: string
+        description: "[Flax] patches (Ex1: 'p1,p2' converted to 'p1\np2'). (Ex2: ',' == empty patchlist). Default: use upstream"
+        default: ''
+        required: false
       PUBLISH:
         type: boolean
         description: Publish dated images and update the 'latest' tag?
@@ -22,6 +32,9 @@ on:
 env:
   BASE_LIBRARY: t5x
   DOCKER_REGISTRY: ghcr.io/nvidia
+  # Relative to rosetta/
+  CUSTOM_T5X_PATCHLIST: patchlist-t5x.txt
+  CUSTOM_FLAX_PATCHLIST: patchlist-flax.txt
 
 permissions:
   contents: read  # to fetch code
@@ -37,6 +50,7 @@ jobs:
       BUILD_DATE: ${{ steps.meta-vars.outputs.BUILD_DATE }}
       BASE_LIBRARY: ${{ steps.meta-vars.outputs.BASE_LIBRARY }}
       BASE_IMAGE: ${{ steps.meta-vars.outputs.BASE_IMAGE }}
+      CUSTOM_PATCHLISTS: ${{ steps.meta-vars.outputs.CUSTOM_PATCHLISTS }}
       PUBLISH: ${{ steps.meta-vars.outputs.PUBLISH }}
     steps:
       - name: Set build metadata
@@ -52,6 +66,12 @@ jobs:
           echo "BUILD_DATE=${BUILD_DATE}" >> $GITHUB_OUTPUT
           echo "BASE_LIBRARY=${{ env.BASE_LIBRARY }}" >> $GITHUB_OUTPUT
           echo "BASE_IMAGE=${BASE_IMAGE}" >> $GITHUB_OUTPUT
+          CUSTOM_PATCHLISTS=()
+          [[ -n "${{ inputs.T5X_PATCHES }}" ]] && CUSTOM_PATCHLISTS+=(${{ env.CUSTOM_T5X_PATCHLIST }}=${{ inputs.T5X_PATCHES }})
+          [[ -n "${{ inputs.FLAX_PATCHES }}" ]] && CUSTOM_PATCHLISTS+=(${{ env.CUSTOM_FLAX_PATCHLIST }}=${{ inputs.FLAX_PATCHES }})
+          IFS=";"
+          echo "CUSTOM_PATCHLISTS=${CUSTOM_PATCHLISTS[*]}" >> $GITHUB_OUTPUT
+          unset IFS
           echo "PUBLISH=${{ inputs.PUBLISH }}" >> $GITHUB_OUTPUT
 
   build:
@@ -62,6 +82,7 @@ jobs:
       BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
       BASE_LIBRARY: ${{ needs.metadata.outputs.BASE_LIBRARY }}
       BASE_IMAGE: ${{ needs.metadata.outputs.BASE_IMAGE }}
+      CUSTOM_PATCHLISTS: ${{ needs.metadata.outputs.CUSTOM_PATCHLISTS }}
       # TODO: Can't build ARM until https://github.com/NVIDIA/JAX-Toolbox/pull/252 is available
       PLATFORMS: '["amd64"]'
     secrets: inherit