update toggle && put some TODO placeholder

pprobe/bootstrap/hook_setup.py

@@ -4,7 +4,6 @@
 from pprobe.toggle.cli import ToggleManager
 
 from pprobe.bootstrap.hooks.pytorch_catch import func_torch_step_count_wrapper
-from pprobe.bootstrap.hooks.pytorch_dist import func_torch_distributed_wrapper
 from pprobe.bootstrap.hooks.pytorch_perf import func_torch_device_conversion_wrapper
 
 
@@ -15,38 +14,94 @@ class PProbeSetup:
     def __init__(self, module, module_fullname):
         self.module = module
         self.pprobe_toggle = ToggleManager()
+        ####################
+        ### ENABLE FIELD
+        ####################
         self.pprobe_enabled = self.pprobe_toggle.get_toggle("PPROBE_ENABLE")
+        self.torch_reproduce_enabled = self.pprobe_toggle.get_toggle("TORCH_REPRODUCE")
+
         self.torch_catch_step_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_CATCH_STEP"
         )
-        self.torch_reproduce_enabled = self.pprobe_toggle.get_toggle("TORCH_REPRODUCE")
+        self.torch_catch_loss_enabled = self.pprobe_toggle.get_toggle(
+            "TORCH_CATCH_LOSS"
+        )
+        self.torch_catch_lr_enabled = self.pprobe_toggle.get_toggle("TORCH_CATCH_LR")
+
         self.torch_dump_op_enabled = self.pprobe_toggle.get_toggle("TORCH_DUMP_OP")
         self.torch_dump_dist_enabled = self.pprobe_toggle.get_toggle("TORCH_DUMP_DIST")
         self.torch_dump_module_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_DUMP_MODULE"
         )
+        self.torch_dump_optim_enabled = self.pprobe_toggle.get_toggle(
+            "TORCH_DUMP_OPTIM"
+        )
+        self.torch_dump_memory_enabled = self.pprobe_toggle.get_toggle(
+            "TORCH_DUMP_MEMORY"
+        )
+        self.torch_test_dump_op_enabled = self.pprobe_toggle.get_toggle(
+            "TORCH_TEST_DUMP_OP"
+        )
+        self.torch_test_dump_dist_enabled = self.pprobe_toggle.get_toggle(
+            "TORCH_TEST_DUMP_DIST"
+        )
+        self.torch_test_dump_module_enabled = self.pprobe_toggle.get_toggle(
+            "TORCH_TEST_DUMP_MODULE"
+        )
         self.torch_perf_issue_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_PERF_ISSUE"
         )
-
+        self.torch_torch_trace_file_enabled = self.pprobe_toggle.get_toggle(
+            "TORCH_TRACE_FILE"
+        )
+        ####################
+        ### INIT HOOK
+        ####################
         self.check_and_run_hook(module_fullname)
 
     def check_and_run_hook(self, module_fullname):
         if self.pprobe_enabled:
             self.run_generic_hook()
             # torch part
             if module_fullname == "torch":
+                if self.torch_reproduce_enabled:
+                    # TODO
+                    pass
                 if self.torch_catch_step_enabled:
                     self.run_torch_catch_step_hook()
-                if self.torch_reproduce_enabled:
-                    self.run_torch_reproduce_hook()
+                if self.torch_catch_loss_enabled:
+                    # TODO
+                    pass
+                if self.torch_catch_lr_enabled:
+                    # TODO
+                    pass
                 if self.torch_dump_op_enabled:
                     self.run_torch_func_hook()
+                if self.torch_dump_dist_enabled:
+                    # TODO
+                    pass
                 if self.torch_dump_module_enabled:
                     self.run_torch_module_hook()
-                if self.torch_dump_dist_enabled:
-                    self.run_torch_func_hook()
+                if self.torch_dump_optim_enabled:
+                    # TODO
+                    pass
+                if self.torch_dump_memory_enabled:
+                    # TODO
+                    pass
+                if self.torch_test_dump_op_enabled:
+                    # TODO
+                    pass
+                if self.torch_test_dump_dist_enabled:
+                    # TODO
+                    pass
+                if self.torch_test_dump_module_enabled:
+                    # TODO
+                    pass
                 if self.torch_perf_issue_enabled:
+                    # TODO
+                    pass
+                if self.torch_torch_trace_file_enabled:
+                    # TODO
                     pass
         else:
             self.print_warning()
@@ -60,18 +115,21 @@ def run_generic_hook(self):
     def run_torch_func_hook(self):
         from pprobe.bootstrap.hooks import pytorch_func_op
 
+        Logger.info(f"[PPROBE] torch function hook executed")
         context = pytorch_func_op.TorchFunctionContext()
         context.__enter__()
-        Logger.info(f"[PPROBE] torch function hook executed")
 
     def run_torch_module_hook(self):
         from pprobe.bootstrap.hooks import pytorch_module
 
+        Logger.info(f"[PPROBE] torch module hook executed")
+
         context = pytorch_module.TorchModuleContext()
         context.__enter__()
-        Logger.info(f"[PPROBE] torch module hook executed")
 
     def run_torch_dist_hook(self):
+        from pprobe.bootstrap.hooks.pytorch_dist import func_torch_distributed_wrapper
+
         Logger.info(f"[PPROBE] torch dist hook executed")
 
         ###################################################

pprobe/toggle/hook.toggle.default

@@ -1,10 +1,10 @@
 #########################
-## PPROBE-CONTROL  
+## PPROBE-CONTROL
 #########################
 PPROBE_ENABLE=False
 
 #########################
-## TORCH-CONTROL  
+## TORCH-CONTROL
 #########################
 TORCH_REPRODUCE=True
 TORCH_CATCH_STEP=False
@@ -17,10 +17,11 @@ TORCH_CATCH_LR=False
 TORCH_DUMP_OP=False
 TORCH_DUMP_MODULE=False
 TORCH_DUMP_DIST=False
+TORCH_DUMP_OPTIM=False
 TORCH_DUMP_MEMORY=False
 
 #########################
-## TORCH-TESTCASE DUMP  
+## TORCH-TESTCASE DUMP
 #########################
 TORCH_TEST_DUMP_OP=False
 TORCH_TEST_DUMP_MODULE=False