Exception Handling Fix for Dynamically Compiled Code

[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]

@JadenFiotto-Kaufman If you have another idea in mind to fix this I can give it a shot.

[AdamBelfki3]

@lmjantsch can you point your PR towards the dev branch. That's the default place for merging first, thanks!

[lmjantsch]

I created a new pull request (#658) that implements the changes on the dev branch and points to it.

If you would like contributions to be pointed to dev in general, it would be good to update the CONTRIBUTING.md:

Submitting Changes

1. Create a branch from the latest main:

git checkout -b my-feature main

2. Make your changes and add tests.

3. Run the test suite:

pytest tests/test_tiny.py tests/test_lm.py --device cpu -x

4. Push your branch and open a pull request against main.