Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/maxtext

NVIDIA_BUILD_ID=51421314 URLREF_MAXTEXT=https://github.com/google/maxtext.git#083293fc373bb26dece7d81901a13f6e9c3e584f SRC_PATH_MAXTEXT=/opt/maxtext pip-finalize.sh

pinned/26.05-devel/maxtext-amd64.txt /opt/pip-tools.d/requirements-pinned.txt

test-maxtext.sh /usr/local/bin

RUN |3 NVIDIA_BUILD_ID=51421314 URLREF_MAXTEXT=https://github.com/google/maxtext.git#083293fc373bb26dece7d81901a13f6e9c3e584f SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"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 # buildkit

RUN |3 NVIDIA_BUILD_ID=51421314 URLREF_MAXTEXT=https://github.com/google/maxtext.git#083293fc373bb26dece7d81901a13f6e9c3e584f SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"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 # MaxText's base_requirements/requirements.txt (above) does not include google-jetstream,
# but inference code paths (maxtext.inference.maxengine, maxtext.input_pipeline.packing.prefill_packing)
# import the `jetstream` package at module load time and DECOUPLE_GCLOUD=TRUE does not
# unblock prefill_packing. Now re-add this dependency
grep -E '^google-jetstream' ${SRC_PATH_MAXTEXT}/src/dependencies/requirements/requirements.txt \
  >> /opt/pip-tools.d/requirements-maxtext.in
EOF # buildkit

RUN |3 NVIDIA_BUILD_ID=51421314 URLREF_MAXTEXT=https://github.com/google/maxtext.git#083293fc373bb26dece7d81901a13f6e9c3e584f SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} EOF # buildkit

SRC_PATH_MAXTEXT=/opt/maxtext

URLREF_MAXTEXT=https://github.com/google/maxtext.git#083293fc373bb26dece7d81901a13f6e9c3e584f

NVIDIA_BUILD_ID=51421314