fix: unify registry implementation for proxy and subclass

Author: Essoz

traincheck/instrumentor/proxy_wrapper/proxy_registry.py

@@ -1,17 +1,13 @@
 import threading
-import typing
-
-from traincheck.utils import typename
-
-if typing.TYPE_CHECKING:
-    from .proxy import Proxy
 
 
 class RegistryEntry:
-    """A class to store the proxy object and its associated metadata"""
+    """A class to store the tracked object and its associated metadata"""
 
-    def __init__(self, proxy: "Proxy", stale: bool):
-        self.proxy = proxy
+    def __init__(self, obj, var_name, var_type, stale):
+        self.var = obj
+        self.var_name = var_name
+        self.var_type = var_type
         self.stale = stale
 
 
@@ -30,14 +26,18 @@ def __init__(self):
         self.registry: dict[str, RegistryEntry] = {}
         self.registry_lock = threading.Lock()
 
-    def add_var(self, var: "Proxy", var_name: str):
+    def add_var(self, var, var_name: str, var_type: str):
         """Add a new proxy variable to the registry"""
         with self.registry_lock:
             if var_name in self.registry:
-                self.registry[var_name].proxy = var
+                self.registry[var_name].var = var
+                self.registry[var_name].var_name = var_name
+                self.registry[var_name].var_type = var_type
                 self.registry[var_name].stale = False
             else:
-                self.registry[var_name] = RegistryEntry(proxy=var, stale=False)
+                self.registry[var_name] = RegistryEntry(
+                    var, var_name, var_type, stale=False
+                )
 
     def dump_sample(self, dump_loc=None):
         """A complete dump of all present proxy objects
@@ -48,7 +48,7 @@ def dump_sample(self, dump_loc=None):
         with self.registry_lock:
             for _, entry in self.registry.items():
                 entry.stale = True
-                entry.proxy.dump_trace(phase="sample", dump_loc=dump_loc)
+                entry.var.dump_trace(phase="sample", dump_loc=dump_loc)
 
     def dump_modified(self, dump_loc=None, dump_config=None):
         """Dump only the proxy variables that might be modified since last dump
@@ -73,16 +73,16 @@ def dump_modified(self, dump_loc=None, dump_config=None):
         """
         to_dump_types = set(dump_config.keys())
         with self.registry_lock:
-            for var_name, entry in self.registry.items():
-                var_type = typename(entry.proxy._obj, is_runtime=True)
+            for _, entry in self.registry.items():
+                var_type = entry.var_type
                 if var_type not in to_dump_types:
                     continue
 
                 if entry.stale:
                     continue
 
                 entry.stale = True
-                entry.proxy.dump_trace(phase="selective-sample", dump_loc=dump_loc)
+                entry.var.dump_trace(phase="selective-sample", dump_loc=dump_loc)
                 if not dump_config[var_type]["dump_unchanged"]:
                     # remove the var from to_dump_types so that we don't dump the same type twice
                     to_dump_types.remove(var_type)

traincheck/instrumentor/proxy_wrapper/subclass.py

@@ -9,7 +9,7 @@
 from traincheck.instrumentor.dumper import dump_trace_VAR
 from traincheck.instrumentor.proxy_wrapper.dumper import dump_attributes, get_meta_vars
 from traincheck.instrumentor.tracer import TraceLineType
-from traincheck.utils import get_timestamp_ns
+from traincheck.utils import get_timestamp_ns, typename
 
 from .proxy_basics import is_fake_tensor
 from .proxy_registry import get_global_registry
@@ -37,6 +37,7 @@ def __new__(
         # TODO
         # recurse=False,
         var_name="",
+        var_type="",
         should_dump_trace=True,
         from_call=False,
         from_iter=False,
@@ -97,6 +98,7 @@ def __init__(
         # TODO
         # recurse=False,
         var_name="",
+        var_type="",
         should_dump_trace=True,
         from_call=False,
         from_iter=False,
@@ -116,6 +118,7 @@ def __init__(
         # TODO
         # self.__dict__["recurse"] = recurse
         self.__dict__["var_name"] = var_name
+        self.__dict__["var_type"] = var_type
         # TODO
         # self.__dict__["old_value"] = None
         # self.__dict__["old_meta_vars"] = None
@@ -165,7 +168,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"], self.__dict__["var_type"]
+        )
 
     def dump_trace(self, phase, dump_loc):
         # print(f"parameter: {self.var_name}, phase = {phase}, dump_loc = {dump_loc}")
@@ -216,11 +221,13 @@ def proxy_parameter(
     if in_dynamo():
         return
     for name, t in list(module.named_parameters(recurse=False)):
+        var_type = typename(t, is_runtime=True)
         module._parameters[name] = ProxyParameter(
             t,
             logdir,
             log_level,
             parent_name + "." + name,
+            var_type,
             should_dump_trace,
             from_call,
             from_iter,