cleanup code

Author: KRRT7

codeflash/benchmarking/function_ranker.py

@@ -2,7 +2,7 @@
 
 from typing import TYPE_CHECKING
 
-from codeflash.cli_cmds.console import console, logger
+from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.config_consts import DEFAULT_IMPORTANCE_THRESHOLD
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.tracing.profile_stats import ProfileStats
@@ -128,7 +128,8 @@ def rank_functions(self, functions_to_optimize: list[FunctionToOptimize]) -> lis
             total_program_time = sum(
                 s["own_time_ns"]
                 for s in self._function_stats.values()
-                if s.get("own_time_ns", 0) > 0 and any(
+                if s.get("own_time_ns", 0) > 0
+                and any(
                     str(s.get("filename", "")).endswith("/" + target_file) or s.get("filename") == target_file
                     for target_file in target_files
                 )

codeflash/benchmarking/replay_test.py

@@ -64,9 +64,7 @@ def get_unique_test_name(module: str, function_name: str, benchmark_name: str, c
 
 
 def create_trace_replay_test_code(
-    trace_file: str,
-    functions_data: list[dict[str, Any]],
-    max_run_count: int = 256,
+    trace_file: str, functions_data: list[dict[str, Any]], max_run_count: int = 256
 ) -> str:
     """Create a replay test for functions based on trace data.
 
@@ -220,9 +218,7 @@ def create_trace_replay_test_code(
     return imports + "\n" + metadata + "\n" + test_template
 
 
-def generate_replay_test(
-    trace_file_path: Path, output_dir: Path, max_run_count: int = 100
-) -> int:
+def generate_replay_test(trace_file_path: Path, output_dir: Path, max_run_count: int = 100) -> int:
     """Generate multiple replay tests from the traced function calls, grouped by benchmark.
 
     Args:
@@ -280,9 +276,7 @@ def generate_replay_test(
                 continue
             # Generate the test code for this benchmark
             test_code = create_trace_replay_test_code(
-                trace_file=trace_file_path.as_posix(),
-                functions_data=functions_data,
-                max_run_count=max_run_count,
+                trace_file=trace_file_path.as_posix(), functions_data=functions_data, max_run_count=max_run_count
             )
             test_code = sort_imports(code=test_code)
             output_file = get_test_file_path(

codeflash/tracing/replay_test.py

@@ -43,11 +43,7 @@ def get_function_alias(module: str, function_name: str) -> str:
     return "_".join(module.split(".")) + "_" + function_name
 
 
-def create_trace_replay_test(
-    trace_file: str,
-    functions: list[FunctionModules],
-    max_run_count: int = 100,
-) -> str:
+def create_trace_replay_test(trace_file: str, functions: list[FunctionModules], max_run_count: int = 100) -> str:
     imports = """import dill as pickle
 from codeflash.tracing.replay_test import get_next_arg_and_return
 """

codeflash/tracing/tracing_new_process.py

@@ -70,7 +70,6 @@ def __init__(
         self,
         config: dict,
         result_pickle_file_path: Path,
-        output: str = "codeflash.trace",
         functions: list[str] | None = None,
         disable: bool = False,  # noqa: FBT001, FBT002
         project_root: Path | None = None,
@@ -80,7 +79,6 @@ def __init__(
     ) -> None:
         """Use this class to trace function calls.
 
-        :param output: The path to the output trace file
         :param functions: List of functions to trace. If None, trace all functions
         :param disable: Disable the tracer if True
         :param max_function_count: Maximum number of times to trace one function
@@ -127,6 +125,7 @@ def __init__(
         self.sanitized_filename = self.sanitize_to_filename(command)
         # Place trace file next to replay tests in the tests directory
         from codeflash.verification.verification_utils import get_test_file_path
+
         function_path = "_".join(functions) if functions else self.sanitized_filename
         test_file_path = get_test_file_path(
             test_dir=Path(config["tests_root"]), function_name=function_path, test_type="replay"
@@ -279,9 +278,7 @@ def __exit__(
         from codeflash.verification.verification_utils import get_test_file_path
 
         replay_test = create_trace_replay_test(
-            trace_file=self.output_file,
-            functions=self.function_modules,
-            max_run_count=self.max_function_count,
+            trace_file=self.output_file, functions=self.function_modules, max_run_count=self.max_function_count
         )
         function_path = "_".join(self.functions) if self.functions else self.sanitized_filename
         test_file_path = get_test_file_path(