[MLGO] Fix logging verbosity in scripts

[boomanaiden154]

This patch fixes issues related to logging verbosity in the MLGO python scripts. This was an oversight when converting from absl.logging to the python logging API as absl natively supports a --verbosity flag to set the desired logging level. This patch adds a flag to support similar functionality in Python's logging library and additionally updates docstrings where relevant to point to the new values.

[llvmbot]

@llvm/pr-subscribers-mlgo

Author: Aiden Grossman (boomanaiden154)

Changes

This patch fixes issues related to logging verbosity in the MLGO python scripts. This was an oversight when converting from absl.logging to the python logging API as absl natively supports a --verbosity flag to set the desired logging level. This patch adds a flag to support similar functionality in Python's logging library and additionally updates docstrings where relevant to point to the new values.

---

Full diff: https://github.com/llvm/llvm-project/pull/107818.diff

2 Files Affected:

• (modified) llvm/utils/mlgo-utils/mlgo/corpus/combine_training_corpus.py (+11)
• (modified) llvm/utils/mlgo-utils/mlgo/corpus/extract_ir.py (+13-4)

diff --git a/llvm/utils/mlgo-utils/mlgo/corpus/combine_training_corpus.py b/llvm/utils/mlgo-utils/mlgo/corpus/combine_training_corpus.py
index 3b2077b4c0e0e6..048876ab6224af 100644
--- a/llvm/utils/mlgo-utils/mlgo/corpus/combine_training_corpus.py
+++ b/llvm/utils/mlgo-utils/mlgo/corpus/combine_training_corpus.py
@@ -24,6 +24,7 @@
 """
 
 import argparse
+import logging
 
 from mlgo.corpus import combine_training_corpus_lib
 
@@ -35,11 +36,21 @@ def parse_args_and_run():
     parser.add_argument(
         "--root_dir", type=str, help="The root dir of module paths to combine."
     )
+    parser.add_argument(
+        "--verbosity",
+        type=str,
+        help="The verbosity level to use for logging",
+        default="INFO",
+        nargs="?",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+    )
     args = parser.parse_args()
     main(args)
 
 
 def main(args):
+    logging.basicConfig(level=args.verbosity)
+
     combine_training_corpus_lib.combine_corpus(args.root_dir)
 
 
diff --git a/llvm/utils/mlgo-utils/mlgo/corpus/extract_ir.py b/llvm/utils/mlgo-utils/mlgo/corpus/extract_ir.py
index a7d52daaedba3b..a8ad79e636a6ca 100644
--- a/llvm/utils/mlgo-utils/mlgo/corpus/extract_ir.py
+++ b/llvm/utils/mlgo-utils/mlgo/corpus/extract_ir.py
@@ -18,10 +18,9 @@
 In a local ThinLTO case, the compilation is assumedto have been performed
 specifying -Wl,--save-temps=import -Wl,--thinlto-emit-index-files
 
-To change the logging verbosity, pass an integer representing the desired
-verbosity to the --verbosity flag. Use 0 for all logs, status information,
-and detailed debug information, -1 for solely warnings, and -2 to not produce
-any output.
+To change the logging verbosity, set the --verbosity flag to the desired level.
+Setting it to a specific level will enable all messages at that level and
+higher. Exact values can be found by invoking the script with --help.
 """
 
 import argparse
@@ -113,11 +112,21 @@ def parse_args_and_run():
         default=".llvmbc",
         nargs="?",
     )
+    parser.add_argument(
+        "--verbosity",
+        type=str,
+        help="The verbosity level to use for logging",
+        default="INFO",
+        nargs="?",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+    )
     args = parser.parse_args()
     main(args)
 
 
 def main(args):
+    logging.basicConfig(level=args.verbosity)
+
     objs = []
     if args.input is not None and args.thinlto_build == "local":
         raise ValueError("--thinlto_build=local cannot be run with --input")