[Build] Added --pdb flag to build.py, drop into pdb on error (mlc-ai#1017)

This commit adds an optional `--pdb` flag to the `build.py` script. If
passed, any exception raised that would otherwise terminate the script
will first enter a pdb post-mortem, allowing the error to be
inspected.

Author: Lunderberg

mlc_llm/build.py

@@ -1,13 +1,46 @@
 """Script for building/compiling models."""
+import contextlib
+import sys
+
 from mlc_llm import core
 
+
+@contextlib.contextmanager
+def debug_on_except():
+    try:
+        yield
+    finally:
+        if sys.exc_info() == (None, None, None):
+            return
+
+        import traceback
+
+        try:
+            import ipdb as pdb
+        except ImportError:
+            import pdb
+
+        traceback.print_exc()
+        pdb.post_mortem()
+
+
 def main():
     """Main method for building model from command line."""
     empty_args = core.convert_build_args_to_argparser()  # Create new ArgumentParser
     parsed_args = empty_args.parse_args()  # Parse through command line
-    # Post processing of arguments
-    parsed_args = core._parse_args(parsed_args)  # pylint: disable=protected-access
-    core.build_model_from_args(parsed_args)
+
+    with contextlib.ExitStack() as stack:
+        # Enter an exception-catching context before post-processing
+        # the arguments, in case the post-processing itself raises an
+        # exception.
+        if parsed_args.pdb:
+            stack.enter_context(debug_on_except())
+
+        # Post processing of arguments
+        parsed_args = core._parse_args(parsed_args)  # pylint: disable=protected-access
+
+        core.build_model_from_args(parsed_args)
+
 
 if __name__ == "__main__":
     main()

mlc_llm/core.py

@@ -236,6 +236,13 @@ class BuildArgs:
             "action": "store_true",
         },
     )
+    pdb: bool = field(
+        default=False,
+        metadata={
+            "help": ("If set, drop into a pdb debugger on error"),
+            "action": "store_true",
+        },
+    )
 
 
 def convert_build_args_to_argparser() -> argparse.ArgumentParser: