Ensure JAX installation is editable #1416
Merged
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
Contributor
|
You'll need to change |
Contributor
Author
done 😁 |
gpupuck
force-pushed
the
fix-editable-build
branch
from
c0d8fdf to
5d77674
Compare
gpupuck
reviewed
Apr 25, 2025
Contributor
There was a problem hiding this comment.
pip --disable-pip-version-check install .... -e ${BUILD_PATH_JAXLIB}/jax <------
...
## after installation (example)
# jax 0.5.4.dev20250325 /opt/jaxlibs/jax <--------
# jax-cuda12-pjrt 0.5.4.dev20250325 /opt/jaxlibs/jax_cuda12_pjrt
Install dir and resulting output needs to be updated too
Contributor
Author
Fixed |
gpupuck
approved these changes
Apr 25, 2025
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.
Currently, the JAX Python source that we installed in the container is created by
build-jax.sh(via python build.py build --editable ... --wheels=jax,jaxlib,jax-cuda-plugin,jax-cuda-pjrt) and placed under/opt/jaxlibs/jax:As a result, two copies of the JAX source exist in the container at
/opt/jaxand/opt/jaxlibs/jax, leading to potential confusions. Also, this defeats the 'editable' purpose of the build, as mounting a JAX working copy from outside the container to/opt/jaxwill not replace the installed JAX package.Given that the editable wheel at
/opt/jaxlibs/jaxis practically identical to the original source (sans a number of build-related files), we should bypass creating this extra copy and directly specify/opt/jaxas the installation target in the Dockerfile and build script. With this change, we will havewhich is the desired usage pattern.