Refactor logging configuration and clean up test cases (#42)

Summary:
- Removed unused log parsing functions from `test_tritonparse.py` to streamline the test code.
- Added calls to `triton.knobs.compilation.listener` and `tritonparse.structured_logging.clear_logging_config()` to ensure proper cleanup and reset of logging configurations after tests.
- Improved test reliability by ensuring that logging settings do not persist between test runs.

These changes enhance the maintainability of the test suite and ensure that logging configurations are correctly managed during testing.

Pull Request resolved: #42

Reviewed By: Sibylau

Differential Revision: D78772087

Pulled By: FindHao

fbshipit-source-id: b44b8093d5f032f21db673766a3a2dc125cef910

Author: FindHao

tests/test_tritonparse.py

@@ -164,6 +164,7 @@ def compile_listener(
         torch.cuda.synchronize()
         assert "python_source" in trace_data
         assert "file_path" in trace_data["python_source"]
+        triton.knobs.compilation.listener = None
 
     @unittest.skipUnless(torch.cuda.is_available(), "CUDA not available")
     def test_whole_workflow(self):
@@ -222,45 +223,6 @@ def run_test_kernel(x):
         )
         print(f"Found {len(log_files)} log files in {temp_dir_logs}: {log_files}")
 
-        def parse_log_line(line: str, line_num: int) -> dict | None:
-            """Parse a single log line and extract event data"""
-            try:
-                return json.loads(line.strip())
-            except json.JSONDecodeError as e:
-                print(f"  Line {line_num}: JSON decode error - {e}")
-                return None
-
-        def process_event_data(
-            event_data: dict, line_num: int, event_counts: dict
-        ) -> None:
-            """Process event data and update counts"""
-            try:
-                event_type = event_data.get("event_type")
-                if event_type is None:
-                    return
-
-                if event_type in event_counts:
-                    event_counts[event_type] += 1
-                    print(
-                        f"  Line {line_num}: event_type = '{event_type}' (count: {event_counts[event_type]})"
-                    )
-                else:
-                    print(
-                        f"  Line {line_num}: event_type = '{event_type}' (not tracked)"
-                    )
-            except (KeyError, TypeError) as e:
-                print(f"  Line {line_num}: Data structure error - {e}")
-
-        def count_events_in_file(file_path: str, event_counts: dict) -> None:
-            """Count events in a single log file"""
-            print(f"Checking event types in: {os.path.basename(file_path)}")
-
-            with open(file_path, "r") as f:
-                for line_num, line in enumerate(f, 1):
-                    event_data = parse_log_line(line, line_num)
-                    if event_data:
-                        process_event_data(event_data, line_num, event_counts)
-
         def check_event_type_counts_in_logs(log_dir: str) -> dict:
             """Count 'launch' and unique 'compilation' events in all log files"""
             event_counts = {"launch": 0}
@@ -326,6 +288,7 @@ def check_event_type_counts_in_logs(log_dir: str) -> dict:
             # Clean up
             shutil.rmtree(temp_dir)
             print("✓ Cleaned up temporary directory")
+            tritonparse.structured_logging.clear_logging_config()
 
     @unittest.skipUnless(torch.cuda.is_available(), "CUDA not available")
     def test_complex_kernels(self):
@@ -628,6 +591,7 @@ def fused_op(a, b, c, scale_factor: float, activation: str):
             # Clean up
             shutil.rmtree(temp_dir)
             print("✓ Cleaned up temporary directory")
+            tritonparse.structured_logging.clear_logging_config()
 
 
 if __name__ == "__main__":

tritonparse/structured_logging.py

@@ -941,12 +941,12 @@ def __call__(self, metadata):
 def maybe_enable_trace_launch():
     global _trace_launch_enabled
     if TRITON_TRACE_LAUNCH and not _trace_launch_enabled:
-        import triton
+        from triton import knobs
 
         launch_hook = LaunchHookImpl()
         jit_hook = JITHookImpl()
-        triton.knobs.runtime.jit_post_compile_hook = jit_hook
-        triton.knobs.runtime.launch_enter_hook = launch_hook
+        knobs.runtime.jit_post_compile_hook = jit_hook
+        knobs.runtime.launch_enter_hook = launch_hook
 
         _trace_launch_enabled = True
 
@@ -964,7 +964,7 @@ def init_basic(trace_folder: Optional[str] = None):
     maybe_enable_debug_logging()
     if triton_trace_folder is not None and trace_folder is not None:
         log.info(
-            "Conflict settings: TRITON_TRACE is already set to %s, we will use provided trace_folder(%s) instead.",
+            "Conflict settings: triton_trace_folder is already set to %s, we will use provided trace_folder(%s) instead.",
             triton_trace_folder,
             trace_folder,
         )
@@ -995,7 +995,43 @@ def init(trace_folder: Optional[str] = None, enable_trace_launch: bool = False):
     global TRITON_TRACE_LAUNCH
     if enable_trace_launch:
         TRITON_TRACE_LAUNCH = True
-    import triton
 
     init_basic(trace_folder)
-    triton.knobs.compilation.listener = maybe_trace_triton
+    from triton import knobs
+
+    knobs.compilation.listener = maybe_trace_triton
+
+
+def clear_logging_config():
+    """
+    Clear all configurations made by init() and init_basic().
+
+    This function resets the logging handlers, global state variables,
+    and Triton knobs to their default states, effectively disabling
+    the custom tracing.
+
+    WARNING: This function is not supposed to be called unless you are sure
+    you want to clear the logging config.
+    """
+    global TRITON_TRACE_HANDLER, triton_trace_folder, _KERNEL_ALLOWLIST_PATTERNS
+    global _trace_launch_enabled
+
+    # 1. Clean up the log handler
+    if TRITON_TRACE_HANDLER is not None:
+        if TRITON_TRACE_HANDLER in triton_trace_log.handlers:
+            triton_trace_log.removeHandler(TRITON_TRACE_HANDLER)
+        TRITON_TRACE_HANDLER.close()
+        TRITON_TRACE_HANDLER = None
+
+    # 2. Reset global state variables
+    triton_trace_folder = None
+    _KERNEL_ALLOWLIST_PATTERNS = None
+    _trace_launch_enabled = False
+
+    # 3. Reset Triton knobs
+    # Check if triton was actually imported and used
+    from triton import knobs
+
+    knobs.compilation.listener = None
+    knobs.runtime.jit_post_compile_hook = None
+    knobs.runtime.launch_enter_hook = None