per module ranking

Author: KRRT7

codeflash/benchmarking/function_ranker.py

@@ -98,7 +98,7 @@ def get_function_ttx_score(self, function_to_optimize: FunctionToOptimize) -> fl
 
     def rank_functions(self, functions_to_optimize: list[FunctionToOptimize]) -> list[FunctionToOptimize]:
         ranked = sorted(functions_to_optimize, key=self.get_function_ttx_score, reverse=True)
-        logger.info(
+        logger.debug(
             f"Function ranking order: {[f'{func.function_name} (ttX={self.get_function_ttx_score(func):.2f})' for func in ranked]}"
         )
         return ranked

codeflash/discovery/functions_to_optimize.py

@@ -158,7 +158,7 @@ def get_functions_to_optimize(
     project_root: Path,
     module_root: Path,
     previous_checkpoint_functions: dict[str, dict[str, str]] | None = None,
-) -> tuple[dict[Path, list[FunctionToOptimize]], int]:
+) -> tuple[dict[Path, list[FunctionToOptimize]], int, Path | None]:
     assert sum([bool(optimize_all), bool(replay_test), bool(file)]) <= 1, (
         "Only one of optimize_all, replay_test, or file should be provided"
     )
@@ -208,27 +208,6 @@ def get_functions_to_optimize(
             functions, test_cfg.tests_root, ignore_paths, project_root, module_root, previous_checkpoint_functions
         )
 
-        if trace_file_path and trace_file_path.exists():
-            from codeflash.benchmarking.function_ranker import FunctionRanker
-
-            ranker = FunctionRanker(trace_file_path)
-
-            all_functions = []
-            for file_functions in filtered_modified_functions.values():
-                all_functions.extend(file_functions)
-
-            if all_functions:
-                ranked_functions = ranker.rank_functions(all_functions)
-                functions_count = len(ranked_functions)
-
-                ranked_dict = {}
-                for func in ranked_functions:
-                    if func.file_path not in ranked_dict:
-                        ranked_dict[func.file_path] = []
-                    ranked_dict[func.file_path].append(func)
-
-                filtered_modified_functions = ranked_dict
-
         logger.info(f"Found {functions_count} function{'s' if functions_count > 1 else ''} to optimize")
         if optimize_all:
             three_min_in_ns = int(1.8e11)
@@ -237,7 +216,7 @@ def get_functions_to_optimize(
                 f"It might take about {humanize_runtime(functions_count * three_min_in_ns)} to fully optimize this project. Codeflash "
                 f"will keep opening pull requests as it finds optimizations."
             )
-        return filtered_modified_functions, functions_count
+        return filtered_modified_functions, functions_count, trace_file_path
 
 
 def get_functions_within_git_diff() -> dict[str, list[FunctionToOptimize]]:

codeflash/optimization/optimizer.py

@@ -109,7 +109,7 @@ def run_benchmarks(
 
         return function_benchmark_timings, total_benchmark_timings
 
-    def get_optimizable_functions(self) -> tuple[dict[Path, list[FunctionToOptimize]], int]:
+    def get_optimizable_functions(self) -> tuple[dict[Path, list[FunctionToOptimize]], int, Path | None]:
         """Discover functions to optimize."""
         from codeflash.discovery.functions_to_optimize import get_functions_to_optimize
 
@@ -255,7 +255,7 @@ def run(self) -> None:
         cleanup_paths(Optimizer.find_leftover_instrumented_test_files(self.test_cfg.tests_root))
 
         function_optimizer = None
-        file_to_funcs_to_optimize, num_optimizable_functions = self.get_optimizable_functions()
+        file_to_funcs_to_optimize, num_optimizable_functions, trace_file_path = self.get_optimizable_functions()
         function_benchmark_timings, total_benchmark_timings = self.run_benchmarks(
             file_to_funcs_to_optimize, num_optimizable_functions
         )
@@ -282,7 +282,21 @@ def run(self) -> None:
 
                 validated_original_code, original_module_ast = module_prep_result
 
-                for function_to_optimize in file_to_funcs_to_optimize[original_module_path]:
+                functions_to_optimize = file_to_funcs_to_optimize[original_module_path]
+                if trace_file_path and trace_file_path.exists() and len(functions_to_optimize) > 1:
+                    try:
+                        from codeflash.benchmarking.function_ranker import FunctionRanker
+
+                        ranker = FunctionRanker(trace_file_path)
+                        functions_to_optimize = ranker.rank_functions(functions_to_optimize)
+                        logger.info(
+                            f"Ranked {len(functions_to_optimize)} functions by performance impact in {original_module_path}"
+                        )
+                        console.rule()
+                    except Exception as e:
+                        logger.debug(f"Could not rank functions in {original_module_path}: {e}")
+
+                for function_to_optimize in functions_to_optimize:
                     function_iterator_count += 1
                     logger.info(
                         f"Optimizing function {function_iterator_count} of {num_optimizable_functions}: "
@@ -322,6 +336,8 @@ def run(self) -> None:
             ph("cli-optimize-run-finished", {"optimizations_found": optimizations_found})
             if self.functions_checkpoint:
                 self.functions_checkpoint.cleanup()
+            if hasattr(self.args, "command") and self.args.command == "optimize":
+                self.cleanup_replay_tests()
             if optimizations_found == 0:
                 logger.info("❌ No optimizations found.")
             elif self.args.all:
@@ -352,6 +368,11 @@ def find_leftover_instrumented_test_files(test_root: Path) -> list[Path]:
             file_path for file_path in test_root.rglob("*") if file_path.is_file() and pattern.match(file_path.name)
         ]
 
+    def cleanup_replay_tests(self) -> None:
+        if self.replay_tests_dir and self.replay_tests_dir.exists():
+            logger.debug(f"Cleaning up replay tests directory: {self.replay_tests_dir}")
+            cleanup_paths([self.replay_tests_dir])
+
     def cleanup_temporary_paths(self) -> None:
         if self.current_function_optimizer:
             self.current_function_optimizer.cleanup_generated_files()