Fix flaky NameError when CompilePackage.install() reuses a global name (#191128)
Summary:
Fixes #186865
Fixes #190664
`output_graph` installs generated functions -- and, on every compile, a builtins-dict global (`install_builtins_dict_in_fglobals`) -- into the frame globals via `CleanupHook.create`, registering the hook in `CleanupManager` keyed on the transformed code object so the global is dropped once that code object is collected.
`CompilePackage` serializes those names deterministically, so a package installed after `torch._dynamo.reset()` binds exactly the names an earlier compile in the same process already used. The pre-reset code object is garbage by that point but not necessarily collected; when it finally is, its hook can delete a binding `install()` is now responsible for, so the next call raises `NameError: name '__compiled_fn_N_<uuid>' is not defined`. Because this hinges on collection timing it surfaces as a flake, and it is likelier on free-threaded builds where deferred refcounting delays collection.
An earlier version of this fix (thanks jansel for reviewing) compared the hook's recorded value against the current binding and skipped the delete on a mismatch. That covers `_install_global`'s writes (`__compiled_fn`/`__resume_at`, always fresh objects), but not `install()`'s builtins-dict path: that name is deliberately left untouched when the existing value already matches (same dict object every compile in a module), so the stale hook's value comparison still succeeded and it deleted a binding nothing had rewritten. `CleanupHook` now tracks ownership with a token instead of a value comparison, and `CompilePackage` explicitly disowns a name's previous hook at every point it takes over that name -- in `_install_global`, and in the builtins-dict path even when it decides no write is needed -- so a stale hook is inert regardless of whether the value it guarded ever changed. This also closes a latent `KeyError` in `CompilePackage.uninstall()`, whose `module.__dict__.pop(name)` had no default and could run after a stale hook had already removed the name; it now passes a default as well, as defense in depth.
On lifetimes: previously the hook held a strong reference to the value it installed; the token design drops that entirely, so there is no change to what stays alive relative to before this fix existed.
### Test plan
The new test pins the pre-reset code objects so their hooks are guaranteed to fire after `install()`, turning the race into a deterministic failure. Its function graph-breaks on a data-dependent branch, and the checked names include the builtins-dict global, so it covers `__compiled_fn`, `__resume_at`, and `__builtins_dict__` -- all three ways `install()` can take over a name:
```
python test/dynamo/test_package.py TestPackage.test_install_survives_stale_cleanup_hooks
```
It fails without the change (including on an earlier version of this fix that only compared values) and passes with it, on CPython 3.12 and on free-threaded CPython 3.14. Full file:
```
python test/dynamo/test_package.py
```
and `test_misc.py` give the same results with and without the change apart from the new test (the pre-existing failures are local toolchain/environment limitations -- e.g. no C++ compiler for inductor codegen -- not regressions).
One caveat worth stating plainly: I did not catch the original flake in the act. Looping the reported test 200x in-process produced no failures, so the race was demonstrated by forcing collection at the critical moment, which reproduces the exact `NameError` from the issue. The mechanism is proven; attributing the specific CI flake to it is inference from the matching signature.
X-link: https://github.com/pytorch/pytorch/pull/191128
Approved by: https://github.com/jansel
Reviewed By: huydhn
Differential Revision: D114415845
fbshipit-source-id: 609c6ed42775735412d8cc75c22e362bb5b26010
Author
generatedunixname499836121