From 5c7351eee08cbe4888c7095671286b4d4e8e4833 Mon Sep 17 00:00:00 2001 From: Terry Kong Date: Tue, 10 Oct 2023 00:29:35 -0700 Subject: [PATCH] Updates rosetta CI to build multiarch (Arm64/amd64) for pax (#284) Example images from run: * pax: `ghcr.io/nvidia/jax-toolbox-internal:6414016731-pax-multiarch` T5x will not be be able to be built multiarch until https://github.com/NVIDIA/JAX-Toolbox/pull/252 is merged, so for now platforms is configurable and t5x is only built for x86 --------- Co-authored-by: Yu-Hang "Maxin" Tang --- .github/workflows/_build_rosetta.yaml | 98 +++++++++++++++---- .../nightly-rosetta-t5x-build-test.yaml | 2 + rosetta/Dockerfile.pax | 14 --- rosetta/Dockerfile.t5x | 1 - 4 files changed, 82 insertions(+), 33 deletions(-) diff --git a/.github/workflows/_build_rosetta.yaml b/.github/workflows/_build_rosetta.yaml index 91a630238..e811e31bf 100644 --- a/.github/workflows/_build_rosetta.yaml +++ b/.github/workflows/_build_rosetta.yaml @@ -17,10 +17,15 @@ on: description: "Build date in YYYY-MM-DD format" required: false default: 'NOT SPECIFIED' + PLATFORMS: + type: string + description: 'JSON list of platforms. Ex: ["amd64"]' + required: false + default: '["arm64", "amd64"]' outputs: DOCKER_TAGS: description: "Tags of the image built" - value: ${{ jobs.build.outputs.DOCKER_TAGS }} + value: ${{ jobs.merge.outputs.DOCKER_TAGS }} env: UPLD_IMAGE: ghcr.io/nvidia/jax-toolbox-internal @@ -34,16 +39,19 @@ permissions: jobs: build: - outputs: - DOCKER_TAGS: ${{ steps.meta.outputs.tags }} - runs-on: [self-hosted, x86, small] + strategy: + fail-fast: false + matrix: + PLATFORM: ${{ fromJSON(inputs.PLATFORMS) }} + runs-on: [self-hosted, "${{ matrix.PLATFORM }}", small] steps: - name: Set default BASE_IMAGE + id: defaults run: | if [[ -n "${{ inputs.BASE_IMAGE }}" ]]; then - echo "BASE_IMAGE=${{ inputs.BASE_IMAGE }}" >> "$GITHUB_ENV" + echo "BASE_IMAGE=${{ inputs.BASE_IMAGE }}" >> "$GITHUB_OUTPUT" else - echo "BASE_IMAGE=${{ env.DOCKER_REGISTRY }}/upstream-${{ inputs.BASE_LIBRARY }}:latest" >> "$GITHUB_ENV" + echo "BASE_IMAGE=${{ env.DOCKER_REGISTRY }}/upstream-${{ inputs.BASE_LIBRARY }}:latest" >> "$GITHUB_OUTPUT" fi - name: Print environment variables @@ -65,23 +73,14 @@ jobs: with: images: ${{ env.UPLD_IMAGE }} flavor: latest=false - tags: type=raw,value=${{ github.run_id }}-${{ inputs.BASE_LIBRARY }} + tags: type=raw,value=${{ github.run_id }}-${{ inputs.BASE_LIBRARY }}-${{ matrix.PLATFORM }} labels: org.opencontainers.image.created=${{ inputs.BUILD_DATE }} - name: Set up Docker Buildx uses: docker/setup-buildx-action@v2 with: driver-opts: | - image=moby/buildkit:v0.10.6 - - - name: Pull and validate BASE_IMAGE=${{ env.BASE_IMAGE }} contains BASE_LIBRARY=${{ inputs.BASE_LIBRARY }} - shell: bash -x -e {0} - run: | - PYTHON_PACKAGE=${{ inputs.BASE_LIBRARY }} - if [[ $PYTHON_PACKAGE == "pax" ]]; then - PYTHON_PACKAGE=paxml - fi - docker run --rm ${{ env.BASE_IMAGE }} python -c "import importlib,sys; sys.exit(importlib.util.find_spec('$PYTHON_PACKAGE') is None)" + image=moby/buildkit:v0.12.1 - name: Build docker images uses: docker/build-push-action@v4 @@ -89,8 +88,71 @@ jobs: context: rosetta/ push: true file: rosetta/Dockerfile.${{ inputs.BASE_LIBRARY }} + platforms: linux/${{ matrix.PLATFORM }} tags: ${{ steps.meta.outputs.tags }} labels: ${{ steps.meta.outputs.labels }} target: rosetta build-args: | - BASE_IMAGE=${{ env.BASE_IMAGE }} + BASE_IMAGE=${{ steps.defaults.outputs.BASE_IMAGE }} + + # Temporary workaround until the following issues are solved: + # https://github.com/orgs/community/discussions/17245 + # https://github.com/actions/runner/pull/2477 + # https://github.com/orgs/community/discussions/26639 + - name: Save image name as text file + shell: bash -x -e {0} + run: | + echo "${{ steps.meta.outputs.tags }}" >> image-name.txt + + - name: Upload image name file as artifact + uses: actions/upload-artifact@v3 + with: + name: image-name-${{ inputs.BASE_LIBRARY }}-${{ matrix.PLATFORM }} + path: image-name.txt + + merge: + runs-on: ubuntu-latest + needs: build + outputs: + DOCKER_TAGS: ${{ steps.meta.outputs.tags }} + steps: + # TODO: currently downloading all artifacts of the entire workflow + # Revise when this request is fulfilled: + # https://github.com/actions/download-artifact/issues/214 + - name: Download image name files into separate folders + uses: actions/download-artifact@v3 + + - name: Login to GitHub Container Registry + uses: docker/login-action@v2 + with: + registry: ghcr.io + username: ${{ github.repository_owner }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Set docker metadata + id: meta + uses: docker/metadata-action@v4 + with: + images: | + ${{ env.UPLD_IMAGE }} + flavor: | + latest=false + tags: | + type=raw,value=${{ github.run_id }}-${{ inputs.BASE_LIBRARY }}-multiarch + labels: + org.opencontainers.image.created=${{ inputs.BUILD_DATE }} + + - name: Combine images into a single multi-arch image + shell: bash -x -e {0} + run: | + docker manifest create ${{ steps.meta.outputs.tags }} $( + for IMAGE in $(cat image-name-${{ inputs.BASE_LIBRARY }}-*/image-name.txt); do + REPO=$(echo $IMAGE | cut -d: -f1) + DIGEST=$( + docker manifest inspect $IMAGE |\ + jq -r '.manifests[] | select(.platform.os == "linux") | .digest' + ) + echo $REPO@${DIGEST} + done + ) + docker manifest push ${{ steps.meta.outputs.tags }} \ No newline at end of file diff --git a/.github/workflows/nightly-rosetta-t5x-build-test.yaml b/.github/workflows/nightly-rosetta-t5x-build-test.yaml index ed0d0352e..efc03fb5c 100644 --- a/.github/workflows/nightly-rosetta-t5x-build-test.yaml +++ b/.github/workflows/nightly-rosetta-t5x-build-test.yaml @@ -61,6 +61,8 @@ jobs: BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }} BASE_LIBRARY: ${{ needs.metadata.outputs.BASE_LIBRARY }} BASE_IMAGE: ${{ needs.metadata.outputs.BASE_IMAGE }} + # TODO: Can't build ARM until https://github.com/NVIDIA/JAX-Toolbox/pull/252 is available + PLATFORMS: '["amd64"]' secrets: inherit publish-build: diff --git a/rosetta/Dockerfile.pax b/rosetta/Dockerfile.pax index 6ce54773b..12d6c77f3 100644 --- a/rosetta/Dockerfile.pax +++ b/rosetta/Dockerfile.pax @@ -37,18 +37,4 @@ bash create-distribution.sh \ -e /opt/praxis-mirror rm -rf $(find /opt -name "__pycache__") EOF - -WORKDIR /opt/rosetta -RUN <