From 7517eba5f349a7a214d492996cf2230f7fd0f630 Mon Sep 17 00:00:00 2001 From: Vladislav Date: Mon, 23 Sep 2024 16:31:24 -0700 Subject: [PATCH 1/8] Add JetStream to MaxText container --- .github/container/Dockerfile.maxtext.amd64 | 14 ++++++++++ .github/container/Dockerfile.maxtext.arm64 | 11 ++++++++ .github/workflows/_ci.yaml | 32 ++++++++++++++++++++++ 3 files changed, 57 insertions(+) diff --git a/.github/container/Dockerfile.maxtext.amd64 b/.github/container/Dockerfile.maxtext.amd64 index 63c6767c0..c3337b9e6 100644 --- 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 index a971d2405..ad0f43b69 100644 --- 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 index fc04b83ab..6974b8c4c 100644 --- 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 + From c952cb320c6e8b319245512fd1f4a6702199ab00 Mon Sep 17 00:00:00 2001 From: Vladislav Date: Wed, 25 Sep 2024 14:51:38 -0700 Subject: [PATCH 2/8] Sandbox test --- .github/workflows/_ci.yaml | 18 ++++----- .github/workflows/_sandbox.yaml | 65 +++++++++++++++------------------ 2 files changed, 39 insertions(+), 44 deletions(-) diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml index fd2525649..0492aa388 100644 --- a/.github/workflows/_ci.yaml +++ b/.github/workflows/_ci.yaml @@ -596,15 +596,15 @@ jobs: 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 + bash <<"EOF" |& tee test-jetstream.log + 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 + 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}') diff --git a/.github/workflows/_sandbox.yaml b/.github/workflows/_sandbox.yaml index 7b90b72ca..e120a39de 100644 --- a/.github/workflows/_sandbox.yaml +++ b/.github/workflows/_sandbox.yaml @@ -4,38 +4,33 @@ on: workflow_dispatch: jobs: - sandbox: - runs-on: ubuntu-22.04 - steps: - - name: Login to GitHub Container Registry - uses: docker/login-action@v3 - with: - registry: ghcr.io - username: ${{ github.repository_owner }} - password: ${{ secrets.GITHUB_TOKEN }} - - - name: Print usage - run: | - cat << EOF - This is an empty workflow file located in the main branch of your - repository. It serves as a testing ground for new GitHub Actions on - development branches before merging them to the main branch. By - defining and overloading this workflow on your development branch, - you can test new actions without affecting your main branch, ensuring - a smooth integration process once the changes are ready to be merged. - - Usage: - - 1. In your development branch, modify the sandbox.yml workflow file - to include the new actions you want to test. Make sure to commit - the changes to the development branch. - 2. Navigate to the 'Actions' tab in your repository, select the - '~Sandbox' workflow, and choose your development branch from the - branch dropdown menu. Click on 'Run workflow' to trigger the - workflow on your development branch. - 3. Once you have tested and verified the new actions in the Sandbox - workflow, you can incorporate them into your main workflow(s) and - merge the development branch into the main branch. Remember to - revert the changes to the sandbox.yml file in the main branch to - keep it empty for future testing. - EOF + 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 ghcr.io/nvidia/jax-toolbox-internal:11038994821-maxtext-amd64 \ + bash <<"EOF" |& tee test-jetstream.log + 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 + 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 \ No newline at end of file From ddc412475180e4eb43a37782cafffed0ae245960 Mon Sep 17 00:00:00 2001 From: Vladislav Date: Wed, 25 Sep 2024 16:31:05 -0700 Subject: [PATCH 3/8] Sandbox fix --- .github/workflows/_sandbox.yaml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/_sandbox.yaml b/.github/workflows/_sandbox.yaml index e120a39de..6d86a44e4 100644 --- a/.github/workflows/_sandbox.yaml +++ b/.github/workflows/_sandbox.yaml @@ -5,8 +5,6 @@ on: jobs: test-jetstream: - needs: build-maxtext - if: inputs.ARCHITECTURE == 'amd64' # no arm64 gpu runners uses: ./.github/workflows/_test_unit.yaml with: TEST_NAME: jetstream From ab21dae070f41c4976691d9960b53c77d0a65701 Mon Sep 17 00:00:00 2001 From: Vladislav Date: Wed, 25 Sep 2024 16:47:18 -0700 Subject: [PATCH 4/8] Revert sandbox --- .github/workflows/_sandbox.yaml | 63 ++++++++++++++++++--------------- 1 file changed, 35 insertions(+), 28 deletions(-) diff --git a/.github/workflows/_sandbox.yaml b/.github/workflows/_sandbox.yaml index 6d86a44e4..7b90b72ca 100644 --- a/.github/workflows/_sandbox.yaml +++ b/.github/workflows/_sandbox.yaml @@ -4,31 +4,38 @@ on: workflow_dispatch: jobs: - test-jetstream: - uses: ./.github/workflows/_test_unit.yaml - with: - TEST_NAME: jetstream - EXECUTE: | - docker run --shm-size=1g --gpus all ghcr.io/nvidia/jax-toolbox-internal:11038994821-maxtext-amd64 \ - bash <<"EOF" |& tee test-jetstream.log - 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 - 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 \ No newline at end of file + sandbox: + runs-on: ubuntu-22.04 + steps: + - name: Login to GitHub Container Registry + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.repository_owner }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Print usage + run: | + cat << EOF + This is an empty workflow file located in the main branch of your + repository. It serves as a testing ground for new GitHub Actions on + development branches before merging them to the main branch. By + defining and overloading this workflow on your development branch, + you can test new actions without affecting your main branch, ensuring + a smooth integration process once the changes are ready to be merged. + + Usage: + + 1. In your development branch, modify the sandbox.yml workflow file + to include the new actions you want to test. Make sure to commit + the changes to the development branch. + 2. Navigate to the 'Actions' tab in your repository, select the + '~Sandbox' workflow, and choose your development branch from the + branch dropdown menu. Click on 'Run workflow' to trigger the + workflow on your development branch. + 3. Once you have tested and verified the new actions in the Sandbox + workflow, you can incorporate them into your main workflow(s) and + merge the development branch into the main branch. Remember to + revert the changes to the sandbox.yml file in the main branch to + keep it empty for future testing. + EOF From 9b7426adb610ba32b5970ce853add921b41f0f08 Mon Sep 17 00:00:00 2001 From: Vladislav Kozlov Date: Thu, 9 Jan 2025 13:44:37 -0800 Subject: [PATCH 5/8] Remove second installation of JetStream --- .github/container/Dockerfile.maxtext | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/.github/container/Dockerfile.maxtext b/.github/container/Dockerfile.maxtext index 19a0dd0f2..5574741c6 100644 --- a/.github/container/Dockerfile.maxtext +++ b/.github/container/Dockerfile.maxtext @@ -75,10 +75,10 @@ for pattern in \ sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/requirements.txt; done echo "tensorflow-metadata>=1.15.0" >> ${SRC_PATH_MAXTEXT}/requirements.txt -# remove outdated version of JetStream in favor of upstream JetStream +# remove outdated version of JetStream in favor of the upstream JetStream sed -i "/google-jetstream/d" ${SRC_PATH_MAXTEXT}/requirements.txt -# Add upstream JetStream +# Add the upstream JetStream git-clone.sh ${URLREF_JETSTREAM} ${SRC_PATH_JETSTREAM} echo "-r ${SRC_PATH_JETSTREAM}/requirements.txt" >> /opt/pip-tools.d/requirements-jetstream.in @@ -90,16 +90,6 @@ EOF ADD test-maxtext.sh /usr/local/bin -############################################################################### -## Add 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 ############################################################################### From 07c5a5498a728e66dd7639c22e232ab1096debdf Mon Sep 17 00:00:00 2001 From: Vladislav Kozlov Date: Fri, 10 Jan 2025 08:50:23 -0800 Subject: [PATCH 6/8] Refactor Dockerfile.maxtext --- .github/container/Dockerfile.maxtext | 49 ---------------------------- 1 file changed, 49 deletions(-) diff --git a/.github/container/Dockerfile.maxtext b/.github/container/Dockerfile.maxtext index 5574741c6..d58593606 100644 --- a/.github/container/Dockerfile.maxtext +++ b/.github/container/Dockerfile.maxtext @@ -3,46 +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_TFTEXT=https://github.com/tensorflow/text.git#master -ARG SRC_PATH_TFTEXT=/opt/tensorflow-text ARG URLREF_JETSTREAM=https://github.com/google/jetstream.git#main ARG SRC_PATH_JETSTREAM=/opt/jetstream - -############################################################################### -## build tensorflow-text and lingvo, which do not have working arm64 pip wheels -############################################################################### - -ARG BASE_IMAGE -FROM ${BASE_IMAGE} as wheel-builder - -#------------------------------------------------------------------------------ -# build tensorflow-text from source -#------------------------------------------------------------------------------ - -# Remove TFTEXT build from source when it has py-3.12 wheels for x86/arm64 -FROM wheel-builder as tftext-builder -ARG URLREF_TFTEXT -ARG SRC_PATH_TFTEXT - -RUN pip install tensorflow_datasets==4.9.2 auditwheel tensorflow==2.18.0 -RUN git-clone.sh ${URLREF_TFTEXT} ${SRC_PATH_TFTEXT} -RUN <<"EOF" bash -exu -o pipefail -cd ${SRC_PATH_TFTEXT} - -# The tftext build script queries GitHub, but these requests are sometimes -# throttled by GH, resulting in a corrupted uri for tensorflow in WORKSPACE. -# A workaround (needs to be updated when the tensorflow version changes): -sed -i "s/# Update TF dependency to installed tensorflow./commit_slug=6550e4bd80223cdb8be6c3afd1f81e86a4d433c3/" oss_scripts/prepare_tf_dep.sh - -# Newer versions of LLVM make lld's --undefined-version check of lld is strict -# by default (https://reviews.llvm.org/D135402), but the tftext build seems to -# rely on this behavior. -echo "write_to_bazelrc \"build --linkopt='-Wl,--undefined-version'\"" >> oss_scripts/configure.sh - -./oss_scripts/run_build.sh -EOF - ############################################################################### ## Download source and add auxiliary scripts ############################################################################### @@ -52,19 +15,11 @@ ARG URLREF_MAXTEXT ARG SRC_PATH_MAXTEXT ARG URLREF_JETSTREAM ARG SRC_PATH_JETSTREAM -ARG URLREF_TFTEXT=https://github.com/tensorflow/text.git#master -ARG SRC_PATH_TFTEXT=/opt/tensorflow-text - - -# Preserve version information of tensorflow-text -COPY --from=tftext-builder ${SRC_PATH_TFTEXT}/tensorflow_text*.whl /opt/ -RUN echo "tensorflow-text @ file://$(ls /opt/tensorflow_text*.whl)" >> /opt/pip-tools.d/requirements-maxtext.in RUN <<"EOF" bash -ex git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} echo "-r ${SRC_PATH_MAXTEXT}/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in - # specify some restrictions to speed up the build and # avoid pip to download and check all available versions of packages for pattern in \ @@ -81,7 +36,6 @@ sed -i "/google-jetstream/d" ${SRC_PATH_MAXTEXT}/requirements.txt # Add the upstream JetStream git-clone.sh ${URLREF_JETSTREAM} ${SRC_PATH_JETSTREAM} echo "-r ${SRC_PATH_JETSTREAM}/requirements.txt" >> /opt/pip-tools.d/requirements-jetstream.in - EOF ############################################################################### @@ -99,6 +53,3 @@ FROM mealkit as final RUN pip-finalize.sh WORKDIR ${SRC_PATH_MAXTEXT} - -# When tftext and lingvo wheels are published on pypi.org, revert this -# Dockerfile to 5c4b687b918e6569bca43758c346ad8e67460154 From 43b2bd003037dcd5c8e1492f82aef80ab408a548 Mon Sep 17 00:00:00 2001 From: Vladislav Kozlov Date: Fri, 10 Jan 2025 08:58:05 -0800 Subject: [PATCH 7/8] Add jetstream to manifest file, and updated main CI file --- .github/container/Dockerfile.maxtext | 2 +- .github/container/manifest.yaml | 5 +++++ .github/workflows/_ci.yaml | 1 + .github/workflows/ci.yaml | 2 +- 4 files changed, 8 insertions(+), 2 deletions(-) diff --git a/.github/container/Dockerfile.maxtext b/.github/container/Dockerfile.maxtext index d58593606..0ccc1d835 100644 --- a/.github/container/Dockerfile.maxtext +++ b/.github/container/Dockerfile.maxtext @@ -3,7 +3,7 @@ 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 URLREF_JETSTREAM=https://github.com/AI-Hypercomputer/JetStream.git#main ARG SRC_PATH_JETSTREAM=/opt/jetstream ############################################################################### diff --git a/.github/container/manifest.yaml b/.github/container/manifest.yaml index 433871c14..36705ae3f 100644 --- a/.github/container/manifest.yaml +++ b/.github/container/manifest.yaml @@ -109,6 +109,11 @@ maxtext: tracking_ref: main latest_verified_commit: 78daad198544def8274dbd656d122fbe6a0e1129 mode: git-clone +jetstream: + url: https://github.com/AI-Hypercomputer/JetStream.git + tracking_ref: main + latest_verified_commit: a3bb30e972a344675c20d56d8f61dc309e6a5f7b + mode: git-clone levanter: url: https://github.com/stanford-crfm/levanter.git tracking_ref: main diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml index 6bfc05e3c..688c93981 100644 --- a/.github/workflows/_ci.yaml +++ b/.github/workflows/_ci.yaml @@ -111,6 +111,7 @@ jobs: DOCKERFILE: .github/container/Dockerfile.maxtext EXTRA_BUILD_ARGS: | URLREF_MAXTEXT=${{ fromJson(inputs.SOURCE_URLREFS).MAXTEXT }} + URLREF_JETSTREAM=${{ fromJson(inputs.SOURCE_URLREFS).JETSTREAM }} secrets: inherit build-levanter: diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 0c3c8bdb0..0d0ffce05 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -37,7 +37,7 @@ on: type: string description: | A comma-separated PACKAGE=URL#REF list to override sources used by build. - PACKAGE∊{JAX,XLA,Flax,transformer-engine,T5X,paxml,praxis,maxtext,levanter,haliax,mujuco,mujuco-mpc,gemma,big-vision,common-loop-utils,flaxformer,panopticapi} (case-insensitive) + PACKAGE∊{JAX,XLA,Flax,transformer-engine,T5X,paxml,praxis,maxtext,jetstream,levanter,haliax,mujuco,mujuco-mpc,gemma,big-vision,common-loop-utils,flaxformer,panopticapi} (case-insensitive) default: '' required: false From a4eb6d81f9286e578209050636c246ee2a3833f3 Mon Sep 17 00:00:00 2001 From: Vladislav Kozlov Date: Wed, 15 Jan 2025 21:28:01 -0800 Subject: [PATCH 8/8] Addressed comments --- .github/container/Dockerfile.maxtext | 4 ++-- .github/workflows/_ci.yaml | 4 +--- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/.github/container/Dockerfile.maxtext b/.github/container/Dockerfile.maxtext index 0ccc1d835..5abea7d6e 100644 --- a/.github/container/Dockerfile.maxtext +++ b/.github/container/Dockerfile.maxtext @@ -18,7 +18,7 @@ ARG SRC_PATH_JETSTREAM RUN <<"EOF" bash -ex git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} -echo "-r ${SRC_PATH_MAXTEXT}/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in +echo "-e ${SRC_PATH_MAXTEXT}" >> /opt/pip-tools.d/requirements-maxtext.in # specify some restrictions to speed up the build and # avoid pip to download and check all available versions of packages @@ -35,7 +35,7 @@ sed -i "/google-jetstream/d" ${SRC_PATH_MAXTEXT}/requirements.txt # Add the upstream JetStream git-clone.sh ${URLREF_JETSTREAM} ${SRC_PATH_JETSTREAM} -echo "-r ${SRC_PATH_JETSTREAM}/requirements.txt" >> /opt/pip-tools.d/requirements-jetstream.in +echo "-e ${SRC_PATH_JETSTREAM}" >> /opt/pip-tools.d/requirements-jetstream.in EOF ############################################################################### diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml index 688c93981..c38a585f4 100644 --- a/.github/workflows/_ci.yaml +++ b/.github/workflows/_ci.yaml @@ -698,11 +698,9 @@ jobs: with: TEST_NAME: jetstream EXECUTE: | - docker run --shm-size=1g --gpus all ${{ needs.build-maxtext.outputs.DOCKER_TAG_FINAL }} \ + docker run --shm-size=1g ${{ 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 python -m unittest -v jetstream.tests.core.test_orchestrator python -m jetstream.engine.mock_engine_test python -m jetstream.core.orchestrator_test