pip-finalize.sh

/opt/rosetta

RUN |5 NVIDIA_BUILD_ID=92034664 GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA UPDATE_PATCHES=false UPDATED_TE_REF= /bin/sh -c <<"EOF" bash -exu MANIFEST_DIR=$(dirname ${MANIFEST_FILE}) if [[ "${UPDATE_PATCHES}" != "true" && "${UPDATE_PATCHES}" != "false" ]]; then echo "UPDATE_PATCHES can only be true or false" exit 1 fi 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 if [[ -n "${UPDATED_TE_REF}" ]]; then TE_INSTALL_DIR=/opt/transformer-engine yq e ".transformer-engine.latest_verified_commit = \"${UPDATED_TE_REF}\"" -i $MANIFEST_FILE # Install from source instead of pre-built wheel sed -i -E 's@( file:///opt/transformer-engine)/dist/[^ ]*@\1@' /opt/pip-tools.d/requirements-te.in git -C $TE_INSTALL_DIR fetch -a if [[ "${UPDATED_TE_REF}" =~ ^pull/ ]]; then PR_ID=$(cut -d/ -f2 <<<"${UPDATED_TE_REF}") git -C $TE_INSTALL_DIR fetch origin ${UPDATED_TE_REF}:PR-${PR_ID} git -C $TE_INSTALL_DIR checkout PR-${PR_ID} else git -C $TE_INSTALL_DIR checkout ${UPDATED_TE_REF} fi fi # Setting the username/email is required to author commits from patches git config --global user.email "${GIT_USER_EMAIL}" git config --global user.name "${GIT_USER_NAME}" bash ${MANIFEST_DIR}/create-distribution.sh \
--manifest ${MANIFEST_FILE} \
--package paxml bash ${MANIFEST_DIR}/create-distribution.sh \
--manifest ${MANIFEST_FILE} \
--package praxis # Remove .gitconfig to avoid end-user authoring commits as the "build user" rm -f ~/.gitconfig EOF # buildkit

NVTE_FUSED_ATTN=1

ENABLE_TE=1

UPDATED_TE_REF

UPDATE_PATCHES

GIT_USER_NAME

GIT_USER_EMAIL

NVIDIA_BUILD_ID=92034664