[tvm] branch main updated: [TVMC] Fix logging in TVMC (#14175)

[lu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

lukhut pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new f6b75792c5 [TVMC] Fix logging in TVMC (#14175)
f6b75792c5 is described below

commit f6b75792c5a3ad88ad4819db657c675a6f982ec7
Author: Elen Kalda <el...@arm.com>
AuthorDate: Wed Mar 15 15:31:31 2023 +0000

    [TVMC] Fix logging in TVMC (#14175)
    
    Three logger related changes in this patch:
    * Currently we don't set the output stream on the Python logger, so
    it defaults to sys.stderr, which means we only get some logger output
    when the command fails. So set the output stream to sys.stdout
    * Currently we can add -v flag to anywhere in the command line for
    tvmc compile, but only between tvmc and run/tune for run and tune.
    Unify the behaviour such that we can add the flag anywhere on the
    command line.
    * Set the effective upper bound of -vs to 3 as 4 could result in
    NOTSET which would not output anything.
---
 python/tvm/driver/tvmc/autotuner.py           |  1 +
 python/tvm/driver/tvmc/main.py                |  7 ++-
 python/tvm/driver/tvmc/runner.py              |  1 +
 tests/python/driver/tvmc/test_command_line.py | 62 +++++++++++++++++++++++++++
 4 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py
index bed829ef6b..9f730cbf32 100644
--- a/python/tvm/driver/tvmc/autotuner.py
+++ b/python/tvm/driver/tvmc/autotuner.py
@@ -87,6 +87,7 @@ def add_tune_parser(subparsers, _, json_params):
         required=True,
         help="output file to store the tuning records for the tuning process",
     )
+    parser.add_argument("-v", "--verbose", action="count", default=0, help="increase verbosity.")
     parser.add_argument(
         "--parallel",
         default=4,
diff --git a/python/tvm/driver/tvmc/main.py b/python/tvm/driver/tvmc/main.py
index 22a5053aee..a324edc1b6 100644
--- a/python/tvm/driver/tvmc/main.py
+++ b/python/tvm/driver/tvmc/main.py
@@ -85,9 +85,12 @@ def _main(argv):
     parser.add_argument("-h", "--help", action="help", help="show this help message and exit.")
 
     args = parser.parse_args(argv)
-    if args.verbose > 4:
-        args.verbose = 4
+    if args.verbose > 3:
+        args.verbose = 3
 
+    # See the meaning of the logging levels at
+    # https://docs.python.org/3/library/logging.html#logging-levels
+    logging.basicConfig(stream=sys.stdout)
     logging.getLogger("TVMC").setLevel(40 - args.verbose * 10)
 
     if args.version:
diff --git a/python/tvm/driver/tvmc/runner.py b/python/tvm/driver/tvmc/runner.py
index 76173da193..6ec033208b 100644
--- a/python/tvm/driver/tvmc/runner.py
+++ b/python/tvm/driver/tvmc/runner.py
@@ -109,6 +109,7 @@ def add_run_parser(subparsers, main_parser, json_params):
         "Profiling may also have an impact on inference time, "
         "making it take longer to be generated. (non-micro devices only)",
     )
+    parser.add_argument("-v", "--verbose", action="count", default=0, help="increase verbosity.")
     parser.add_argument(
         "--end-to-end",
         action="store_true",
diff --git a/tests/python/driver/tvmc/test_command_line.py b/tests/python/driver/tvmc/test_command_line.py
index af45f0bb7e..6ec1fecd97 100644
--- a/tests/python/driver/tvmc/test_command_line.py
+++ b/tests/python/driver/tvmc/test_command_line.py
@@ -18,6 +18,8 @@ import os
 import platform
 import pytest
 import shutil
+import logging
+import sys
 
 from pytest_lazyfixture import lazy_fixture
 from unittest import mock
@@ -26,6 +28,7 @@ import tvm
 from tvm.driver.tvmc.main import _main
 from tvm.driver.tvmc.model import TVMCException
 from tvm.driver.tvmc import compiler
+from unittest.mock import MagicMock
 
 
 @pytest.mark.skipif(
@@ -210,3 +213,62 @@ def test_tvmc_compile_input_model(mock_compile_model, tmpdir_factory, model):
     _main(run_arg)
 
     mock_compile_model.assert_called_once()
+
+
+def test_tvmc_logger(caplog, tmpdir_factory, keras_simple):
+    pytest.importorskip("tensorflow")
+    tmpdir = tmpdir_factory.mktemp("out")
+
+    # TUNE
+    log_path = os.path.join(tmpdir, "records.json")
+    tune_cmd = f"tvmc tune --target llvm -vvvv --output {log_path} " f"--trials 2 {keras_simple}"
+
+    tuning_args = tune_cmd.split(" ")[1:]
+    _main(tuning_args)
+
+    # Check that we log during tvmc tune
+    for log_str in ("DEBUG", "INFO", "WARNING", "TVMC"):
+        assert log_str in caplog.text
+
+    caplog.clear()
+
+    # COMPILE
+    module_file = os.path.join(tmpdir, "m.tar")
+    compile_cmd = f"tvmc compile --target 'llvm' {keras_simple} -vvvv --output {module_file}"
+
+    compile_args = compile_cmd.split(" ")[1:]
+    _main(compile_args)
+
+    # Check that we log during tvmc compile
+    for log_str in ("DEBUG", "WARNING", "TVMC"):
+        assert log_str in caplog.text
+
+    caplog.clear()
+
+    # RUN
+    run_cmd = f"tvmc run -vvvv {module_file}"
+
+    run_args = run_cmd.split(" ")[1:]
+    _main(run_args)
+
+    # Check that we log during tvmc run
+    for log_str in ("DEBUG", "TVMC"):
+        assert log_str in caplog.text
+
+
+# Unfortunately pytest seems to intercept the logging output, so we can't test whether it
+# actually writes the logging output to sys.stdout, but we can test that we call
+# logging.basicConfig with the correct arguments
+def test_tvmc_logger_set_basicConfig(monkeypatch, tmpdir_factory, keras_simple):
+    pytest.importorskip("tensorflow")
+    mock_basicConfig = MagicMock()
+    monkeypatch.setattr(logging, "basicConfig", mock_basicConfig)
+
+    # Run a random tvmc command
+    tmpdir = tmpdir_factory.mktemp("out")
+    module_file = os.path.join(tmpdir, "m.tar")
+    compile_cmd = f"tvmc compile --target 'llvm' {keras_simple} -vvvv --output {module_file}"
+    compile_args = compile_cmd.split(" ")[1:]
+    _main(compile_args)
+
+    mock_basicConfig.assert_called_with(stream=sys.stdout)