diff --git a/codeflash/api/aiservice.py b/codeflash/api/aiservice.py index f337a985..00585f9c 100644 --- a/codeflash/api/aiservice.py +++ b/codeflash/api/aiservice.py @@ -353,6 +353,59 @@ def get_new_explanation( # noqa: D417 console.rule() return "" + def generate_ranking( # noqa: D417 + self, trace_id: str, diffs: list[str], optimization_ids: list[str], speedups: list[float] + ) -> list[int] | None: + """Optimize the given python code for performance by making a request to the Django endpoint. + + Parameters + ---------- + - source_code (str): The python code to optimize. + - optimized_code (str): The python code generated by the AI service. + - dependency_code (str): The dependency code used as read-only context for the optimization + - original_line_profiler_results: str - line profiler results for the baseline code + - optimized_line_profiler_results: str - line profiler results for the optimized code + - original_code_runtime: str - runtime for the baseline code + - optimized_code_runtime: str - runtime for the optimized code + - speedup: str - speedup of the optimized code + - annotated_tests: str - test functions annotated with runtime + - optimization_id: str - unique id of opt candidate + - original_explanation: str - original_explanation generated for the opt candidate + + Returns + ------- + - List[OptimizationCandidate]: A list of Optimization Candidates. + + """ + payload = { + "trace_id": trace_id, + "diffs": diffs, + "speedups": speedups, + "optimization_ids": optimization_ids, + "python_version": platform.python_version(), + } + logger.info("Generating ranking") + console.rule() + try: + response = self.make_ai_service_request("/rank", payload=payload, timeout=60) + except requests.exceptions.RequestException as e: + logger.exception(f"Error generating ranking: {e}") + ph("cli-optimize-error-caught", {"error": str(e)}) + return None + + if response.status_code == 200: + ranking: list[int] = response.json()["ranking"] + console.rule() + return ranking + try: + error = response.json()["error"] + except Exception: + error = response.text + logger.error(f"Error generating ranking: {response.status_code} - {error}") + ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error}) + console.rule() + return None + def log_results( # noqa: D417 self, function_trace_id: str, diff --git a/codeflash/code_utils/code_utils.py b/codeflash/code_utils/code_utils.py index dfd79a76..4ff01004 100644 --- a/codeflash/code_utils/code_utils.py +++ b/codeflash/code_utils/code_utils.py @@ -20,6 +20,23 @@ ImportErrorPattern = re.compile(r"ModuleNotFoundError.*$", re.MULTILINE) +def unified_diff_strings(code1: str, code2: str, fromfile: str = "original", tofile: str = "modified") -> str: + """Return the unified diff between two code strings as a single string. + + :param code1: First code string (original). + :param code2: Second code string (modified). + :param fromfile: Label for the first code string. + :param tofile: Label for the second code string. + :return: Unified diff as a string. + """ + code1_lines = code1.splitlines(keepends=True) + code2_lines = code2.splitlines(keepends=True) + + diff = difflib.unified_diff(code1_lines, code2_lines, fromfile=fromfile, tofile=tofile, lineterm="") + + return "".join(diff) + + def diff_length(a: str, b: str) -> int: """Compute the length (in characters) of the unified diff between two strings. diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index c523dcbc..4e624bbc 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -39,6 +39,7 @@ has_any_async_functions, module_name_from_file_path, restore_conftest, + unified_diff_strings, ) from codeflash.code_utils.config_consts import ( INDIVIDUAL_TESTCASE_TIMEOUT, @@ -170,9 +171,10 @@ def _process_refinement_results(self) -> OptimizedCandidate | None: self.candidate_queue.put(candidate) self.candidate_len += len(refinement_response) - logger.info( - f"Added {len(refinement_response)} candidates from refinement, total candidates now: {self.candidate_len}" - ) + if len(refinement_response) > 0: + logger.info( + f"Added {len(refinement_response)} candidates from refinement, total candidates now: {self.candidate_len}" + ) self.refinement_done = True return self.get_next_candidate() @@ -659,6 +661,9 @@ def determine_best_candidate( # reassign the shorter code here valid_candidates_with_shorter_code = [] diff_lens_list = [] # character level diff + speedups_list = [] + optimization_ids = [] + diff_strs = [] runtimes_list = [] for valid_opt in valid_optimizations: valid_opt_normalized_code = ast.unparse(ast.parse(valid_opt.candidate.source_code.flat.strip())) @@ -682,12 +687,35 @@ def determine_best_candidate( diff_lens_list.append( diff_length(new_best_opt.candidate.source_code.flat, code_context.read_writable_code.flat) ) # char level diff + diff_strs.append( + unified_diff_strings(code_context.read_writable_code.flat, new_best_opt.candidate.source_code.flat) + ) + speedups_list.append( + 1 + + performance_gain( + original_runtime_ns=original_code_baseline.runtime, optimized_runtime_ns=new_best_opt.runtime + ) + ) + optimization_ids.append(new_best_opt.candidate.optimization_id) runtimes_list.append(new_best_opt.runtime) - diff_lens_ranking = create_rank_dictionary_compact(diff_lens_list) - runtimes_ranking = create_rank_dictionary_compact(runtimes_list) - # TODO: better way to resolve conflicts with same min ranking - overall_ranking = {key: diff_lens_ranking[key] + runtimes_ranking[key] for key in diff_lens_ranking.keys()} # noqa: SIM118 - min_key = min(overall_ranking, key=overall_ranking.get) + future_ranking = self.executor.submit( + ai_service_client.generate_ranking, + diffs=diff_strs, + optimization_ids=optimization_ids, + speedups=speedups_list, + trace_id=self.function_trace_id[:-4] + exp_type if self.experiment_id else self.function_trace_id, + ) + concurrent.futures.wait([future_ranking]) + ranking = future_ranking.result() + if ranking: + ranking = [x - 1 for x in ranking] + min_key = ranking[0] + else: + diff_lens_ranking = create_rank_dictionary_compact(diff_lens_list) + runtimes_ranking = create_rank_dictionary_compact(runtimes_list) + # TODO: better way to resolve conflicts with same min ranking + overall_ranking = {key: diff_lens_ranking[key] + runtimes_ranking[key] for key in diff_lens_ranking.keys()} # noqa: SIM118 + min_key = min(overall_ranking, key=overall_ranking.get) best_optimization = valid_candidates_with_shorter_code[min_key] # reassign code string which is the shortest ai_service_client.log_results(