Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Don't set $TPU_LIBRARY_PATH during import #5698

Merged
merged 2 commits into from
Oct 16, 2023
Merged

Conversation

will-cromar
Copy link
Collaborator

@will-cromar will-cromar commented Oct 11, 2023

JAX and PyTorch/XLA are both overriding $TPU_LIBRARY_PATH during import, leading to confusing issues that depend on import order, e.g. #5625

  • Use the new get_library_path function to get that path from the libtpu package.
  • Keep $TPU_LIBRARY_PATH as an override for compatibility

Corresponding fix in JAX: jax-ml/jax@b81a3e1

The issue will be resolved when both JAX and PyTorch/XLA release new packages that include these two changes.

$PTXLA_TPU_LIBRARY_PATH is a temporary hack. I expect we'll have a better way to get the PJRT Plugin path in the near future.

cc @jyingl3

@JackCaoG
Copy link
Collaborator

I am a bit confuse after this pr, what will happen if user still set TPU_LIBRARY_PATH ?

@will-cromar
Copy link
Collaborator Author

will-cromar commented Oct 11, 2023

We'll still take TPU_LIBRARY_PATH as an override. The change here is that we will not set it. Right now, if you import us first, we'll override JAX's libtpu path, since they also treat TPU_LIBRARY_PATH as a user override.

Internally, we can use PTXLA_... so we stop conflicting with JAX, but we need a better way to plumb the plugin path.


We load libtpu.so in the following order of precedence:

1. User-set $TPU_LIBRARY_PATH
2. libtpu.so included in torch_xla/lib
3. libtpu-nightly pip package

Sets $PTXLA_TPU_LIBRARY_PATH if path is inferred by us to prevent conflicts
with other frameworks. This env var will be removed in a future version.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder under what condition can we remove this env var.

@will-cromar will-cromar merged commit 146f2a0 into master Oct 16, 2023
18 checks passed
zpcore pushed a commit that referenced this pull request Oct 19, 2023
* Don't set $TPU_LIBRARY_PATH during import

* remove chekck for new env var so people don't use it
@alanwaketan
Copy link
Collaborator

This PR regress the profiler. I cannot take multi-host profilers after this change. I'm going to revert it.

alanwaketan added a commit that referenced this pull request Oct 25, 2023
alanwaketan added a commit that referenced this pull request Oct 25, 2023
jonb377 pushed a commit that referenced this pull request Oct 31, 2023
ghpvnist pushed a commit to ghpvnist/xla that referenced this pull request Oct 31, 2023
* Don't set $TPU_LIBRARY_PATH during import

* remove chekck for new env var so people don't use it
will-cromar added a commit that referenced this pull request Nov 13, 2023
will-cromar added a commit that referenced this pull request Nov 14, 2023
mbzomowski pushed a commit to mbzomowski-test-org/xla that referenced this pull request Nov 16, 2023
* Don't set $TPU_LIBRARY_PATH during import

* remove chekck for new env var so people don't use it
mbzomowski pushed a commit to mbzomowski-test-org/xla that referenced this pull request Nov 16, 2023
chunnienc pushed a commit to chunnienc/xla that referenced this pull request Dec 14, 2023
* Don't set $TPU_LIBRARY_PATH during import

* remove chekck for new env var so people don't use it
chunnienc pushed a commit to chunnienc/xla that referenced this pull request Dec 14, 2023
golechwierowicz pushed a commit that referenced this pull request Jan 12, 2024
* Don't set $TPU_LIBRARY_PATH during import

* remove chekck for new env var so people don't use it
golechwierowicz pushed a commit that referenced this pull request Jan 12, 2024
bhavya01 pushed a commit that referenced this pull request Apr 22, 2024
* Don't set $TPU_LIBRARY_PATH during import

* remove chekck for new env var so people don't use it
bhavya01 pushed a commit that referenced this pull request Apr 22, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants