diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index d74f59ecc..fe4357839 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/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