mirror of
https://github.com/pytorch/pytorch.git
synced 2026-10-25 18:53:39 +08:00
Ryan Murray
78589bea5b
Summary: # Author Notes Our training logs showed `torch.jit` deprecation warnings that all pointed at lines inside `torch/jit/`, so we couldn't tell which of our code was calling TorchScript. This PR adds `stacklevel=2` so each warning names the line that called it, and separate call sites no longer collapse into one message. # Agent Notes > The TorchScript deprecation `FutureWarning`s (made visible in #189914) don't pass `stacklevel`, so Python attributes every one of them to the `warnings.warn` line inside `torch/jit/` rather than to the user's call site, e.g. > > ``` > torch/jit/_script.py:1485: FutureWarning: `torch.jit.script` is deprecated. ... > warnings.warn( > ``` > > Two consequences: > - Users can't tell which of their call sites needs migrating. > - Because the default filter dedups per location, every distinct caller in a process collapses into a single message. > > This adds `stacklevel=2` to all 17 such sites in `torch/jit/`: > - `_script.py`: `script`, `script_method`, `interface` > - `_trace.py`: `trace`, `trace_module` > - `_serialization.py`: `save`, `load` > - `_freeze.py`: `freeze`, `optimize_for_inference` > - `_async.py`: `fork`, `wait` > > All of these are public entry points called directly (none are wrapped by another decorator), so frame 2 is the user's code. This matches the existing warnings in these files that already pass `stacklevel=2`, such as the `optimize=` warnings in `_script.py` and `_trace.py`. > > Side effect: processes with many distinct callers now print one warning per call site instead of one in total. That is intended for an actionable deprecation. Internal re-entrant calls (e.g. `torch.jit.trace`'s `check_trace` path calling `torch.jit.trace` again at `_trace.py:375`) now show up as their own location too. > > This PR was authored with the assistance of Claude (an AI assistant). Test Plan: Ran this script (`demo.py`) against a build with and without this change, with `warnings.simplefilter("default")`: ```python import warnings import torch def main() -> None: warnings.simplefilter("default") def f(x: torch.Tensor) -> torch.Tensor: return x + 1 def g(x: torch.Tensor) -> torch.Tensor: return x * 2 torch.jit.script(f) # line 15 torch.jit.script(g) # line 16 class M(torch.jit.ScriptModule): @torch.jit.script_method # line 19 def forward(self, x: torch.Tensor) -> torch.Tensor: return x @torch.jit.interface # line 23 class I(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: pass torch.jit.trace(f, torch.ones(1)) # line 28 fut = torch.jit.fork(f, torch.ones(1)) # line 30 torch.jit.wait(fut) # line 31 main() ``` **Before** (run without the `fork`/`wait` lines): every warning points inside `torch/jit/`, and the two `torch.jit.script` calls on lines 15 and 16 are deduplicated into one: ``` torch/jit/_script.py:1485: FutureWarning: `torch.jit.script` is not supported in Python 3.14+ and may break. ... warnings.warn( torch/jit/_script.py:359: FutureWarning: `torch.jit.script_method` is not supported in Python 3.14+ and may break. ... warnings.warn( torch/jit/_script.py:1643: FutureWarning: `torch.jit.interface` is deprecated. Please use `torch.compile` instead. warnings.warn( torch/jit/_trace.py:1000: FutureWarning: `torch.jit.trace` is not supported in Python 3.14+ and may break. ... warnings.warn( ``` **After**: each warning points at the caller's line, and lines 15 and 16 each get their own warning: ``` demo.py:15: FutureWarning: `torch.jit.script` is not supported in Python 3.14+ and may break. Please switch to `torch.compile` or `torch.export`. torch.jit.script(f) demo.py:16: FutureWarning: `torch.jit.script` is not supported in Python 3.14+ and may break. Please switch to `torch.compile` or `torch.export`. torch.jit.script(g) demo.py:19: FutureWarning: `torch.jit.script_method` is not supported in Python 3.14+ and may break. Please switch to `torch.compile` or `torch.export`. @torch.jit.script_method demo.py:23: FutureWarning: `torch.jit.interface` is deprecated. Please use `torch.compile` instead. @torch.jit.interface demo.py:28: FutureWarning: `torch.jit.trace` is not supported in Python 3.14+ and may break. Please switch to `torch.compile` or `torch.export`. torch.jit.trace(f, torch.ones(1)) torch/jit/_trace.py:375: FutureWarning: `torch.jit.trace` is not supported in Python 3.14+ and may break. Please switch to `torch.compile` or `torch.export`. check_mod = torch.jit.trace( demo.py:30: FutureWarning: `torch.jit.fork` is deprecated. Please use `torch.compile` instead. fut = torch.jit.fork(f, torch.ones(1)) demo.py:31: FutureWarning: `torch.jit.wait` is deprecated. Please use `torch.compile` instead. torch.jit.wait(fut) ``` The `_trace.py:375` line is `torch.jit.trace`'s own `check_trace` path calling `torch.jit.trace` again, so it is correctly attributed to that internal call site. Existing `torch.jit` tests run in CI. Pull Request resolved: https://github.com/pytorch/pytorch/pull/200204 Approved by: https://github.com/janeyx99