diff --git a/.github/container/Dockerfile.maxtext b/.github/container/Dockerfile.maxtext index 60f0ee4c6..1da11247a 100644 --- a/.github/container/Dockerfile.maxtext +++ b/.github/container/Dockerfile.maxtext @@ -1,12 +1,8 @@ # syntax=docker/dockerfile:1-labs -# TODO: remove any JetStream references if not needed anymore - ARG BASE_IMAGE=ghcr.io/nvidia/jax-mealkit:jax ARG URLREF_MAXTEXT=https://github.com/google/maxtext.git#main -ARG URLREF_JETSTREAM=https://github.com/AI-Hypercomputer/JetStream.git#main ARG SRC_PATH_MAXTEXT=/opt/maxtext -ARG SRC_PATH_JETSTREAM=/opt/jetstream ############################################################################### ## Download source and add auxiliary scripts @@ -14,54 +10,30 @@ ARG SRC_PATH_JETSTREAM=/opt/jetstream FROM ${BASE_IMAGE} AS mealkit ARG URLREF_MAXTEXT -ARG URLREF_JETSTREAM ARG SRC_PATH_MAXTEXT -ARG SRC_PATH_JETSTREAM RUN <<"EOF" bash -ex -o pipefail git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} -git-clone.sh ${URLREF_JETSTREAM} ${SRC_PATH_JETSTREAM} -EOF - -# WAR: tunix package is not needed for Nvidia GPU support -RUN <<"EOF" bash -ex -o pipefail -sed -i '/tunix/d' ${SRC_PATH_MAXTEXT}/src/install_maxtext_extra_deps/extra_deps_from_github.txt 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/install_maxtext_extra_deps/extra_deps_from_github.txt" >> /opt/pip-tools.d/requirements-maxtext.in -#echo "-e file://${SRC_PATH_JETSTREAM}" >> /opt/pip-tools.d/requirements-maxtext.in EOF -# # remove GitHub direct-reference of JetStream in MaxText requirements -# RUN <<"EOF" bash -ex -o pipefail -# sed -i '/^google-jetstream/d' ${SRC_PATH_MAXTEXT}/requirements.txt -# EOF - # add version constraints to avoid eternal dependency resolution RUN <<"EOF" bash -ex -o pipefail for pattern in \ "s|tensorflow>=2.19.1|tensorflow==2.18.1|g" \ "s|tensorboard>=2.19.0|tensorboard>=2.18,<2.19|g" \ "s|tensorflow-text>=2.19.0|tensorflow-text==2.18.1|g" \ + "/tunix/d" \ ; do # tensorflow-cpu,tensorboard,tensorflow-text>=2.19.0 is incompatible with tensorflow==2.18.1 sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/pyproject.toml done EOF -# # add extra dependencies -# RUN <<"EOF" bash -ex -o pipefail -# echo >> ${SRC_PATH_MAXTEXT}/requirements.txt # add new line -# for requirement in \ -# "tensorflow-metadata>=1.15.0" \ -# "seqio@git+https://github.com/google/seqio.git" \ -# ; do -# echo "${requirement}" >> ${SRC_PATH_MAXTEXT}/requirements.txt -# done -# EOF - ############################################################################### ## Add test script to the path ############################################################################### diff --git a/.github/container/manifest.yaml b/.github/container/manifest.yaml index 9f858c47e..0103df6d1 100644 --- a/.github/container/manifest.yaml +++ b/.github/container/manifest.yaml @@ -70,11 +70,6 @@ seqio: tracking_ref: main latest_verified_commit: 11706e4a1e01a81ea6b3e02c5ad147028d5b94bb mode: pip-vcs -google-jetstream: - url: https://github.com/AI-Hypercomputer/JetStream.git - tracking_ref: main - latest_verified_commit: b8b9cb2ea4668da2c5012fc4c7ba958424d82ac9 - mode: git-clone maxtext: url: https://github.com/google/maxtext.git tracking_ref: main diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml index aafc32fdb..b239cebf3 100644 --- a/.github/workflows/_ci.yaml +++ b/.github/workflows/_ci.yaml @@ -147,7 +147,6 @@ jobs: bazel-remote-cache-url: ${{ vars.BAZEL_REMOTE_CACHE_URL }} EXTRA_BUILD_ARGS: | URLREF_MAXTEXT=${{ fromJson(inputs.SOURCE_URLREFS).MAXTEXT }} - URLREF_JETSTREAM=${{ fromJson(inputs.SOURCE_URLREFS).GOOGLE_JETSTREAM }} build-upstream-t5x: needs: build-jax diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 1a6f53ec4..1031f1e25 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -38,7 +38,7 @@ on: type: string description: | A comma-separated PACKAGE=URL#REF list to override sources used by build. - PACKAGE∊{JAX,XLA,Flax,transformer-engine,airio,axlearn,equinox,T5X,maxtext,google-jetstream} (case-insensitive) + PACKAGE∊{JAX,XLA,Flax,transformer-engine,airio,axlearn,equinox,T5X,maxtext} (case-insensitive) default: '' required: false MODE: