Vllm

src/nnsight/__init__.py

@@ -1,4 +1,4 @@
-# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
+# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
 #                                                                                                                       #
 #      ::::    ::: ::::    :::  :::::::: ::::::::::: ::::::::  :::    ::: :::::::::::       :::::::       ::::::::      #
 #      :+:+:   :+: :+:+:   :+: :+:    :+:    :+:    :+:    :+: :+:    :+:     :+:          :+:   :+:     :+:    :+:     #
@@ -8,10 +8,10 @@
 #      #+#   #+#+# #+#   #+#+# #+#    #+#    #+#    #+#    #+# #+#    #+#     #+#          #+#   #+# #+# #+#    #+#     #
 #      ###    #### ###    ####  ######## ########### ########  ###    ###     ###           #######  ###  ########      #
 #                                                                                                                       #
-# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
+# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
 import os
 from functools import wraps
-from typing import Dict, Union
+from typing import Callable, Dict, Union
 
 from importlib.metadata import version, PackageNotFoundError
 
@@ -56,11 +56,11 @@
 from torch._subclasses.fake_tensor import FakeTensor
 
 
-def _bool(self):
+def fake_bool(self):
     return True
 
 
-DEFAULT_PATCHER.add(Patch(FakeTensor, _bool, "__bool__"))
+DEFAULT_PATCHER.add(Patch(FakeTensor, fake_bool, "__bool__"))
 
 
 def fake_tensor_new_wrapper(fn):
@@ -118,10 +118,11 @@ def noop(input: torch.Tensor, *args, **kwargs):
 )
 
 import warnings
+
 _str = str
+_bool = bool
+
 try:
-
-
 
     from torch.amp.autocast_mode import autocast, is_autocast_available
 
@@ -555,3 +556,88 @@ def set_module_tensor_to_device(
 apply = GlobalTracingContext.GLOBAL_TRACING_CONTEXT.apply
 log = GlobalTracingContext.GLOBAL_TRACING_CONTEXT.log
 cond = GlobalTracingContext.GLOBAL_TRACING_CONTEXT.cond
+
+import inspect
+
+from . import util
+from .intervention import InterventionProxy
+
+
+def trace(fn: Callable):
+    """Helper decorator to add a function to the intervention graph via `.apply(...)`.
+    This is opposed to entering the function during tracing and tracing all inner operations.
+
+    Args:
+        fn (Callable): Function to apply.
+
+    Returns:
+        Callable: Traceable function.
+    """
+
+    @wraps(fn)
+    def inner(*args, **kwargs):
+
+        return apply(fn, *args, **kwargs)
+
+    return inner
+
+
+def local(object: Callable | InterventionProxy):
+    """Helper decorator to add a function to the intervention graph via `.apply(...)`
+    AND convert all input Proxies to local ones via `.local()`.
+
+    If a non-function is passed in, its assumed to be an `InterventionProxy` and `.local()` is called and returned.
+
+    Args:
+        object ( Callable | InterventionProxy): Function to apply or Proxy to make local.
+
+    Returns:
+        Callable | InterventionProxy: Traceable local function or local Proxy.
+    """
+
+    if inspect.isroutine(object):
+
+        fn = trace(object)
+
+        @wraps(fn)
+        def inner(*args, **kwargs):
+
+            args, kwargs = util.apply(
+                (args, kwargs), lambda x: x.local(), InterventionProxy
+            )
+
+            return fn(*args, **kwargs)
+
+        return inner
+
+    return object.local()
+
+
+def remote(object: Callable | Any):
+    """Helper decorator to add a function to the intervention graph via `.apply(...)`
+    AND convert all input Proxies to downloaded local ones via `.local()`
+    AND convert the output to an uploaded remote one via `remote()`.
+
+    If a non-function is passed in, `remote(object)` is called and returned.
+
+    Args:
+        object ( Callable | Any): Function to apply or object to make remote.
+
+    Returns:
+        Callable | InterventionProxy: Traceable local -> remote function or remote Proxy.
+    """
+
+    if inspect.isroutine(object):
+
+        fn = local(object)
+
+        @wraps(fn)
+        def inner(*args, **kwargs):
+
+            return GlobalTracingContext.GLOBAL_TRACING_CONTEXT.remote(
+                fn(*args, **kwargs)
+            )
+
+        return inner
+
+    return GlobalTracingContext.GLOBAL_TRACING_CONTEXT.remote(object)

src/nnsight/contexts/GraphBasedContext.py

@@ -136,6 +136,18 @@ def log(self, *data: Any) -> None:
             data (Any): Data to print.
         """
         self.apply(print, *data)
+
+    def remote(self, data:Any) -> InterventionProxy:
+        """Streams data remotely when it becomes available locally.
+        The remote service will block until the local value is uploaded and received.
+
+        Is a no-op when not executing remotely.
+
+        Returns:
+            InterventionProxy: Proxy.
+        """
+
+        return protocols.StreamingUploadProtocol.add(self.graph, data)
 
     def bool(self, *args, **kwargs) -> InterventionProxy:
         """NNsight helper method to create a traceable bool."""
@@ -208,13 +220,16 @@ def __enter__(self) -> Self:
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
 
         GlobalTracingContext.try_deregister(self)
-
+        
         if isinstance(exc_val, BaseException):
-            self.graph.alive = False
             self.graph = None
             raise exc_val
 
         self.backend(self)
+
+        if not isinstance(self.graph, weakref.ProxyType):
+            self.graph = weakref.proxy(self.graph)
+
 
     ### BACKENDS ########
 
@@ -225,22 +240,13 @@ def local_backend_execute(self) -> None:
             self.graph.execute()
         except protocols.EarlyStopProtocol.EarlyStopException as e:
             raise e
-        finally:
-            graph = self.graph
-            graph.alive = False
-
-            if not isinstance(graph, weakref.ProxyType):
-                self.graph = weakref.proxy(graph)
 
     def bridge_backend_handle(self, bridge: Bridge) -> None:
 
         bridge.pop_graph()
 
         protocols.LocalBackendExecuteProtocol.add(self, bridge.peek_graph())
 
-        self.graph = weakref.proxy(self.graph)
-
-
 from inspect import getmembers, isclass
 
 from torch.utils import data

src/nnsight/contexts/Invoker.py

@@ -37,15 +37,15 @@ class Invoker(AbstractContextManager):
     def __init__(
         self,
         tracer: "Tracer",
-        *inputs: Any,
+        *args,
         scan: bool = False,
         **kwargs,
     ) -> None:
 
         self.tracer = tracer
-        self.inputs = inputs
+        self.input = (args, kwargs)
+
         self.scan = scan
-        self.kwargs = kwargs
 
         self.scanning = False
 
@@ -71,28 +71,28 @@ def __enter__(self) -> Invoker:
         # Set self.inputs to be the proxy_value so we can prepare_inputs, get the batch size, and scan.
         if self.tracer.model._session is not None:
 
-            self.inputs, has_proxies_in_inputs = check_for_dependencies(
-                self.inputs
+            self.input, has_proxies_in_inputs = check_for_dependencies(
+                self.input
             )
 
         with GlobalTracingContext.exit_global_tracing_context():
 
             if not has_proxies_in_inputs:
 
-                self.inputs, batch_size = self.tracer.model._prepare_inputs(
-                    *self.inputs, **self.kwargs
+                self.input, batch_size = self.tracer.model._prepare_input(
+                    *self.input[0], **self.input[1]
                 )
 
             if self.scan:
 
-                inputs = self.inputs
+                input = self.input
 
                 if has_proxies_in_inputs:
 
-                    inputs = util.apply(inputs, lambda x: x.proxy_value, Node)
+                    input = util.apply(input, lambda x: x.proxy_value, Node)
 
-                    inputs, batch_size = self.tracer.model._prepare_inputs(
-                        *inputs, **self.kwargs
+                    input, batch_size = self.tracer.model._prepare_input(
+                        *input[0], **input[1]
                     )
 
                 self.tracer.model._envoy._clear()
@@ -112,8 +112,10 @@ def __enter__(self) -> Invoker:
                         shape_env=ShapeEnv(assume_static_by_default=True),
                     ) as fake_mode:
                         with FakeCopyMode(fake_mode):
-                            self.tracer.model._execute(
-                                *copy.deepcopy(inputs),
+                            fn = self.tracer.model._execute if self.tracer.method is None else getattr(self.tracer.model, self.tracer.method)
+                            fn(
+                                *copy.deepcopy(input[0]),
+                                **copy.deepcopy(input[1]),
                                 **copy.deepcopy(self.tracer._kwargs),
                             )
 
@@ -122,7 +124,7 @@ def __enter__(self) -> Invoker:
             else:
                 self.tracer.model._envoy._reset()
 
-            self.tracer._invoker_inputs.append(self.inputs)
+            self.tracer._invoker_inputs.append(self.input)
 
         return self