registration¶
Interventions that stay on the engine and run for every request.
A trace ships its block to the worker on the request it rides, so a sweep of a thousand prompts serializes and deserializes the same block a thousand times, and only requests that are nnsight traces are touched at all.
A registration inverts that. The block goes to the worker once and stays there;
from then on every request the engine runs gets its own copy of it — including
requests submitted by something that has never heard of nnsight, an OpenAI-API
client on the same server. Each copy keeps its own scope, so what it saves is
that request's own, and it comes back on that request's RequestOutput —
output.saves — by the same collect a traced value uses, then is dropped.
Example
model = VLLM("meta-llama/Llama-3.1-8B", dispatch=True, ... enable_prefix_caching=False)
with model.edit() as (tracer, edit): # doctest: +SKIP ... hidden = model.model.layers[16].output[0].save()
outputs = model.vllm_entrypoint.generate(prompts, sampling) # doctest: +SKIP outputs[5].saves["hidden"].shape # doctest: +SKIP
edit.clear() # doctest: +SKIP
The block is written exactly like a trace body — the same envoy tree, the same
.save(). What it cannot do is anything that belongs to one particular
request: there is no prompt to invoke, so tracer.invoke(...) has no meaning
here, and the block applies to whatever the engine happens to run.
The engine has to be built with enable_prefix_caching=False. A prefix-cached
token is served without a forward pass, so no hook fires for it and the block
sees fewer rows than the prompt has, with nothing to say so. A trace asks for its
own request to be recomputed; a registration rides requests it did not create and
cannot, so the cache has to be off at the engine — registering against one that
has it on warns.
A handle on a block the engine is running for every request.
Returned by VLLM.edit and live
until clear.
There is nothing to read here. What a registered block saves comes back on
the RequestOutput of the request it ran on, as output.saves, by the
same collect that carries a traced value — so the values arrive where the
request that produced them already is, and are dropped as they go. Nothing
accumulates for as long as somebody is reading the outputs, which on the
synchronous engine is every request there is.
A plain with is synchronous-engine only: installing the block is a
collective_rpc, which on mode="async" is a coroutine there is no safe
way to finish from a plain statement, so the exit raises there rather than
silently not installing it. On an async engine use async with and
aclear, which can await that trip from inside the loop it is already on.
| ATTRIBUTE | DESCRIPTION |
|---|---|
model |
The engine this is registered on.
|
id |
The engine-wide id for this registration, used to address it in the worker.
|
name |
The name it was installed under, or
|
ServeRegistration(model: 'VLLM', id: str, host: str, api_key: str | None = None, name: str | None = None)
Bases: Registration
A block installed on an nnsight-serve engine.
The same handle, over HTTP: the client has no engine to collective_rpc
into — it holds a meta model and no weights — so the block goes to the
server, which installs it on the engine it holds. What it saves rides the
request it ran on, which for a serve client means tracer.result.saves of
a trace sent to that same server.
RegisteringTracer(model: 'VLLM', *, backend: Backend | None = None, serve: str | None = None, api_key: str | None = None, name: str | None = None)
Bases: VLLMTracer
Capture a block and leave it on the engine instead of running it once.
The counterpart of EditingTracer
for a runtime whose model lives in another process: an edit is replayed by
the envoy that stores it, which on vLLM would leave it in the client, where
there are no weights. This sends the block across instead, and hands back a
Registration to switch it
off again with.
__enter__() -> tuple['RegisteringTracer', Registration]
Enter the block, binding the tracer and the handle that ends it.
Both, the way Envoy.edit binds
(tracer, edited): a registered block is still a trace body, and the
tracer is what carries iter/all, without which the block could
only ever see a request's first forward — its prefill — and never the
steps it generates::
with model.edit() as (tracer, edit):
readouts = nnsight.save([])
for step in tracer.all():
readouts.append(model.model.layers[16].output[0][-1])
async
¶
__aenter__() -> tuple['RegisteringTracer', Registration]
__enter__, for async with — the form an async engine needs.
Spelled out rather than delegating to __enter__: both capture and the
skip guard read the caller's frame at a fixed depth, so going through
another call would hand them this method's frame — capture would find no
with block at all, and skip_context would arm the wrong frame and
let the body run here as well as on the workers.
async
¶
Capture as usual, then await the send.
The block is prepared by execute either way; only the trip to the
workers differs, and on an async engine that trip has to be awaited from
inside the loop.
Prepare the block for the workers rather than running it here.
Only prepared: the send happens in __exit__/__aexit__, because
whether it can be awaited depends on which of the two ran.
It goes to every rank, so each builds the same per-request copies and their reads stay in step — which is what a sharded model's gathers need.
copy is left off: the worker builds a fresh mediator per request from
this template, so each request already has a scope of its own, and the
block's saves have to land in it for the collect to find them.