Remove tunix from pyproject.toml and JetStream

.github/container/Dockerfile.maxtext

@@ -1,67 +1,39 @@
 # 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
 ###############################################################################
 
 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
 ###############################################################################

.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

.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

.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: