[RFC] Add support for device extension autoloading by shink · Pull Request #127074 · pytorch/pytorch

shink · 2024-05-24T09:31:19Z

Load device extensions at the end of torch/__init__.py
Enabled by default, or you can disable it with TORCH_DEVICE_BACKEND_AUTOLOAD=0

run test:

python test/run_test.py -i test_autoload_enable
python test/run_test.py -i test_autoload_disable

doc:

https://docs-preview.pytorch.org/pytorch/pytorch/127074/miscellaneous_environment_variables.html

co-author: @jgong5 @bsochack @bkowalskiINTEL @jczaja @fffrog @hipudding

cc @albanD

pytorch-bot · 2024-05-24T09:31:22Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/127074

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 45fd7ad with merge base 6b5fbc5 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

jgong5 · 2024-05-24T09:47:01Z

Thanks for the pull request. Could you please add some tests?

jgong5 · 2024-05-24T09:49:11Z

cc @bsochack

shink · 2024-05-24T09:56:24Z

Thanks for the pull request. Could you please add some tests?

Sure, I'm willing to do this!

bsochack · 2024-05-27T14:45:58Z

@shink we pushed a draft PR of a change here: #127228
A few changes (as in review) + unit tests are still required.

jczaja · 2024-05-27T14:53:28Z

@shink , @jgong5 This is a nice work that directly reflect proposal (RFC). We have tried to test it against our device_extension and there was a crash :

Problem seems to be that torch._dynamo is having a code checking all symbols exported from torch/init.py .
So to avoid that it would be good to hide changes proposed here inside a function. This issue may be observed with extensions that are directly importing torch._dynamo . Here is other PR implementing mentioned RFC that works for us: #127228 .

shink · 2024-05-28T01:23:06Z

@bsochack @jczaja Nice work! Thanks so much!

@shink , @jgong5 This is a nice work that directly reflect proposal (RFC). We have tried to test it against our device_extension and there was a crash :

Problem seems to be that torch._dynamo is having a code checking all symbols exported from torch/init.py . So to avoid that it would be good to hide changes proposed here inside a function. This issue may be observed with extensions that are directly importing torch._dynamo . Here is other PR implementing mentioned RFC that works for us: #127228 .

Yes, we alse met this crash.

shink · 2024-05-31T08:47:52Z

@jgong5 @bsochack @bkowalskiINTEL @fffrog @hipudding Please have a look at this unit test. See b0e08d3

shink · 2024-05-31T08:52:52Z

in fact, this file should be ignored, and the reason it is here is to declare the entry point of the backend package

fffrog · 2024-05-31T09:50:02Z

I think it would be better to show failure information when an exception is encountered.

fffrog · 2024-05-31T09:51:05Z

getenv with default value shoule be better.

fffrog · 2024-05-31T09:52:48Z

It would be better to think of this path operation as context.

fxied, please take a look again

bsochack · 2024-05-31T10:14:03Z

@jgong5 @bsochack @bkowalskiINTEL @fffrog @hipudding Please have a look at this unit test. See b0e08d3

Please add @bkowalskiINTEL and @jczaja as co-authors. Lets also continue all addressing code review here.

shink · 2024-05-31T10:41:40Z

@jgong5 @bsochack @bkowalskiINTEL @fffrog @hipudding Please have a look at this unit test. See b0e08d3

Please add @bkowalskiINTEL and @jczaja as co-authors. Lets also continue all addressing code review here.

Sure, thanks so much, and any suggestions are welcome.

jgong5 · 2024-06-03T02:45:57Z

Besides checking the environment variable, are you going to add test to autoload an extension (maybe a mock one)?

Thanks for your review. Please have a look at the code I just updated.

albanD · 2024-06-03T18:35:52Z

Would it be possible to reuse the existing modules we already have for testing in test/cpp_extension/setup.py to test this without having to create a brand new extension?

Thank you so much for your review. I'm testing this idea now.

Could you please have a look at this commit 1e7cf76

albanD · 2024-06-03T18:37:04Z

I'm not sure what this is testing?

it's testing a function defined in the backend extension

jgong5 · 2024-06-07T05:12:28Z

@shink Seems there is still lint errors. You may repro with "lintrunner -a" locally.

jgong5 · 2024-06-07T05:23:40Z

Should we do things inside this entrypoint instead to make sure it is called?

jgong5 · 2024-06-07T05:26:28Z

Not sure if I understand it correctly but when the cpp extension is explicitly imported from this test file, not autoloaded implicitly. Is that correct?

The reason for putting this test case here is that I want to reuse the existing modules in test/cpp_extension/setup.py.

The extension is installed here:

pytorch/test/run_test.py

Lines 667 to 671 in 7efaeb1

# Build the test cpp extensions modules

shell_env = os.environ.copy()

shell_env["USE_NINJA"] = str(1 if use_ninja else 0)

cmd = [sys.executable, "setup.py", "install", "--root", "./install"]

return_code = shell(cmd, cwd=cpp_extensions_test_dir, env=shell_env)

albanD · 2024-06-25T19:46:29Z

+    """
+    Whether autoloading out-of-the-tree device extensions is enabled.
+    The switch depends on the value of the environment variable
+    `TORCH_DEVICE_BACKEND_AUTOLOAD`.


@drisspg which .rst file should we add this new env variable to for proper documentation?

The main env var RST is here:https://github.com/pytorch/pytorch/blob/main/docs/source/torch_environment_variables.rst

and they branch out to others

Co-authored-by: albanD <desmaison.alban@gmail.com>

bsochack · 2024-07-02T17:28:41Z

@albanD any comments to this PR?

albanD

Small nit on moving the doc to another file, but SGTM otherwise!

albanD · 2024-07-02T18:13:54Z

   miscellaneous_environment_variables
   logging
   torch_nccl_environment_variables
+   privateuse1_environment_variables


Could you please move this to the miscellaneous section? I do expect all users to be interested in this env variable, not just privateuse1 devs.

pytorch/docs/source/miscellaneous_environment_variables.rst

Line 13 in 34e94c5

- Under some conditions, autograd threads can hang on shutdown, therefore we do not wait for them to shutdown indefinitely but rely on timeout that is default set to ``10`` seconds. This environment variable can be used to set the timeout in seconds.

@albanD Sure, it has been moved. Please take a look again. Thanks so much for your patient review.

jgong5 · 2024-07-09T00:06:58Z

@pytorchbot merge

pytorchmergebot · 2024-07-09T00:09:07Z

Merge failed

Reason: This PR needs a release notes: label
If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Details for Dev Infra team

Raised by workflow job

jgong5 · 2024-07-09T00:33:22Z

@pytorchbot merge

pytorchmergebot · 2024-07-09T00:35:05Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

@jgong5

Fixes pytorch#122468 - Load device extensions at the end of `torch/__init__.py` - Enabled by default, or you can disable it with `TORCH_DEVICE_BACKEND_AUTOLOAD=0` run test: ```python python test/run_test.py -i test_autoload_enable python test/run_test.py -i test_autoload_disable ``` doc: https://docs-preview.pytorch.org/pytorch/pytorch/127074/miscellaneous_environment_variables.html co-author: @jgong5 @bsochack @bkowalskiINTEL @jczaja @fffrog @hipudding Co-authored-by: albanD <desmaison.alban@gmail.com> Co-authored-by: Jiong Gong <jiong.gong@intel.com> Pull Request resolved: pytorch#127074 Approved by: https://github.com/albanD, https://github.com/jgong5

This commit adds support of out-of-tree plugins autoloading using python package entry point specification: * https://packaging.python.org/en/latest/specifications/entry-points/#entry-points * https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/#using-package-metadata Out-of-tree plugins must register entry points as `torchcodec.backends`: ``` [project.entry-points.'torchcodec.backends'] device_backend = 'torchcodec_plugin:load_plugin' ``` Torchcodec will automatically load plugins if discovered. Loading can be explicitly suppressed with `TORCHCODEC_DEVICE_BACKEND_AUTOLOAD=1` environment variable. The same approach is being used to load PyTorch device backends. See: * pytorch/pytorch#127074 Signed-off-by: Dmitry Rogozhkin <dmitry.v.rogozhkin@intel.com>

shink marked this pull request as draft May 24, 2024 09:35

pytorchbot added the open source label May 24, 2024

shink marked this pull request as ready for review May 31, 2024 08:44

shink commented May 31, 2024

View reviewed changes

shink mentioned this pull request May 31, 2024

[from RFC #122468] Implement autoload device extension mechanism #127386

Closed

fffrog reviewed May 31, 2024

View reviewed changes

Comment thread torch/__init__.py Outdated

Copy link
Copy Markdown

Collaborator

fffrog May 31, 2024

Choose a reason for hiding this comment

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

getenv with default value shoule be better.

fffrog reviewed May 31, 2024

View reviewed changes

jgong5 reviewed Jun 3, 2024

View reviewed changes

cpuhrsch requested a review from albanD June 3, 2024 18:10

cpuhrsch added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Jun 3, 2024

cpuhrsch requested a review from ezyang June 3, 2024 18:10

cpuhrsch added the module: python frontend For issues relating to PyTorch's Python frontend label Jun 3, 2024

albanD reviewed Jun 3, 2024

View reviewed changes

shink requested a review from albanD June 5, 2024 10:04

jgong5 reviewed Jun 7, 2024

View reviewed changes

shink requested a review from a team as a code owner June 7, 2024 12:25

albanD reviewed Jun 25, 2024

View reviewed changes

shink and others added 4 commits June 26, 2024 09:07

Update test/test_autoload.py [skip ci]

d9682dd

Co-authored-by: albanD <desmaison.alban@gmail.com>

remove redundant import

5ee36dd

docs: add env var rst

53e2498

docs: fix format

6ae2d01

shink requested a review from albanD June 28, 2024 01:20

albanD approved these changes Jul 2, 2024

View reviewed changes

docs: move the env to miscellaneous_environment_variables.rst

45fd7ad

jgong5 approved these changes Jul 9, 2024

View reviewed changes

pytorch-bot Bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jul 9, 2024

pytorchmergebot added the merging label Jul 9, 2024

pytorchmergebot removed the merging label Jul 9, 2024

jgong5 added the release notes: python_frontend python frontend release notes category label Jul 9, 2024

pytorchmergebot added the merging label Jul 9, 2024

pytorchmergebot added the Merged label Jul 9, 2024

pytorchmergebot closed this in 312652c Jul 9, 2024

pytorchmergebot removed the merging label Jul 9, 2024

nicholasw-gc mentioned this pull request Jul 26, 2024

Allow device specific imports for inductor wrapper #131190

Closed

shink mentioned this pull request Sep 5, 2024

Add autoloading tutorial pytorch/tutorials#3037

Merged

MengqingCao mentioned this pull request Oct 16, 2024

Add NPU backend support for val and inference huggingface/pytorch-image-models#2109

Merged

fffrog mentioned this pull request Aug 7, 2025

[OpenReg] Implement device autoload mechanism #158555

Closed

dvrogozh mentioned this pull request Jan 6, 2026

Support out-of-tree plugins autoloading meta-pytorch/torchcodec#1151

Open

	# Build the test cpp extensions modules
	shell_env = os.environ.copy()
	shell_env["USE_NINJA"] = str(1 if use_ninja else 0)
	cmd = [sys.executable, "setup.py", "install", "--root", "./install"]
	return_code = shell(cmd, cwd=cpp_extensions_test_dir, env=shell_env)

Conversation

shink commented May 24, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented May 24, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/127074

✅ No Failures

Uh oh!

jgong5 commented May 24, 2024

Uh oh!

jgong5 commented May 24, 2024

Uh oh!

shink commented May 24, 2024

Uh oh!

bsochack commented May 27, 2024

Uh oh!

jczaja commented May 27, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shink commented May 28, 2024

Uh oh!

shink commented May 31, 2024

Uh oh!

shink May 31, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bsochack commented May 31, 2024

Uh oh!

shink commented May 31, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jgong5 commented Jun 7, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shink Jun 7, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bsochack commented Jul 2, 2024

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

shink commented May 24, 2024 •

edited

Loading

pytorch-bot Bot commented May 24, 2024 •

edited

Loading

jczaja commented May 27, 2024 •

edited

Loading

shink May 31, 2024 •

edited

Loading

shink commented May 31, 2024 •

edited

Loading

shink Jun 7, 2024 •

edited

Loading