Maxtext unit tests with Pathways backend. #1211
Open
+187
−2
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
RoshaniN
requested review from
gobbleturk,
khatwanimohit,
bvandermoon,
vipannalla and
RissyRan
as code owners
RoshaniN
force-pushed
the
roshanin_unit
branch
13 times, most recently
from
3901ccb to
27fbae0
Compare
RoshaniN
force-pushed
the
roshanin_unit
branch
12 times, most recently
from
0fdec67 to
a688c60
Compare
RoshaniN
force-pushed
the
roshanin_unit
branch
6 times, most recently
from
c8d3e1c to
5cfb878
Compare
RoshaniN
force-pushed
the
roshanin_unit
branch
from
5cfb878 to
d6a6e21
Compare
RoshaniN
changed the title
shauryagup
reviewed
Feb 27, 2025
RoshaniN
force-pushed
the
roshanin_unit
branch
from
070f151 to
e016673
Compare
gobbleturk
approved these changes
Feb 28, 2025
gagika
approved these changes
Feb 28, 2025
RoshaniN
force-pushed
the
roshanin_unit
branch
5 times, most recently
from
a8ed8d1 to
7393067
Compare
Modified Pathways workflow to run specifically on Pathways runner. Final tests on Pathways. New changes to help run tests. Trial to get Pathways working. run Pathways containers and Maxtext as part of the same job. Move installation to script More changes. Installing docker also as part of the script. Simplified flow test. Few more changes. New way of installation Docker compose with Maxtext and Pathways containers. Other changes to use maxtext container in docker compose YAML. Trying to merge all services together. Adding everything as services again. Directly running docker compose. Elegant solution with Pathways containers as services. Elegant solution with Pathways containers as services. Reverting to Maxtext in docker compose. Changes to run with latest JAX SS image. Run tests with correct markers. Move installation to Github runner step, simplify command. Formalize the changes. Formalize the changes. Formalize the changes. Notify on Pathways test failures, modify test script. Error handling.
RoshaniN
requested review from
shralex,
yangyuwei,
SurbhiJainUSC,
hengtaoguo and
A9isha
as code owners
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR enables TPU unit tests to also run with Pathways backend. Essentially, we will have two sets of tests - one with McJAX and one with Pathways.
For more details, please read the doc on b/397475777 . Note that extra self-hosted runners have been added so that tests can be executed in parallel and complete faster overall.
Tests
Please describe how you tested this change -
bash docker_run_pathways_containers.sh maxtext_image=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest command="cd MaxText ; python3 -m pytest tests -m 'not gpu_only and not integration_test' -s"Example runs -
Checklist
Before submitting this PR, please make sure (put X in square brackets):