Merge branch 'main' into disable-lp-jit

Author: aseembits93

code_to_optimize/code_directories/simple_tracer_e2e/codeflash.trace

@@ -1,74 +1,39 @@
 from concurrent.futures import ThreadPoolExecutor
-from time import sleep
 
 
 def funcA(number):
-    number = number if number < 1000 else 1000
+    number = number if number < 100 else 100
     k = 0
-    for i in range(number * 100):
+    for i in range(number * 10):
         k += i
-    # Simplify the for loop by using sum with a range object
     j = sum(range(number))
-
-    # Use a generator expression directly in join for more efficiency
     return " ".join(str(i) for i in range(number))
 
 
 def test_threadpool() -> None:
-    pool = ThreadPoolExecutor(max_workers=3)
-    args = list(range(10, 31, 10))
+    pool = ThreadPoolExecutor(max_workers=2)
+    args = [5, 10, 15]
     result = pool.map(funcA, args)
 
     for r in result:
         print(r)
 
 class AlexNet:
-    def __init__(self, num_classes=1000):
+    def __init__(self, num_classes=10):
         self.num_classes = num_classes
-        self.features_size = 256 * 6 * 6
 
     def forward(self, x):
-        features = self._extract_features(x)
-        
-        output = self._classify(features)
-        return output
-
-    def _extract_features(self, x):
-        result = []
-        for i in range(len(x)):
-            pass
-
-        return result
-
-    def _classify(self, features):
-        total = sum(features)
-        return [total % self.num_classes for _ in features]
-
-class SimpleModel:
-    @staticmethod
-    def predict(data):
-        result = []
-        sleep(0.1) # can be optimized away
-        for i in range(500):
-            for x in data:
-                computation = 0
-                computation += x * i ** 2
-                result.append(computation)
-        return result
-    
-    @classmethod
-    def create_default(cls):
-        return cls()
+        result = 0
+        for val in x:
+            result += val * val
+        return result % self.num_classes
 
 
 def test_models():
     model = AlexNet(num_classes=10)
     input_data = [1, 2, 3, 4, 5]
     result = model.forward(input_data)
 
-    model2 = SimpleModel.create_default()
-    prediction = model2.predict(input_data)
-
 if __name__ == "__main__":
     test_threadpool()
     test_models()

code_to_optimize/code_directories/simple_tracer_e2e/workload.py

@@ -21,6 +21,7 @@
 from codeflash.models.models import (
     AIServiceRefinerRequest,
     CodeStringsMarkdown,
+    OptimizationReviewResult,
     OptimizedCandidate,
     OptimizedCandidateSource,
 )
@@ -652,7 +653,7 @@ def get_optimization_review(
         replay_tests: str,
         concolic_tests: str,  # noqa: ARG002
         calling_fn_details: str,
-    ) -> str:
+    ) -> OptimizationReviewResult:
         """Compute the optimization review of current Pull Request.
 
         Args:
@@ -670,7 +671,7 @@ def get_optimization_review(
 
         Returns:
         -------
-        - 'high', 'medium' or 'low' optimization review
+        OptimizationReviewResult with review ('high', 'medium', 'low', or '') and explanation
 
         """
         diff_str = "\n".join(
@@ -706,18 +707,21 @@ def get_optimization_review(
         except requests.exceptions.RequestException as e:
             logger.exception(f"Error generating optimization refinements: {e}")
             ph("cli-optimize-error-caught", {"error": str(e)})
-            return ""
+            return OptimizationReviewResult(review="", explanation="")
 
         if response.status_code == 200:
-            return cast("str", response.json()["review"])
+            data = response.json()
+            return OptimizationReviewResult(
+                review=cast("str", data["review"]), explanation=cast("str", data.get("review_explanation", ""))
+            )
         try:
             error = cast("str", response.json()["error"])
         except Exception:
             error = response.text
         logger.error(f"Error generating optimization review: {response.status_code} - {error}")
         ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
         console.rule()
-        return ""
+        return OptimizationReviewResult(review="", explanation="")
 
     def generate_workflow_steps(
         self,

codeflash/api/aiservice.py

@@ -7,6 +7,7 @@
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE
+from codeflash.code_utils.shell_utils import get_cross_platform_subprocess_run_args
 
 
 def trace_benchmarks_pytest(
@@ -17,20 +18,18 @@ def trace_benchmarks_pytest(
         benchmark_env["PYTHONPATH"] = str(project_root)
     else:
         benchmark_env["PYTHONPATH"] += os.pathsep + str(project_root)
-    result = subprocess.run(
+    run_args = get_cross_platform_subprocess_run_args(
+        cwd=project_root, env=benchmark_env, timeout=timeout, check=False, text=True, capture_output=True
+    )
+    result = subprocess.run(  # noqa: PLW1510
         [
             SAFE_SYS_EXECUTABLE,
             Path(__file__).parent / "pytest_new_process_trace_benchmarks.py",
             benchmarks_root,
             tests_root,
             trace_file,
         ],
-        cwd=project_root,
-        check=False,
-        capture_output=True,
-        text=True,
-        env=benchmark_env,
-        timeout=timeout,
+        **run_args,
     )
     if result.returncode != 0:
         if "ERROR collecting" in result.stdout:

codeflash/benchmarking/trace_benchmarks.py

@@ -58,6 +58,19 @@
     )
 
 
+class DummyTask:
+    def __init__(self) -> None:
+        self.id = 0
+
+
+class DummyProgress:
+    def __init__(self) -> None:
+        pass
+
+    def advance(self, task_id: TaskID, advance: int = 1) -> None:
+        pass
+
+
 def lsp_log(message: LspMessage) -> None:
     if not is_LSP_enabled():
         return
@@ -120,10 +133,6 @@ def progress_bar(
         logger.info(message)
 
         # Create a fake task ID since we still need to yield something
-        class DummyTask:
-            def __init__(self) -> None:
-                self.id = 0
-
         yield DummyTask().id
     else:
         progress = Progress(
@@ -141,6 +150,13 @@ def __init__(self) -> None:
 @contextmanager
 def test_files_progress_bar(total: int, description: str) -> Generator[tuple[Progress, TaskID], None, None]:
     """Progress bar for test files."""
+    if is_LSP_enabled():
+        lsp_log(LspTextMessage(text=description, takes_time=True))
+        dummy_progress = DummyProgress()
+        dummy_task = DummyTask()
+        yield dummy_progress, dummy_task.id
+        return
+
     with Progress(
         SpinnerColumn(next(spinners)),
         TextColumn("[progress.description]{task.description}"),

codeflash/cli_cmds/console.py

@@ -447,7 +447,7 @@ def replace_function_definitions_in_module(
 
     new_code: str = replace_functions_and_add_imports(
         # adding the global assignments before replacing the code, not after
-        # becuase of an "edge case" where the optimized code intoduced a new import and a global assignment using that import
+        # because of an "edge case" where the optimized code intoduced a new import and a global assignment using that import
         # and that import wasn't used before, so it was ignored when calling AddImportsVisitor.add_needed_import inside replace_functions_and_add_imports (because the global assignment wasn't added yet)
         # this was added at https://github.com/codeflash-ai/codeflash/pull/448
         add_global_assignments(code_to_apply, source_code) if should_add_global_assignments else source_code,

codeflash/code_utils/code_replacer.py

@@ -34,7 +34,7 @@ def get_git_diff(
             only_this_commit + "^1", only_this_commit, ignore_blank_lines=True, ignore_space_at_eol=True
         )
     elif uncommitted_changes:
-        uni_diff_text = repository.git.diff(None, "HEAD", ignore_blank_lines=True, ignore_space_at_eol=True)
+        uni_diff_text = repository.git.diff("HEAD", ignore_blank_lines=True, ignore_space_at_eol=True)
     else:
         uni_diff_text = repository.git.diff(
             commit.hexsha + "^1", commit.hexsha, ignore_blank_lines=True, ignore_space_at_eol=True

codeflash/code_utils/git_utils.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 import configparser
+import shutil
+import stat
 import subprocess
 import tempfile
 import time
 from pathlib import Path
-from typing import Optional
+from typing import Any, Callable, Optional
 
 import git
 
@@ -95,10 +97,24 @@ def create_detached_worktree(module_root: Path) -> Optional[Path]:
         return worktree_dir
 
 
+def _handle_remove_readonly(
+    func: Callable[[str], None], path: str, exc_info: tuple[type[BaseException], BaseException, Any]
+) -> None:
+    """Error handler for shutil.rmtree to handle read-only files on Windows."""
+    if isinstance(exc_info[1], PermissionError):
+        Path(path).chmod(stat.S_IWUSR | stat.S_IRUSR | stat.S_IXUSR)
+        func(path)
+    else:
+        raise exc_info[1]
+
+
 def remove_worktree(worktree_dir: Path) -> None:
+    """Remove a git worktree directory."""
+    if not worktree_dir.exists():
+        return
     try:
-        repository = git.Repo(worktree_dir, search_parent_directories=True)
-        repository.git.worktree("remove", "--force", worktree_dir)
+        shutil.rmtree(worktree_dir, onerror=_handle_remove_readonly)
+        logger.debug(f"Removed worktree: {worktree_dir}")
     except Exception:
         logger.exception(f"Failed to remove worktree: {worktree_dir}")
 

codeflash/code_utils/git_worktree_utils.py

@@ -3,6 +3,8 @@
 import contextlib
 import os
 import re
+import subprocess
+import sys
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
@@ -11,8 +13,11 @@
 from codeflash.either import Failure, Success
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
+
     from codeflash.either import Result
 
+
 # PowerShell patterns and prefixes
 POWERSHELL_RC_EXPORT_PATTERN = re.compile(
     r'^\$env:CODEFLASH_API_KEY\s*=\s*(?:"|\')?(cf-[^\s"\']+)(?:"|\')?\s*$', re.MULTILINE
@@ -231,3 +236,24 @@ def save_api_key_to_rc(api_key: str) -> Result[str, str]:
             f"To ensure your Codeflash API key is automatically loaded into your environment at startup, you can create {shell_rc_path} and add the following line:{LF}"
             f"{LF}{api_key_line}{LF}"
         )
+
+
+def get_cross_platform_subprocess_run_args(
+    cwd: Path | str | None = None,
+    env: Mapping[str, str] | None = None,
+    timeout: Optional[float] = None,
+    check: bool = False,  # noqa: FBT001, FBT002
+    text: bool = True,  # noqa: FBT001, FBT002
+    capture_output: bool = True,  # noqa: FBT001, FBT002 (only for non-Windows)
+) -> dict[str, str]:
+    run_args = {"cwd": cwd, "env": env, "text": text, "timeout": timeout, "check": check}
+    if sys.platform == "win32":
+        creationflags = subprocess.CREATE_NEW_PROCESS_GROUP
+        run_args["creationflags"] = creationflags
+        run_args["stdout"] = subprocess.PIPE
+        run_args["stderr"] = subprocess.PIPE
+        run_args["stdin"] = subprocess.DEVNULL
+    else:
+        run_args["capture_output"] = capture_output
+
+    return run_args