From d13ec6a9b26fd09bfda685f5589d8c32fd4ec908 Mon Sep 17 00:00:00 2001
From: Bangtian Liu <liubangtian@gmail.com>
Date: Tue, 14 Jan 2025 10:41:44 -0600
Subject: [PATCH] clean up the code

Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
---
 tuner/examples/simple/simple_tuner.py | 17 ++++++++---------
 tuner/tuner/common.py                 |  3 ---
 tuner/tuner/libtuner.py               |  1 -
 3 files changed, 8 insertions(+), 13 deletions(-)

diff --git a/tuner/examples/simple/simple_tuner.py b/tuner/examples/simple/simple_tuner.py
index b98b083ae..800cd0f34 100644
--- a/tuner/examples/simple/simple_tuner.py
+++ b/tuner/examples/simple/simple_tuner.py
@@ -102,7 +102,7 @@ def main():
     )
     print("Generating candidate tuning specs...")
     with TunerContext(logger=root_logger) as tuner_context:
-        tuner_context.add_logging_handler(summary_handler)
+        tuner_context.logger.addHandler(summary_handler)
         simple_tuner = SimpleTuner(tuner_context)
         candidates = libtuner.generate_candidate_specs(
             args, path_config, candidate_trackers, simple_tuner
@@ -121,11 +121,11 @@ def main():
         if stop_after_phase == libtuner.ExecutionPhases.compile_dispatches:
             return
 
-        print("Benchmarking compiled dispatch candidates...")
-        logging.info(f"Summarization about top dispatch candidates:")
+        message = "Benchmarking compiled dispatch candidates..."
+        print(message)
+        logging.info(message)
         simple_tuner.benchmark_flags = ["--input=1", "--benchmark_repetitions=3"]
         top_candidates = libtuner.benchmark(
-            tuner_context,
             args,
             path_config,
             compiled_candidates,
@@ -153,12 +153,12 @@ def main():
         if stop_after_phase == libtuner.ExecutionPhases.compile_models:
             return
 
-        print("Benchmarking compiled model candidates...")
-        logging.info(f"Summarization about top model candidates:")
+        message = "Benchmarking compiled model candidates..."
+        print(message)
+        logging.info(message)
         simple_tuner.benchmark_flags = model_benchmark_flags
         simple_tuner.benchmark_timeout = 60
         top_model_candidates = libtuner.benchmark(
-            tuner_context,
             args,
             path_config,
             compiled_model_candidates,
@@ -173,6 +173,5 @@ def main():
 
         print("Check the detailed execution logs in:")
         print(path_config.run_log.resolve())
-
-        print("Check the tuning results in:")
+        print("Check the summary in:")
         print(path_config.summary_log.resolve())
diff --git a/tuner/tuner/common.py b/tuner/tuner/common.py
index 5719c9304..05820526b 100644
--- a/tuner/tuner/common.py
+++ b/tuner/tuner/common.py
@@ -46,9 +46,6 @@ def __enter__(self) -> "TunerContext":
         self.mlir_ctx.__enter__()
         return self
 
-    def add_logging_handler(self, handler: logging.Handler) -> None:
-        self.logger.addHandler(handler)
-
     def __exit__(
         self,
         exc_type: type[BaseException] | None,
diff --git a/tuner/tuner/libtuner.py b/tuner/tuner/libtuner.py
index 90fb23467..0642daa2c 100644
--- a/tuner/tuner/libtuner.py
+++ b/tuner/tuner/libtuner.py
@@ -895,7 +895,6 @@ def get_speedup(result: BenchmarkResult) -> float:
 
 
 def benchmark(
-    tuner_context: TunerContext,
     args: argparse.Namespace,
     path_config: PathConfig,
     compiled_candidates: list[int],