source

source

Source-level operation tracing.

Architecture

Each nnsight-wrapped module has one global :class:SourceAccessor that is lazily built on the first .source access (by any Envoy / Interleaver / Mediator). The SourceAccessor is cached on the module itself as module.__source_accessor__ so it survives any later replacement of module.forward (e.g. torch.compile, accelerate hot-swap). It owns:

• The injected version of the module's forward — its AST has been rewritten so every call site is wrapped by a wrap(fn, name=...) lookup that consults the per-operation hook state.
• A dict {op_name: OperationAccessor} — one entry per call site.

Each :class:OperationAccessor is also global per (module, op): it owns the hook lists (pre_hooks, post_hooks, fn_hooks, fn_replacement) and, for recursive source tracing, a nested :class:SourceAccessor for the operation's own fn.

Per-Envoy wrappers — :class:SourceEnvoy and :class:OperationEnvoy — sit on top of the accessors and provide the user-facing API (eproperties for .input/.output, pretty-printed .source). Multiple Envoys or Interleavers wrapping the same module share the same underlying accessors; only the per-Envoy wrappers are duplicated.

Forward routing

nnsight_forward (installed by :meth:Interleaver.wrap_module) checks module.__source_accessor__ on each call:

• If a SourceAccessor exists and .hooked is True (any OperationAccessor under it has any active hook), it invokes source_accessor(module, *args, **kwargs) which runs the injected forward.
• Otherwise it calls the unwrapped original via module.__nnsight_forward__ — zero overhead for modules that nobody is source-tracing.

Lifetimes

• SourceAccessor and OperationAccessor live as long as the module's nnsight_forward wrapper does (effectively the lifetime of the model).
• SourceEnvoy / OperationEnvoy live as long as their owning Envoy.
• Hooks on an OperationAccessor are one-shot and self-remove when they fire; they are also tracked on mediator.hooks so session cleanup drains them.
• fn_replacement is one-shot too: cleared after :func:wrap_operation runs once. Re-accessing .source on an OperationEnvoy reinstalls it.

FunctionCallWrapper

FunctionCallWrapper(name: str)

Bases: NodeTransformer

name_index instance-attribute

name_index = defaultdict(int)

line_numbers instance-attribute

line_numbers = {}

name instance-attribute

name = name

get_name

get_name(node: Call)

Extract and index function name from a Call node.

visit_Call

visit_Call(node)

visit_FunctionDef

visit_FunctionDef(node)

visit_AsyncFunctionDef

visit_AsyncFunctionDef(node)

OperationAccessor

OperationAccessor(name: str, source: str, line_number: int)

Global hook state for a single call site inside a module's forward.

Exactly one instance per (module, op) pair, owned by the parent :class:SourceAccessor. Hook registrations from any Envoy / Interleaver / Mediator that touches this operation land on the same accessor — by design, so multiple consumers can coexist over the module's lifetime.

Hook lists (read live by :func:wrap_operation at call time):

• pre_hooks — appended by :func:operation_input_hook. Each is called with (args, kwargs); non-None return replaces them.
• post_hooks — appended by :func:operation_output_hook. Each is called with the return value; non-None return replaces it.
• fn_hooks — appended by :func:operation_fn_hook for recursive source tracing. Each receives the current fn and returns a (possibly replaced) fn.
• fn_replacement — a one-shot fn replacement installed by :attr:OperationEnvoy.source. When set, :func:wrap_operation uses it in place of the original fn for one call, then clears it.

All input/output/fn hooks are one-shot and self-remove when they fire. The hooked property is True if any list is non-empty; :class:SourceAccessor.wrap checks it to take the zero-overhead fast path for unhooked sites.

PARAMETER	DESCRIPTION
name	Fully-qualified path of the operation (e.g. "model.transformer.h.0.attn.split_1"). TYPE: str
source	Source code of the enclosing module's forward (used by __str__ for pretty-printing). TYPE: str
line_number	Line number of the operation in source. TYPE: int

path instance-attribute

path = name

source_code instance-attribute

source_code = source

line_number instance-attribute

line_number = line_number

pre_hooks instance-attribute

pre_hooks: List[Callable] = []

post_hooks instance-attribute

post_hooks: List[Callable] = []

fn_hooks instance-attribute

fn_hooks: List[Callable] = []

fn_replacement instance-attribute

fn_replacement: Optional[Callable] = None

hooked property

hooked: bool

True if the op has any active hook or a pending fn replacement.

__str__

__str__()

SourceAccessor

SourceAccessor(fn: Callable, path: str)

Global injected-forward + operation accessor map for one fn.

Built once on first .source access for a module (or for an operation's fn, in the recursive case) and cached. Subsequent accesses — even from different Envoys / Interleavers — reuse the same accessor.

The injected forward is not written onto the module. Instead, nnsight_forward (installed by :meth:Interleaver.wrap_module) branches on module.__source_accessor__: if present, it calls the accessor; otherwise it calls the original __nnsight_forward__ directly. This keeps the non-source path zero-overhead.

PARAMETER	DESCRIPTION
fn	Unwrapped original function whose AST should be rewritten. For modules this is found by :func:resolve_true_forward; for recursive source it is the op's own fn (delivered via :func:operation_fn_hook). TYPE: Callable
path	Dotted prefix for operation names (e.g. "model.transformer.h.0.attn"). TYPE: str

path instance-attribute

path = path

source instance-attribute

source = source

line_numbers instance-attribute

line_numbers = line_numbers

operations instance-attribute

operations: Dict[str, OperationAccessor] = {}

hooked property

hooked: bool

True if any OperationAccessor under this SourceAccessor is hooked.

Provided for introspection; not load-bearing on the forward routing path. nnsight_forward routes through the SourceAccessor whenever it exists (regardless of hooked), since per-op hooks may register mid-forward and the injected wrap closure already has a per-op fast path for unhooked sites.

wrap

wrap(fn: Callable, **kwargs) -> Callable

Per-call-site dispatcher baked into the injected forward.

Fast path: return fn unchanged when the op's accessor has no hooks. Lazy path: build a wrapper via :func:wrap_operation that runs the hook lists at call time.

rebind

rebind(fn: Callable) -> None

Re-inject against fn while preserving OperationAccessor state.

Called by :meth:Envoy._update when a module is replaced (typically on dispatch — meta-tensor module swapped for the loaded one). The new fn is expected to share the source code of the old (same class), so operation names line up; their hook lists, fn_replacement, and nested SourceAccessors are kept intact, so any pre-existing OperationEnvoy / SourceEnvoy references remain valid.

__call__

__call__(*args, **kwargs)

Invoke the injected forward.

Called by nnsight_forward when hooked is True. The first positional argument should be the module (self of the original forward method), since the injected fn is unbound.

__iter__

__iter__()

Yield each operation's short name, deduplicated by line number.

__str__

__str__()

Pretty-print the source with operation names at their call sites.

OperationEnvoy

OperationEnvoy(accessor: OperationAccessor, interleaver: Optional[Interleaver] = None)

Per-Envoy proxy for a single call site.

Implements :class:IEnvoy so it can back eproperty descriptors for .output, .input, .inputs, and .source (recursive source tracing). All hook state lives on the underlying :class:OperationAccessor; this class is a thin per-Envoy view that routes hook registration to the shared accessor.

PARAMETER	DESCRIPTION
accessor	The shared :class:OperationAccessor that owns this operation's hook state. Multiple OperationEnvoys (from different parent Envoys) may share the same accessor. TYPE: OperationAccessor
interleaver	The :class:Interleaver whose current mediator should receive the registered hook callbacks. TYPE: Optional[Interleaver] DEFAULT: None

accessor instance-attribute

accessor = accessor

path instance-attribute

path = path

interleaver instance-attribute

interleaver = interleaver

source property

source: 'SourceEnvoy'

Get the source of this operation's fn for recursive source tracing.

On first access for the underlying op (any Envoy), builds a nested :class:SourceAccessor on the :class:OperationAccessor and uses the fn-hook + swap dance to substitute the injected fn into the currently-running operation. The injected fn is also installed as fn_replacement so the operation wrapper picks it up.

Because fn_replacement is one-shot (see :func:wrap_operation), subsequent .source accesses re-install it from the cached nested accessor.

__str__

__str__()

output

output() -> Any

Get the output of this operation.

Examples:

>>> with model.trace("Hello World"):
...     attn = model.transformer.h[0].attn.source.attention_interface_0.output.save()

inputs

inputs() -> Tuple[Tuple[Any, ...], Dict[str, Any]]

Get the inputs to this operation as (args, kwargs).

input

input(value)

SourceEnvoy

SourceEnvoy(accessor: SourceAccessor, interleaver: Optional[Interleaver] = None)

Per-Envoy user-facing wrapper around a :class:SourceAccessor.

Provides named attribute access to :class:OperationEnvoy instances (one per call site) and pretty-prints by delegating to the accessor. Multiple SourceEnvoys may wrap the same accessor — they share hook state via the underlying :class:OperationAccessors.

PARAMETER	DESCRIPTION
accessor	The shared :class:SourceAccessor for the underlying module. Several SourceEnvoys may wrap the same accessor. TYPE: SourceAccessor
interleaver	The :class:Interleaver whose current mediator will receive operation-level hook callbacks. TYPE: Optional[Interleaver] DEFAULT: None

accessor instance-attribute

accessor = accessor

path instance-attribute

path = path

interleaver instance-attribute

interleaver = interleaver

operations instance-attribute

operations: List[OperationEnvoy] = []

__str__

__str__()

__iter__

__iter__()

__getattribute__

__getattribute__(name: str) -> Union[OperationEnvoy, Any]

convert

convert(fn: Callable, wrap: Callable, name: str)

Rewrite fn's AST so every call site is wrapped by wrap(fn, name=...).

Returns the source string, a {op_name: line_number} map, and the compiled & executed wrapped function.

wrap_operation

wrap_operation(fn: Callable, name: str, bound_obj: Optional[Any] = None, op_accessor: Optional['OperationAccessor'] = None) -> Callable

Wrap fn so it processes the hook lists on op_accessor.

Installed by :meth:SourceAccessor.wrap only when the OperationAccessor has at least one active hook (otherwise the call site uses fn directly). Hook lists are read live at call time so hooks registered after wrapper creation are still seen.

fn_replacement is one-shot: cleared after the operation completes so subsequent forward passes fall back to the original fn unless a new replacement is installed (typically by re-accessing .source on the matching :class:OperationEnvoy).

resolve_true_forward

resolve_true_forward(module: Module) -> Callable

Find the unwrapped fn whose AST should be injected.

A module's forward may have been wrapped by accelerate (module.forward = partial(new_forward, module), which calls module._old_forward(*args, **kwargs)) or by nnsight (module.forward = nnsight_forward, which calls module.__nnsight_forward__(module, *args, **kwargs)). In both cases the true forward — the user's actual compute — lives one level deeper.

• Accelerate: module._old_forward (often a bound method).
• nnsight: module.__nnsight_forward__ (set by :meth:Interleaver.wrap_module).
• Plain module: type(module).forward.

Returns an unbound function suitable for re-execution with the module as the first positional argument (which is how the injected fn is called by :class:SourceAccessor.__call__).

get_or_create_source_accessor

get_or_create_source_accessor(module: Module) -> SourceAccessor

Return the module's :class:SourceAccessor, building it on first access.

The accessor is cached on module.__source_accessor__ directly so it survives any replacement of module.forward (e.g. by torch.compile, accelerate's hot-swap, or other wrappers that re-bind forward after nnsight has wrapped the module). Subsequent calls — even from different Envoys / Interleavers / Mediators — return the same instance.