[Docker] Update TPU Docker image base deps#63006
Merged
elliot-barn merged 9 commits intoMay 7, 2026
Merged
Conversation
Contributor
Author
|
This is what the |
![gemini-code-assist[bot]](https://cdn.statically.io/img/avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request updates the dependency configuration for TPU environments, including switching to --find-links in the depset configuration and expanding the requirements_tpu.in file with necessary JAX ecosystem packages. Feedback was provided regarding a potential version typo for JAX and the need for a lower bound on the numpy version constraint to ensure compatibility.
ray-gardener
Bot
added
core
elliot-barn
reviewed
Apr 29, 2026
Collaborator
|
For some more context: We have a custom python dependency management tool called raydepsets. Its a uv wrapper that uses configs to determine requirements, constraints for rays dependencies and expands them for different workloads/use cases |
![cursor[bot]](https://cdn.statically.io/img/avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Contributor
Author
Sounds good, ran the depset file and changed it to a lower bound in 9346d34, Now this version gets installed: |
Contributor
Author
|
@elliot-barn should be updated now accordingly, I added a python dependency check in the TPU requirements file since the JAX version for Ironwood TPU requires Python |
elliot-barn
approved these changes
May 2, 2026
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Reviewed by Cursor Bugbot for commit d633d6ad0093bcd9d2d87d3189af9ec6bd7bacdd. Configure here.
Contributor
Author
|
had to change the tpu image to build from in d633d6a |
Contributor
Author
|
@elliot-barn wondering if there's any concern with the change here: #63006 (comment), if not I'll go ahead and merge this. Thanks again!! |
Collaborator
This looks great can you fix the merge conflicts and i can merge for you |
ryanaoleary
force-pushed
the
update-tpu-dependencies
branch
from
fa3101f to
41c0e53
Compare
Signed-off-by: ryanaoleary <ryanaoleary@google.com>
Signed-off-by: ryanaoleary <ryanaoleary@google.com>
… python version Signed-off-by: ryanaoleary <ryanaoleary@google.com>
Signed-off-by: ryanaoleary <ryanaoleary@google.com>
Signed-off-by: ryanaoleary <ryanaoleary@google.com>
Signed-off-by: ryanaoleary <ryanaoleary@google.com>
Signed-off-by: ryanaoleary <ryanaoleary@google.com>
Signed-off-by: ryanaoleary <ryanaoleary@google.com>
Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
ryanaoleary
force-pushed
the
update-tpu-dependencies
branch
from
41c0e53 to
6b584ea
Compare
Contributor
Author
Sounds good thanks!! rebased on master and re-ran the build depsets script |
Lucas61000
pushed a commit
to Lucas61000/ray
that referenced
this pull request
May 15, 2026
## Description This PR resolves ray-project#62402 by ensuring a `jax[tpu]` version >= `0.8.2` is installed. More context for why this change is necessary is provided in the linked issue. This PR also adds several lightweight packages that are commonly used with `jax[tpu]` and `libtpu`. The list of updated dependencies is based on common packages seen in JAX-based frameworks like: [Maxtext](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/dependencies/requirements/requirements.txt) and [tpu_inference](https://github.com/vllm-project/tpu-inference/blob/main/requirements.txt). ## Related issues ray-project#62402 --------- Signed-off-by: ryanaoleary <ryanaoleary@google.com> Signed-off-by: Ryan O'Leary <ryanaoleary@google.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
2 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
This PR resolves #62402 by ensuring a
jax[tpu]version >=0.8.2is installed. More context for why this change is necessary is provided in the linked issue.This PR also adds several lightweight packages that are commonly used with
jax[tpu]andlibtpu. The list of updated dependencies is based on common packages seen in JAX-based frameworks like: Maxtext and tpu_inference.Related issues
#62402