diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py index bed829ef6b29..9f730cbf3205 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 22a5053aee5a..a324edc1b63d 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 76173da1939f..6ec033208bd4 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 af45f0bb7e00..6ec1fecd9754 100644 --- a/tests/python/driver/tvmc/test_command_line.py +++ b/tests/python/driver/tvmc/test_command_line.py @@ -18,6 +18,8 @@ import platform import pytest import shutil +import logging +import sys from pytest_lazyfixture import lazy_fixture from unittest import mock @@ -26,6 +28,7 @@ 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)