Deeply rework WeakIdKeyDictionary #90825

ezyang · 2022-12-14T06:58:04Z

Stack from ghstack (oldest at bottom):

In the prior patch, I just YOLOed a mutable mapping implementation.
Many edge cases were not handled correctly. In this PR, I just
copy paste the WeakKeyDictionary from CPython and the hacked it up
to use WeakIdRef instead of weakref.ref. You can see each line
I changed with the comment CHANGED; there aren't many.

Being exactly API compatible with WeakKeyDictionary means I can also
rob all of the tests from CPython, which I also did for
test/test_weak.py

How to review? You could either try taking the delta from CPython
(recommended), or review everything from scratch (not recommended).
Can post diff representing delta on request.

Signed-off-by: Edward Z. Yang ezyang@fb.com

In the prior patch, I just YOLOed a mutable mapping implementation. Many edge cases were not handled correctly. In this PR, I just copy paste the WeakKeyDictionary from CPython and the hacked it up to use WeakIdRef instead of weakref.ref. You can see each line I changed with the comment CHANGED; there aren't many. Being exactly API compatible with WeakKeyDictionary means I can also rob all of the tests from CPython, which I also did for test/test_weak.py How to review? You could either try taking the delta from CPython (recommended), or review everything from scratch (not recommended). Can post diff representing delta on request. Signed-off-by: Edward Z. Yang <ezyang@fb.com> [ghstack-poisoned]

pytorch-bot · 2022-12-14T06:58:06Z

🔗 Helpful Links

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

📄 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 ea2a4d7:
💚 Looks good so far! There are no failures yet. 💚

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

In the prior patch, I just YOLOed a mutable mapping implementation. Many edge cases were not handled correctly. In this PR, I just copy paste the WeakKeyDictionary from CPython and the hacked it up to use WeakIdRef instead of weakref.ref. You can see each line I changed with the comment CHANGED; there aren't many. Being exactly API compatible with WeakKeyDictionary means I can also rob all of the tests from CPython, which I also did for test/test_weak.py How to review? You could either try taking the delta from CPython (recommended), or review everything from scratch (not recommended). Can post diff representing delta on request. Signed-off-by: Edward Z. Yang <ezyangfb.com> [ghstack-poisoned]

In the prior patch, I just YOLOed a mutable mapping implementation. Many edge cases were not handled correctly. In this PR, I just copy paste the WeakKeyDictionary from CPython and the hacked it up to use WeakIdRef instead of weakref.ref. You can see each line I changed with the comment CHANGED; there aren't many. Being exactly API compatible with WeakKeyDictionary means I can also rob all of the tests from CPython, which I also did for test/test_weak.py How to review? You could either try taking the delta from CPython (recommended), or review everything from scratch (not recommended). Can post diff representing delta on request. Signed-off-by: Edward Z. Yang <ezyangfb.com> ghstack-source-id: d2498245435465c08cce6b9cbd584a4ca01271ae Pull Request resolved: #90825

albanD

This solves #81597 ?

Sounds good!

albanD · 2022-12-14T15:03:41Z

test/test_weak.py

+        dict2 = WeakIdKeyDictionary(dict)
+        self.assertEqual(dict[o], 364)
+
+    def make_weak_keyed_dict(self):


This is not used?

albanD · 2022-12-14T15:10:24Z

test/test_weak.py

+
+# These tests are ported from cpython/Lib/test/test_weakref.py,
+# but adapted to use tensor rather than object
+class WeakTest(TestCase):


There is no test in there that ensures that we only hold a weak reference and not a strong one?
Also any test to ensure we don't introduce ref cycles (which would delay deletion in practice)?

Weak reference is indirectly tested, as we do test that entries are promptly removed from dictionary when they die (this is how I found _fix_weakref leak).

There were some ref cycle tests in cpython that seemed complicated, but I could try to intro them

leaving this as follow up, I think prompt deletion also covers the ref cycle case

eellison

this is not the same as the prior implementation, because the prior implementation was weak on both key and value

eellison · 2022-12-14T16:44:35Z

torch/utils/weak.py

+        del self.data[WeakIdRef(key)]  # CHANGED
+
+    def __getitem__(self, key):
+        return self.data[WeakIdRef(key)]  # CHANGED


we need to do _fix_weakref here. You're calling it elsewhere but if we're going to expose this as a utility we shouldn't require users to call it

I disagree. I never take out a strong reference to the key here, so I don't need to fix weakref

You're right. Was thinking that data was a WeakValueDictionary

eellison · 2022-12-14T16:51:17Z

torch/utils/weak.py

+
+    def pop(self, key, *args):
+        self._dirty_len = True
+        return self.data.pop(WeakIdRef(key), *args)  # CHANGED


1, _fix_weakref, right ?

ezyang · 2022-12-14T20:15:05Z

I don't think so, pretty sure old impl was strong on value

In the prior patch, I just YOLOed a mutable mapping implementation. Many edge cases were not handled correctly. In this PR, I just copy paste the WeakKeyDictionary from CPython and the hacked it up to use WeakIdRef instead of weakref.ref. You can see each line I changed with the comment CHANGED; there aren't many. Being exactly API compatible with WeakKeyDictionary means I can also rob all of the tests from CPython, which I also did for test/test_weak.py How to review? You could either try taking the delta from CPython (recommended), or review everything from scratch (not recommended). Can post diff representing delta on request. Signed-off-by: Edward Z. Yang <ezyangfb.com> [ghstack-poisoned]

ezyang · 2022-12-15T02:24:57Z

@pytorchbot merge -f "known master failure only"

pytorchmergebot · 2022-12-15T02:28:11Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

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

pytorchmergebot · 2022-12-15T02:28:19Z

Merge failed

Reason: Command git -C /home/runner/work/pytorch/pytorch cherry-pick -x b21ccd99a845d056d917a16aae97e44cbe7b2fc3 returned non-zero exit code 1

Auto-merging .lintrunner.toml
Auto-merging torch/_subclasses/fake_tensor.py
CONFLICT (content): Merge conflict in torch/_subclasses/fake_tensor.py
error: could not apply b21ccd99a8... Deeply rework WeakIdKeyDictionary
hint: After resolving the conflicts, mark them with
hint: "git add/rm <pathspec>", then run
hint: "git cherry-pick --continue".
hint: You can instead skip this commit with "git cherry-pick --skip".
hint: To abort and get back to the state before "git cherry-pick",
hint: run "git cherry-pick --abort".

Details for Dev Infra team

Raised by workflow job

In the prior patch, I just YOLOed a mutable mapping implementation. Many edge cases were not handled correctly. In this PR, I just copy paste the WeakKeyDictionary from CPython and the hacked it up to use WeakIdRef instead of weakref.ref. You can see each line I changed with the comment CHANGED; there aren't many. Being exactly API compatible with WeakKeyDictionary means I can also rob all of the tests from CPython, which I also did for test/test_weak.py How to review? You could either try taking the delta from CPython (recommended), or review everything from scratch (not recommended). Can post diff representing delta on request. Signed-off-by: Edward Z. Yang <ezyangfb.com> [ghstack-poisoned]

ezyang · 2022-12-15T05:24:57Z

@pytorchbot merge

In the prior patch, I just YOLOed a mutable mapping implementation. Many edge cases were not handled correctly. In this PR, I just copy paste the WeakKeyDictionary from CPython and the hacked it up to use WeakIdRef instead of weakref.ref. You can see each line I changed with the comment CHANGED; there aren't many. Being exactly API compatible with WeakKeyDictionary means I can also rob all of the tests from CPython, which I also did for test/test_weak.py How to review? You could either try taking the delta from CPython (recommended), or review everything from scratch (not recommended). Can post diff representing delta on request. Signed-off-by: Edward Z. Yang <ezyangfb.com> ghstack-source-id: 32b861c2462c3a686269882d4c6dae77ac5f0083 Pull Request resolved: #90825

pytorchmergebot · 2022-12-15T05:26:31Z

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

pytorchmergebot · 2022-12-15T08:23:11Z

Merge failed

Reason: Command git -C /home/runner/work/pytorch/pytorch cherry-pick -x b21ccd99a845d056d917a16aae97e44cbe7b2fc3 returned non-zero exit code 1

Auto-merging .lintrunner.toml
Auto-merging torch/_subclasses/fake_tensor.py
CONFLICT (content): Merge conflict in torch/_subclasses/fake_tensor.py
error: could not apply b21ccd99a8... Deeply rework WeakIdKeyDictionary
hint: After resolving the conflicts, mark them with
hint: "git add/rm <pathspec>", then run
hint: "git cherry-pick --continue".
hint: You can instead skip this commit with "git cherry-pick --skip".
hint: To abort and get back to the state before "git cherry-pick",
hint: run "git cherry-pick --abort".

Details for Dev Infra team

Raised by workflow job

In the prior patch, I just YOLOed a mutable mapping implementation. Many edge cases were not handled correctly. In this PR, I just copy paste the WeakKeyDictionary from CPython and the hacked it up to use WeakIdRef instead of weakref.ref. You can see each line I changed with the comment CHANGED; there aren't many. Being exactly API compatible with WeakKeyDictionary means I can also rob all of the tests from CPython, which I also did for test/test_weak.py How to review? You could either try taking the delta from CPython (recommended), or review everything from scratch (not recommended). Can post diff representing delta on request. Signed-off-by: Edward Z. Yang <ezyangfb.com> [ghstack-poisoned]

In the prior patch, I just YOLOed a mutable mapping implementation. Many edge cases were not handled correctly. In this PR, I just copy paste the WeakKeyDictionary from CPython and the hacked it up to use WeakIdRef instead of weakref.ref. You can see each line I changed with the comment CHANGED; there aren't many. Being exactly API compatible with WeakKeyDictionary means I can also rob all of the tests from CPython, which I also did for test/test_weak.py How to review? You could either try taking the delta from CPython (recommended), or review everything from scratch (not recommended). Can post diff representing delta on request. Signed-off-by: Edward Z. Yang <ezyangfb.com> ghstack-source-id: a2c7d6869804fc139233b8dd03e5974e5728adb0 Pull Request resolved: #90825

ezyang · 2022-12-15T08:38:45Z

@pytorchbot merge -f "idk doesn't look conflicting to me"

pytorchmergebot · 2022-12-15T08:43:01Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

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

danthe3rd · 2022-12-19T17:50:39Z

torch/utils/weak.py

+        if hasattr(r, '_fix_weakref'):
+            r._fix_weakref()  # type: ignore[union-attr]


This class seems to still keep a strong reference to the object. Repro:

import torch import gc import weakref x = torch.zeros([1]) y = torch.zeros([1]) refx = weakref.ref(x) refy = weakref.ref(y) WeakIdRef(x) == WeakIdRef(torch.ones([1])) del x del y gc.collect() print("refX: ", type(refx())) # refX: <class 'torch.Tensor'> print("refY: ", type(refy())) # refY: <class 'NoneType'>

Commenting out "r._fix_weakref()" here fixes it

Do you have the first PR in this stack? There was a memory leak that I fixed

yeah on master, fixing your example, it works

(/Users/ezyang/Dev/pytorch-cpu-env) ezyang-mbp:pytorch-cpu ezyang$ cat foo.py import torch import gc import weakref from torch.utils.weak import WeakIdRef x = torch.zeros([1]) y = torch.zeros([1]) refx = weakref.ref(x) refy = weakref.ref(y) WeakIdRef(x) == WeakIdRef(torch.ones([1])) del x del y gc.collect() print("refX: ", type(refx())) # refX: <class 'torch.Tensor'> print("refY: ", type(refy())) # refY: <class 'NoneType'> (/Users/ezyang/Dev/pytorch-cpu-env) ezyang-mbp:pytorch-cpu ezyang$ python foo.py refX: <class 'NoneType'> refY: <class 'NoneType'>

Oh right, thanks a lot! I just imported the weak.py file (as I wanted it to work with older versions of pytorch as well). I'll look into the other PR. Thanks!

For old versions just get rid of the fix weakref call, it's super busted lol. I suppose if you are willing to write a little C extension code, manually decreffing the Tensor after fixing weakref would work too

This was referenced Dec 14, 2022

Fix _fix_weakref memory leak #90823

Closed

Don't put tracing state on Tensor #90628

Closed

github-actions bot added the ciflow/inductor label Dec 14, 2022

github-actions bot requested review from albanD, anjali411, antoniojkim, bdhirsh, Chillee, miladm, SherlockNoMad, voznesenskym and wconstab December 14, 2022 06:58

ezyang added release notes: composability release notes category topic: new features topic category labels Dec 14, 2022

ezyang requested a review from eellison December 14, 2022 07:04

ezyang added the ciflow/trunk Trigger trunk jobs on your pull request label Dec 14, 2022

albanD approved these changes Dec 14, 2022

View reviewed changes

eellison reviewed Dec 14, 2022

View reviewed changes

ezyang mentioned this pull request Feb 1, 2023

Umbrella issue for weakref related Dynamo PyTorch test suite failures #93481

Closed

pytorchmergebot added the Merged label Dec 15, 2022

pytorchmergebot closed this in 0b22f5a Dec 15, 2022

danthe3rd reviewed Dec 19, 2022

View reviewed changes

facebook-github-bot deleted the gh/ezyang/1655/head branch June 8, 2023 16:39

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Deeply rework WeakIdKeyDictionary #90825

Deeply rework WeakIdKeyDictionary #90825

ezyang commented Dec 14, 2022 •

edited

Loading

pytorch-bot bot commented Dec 14, 2022 •

edited

Loading

albanD left a comment

albanD Dec 14, 2022

albanD Dec 14, 2022

ezyang Dec 14, 2022

ezyang Dec 14, 2022

eellison left a comment

eellison Dec 14, 2022

ezyang Dec 14, 2022

eellison Dec 14, 2022

eellison Dec 14, 2022

ezyang commented Dec 14, 2022

ezyang commented Dec 15, 2022

pytorchmergebot commented Dec 15, 2022

pytorchmergebot commented Dec 15, 2022

ezyang commented Dec 15, 2022

pytorchmergebot commented Dec 15, 2022

pytorchmergebot commented Dec 15, 2022

ezyang commented Dec 15, 2022

pytorchmergebot commented Dec 15, 2022

danthe3rd Dec 19, 2022 •

edited

Loading

ezyang Dec 19, 2022

ezyang Dec 19, 2022

danthe3rd Dec 20, 2022

ezyang Dec 20, 2022

		if hasattr(r, '_fix_weakref'):
		r._fix_weakref() # type: ignore[union-attr]

Deeply rework WeakIdKeyDictionary #90825

Deeply rework WeakIdKeyDictionary #90825

Conversation

ezyang commented Dec 14, 2022 • edited Loading

pytorch-bot bot commented Dec 14, 2022 • edited Loading

🔗 Helpful Links

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

✅ No Failures

albanD left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eellison left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ezyang commented Dec 14, 2022

ezyang commented Dec 15, 2022

pytorchmergebot commented Dec 15, 2022

Merge started

pytorchmergebot commented Dec 15, 2022

Merge failed

ezyang commented Dec 15, 2022

pytorchmergebot commented Dec 15, 2022

Merge started

pytorchmergebot commented Dec 15, 2022

Merge failed

ezyang commented Dec 15, 2022

pytorchmergebot commented Dec 15, 2022

Merge started

danthe3rd Dec 19, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ezyang commented Dec 14, 2022 •

edited

Loading

pytorch-bot bot commented Dec 14, 2022 •

edited

Loading

danthe3rd Dec 19, 2022 •

edited

Loading