Added Error Logging

BackendBench/agent_errors.py

@@ -0,0 +1,16 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD 3-Clause license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+class AgentError(Exception):
+    """
+    Exception raised for errors related to LLM/agent failures,
+    such as rate limits, empty code, bad formatting, or API issues.
+    """
+
+    def __init__(self, message: str):
+        super().__init__(message)
+        self.message = message

BackendBench/backends/kernel_agent.py

@@ -8,6 +8,7 @@
 import os
 from typing import Callable, Dict
 
+from BackendBench.agent_errors import AgentError
 from BackendBench.utils import compile_kernel_from_string
 
 from .base import Backend
@@ -222,28 +223,27 @@ def generate_kernel_with_agent(self, op, op_name: str) -> tuple[str, bool]:
         """
         try:
             agent = self._get_kernel_agent()
-
-            # Create problem description
             problem_description = self._create_problem_description_from_op(op, op_name)
-
             print(
                 f"🚀 Generating {op_name} kernel with KernelAgent (parallel workers + refinement)"
             )
 
-            # Generate kernel using KernelAgent
             result = agent.generate_kernel(
                 problem_description=problem_description,
-                test_code=None,  # Let KernelAgent auto-generate the test
+                test_code=None,
             )
 
+            # Only raise AgentError if kernel_code is missing or malformed
+            kernel_code = result.get("kernel_code")
+            if not kernel_code or not isinstance(kernel_code, str):
+                raise AgentError(f"Agent error: No kernel code produced for {op_name}.")
+
             if result["success"]:
                 print(f"✅ KernelAgent succeeded for {op_name}!")
                 print(
                     f"   Worker {result['worker_id']} found solution in {result['rounds']} rounds"
                 )
                 print(f"   Session: {result['session_dir']}")
-
-                # Copy the session directory to our kernels directory for preservation
                 import shutil
 
                 session_name = os.path.basename(result["session_dir"])
@@ -255,15 +255,20 @@ def generate_kernel_with_agent(self, op, op_name: str) -> tuple[str, bool]:
                     print(f"   Session preserved: {preserved_session}")
                 except Exception as e:
                     print(f"   Warning: Could not preserve session: {e}")
-
-                return result["kernel_code"], True
+                return kernel_code, True
             else:
-                print(f"❌ KernelAgent failed for {op_name}: {result['message']}")
-                return "", False
+                # This is an agent output error, so raise AgentError
+                raise AgentError(
+                    f"Agent error: ❌ KernelAgent failed for {op_name}: {result['message']}"
+                )
 
-        except Exception as e:
+        except AgentError as e:
             print(f"❌ KernelAgent error for {op_name}: {e}")
             return "", False
+        except Exception as e:
+            # API/provider errors are not actionable by the agent
+            print(f"❌ API/provider error for {op_name}: {e}")
+            return "", False
 
     def __getitem__(self, key):
         if key in self.compiled_kernels:

BackendBench/backends/llm.py

@@ -14,6 +14,7 @@
 
 import torch
 
+from BackendBench.agent_errors import AgentError
 from BackendBench.eval import eval_performance
 from BackendBench.llm_client import LLMKernelGenerator
 from BackendBench.multiprocessing_eval import MultiprocessingEvaluator
@@ -167,9 +168,15 @@ def test_kernel_correctness(
             "compilation_error": None,
             "test_errors": [],
             "summary": None,
+            "agent_error": "",
         }
 
         try:
+            # Only raise AgentError if kernel_code is missing or malformed
+            if not kernel_code or not isinstance(kernel_code, str):
+                raise AgentError(
+                    "Kernel code is empty or not a string (agent failed to produce a kernel)."
+                )
             kernel_file = self._generate_kernel_file_path(op_name, attempt)
             if not os.path.exists(kernel_file):
                 save_kernel_to_file(kernel_code, kernel_file)
@@ -178,16 +185,12 @@ def test_kernel_correctness(
                 f"{op_name}_implementation_v{attempt}", kernel_file
             )
             module = importlib.util.module_from_spec(spec)
-
-            # Add to sys.modules so triton can find it
             sys.modules[f"{op_name}_implementation_v{attempt}"] = module
 
             try:
                 spec.loader.exec_module(module)
-
                 expected_name = f"{op_name}_kernel_impl"
                 if hasattr(module, expected_name):
-                    # check if the kernel compile / is loadable
                     _ = getattr(module, expected_name)
                 else:
                     available_functions = [
@@ -198,12 +201,9 @@ def test_kernel_correctness(
                     raise ValueError(
                         f"Expected function '{expected_name}' not found. Available: {available_functions}"
                     )
-
             finally:
                 if f"test_kernel_{op_name}_{attempt}" in sys.modules:
                     del sys.modules[f"test_kernel_{op_name}_{attempt}"]
-
-                # Clear CUDA cache and synchronize to prevent memory buildup
                 if torch.cuda.is_available():
                     torch.cuda.empty_cache()
                     torch.cuda.synchronize()
@@ -220,10 +220,7 @@ def test_kernel_correctness(
                     test_cases,
                     [],
                 )
-
-                # Start evaluation
                 evaluator.start_evaluation()
-                # Get results
                 results = evaluator.get_results()
 
             for result in results:
@@ -248,6 +245,10 @@ def test_kernel_correctness(
 
             return is_correct, feedback_info
 
+        except AgentError as e:
+            feedback_info["agent_error"] = str(e)
+            feedback_info["summary"] = f"Agent error: {str(e)}"
+            return False, feedback_info
         except Exception as e:
             logger.error("    ✗ Compilation failed:")
             logger.error(f"      Error: {str(e)}")

BackendBench/llm_client.py

@@ -12,6 +12,8 @@
 from tenacity import retry
 from tenacity.wait import wait_random_exponential
 
+from BackendBench.agent_errors import AgentError
+
 from .kernel_templates import KernelTemplateManager
 
 
@@ -60,15 +62,26 @@ def readme_setup_section(self) -> str:
 
     @retry(wait=wait_random_exponential(multiplier=2, min=1, max=60, exp_base=2))
     def call_llm(self, prompt: str) -> str:
-        response = self.client.messages.create(
-            model=self.model,
-            max_tokens=8000,
-            temperature=0.2,
-            timeout=120.0,
-            messages=[{"role": "user", "content": prompt}],
-        )
-        content = response.content[0].text
-        return content
+        try:
+            response = self.client.messages.create(
+                model=self.model,
+                max_tokens=8000,
+                temperature=0.2,
+                timeout=120.0,
+                messages=[{"role": "user", "content": prompt}],
+            )
+            content = response.content[0].text
+            if not content:
+                raise ConnectionError(
+                    "API error: Empty response from LLM API (possible rate limit or outage)."
+                )
+            if "rate limit" in content.lower():
+                raise ConnectionError("API error: Rate limit encountered from LLM API.")
+            return content
+        except anthropic.AnthropicError as e:
+            raise ConnectionError(f"API error: Anthropic API error: {e}")
+        except Exception as e:
+            raise ConnectionError(f"API error: Unexpected error: {e}")
 
     def generate_kernel(
         self,
@@ -93,9 +106,7 @@ def generate_kernel(
 
         try:
             content = self.call_llm(prompt)
-            if not content:
-                raise RuntimeError("Empty response from LLM relay server")
-
+            # Only raise AgentError if kernel extraction fails
             extracted_code = self._extract_code_from_response(content)
 
             print("\n=== DEBUG: RAW LLM RELAY RESPONSE ===")
@@ -106,12 +117,10 @@ def generate_kernel(
 
             return extracted_code
 
-        except requests.exceptions.RequestException as e:
-            raise RuntimeError(
-                f"Failed to communicate with LLM relay server for {op_name}: {str(e)}"
-            )
+        except AgentError:
+            raise
         except Exception as e:
-            raise RuntimeError(f"Failed to generate kernel for {op_name}: {str(e)}")
+            raise AgentError(f"Agent error: Failed to generate kernel for {op_name}: {str(e)}")
 
     def generate_kernel_with_retry(
         self,
@@ -184,16 +193,11 @@ def _format_feedback(self, feedback_info: Dict) -> str:
 
     def _extract_code_from_response(self, response: str) -> str:
         if "```python" not in response:
-            raise ValueError(
-                "No Python code block found in LLM response. Response should contain ```python...``` block."
-            )
-
+            raise AgentError("Agent error: No Python code block found in LLM response.")
         start = response.find("```python") + len("```python")
         end = response.find("```", start)
-
         if end == -1:
-            raise ValueError("Unclosed Python code block in LLM response.")
-
+            raise AgentError("Agent error: Unclosed Python code block in LLM response.")
         return response[start:end].strip()
 
 
@@ -249,17 +253,24 @@ def call_llm(self, prompt: str) -> str:
             else None
         )
 
-        response = requests.post(
-            self.server_url,
-            json=request_data,
-            headers={"Content-Type": "application/json"},
-            timeout=120.0,
-            proxies=proxies,
-        )
-
-        if response.status_code != 200:
-            raise RuntimeError(f"Server returned status {response.status_code}: {response.text}")
-
-        response_data = response.json()
-        content = response_data.get("output", "")
-        return content
+        try:
+            response = requests.post(
+                self.server_url,
+                json=request_data,
+                headers={"Content-Type": "application/json"},
+                timeout=120.0,
+                proxies=proxies,
+            )
+            if response.status_code != 200:
+                raise AgentError(
+                    f"Agent error: Server returned status {response.status_code}: {response.text}"
+                )
+            response_data = response.json()
+            content = response_data.get("output", "")
+            if not content or "rate limit" in content.lower():
+                raise AgentError("Agent error: Empty response or rate limit encountered.")
+            return content
+        except requests.exceptions.RequestException as e:
+            raise AgentError(f"Agent error: Failed to communicate with LLM relay server: {str(e)}")
+        except Exception as e:
+            raise AgentError(f"Agent error: Unexpected error in LLM relay call: {e}")