Update rosetta/Dockerfile.maxtext

Co-authored-by: Terry Kong <[email protected]>

.github/workflows/_ci.yaml

@@ -86,7 +86,7 @@ jobs:
       BADGE_FILENAME: badge-maxtext-build
       BUILD_DATE: ${{ inputs.BUILD_DATE }}
       BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
-      CONTAINER_NAME: maxtext
+      CONTAINER_NAME: upstream-maxtext
       DOCKERFILE: .github/container/Dockerfile.maxtext.amd64
     secrets: inherit
 
@@ -184,8 +184,10 @@ jobs:
       - build-levanter
       - build-upstream-t5x
       - build-upstream-pax
+      - build-upstream-maxtext
       - build-rosetta-t5x
       - build-rosetta-pax
+      - build-rosetta-maxtext
       - build-grok
     outputs:
       TAGS: ${{ steps.collect-tags.outputs.TAGS }}

rosetta/Dockerfile.maxtext

@@ -1,5 +1,5 @@
 # syntax=docker/dockerfile:1-labs
-ARG BASE_IMAGE=ghcr.io/nvidia/jax-mealkit:upstream-pax
+ARG BASE_IMAGE=ghcr.io/nvidia/jax-mealkit:upstream-maxtext
 ARG [email protected]
 ARG GIT_USER_NAME=NVIDIA
 # If set to "true", then will pull new local patches, the manifest.yaml and create-distribution.sh (in case it was updated).
@@ -35,8 +35,6 @@ if [[ "${UPDATE_PATCHES}" == "true" ]]; then
   cp -r /mnt/jax-toolbox/.github/container/patches ${MANIFEST_DIR}/
   cp /mnt/jax-toolbox/.github/container/manifest.yaml ${MANIFEST_DIR}/manifest.yaml
   cp /mnt/jax-toolbox/.github/container/create-distribution.sh ${MANIFEST_DIR}/create-distribution.sh
-  # TODO: remove
-  cp /mnt/jax-toolbox/.github/container/pip-finalize.sh /usr/local/bin/
 fi
 cp -r /mnt/jax-toolbox/rosetta /opt/rosetta