WIP: strengthen instrumentation logic

Author: Essoz

docs/assets/code/mnist.py

@@ -8,9 +8,9 @@
 from torchvision import datasets, transforms
 
 from traincheck import annotate_stage
-from traincheck.instrumentor import meta_vars
+from traincheck.instrumentor import META_VARS
 
-meta_vars["step"] = -1
+META_VARS["step"] = -1
 
 
 class Net(nn.Module):
@@ -43,7 +43,7 @@ def train(args, model, device, train_loader, optimizer, epoch):
     annotate_stage("training")  # ML_DAIKON: stage annotation
     model.train()
     for batch_idx, (data, target) in enumerate(train_loader):
-        meta_vars["step"] += 1
+        META_VARS["step"] += 1
         data, target = data.to(device), target.to(device)
         optimizer.zero_grad()
         output = model(data)

docs/assets/examples/traincheck-collect/mnist-config/mnist.py

@@ -8,9 +8,9 @@
 from torchvision import datasets, transforms
 
 from traincheck import annotate_stage
-from traincheck.instrumentor import meta_vars
+from traincheck.instrumentor import META_VARS
 
-meta_vars["step"] = -1
+META_VARS["step"] = -1
 
 
 class Net(nn.Module):
@@ -43,7 +43,7 @@ def train(args, model, device, train_loader, optimizer, epoch):
     annotate_stage("training")  # ML_DAIKON: stage annotation
     model.train()
     for batch_idx, (data, target) in enumerate(train_loader):
-        meta_vars["step"] += 1
+        META_VARS["step"] += 1
         data, target = data.to(device), target.to(device)
         optimizer.zero_grad()
         output = model(data)

traincheck/developer/annotations.py

@@ -1,6 +1,6 @@
 import traincheck.instrumentor.tracer as tracer
 from traincheck.config.config import ALL_STAGE_NAMES
-from traincheck.instrumentor import meta_vars
+from traincheck.instrumentor import META_VARS
 
 
 def annotate_stage(stage_name: str):
@@ -16,7 +16,7 @@ def annotate_stage(stage_name: str):
         stage_name in ALL_STAGE_NAMES
     ), f"Invalid stage name: {stage_name}, valid ones are {ALL_STAGE_NAMES}"
 
-    meta_vars["stage"] = stage_name
+    META_VARS["stage"] = stage_name
 
 
 def annotate_answer_start_token_ids(

traincheck/instrumentor/VFProxy.py

@@ -1,4 +1,4 @@
-from .caches import meta_vars  # noqa: F401
+from .caches import META_VARS  # noqa: F401
 from .source_file import *  # noqa: F403
 from .tracer import *  # noqa: F403
 from .tracer import VarSampler  # noqa: F401

traincheck/instrumentor/__init__.py

@@ -1,8 +1,3 @@
-from collections import defaultdict
-
-from traincheck.instrumentor.types import PTID
-
-cache_meta_vars: dict[PTID, dict[str, dict]] = defaultdict(lambda: defaultdict(dict))
-meta_vars: dict[str, object] = {
+META_VARS: dict[str, object] = {
     "step": 0,
 }

traincheck/instrumentor/caches.py

@@ -93,24 +93,23 @@ def unproxy_arg(arg, inspect_torch_module=False):
         return arg
 
 
+def unproxy_args_kwargs(args, kwargs, inspect_torch_module=False):
+    args = [unproxy_arg(arg, inspect_torch_module) for arg in args]
+    kwargs = {k: unproxy_arg(v) for k, v in kwargs.items()}
+    return args, kwargs
+
+
 def unproxy_func(func, inspect_torch_module=False):
     original_func = func
 
     @functools.wraps(original_func)
     def wrapper(*args, **kwargs):
-        args = [unproxy_arg(arg, inspect_torch_module) for arg in args]
-        kwargs = {k: unproxy_arg(v) for k, v in kwargs.items()}
+        args, kwargs = unproxy_args_kwargs(args, kwargs, inspect_torch_module)
         return original_func(*args, **kwargs)
 
     return wrapper
 
 
-def unproxy_args_kwargs(args, kwargs, inspect_torch_module=False):
-    args = [unproxy_arg(arg, inspect_torch_module) for arg in args]
-    kwargs = {k: unproxy_arg(v) for k, v in kwargs.items()}
-    return args, kwargs
-
-
 def type_handle_traincheck_proxy(x):
     if hasattr(x, "is_traincheck_proxied_obj"):
         return type(x._obj)

traincheck/instrumentor/proxy_wrapper/proxy_basics.py

@@ -61,7 +61,7 @@ def wrapper(*args, **kwargs):
         result = processed_function(*args, **kwargs)
 
         # post observe
-        for i, var in enumerate(proxied_vars):
+        for var in proxied_vars:
             observe_proxy_var(
                 var,
                 "post_observe",

traincheck/instrumentor/proxy_wrapper/proxy_observer.py

@@ -15,11 +15,15 @@ def __init__(self, proxy: "Proxy", stale: bool):
         self.stale = stale
 
 
-class ProxyRegistry:
-    """A helper class managing all proxy variables being tracked and allow for controlled dumps of
+class VarRegistry:
+    """A helper class managing all variables being tracked and allow for controlled dumps of
     the variable states.
 
     A variable is uniquely identified by its "name"
+    When a variable is added to the registry, it is marked as "not stale".
+    When a variable is dumped through `dump_sample` or `dump_modified`, it is marked as "stale".
+    A variable is only dumped through `dump_modified` if it is not stale.
+
     """
 
     def __init__(self):
@@ -29,20 +33,24 @@ def __init__(self):
     def add_var(self, var: "Proxy", var_name: str):
         """Add a new proxy variable to the registry"""
         with self.registry_lock:
-            self.registry[var_name] = RegistryEntry(proxy=var, stale=False)
+            if var_name in self.registry:
+                self.registry[var_name].proxy = var
+                self.registry[var_name].stale = False
+            else:
+                self.registry[var_name] = RegistryEntry(proxy=var, stale=False)
 
     def dump_sample(self, dump_loc=None):
         """A complete dump of all present proxy objects
 
         Calling this API mark all proxy objects as stale which
-        will affect the `dump_only_modified` API.
+        will affect the `dump_modified` API.
         """
         with self.registry_lock:
-            for var_name, entry in self.registry.items():
+            for _, entry in self.registry.items():
                 entry.stale = True
                 entry.proxy.dump_trace(phase="sample", dump_loc=dump_loc)
 
-    def dump_only_modified(self, dump_loc=None, dump_config=None):
+    def dump_modified(self, dump_loc=None, dump_config=None):
         """Dump only the proxy variables that might be modified since last dump
 
         args:
@@ -81,7 +89,7 @@ def dump_only_modified(self, dump_loc=None, dump_config=None):
 
 
 # Global dictionary to store registered objects
-global_registry = ProxyRegistry()
+global_registry = VarRegistry()
 
 
 def get_global_registry():

traincheck/instrumentor/proxy_wrapper/proxy_registry.py

@@ -12,8 +12,8 @@
 from traincheck.utils import get_timestamp_ns
 
 from .proxy_basics import is_fake_tensor
+from .proxy_registry import get_global_registry
 
-# from .proxy_registry import get_global_registry
 # from .utils import print_debug
 
 
@@ -165,9 +165,9 @@ def update_timestamp(self):
         # Proxy.var_dict[self.__dict__["var_name"]].last_update_timestamp = current_time
 
     def register_object(self):
-        # get_global_registry().add_var(self, self.__dict__["var_name"])
+        get_global_registry().add_var(self, self.__dict__["var_name"])
         # TODO: implement the registry, we will need to make sure the registerred timestamp is updated and is consistent with the timestamp in the object
-        pass
+        # pass
 
     def dump_trace(self, phase, dump_loc):
         # print(f"parameter: {self.var_name}, phase = {phase}, dump_loc = {dump_loc}")