Configure logging for CLI.

Author: superdosh

src/modelbench/cli.py

@@ -2,7 +2,6 @@
 import faulthandler
 import io
 import json
-import logging
 import os
 
 # silence Together's upgrade message, as the new library is not out of beta
@@ -14,25 +13,23 @@
 import sys
 from datetime import datetime, timezone
 from functools import wraps
+
 import click
 import termcolor
 from click import echo
+from modellogger.log_config import configure_logging
 from rich.console import Console
 from rich.table import Table
 
 import modelgauge.annotators.cheval.registration  # noqa: F401
 from modelbench.benchmark_runner import BenchmarkRun, BenchmarkRunner, JsonRunTracker, TqdmRunTracker
 from modelbench.benchmarks import GeneralPurposeAiChatBenchmarkV1, SecurityBenchmark
-from modelbench.consistency_checker import (
-    ConsistencyChecker,
-    summarize_consistency_check_results,
-)
+from modelbench.consistency_checker import ConsistencyChecker, summarize_consistency_check_results
 from modelbench.record import dump_json
 from modelbench.standards import Standards
 from modelgauge.config import load_secrets_from_config, write_default_config
 from modelgauge.load_namespaces import load_namespaces
 from modelgauge.locales import DEFAULT_LOCALE, LOCALES
-from modellogger.log_config import get_logger
 from modelgauge.monitoring import PROMETHEUS
 from modelgauge.preflight import check_secrets, make_sut
 from modelgauge.prompt_sets import GENERAL_PROMPT_SETS, SECURITY_JAILBREAK_PROMPT_SETS
@@ -125,7 +122,7 @@ def cli() -> None:
     log_dir = pathlib.Path("run/logs")
     log_dir.mkdir(exist_ok=True, parents=True)
     filename = log_dir / f'modelbench-{datetime.now().strftime("%y%m%d-%H%M%S")}.log'
-    logging.basicConfig(level=logging.DEBUG, handlers=[get_file_logging_handler(filename)], force=True)
+    configure_logging(app_name="modelbench", log_file=filename)
     write_default_config()
     load_namespaces(disable_progress_bar=True)
 

src/modelgauge/cli.py

@@ -1,10 +1,10 @@
-import logging
 import os
 import pathlib
 import warnings
 from typing import Optional
 
 import click
+from modellogger.log_config import get_logger
 
 from modelgauge.annotator import Annotator
 from modelgauge.annotator_registry import ANNOTATORS
@@ -35,7 +35,7 @@
 from modelgauge.sut_registry import SUTS
 from modelgauge.test_registry import TESTS
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 @cli.command(name="list")
@@ -338,8 +338,6 @@ def run_job(
     If running a SUT, the file must have 'UID' and 'Text' columns. The output will be saved to a CSV file.
     If running ONLY annotators, the file must have 'UID', 'Prompt', 'SUT', and 'Response' columns. The output will be saved to a json lines file.
     """
-    logging.basicConfig(level=logging.DEBUG if debug else logging.INFO)
-
     # TODO: break this function up. It's branching too much
     # make sure the job has everything it needs to run
     secrets = load_secrets_from_config()

src/modelgauge/command_line.py

@@ -3,6 +3,8 @@
 import sys
 
 import click
+from modellogger.log_config import configure_logging
+
 from modelgauge.annotator_registry import ANNOTATORS
 from modelgauge.cli_lazy import LOAD_ALL, LazyModuleImportGroup
 from modelgauge.config import write_default_config
@@ -26,7 +28,7 @@
 def cli():
     """Run the ModelGauge library from the command line."""
     # To add a command, decorate your function with @cli.command().
-
+    configure_logging(app_name="modelgauge")
     # Always create the config directory if it doesn't already exist.
     write_default_config()