Dockerfile.maxtext

# syntax=docker/dockerfile:1-labs

ARG BASE_IMAGE=ghcr.io/nvidia/jax-mealkit:jax
ARG URLREF_MAXTEXT=https://github.com/google/maxtext.git#main
ARG SRC_PATH_MAXTEXT=/opt/maxtext

###############################################################################
## Download source and add auxiliary scripts
###############################################################################

FROM ${BASE_IMAGE} AS mealkit
ARG URLREF_MAXTEXT
ARG SRC_PATH_MAXTEXT

RUN <<"EOF" bash -ex -o pipefail
git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT}
EOF

RUN <<"EOF" bash -ex -o pipefail
echo "-e file://${SRC_PATH_MAXTEXT}" >> /opt/pip-tools.d/requirements-maxtext.in
echo "-r ${SRC_PATH_MAXTEXT}/src/dependencies/requirements/base_requirements/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in
EOF

RUN <<"EOF" bash -exuo pipefail
# Avoid old tensorboard constraint from google-cloud-aiplatform[tensorboard].
sed -i "/^cloud-accelerator-diagnostics/d" \
  ${SRC_PATH_MAXTEXT}/src/dependencies/requirements/base_requirements/requirements.txt

# remove custom build hook that conflicts with our dependencies introduced in
# https://github.com/AI-Hypercomputer/maxtext/commit/de886f9a918d17cfb025a353e7da2bc4d64c26b5
sed -i -e '/^\[tool\.hatch\.build\.targets\.wheel\.hooks\.custom\]$/d' \
       -e '/^path = "build_hooks\.py"$/d' \
       ${SRC_PATH_MAXTEXT}/pyproject.toml
EOF

###############################################################################
## Add test script to the path
###############################################################################

ADD test-maxtext.sh /usr/local/bin

###############################################################################
## Install accumulated packages from the base image and the previous stage
###############################################################################

FROM mealkit AS final

RUN pip-finalize.sh

WORKDIR ${SRC_PATH_MAXTEXT}