KUNLUNXIN][BUF-FIX]fix that transformer engine disable benchmark logger

the transformer_engine may disable flagGems benchMark logging, which lead to "pytest --record log" failure

Author: dongjibin1996

benchmark/conftest.py

@@ -22,6 +22,8 @@
 
 device = flag_gems.device
 vendor_name = flag_gems.vendor_name
+recordLogger = logging.getLogger("flag_gems.benchmark.record")
+recordLogger.propagate = False
 
 
 class BenchConfig:
@@ -165,12 +167,21 @@ def pytest_configure(config):
             for arg in config.invocation_params.args
         ]
 
-        logging.basicConfig(
-            filename="result_{}.log".format("_".join(cmd_args)).replace("_-", "-"),
-            filemode="w",
-            level=logging.INFO,
-            format="[%(levelname)s] %(message)s",
-        )
+        log_file = "result_{}.log".format("_".join(cmd_args)).replace("_-", "-")
+
+        for h in list(recordLogger.handlers):
+            recordLogger.removeHandler(h)
+            try:
+                h.close()
+            except Exception:
+                pass
+
+        handler = logging.FileHandler(log_file, mode="w", encoding="utf-8")
+        handler.setLevel(logging.INFO)
+        handler.setFormatter(logging.Formatter("[%(levelname)s] %(message)s"))
+        recordLogger.addHandler(handler)
+        recordLogger.setLevel(logging.INFO)
+        recordLogger.info("Benchmark record logger enabled")
 
 
 BUILTIN_MARKS = {

benchmark/performance_utils.py

@@ -1,6 +1,5 @@
 import gc
 import importlib
-import logging
 import os
 import time
 from typing import Any, Generator, List, Optional, Tuple
@@ -26,7 +25,7 @@
     OperationAttribute,
     check_metric_dependencies,
 )
-from .conftest import Config
+from .conftest import Config, recordLogger
 
 torch_backend_device = flag_gems.runtime.torch_backend_device
 torch_device_fn = flag_gems.runtime.torch_device_fn
@@ -372,7 +371,7 @@ def run(self):
                 shape_desc=self.shape_desc,
             )
             print(attri)
-            logging.info(attri.to_dict())
+            recordLogger.info(attri.to_dict())
             return
         self.init_user_config()
         for dtype in self.to_bench_dtypes:
@@ -425,7 +424,7 @@ def run(self):
                 result=metrics,
             )
             print(result)
-            logging.info(result.to_json())
+            recordLogger.info(result.to_json())
 
 
 class GenericBenchmark(Benchmark):

benchmark/test_transformer_engine_perf.py

@@ -8,6 +8,8 @@
 try:
     from transformer_engine.pytorch import cpp_extensions as tex
 
+    # note: the transformer_engine has its' own logging which may have conflict with flag_gems logging
+
     TE_AVAILABLE = True
 except ImportError:
     TE_AVAILABLE = False

tests/test_DSA/test_bin_topk.py

@@ -121,7 +121,7 @@ def debug_topk_results(actual, expected, inputs, test_name=""):
         print(f"  Actual indices: {sorted(actual_set)[:m]}...")  # Only show first 10
         print(f"  Expected indices: {sorted(expected_set)[:m]}...")
         print(
-            f"  Intersection: {len(intersection)}/{len(expected_set)} = {len(intersection)/len(expected_set):.4f}"
+            f"  Intersection: {len(intersection)}/{len(expected_set)} = {len(intersection) / len(expected_set):.4f}"
         )
 
         # Check quality of actually selected values