[dynamo] raise hard error if error is encountered while tracing resume function prologue (#154564)

This should prevent bad resume function prologues from slipping by. In particular, graph breaks in resume function prologues will now hard error.

Implementation details:
- The resume function prologue is surrounded by `LOAD_CONST arg, STORE_FAST __is_tracing_resume_prologue` instructions. The first sequence has `arg=True` and the second sequence has `arg=False`.
- InstructionTranslator will know when it is tracing a resume function prologue when it detects `STORE_FAST __is_tracing_resume_prologue`. The top of stack will be True to mark the start of the prologue, False to mark the end.
- When `convert_frame.py` detects that an error occurred while the InstructionTranslator was tracing a resume function prologue, we will wrap the exception and hard error

Pull Request resolved: https://github.com/pytorch/pytorch/pull/154564
Approved by: https://github.com/jansel
ghstack dependencies: #154283, #154289, #154782, #155166
This commit is contained in:
William Wen
2025-06-17 17:39:46 -07:00
committed by PyTorch MergeBot
parent 614a415145
commit a6a3a44144
5 changed files with 107 additions and 13 deletions

View File

@ -70,6 +70,10 @@ class InternalTorchDynamoError(TorchDynamoException):
pass
class ResumePrologueTracingError(TorchDynamoException):
pass
class RestartAnalysis(TorchDynamoException):
restart_reason: Optional[str]