PR #180636: [dynamo] restore Python dispatch TLS across graph breaks

• Repository: pytorch/pytorch
• Author: jansel
• State: open | merged: False
• Link: https://github.com/pytorch/pytorch/pull/180636

Description (problem / solution / changelog)

Stack from ghstack (oldest at bottom):

• -> #180636

Fix a flaky functorch+dynamo failure where Python fallback can be entered without a valid TLS snapshot, or with Python dispatch TLS leaked past the compiled frame. Use MaybeSetTLSOnEntryGuard in Python fallback, and restore the saved dispatch TLS around the compiled call so later eager ops like assertEqual/isclose/module setup do not trip the same failure after a graph break.

Add a regression test covering the graph-break path that previously left Python=True and PythonTLSSnapshot=False in local dispatch TLS.

This appears related to the flaky test disables in #180299, #180300, #180320, #180321, #180322, #180336, #180337, and likely the same failure family as #180436, #180555, #180590, #180591, and #180592. Though I could not repro any of them locally.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @kadeng @chauhang @amjames @Lucaskabela @jataylo @azahed98

Changed files

• aten/src/ATen/core/PythonFallbackKernel.cpp (modified, +1/-1)
• test/dynamo/test_error_messages.py (modified, +35/-4)
• test/dynamo/test_python_dispatcher.py (modified, +29/-0)
• torch/_dynamo/eval_frame.py (modified, +46/-39)