Support out-of-tree plugins autoloading (#1151)

Co-authored-by: Nicolas Hug <contact@nicolas-hug.com>
Signed-off-by: Dmitry Rogozhkin <dmitry.v.rogozhkin@intel.com>

Author: dvrogozh

.gitignore

@@ -1,6 +1,7 @@
 build/
 dist/
 src/TorchCodec.egg-info/
+test/plugin/torchcodec_test_plugin.egg-info/
 */**/__pycache__
 */__pycache__
 */*.pyc

README.md

@@ -170,12 +170,6 @@ Intel GPUs (XPU) support requires a stand-alone plugin for TorchCodec:
 pip install torchcodec-xpu --extra-index-url=https://download.pytorch.org/whl/xpu
 ```
 
-Upon installation, the plugin needs to be imported:
-
-```python
-import torchcodec_xpu
-```
-
 For any XPU-related support, please refer to
 https://github.com/intel/torchlib-xpu.
 

src/torchcodec/__init__.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import os
 from pathlib import Path
 
 # Note: usort wants to put Frame and FrameBatch after decoders and samplers,
@@ -25,3 +26,25 @@
 # Similarly, these are exposed for downstream builds that use torchcodec as a
 # dependency.
 from ._core import core_library_path, ffmpeg_major_version  # usort:skip
+
+
+# Leverage the Python plugin mechanism to load out-of-the-tree device extensions.
+# See https://github.com/pytorch/pytorch/pull/127074 that enabled the same
+# plugin support in torch.
+# This block should be kept at the end to ensure all the other functions in this
+# module that may be accessed by an autoloaded backend are defined.
+_TORCHCODEC_DEVICE_BACKEND_AUTOLOAD_VAR_NAME = "TORCHCODEC_DEVICE_BACKEND_AUTOLOAD"
+if os.getenv(_TORCHCODEC_DEVICE_BACKEND_AUTOLOAD_VAR_NAME, "1") == "1":
+    from importlib.metadata import entry_points
+
+    _backend_extensions = entry_points(group="torchcodec.backends")
+
+    for _backend_extension in _backend_extensions:
+        try:
+            _entrypoint = _backend_extension.load()
+            _entrypoint()
+        except Exception as _err:
+            raise RuntimeError(
+                f"Failed to load the backend extension: {_backend_extension.name}. "
+                f"You can disable extension auto-loading with {_TORCHCODEC_DEVICE_BACKEND_AUTOLOAD_VAR_NAME}=0."
+            ) from _err

test/plugin/dummy_plugin/__init__.py

@@ -0,0 +1,9 @@
+"""
+This is a device backend extension used for testing.
+"""
+
+import os
+
+
+def _autoload():
+    os.environ["IS_DUMMY_PLUGIN_LOADED"] = "1"

test/plugin/pyproject.toml

@@ -0,0 +1,12 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torchcodec-test-plugin"
+description = "Test extension for torchcodec"
+requires-python = ">=3.8"
+dynamic = ["version"]
+
+[project.entry-points.'torchcodec.backends']
+device_backend = 'dummy_plugin:_autoload'

test/plugin/test_plugins.py

@@ -0,0 +1,51 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+
+@pytest.mark.parametrize("enable_autoload", [True, False])
+def test_plugin_autoload(tmp_path, enable_autoload):
+    plugin_dir = Path(__file__).parent
+    install_dir = str(tmp_path / "install")
+
+    # Install the dummy plugin into a temp directory. We use --target so that
+    # the .dist-info metadata and the plugin package are installed flat into a
+    # separate directory to avoid polluting the current env's site-packages.
+    subprocess.run(
+        [sys.executable, "-m", "pip", "install", "--target", install_dir, plugin_dir],
+        check=True,
+    )
+
+    env = {
+        **os.environ,
+        "PYTHONPATH": os.pathsep.join([install_dir, os.environ.get("PYTHONPATH", "")]),
+        "TORCHCODEC_DEVICE_BACKEND_AUTOLOAD": str(int(enable_autoload)),
+    }
+
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            """
+import os
+import torchcodec
+
+loaded = os.environ.get("IS_DUMMY_PLUGIN_LOADED") == "1"
+autoload = os.environ.get("TORCHCODEC_DEVICE_BACKEND_AUTOLOAD") == "1"
+assert loaded == autoload, f"{loaded=} {autoload=}"
+""",
+        ],
+        env=env,
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 0, result.stderr