fix: simplify parallel evaluator, add line profiler and immediate refinement dispatch

Move refinement/repair dispatch from lazy carry-over pattern to immediate
ThreadPoolExecutor submission in function_optimizer.py. Add
_run_line_profiler_for_winner() that runs real line profiling on the
winning candidate after selection. Remove AsyncAiServiceClient dependency
and dead code (PendingRefinement, _dispatch_refinements, etc.) from
parallel_evaluator.py — it now only handles worktree-based evaluation.

Author: KRRT7

codeflash/languages/function_optimizer.py

@@ -133,7 +133,6 @@
         TestDiff,
         TestFileReview,
     )
-    from codeflash.optimization.parallel_evaluator import PendingRepair
     from codeflash.verification.verification_utils import TestConfig
 
 
@@ -1021,6 +1020,34 @@ def handle_successful_candidate(
 
         return best_optimization, benchmark_tree
 
+    def _run_line_profiler_for_winner(
+        self,
+        best_optimization: BestOptimization,
+        code_context: CodeOptimizationContext,
+        original_helper_code: dict[Path, str],
+        eval_ctx: CandidateEvaluationContext,
+    ) -> BestOptimization:
+        """Run line profiler on the winning candidate from parallel evaluation."""
+        try:
+            self.replace_function_and_helpers_with_optimized_code(
+                code_context=code_context,
+                optimized_code=best_optimization.candidate.source_code,
+                original_helper_code=original_helper_code,
+            )
+            with progress_bar("Running line-by-line profiling"):
+                lp_results = self.line_profiler_step(
+                    code_context=code_context, original_helper_code=original_helper_code, candidate_index=0
+                )
+            eval_ctx.record_line_profiler_result(best_optimization.candidate.optimization_id, lp_results["str_out"])
+            best_optimization.line_profiler_test_results = lp_results
+        except (ValueError, SyntaxError, AttributeError) as e:
+            logger.warning(f"Line profiler failed for winning candidate: {e}")
+        finally:
+            self.write_code_and_helpers(
+                self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
+            )
+        return best_optimization
+
     def select_best_optimization(
         self,
         eval_ctx: CandidateEvaluationContext,
@@ -1439,6 +1466,11 @@ def determine_best_candidate(
         )
 
         if best_optimization:
+            if parallel_pool_size > 1:
+                best_optimization = self._run_line_profiler_for_winner(
+                    best_optimization, code_context, original_helper_code, eval_ctx
+                )
+
             self.log_evaluation_results(
                 eval_ctx=eval_ctx,
                 best_optimization=best_optimization,
@@ -1463,18 +1495,14 @@ def _evaluate_candidates_parallel(
         pool_size: int,
     ) -> None:
         """Evaluate candidates in parallel using git worktrees and async subprocess execution."""
-        from codeflash.optimization.parallel_evaluator import (
-            PendingBatchRefinement,
-            PendingRefinement,
-            run_parallel_evaluation,
-        )
+        from codeflash.optimization.parallel_evaluator import run_parallel_evaluation
+
+        ai_service_client = self.aiservice_client if exp_type == "EXP0" else self.local_aiservice_client
+        assert ai_service_client is not None
 
         candidate_index = 0
-        carry_refinements: list[PendingRefinement | PendingBatchRefinement] = []
-        carry_repairs: list[PendingRepair] = []
 
         while not processor.is_done():
-            # Drain available candidates from the queue into a batch
             batch: list[tuple[CandidateNode, int, str | None]] = []
             while len(batch) < pool_size:
                 candidate_node = processor.get_next_candidate()
@@ -1483,7 +1511,6 @@ def _evaluate_candidates_parallel(
                 candidate_index += 1
                 cached = processor.normalized_cache.get(candidate_node.candidate.optimization_id)
 
-                # Pre-filter: skip duplicates and identical-to-original
                 normalized_code = cached or self.language_support.normalize_code(
                     candidate_node.candidate.source_code.flat.strip()
                 )
@@ -1502,13 +1529,12 @@ def _evaluate_candidates_parallel(
                 )
                 batch.append((candidate_node, candidate_index, cached))
 
-            if not batch and not carry_refinements and not carry_repairs:
+            if not batch:
                 break
 
-            if batch:
-                logger.info(f"Evaluating batch of {len(batch)} candidates in parallel…")
+            logger.info(f"Evaluating batch of {len(batch)} candidates in parallel…")
 
-            results, refinement_futures, repair_futures = run_parallel_evaluation(
+            results, _, _ = run_parallel_evaluation(
                 optimizer=self,
                 candidates=batch,
                 code_context=code_context,
@@ -1518,29 +1544,21 @@ def _evaluate_candidates_parallel(
                 eval_ctx=eval_ctx,
                 exp_type=exp_type,
                 pool_size=pool_size,
-                pending_refinements=carry_refinements if carry_refinements else None,
-                pending_repairs=carry_repairs if carry_repairs else None,
             )
 
-            # Append resolved futures from the async dispatch
-            self.future_all_refinements.extend(refinement_futures)
-            self.future_all_code_repair.extend(repair_futures)
-            carry_refinements = []
-            carry_repairs = []
-
-            # Process results: build refinement/repair requests for the next async pass
+            # Process results and dispatch refinement/repair futures immediately
             batch_refiner_candidates: list[AIServiceBatchRefinerCandidate] = []
             for (candidate_node, _idx, _), (_, run_result) in zip(batch, results):
                 candidate = candidate_node.candidate
 
                 if run_result is None or not is_successful(run_result):
                     eval_ctx.record_failed_candidate(candidate.optimization_id)
-                    if run_result is not None and hasattr(run_result, "error"):
-                        repair_pending = self._build_repair_request_if_possible(
-                            candidate, [], eval_ctx, code_context, 0, exp_type
+                    if run_result is not None and isinstance(run_result, Failure):
+                        repair_future = self._dispatch_repair_if_possible(
+                            candidate, eval_ctx, code_context, exp_type, ai_service_client
                         )
-                        if repair_pending is not None:
-                            carry_repairs.append(repair_pending)
+                        if repair_future is not None:
+                            self.future_all_code_repair.append(repair_future)
                     continue
 
                 candidate_result = run_result.unwrap()
@@ -1563,7 +1581,6 @@ def _evaluate_candidates_parallel(
                 ) and quantity_of_tests_critic(candidate_result)
 
                 if is_successful_opt:
-                    # Defer line profiling — use empty placeholder for now
                     empty_lp = {"timings": {}, "unit": 0, "str_out": ""}
                     best_optimization = BestOptimization(
                         candidate=candidate,
@@ -1591,67 +1608,96 @@ def _evaluate_candidates_parallel(
                         )
                     )
 
-            # Build pending refinement for the next batch's async boundary
+            # Dispatch refinement immediately so CandidateProcessor sees it
             if batch_refiner_candidates:
-                if len(batch_refiner_candidates) > 1:
-                    carry_refinements.append(
-                        PendingBatchRefinement(
-                            original_source_code=code_context.read_writable_code.markdown,
-                            read_only_dependency_code=code_context.read_only_context_code,
-                            original_line_profiler_results=original_code_baseline.line_profile_results["str_out"],
-                            trace_id=self.get_trace_id(exp_type),
-                            language=self.function_to_optimize.language,
-                            language_version=self.language_support.language_version,
-                            function_references=function_references,
-                            candidates=batch_refiner_candidates,
-                            rerun_trace_id=self.rerun_trace_id,
-                        )
-                    )
-                else:
-                    c = batch_refiner_candidates[0]
-                    carry_refinements.append(
-                        PendingRefinement(
-                            request=[
-                                AIServiceRefinerRequest(
-                                    optimization_id=c.optimization_id,
-                                    original_source_code=code_context.read_writable_code.markdown,
-                                    read_only_dependency_code=code_context.read_only_context_code,
-                                    original_code_runtime=c.original_code_runtime,
-                                    optimized_source_code=c.optimized_source_code,
-                                    optimized_explanation=c.optimized_explanation,
-                                    optimized_code_runtime=c.optimized_code_runtime,
-                                    speedup=c.speedup,
-                                    trace_id=self.get_trace_id(exp_type),
-                                    original_line_profiler_results=original_code_baseline.line_profile_results[
-                                        "str_out"
-                                    ],
-                                    optimized_line_profiler_results=c.optimized_line_profiler_results,
-                                    function_references=function_references,
-                                    language=self.function_to_optimize.language,
-                                    language_version=self.language_support.language_version,
-                                )
-                            ],
-                            rerun_trace_id=self.rerun_trace_id,
-                        )
-                    )
+                self._dispatch_refinement(
+                    batch_refiner_candidates,
+                    code_context,
+                    original_code_baseline,
+                    exp_type,
+                    function_references,
+                    ai_service_client,
+                )
 
-        # Flush any remaining pending requests in a final async pass
-        if carry_refinements or carry_repairs:
-            _, final_ref_futures, final_rep_futures = run_parallel_evaluation(
-                optimizer=self,
-                candidates=[],
-                code_context=code_context,
-                original_code_baseline=original_code_baseline,
-                original_helper_code=original_helper_code,
-                file_path_to_helper_classes=file_path_to_helper_classes,
-                eval_ctx=eval_ctx,
-                exp_type=exp_type,
-                pool_size=pool_size,
-                pending_refinements=carry_refinements,
-                pending_repairs=carry_repairs,
+    def _dispatch_refinement(
+        self,
+        batch_refiner_candidates: list[AIServiceBatchRefinerCandidate],
+        code_context: CodeOptimizationContext,
+        original_code_baseline: OriginalCodeBaseline,
+        exp_type: str,
+        function_references: str,
+        ai_service_client: AiServiceClient,
+    ) -> None:
+        """Submit refinement request to thread pool so CandidateProcessor can consume results."""
+        if len(batch_refiner_candidates) > 1:
+            future = self.executor.submit(
+                ai_service_client.optimize_code_refinement_batch,
+                original_source_code=code_context.read_writable_code.markdown,
+                read_only_dependency_code=code_context.read_only_context_code,
+                original_line_profiler_results=original_code_baseline.line_profile_results["str_out"],
+                trace_id=self.get_trace_id(exp_type),
+                language=self.function_to_optimize.language,
+                language_version=self.language_support.language_version,
+                function_references=function_references,
+                candidates=batch_refiner_candidates,
+                rerun_trace_id=self.rerun_trace_id,
             )
-            self.future_all_refinements.extend(final_ref_futures)
-            self.future_all_code_repair.extend(final_rep_futures)
+        else:
+            c = batch_refiner_candidates[0]
+            future = self.executor.submit(
+                ai_service_client.optimize_code_refinement,
+                request=[
+                    AIServiceRefinerRequest(
+                        optimization_id=c.optimization_id,
+                        original_source_code=code_context.read_writable_code.markdown,
+                        read_only_dependency_code=code_context.read_only_context_code,
+                        original_code_runtime=c.original_code_runtime,
+                        optimized_source_code=c.optimized_source_code,
+                        optimized_explanation=c.optimized_explanation,
+                        optimized_code_runtime=c.optimized_code_runtime,
+                        speedup=c.speedup,
+                        trace_id=self.get_trace_id(exp_type),
+                        original_line_profiler_results=original_code_baseline.line_profile_results["str_out"],
+                        optimized_line_profiler_results=c.optimized_line_profiler_results,
+                        function_references=function_references,
+                        language=self.function_to_optimize.language,
+                        language_version=self.language_support.language_version,
+                    )
+                ],
+                rerun_trace_id=self.rerun_trace_id,
+            )
+        self.future_all_refinements.append(future)
+
+    def _dispatch_repair_if_possible(
+        self,
+        candidate: OptimizedCandidate,
+        eval_ctx: CandidateEvaluationContext,
+        code_context: CodeOptimizationContext,
+        exp_type: str,
+        ai_service_client: AiServiceClient,
+    ) -> concurrent.futures.Future | None:
+        """Submit a code repair request if the candidate is eligible."""
+        max_repairs = get_effort_value(EffortKeys.MAX_CODE_REPAIRS_PER_TRACE, self.effort)
+        if self.repair_counter >= max_repairs:
+            return None
+
+        successful_candidates_count = sum(1 for is_correct in eval_ctx.is_correct.values() if is_correct)
+        if successful_candidates_count >= MIN_CORRECT_CANDIDATES:
+            return None
+
+        if candidate.source not in (OptimizedCandidateSource.OPTIMIZE, OptimizedCandidateSource.OPTIMIZE_LP):
+            return None
+
+        self.repair_counter += 1
+        request = AIServiceCodeRepairRequest(
+            optimization_id=candidate.optimization_id,
+            original_source_code=code_context.read_writable_code.markdown,
+            modified_source_code=candidate.source_code.markdown,
+            test_diffs=[],
+            trace_id=self.function_trace_id[:-4] + exp_type if self.experiment_id else self.function_trace_id,
+            language=self.function_to_optimize.language,
+        )
+        return self.executor.submit(ai_service_client.code_repair, request=request, rerun_trace_id=self.rerun_trace_id)
 
     def call_adaptive_optimize(
         self,
@@ -3181,45 +3227,6 @@ def repair_if_possible(
             )
         )
 
-    def _build_repair_request_if_possible(
-        self,
-        candidate: OptimizedCandidate,
-        diffs: list[TestDiff],
-        eval_ctx: CandidateEvaluationContext,
-        code_context: CodeOptimizationContext,
-        test_results_count: int,
-        exp_type: str,
-    ) -> PendingRepair | None:
-        """Like repair_if_possible but returns a PendingRepair for async dispatch instead of submitting directly."""
-        from codeflash.optimization.parallel_evaluator import PendingRepair
-
-        max_repairs = get_effort_value(EffortKeys.MAX_CODE_REPAIRS_PER_TRACE, self.effort)
-        if self.repair_counter >= max_repairs:
-            return None
-
-        successful_candidates_count = sum(1 for is_correct in eval_ctx.is_correct.values() if is_correct)
-        if successful_candidates_count >= MIN_CORRECT_CANDIDATES:
-            return None
-
-        if candidate.source not in (OptimizedCandidateSource.OPTIMIZE, OptimizedCandidateSource.OPTIMIZE_LP):
-            return None
-        if not diffs:
-            return None
-        result_unmatched_perc = len(diffs) / test_results_count if test_results_count > 0 else 1.0
-        if result_unmatched_perc > get_effort_value(EffortKeys.REPAIR_UNMATCHED_PERCENTAGE_LIMIT, self.effort):
-            return None
-
-        self.repair_counter += 1
-        request = AIServiceCodeRepairRequest(
-            optimization_id=candidate.optimization_id,
-            original_source_code=code_context.read_writable_code.markdown,
-            modified_source_code=candidate.source_code.markdown,
-            test_diffs=diffs,
-            trace_id=self.function_trace_id[:-4] + exp_type if self.experiment_id else self.function_trace_id,
-            language=self.function_to_optimize.language,
-        )
-        return PendingRepair(request=request, rerun_trace_id=self.rerun_trace_id)
-
     def run_optimized_candidate(
         self,
         *,