add: adding excepthooks to instrumented program when --debug-mode is set

Author: Essoz

mldaikon/collect_trace.py

@@ -306,7 +306,7 @@ def is_path_md_output_dir(output_dir: str) -> bool:
         "-d",
         "--debug-mode",
         action="store_true",
-        help="Enable debug mode for the program",
+        help="Enable debug mode for the program, insert an exception hook into the program to log the entire stack trace",
     )
 
     ## instrumentor configs

mldaikon/instrumentor/source_file.py

@@ -442,6 +442,13 @@ def instrument_file(
 os.environ['ML_DAIKON_OUTPUT_DIR'] = "{output_dir}"
 """
 
+    debug_hook_code = """
+from mldaikon.utils import register_custom_excepthook
+if os.environ.get("ML_DAIKON_DEBUG") == "1":
+    print("ML_DAIKON_DEBUG is set to 1, registering custom excepthook")
+    register_custom_excepthook(True)
+"""
+
     # general config update
     general_config_update = f"""
 import mldaikon.config.config as general_config
@@ -469,7 +476,11 @@ def instrument_file(
     # HACK: this is a hack to attach the logging code to the instrumented source after the __future__ imports
     code_head, code_tail = get_code_head_and_tail(instrumented_source)
     instrumented_source = (
-        code_head + logging_start_code + general_config_update + code_tail
+        code_head
+        + logging_start_code
+        + debug_hook_code
+        + general_config_update
+        + code_tail
     )
 
     return instrumented_source

mldaikon/utils.py

@@ -103,7 +103,17 @@ def thread_excepthook(args):
     raise exc_type(exc_value) from None
 
 
-def register_custom_excepthook():
+def register_custom_excepthook(add_file_handler=False):
+    if add_file_handler:
+        file_handler = logging.FileHandler("debug.log")
+        file_handler.setLevel(logging.DEBUG)
+        formatter = logging.Formatter(
+            "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+        )
+        file_handler.setFormatter(formatter)
+        # add to the "mldaikon" logger and the "threading" logger
+        logging.getLogger("mldaikon").addHandler(file_handler)
+        logging.getLogger("threading").addHandler(file_handler)
     sys.excepthook = handle_excepthook
     threading.excepthook = thread_excepthook