Reference cycles involving code -> co_extra -> compiled output -> reference to code

[ezyang]

🐛 Describe the bug

By default, libraries which hook into the frame evaluation API are encouraged to store compiled products on co_extra. However, storing arbitrary Python objects on code objects is a recipe for reference cycles, because most people assume that holding a reference to a code object is "safe" (unlikely to result in a cycle); e.g., that it is safe to store a CapturedTraceback on a compiled code object. I recently ran into this on the stack at #107457 where I wanted to store debugging tracebacks on internal data structures in PT2; I only noticed because we had two bugs in GC traverse/clear and this resulted in code getting leaked entirely.

I'm not sure how to 100% solve this problem.

• Stop storing compiled products in co_extra, store them somewhere else. However, if you torch.compile a global function that takes only Tensor inputs, there is no where else to store it, the code object really is your only logical choice.
• Stop storing so much stuff in our compiled products. This might be plausible but we'll need to take a close look at our object graph on the compiled code and see what can be dropped and what can't
• Something else?

cc @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov

Versions

main

[colesbury]

Code objects don't support cyclic garbage collection but they do support weak references:

https://github.com/python/cpython/blob/9bb576cb07b42f34fd882b8502642b024f771c62/Objects/codeobject.c#L2121-L2160

So you might be able to use something like weakref.WeakKeyDictionary with the code object as the key and whatever we stick in co_extra as the value. I'd expect lookups to be slower than the current approach.

EDIT: Nevermind, I don't think the above works. Our compiled products would still keep the code object alive preventing collection.

[Fidget-Spinner]

One possible alternative is to implement weakrefs for compiled products and have code objects be referenced by weakref.

So the cycle becomes code -> co_extra -> compiled output --weakref--> code.

That way, no "strong" reference cycles are formed, and the refcounting mechanism can easily handle this.

[Fidget-Spinner]

The weakref way might not be ideal either. As it's still easy to create a strong cycle. E.g. if the compiled output -> external object -> code

I think a possible way forward would be to create a global weakref.WeakKeyDictionary mapping code objects to their cached compiled outputs.

[ezyang]

Can you spell it out for me why this works? Naively I tend to treat weak key map and storing an object directly on the key as equivalent

[Fidget-Spinner]

Woops I forgot to mention, I'm proposing to move all the compiled artefacts out of co_extra, and place them in the global weak key map mapping code objects to the compiled cache.

Naively I tend to treat weak key map and storing an object directly on the key as equivalent

I'm not too sure I understand the second part of your question, so I'll answer both of my interpretations:

1. If you meant using a normal dictionary, we shouldn't do that here as being in the dictionary key will hold a strong reference to the code object. So code objects will always be kept alive by our map.
2. If you meant directly storing a borrowed reference (so a pointer without a Py_INCREF), yes that is functionally equivalent to a weak key map. However, borrowed references are generally unsafe. The weakref is a thin wrapper over the actual borrowed reference, so it returns None if the object is evaporated, while accessing a borrowed reference could just segfault.

The assumption here is that nothing except the global/module state must hold a reference to this weak map. Assuming nothing holds a reference to the global/module state (to be fair if you do that, then an unclearable cycle is unavoidable no matter what we do), there is no way to have a reference cycle back to this weak key map. Thus the references will be in the linear fashion:

global/module state -> weak key map -> cached entries

The one downside of this is that cache lookups will be slightly slower as indexing a weak key map is slower than directly accessing co_extra (one needs hashing, the other is just a pointer + offset). However, IMO that is better than leaking memory.

[ezyang]

No I don't understand. Previously:

code object -> compiled code -> traceback -> code object (cycle)

Under your proposal, we have:

ROOT -> global code map -> compiled code -> traceback -> code object

But if the traceback is a string reference to code object, that will prevent us from deallocating the code object (since it's live via the global code map, which is a root) and then you'll never clear the entry from the global code map.

[Fidget-Spinner]

ROOT -> global code map -> compiled code -> traceback -> code object

I don't fully understand the reference the compiled code holds on the traceback (I will take a look soon). Does it call Py_DECREF on it? If so that would be enough to trigger it to GC properly. If the compiled code never decrefs the strong reference it has on the traceback, then the code object will leak.

(since it's live via the global code map, which is a root) and then you'll never clear the entry from the global code map.

It's live via a weak reference. Which does not impact reference counting nor keeping it alive in GC. If you have a weak reference to something and its refcount drops to zero or participates in GC, it will be cleared. "A weak reference to an object is not enough to keep the object alive: when the only remaining references to a referent are weak references, garbage collection is free to destroy the referent and reuse its memory for something else.". - (https://docs.python.org/3/library/weakref.html)

[ezyang]

@zdevito's fast traceback format involves avoiding formatting the stack trace, and instead just holding a strong reference to the code object so we can extract out the filename later. The cycle occurs when I accidentally have some tracebacks stored for debugging that somehow are referenced by the compiled product (I accidentally did this, and then I undid it to unblock--so I don't think this happens today but it's easy to accidentally cause.) So in the cycle case, the compiled code validly holds a strong ref to the code object, and is NOT clearing it.