Make exceptions from user code more readable

[apaz-cli]

Code:

import thunder

def f1():
    raise ValueError

def f2():
    f1()

def f3():
    f2()

jfn = thunder.jit(f3)

jfn()

After:

⚡ ap-readable_exceptions ~/lightning-thunder python ti.py                       
Traceback (most recent call last):
  File "/teamspace/studios/this_studio/lightning-thunder/ti.py", line 14, in <module>
    jfn()
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/__init__.py", line 849, in _thunder_unwrap_inner_exception
    raise exc
  File "/teamspace/studios/this_studio/lightning-thunder/ti.py", line 10, in f3
    f2()
  File "/teamspace/studios/this_studio/lightning-thunder/ti.py", line 7, in f2
    f1()
  File "/teamspace/studios/this_studio/lightning-thunder/ti.py", line 4, in f1
    raise ValueError
ValueError

Before:

⚡ main ~/lightning-thunder python ti.py
Traceback (most recent call last):
  File "/teamspace/studios/this_studio/lightning-thunder/ti.py", line 14, in <module>
    jfn()
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/__init__.py", line 774, in wrapped
    return fn(*args, **kwargs)
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/__init__.py", line 824, in fn_
    cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/__init__.py", line 756, in wrapped
    cache_entry, inps, pro_to_epi = get_computation_and_inputs_fn(*args, **kwargs)
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/core/langctxs.py", line 136, in _fn
    result = fn(*args, **kwargs)
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/__init__.py", line 236, in cache_info_wrapper
    res = fn(*args, **kwargs)
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/__init__.py", line 529, in get_computation_and_inputs
    jit_results: TraceResults = thunder_general_jit(
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/core/jit_ext.py", line 1744, in thunder_general_jit
    result = jfn(*args, **kwargs)
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/core/interpreter.py", line 7231, in fn_
    raise e
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/core/interpreter.py", line 7191, in fn_2
    return fn(*args, **kwargs)
  File "/teamspace/studios/this_studio/lightning-thunder/ti.py", line 10, in f3
    f2()
  File "/teamspace/studios/this_studio/lightning-thunder/ti.py", line 7, in f2
    f1()
  File "/teamspace/studios/this_studio/lightning-thunder/ti.py", line 4, in f1
    raise ValueError
ValueError

[t-vi]

Thank you @apaz-cli , the general idea seems OK, but we would

• need a flag that gives the full traceback (register DebugOptions.interpreter_traceback),
• need a test for both full traceback and reduced..
  Could you add these please?
  Also the code comments.

[t-vi]

• Why just fn_ and fn_2 and not all of the methods in the interpreter?
• Please don't use the filename/code object here. Instead check whether that is in f_globals of the frame has __name__ == 'thunder.core.interpreter'.

If we drop all interpreter methods, we likely want to keep internal frames at the top of the stacktrace.

[apaz-cli]

Excluding everything in thunder.core.interpreter sounds good. I was trying to be conservative by default.

But using f_globals to do that instead of co_name doesn't work, because not every function has a f_globals with a __name__ in it. Certainly it would not work with some uses of the exec() lookaside. It also depends on how exactly the module is imported. There are a bunch of reasons why it doesn't work.

As for the functions filtered out by code object, those live outside of interpreter.py, in files that I don't want to exclude because some of their functions could be called by user code. And using the code objects is very convenient because it exactly matches the functions to exclude, and we wouldn't want to exclude all of thunder.__init__.

[t-vi]

Which functions we want to filter does not have an f_globals with __name__?

[apaz-cli]

thunder_general_jit() does not. When I add an assertion that f_globals has a __name__, it fails.

⚡ ap-readable_exceptions ~/lightning-thunder python exctest.py 
Traceback (most recent call last):
  File "/teamspace/studios/this_studio/lightning-thunder/exctest.py", line 14, in <module>
    jfn()
  File "/teamspace/studios/this_studio/lightning-thunder/thunder/__init__.py", line 836, in _thunder_unwrap_inner_exception
    assert hasattr(tb.tb_frame.f_globals, "__name__"), tb.tb_frame.f_code.co_name
AssertionError: thunder_general_jit

[t-vi]

That's because f_globals is a mapping, so it's members, not attributes.

[t-vi]

I wonder if this could live in the interpreter instead.

[apaz-cli]

The function can. The unwrapping actually can't, tried that already. Traceback frames are pushed as the stack unwinds, not when the exception is thrown.

[apaz-cli]

@t-vi I can't think of any situation where I would want to see the full traceback. This change is only triggered in situations where the exception comes from user code inside the interpreter.

If there's an interpreter bug where the path is relevant, this is not triggered. If there's a bug, then the bug is in their code, and not the interpreter, so the path through the interpreter is known to be irrelevant.

[t-vi]

Can you please add the flag to get the full traceback? Thank you.