From 64608e365b98a0c188a3d65b091983adc8dcafbd Mon Sep 17 00:00:00 2001
From: Patrick Wang <wang.patrick57@gmail.com>
Date: Sat, 19 Oct 2024 19:50:42 +0000
Subject: [PATCH] now configuring root logger globally in task.py and making it
 output different levels of stuff to the console and to the output file

---
 scripts/pat_test.sh                 |  4 +---
 task.py                             | 29 ++++++++++++++++++++++++-----
 tune/protox/env/artifact_manager.py | 19 ++++---------------
 3 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/scripts/pat_test.sh b/scripts/pat_test.sh
index 4353a5dc..f1074cd0 100755
--- a/scripts/pat_test.sh
+++ b/scripts/pat_test.sh
@@ -7,9 +7,7 @@ INTENDED_DBDATA_HARDWARE=ssd
 . ./experiments/load_per_machine_envvars.sh
 
 # space for testing. uncomment this to run individual commands from the script (copy pasting is harder because there are envvars)
-python3 task.py tune protox agent hpo tpch --scale-factor $SCALE_FACTOR --num-samples 2 --max-concurrent 2 --workload-timeout 15 --query-timeout 1 --tune-duration-during-hpo 0.01  --intended-dbdata-hardware $INTENDED_DBDATA_HARDWARE --dbdata-parent-dpath $DBDATA_PARENT_DPATH --build-space-good-for-boot
-python3 task.py tune protox agent tune tpch --scale-factor $SCALE_FACTOR --tune-duration-during-tune 0.02
-python3 task.py tune protox agent replay tpch --scale-factor $SCALE_FACTOR
+python3 task.py benchmark tpch data $SCALE_FACTOR
 exit 0
 
 # benchmark
diff --git a/task.py b/task.py
index 37ac3a69..640839a5 100644
--- a/task.py
+++ b/task.py
@@ -19,14 +19,33 @@
 def task(ctx: click.Context) -> None:
     """💩💩💩 CMU-DB Database Gym: github.com/cmu-db/dbgym 💩💩💩"""
     dbgym_config_path = Path(os.getenv("DBGYM_CONFIG_PATH", "dbgym_config.yaml"))
-    ctx.obj = DBGymConfig(dbgym_config_path)
+    dbgym_cfg = DBGymConfig(dbgym_config_path)
+    ctx.obj = dbgym_cfg
 
+    # The root logger is set up globally here. Do not reconfigure the root logger anywhere else.
+    _set_up_root_logger(dbgym_cfg)
 
-if __name__ == "__main__":
-    logging.basicConfig(
-        format="%(asctime)s:%(name)s:%(levelname)s - %(message)s", level=logging.INFO
-    )
 
+def _set_up_root_logger(dbgym_cfg: DBGymConfig) -> None:
+    format = "%(levelname)s:%(asctime)s [%(filename)s:%(lineno)s]  %(message)s"
+    
+    # Set this so that the root logger captures everything.
+    logging.getLogger().setLevel(logging.DEBUG)
+
+    # Only make it output warnings or higher to the console.
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.WARNING)
+    console_handler.setFormatter(logging.Formatter(format))
+    logging.getLogger().addHandler(console_handler)
+
+    # Let it output everything to the output file.
+    file_handler = logging.FileHandler(dbgym_cfg.cur_task_runs_artifacts_path(mkdir=True) / "output.log")
+    file_handler.setFormatter(logging.Formatter(format))
+    file_handler.setLevel(logging.DEBUG)
+    logging.getLogger().addHandler(file_handler)
+
+
+if __name__ == "__main__":
     task.add_command(benchmark_group)
     task.add_command(manage_group)
     task.add_command(dbms_group)
diff --git a/tune/protox/env/artifact_manager.py b/tune/protox/env/artifact_manager.py
index 73ae0274..1bbbafda 100644
--- a/tune/protox/env/artifact_manager.py
+++ b/tune/protox/env/artifact_manager.py
@@ -62,12 +62,10 @@ class ArtifactManager(object):
     """
     This class manages the following artifacts of Proto-X: info for replaying and TensorBoard output.
 
-    Initializing this class sets up the root logger. However, to use the root logger, you should
-    directly use the logging library.
+    Note that the root logger is set up globally inside the upper-most task.py file. Do not reconfigure it here.
     """
 
     # The output log is the file that the root logger writes to
-    OUTPUT_LOG_FNAME = "output.log"
     REPLAY_INFO_LOG_FNAME = "replay_info.log"
     REPLAY_LOGGER_NAME = "replay_logger"
 
@@ -82,22 +80,13 @@ def __init__(
         self.tuning_steps_dpath = self.log_dpath / "tuning_steps"
         self.tuning_steps_dpath.mkdir(parents=True, exist_ok=True)
 
-        # Setup the root and replay loggers
-        formatter = "%(levelname)s:%(asctime)s [%(filename)s:%(lineno)s]  %(message)s"
-        logging.basicConfig(format=formatter, level=logging.DEBUG, force=True)
-        output_log_handler = logging.FileHandler(
-            self.log_dpath / ArtifactManager.OUTPUT_LOG_FNAME
-        )
-        output_log_handler.setFormatter(logging.Formatter(formatter))
-        output_log_handler.setLevel(logging.DEBUG)
-        logging.getLogger().addHandler(output_log_handler)
-
+        # Create replay logger
         replay_info_log_handler = logging.FileHandler(
             self.tuning_steps_dpath / ArtifactManager.REPLAY_INFO_LOG_FNAME
         )
-        replay_info_log_handler.setFormatter(logging.Formatter(formatter))
+        replay_info_log_handler.setFormatter(logging.Formatter("%(levelname)s:%(asctime)s [%(filename)s:%(lineno)s]  %(message)s"))
         replay_info_log_handler.setLevel(logging.INFO)
-        logging.getLogger(ArtifactManager.REPLAY_LOGGER_NAME)
+        logging.getLogger(ArtifactManager.REPLAY_LOGGER_NAME).addHandler(replay_info_log_handler)
 
         # Setup the writer.
         self.writer: Union[SummaryWriter, None] = None