From 33e72a3a48f55db165389b9e4c614f0939140ffc Mon Sep 17 00:00:00 2001
From: clemente0731 <clemente0620@gmail.com>
Date: Wed, 3 Jul 2024 17:41:43 +0800
Subject: [PATCH] refine && add optim placeholder

---
 pprobe/bootstrap/hook_setup.py          | 23 +++++++++++++++++++----
 pprobe/bootstrap/hooks/pytorch_optim.py |  0
 pprobe/toggle/hook.toggle.default       |  2 +-
 3 files changed, 20 insertions(+), 5 deletions(-)
 create mode 100644 pprobe/bootstrap/hooks/pytorch_optim.py

diff --git a/pprobe/bootstrap/hook_setup.py b/pprobe/bootstrap/hook_setup.py
index a86f746..95cfd4a 100644
--- a/pprobe/bootstrap/hook_setup.py
+++ b/pprobe/bootstrap/hook_setup.py
@@ -18,8 +18,11 @@ def __init__(self, module, module_fullname):
         ### ENABLE FIELD
         ####################
         self.pprobe_enabled = self.pprobe_toggle.get_toggle("PPROBE_ENABLE")
-        self.torch_reproduce_enabled = self.pprobe_toggle.get_toggle("TORCH_REPRODUCE")
 
+        #########################
+        ## TORCH-CONTROL
+        #########################
+        self.torch_reproduce_enabled = self.pprobe_toggle.get_toggle("TORCH_REPRODUCE")
         self.torch_catch_step_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_CATCH_STEP"
         )
@@ -27,7 +30,13 @@ def __init__(self, module, module_fullname):
             "TORCH_CATCH_LOSS"
         )
         self.torch_catch_lr_enabled = self.pprobe_toggle.get_toggle("TORCH_CATCH_LR")
+        self.torch_catch_memory_enabled = self.pprobe_toggle.get_toggle(
+            "TORCH_CATCH_MEMORY"
+        )
 
+        #########################
+        ## TORCH-SNAPSHOT
+        #########################
         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(
@@ -36,9 +45,10 @@ def __init__(self, module, module_fullname):
         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"
-        )
+
+        #########################
+        ## TORCH-TESTCASE DUMP
+        #########################
         self.torch_test_dump_op_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_TEST_DUMP_OP"
         )
@@ -48,12 +58,17 @@ def __init__(self, module, module_fullname):
         self.torch_test_dump_module_enabled = self.pprobe_toggle.get_toggle(
             "TORCH_TEST_DUMP_MODULE"
         )
+
+        #########################
+        ## TORCH-SCAN
+        #########################
         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
         ####################
diff --git a/pprobe/bootstrap/hooks/pytorch_optim.py b/pprobe/bootstrap/hooks/pytorch_optim.py
new file mode 100644
index 0000000..e69de29
diff --git a/pprobe/toggle/hook.toggle.default b/pprobe/toggle/hook.toggle.default
index 6a097bf..918825f 100644
--- a/pprobe/toggle/hook.toggle.default
+++ b/pprobe/toggle/hook.toggle.default
@@ -10,6 +10,7 @@ TORCH_REPRODUCE=True
 TORCH_CATCH_STEP=False
 TORCH_CATCH_LOSS=False
 TORCH_CATCH_LR=False
+TORCH_CATCH_MEMORY=False
 
 #########################
 ## TORCH-SNAPSHOT
@@ -18,7 +19,6 @@ TORCH_DUMP_OP=False
 TORCH_DUMP_MODULE=False
 TORCH_DUMP_DIST=False
 TORCH_DUMP_OPTIM=False
-TORCH_DUMP_MEMORY=False
 
 #########################
 ## TORCH-TESTCASE DUMP