Implement autoload device extension mechanism (from RFC #122468) #127228
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/127228
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 2f938df with merge base 1110edb ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
|
| @@ -2065,6 +2065,28 @@ def _constrain_as_size(symbol, min: Optional[builtins.int] = None, max: Optional | |||
| """ | |||
| torch.sym_constrain_range_for_size(symbol, min=min, max=max) | |||
|
|
|||
| def import_device_backends(): | |||
There was a problem hiding this comment.
Please move the change to the end of file
| for backend in entry_points(group='torch.backends'): | ||
| try: | ||
| backend_hook = backend.load() | ||
| if not hasattr(backend_hook, 'init_custom_backend'): |
There was a problem hiding this comment.
Please change a name of function to init_device_backend for consistency
There was a problem hiding this comment.
Why don't we use the official load() function of EntryPoint to get the init function to call?
There was a problem hiding this comment.
Sure, we can simplify to initialize at load
| for backend in entry_points(group='torch.backends'): | ||
| try: | ||
| backend_hook = backend.load() | ||
| if not hasattr(backend_hook, 'init_custom_backend'): |
There was a problem hiding this comment.
Why don't we use the official load() function of EntryPoint to get the init function to call?
|
|
||
| def is_device_backend_autoload_enabled() -> bool: | ||
| var = os.getenv("TORCH_DISABLE_DEVICE_BACKEND_AUTOLOAD") | ||
| return not var in (1, 'True', 'true', 'Yes', 'yes') |
There was a problem hiding this comment.
nit: suggest to do var.upper() to simplify the code.
There was a problem hiding this comment.
Since os.getenv() returns None, when the variable isn't set, that's how I would see it then:
def is_device_backend_autoload_enabled() -> bool:
var = os.getenv("TORCH_DISABLE_DEVICE_BACKEND_AUTOLOAD")
return not str(var).upper() in ('1', 'TRUE', 'YES')What do You think?
There was a problem hiding this comment.
def is_device_backend_autoload_enabled() -> bool:
var = os.getenv("TORCH_DISABLE_DEVICE_BACKEND_AUTOLOAD")
return var is None or not str(var).upper() in ('1', 'TRUE', 'YES')| for backend in entry_points(group='torch.backends'): | ||
| try: | ||
| backend_hook = backend.load() | ||
| if not hasattr(backend_hook, 'init_custom_backend'): | ||
| print(f"No explicit backend init function for \'{backend_hook.__name__}\' has been found, custom backend can't be loaded.") | ||
| continue | ||
| backend_hook.init_custom_backend() | ||
| except IndexError: | ||
| pass |
There was a problem hiding this comment.
I think we don't need to call the init function, only load it.
Since the init function is implemented in your extension, your extension will be imported automatically after backend.load() is called. And calling it directly might be risky.
Here is a simple example and any suggestions are welcome.
in torch_npu implementation:
# setup.py
setup(
entry_points={
'torch.backends': [
'torch_npu = torch_npu:_init_device_backend',
],
}
)
# torch_npu/__init__.py
def _init_device_backend():
passin pytorch:
# torch/__init__.py
for plugin in discovered_plugins:
try:
# just loads the plugin without calling
plugin.load()
except Exception:
# keep quiet
pass
These changes refer to RFC from #122468 issue. Early testing has been done against habana_frameworks extension, with the entry point exported according to the RFC discussion.
Things to do:
Looking forward for your review.