add: monkey patch __setattr__ at the module level when using subclass, to ensure submodule assignments are captured

Author: Essoz

traincheck/instrumentor/proxy_wrapper/subclass.py

@@ -1,3 +1,4 @@
+import functools
 import logging
 import os
 import threading
@@ -14,7 +15,9 @@
 from .proxy_basics import is_fake_tensor
 from .proxy_registry import get_global_registry
 
-# from .utils import print_debug
+SUBCLASS_HOOK_KEY = "_tc_setattr_hook"
+
+logger = logging.getLogger(__name__)
 
 
 def in_dynamo() -> bool:
@@ -126,9 +129,9 @@ def __init__(
         current_time = get_timestamp_ns()
 
         self.__dict__["last_update_timestamp"] = current_time
+        logger.debug(f"[ProxyParameter] Created ProxyParameter: {self.var_name}")
         self.register_object()
 
-        # print(f"init: {self.var_name}")
         if should_dump_trace and not should_disable_proxy_dumping():
             if from_call:
                 phase = "call"
@@ -141,7 +144,7 @@ def __init__(
             self.dump_trace(phase=phase, dump_loc="initing")
 
     def __setattr__(self, name, value):
-        # print(f"paremeter: {self.var_name}, name = {name}, value = {value}")
+
         super().__setattr__(name, value)
         self.update_timestamp()
         self.register_object()
@@ -175,7 +178,6 @@ def register_object(self):
         )
 
     def dump_trace(self, phase, dump_loc):
-        # print(f"parameter: {self.var_name}, phase = {phase}, dump_loc = {dump_loc}")
         # TODO
         var_name = self.__dict__["var_name"]
         # assert var_name is not None  # '' is allowed as a var_name (root object)
@@ -224,6 +226,9 @@ def proxy_parameter(
         return
     for name, t in list(module.named_parameters(recurse=False)):
         var_type = typename(t, is_runtime=True)
+        logger.debug(
+            f"[ProxyParameter] Proxying parameter: {parent_name}.{name} of type {var_type}"
+        )
         module._parameters[name] = ProxyParameter(
             t,
             logdir,
@@ -244,3 +249,37 @@ def proxy_parameter(
             from_call,
             from_iter,
         )
+
+    # we need to instrument the __setattr__ of the module to capture parameter updates
+    def subclass_setattr_hook(self, name, value):
+        logger.debug(
+            f"[ProxyParameter] Module __setattr__ called: {parent_name}.{name} = {type(value)}"
+        )
+        if isinstance(value, torch.Tensor) or isinstance(value, torch.nn.Module):
+            proxy_parameter(
+                value,
+                logdir,
+                log_level,
+                parent_name + "." + name,
+                should_dump_trace,
+                from_call,
+                from_iter,
+            )
+
+    module.__dict__[SUBCLASS_HOOK_KEY] = subclass_setattr_hook
+
+
+# instrument torch.nn.Module's setattr
+orig_setattr = torch.nn.Module.__setattr__
+
+
+@functools.wraps(orig_setattr)
+def wrapped_setattr(self, name, value):
+    hook = getattr(self, SUBCLASS_HOOK_KEY, None)
+    if hook is not None:
+        # If hook returns True, skip the original setattr; otherwise continue.
+        hook(self, name, value)
+    return orig_setattr(self, name, value)
+
+
+torch.nn.Module.__setattr__ = wrapped_setattr