Merge pull request #226 from codeflash-ai/fix-lineprofiler

Fix loop ending when all candidates have been tested but line profiler hasn't returned new candidates

Author: aseembits93

codeflash/optimization/function_optimizer.py

@@ -393,7 +393,7 @@ def determine_best_candidate(
             try:
                 candidate_index = 0
                 original_len = len(candidates)
-                while candidates:
+                while True:
                     done = True if future_line_profile_results is None else future_line_profile_results.done()
                     if done and (future_line_profile_results is not None):
                         line_profile_results = future_line_profile_results.result()
@@ -403,8 +403,13 @@ def determine_best_candidate(
                             f"Added results from line profiler to candidates, total candidates now: {original_len}"
                         )
                         future_line_profile_results = None
+                    try:
+                        candidate = candidates.popleft()
+                    except IndexError as e:
+                        if done:
+                            break
+                        continue
                     candidate_index += 1
-                    candidate = candidates.popleft()
                     get_run_tmp_file(Path(f"test_return_values_{candidate_index}.bin")).unlink(missing_ok=True)
                     get_run_tmp_file(Path(f"test_return_values_{candidate_index}.sqlite")).unlink(missing_ok=True)
                     logger.info(f"Optimization candidate {candidate_index}/{original_len}:")
@@ -512,7 +517,8 @@ def determine_best_candidate(
                     self.write_code_and_helpers(
                         self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
                     )
-
+                    if done and not candidates:
+                        break
             except KeyboardInterrupt as e:
                 self.write_code_and_helpers(
                     self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path