util

util

Helpers for the tracer: frame serialization, traceback cleanup, and writing variables back into a live frame.

These support two jobs that the Tracer can't do inline:

• Serialization — a real Python frame can't be pickled, so SerializedFrame stands in for one when a trace is shipped to a remote worker, keeping just the code metadata needed for tracebacks.
• Traceback surgery — user intervention code runs deep inside the model's forward pass, so exceptions come back wrapped in nnsight frames. clean_traceback strips nnsight internals away so the error points at the user's own code (across whatever files it spans).
• Frame write-back — push copies the results of an executed block back into the frame the with statement lived in, so assignments made inside the block are visible afterwards.
• Block scope — Scope is the namespace a captured block runs in, and the reason a block can read the names around it even though it runs later, and somewhere else, than the line it was written on.

Scope

Scope(own: dict, shared: dict, glbls: dict)

Bases: dict

The namespace a captured block runs in.

A block runs later than the line it was written on — and, once interleaving, somewhere else entirely — so the names it reaches have to come from three places, in order:

1. This dict, a snapshot of the frame's locals taken when the block was captured. A name the block only reads means what it meant where the block was written: a for prompt in prompts: variable has moved on (or gone) by the time the block runs, so the snapshot is what keeps it right.
2. shared — the store the blocks written in that frame share (shared_locals). A name bound by a sibling block (an earlier tracer.invoke(...)) isn't in the snapshot, because nothing had bound it when this block was captured. Blocks written in one frame share these; blocks written in different functions have different frames, so they stay as separate as their code.
3. glbls — the frame's globals, reached by fallback rather than copied.

Passing this as exec's globals (not just its locals) is what lets a lambda or nested def inside a block reach the block's own names: their free variables compile to LOAD_GLOBAL, which never consults a locals mapping — but does honor __missing__ on a dict subclass.

Writes land here and in shared, so the block sees its own assignments and so do the blocks written beside it — but not the frame, which only push writes to. Because only the snapshot and the block's writes are ever stored, iterating a scope yields the block's own names — what push and the block reducers want — and never the whole global namespace.

shared instance-attribute

shared = shared

glbls instance-attribute

glbls = glbls

__missing__

__missing__(key: str) -> Any

__setitem__

__setitem__(key: str, value: Any) -> None

copy

copy() -> 'Scope'

A scope with this one's names, sharing the same frame and globals.

SerializedFrame

SerializedFrame(co_filename: str, co_firstlineno: int, co_name: str)

A picklable stand-in for a real frame.

Carries only the code metadata (filename/lineno/name) needed remotely for tracebacks and source lookup; f_locals/f_globals are empty because the replay closure already captured the real ones.

f_locals instance-attribute

f_locals: dict = {}

f_globals instance-attribute

f_globals: dict = {}

f_code instance-attribute

f_code = types.SimpleNamespace(co_filename=co_filename, co_firstlineno=co_firstlineno, co_name=co_name)

of classmethod

of(frame: FrameType) -> SerializedFrame

Build a serializable stand-in from a live frame.

filter_traceback

filter_traceback(traceback: TracebackType | None, keep: Callable[[FrameType], bool]) -> TracebackType | None

Rebuild a traceback keeping only frames for which keep(frame) is true.

Tracebacks are immutable and singly linked, so we can't delete entries in place — we collect the ones to keep and re-link a fresh chain (from the tail up, since each TracebackType points to the next frame).

PARAMETER	DESCRIPTION
traceback	The head of the traceback chain to filter (may be None). TYPE: TracebackType | None
keep	Predicate called with each frame; the entry is kept when it returns true. TYPE: Callable[[FrameType], bool]

RETURNS	DESCRIPTION
TracebackType | None	The head of the filtered chain, or None if nothing was kept.

clean_traceback

clean_traceback(traceback: TracebackType | None) -> TracebackType | None

Drop nnsight-internal frames, leaving only the user's frames.

Used as the fallback when we can't pin the error to the traced block's own file: at least hide the library's plumbing. In debug mode (CONFIG.APP.DEBUG) nothing is stripped, so the full stack — nnsight internals included — is shown.

push

push(frame: FrameType, variables: dict[str, Any]) -> None

Write variables back into a live frame's locals.

The traced block runs in a scratch namespace, not the original frame, so its results have to be copied back for assignments to be visible after the with statement.

On Python < 3.13, frame.f_locals is a snapshot and updating it doesn't reach the interpreter's fast-locals array, so we call the C-API PyFrame_LocalsToFast to flush the change through. From 3.13 on, f_locals is a live write-through mapping (PEP 667) and the plain update is enough.

This is the only route from a block's namespace back into the frame — see shared_locals, which keeps a block's other writes out of it.

PARAMETER	DESCRIPTION
frame	The frame to write into (the one the with block lived in), or a SerializedFrame standing in for one, whose f_locals is an ordinary dict. TYPE: FrameType
variables	Names and values to set on that frame. TYPE: dict[str, Any]