Add option to register installed plugins automatically

setup.py

@@ -346,14 +346,18 @@ def run(self):
         'absl-py>=1.0.0',
         'cloud-tpu-client>=0.10.0',
         'pyyaml',
+        # importlib.metadata backport required for PJRT plugin discovery prior
+        # to Python 3.10
+        'importlib_metadata>=4.6;python_version<"3.10"',
     ],
     package_data={
         'torch_xla': ['lib/*.so*',],
     },
     entry_points={
         'console_scripts': [
             'stablehlo-to-saved-model = torch_xla.tf_saved_model_integration:main'
-        ]
+        ],
+        'torch_xla.plugins': ['tpu = torch_xla._internal.tpu:TpuPlugin',],
     },
     extras_require={
         # On Cloud TPU VM install with:

torch_xla/__init__.py

@@ -171,3 +171,9 @@ def _init_xla_lazy_backend():
 torch._dynamo.config.automatic_dynamic_shapes = False
 
 from .stablehlo import save_as_stablehlo, save_torch_model_as_stablehlo
+
+from .experimental import plugins
+
+if os.getenv('XLA_REGISTER_INSTALLED_PLUGINS') == '1':
+  plugins.use_dynamic_plugins()
+  plugins.register_installed_plugins()

torch_xla/experimental/plugins.py

@@ -1,4 +1,13 @@
+import logging
 import os
+import sys
+
+# TODO: delete this and just use importlib.metadata after we drop Python 3.9
+# support.
+if sys.version_info < (3, 10):
+  import importlib_metadata
+else:
+  import importlib.metadata as importlib_metadata
 
 import torch_xla
 import torch_xla.core.xla_env_vars as xenv
@@ -65,3 +74,18 @@ def default() -> DevicePlugin:
 def register_plugin(name: str, device_plugin: DevicePlugin):
   _plugin_registry[name.upper()] = device_plugin
   torch_xla._XLAC._register_pjrt_plugin(name, device_plugin.library_path())
+
+
+def register_installed_plugins():
+  pjrt_entry_points = importlib_metadata.entry_points(group='torch_xla.plugins')
+  for ep in pjrt_entry_points:
+    device_plugin_class = ep.load()
+
+    # HACK: TpuPlugin raises EnvironmentError if libtpu is not installed.
+    # TODO(wcromar): Decide if catching `EnvironmentError` is a permanent
+    # behavior or temporary hack.
+    try:
+      register_plugin(ep.name.upper(), device_plugin_class())
+    except EnvironmentError as e:
+      logging.warning(
+          "Failed to register plugin {}".format(ep.name), exc_info=e)