[dynamo] delete dynamo cache entry when guard function is invalidated #117875

williamwen42 · 2024-01-19T19:51:47Z

Stack from ghstack (oldest at bottom):

-> [dynamo] delete dynamo cache entry when guard function is invalidated #117875

Fixes #112090.

Summary of changes:

Changed CacheEntry linked list into a doubly-linked list structure to support deletion.
Added CacheEntry weakref to GuardFn so that GuardFn can tell CacheEntry to delete itself when the GuardFn is invalidated.
Added ExtraState (borrowed) reference to CacheEntry so that we can get ExtraState to correctly point to the first CacheEntry if it gets deleted.
ExtraState destructor needs to nullify CacheEntry pointers to it, in order to prevent use-after-free.
code_context values that are nn.GraphModules need to be weakrefs in order to prevent circular references.
Added tests that check for memory leaks and cache deletion operations.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng

[ghstack-poisoned]

pytorch-bot · 2024-01-19T19:51:50Z

🔗 Helpful Links

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

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 7efa508 with merge base 01abb5a ():

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

pull / linux-focal-py3_8-clang9-xla / test (xla, 1, 1, linux.12xlarge, unstable) (gh)
Process completed with exit code 128.

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

…invalidated" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

ghstack-source-id: 2af0ddfb91528aeaef067bf40ae8f599783d173b Pull Request resolved: #117875

…invalidated" Fixes #112090. Summary of changes: - Changed CacheEntry linked list into a doubly-linked list structure to support deletion. - Added CacheEntry weakref to GuardFn so that GuardFn can tell CacheEntry to delete itself when the GuardFn is invalidated. - Added ExtraState (borrowed) reference to CacheEntry so that ExtraState properly points to the first CacheEntry, especially when a CacheEntry is deleted. - ExtraState destructor needs to nullify CacheEntry pointers to it in order to prevent use-after-free. - code_context values that are nn.GraphModules need to be weakrefs in order to prevent circular references. - Added tests that check for memory leaks and cache deletion operations. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

anijain2305 · 2024-01-22T20:26:09Z

torch/csrc/dynamo/eval_frame.c

        e->next = old_cache_entry;
+        e->prev = (CacheEntry*)Py_None;
+        if (old_cache_entry != (CacheEntry*)Py_None) {


I think old_cache_entry will never be Py_None. So , maybe add a DEBUG_CHECK.

anijain2305 · 2024-01-22T20:28:24Z

torch/csrc/dynamo/eval_frame.c

+    Py_XDECREF(check_fn_cache_entry);
+    #endif
+
+    // Use a weakref to avoid circular references


Add a comment on how this would be used to setup finalizer in guards.py, point to function invalidate

torch/csrc/dynamo/eval_frame.c

anijain2305 · 2024-01-22T20:35:29Z

torch/csrc/dynamo/eval_frame.c

+  if (first_cache_entry != NULL && first_cache_entry == self) {
+    DEBUG_CHECK((PyObject*)self->prev == Py_None);
+    // will be decremented later on, so increment now
+    Py_INCREF(self->next);


Curious: Why do we need to INCREF and DECREF self->next here?

torch/_dynamo/guards.py

anijain2305 · 2024-01-22T20:38:08Z

torch/_dynamo/guards.py

+        if (
+            self.valid
+            and hasattr(self, "check_fn")
+            and self.check_fn is not DeletedGuardFn


if self.valid is already checked, why do we need DeletedGuardFn?

We want to clean up check_fn and prevent it from being used again. DeletedGuardFn could be replaced with None - I think that using a class here makes it more clear (for debugging/error logging purposes) that we intentionally removed the reference to the actual guard function.

torch/csrc/dynamo/eval_frame.c

…invalidated" Fixes #112090. Summary of changes: - Changed CacheEntry linked list into a doubly-linked list structure to support deletion. - Added CacheEntry weakref to GuardFn so that GuardFn can tell CacheEntry to delete itself when the GuardFn is invalidated. - Added ExtraState (borrowed) reference to CacheEntry so that we can get ExtraState to correctly point to the first CacheEntry if it gets deleted. - ExtraState destructor needs to nullify CacheEntry pointers to it in order to prevent use-after-free. - code_context values that are nn.GraphModules need to be weakrefs in order to prevent circular references. - Added tests that check for memory leaks and cache deletion operations. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

ghstack-source-id: c76825e05b5a08679ca1ce4fcf5d19a8c01a5ec0 Pull Request resolved: #117875

ezyang · 2024-01-23T15:00:02Z

Letting Animesh captain this review, but let me know if there's anything specific you want me to look at

jansel

I'd suggest one of two things:

Remove cache entries lazily if .valid is false (and move .valid to C, allowing us to delete check_fn and code right away to free almost all the memory)
Move some of this code to C++ so we can use THObject to manage reference counts.

jansel · 2024-01-23T17:04:46Z

torch/_dynamo/guards.py

-        # list.
+        if (
+            self.valid
+            and hasattr(self, "check_fn")


When will this be false?

jansel · 2024-01-23T17:11:16Z

torch/csrc/dynamo/eval_frame.c

@@ -337,6 +367,50 @@ static struct PyGetSetDef CacheEntry_properties[] = {
    {NULL}};


+static CacheEntry* extract_cache_entry(ExtraState* extra_state);
+
+static PyObject* CacheEntry_invalidate(CacheEntry* self, PyObject* args) {


How are we validating all the reference counting logic here? I worry we will be leaking somewhere. Could we refactor things such that the ownership was more clear?

williamwen42 · 2024-01-24T22:22:48Z

@jansel I tried 1. locally by making check_fn in CacheEntry a weakref, but this doesn't work because check_fn can be deleted without going through CheckFunctionManager.invalidate. It seems CacheEntry needs to have an owning reference to check_fn, so we can't take a lazy approach.

jansel · 2024-01-25T17:13:53Z

@williamwen42 you misunderstand me. Both check_fn and the code a strong refs, but get cleared by invalidate():

pytorch/torch/_dynamo/guards.py

Lines 1192 to 1197 in 817debe

    
           def invalidate(self): 
        
               # A weakref is no longer valid, self.check_fn should return false 
        
               # TODO(janimesh) - Free up cache entry after the cache entry formation 
        
               # is in python, and the underlying data structure is a doubly linked 
        
               # list. 
        
               self.valid = False

For this to work the .valid check would need to be moved into C++, so check_fn would never get called if .valid is false and it can be safely deleted.

williamwen42 · 2024-02-03T00:19:44Z

Abandoning due to attempt 2

…ard function is invalidated [attempt 2]" Attempt #2 for #117875 to fix #112090. Summary of changes: - ~Changed CacheEntry linked list into a doubly-linked list structure to support deletion.~ (done by C++ refactor) - Added CacheEntry and ExtraState borrowed references to GuardFn so that GuardFn can tell ExtraState to delete CacheEntry when the GuardFn is invalidated. - ~Added ExtraState raw reference to CacheEntry so that we can get ExtraState to correctly point to the first CacheEntry if it gets deleted.~ (done by C++ refactor) - CacheEntry destructor needs to reset GuardFn refs to ExtraState/CacheEntry in order to prevent use-after-free. - code_context values that are nn.GraphModules need to be weakrefs in order to prevent circular references. - Added tests that check for memory leaks and cache deletion operations. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]

…invalidated [attempt 2]" Attempt #2 for #117875 to fix #112090. Summary of changes: - ~Changed CacheEntry linked list into a doubly-linked list structure to support deletion.~ (done by C++ refactor) - Added CacheEntry and ExtraState borrowed references to GuardFn so that GuardFn can tell ExtraState to delete CacheEntry when the GuardFn is invalidated. - ~Added ExtraState raw reference to CacheEntry so that we can get ExtraState to correctly point to the first CacheEntry if it gets deleted.~ (done by C++ refactor) - CacheEntry destructor needs to reset GuardFn refs to ExtraState/CacheEntry in order to prevent use-after-free. - code_context values that are nn.GraphModules need to be weakrefs in order to prevent circular references. - Added tests that check for memory leaks and cache deletion operations. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]