handle exceptions in converted functions

[lmjantsch]

The Problem

Exceptions in code compiled by convert() (which injects operation wrapping into forward methods using generic <nnsight> filename) crashed ExceptionWrapper._collect_frames() with a KeyError. The exception handler only knew how to map Tracer.Info frames, not convert() frames.

The Solution

Three files were modified:

1. globals.py

Added a metadata registry:

class Globals:
    converted_fn_files = {}  # Maps <nnsight {metadata_key}> filenames to (co_filename, co_firstlineno, co_name)

2. inject.py

Modified convert() to register metadata before compiling:

metadata_key = hash((fn.__code__.co_filename, fn.__code__.co_firstlineno, fn.__code__.co_name))
filename = f"<nnsight {metadata_key}>"
Globals.converted_fn_files[filename] = (
    fn.__code__.co_filename,
    fn.__code__.co_firstlineno,
    fn.__code__.co_name
)

3. util.py

Updated ExceptionWrapper._collect_frames() to check both mappings:

if fname in filename_mapping:
    # Handle Tracer.Info frames (existing logic)
    ...
if fname in Globals.converted_fn_files:
    # Handle convert() frames (new logic)
    co_filename, co_firstlineno, co_name = Globals.converted_fn_files[fname]
    line_number = co_firstlineno - 1 + lineno
    # Reconstruct original location
    ...
else:
    # Unknown frame - show as-is
    frames.append((fname, lineno, name, "", False))

Result

• Exceptions in convert() compiled code no longer crash
• Tracebacks now show real file paths and correct line numbers
• Both Tracer and convert() frames are handled

[lmjantsch]

I am not sure if saving the file metadata to tracing.Globals is the right way to go. If you have a better idea on where to attach the metadata to let me know and I will give it a shot.

[lmjantsch]

@AdamBelfki3 @JadenFiotto-Kaufman It would be nice, if this could make its way in the next version :)