Support with statement on torch.Stream #140138

guangyey · 2024-11-08T13:59:33Z

Stack from ghstack (oldest at bottom):

-> Support with statement on torch.Stream #140138

Motivation

We propose to support Python with statement on torch.Stream. This is a benefit for all accelerators when writing device-agnostic code. The device-specific stream will also be supported because they are generally derived from torch.Stream.

With this PR, we can do like this

s1= torch.Stream()
# Set s1 to the current stream
torch.accelerator.set_stream(s1)
with torch.Stream() as s2:
    # Inside with statement, we set s2 to the current stream
    assert torch.accelerator.current_stream() == s2
# Here the current stream should be s1
assert torch.accelerator.current_stream() == s1

cc @albanD @EikanWang

pytorch-bot · 2024-11-08T13:59:37Z

🔗 Helpful Links

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

📄 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.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

Failures because of deprecated version of actions/download-artifact: v3

❌ 1 New Failure, 1 Unrelated Failure

As of commit fb3f110 with merge base 3beb700 ():

NEW FAILURE - The following job has failed:

xpu / linux-jammy-xpu-2025.0-py3.9 / test (default, 4, 4, linux.idc.xpu) (gh)
inductor/test_aot_inductor.py::AOTInductorTestABICompatibleGpu::test_fft_c2c_xpu

FLAKY - The following job failed but was likely due to flakiness present on trunk:

xpu / linux-jammy-xpu-2025.0-py3.9 / test (default, 3, 4, linux.idc.xpu) (gh) (similar failure)
inductor/test_pattern_matcher.py::TestPatternMatcher::test_duplicate_search

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

ghstack-source-id: cd25eb6 Pull Request resolved: #140138

[ghstack-poisoned]

ghstack-source-id: 4c39e55 Pull Request resolved: #140138

[ghstack-poisoned]

ghstack-source-id: c51a38d Pull Request resolved: #140138 g# modified: test/test_accelerator.py

albanD

This would still leak if the second raises an error, here is an updated version.

albanD · 2025-01-06T14:51:44Z

torch/csrc/Stream.cpp

+  PyObject* ctx_stream = nullptr;
+  if (PyDict_GetItemStringRef(self->context, "_ctx_stream", &ctx_stream) < 0) {
+    throw python_error();
+  }
+  TORCH_CHECK(ctx_stream, "ctx_stream should be initialized.");
+  PyObject* ctx_device_index = nullptr;
+  if (PyDict_GetItemStringRef(
+          self->context, "_ctx_device_index", &ctx_device_index) < 0) {
+    throw python_error();
+  }
+  TORCH_CHECK(ctx_device_index, "ctx_device_index should be initialized.");
+  auto prev_stream = (THPStream*)(THPObjectPtr(ctx_stream).get());
+  auto prev_device_index =
+      THPUtils_unpackDeviceIndex(THPObjectPtr(ctx_device_index).get());


Suggested change

PyObject* ctx_stream = nullptr;

if (PyDict_GetItemStringRef(self->context, "_ctx_stream", &ctx_stream) < 0) {

throw python_error();

}

TORCH_CHECK(ctx_stream, "ctx_stream should be initialized.");

PyObject* ctx_device_index = nullptr;

if (PyDict_GetItemStringRef(

self->context, "_ctx_device_index", &ctx_device_index) < 0) {

throw python_error();

}

TORCH_CHECK(ctx_device_index, "ctx_device_index should be initialized.");

auto prev_stream = (THPStream*)(THPObjectPtr(ctx_stream).get());

auto prev_device_index =

THPUtils_unpackDeviceIndex(THPObjectPtr(ctx_device_index).get());

THPObjectPtr ctx_stream;

if (PyDict_GetItemStringRef(self->context, "_ctx_stream", &ctx_stream.get()) < 0) {

throw python_error();

}

TORCH_INTERNAL_ERROR(ctx_stream.get(), "ctx_stream should be present on the context dict.");

THPObjectPtr ctx_device_index;

if (PyDict_GetItemStringRef(

self->context, "_ctx_device_index", &ctx_device_index.get()) < 0) {

throw python_error();

}

TORCH_CHECK(ctx_device_index.get(), "ctx_device_index should be present on the context dict.");

auto prev_stream = (THPStream*)(ctx_stream.get());

auto prev_device_index =

THPUtils_unpackDeviceIndex(ctx_device_index.get());

Good catch! Thanks~

ctx_stream.get() will raise an error lvalue required as unary & operand because it returns a rvalue. So I changed a minor code to use a lvalue py_stream and passed its reference to PyDict_GetItemStringRef.

ghstack-source-id: e1e2c55 Pull Request resolved: #140138 g# modified: test/test_accelerator.py

[ghstack-poisoned]

ghstack-source-id: cfbd12a Pull Request resolved: #140138 g# modified: test/test_accelerator.py

[ghstack-poisoned]

ghstack-source-id: 50ef4dc Pull Request resolved: #140138 g# modified: test/test_accelerator.py

[ghstack-poisoned]

albanD

One minor nit on the clearing, sounds good otherwise!
Thanks for your patience working through cpython api fun!

albanD · 2025-01-09T14:06:40Z

torch/csrc/Stream.cpp

+  Py_DECREF(self->context);
+  self->context = nullptr;


Suggested change

Py_DECREF(self->context);

self->context = nullptr;

Py_CLEAR(self->context);

Very interesting read about that in https://github.com/python/cpython/blob/ea39c8b08d8f025273bfa5b7a95f7b5984dc1e86/Include/refcount.h#L416

Thanks for your suggestions. I am now more familiar with CPython:)

ghstack-source-id: ac1c365 Pull Request resolved: #140138 g# modified: test/test_accelerator.py

[ghstack-poisoned]

guangyey · 2025-01-10T01:57:11Z

"So happy to try to land this PR, the failure is irrelevant."
@pytorchbot merge -i

pytorchmergebot · 2025-01-10T01:59:16Z

Merge started

Your change will be merged while ignoring the following 2 checks: xpu / linux-jammy-xpu-2025.0-py3.9 / test (default, 3, 4, linux.idc.xpu), xpu / linux-jammy-xpu-2025.0-py3.9 / test (default, 4, 4, linux.idc.xpu)

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

# Motivation In #137678, we help use the device-agnostic APIs to generalize distributed module. As this [comment](#137678 (comment)) said, we will use the with statement of `torch.Stream` once #140138 is landed. Pull Request resolved: #144951 Approved by: https://github.com/kwen2501, https://github.com/albanD

guangyey mentioned this pull request Nov 8, 2024

[RELAND] Add UTs for accelerator device-agnostic runtime APIs #133572

Closed

This was referenced Nov 8, 2024

[RELAND] Add device-agnostic runtime Device/Stream C++ API #138677

Closed

[WIP] Clean up accelerator APIs #137493

Closed

guangyey added a commit that referenced this pull request Nov 8, 2024

Support with statement on torch.Stream

eb95c61

ghstack-source-id: cd25eb6 Pull Request resolved: #140138

guangyey changed the title ~~Support with statement on torch.Stream~~ [WIP] Support with statement on torch.Stream Nov 8, 2024

guangyey marked this pull request as draft November 8, 2024 14:00

pytorchbot added the open source label Nov 8, 2024

Update

e5d2ddd

[ghstack-poisoned]

zhangxiaoli73 mentioned this pull request Nov 11, 2024

Use device-agnostic runtime API in distributed DDP/FSDP instead of cuda device specific. #137678

Closed

guangyey added ciflow/xpu Run XPU CI tasks release notes: python_frontend python frontend release notes category labels Nov 28, 2024

guangyey added a commit that referenced this pull request Nov 28, 2024

Support with statement on torch.Stream

b04ae23

ghstack-source-id: 4c39e55 Pull Request resolved: #140138

guangyey added 4 commits November 28, 2024 13:16

Update

544e78a

[ghstack-poisoned]

Update

c2843e0

[ghstack-poisoned]

Update

7636b68

[ghstack-poisoned]

Update

92c1122

[ghstack-poisoned]

guangyey changed the title ~~[WIP] Support with statement on torch.Stream~~ Support with statement on torch.Stream Nov 29, 2024

guangyey marked this pull request as ready for review November 29, 2024 07:55

guangyey requested review from EikanWang, eqy, gujinghui and syed-ahmed as code owners November 29, 2024 07:55

guangyey added 2 commits November 29, 2024 09:39

Update

7bc60c1

[ghstack-poisoned]

Update

217207e

[ghstack-poisoned]

guangyey requested a review from albanD January 6, 2025 10:51

guangyey added a commit that referenced this pull request Jan 6, 2025

Support with statement on torch.Stream

4872263

ghstack-source-id: c51a38d Pull Request resolved: #140138 g# modified: test/test_accelerator.py

albanD reviewed Jan 6, 2025

View reviewed changes

guangyey added a commit that referenced this pull request Jan 6, 2025

Support with statement on torch.Stream

ae71d5a

ghstack-source-id: e1e2c55 Pull Request resolved: #140138 g# modified: test/test_accelerator.py

guangyey requested a review from albanD January 6, 2025 17:51

guangyey added 3 commits January 6, 2025 18:12

Update

7cb45e3

[ghstack-poisoned]

Update

96ef9ca

[ghstack-poisoned]

Update

9511a38

[ghstack-poisoned]

guangyey added a commit that referenced this pull request Jan 7, 2025

Support with statement on torch.Stream

37be73d

ghstack-source-id: cfbd12a Pull Request resolved: #140138 g# modified: test/test_accelerator.py

guangyey added 3 commits January 7, 2025 09:23

Update

679716c

[ghstack-poisoned]

Update

5694a35

[ghstack-poisoned]

Update

760941a

[ghstack-poisoned]

guangyey added a commit that referenced this pull request Jan 8, 2025

Support with statement on torch.Stream

2d74f7a

ghstack-source-id: 50ef4dc Pull Request resolved: #140138 g# modified: test/test_accelerator.py

Update

24353b9

[ghstack-poisoned]

albanD approved these changes Jan 9, 2025

View reviewed changes

guangyey added a commit that referenced this pull request Jan 9, 2025

Support with statement on torch.Stream

53da844

ghstack-source-id: ac1c365 Pull Request resolved: #140138 g# modified: test/test_accelerator.py

Update

fb3f110

[ghstack-poisoned]

pytorchmergebot added the merging label Jan 10, 2025

pytorchmergebot added the Merged label Jan 10, 2025

pytorchmergebot closed this in 6de110b Jan 10, 2025

pytorchmergebot removed the merging label Jan 10, 2025

guangyey mentioned this pull request Jan 16, 2025

Use torch with statement in torch distributed module #144951

Closed

github-actions bot deleted the gh/guangyey/90/head branch February 12, 2025 02:05

yao-matrix mentioned this pull request Mar 12, 2025

enable OffloadedCache on XPU from PyTorch 2.7 huggingface/transformers#36654

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Support with statement on torch.Stream #140138

Support with statement on torch.Stream #140138

Uh oh!

guangyey commented Nov 8, 2024 •

edited by pytorch-bot bot

Loading

Uh oh!

pytorch-bot bot commented Nov 8, 2024 •

edited

Loading

Uh oh!

albanD left a comment

Uh oh!

albanD Jan 6, 2025

Uh oh!

guangyey Jan 6, 2025

Uh oh!

guangyey Jan 7, 2025 •

edited

Loading

Uh oh!

albanD left a comment

Uh oh!

albanD Jan 9, 2025

Uh oh!

guangyey Jan 9, 2025

Uh oh!

guangyey commented Jan 10, 2025

Uh oh!

pytorchmergebot commented Jan 10, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	Py_DECREF(self->context);
	self->context = nullptr;
	Py_CLEAR(self->context);

Support with statement on torch.Stream #140138

Support with statement on torch.Stream #140138

Uh oh!

Conversation

guangyey commented Nov 8, 2024 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Uh oh!

pytorch-bot bot commented Nov 8, 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/140138

❗ 1 Active SEVs

❌ 1 New Failure, 1 Unrelated Failure

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

albanD Jan 6, 2025

Choose a reason for hiding this comment

Uh oh!

guangyey Jan 6, 2025

Choose a reason for hiding this comment

Uh oh!

guangyey Jan 7, 2025 • 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!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

albanD Jan 9, 2025

Choose a reason for hiding this comment

Uh oh!

guangyey Jan 9, 2025

Choose a reason for hiding this comment

Uh oh!

guangyey commented Jan 10, 2025

Uh oh!

pytorchmergebot commented Jan 10, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

guangyey commented Nov 8, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Nov 8, 2024 •

edited

Loading

guangyey Jan 7, 2025 •

edited

Loading