Fix logging verbosity in BackendSimulator

ax/utils/testing/backend_simulator.py

@@ -6,7 +6,6 @@
 
 # pyre-strict
 
-import logging
 import random
 import time
 from dataclasses import dataclass
@@ -140,9 +139,6 @@ def __init__(
                 (only used for testing particular initialization cases)
             verbose_logging: If False, sets the logging level to WARNING.
         """
-        if not verbose_logging:
-            logger.setLevel(logging.WARNING)
-
         self.options: BackendSimulatorOptions = (
             BackendSimulatorOptions() if options is None else options
         )
@@ -153,6 +149,13 @@ def __init__(
         self._verbose_logging = verbose_logging
         self._create_index_to_trial_map()
 
+    def log(self, msg: str) -> None:
+        """Log at INFO level if `verbose_logging`, otherwise DEBUG."""
+        if self._verbose_logging:
+            logger.info(msg)
+        else:
+            logger.debug(msg)
+
     @property
     def num_queued(self) -> int:
         """The number of queued trials (to run as soon as capacity is available)."""
@@ -216,8 +219,8 @@ def update(self) -> None:
             self.options.internal_clock = none_throws(self.options.internal_clock) + 1
         self._update(self.time)
         state = self.state()
-        logger.info(
-            "\n-----------\n"
+        self.log(
+            msg="\n-----------\n"
             f"Updated backend simulator state (time = {self.time}):\n"
             f"** Queued:\n{format(state.queued)}\n"
             f"** Running:\n{format(state.running)}\n"
@@ -308,15 +311,15 @@ def stop_trial(self, trial_index: int) -> None:
         """
         trial_status = self.lookup_trial_index_status(trial_index)
         if trial_status is not TrialStatus.RUNNING:
-            logger.info(
-                f"Trial {trial_index} is not currently running (has status "
+            self.log(
+                msg=f"Trial {trial_index} is not currently running (has status "
                 f"{trial_status}) and cannot be stopped."
             )
         else:
             trial = self._index_to_trial_map[trial_index]
             trial.sim_completed_time = self.time
-            logger.info(
-                f"Trial {trial_index} stopped at time {trial.sim_completed_time}."
+            self.log(
+                msg=f"Trial {trial_index} stopped at time {trial.sim_completed_time}."
             )
 
     def status(self) -> SimStatus:

ax/utils/testing/tests/test_backend_simulator.py

@@ -9,6 +9,7 @@
 from unittest.mock import Mock, patch
 
 from ax.core.trial_status import TrialStatus
+from ax.utils.common.logger import get_logger
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.backend_simulator import BackendSimulator, BackendSimulatorOptions
 from ax.utils.testing.utils_testing_stubs import get_backend_simulator_with_trials
@@ -73,6 +74,20 @@ def test_backend_simulator(self, time_mock: Mock) -> None:
         self.assertEqual(sim3.num_failed, 1)
         self.assertEqual(sim3.num_completed, 0)
 
+        with self.subTest("Test logging"):
+            with self.assertLogs(
+                logger=get_logger("utils.testing.backend_simulator"), level="INFO"
+            ):
+                sim3.update()
+
+            non_verbose_simulator = BackendSimulator(
+                options=options, verbose_logging=False
+            )
+            with self.assertLogs(
+                logger=get_logger("utils.testing.backend_simulator"), level="DEBUG"
+            ):
+                non_verbose_simulator.update()
+
     def test_backend_simulator_internal_clock(self) -> None:
         sim = get_backend_simulator_with_trials()
         self.assertEqual(len(sim.all_trials), 3)