Merge pull request #121 from OrderLab/fix_proxy_changing_model

Fix Proxy Changing Wrapped Model #120

Author: Essoz

mldaikon/collect_trace.py

@@ -352,12 +352,6 @@ def is_path_md_output_dir(output_dir: str) -> bool:
         choices=["sampler", "proxy"],
         default="proxy",
     )
-    parser.add_argument(
-        "--proxy-update-limit",
-        type=float,
-        default=proxy_config.proxy_update_limit,
-        help="The threshold for updating the proxy object",
-    )
     parser.add_argument(
         "--tensor-dump-format",
         choices=["hash", "stats", "full"],
@@ -431,7 +425,6 @@ def is_path_md_output_dir(output_dir: str) -> bool:
     # set up adjusted proxy_config
     proxy_basic_config: dict[str, int | bool | str] = {}
     for configs in [
-        "proxy_update_limit",
         "debug_mode",
         "enable_C_level_observer",
     ]:

mldaikon/config/config.py

@@ -198,7 +198,7 @@ def should_disable_proxy_dumping() -> bool:
     "__doc__",
     "logdir",
     "log_level",
-    "is_root",
+    "recurse",
     "var_name",
     "mode",
     "process_id",

mldaikon/instrumentor/source_file.py

@@ -178,7 +178,7 @@ def visit_Assign(self, node):
                         func=ast.Name(id="Proxy", ctx=ast.Load()),
                         args=[node.value],
                         keywords=[
-                            ast.keyword(arg="is_root", value=ast.Constant(value=True)),
+                            ast.keyword(arg="recurse", value=ast.Constant(value=True)),
                             ast.keyword(
                                 arg="logdir",
                                 value=ast.Attribute(
@@ -187,6 +187,10 @@ def visit_Assign(self, node):
                                     ctx=ast.Load(),
                                 ),
                             ),
+                            ast.keyword(
+                                arg="var_name",
+                                value=ast.Constant(value=self.model_name),
+                            ),
                         ],
                     )
             return node

mldaikon/proxy_wrapper/README.md

@@ -6,19 +6,19 @@
 
 `import src.proxy_wrapper.proxy as Proxy`
 
-- wrap the machine learning model with `is_root = True`
+- wrap the machine learning model with `recurse = True`
 
-`model = Proxy(model, is_root = True)`
+`model = Proxy(model, recurse = True)`
 
 - Examples:
 
 As shown in line 140 in `./proxyclass_tracer_result/instrumented_mnist.py`
 
-`model=Proxy.Proxy(model, is_root = True, logdir='log-model-proxy-example.log', log_level=logging.INFO)`
+`model=Proxy.Proxy(model, recurse = True, logdir='log-model-proxy-example.log', log_level=logging.INFO)`
 
 line 99 in `./proxyclass_tracer_result/instrumented_84911.py`
 
-`model_transfer=Proxy.Proxy(model_transfer, is_root = True, "model_transfer-example.log", log_level = logging.INFO)`
+`model_transfer=Proxy.Proxy(model_transfer, recurse = True, "model_transfer-example.log", log_level = logging.INFO)`
 
 Note: Initially, we want to achieve automatic instrumentation via `__new__ wrapper`, wrapping all `torch.nn` modules and add a proxy to the targeted modules inherently. However, it is discovered that this instrumentation would interfere with inherent torch functionality, such as `torch.autograd` behavior.
 

mldaikon/proxy_wrapper/dumper.py

@@ -22,14 +22,6 @@ def __call__(cls, *args, **kwargs):
         return cls._instances[cls]
 
 
-class SkippedDumpingObj:
-    def __init__(self, obj):
-        self._obj = obj
-
-    def __repr__(self):
-        return f"Skipped Dumping Object: ({self._obj})"
-
-
 class json_dumper(metaclass=Singleton):
     # singleton pattern for shared state
     _shared_state = False
@@ -58,12 +50,10 @@ def dump_json(
             return
 
         data = {
-            # "value": var_value,
             "var_name": var_name,
             "var_type": var_type,
             "mode": change_type,  # "new", "update"
             "dump_loc": dump_loc,
-            # "stack_trace": stack_trace,
             "process_id": process_id,
             "thread_id": thread_id,
             "time": time,