Fixes to kernel agent backend tests (#46)

Author: bertmaher

BackendBench/backends.py

@@ -1,7 +1,8 @@
-import os
 import importlib.util
 import logging
-from typing import Dict, Callable, List
+import os
+from typing import Callable, Dict, List
+
 import torch
 
 logger = logging.getLogger(__name__)
@@ -397,7 +398,8 @@ def __init__(self) -> None:
         # Create README for this run
         readme_path = os.path.join(self.kernels_dir, "README.md")
         with open(readme_path, "w") as f:
-            f.write(f"""# Generated Kernels - {timestamp}
+            f.write(
+                f"""# Generated Kernels - {timestamp}
 
 This directory contains PyTorch/Triton kernels generated by the LLM Backend.
 
@@ -413,7 +415,8 @@ def __init__(self) -> None:
 
 ## Usage
 You can inspect these files to debug kernel generation, manually test implementations, or understand what the LLM produced.
-""")
+"""
+            )
 
         print(f"Saving generated kernels to: {self.kernels_dir}")
 
@@ -521,8 +524,8 @@ def test_kernel_correctness(
                     f.write(full_code)
                 print(f"Saved kernel to: {kernel_file}")
 
-            import sys
             import importlib.util
+            import sys
 
             spec = importlib.util.spec_from_file_location(
                 f"test_kernel_{op_name}_{attempt}", kernel_file
@@ -633,7 +636,8 @@ def __init__(self) -> None:
         # Create README for this run
         readme_path = os.path.join(self.kernels_dir, "README.md")
         with open(readme_path, "w") as f:
-            f.write(f"""# Generated Kernels - KernelAgent - {timestamp}
+            f.write(
+                f"""# Generated Kernels - KernelAgent - {timestamp}
 
 This directory contains PyTorch/Triton kernels generated by the KernelAgent Backend.
 
@@ -656,7 +660,8 @@ def __init__(self) -> None:
 ## Usage
 You can inspect these files to debug kernel generation, analyze the parallel worker outputs,
 or understand the sophisticated generation process used by KernelAgent.
-""")
+"""
+            )
 
         print(f"Saving KernelAgent generated kernels to: {self.kernels_dir}")
 
@@ -688,7 +693,9 @@ def _get_kernel_agent(self):
                 os.makedirs(agent_log_dir, exist_ok=True)
 
                 self.kernel_agent = TritonKernelAgent(
-                    log_dir=agent_log_dir, num_workers=self.num_workers, max_rounds=self.max_rounds
+                    log_dir=agent_log_dir,
+                    num_workers=self.num_workers,
+                    max_rounds=self.max_rounds,
                 )
 
                 print(f"✓ KernelAgent initialized with log directory: {agent_log_dir}")

test/test_backends.py

@@ -1,7 +1,13 @@
+from unittest.mock import Mock, patch
+
 import pytest
 import torch
-from unittest.mock import Mock, patch
-from BackendBench.backends import AtenBackend, FlagGemsBackend, LLMBackend, KernelAgentBackend
+from BackendBench.backends import (
+    AtenBackend,
+    FlagGemsBackend,
+    KernelAgentBackend,
+    LLMBackend,
+)
 
 try:
     import importlib.util
@@ -11,9 +17,9 @@
     HAS_FLAG_GEMS = False
 
 try:
-    import sys
-    import os
     import importlib.util
+    import os
+    import sys
 
     kernel_agent_path = os.path.join(os.path.dirname(__file__), "..", "KernelAgent")
     sys.path.insert(0, os.path.abspath(kernel_agent_path))
@@ -159,35 +165,38 @@ def generated_kernel(x):
 class TestKernelAgentBackend:
     @pytest.mark.skipif(not HAS_KERNEL_AGENT, reason="KernelAgent not available")
     def test_kernel_agent_backend_initialization(self):
-        with patch("os.makedirs"):
-            backend = KernelAgentBackend()
-            assert backend.name == "kernel_agent"
-            assert "kernel_agent_run_" in backend.kernels_dir
-            assert backend.num_workers == 4  # default value
-            assert backend.max_rounds == 10  # default value
+        backend = KernelAgentBackend()
+        assert backend.name == "kernel_agent"
+        assert "kernel_agent_run_" in backend.kernels_dir
+        assert backend.num_workers == 4  # default value
+        assert backend.max_rounds == 10  # default value
 
     @pytest.mark.skipif(not HAS_KERNEL_AGENT, reason="KernelAgent not available")
     def test_kernel_agent_backend_set_config(self):
-        with patch("os.makedirs"):
-            backend = KernelAgentBackend()
+        backend = KernelAgentBackend()
 
-            backend.set_config(num_workers=8, max_rounds=20)
+        backend.set_config(num_workers=8, max_rounds=20)
 
-            assert backend.num_workers == 8
-            assert backend.max_rounds == 20
+        assert backend.num_workers == 8
+        assert backend.max_rounds == 20
 
     @pytest.mark.skipif(not HAS_KERNEL_AGENT, reason="KernelAgent not available")
     def test_kernel_agent_backend_generate_kernel(self):
         with (
-            patch("os.makedirs"),
             patch("triton_kernel_agent.TritonKernelAgent") as mock_kernel_agent_class,
         ):
             backend = KernelAgentBackend()
 
             mock_agent = Mock()
             mock_kernel_agent_class.return_value = mock_agent
 
-            mock_agent.generate_kernel.return_value = (True, "def kernel(): pass")
+            mock_agent.generate_kernel.return_value = {
+                "success": True,
+                "kernel_code": "def kernel(): pass",
+                "rounds": 1,
+                "session_dir": "test_session_dir",
+                "worker_id": 0,
+            }
 
             mock_op = Mock()
             mock_op.__str__ = Mock(return_value="test_op")
@@ -205,9 +214,8 @@ def test_backend_polymorphism(self):
         backends.append(AtenBackend())
         with patch("BackendBench.backends.flag_gems"):
             backends.append(FlagGemsBackend())
-        with patch("os.makedirs"):
-            backends.append(LLMBackend())
-            backends.append(KernelAgentBackend())
+        backends.append(LLMBackend())
+        backends.append(KernelAgentBackend())
         for backend in backends:
             assert hasattr(backend, "name")
             assert hasattr(backend, "__contains__")