Add JetStream to MaxText container

NVIDIA · Sep 23, 2024 · 7517eba · 7517eba
1 parent 056a3b0
commit 7517eba
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 0 deletions.
diff --git a/.github/container/Dockerfile.maxtext.amd64 b/.github/container/Dockerfile.maxtext.amd64
@@ -3,6 +3,9 @@
 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
+ARG URLREF_JETSTREAM=https://github.com/google/jetstream.git#main
+ARG SRC_PATH_JETSTREAM=/opt/jetstream
+
 
 ###############################################################################
 ## Download source and add auxiliary scripts
@@ -30,6 +33,17 @@ RUN cd "${SRC_PATH_MAXTEXT}" && patch -p1 < /opt/maxtext-mha.patch && git diff
 
 ADD test-maxtext.sh /usr/local/bin
 
+###############################################################################
+## Add JetStream
+###############################################################################
+ARG URLREF_JETSTREAM
+ARG SRC_PATH_JETSTREAM
+
+RUN <<"EOF" bash -ex
+git-clone.sh ${URLREF_JETSTREAM} ${SRC_PATH_JETSTREAM}
+echo "-r ${SRC_PATH_JETSTREAM}/requirements.txt" >> /opt/pip-tools.d/requirements-jetstream.in
+EOF
+
 ###############################################################################
 ## Install accumulated packages from the base image and the previous stage
 ###############################################################################

diff --git a/.github/container/Dockerfile.maxtext.arm64 b/.github/container/Dockerfile.maxtext.arm64
@@ -71,6 +71,17 @@ RUN cd "${SRC_PATH_MAXTEXT}" && patch -p1 < /opt/maxtext-mha.patch && git diff
 
 ADD test-maxtext.sh /usr/local/bin
 
+###############################################################################
+## Add JetStream
+###############################################################################
+ARG URLREF_JETSTREAM
+ARG SRC_PATH_JETSTREAM
+
+RUN <<"EOF" bash -ex
+git-clone.sh ${URLREF_JETSTREAM} ${SRC_PATH_JETSTREAM}
+echo "-r ${SRC_PATH_JETSTREAM}/requirements.txt" >> /opt/pip-tools.d/requirements-jetstream.in
+EOF
+
 ###############################################################################
 ## Install accumulated packages from the base image and the previous stage
 ###############################################################################

diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml
@@ -581,3 +581,35 @@ jobs:
     with:
       MAXTEXT_IMAGE: ${{ needs.build-maxtext.outputs.DOCKER_TAG_FINAL }}
     secrets: inherit
+
+  test-jetstream:
+    needs: build-maxtext
+    if: inputs.ARCHITECTURE == 'amd64' # no arm64 gpu runners
+    uses: ./.github/workflows/_test_unit.yaml
+    with:
+      TEST_NAME: jetstream
+      EXECUTE: |
+        docker run --shm-size=1g --gpus all ${{ needs.build-maxtext.outputs.DOCKER_TAG_FINAL }} \
+        bash -ec \
+        "cd /opt/jetstream && " \
+        "pip install -r requirements.txt && " \
+        "export CUDA_VISIBLE_DEVICES=0 && " \
+        "python -m unittest -v jetstream.tests.core.test_orchestrator && " \
+        "python -m jetstream.engine.mock_engine_test && " \
+        "python -m jetstream.core.orchestrator_test && " \
+        "python -m jetstream.core.server_test" \
+        | tee test-jetstream.log
+      STATISTICS_SCRIPT: |
+        summary_line=$(tail -n1 test-jetstream.log)
+        errors=$(echo $summary_line | grep -oE '[0-9]+ error' | awk '{print $1} END { if (!NR) print 0}')
+        failed_tests=$(echo $summary_line | grep -oE '[0-9]+ failed' | awk '{print $1} END { if (!NR) print 0}')
+        passed_tests=$(echo $summary_line | grep -oE '[0-9]+ passed' | awk '{print $1} END { if (!NR) print 0}')
+        total_tests=$((failed_tests + passed_tests))
+        echo "TOTAL_TESTS=${total_tests}" >> $GITHUB_OUTPUT
+        echo "ERRORS=${errors}" >> $GITHUB_OUTPUT
+        echo "PASSED_TESTS=${passed_tests}" >> $GITHUB_OUTPUT
+        echo "FAILED_TESTS=${failed_tests}" >> $GITHUB_OUTPUT
+      ARTIFACTS: |
+        test-jetstream.log
+    secrets: inherit
+