nnsight.intervention

exception nnsight.intervention.interleaver.Cancelation

Exception raised when a request is canceled.

exception nnsight.intervention.interleaver.EarlyStopException

Exception raised to stop the execution of the model.

class nnsight.intervention.interleaver.Events(value)

Enum for different types of events in the interleaving process.

class nnsight.intervention.interleaver.Interleaver(invokers: List[Mediator] = None, tracer: InterleavingTracer = None, batcher: Batcher = None, user_cache: Cache | None = None)

Manages the interleaving of model execution and interventions.

This class coordinates the flow between the model’s forward pass and user-defined intervention functions, allowing for inspection and modification of intermediate values.

cancel()

Cancel all intervention threads.

check_cache_full()

Print a warning if a module to be cached was missed.

property current: Mediator

Get the current mediator.

handle(provider: Any | None = None, value: Any | None = None, iterate: bool = False)

Handle a provider’s value, allowing mediators to consume and modify it.

Parameters:

• provider – The identifier of the provider

• value – The value being provided

Returns:

The original or modified value

wrap_operation(fn: Callable, name: str, bound_obj: Any | None = None)

Wrap an operation to intercept inputs and outputs for intervention.

Parameters:

• fn – The function to wrap

• name – The name of the operation

• bound_obj – The object fn is bound to if it is a method

Returns:

A wrapped version of the function

class nnsight.intervention.interleaver.Mediator(intervention: Callable, info: Tracer.Info, name: str | None = None, batch_group: int | None = 0, stop: int | None = None)

Mediates between the model execution and intervention functions.

This class handles the communication between the model’s forward pass and user-defined intervention functions, allowing for inspection and modification of intermediate values.

exception OutOfOrderError

Exception raised when interventions are defined out of order.

cancel()

Cancel the intervention thread and clear caches.

end()

Signal that execution should continue without further intervention.

exception(exception: Exception)

Signal that an exception occurred during intervention.

Parameters:

exception – The exception that occurred

property frame: FrameType

Get the frame of the intervention function.

Returns:

The frame of the intervention function

handle(provider: Any | None = None)

Handle events in the event queue and process provider values.

Parameters:

provider – The identifier of the provider

Returns:

The original or modified value

handle_barrier_event(provider: Any, participants: Set[str])

Handle a barrier event by setting a barrier.

handle_end_event()

Handle an end event by stopping the mediator.

handle_exception_event(exception: Exception)

Handle an exception event by raising the exception.

Parameters:

exception – The exception to raise

Returns:

Boolean indicating whether to continue processing events

handle_swap_event(provider: Any, requester: Any, swap_value: Any)

Handle a swap event by swapping the value if the provider matches the requester.

Parameters:

• requester – The identifier of the requester

• provider – The identifier of the provider

• swap_value – The value to swap in

Returns:

Boolean indicating whether to continue processing events

handle_value_event(requester: Any, provider: Any, value: Any)

Handle a value event by providing the requested value or recording a missed provider.

Parameters:

• requester – The identifier of the requester

• provider – The identifier of the provider

• value – The value being provided

Returns:

Boolean indicating whether to continue processing events

iter(mediator: Mediator, iteration: int | slice)

Iterate a mediator a specified number of times.

Parameters:

• mediator – The mediator to iterate

• iteration – The number of iterations

pull()

Pull variables from the interleaver state to the frame globals.

push()

Push local variables to the interleaver state.

request(requester: Any)

Request a value from a specific provider.

Parameters:

requester – The identifier of the provider to request a value from

Returns:

The requested value

respond(value: Any | None = None)

Set the value for a pending value request.

Parameters:

value – The value to provide

send(event: Events, requester: Any)

Send an event to the event queue and wait for a response.

Parameters:

• event – The event to send

• requester – The identifier of the requester

Returns:

The response from the provider

set_user_cache(cache: Cache)

Set the user cache for this mediator.

Parameters:

cache – The cache to set

start(interleaver: Interleaver)

Start the mediator’s intervention thread.

Parameters:

interleaver – The interleaver managing this mediator

stop()

Stop the execution of the model by raising an EarlyStopException.

swap(requester: Any, value: Any)

Set a value to swap during execution.

Parameters:

• requester – The identifier of the requester

• value – The value to swap in

wait()

Wait for the next event to be set in the event queue.

exception nnsight.intervention.interleaver.SkipException(value: Any)

Exception raised to skip the execution of the model.

class nnsight.intervention.serialization.CustomCloudPickler(file, protocol=None, buffer_callback=None)

class nnsight.intervention.serialization.CustomCloudUnpickler(file, root: Envoy, frame: FrameType)

class nnsight.intervention.inject.FunctionCallWrapper(name: str)