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

Add option to register installed plugins automatically #6328

Merged
merged 4 commits into from
Jan 19, 2024

Conversation

will-cromar
Copy link
Collaborator

See #6242 for context

  • Add XLA_REGISTER_INSTALLED_PLUGINS to automatically register any plugin with the torch_xla.plugins entry point.
  • Requires recent additions to importlib.metadata, which are backported in importlib_metadata.
  • The default behavior remains the same for now. After this and related PRs are all merged, I'll try switching our CI over, then changing the default behavior if all tests pass. For now, I tested locally.

@will-cromar will-cromar requested review from jonb377, JackCaoG and vanbasten23 and removed request for jonb377 January 18, 2024 23:09
@will-cromar will-cromar marked this pull request as draft January 18, 2024 23:26
@will-cromar
Copy link
Collaborator Author

Marking this a draft for now while I figure out how to register TpuPlugin from within torch_xla too.



def register_installed_plugins():
pjrt_entry_points = importlib_metadata.entry_points(group='torch_xla.plugins')
Copy link
Collaborator

@vanbasten23 vanbasten23 Jan 19, 2024

Choose a reason for hiding this comment

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

I see that you have 'torch_xla.plugins': ['tpu = torch_xla._internal.tpu:TpuPlugin',], in the setup.py. I wonder what this line does under the hood

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Entry points are weird. setuptools' docs have the best explanation that I've read.

@will-cromar will-cromar marked this pull request as ready for review January 19, 2024 17:51
@will-cromar
Copy link
Collaborator Author

Tested manually on TPU to confirm that this works.

@will-cromar will-cromar merged commit d7ddecd into master Jan 19, 2024
2 of 3 checks passed
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.

3 participants