Scope DeepCompile compiler state to graph and engine lifecycles

[tohtana]

DeepCompile uses several process-global PyTorch compiler mechanisms while building graph-specific state. Scheduled recompilation also replaces previously compiled graphs.

These lifetimes did not have explicit ownership boundaries.

• Forward inputs could be consumed by a different graph
• Dynamo frame IDs could collide across engines
• Cleanup could either leave graph-specific patches installed or release shared compiler state that another engine still needed

This PR gives those existing mechanisms explicit graph and engine lifecycle ownership.

Fixes:

• Keep forward inputs in a graph-local one-shot queue and InputStorage.
• Qualify compiled-backward frames with an owner token and track them per engine.
• Keep the torch.autograd.Function patch active until the last owned backward frame is released, then clear captured backward inputs.
• Reference-count the ZeRO-3 Dynamo configuration overrides across engines and restore the original values after the last owner exits.
• Limit the graph-specific AotAutograd.__init__ patch to the Inductor compilation call and restore it in finally.
• Release engine-owned compiler state during scheduled recompilation, deactivation, and destruction.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8f8f69327a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[eternalNight]

The changes look good to me.

One question about the assumptions of use. This PR seems relevant to use cases where multiple engines are created within a single process. Do we assume that the forward of two engines do not overlap? If they do, the patched torch.autograd.Function will be active when another model has completed compilation, probably leading to recompilation issues similar to what #7980 had seen.

I have not seen any use case running multiple engines in an overlapping manner, though. So the question is just out of curiosity.

[tohtana]

Hi @eternalNight,
Thank you for your comment! Good point, I think your concern is valid.

This PR only makes captured backward inputs and cleanup owner-specific, but the torch.autograd.Function patch is still process-global. This can cause the same kind of repeated recompilation that #7980 fixed. I haven’t seen a real workload relying on that pattern either, but it is a real limitation.

To address the issue, I think we could limit the behavior to the owning engine/frame. For example, only resetting compiled_bw while that engine's frame_key is still pending. Since we are not aware of a workload that interleaves multiple engines in this way, I would prefer to avoid adding it in this PR and revisit it if such a use case arises.