nnsight.tracing#

The tracing module acts as a standalone library to trace and execute Python based deferred computation graphs.

The graph sub-module defines the computation graph primitives. The protocol sub-module contains logic for adding custom operations to the computation graph. The contexts sub-module contains logic for defining scoped sub-graphs that handle execution of their piece of the computation graph. The backends sub-module contains logic for executing the traced computation graph.

class nnsight.tracing.graph.graph.Graph(node_class: ~typing.Type[~nnsight.tracing.graph.node.NodeType] = <class 'nnsight.tracing.graph.node.Node'>, proxy_class: ~typing.Type[~nnsight.tracing.graph.proxy.ProxyType] = <class 'nnsight.tracing.graph.proxy.Proxy'>, debug: bool = False)[source]#

The Graph class represents a computation graph composed of individual Node`s (operations). It contains logic to both trace/build the computation graph, as well as how to execute it. Sections of the graph can be divided into `SubGraphs, but there will always be one root Graph. The final Node of the graph (graph[-1]) should be the root Node which when executed, downstream executes the entire Graph.

node_class#

Class used to create Node`s. Can be changed to add additional functionality to `Node’s. Defaults to `Node.

Type:: Type[NodeType]

proxy_class#

Class used to create Proxy`s for ‘Node’s. Can be changed to add additional functionality to `Proxy’s. Defaults to `Proxy.

Type:: Type[ProxyType]

nodes#

Ordered list of all Node`s. Used to access `Nodes via their index.

Type:: List[Node]

stack#

List of Graph`s as a stack. Used to move `Node`s onto the most recent graph, as opposed to the `Graph used to create the Node. Managed outside the Graph class by the Context objects.

Type:: List[Graph]

defer_stack#

List of Node indexes as a stack. Used to prevent destruction/memory cleanup of Node`s whose index is less than the most recent index on the stack. This happens when you have `Node`s that will be executed more than once. In a loop for example, you only want to destroy a `Node`s dependencies on the final iteration. Also managed outside the `Graph object.

Type:: List[int]

alive#

If the Graph is “alive”. Alive meaning its still open for tracing (adding new Node`s). Set to False before executing the `Graph.

Type:: bool

add(node: NodeType) → None[source]#

Adds a Node to this Graph. Sets the Node’s .index attribute so it knows its own index within the entire computation graph.

Parameters:: node (NodeType) – Node to add.

clean(start: int | None = None)[source]#

Cleans up dependencies of Node`s so their values are appropriately memory managed. Cleans all `Node`s from start to end regardless if they are on this `Graph.

Parameters:: start (Optional[int], optional) – Node index to start cleaning up from. Defaults to None.

copy(new_graph: Graph[NodeType, ProxyType] | None = None) → Graph[source]#

Creates a shallow copy of the root Graph object.

Parameters:: new_graph (Optional[Graph[NodeType, ProxyType]], optional) – Graph to copy into. Defaults to None and creates a new Graph.
Returns:: New Graph.
Return type:: Graph

create(target: Callable | Protocol, *args, redirect: bool = True, **kwargs) → ProxyType[source]#

Creates a new Node using this Graph’s node_class and returns a Proxy for it with this Graph’s proxy_class.

Parameters:

target (Union[Callable, protocols.Protocol]) – Target for the new Node.
redirect (bool, optional) – If to move the newly created Node to the most recent Graph on the Graph.stack. Defaults to True.

Returns:

Proxy for newly created Node.

Return type:

ProxyType

execute() → None[source]#

Executes all Node`s (operations) in this `Graph.

Raises:: exception – If there is an exception during executing a Node. If so, we need to clean up the dependencies of `Node`s yet to be executed.

reset() → None[source]#: Resets the Graph to prepare for execution. Simply resets all Node`s in the `Graph.

class nnsight.tracing.graph.graph.SubGraph(parent: GraphType, subset: List[int] | None = None)[source]#

Represents a slice of the greater computation graph. It has a reference to the same underlying list of nodes and simply maintains a subset of node indexes.

subset#

Node indexes for `Node`s contained within this subgraph.

Type:: List[int]

class Iterator(subgraph: SubGraph[GraphType])[source]#

add(node: NodeType) → None[source]#

Adds a Node to this Graph. Sets the Node’s .index attribute so it knows its own index within the entire computation graph.

Parameters:: node (NodeType) – Node to add.

copy(new_graph: SubGraph[NodeType, ProxyType] | None = None, parent: Graph[NodeType, ProxyType] | None = None, memo: Dict[int, NodeType] | None = None) → Self[source]#

Creates a shallow copy of this SubGraph.

Parameters:

new_graph (Optional[SubGraph[NodeType, ProxyType]], optional) – SubGraph to copy into. Defaults to None and creates a new SubGraph of the same type.
parent (Optional[Graph[NodeType, ProxyType]], optional) – Parent graph. Defaults to None and will create a root Graph as the parent.

Returns:

New graph.

Return type:

Self

class nnsight.tracing.graph.node.Node(target: Callable | Protocol, *args, graph: Graph = None, **kwargs)[source]#

A computation Graph is made up of individual Node`s which represent a single operation. It has a `target which the operation this Node will execute. It has args and kwargs to execute its target with. These may contain other Node`s and are therefore `dependencies of this Node. Conversely this Node is a listener of its dependencies.

During execution of the computation graph and therefore the `Node`s, each

index#

Integer index of this Node within its greater computation graph.

Type:: Optional[int]

graph#

Type:: Graph

target#

Callable to execute as this Node’s operation. Might be a Protocol which is handled differently in node execution.

Type:: Union[Callable, Protocol]

property attached: bool#

Checks to see if the Graph this Node is a part of is alive.. Alive meaning the Graph is still open to tracing new Nodes.

Returns:: Is Node attached.
Return type:: bool

create(*args, **kwargs) → NodeType | Any[source]#

We use Node.create vs Graph.create in case graph is dead. If the graph is dead, we first check the GlobalTracing Context to add assume this node is ready to execute and therefore we try and execute it and then return its value.

Returns:: Proxy or value
Return type:: Union[NodeType, Any]

property dependencies: List[Self]#

Iterator from index to Node.

Returns:: List of dependency `Node`s.
Return type:: List[Self]

destroy() → None[source]#: Removes the reference to the node’s value and logs it’s destruction.

property done: bool#

Returns true if the value of this node has been set.

Returns:: If done.
Return type:: bool

execute() → None[source]#: Actually executes this node. Lets protocol execute if target is Protocol. Else prepares args and kwargs and passes them to target. Gets output of target and sets the Node’s value to it.

property fulfilled: bool#

Returns true if remaining_dependencies is 0.

Returns:: If fulfilled.
Return type:: bool

property listeners: List[Self]#

Iterator from index to Node.

Returns:: List of listener `Node`s.
Return type:: List[Self]

classmethod prepare_inputs(inputs: Any) → Any[source]#

Prepare arguments for executing this node’s target. Converts Nodes in args and kwargs to their value.

Returns:: Prepared inputs.
Return type:: Any

preprocess() → None[source]#: Preprocess Node.args and Node.kwargs. Converts Proxies to their Node. Converts Nodes that are done to their value. Adds Node arguments to self dependencies. Add self to Node argument listeners.

property redundant: bool#

Returns true if remaining_listeners is 0.

Returns:: If redundant.
Return type:: bool

reset() → None[source]#: Resets this Nodes remaining_listeners and remaining_dependencies.

set_value(value: Any) → None[source]#

Sets the value of this Node and logs the event. Updates remaining_dependencies of listeners. If they are now fulfilled, execute them. Updates remaining_listeners of dependencies. If they are now redundant, destroy them.

Parameters:: value (Any) – Value.

subgraph(subgraph: Set[int] | None = None) → Set[int][source]#

Returns a Set of indexes starting from this node, and recursively iterating over all the Node’s listeners.

Parameters:: subgraph (Optional[Set[int]], optional) – Current subgraph. Defaults to None.
Returns:: Set of Node indexes.
Return type:: Set[int]

update_dependencies()[source]#: Updates remaining_listeners of dependencies. If they are now redundant, destroy them.

update_listeners()[source]#: Updates remaining_dependencies of listeners.

property value: Any#

Property to return the value of this node.

Returns:: The stored value of the node, populated during execution.
Return type:: Any
Raises:: ValueError – If the underlying ._value is inspect._empty (therefore never set or was destroyed).

class nnsight.tracing.graph.proxy.Proxy(node: Node)[source]#

Proxy objects are the actual objects that interact with operations in order to update the graph to create new Nodes.

The operations that are traceable on base Proxy objects are many python built-in and magic methods.

node#

This proxy’s Node.

Type:: NodeType

save() → Self[source]#

Adds a lock Node to prevent its value from being cleared where normally it would be cleared when its no longer needed to save memory. Used to access values outside of the tracing context, after execution.

Returns:: Proxy.
Return type:: InterventionProxy

property value: Any#

Property to return the value of this proxy’s node.

Returns:: The stored value of the proxy, populated during execution of the model.
Return type:: Any

class nnsight.tracing.contexts.base.Context(*args, backend: ~nnsight.tracing.backends.base.Backend | None = None, parent: ~nnsight.tracing.graph.graph.GraphType | None = None, graph: ~nnsight.tracing.graph.graph.GraphType | None = None, graph_class: ~typing.Type[~nnsight.tracing.graph.graph.SubGraph] = <class 'nnsight.tracing.graph.graph.SubGraph'>, node_class: ~typing.Type[~nnsight.tracing.graph.node.NodeType] = <class 'nnsight.tracing.graph.node.Node'>, proxy_class: ~typing.Type[~nnsight.tracing.graph.proxy.ProxyType] = <class 'nnsight.tracing.graph.proxy.Proxy'>, debug: bool = False, **kwargs)[source]#

A Context represents a scope (or slice) of a computation graph with specific logic for adding and executing nodes defined within it. It has a SubGraph which contains the nodes that make up the operations of the context. As an AbstractContextManager, entering adds its sub-graph to the stack, making new nodes created while within this context added to it’s sub-graph.

Exiting pops its sub-graph off the stack, allowing nodes to be added to its parent, and adds itself as a node to its parent Context/SubGraph. ( To say, “execute me”) If the Context has a backend, it pops its parent off the stack and passes it to the Backend object to execute. (This only happens if the context is the root-most context, and its parent is therefore the root Graph)

As a Context is itself a Protocol, it defines how to execute it’s sub-graph in the execute method.

backend#

Backend to execute the deferred root computation graph

Type:: Backend

class nnsight.tracing.contexts.conditional.Condition(condition: NodeType | None, branch: NodeType | None = None, *args, **kwargs)[source]#

classmethod style() → Dict[str, Any][source]#

Visualization style for this protocol node.

Returns:

dictionary style.

Return type:

Dict

class nnsight.tracing.contexts.globals.GlobalTracingContext[source]#

The Global Tracing Context handles adding tracing operations globally without reference to a given GraphBasedContext. There should only be one of these and that is GlobalTracingContext.GLOBAL_TRACING_CONTEXT. GlobalTracingContext.TORCH_HANDLER handles adding torch functions without reference to a given GraphBasedContext.

class GlobalTracingExit[source]#

static deregister() → None[source]#

Deregister GraphBasedContext globally.

Parameters:: graph_based_context (GraphBasedContext) – GraphBasedContext to deregister.

static register(graph_based_context: Tracer) → None[source]#

Register GraphBasedContext globally.

Parameters:: graph_based_context (GraphBasedContext) – GraphBasedContext to register.

static try_deregister(graph_based_context: Tracer) → bool[source]#

Attempts to deregister a Graph globally. Will not if graph_based_context does not have the same Graph as the currently registered one.

Parameters:: graph_based_context (GraphBasedContext) – GraphBasedContext to deregister.
Returns:: True if deregistering ws successful, False otherwise.
Return type:: bool

static try_register(graph_based_context: Tracer) → bool[source]#

Attempts to register a Graph globally.] Will not if one is already registered.

Parameters:: graph_based_context (GraphBasedContext) – GraphBasedContext to register.
Returns:: True if registering ws successful, False otherwise.
Return type:: bool

class nnsight.tracing.contexts.iterator.Iterator(collection: Collection, *args, **kwargs)[source]#

classmethod style() → Dict[str, Any][source]#

Visualization style for this protocol node.

Returns:

dictionary style.

Return type:

Dict

class nnsight.tracing.contexts.tracer.Tracer(*args, backend: ~nnsight.tracing.backends.base.Backend | None = None, parent: ~nnsight.tracing.graph.graph.GraphType | None = None, graph: ~nnsight.tracing.graph.graph.GraphType | None = None, graph_class: ~typing.Type[~nnsight.tracing.graph.graph.SubGraph] = <class 'nnsight.tracing.graph.graph.SubGraph'>, node_class: ~typing.Type[~nnsight.tracing.graph.node.NodeType] = <class 'nnsight.tracing.graph.node.Node'>, proxy_class: ~typing.Type[~nnsight.tracing.graph.proxy.ProxyType] = <class 'nnsight.tracing.graph.proxy.Proxy'>, debug: bool = False, **kwargs)[source]#