tracer¶
Turn a with block into code we run on our own terms.
The core trick: when the user writes
with Tracer():
...their code...
we don't want the block's body to run normally. Instead we want to capture it,
hand it to a backend, and execute it in a controlled environment (for nnsight,
that means interleaved with a model's forward pass). Tracer pulls this
off with a small dance across the context-manager protocol:
- Capture (
Tracer.capture) — from the frame that ran thewith, read the surrounding source, parse it, and find thewithnode at that line. Compile the block's body into a standalone code object. Results are memoized per site (seennsight.tracing.globals). - Skip (
Tracer.__enter__→skip_context) — install a trace function that raisesExitTracingExceptionthe instant the body is about to run, so the body never executes inline. - Execute (
Tracer.__exit__) — swallow that exception and instead run the captured code through the backend.
A block that contains only pass / a docstring has nothing to intercept, so
step 2 is skipped and it just runs (harmlessly) before the backend is invoked.
module-attribute
¶
TRY_STATEMENTS = tuple(node for node in (ast.Try, getattr(ast, 'TryStar', None)) if node is not None)
Bases: Exception
Control-flow signal: bail out of the with body before it runs inline.
Raised by the installed trace function the moment execution reaches the
block, and caught in Tracer.__exit__, which then runs the captured
code through the backend instead. Not an error — never surfaces to the user.
Bases: Exception
No with block was found at the call site, so there's nothing to trace.
Either the tracer wasn't used as a with block, or the source it was
written in couldn't be read back as one; _not_found_message tells the
two apart for the user.
Tracer(backend: Backend | None = None)
Context manager that captures a with block and runs it via a backend.
See the module docstring for the capture → skip → execute flow. Subclasses
override the small steps (parse, build, compile,
execute) or supply a different Backend
to change what "run the block" means.
Defaults to a plain [`Backend`][nnsight.tracing.backend.Backend], which
just executes it in place.
Serialize the tracer, minus what can't or shouldn't travel.
Drops the backend (transport-side, may hold credentials), the AST node
(only needed locally), and the batcher (per-run state, rebuilt by execute
server-side); Info handles its own frame swap.
Capture the with block at the call site into info.
Looks two frames up (past __enter__) to find the user's frame,
reads and parses its source, locates the with node at that line, and
compiles the block body. The parsed node and compiled code are memoized
per site in BLOCKS, so re-entering the
same block is cheap.
Idempotent: safe to call early (to sniff whether a with block exists)
and again on __enter__.
| RAISES | DESCRIPTION |
|---|---|
WithBlockNotFoundError
|
if the call site isn't a |
staticmethod
¶
Return the source text for filename, reading it at most once.
Cached in SOURCES and never
re-validated, so a file edited mid-run is still traced as it was first
seen (we deliberately skip linecache.checkcache). Handles files,
IPython/Jupyter cells, and python -c programs; see the inline notes.
Find the with/async with node that starts on lineno.
Returns the matching AST node, or None if there's no with
statement there (the call site isn't a trace block).
Parsing the whole file is O(its AST) and dominates a cold capture, so try
slicing just the block out first (_parse_block); parse the whole
file only if the slice can't be isolated cleanly.
Wrap the block's body (not the with line) in a compilable module.
fix_missing_locations backfills line/column info so the synthesized
module can be compiled.
Compile the body module into a code object.
Compiled under the original frame's filename, and its co_name is set
to the frame's so tracebacks read as if the body ran where it was written.
Run the captured body and write its results back to the user's frame.
Executes the block against a copy of the caller's locals, then
push_result writes results back — a nested
trace passes all its locals up to the enclosing block, while the outermost
trace returns only the values marked with nnsight.save.
Subclasses override this to change how the block runs (interleaved with a
forward pass, a backward pass, ...), ending with the same
push_result; the trace-scope depth it reads
is managed by __exit__.
__exit__(exc_type: type[BaseException] | None, exc_value: BaseException | None, exc_traceback: TracebackType | None) -> bool
Run the captured block through the backend, or propagate a real error.
Tears down the trace hook, then — for the expected
ExitTracingException (body was skipped) or a clean exit (nothing
to skip) — invokes the backend. An exception from the backend is
re-raised with a clean_traceback so nnsight
internals are dropped, leaving the user's own frames (across any files).
Any other exception from the body propagates normally (returns False).
skip_context(tracer: Tracer) -> None
Arm the trace function that skips the block body so the backend runs it.
Sets a per-frame trace hook on the user's frame that raises
ExitTracingException as soon as the body is about to execute. The
global sys.settrace is set to a no-op so Python actually delivers
per-frame trace events (it won't call f_trace otherwise).
Python delivers those events per line, so the body has to start on a line of
its own for the hook to reach it first. A body written on the with line
would run where it stands — and then again through the backend — so refuse it
rather than execute it twice.
Which line the body starts on is the hook's cue to fire, rather than the next
line to come along: a with header written over several lines reaches its
own closing line before the body, and raising there would skip the block
before the as target is bound.
Where the raise lands has to be somewhere the with can catch it, which
rules out a body whose first statement is a try — refuse that shape too.
Whether the block has real code worth intercepting.
Returns False for a body made up entirely of pass statements and bare
constant expressions (e.g. a docstring or ...) — nothing to run through
the backend, so we let it execute normally instead of installing the skip
hook. Returns True as soon as any other statement is found.
Add value to this thread's saved set (by identity), unconditionally.
The mechanism behind save, without its in-a-trace guard — for internal
callers that mark values to return outside a running trace (e.g. a remote
backend recording the values a finished request sent back).
Mark value to be returned to the user after the outermost trace.
Marks the concrete object by identity and returns it unchanged, so the idiom
is to save the value you bind: h = model.layer1.output.save(). A value
built from a saved one is not itself saved — (x.save() * 2) returns x,
not the product; write (x * 2).save() instead.
Only meaningful inside a trace — a saved value is what the with
model.trace(...): block hands back — so calling it with no trace running
(e.g. x = nnsight.save([]) before the block) is an error, not a silent
no-op whose mark is cleared before anything reads it.
| PARAMETER | DESCRIPTION |
|---|---|
value
|
The object to mark. Returned unchanged so it can be bound in the same expression that saves it.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Any
|
|
| RAISES | DESCRIPTION |
|---|---|
ValueError
|
If called outside a |
Examples:
>>> with model.trace("The Eiffel Tower is in"):
... hidden = model.transformer.h[0].output.save()
>>> hidden.shape # readable after the block
torch.Size([1, 7, 768])
The function form ``nnsight.save(x)`` is equivalent to ``x.save()``.
Enter a trace scope (increment this thread's nesting depth).
Called around the point a trace actually runs its body — the backend call in
Tracer.__exit__. A nested trace's whole with runs inside that, so
it sees a depth greater than 1 and is treated as inner.
Write a trace body's variables back into frame.
An inner trace pushes everything, so values flow up to the enclosing block.
The outermost trace — depth 1 here, since every nested trace has already left
— pushes only the values marked with save.