[Dynamo] Error in speculate_subgraph doesn't report inner user stack trace

[ezyang]

🐛 Describe the bug

self explanatory

Versions

main

cc @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @ydwu4 @aakhundov @wconstab @Xia-Weiwen

[cpuhrsch]

@ezyang - I assume adding the pt2 oncall label is not desired?

[zou3519]

Do you have an example?

[ezyang]

it's just this code

    except Unsupported as ex:
        f_name = f"{type(f).__name__}"
        if isinstance(f, UserFunctionVariable):
            f_name = f.get_name()
        msg = (
            f"speculate_subgraph: while introspecting {description}, we were unable "
            f"to trace function `{f_name}` into a single graph. This means "
            f"that Dynamo was unable to prove safety for this API and will "
            f"fall back to eager-mode PyTorch, which could lead to a slowdown."
        )
        log.warning(msg)
        log.exception(ex)
        tx.output.graph = graph_checkpoint
        tx.restore_graphstate(checkpoint)
        raise Unsupported(
            f"{msg} Scroll up for the stack trace "
            f"of the initial exception. The reason was: {ex.msg}"
        ) from ex

doesn't preserve the user stack trace from ex

[zou3519]

Just to check, are you saying that the stack trace should should directly point to a line from user code?

For example, the following code

import torch
import torch._higher_order_ops.wrap

glob = []


def f(x):
    glob.append(x)
    return x.clone()

@torch.compile(backend='eager', fullgraph=True)
def g(x):
    return torch.ops.higher_order.wrap(f, x)

x = torch.randn(3)
g(x)

currently gives me the the following if I scroll far enough up:

but it should instead have the glob.append(x) somewhere in it?

[zou3519]

Okay, yeah, I understand now. The current error message is the following

but it should really point to the glob.append(...) line (or be something like a nested stack trace)