tidy up

Author: KRRT7

codeflash/benchmarking/codeflash_trace.py

@@ -6,6 +6,7 @@
 import time
 from typing import Callable
 
+from codeflash.cli_cmds.cli import logger
 from codeflash.picklepatch.pickle_patcher import PicklePatcher
 
 
@@ -42,10 +43,8 @@ def setup(self, trace_path: str) -> None:
             )
             self._connection.commit()
         except Exception as e:
-            print(f"Database setup error: {e}")
-            if self._connection:
-                self._connection.close()
-                self._connection = None
+            logger.error(f"Database setup error: {e}")
+            self.close()
             raise
 
     def write_function_timings(self) -> None:
@@ -63,7 +62,6 @@ def write_function_timings(self) -> None:
 
         try:
             cur = self._connection.cursor()
-            # Insert data into the benchmark_function_timings table
             cur.executemany(
                 "INSERT INTO benchmark_function_timings"
                 "(function_name, class_name, module_name, file_path, benchmark_function_name, "
@@ -74,7 +72,7 @@ def write_function_timings(self) -> None:
             self._connection.commit()
             self.function_calls_data = []
         except Exception as e:
-            print(f"Error writing to function timings database: {e}")
+            logger.error(f"Error writing to function timings database: {e}")
             if self._connection:
                 self._connection.rollback()
             raise
@@ -103,7 +101,7 @@ def __call__(self, func: Callable) -> Callable:
         func_id = (func.__module__, func.__name__)
 
         @functools.wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: tuple, **kwargs: dict) -> object:
             # Initialize thread-local active functions set if it doesn't exist
             if not hasattr(self._thread_local, "active_functions"):
                 self._thread_local.active_functions = set()
@@ -124,19 +122,17 @@ def wrapper(*args, **kwargs):
             if os.environ.get("CODEFLASH_BENCHMARKING", "False") == "False":
                 self._thread_local.active_functions.remove(func_id)
                 return result
-            # Get benchmark info from environment
+
             benchmark_function_name = os.environ.get("CODEFLASH_BENCHMARK_FUNCTION_NAME", "")
             benchmark_module_path = os.environ.get("CODEFLASH_BENCHMARK_MODULE_PATH", "")
             benchmark_line_number = os.environ.get("CODEFLASH_BENCHMARK_LINE_NUMBER", "")
-            # Get class name
             class_name = ""
             qualname = func.__qualname__
             if "." in qualname:
                 class_name = qualname.split(".")[0]
 
-            # Limit pickle count so memory does not explode
             if self.function_call_count > self.pickle_count_limit:
-                print("Pickle limit reached")
+                logger.debug("CodeflashTrace: Pickle limit reached")
                 self._thread_local.active_functions.remove(func_id)
                 overhead_time = time.thread_time_ns() - end_time
                 self.function_calls_data.append(
@@ -161,7 +157,7 @@ def wrapper(*args, **kwargs):
                 pickled_args = PicklePatcher.dumps(args, protocol=pickle.HIGHEST_PROTOCOL)
                 pickled_kwargs = PicklePatcher.dumps(kwargs, protocol=pickle.HIGHEST_PROTOCOL)
             except Exception as e:
-                print(f"Error pickling arguments for function {func.__name__}: {e}")
+                logger.debug(f"CodeflashTrace: Error pickling arguments for function {func.__name__}: {e}")
                 # Add to the list of function calls without pickled args. Used for timing info only
                 self._thread_local.active_functions.remove(func_id)
                 overhead_time = time.thread_time_ns() - end_time
@@ -181,7 +177,6 @@ def wrapper(*args, **kwargs):
                     )
                 )
                 return result
-            # Flush to database every 100 calls
             if len(self.function_calls_data) > 100:
                 self.write_function_timings()
 
@@ -208,5 +203,4 @@ def wrapper(*args, **kwargs):
         return wrapper
 
 
-# Create a singleton instance
 codeflash_trace = CodeflashTrace()

codeflash/benchmarking/instrument_codeflash_trace.py

@@ -93,8 +93,5 @@ def instrument_codeflash_trace_decorator(file_to_funcs_to_optimize: dict[Path, l
     for file_path, functions_to_optimize in file_to_funcs_to_optimize.items():
         original_code = file_path.read_text(encoding="utf-8")
         new_code = add_codeflash_decorator_to_code(original_code, functions_to_optimize)
-        # Modify the code
         modified_code = isort.code(code=new_code, float_to_top=True)
-
-        # Write the modified code back to the file
         file_path.write_text(modified_code, encoding="utf-8")

codeflash/benchmarking/replay_test.py

@@ -28,13 +28,11 @@ def get_next_arg_and_return(
     limit = num_to_get
 
     if class_name is not None:
-        cursor = cur.execute(
-            "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = ? LIMIT ?",
-            (benchmark_function_name, function_name, file_path, class_name, limit),
-        )
+        query = "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = ? LIMIT ?"  # noqa: E501
+        cursor = cur.execute(query(benchmark_function_name, function_name, file_path, class_name, limit))
     else:
         cursor = cur.execute(
-            "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = '' LIMIT ?",
+            "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = '' LIMIT ?",  # noqa: E501
             (benchmark_function_name, function_name, file_path, limit),
         )
 
@@ -98,7 +96,7 @@ def create_trace_replay_test_code(
             args = pickle.loads(args_pkl)
             kwargs = pickle.loads(kwargs_pkl)
             ret = {function_name}(*args, **kwargs)
-            """
+            """  # noqa: E501
     )
 
     test_method_body = textwrap.dedent(
@@ -114,7 +112,7 @@ def create_trace_replay_test_code(
             else:
                 instance = args[0] # self
                 ret = instance{method_name}(*args[1:], **kwargs)
-            """
+            """  # noqa: E501
     )
 
     test_class_method_body = textwrap.dedent(
@@ -125,15 +123,15 @@ def create_trace_replay_test_code(
             if not args:
                 raise ValueError("No arguments provided for the method.")
             ret = {class_name_alias}{method_name}(*args[1:], **kwargs)
-            """
+            """  # noqa: E501
     )
     test_static_method_body = textwrap.dedent(
         """\
         for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}):
             args = pickle.loads(args_pkl)
             kwargs = pickle.loads(kwargs_pkl){filter_variables}
             ret = {class_name_alias}{method_name}(*args, **kwargs)
-            """
+            """  # noqa: E501
     )
 
     # Create main body
@@ -166,7 +164,7 @@ def create_trace_replay_test_code(
             alias = get_function_alias(module_name, class_name + "_" + function_name)
 
             filter_variables = ""
-            # filter_variables = '\n    args.pop("cls", None)'
+            # filter_variables = '\n    args.pop("cls", None)'  # noqa: ERA001
             method_name = "." + function_name if function_name != "__init__" else ""
             if function_properties.is_classmethod:
                 test_body = test_class_method_body.format(
@@ -227,11 +225,9 @@ def generate_replay_test(
     """
     count = 0
     try:
-        # Connect to the database
         conn = sqlite3.connect(trace_file_path.as_posix())
         cursor = conn.cursor()
 
-        # Get distinct benchmark file paths
         cursor.execute("SELECT DISTINCT benchmark_module_path FROM benchmark_function_timings")
         benchmark_files = cursor.fetchall()
 
@@ -240,15 +236,14 @@ def generate_replay_test(
             benchmark_module_path = benchmark_file[0]
             # Get all benchmarks and functions associated with this file path
             cursor.execute(
-                "SELECT DISTINCT benchmark_function_name, function_name, class_name, module_name, file_path, benchmark_line_number FROM benchmark_function_timings "
+                "SELECT DISTINCT benchmark_function_name, function_name, class_name, module_name, file_path, benchmark_line_number FROM benchmark_function_timings "  # noqa: E501
                 "WHERE benchmark_module_path = ?",
                 (benchmark_module_path,),
             )
 
             functions_data = []
             for row in cursor.fetchall():
                 benchmark_function_name, function_name, class_name, module_name, file_path, benchmark_line_number = row
-                # Add this function to our list
                 functions_data.append(
                     {
                         "function_name": function_name,
@@ -278,13 +273,12 @@ def generate_replay_test(
             output_file = get_test_file_path(
                 test_dir=Path(output_dir), function_name=benchmark_module_path, test_type="replay"
             )
-            # Write test code to file, parents = true
             output_dir.mkdir(parents=True, exist_ok=True)
             output_file.write_text(test_code, "utf-8")
             count += 1
 
         conn.close()
-    except Exception as e:
-        logger.info(f"Error generating replay tests: {e}")
+    except Exception as e:  # noqa: BLE001
+        logger.error(f"Error generating replay test: {e}")
 
     return count