working demo of new opt candidates with lineprof info

Author: aseembits93

codeflash/api/aiservice.py

@@ -135,6 +135,72 @@ def optimize_python_code(
         console.rule()
         return []
 
+    def optimize_python_code_line_profiler(
+        self,
+        source_code: str,
+        dependency_code: str,
+        trace_id: str,
+        line_profiler_results: str,
+        num_candidates: int = 10,
+        experiment_metadata: ExperimentMetadata | None = None,
+    ) -> list[OptimizedCandidate]:
+        """Optimize the given python code for performance by making a request to the Django endpoint.
+
+        Parameters
+        ----------
+        - source_code (str): The python code to optimize.
+        - dependency_code (str): The dependency code used as read-only context for the optimization
+        - trace_id (str): Trace id of optimization run
+        - num_candidates (int): Number of optimization variants to generate. Default is 10.
+        - experiment_metadata (Optional[ExperimentalMetadata, None]): Any available experiment metadata for this optimization
+
+        Returns
+        -------
+        - List[OptimizationCandidate]: A list of Optimization Candidates.
+
+        """
+        payload = {
+            "source_code": source_code,
+            "dependency_code": dependency_code,
+            "num_variants": num_candidates,
+            "line_profiler_results": line_profiler_results,
+            "trace_id": trace_id,
+            "python_version": platform.python_version(),
+            "experiment_metadata": experiment_metadata,
+            "codeflash_version": codeflash_version,
+        }
+
+        logger.info("Generating optimized candidates…")
+        console.rule()
+        try:
+            response = self.make_ai_service_request("/optimize-line-profiler", payload=payload, timeout=600)
+        except requests.exceptions.RequestException as e:
+            logger.exception(f"Error generating optimized candidates: {e}")
+            ph("cli-optimize-error-caught", {"error": str(e)})
+            return []
+
+        if response.status_code == 200:
+            optimizations_json = response.json()["optimizations"]
+            logger.info(f"Generated {len(optimizations_json)} candidates.")
+            console.rule()
+            return [
+                OptimizedCandidate(
+                    source_code=opt["source_code"],
+                    explanation=opt["explanation"],
+                    optimization_id=opt["optimization_id"],
+                )
+                for opt in optimizations_json
+            ]
+        try:
+            error = response.json()["error"]
+        except Exception:
+            error = response.text
+        logger.error(f"Error generating optimized candidates: {response.status_code} - {error}")
+        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        console.rule()
+        return []
+
+
     def log_results(
         self,
         function_trace_id: str,

codeflash/models/models.py

@@ -218,6 +218,7 @@ class OriginalCodeBaseline(BaseModel):
     benchmarking_test_results: TestResults
     runtime: int
     coverage_results: Optional[CoverageData]
+    lprof_results: str
 
 
 class CoverageStatus(Enum):

codeflash/optimization/function_optimizer.py

@@ -21,7 +21,6 @@
 from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient
 from codeflash.cli_cmds.console import code_print, console, logger, progress_bar
 from codeflash.code_utils import env_utils
-from codeflash.code_utils.code_extractor import add_needed_imports_from_module, extract_code
 from codeflash.code_utils.code_replacer import replace_function_definitions_in_module, add_decorator_imports
 from codeflash.code_utils.code_utils import (
     cleanup_paths,
@@ -208,7 +207,7 @@ def optimize_function(self) -> Result[BestOptimization, str]:
                 and "." in function_source.qualified_name
             ):
                 file_path_to_helper_classes[function_source.file_path].add(function_source.qualified_name.split(".")[0])
-
+        pass
         baseline_result = self.establish_original_code_baseline(  # this needs better typing
             code_context=code_context,
             original_helper_code=original_helper_code,
@@ -232,7 +231,29 @@ def optimize_function(self) -> Result[BestOptimization, str]:
             return Failure("The threshold for test coverage was not met.")
 
         best_optimization = None
-
+        logger.info(f"Adding more candidates based on lineprof info, calling ai service")
+        with progress_bar(
+            f"Generating new optimizations for function {self.function_to_optimize.function_name} with line profiler information",
+            transient=True,
+        ):
+            pass
+            lprof_generated_results = self.aiservice_client.optimize_python_code_line_profiler(
+
+            source_code=code_context.read_writable_code,
+            dependency_code=code_context.read_only_context_code,
+            trace_id=self.function_trace_id,
+            line_profiler_results=original_code_baseline.lprof_results,
+            num_candidates = 10,
+            experiment_metadata = None)
+
+        if len(lprof_generated_results)==0:
+            logger.info(f"Generated tests with line profiler failed.")
+        else:
+            logger.info(f"Generated tests with line profiler succeeded. Appending to optimization candidates.")
+            print("initial optimization candidates",len(optimizations_set.control))
+            optimizations_set.control.extend(lprof_generated_results)
+            print("after adding optimization candidates",len(optimizations_set.control))
+            #append to optimization candidates
         for _u, candidates in enumerate([optimizations_set.control, optimizations_set.experiment]):
             if candidates is None:
                 continue
@@ -813,7 +834,7 @@ def establish_original_code_baseline(
                    files_to_instrument.append(helper_obj.file_path)
                    fns_to_instrument.append(helper_obj.qualified_name)
                add_decorator_imports(files_to_instrument,fns_to_instrument, lprofiler_database_file)
-               behavioral_results, coverage_results = self.run_and_parse_tests(
+               lprof_results, _ = self.run_and_parse_tests(
                     testing_type=TestingMode.BEHAVIOR,
                     test_env=test_env,
                     test_files=self.test_files,
@@ -822,7 +843,9 @@ def establish_original_code_baseline(
                     enable_coverage=False,
                     enable_lprofiler=test_framework == "pytest",
                     code_context=code_context,
-                )
+                    lprofiler_database_file=lprofiler_database_file,
+               )
+               pass
             except Exception as e:
                 logger.warning(f"Failed to run lprof for {self.function_to_optimize.function_name}. SKIPPING OPTIMIZING THIS FUNCTION.")
                 console.rule()
@@ -905,6 +928,7 @@ def establish_original_code_baseline(
                         benchmarking_test_results=benchmarking_results,
                         runtime=total_timing,
                         coverage_results=coverage_results,
+                        lprof_results=lprof_results,
                     ),
                     functions_to_remove,
                 )
@@ -1041,6 +1065,7 @@ def run_and_parse_tests(
         pytest_max_loops: int = 100_000,
         code_context: CodeOptimizationContext | None = None,
         unittest_loop_index: int | None = None,
+        lprofiler_database_file: str | None = None,
     ) -> tuple[TestResults, CoverageData | None]:
         coverage_database_file = None
         coverage_config_file = None
@@ -1083,20 +1108,24 @@ def run_and_parse_tests(
                 f"stdout: {run_result.stdout}\n"
                 f"stderr: {run_result.stderr}\n"
             )
-
-        results, coverage_results = parse_test_results(
-            test_xml_path=result_file_path,
-            test_files=test_files,
-            test_config=self.test_cfg,
-            optimization_iteration=optimization_iteration,
-            run_result=run_result,
-            unittest_loop_index=unittest_loop_index,
-            function_name=self.function_to_optimize.function_name,
-            source_file=self.function_to_optimize.file_path,
-            code_context=code_context,
-            coverage_database_file=coverage_database_file,
-            coverage_config_file=coverage_config_file,
-        )
+        if not enable_lprofiler:
+            results, coverage_results = parse_test_results(
+                test_xml_path=result_file_path,
+                test_files=test_files,
+                test_config=self.test_cfg,
+                optimization_iteration=optimization_iteration,
+                run_result=run_result,
+                unittest_loop_index=unittest_loop_index,
+                function_name=self.function_to_optimize.function_name,
+                source_file=self.function_to_optimize.file_path,
+                code_context=code_context,
+                coverage_database_file=coverage_database_file,
+                coverage_config_file=coverage_config_file,
+            )
+        else:
+            pass
+            file_contents = Path(str(lprofiler_database_file)+".txt").read_text("utf-8")
+            return file_contents, None
         return results, coverage_results
 
     def generate_and_instrument_tests(

codeflash/verification/test_runner.py

@@ -131,7 +131,8 @@ def run_behavioral_tests(
     else:
         msg = f"Unsupported test framework: {test_framework}"
         raise ValueError(msg)
-
+    if enable_lprofiler:
+        pass
     return result_file_path, results, coverage_database_file if enable_coverage else None, coverage_config_file if enable_coverage else None