[llvm] 99ea357 - [MLGO] Fix logging verbosity in scripts (#107818)

[via llvm-commits]

Author: Aiden Grossman
Date: 2024-09-09T11:34:53-07:00
New Revision: 99ea357f7b5e7e01e42b8d68dd211dc304b3115b

URL: https://github.com/llvm/llvm-project/commit/99ea357f7b5e7e01e42b8d68dd211dc304b3115b
DIFF: https://github.com/llvm/llvm-project/commit/99ea357f7b5e7e01e42b8d68dd211dc304b3115b.diff

LOG: [MLGO] Fix logging verbosity in scripts (#107818)

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.

Added: 
    

Modified: 
    llvm/utils/mlgo-utils/mlgo/corpus/combine_training_corpus.py
    llvm/utils/mlgo-utils/mlgo/corpus/extract_ir.py

Removed: 
    


################################################################################
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..90cf51cad75ed5 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,22 @@ def parse_args_and_run():
     parser.add_argument(
         "--root_dir", type=str, help="The root dir of module paths to combine."
     )
+    # TODO(#107898): Refactor this into a common location.
+    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 8900ee99638387..5edb429241d0c0 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
@@ -112,11 +111,22 @@ def parse_args_and_run():
         default=".llvmbc",
         nargs="?",
     )
+    # TODO(#107898): Refactor this into a common location.
+    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")