[JAX] Fixes for L0_jax_distributed_unittest#1884
Merged
phu0ngng merged 11 commits intoNVIDIA:mainfrom Jun 17, 2025
Merged
Conversation
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Collaborator
Author
|
/te-ci JAX L0 |
jaro-sevcik
reviewed
Jun 16, 2025
| python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_multigpu_encoder.xml $TE_PATH/examples/jax/encoder/test_multigpu_encoder.py || test_fail "test_multigpu_encoder.py" | ||
| wait | ||
| python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_model_parallel_encoder.xml $TE_PATH/examples/jax/encoder/test_model_parallel_encoder.py || test_fail "test_model_parallel_encoder.py" | ||
| # wait |
There was a problem hiding this comment.
Any reason to not uncomment this wait?
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Collaborator
Author
|
/te-ci JAX L0 |
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Collaborator
Author
|
/te-ci JAX L0 |
Comment on lines
+27
to
+31
| python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_multigpu_encoder.xml $TE_PATH/examples/jax/encoder/test_multigpu_encoder.py || test_fail "test_multigpu_encoder.py" | ||
| wait | ||
| python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_model_parallel_encoder.xml $TE_PATH/examples/jax/encoder/test_model_parallel_encoder.py || test_fail "test_model_parallel_encoder.py" | ||
| wait | ||
| bash $TE_PATH/examples/jax/encoder/run_test_multiprocessing_encoder.sh || test_fail "run_test_multiprocessing_encoder.sh" |
Collaborator
There was a problem hiding this comment.
Looks good to me !
Quick question:
- Do we know if these were ever enabled in the past ? If yes, is there a reason we disabled them ?
- Also, interested in knowing if we know approximately how much longer will this test run due to us uncommenting these tests ? (asking from a CI/QA budget point of view)
- Is it okay to explicitly uses
bash- what if the system running the script has a shell is notbash?
Thanks !
Collaborator
Author
There was a problem hiding this comment.
- I accidentally left out in one of the PR that merged recently.
- We have always been running all of these tests, and the whole suite takes ~25 mins, I think so we should not have any issues.
- The
qa/L0_jax_distributed_unittest/test.shis a bash script itself.
Collaborator
|
These tests were all running recently, I assume they were recently disabled when debugging a PR and accidentally merged, right? If so, then LGTM pending one CI failure |
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Collaborator
Author
|
/te-ci JAX L0 |
Collaborator
Author
|
/te-ci JAX L0 |
1 similar comment
Collaborator
|
/te-ci JAX L0 |
jberchtold-nvidia
approved these changes
Jun 17, 2025
| # python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_model_parallel_encoder.xml $TE_PATH/examples/jax/encoder/test_model_parallel_encoder.py || test_fail "test_model_parallel_encoder.py" | ||
| # wait | ||
| . $TE_PATH/examples/jax/encoder/run_test_multiprocessing_encoder.sh || test_fail "run_test_multiprocessing_encoder.sh" | ||
| python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_multigpu_encoder.xml $TE_PATH/examples/jax/encoder/test_multigpu_encoder.py || test_fail "test_multigpu_encoder.py" |
Collaborator
There was a problem hiding this comment.
Might be unrelated to this PR, but TE_PATH doesn't seem to be set correctly at least in one configuration. CI is failing here with
FileNotFoundError: [Errno 2] No such file or directory: '/tests/jax/pytest.ini'
which I'd assume is because TE_PATH is empty.
Collaborator
Author
There was a problem hiding this comment.
Yeah, that was not the latest pipeline.
The latest pipeline (#30209220) shows that all tests have passed.
KshitijLakhani
pushed a commit
that referenced
this pull request
Jun 17, 2025
* include previously accidentally excluded tests * Execute run_test_multiprocessing_encoder with nested bash + exit code for inner bash shell * Adapt run_test_multiprocessing to handle segfault Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
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
run_test_multiprocessing_encoder.shto avoid premature exit in the test suite.Type of change
Checklist: