Skip to content

Add JetStream to MaxText container #1058

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 12 commits into from
Closed

Add JetStream to MaxText container #1058

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
7517eba
Add JetStream to MaxText container
DwarKapex Sep 23, 2024
333f231
Merge branch 'main' into vkozlov/jetstream-4-maxtext
DwarKapex Sep 25, 2024
c952cb3
Sandbox test
DwarKapex Sep 25, 2024
ddc4124
Sandbox fix
DwarKapex Sep 25, 2024
ab21dae
Revert sandbox
DwarKapex Sep 25, 2024
d9f5c18
Merge branch 'main' into vkozlov/jetstream-4-maxtext
yhtang Oct 10, 2024
9d59a81
Refactor MaxText docker and add upstream JetStream to the container
DwarKapex Jan 9, 2025
9b7426a
Remove second installation of JetStream
DwarKapex Jan 9, 2025
07c5a54
Refactor Dockerfile.maxtext
DwarKapex Jan 10, 2025
d726b1b
Merge main
DwarKapex Jan 10, 2025
43b2bd0
Add jetstream to manifest file, and updated main CI file
DwarKapex Jan 10, 2025
a4eb6d8
Addressed comments
DwarKapex Jan 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions .github/container/Dockerfile.maxtext.amd64
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -23,6 +26,17 @@ EOF

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
###############################################################################
Expand Down
11 changes: 11 additions & 0 deletions .github/container/Dockerfile.maxtext.arm64
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,17 @@ EOF

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
###############################################################################
Expand Down
32 changes: 32 additions & 0 deletions .github/workflows/_ci.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -588,3 +588,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 <<"EOF" |& tee test-jetstream.log
cd /opt/jetstream
pip install -r requirements.txt
export CUDA_VISIBLE_DEVICES=0
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we test multi-GPU functionality?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did not find JetStream unit-test examples on GitHub for multi-GPU. Even more: afaik, it does not support GPUs so far. But it will start GPU support in the nearest future.

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
EOF
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

Loading