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

Safer `while` loop when processing optimization candidates

Author: aseembits93

codeflash/optimization/function_optimizer.py

@@ -392,7 +392,7 @@ def determine_best_candidate(
             try:
                 candidate_index = 0
                 original_len = len(candidates)
-                while True:
+                while candidates:
                     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()
@@ -402,13 +402,7 @@ 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:
-                        if done:
-                            break
-                        time.sleep(0.1)
-                        continue
+                    candidate = candidates.popleft()
                     candidate_index += 1
                     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)
@@ -517,8 +511,17 @@ 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
+                    if (not len(candidates)) and (
+                        not done
+                    ):  # all original candidates processed but lp results haven't been processed
+                        concurrent.futures.wait([future_line_profile_results])
+                        line_profile_results = future_line_profile_results.result()
+                        candidates.extend(line_profile_results)
+                        original_len += len(line_profile_results)
+                        logger.info(
+                            f"Added results from line profiler to candidates, total candidates now: {original_len}"
+                        )
+                        future_line_profile_results = None
             except KeyboardInterrupt as e:
                 self.write_code_and_helpers(
                     self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path