Merge branch 'main' into tracer-optimization

Author: misrasaurabh1

codeflash/cli_cmds/cmd_init.py

@@ -22,7 +22,7 @@
 from codeflash.cli_cmds.console import console, logger
 from codeflash.code_utils.compat import LF
 from codeflash.code_utils.config_parser import parse_config_file
-from codeflash.code_utils.env_utils import get_codeflash_api_key
+from codeflash.code_utils.env_utils import check_formatter_installed, get_codeflash_api_key
 from codeflash.code_utils.git_utils import get_git_remotes, get_repo_owner_and_name
 from codeflash.code_utils.github_utils import get_github_secrets_page_url
 from codeflash.code_utils.shell_utils import get_shell_rc_path, save_api_key_to_rc
@@ -201,7 +201,7 @@ def collect_setup_info() -> SetupInfo:
             path_type=inquirer.Path.DIRECTORY,
         )
         if custom_module_root_answer:
-            module_root = Path(curdir) / Path(custom_module_root_answer["path"])
+            module_root = Path(custom_module_root_answer["path"])
         else:
             apologize_and_exit()
     else:
@@ -514,7 +514,8 @@ def install_github_actions(override_formatter_check: bool = False) -> None:  # n
         from importlib.resources import files
 
         benchmark_mode = False
-        if "benchmarks_root" in config:
+        benchmarks_root = config.get("benchmarks_root", "").strip()
+        if benchmarks_root and benchmarks_root != "":
             benchmark_mode = inquirer_wrapper(
                 inquirer.confirm,
                 message="⚡️It looks like you've configured a benchmarks_root in your config. Would you like to run the Github action in benchmark mode? "
@@ -537,7 +538,7 @@ def install_github_actions(override_formatter_check: bool = False) -> None:  # n
             existing_api_key = None
         click.prompt(
             f"Next, you'll need to add your CODEFLASH_API_KEY as a secret to your GitHub repo.{LF}"
-            f"Press Enter to open your repo's secrets page at {get_github_secrets_page_url(repo)}…{LF}"
+            f"Press Enter to open your repo's secrets page at {get_github_secrets_page_url(repo)} {LF}"
             f"Then, click 'New repository secret' to add your api key with the variable name CODEFLASH_API_KEY.{LF}"
             f"{'Here is your CODEFLASH_API_KEY: ' + existing_api_key + ' ' + LF}"
             if existing_api_key
@@ -720,11 +721,7 @@ def configure_pyproject_toml(setup_info: SetupInfo) -> None:
         )
     elif formatter == "don't use a formatter":
         formatter_cmds.append("disabled")
-    if formatter in ["black", "ruff"]:
-        try:
-            subprocess.run([formatter], capture_output=True, check=False)
-        except (FileNotFoundError, NotADirectoryError):
-            click.echo(f"⚠️ Formatter not found: {formatter}, please ensure it is installed")
+    check_formatter_installed(formatter_cmds)
     codeflash_section["formatter-cmds"] = formatter_cmds
     # Add the 'codeflash' section, ensuring 'tool' section exists
     tool_section = pyproject_data.get("tool", tomlkit.table())
@@ -750,7 +747,7 @@ def install_github_app() -> None:
 
     else:
         click.prompt(
-            f"Finally, you'll need install the Codeflash GitHub app by choosing the repository you want to install Codeflash on.{LF}"
+            f"Finally, you'll need to install the Codeflash GitHub app by choosing the repository you want to install Codeflash on.{LF}"
             f"I will attempt to open the github app page - https://github.com/apps/codeflash-ai/installations/select_target {LF}"
             f"Press Enter to open the page to let you install the app…{LF}",
             default="",
@@ -924,6 +921,14 @@ def test_sort():
 
 
 def run_end_to_end_test(args: Namespace, bubble_sort_path: str, bubble_sort_test_path: str) -> None:
+    try:
+        check_formatter_installed(args.formatter_cmds)
+    except Exception:
+        logger.error(
+            "Formatter not found. Review the formatter_cmds in your pyproject.toml file and make sure the formatter is installed."
+        )
+        return
+
     command = ["codeflash", "--file", "bubble_sort.py", "--function", "sorter"]
     if args.no_pr:
         command.append("--no-pr")

codeflash/code_utils/code_utils.py

@@ -2,13 +2,87 @@
 
 import ast
 import os
+import re
 import shutil
 import site
+from contextlib import contextmanager
 from functools import lru_cache
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
+import tomlkit
+
 from codeflash.cli_cmds.console import logger
+from codeflash.code_utils.config_parser import find_pyproject_toml
+
+
+@contextmanager
+def custom_addopts() -> None:
+    pyproject_file = find_pyproject_toml()
+    original_content = None
+    non_blacklist_plugin_args = ""
+
+    try:
+        # Read original file
+        if pyproject_file.exists():
+            with Path.open(pyproject_file, encoding="utf-8") as f:
+                original_content = f.read()
+                data = tomlkit.parse(original_content)
+            # Backup original addopts
+            original_addopts = data.get("tool", {}).get("pytest", {}).get("ini_options", {}).get("addopts", "")
+            # nothing to do if no addopts present
+            if original_addopts != "":
+                original_addopts = [x.strip() for x in original_addopts]
+                non_blacklist_plugin_args = re.sub(r"-n(?: +|=)\S+", "", " ".join(original_addopts)).split(" ")
+                non_blacklist_plugin_args = [x for x in non_blacklist_plugin_args if x != ""]
+                if non_blacklist_plugin_args != original_addopts:
+                    data["tool"]["pytest"]["ini_options"]["addopts"] = non_blacklist_plugin_args
+                    # Write modified file
+                    with Path.open(pyproject_file, "w", encoding="utf-8") as f:
+                        f.write(tomlkit.dumps(data))
+
+        yield
+
+    finally:
+        # Restore original file
+        if (
+            original_content
+            and pyproject_file.exists()
+            and tuple(original_addopts) not in {(), tuple(non_blacklist_plugin_args)}
+        ):
+            with Path.open(pyproject_file, "w", encoding="utf-8") as f:
+                f.write(original_content)
+
+
+@contextmanager
+def add_addopts_to_pyproject() -> None:
+    pyproject_file = find_pyproject_toml()
+    original_content = None
+    try:
+        # Read original file
+        if pyproject_file.exists():
+            with Path.open(pyproject_file, encoding="utf-8") as f:
+                original_content = f.read()
+                data = tomlkit.parse(original_content)
+            data["tool"]["pytest"] = {}
+            data["tool"]["pytest"]["ini_options"] = {}
+            data["tool"]["pytest"]["ini_options"]["addopts"] = [
+                "-n=auto",
+                "-n",
+                "1",
+                "-n 1",
+                "-n      1",
+                "-n      auto",
+            ]
+            with Path.open(pyproject_file, "w", encoding="utf-8") as f:
+                f.write(tomlkit.dumps(data))
+
+        yield
+
+    finally:
+        # Restore original file
+        with Path.open(pyproject_file, "w", encoding="utf-8") as f:
+            f.write(original_content)
 
 
 def encoded_tokens_len(s: str) -> int:

codeflash/code_utils/env_utils.py

@@ -1,13 +1,52 @@
 from __future__ import annotations
 
 import os
+import shlex
+import subprocess
+import tempfile
 from functools import lru_cache
+from pathlib import Path
 from typing import Optional
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.shell_utils import read_api_key_from_shell_config
 
 
+class FormatterNotFoundError(Exception):
+    """Exception raised when a formatter is not found."""
+
+    def __init__(self, formatter_cmd: str) -> None:
+        super().__init__(f"Formatter command not found: {formatter_cmd}")
+
+
+def check_formatter_installed(formatter_cmds: list[str]) -> bool:
+    return_code = True
+    if formatter_cmds[0] == "disabled":
+        return return_code
+    tmp_code = """print("hello world")"""
+    with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", suffix=".py") as f:
+        f.write(tmp_code)
+        f.flush()
+        tmp_file = Path(f.name)
+        file_token = "$file"  # noqa: S105
+        for command in set(formatter_cmds):
+            formatter_cmd_list = shlex.split(command, posix=os.name != "nt")
+            formatter_cmd_list = [tmp_file.as_posix() if chunk == file_token else chunk for chunk in formatter_cmd_list]
+            try:
+                result = subprocess.run(formatter_cmd_list, capture_output=True, check=False)
+            except (FileNotFoundError, NotADirectoryError):
+                return_code = False
+                break
+            if result.returncode:
+                return_code = False
+                break
+    tmp_file.unlink(missing_ok=True)
+    if not return_code:
+        msg = f"Error running formatter command: {command}"
+        raise FormatterNotFoundError(msg)
+    return return_code
+
+
 @lru_cache(maxsize=1)
 def get_codeflash_api_key() -> str:
     api_key = os.environ.get("CODEFLASH_API_KEY") or read_api_key_from_shell_config()

codeflash/discovery/discover_unit_tests.py

@@ -17,7 +17,7 @@
 from pydantic.dataclasses import dataclass
 
 from codeflash.cli_cmds.console import console, logger, test_files_progress_bar
-from codeflash.code_utils.code_utils import get_run_tmp_file, module_name_from_file_path
+from codeflash.code_utils.code_utils import custom_addopts, get_run_tmp_file, module_name_from_file_path
 from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE, codeflash_cache_db
 from codeflash.models.models import CodePosition, FunctionCalledInTest, TestsInFile, TestType
 
@@ -150,19 +150,20 @@ def discover_tests_pytest(
     project_root = cfg.project_root_path
 
     tmp_pickle_path = get_run_tmp_file("collected_tests.pkl")
-    result = subprocess.run(
-        [
-            SAFE_SYS_EXECUTABLE,
-            Path(__file__).parent / "pytest_new_process_discovery.py",
-            str(project_root),
-            str(tests_root),
-            str(tmp_pickle_path),
-        ],
-        cwd=project_root,
-        check=False,
-        capture_output=True,
-        text=True,
-    )
+    with custom_addopts():
+        result = subprocess.run(
+            [
+                SAFE_SYS_EXECUTABLE,
+                Path(__file__).parent / "pytest_new_process_discovery.py",
+                str(project_root),
+                str(tests_root),
+                str(tmp_pickle_path),
+            ],
+            cwd=project_root,
+            check=False,
+            capture_output=True,
+            text=True,
+        )
     try:
         with tmp_pickle_path.open(mode="rb") as f:
             exitcode, tests, pytest_rootdir = pickle.load(f)

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

codeflash/optimization/optimizer.py

@@ -82,6 +82,8 @@ def run(self) -> None:
         console.rule()
         if not env_utils.ensure_codeflash_api_key():
             return
+        if not env_utils.check_formatter_installed(self.args.formatter_cmds):
+            return
         function_optimizer = None
         file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]]
         num_optimizable_functions: int

codeflash/verification/test_runner.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING
 
 from codeflash.cli_cmds.console import logger
-from codeflash.code_utils.code_utils import get_run_tmp_file
+from codeflash.code_utils.code_utils import custom_addopts, get_run_tmp_file
 from codeflash.code_utils.compat import IS_POSIX, SAFE_SYS_EXECUTABLE
 from codeflash.code_utils.config_consts import TOTAL_LOOPING_TIME
 from codeflash.code_utils.coverage_utils import prepare_coverage_files
@@ -23,8 +23,9 @@ def execute_test_subprocess(
     cmd_list: list[str], cwd: Path, env: dict[str, str] | None, timeout: int = 600
 ) -> subprocess.CompletedProcess:
     """Execute a subprocess with the given command list, working directory, environment variables, and timeout."""
-    logger.debug(f"executing test run with command: {' '.join(cmd_list)}")
-    return subprocess.run(cmd_list, capture_output=True, cwd=cwd, env=env, text=True, timeout=timeout, check=False)
+    with custom_addopts():
+        logger.debug(f"executing test run with command: {' '.join(cmd_list)}")
+        return subprocess.run(cmd_list, capture_output=True, cwd=cwd, env=env, text=True, timeout=timeout, check=False)
 
 
 def run_behavioral_tests(
@@ -97,6 +98,7 @@ def run_behavioral_tests(
                 coverage_cmd.extend(shlex.split(pytest_cmd, posix=IS_POSIX)[1:])
 
             blocklist_args = [f"-p no:{plugin}" for plugin in BEHAVIORAL_BLOCKLISTED_PLUGINS if plugin != "cov"]
+
             results = execute_test_subprocess(
                 coverage_cmd + common_pytest_args + blocklist_args + result_args + test_files,
                 cwd=cwd,
@@ -252,7 +254,6 @@ def run_benchmarking_tests(
         pytest_test_env = test_env.copy()
         pytest_test_env["PYTEST_PLUGINS"] = "codeflash.verification.pytest_plugin"
         blocklist_args = [f"-p no:{plugin}" for plugin in BENCHMARKING_BLOCKLISTED_PLUGINS]
-
         results = execute_test_subprocess(
             pytest_cmd_list + pytest_args + blocklist_args + result_args + test_files,
             cwd=cwd,
@@ -278,7 +279,6 @@ def run_unittest_tests(
     log_level = ["-v"] if verbose else []
     files = [str(file) for file in test_file_paths]
     output_file = ["--output-file", str(result_file_path)]
-
     results = execute_test_subprocess(
         unittest_cmd_list + log_level + files + output_file, cwd=cwd, env=test_env, timeout=600
     )