diff --git a/.github/workflows/_build_rosetta.yaml b/.github/workflows/_build_rosetta.yaml index e811e31bf..16a6bd28a 100644 --- a/.github/workflows/_build_rosetta.yaml +++ b/.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 diff --git a/.github/workflows/nightly-rosetta-pax-build.yaml b/.github/workflows/nightly-rosetta-pax-build.yaml index 537d8a78f..87c8c899e 100644 --- a/.github/workflows/nightly-rosetta-pax-build.yaml +++ b/.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: diff --git a/.github/workflows/nightly-rosetta-t5x-build-test.yaml b/.github/workflows/nightly-rosetta-t5x-build-test.yaml index 09d39867c..eb950cb65 100644 --- a/.github/workflows/nightly-rosetta-t5x-build-test.yaml +++ b/.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