boilerplate ready, need to fill in the gaps

Author: aseembits93

codeflash/optimization/function_optimizer.py

@@ -828,6 +828,7 @@ def establish_original_code_baseline(
                 test_env["PYTHONPATH"] += os.pathsep + str(self.args.project_root)
 
             coverage_results = None
+            lprofiler_results = None
             # Instrument codeflash capture
             try:
                 instrument_codeflash_capture(
@@ -840,6 +841,7 @@ def establish_original_code_baseline(
                     optimization_iteration=0,
                     testing_time=TOTAL_LOOPING_TIME,
                     enable_coverage=test_framework == "pytest",
+                    enable_profiler=test_framework == "pytest",
                     code_context=code_context,
                 )
             finally:
@@ -1058,6 +1060,7 @@ def run_and_parse_tests(
         testing_time: float = TOTAL_LOOPING_TIME,
         *,
         enable_coverage: bool = False,
+        enable_profiler: bool = False,
         pytest_min_loops: int = 5,
         pytest_max_loops: int = 100_000,
         code_context: CodeOptimizationContext | None = None,
@@ -1074,6 +1077,7 @@ def run_and_parse_tests(
                     pytest_timeout=INDIVIDUAL_TESTCASE_TIMEOUT,
                     verbose=True,
                     enable_coverage=enable_coverage,
+                    enable_profiler=enable_profiler,
                 )
             elif testing_type == TestingMode.PERFORMANCE:
                 result_file_path, run_result = run_benchmarking_tests(

codeflash/optimization/optimizer.py

@@ -89,7 +89,7 @@ def run(self) -> None:
                 tempfile.mkdtemp(dir=self.args.tests_root, prefix="codeflash_concolic_")
             )
         try:
-            ph("cli-optimize-functions-to-optimize", {"num_functions": num_optimizable_functions})
+            ph('cli-optimize-functions-to-optimize", {"num_functions": num_optimizable_functions})
             if num_optimizable_functions == 0:
                 logger.info("No functions found to optimize. Exiting…")
                 return
@@ -102,7 +102,7 @@ def run(self) -> None:
             console.rule()
             logger.info(f"Discovered {num_discovered_tests} existing unit tests in {self.test_cfg.tests_root}")
             console.rule()
-            ph("cli-optimize-discovered-tests", {"num_tests": num_discovered_tests})
+            ph('cli-optimize-discovered-tests", {"num_tests": num_discovered_tests})
 
             for original_module_path in file_to_funcs_to_optimize:
                 logger.info(f"Examining file {original_module_path!s}…")

codeflash/verification/test_runner.py

@@ -36,6 +36,7 @@ def run_behavioral_tests(
     verbose: bool = False,
     pytest_target_runtime_seconds: int = TOTAL_LOOPING_TIME,
     enable_coverage: bool = False,
+    enable_profiler: bool = False,
 ) -> tuple[Path, subprocess.CompletedProcess, Path | None]:
     if test_framework == "pytest":
         test_files: list[str] = []
@@ -109,6 +110,28 @@ def run_behavioral_tests(
             logger.debug(
                 f"""Result return code: {results.returncode}, {"Result stderr:" + str(results.stderr) if results.stderr else ""}"""
             )
+        if enable_profiler:
+            coverage_database_file, coveragercfile = prepare_coverage_files()
+
+            cov_erase = execute_test_subprocess(
+                shlex.split(f"{SAFE_SYS_EXECUTABLE} -m coverage erase"), cwd=cwd, env=pytest_test_env
+            )  # this cleanup is necessary to avoid coverage data from previous runs, if there are any,
+            # then the current run will be appended to the previous data, which skews the results
+            logger.debug(cov_erase)
+            coverage_cmd = [SAFE_SYS_EXECUTABLE, "-m", "coverage", "run", f"--rcfile={coveragercfile.as_posix()}", "-m"]
+
+            if pytest_cmd == "pytest":
+                coverage_cmd.extend(["pytest"])
+            else:
+                coverage_cmd.extend(shlex.split(pytest_cmd, posix=IS_POSIX)[1:])
+
+            results = execute_test_subprocess(
+                coverage_cmd + common_pytest_args + result_args + test_files, cwd=cwd, env=pytest_test_env, timeout=600
+            )
+            logger.debug(
+                f"Result return code: {results.returncode}, "
+                f"{'Result stderr:' + str(results.stderr) if results.stderr else ''}"
+            )
     elif test_framework == "unittest":
         if enable_coverage:
             msg = "Coverage is not supported yet for unittest framework"
@@ -127,7 +150,6 @@ def run_behavioral_tests(
 
     return result_file_path, results, coverage_database_file if enable_coverage else None
 
-
 def run_benchmarking_tests(
     test_paths: TestFiles,
     pytest_cmd: str,