Merge remote-tracking branch 'origin/main' into opt-impact-aseem

Author: aseembits93

.github/workflows/e2e-async.yaml

@@ -0,0 +1,69 @@
+name: E2E - Async
+
+on:
+  pull_request:
+    paths:
+      - '**'  # Trigger for all paths
+
+  workflow_dispatch:
+
+jobs:
+  async-optimization:
+    # Dynamically determine if environment is needed only when workflow files change and contributor is external
+    environment: ${{ (github.event_name == 'workflow_dispatch' || (contains(toJSON(github.event.pull_request.files.*.filename), '.github/workflows/') && github.event.pull_request.user.login != 'misrasaurabh1' && github.event.pull_request.user.login != 'KRRT7')) && 'external-trusted-contributors' || '' }}
+
+    runs-on: ubuntu-latest
+    env:
+      CODEFLASH_AIS_SERVER: prod
+      POSTHOG_API_KEY: ${{ secrets.POSTHOG_API_KEY }}
+      CODEFLASH_API_KEY: ${{ secrets.CODEFLASH_API_KEY }}
+      COLUMNS: 110
+      MAX_RETRIES: 3
+      RETRY_DELAY: 5
+      EXPECTED_IMPROVEMENT_PCT: 10
+      CODEFLASH_END_TO_END: 1
+    steps:
+      - name: 🛎️ Checkout
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.pull_request.head.ref }}
+          repository: ${{ github.event.pull_request.head.repo.full_name }}
+          fetch-depth: 0
+          token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Validate PR
+        run: |
+          # Check for any workflow changes
+          if git diff --name-only "${{ github.event.pull_request.base.sha }}" "${{ github.event.pull_request.head.sha }}" | grep -q "^.github/workflows/"; then
+            echo "⚠️ Workflow changes detected."
+
+            # Get the PR author
+            AUTHOR="${{ github.event.pull_request.user.login }}"
+            echo "PR Author: $AUTHOR"
+
+            # Allowlist check
+            if [[ "$AUTHOR" == "misrasaurabh1" || "$AUTHOR" == "KRRT7" ]]; then
+              echo "✅ Authorized user ($AUTHOR). Proceeding."
+            elif [[ "${{ github.event.pull_request.state }}" == "open" ]]; then
+              echo "✅ PR triggered by 'pull_request_target' and is open. Assuming protection rules are in place. Proceeding."
+            else
+              echo "⛔ Unauthorized user ($AUTHOR) attempting to modify workflows. Exiting."
+              exit 1
+            fi
+          else
+            echo "✅ No workflow file changes detected. Proceeding."
+          fi
+
+      - name: Set up Python 3.11 for CLI
+        uses: astral-sh/setup-uv@v5
+        with:
+          python-version: 3.11.6
+
+      - name: Install dependencies (CLI)
+        run: |
+          uv sync
+          
+      - name: Run Codeflash to optimize async code
+        id: optimize_async_code
+        run: |
+          uv run python tests/scripts/end_to_end_test_async.py

.github/workflows/e2e-bubblesort-pytest-nogit.yaml

@@ -20,7 +20,7 @@ jobs:
       COLUMNS: 110
       MAX_RETRIES: 3
       RETRY_DELAY: 5
-      EXPECTED_IMPROVEMENT_PCT: 300
+      EXPECTED_IMPROVEMENT_PCT: 70
       CODEFLASH_END_TO_END: 1
     steps:
       - name: 🛎️ Checkout

.github/workflows/e2e-bubblesort-unittest.yaml

@@ -20,7 +20,7 @@ jobs:
       COLUMNS: 110
       MAX_RETRIES: 3
       RETRY_DELAY: 5
-      EXPECTED_IMPROVEMENT_PCT: 300
+      EXPECTED_IMPROVEMENT_PCT: 40
       CODEFLASH_END_TO_END: 1
     steps:
       - name: 🛎️ Checkout

code_to_optimize/async_bubble_sort.py

@@ -0,0 +1,43 @@
+import asyncio
+from typing import List, Union
+
+
+async def async_sorter(lst: List[Union[int, float]]) -> List[Union[int, float]]:
+    """
+    Async bubble sort implementation for testing.
+    """
+    print("codeflash stdout: Async sorting list")
+    
+    await asyncio.sleep(0.01)
+    
+    n = len(lst)
+    for i in range(n):
+        for j in range(0, n - i - 1):
+            if lst[j] > lst[j + 1]:
+                lst[j], lst[j + 1] = lst[j + 1], lst[j]
+    
+    result = lst.copy()
+    print(f"result: {result}")
+    return result
+
+
+class AsyncBubbleSorter:
+    """Class with async sorting method for testing."""
+    
+    async def sorter(self, lst: List[Union[int, float]]) -> List[Union[int, float]]:
+        """
+        Async bubble sort implementation within a class.
+        """
+        print("codeflash stdout: AsyncBubbleSorter.sorter() called")
+        
+        # Add some async delay
+        await asyncio.sleep(0.005)
+        
+        n = len(lst)
+        for i in range(n):
+            for j in range(0, n - i - 1):
+                if lst[j] > lst[j + 1]:
+                    lst[j], lst[j + 1] = lst[j + 1], lst[j]
+        
+        result = lst.copy()
+        return result

code_to_optimize/code_directories/async_e2e/main.py

@@ -0,0 +1,16 @@
+import time
+import asyncio
+
+
+async def retry_with_backoff(func, max_retries=3):
+    if max_retries < 1:
+        raise ValueError("max_retries must be at least 1")
+    last_exception = None
+    for attempt in range(max_retries):
+        try:
+            return await func()
+        except Exception as e:
+            last_exception = e
+            if attempt < max_retries - 1:
+                time.sleep(0.0001 * attempt)
+    raise last_exception

code_to_optimize/code_directories/async_e2e/pyproject.toml

@@ -0,0 +1,6 @@
+[tool.codeflash]
+disable-telemetry = true
+formatter-cmds = ["ruff check --exit-zero --fix $file", "ruff format $file"]
+module-root = "."
+test-framework = "pytest"
+tests-root = "tests"

code_to_optimize/code_directories/async_e2e/tests/__init__.py

@@ -105,6 +105,8 @@ def optimize_python_code(  # noqa: D417
         trace_id: str,
         num_candidates: int = 10,
         experiment_metadata: ExperimentMetadata | None = None,
+        *,
+        is_async: bool = False,
     ) -> list[OptimizedCandidate]:
         """Optimize the given python code for performance by making a request to the Django endpoint.
 
@@ -136,6 +138,7 @@ def optimize_python_code(  # noqa: D417
             "repo_owner": git_repo_owner,
             "repo_name": git_repo_name,
             "n_candidates": N_CANDIDATES_EFFECTIVE,
+            "is_async": is_async,
         }
 
         logger.info("!lsp|Generating optimized candidates…")
@@ -302,6 +305,9 @@ def get_new_explanation(  # noqa: D417
         annotated_tests: str,
         optimization_id: str,
         original_explanation: str,
+        original_throughput: str | None = None,
+        optimized_throughput: str | None = None,
+        throughput_improvement: str | None = None,
     ) -> str:
         """Optimize the given python code for performance by making a request to the Django endpoint.
 
@@ -318,6 +324,9 @@ def get_new_explanation(  # noqa: D417
         - 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
+        - original_throughput: str | None - throughput for the baseline code (operations per second)
+        - optimized_throughput: str | None - throughput for the optimized code (operations per second)
+        - throughput_improvement: str | None - throughput improvement percentage
 
         Returns
         -------
@@ -337,6 +346,9 @@ def get_new_explanation(  # noqa: D417
             "optimization_id": optimization_id,
             "original_explanation": original_explanation,
             "dependency_code": dependency_code,
+            "original_throughput": original_throughput,
+            "optimized_throughput": optimized_throughput,
+            "throughput_improvement": throughput_improvement,
         }
         logger.info("loading|Generating explanation")
         console.rule()
@@ -491,6 +503,7 @@ def generate_regression_tests(  # noqa: D417
             "test_index": test_index,
             "python_version": platform.python_version(),
             "codeflash_version": codeflash_version,
+            "is_async": function_to_optimize.is_async,
         }
         try:
             response = self.make_ai_service_request("/testgen", payload=payload, timeout=600)

codeflash/api/aiservice.py

@@ -26,7 +26,7 @@
 
 from packaging import version
 
-if os.environ.get("CODEFLASH_CFAPI_SERVER", default="prod").lower() == "local":
+if os.environ.get("CODEFLASH_CFAPI_SERVER", "prod").lower() == "local":
     CFAPI_BASE_URL = "http://localhost:3001"
     logger.info(f"Using local CF API at {CFAPI_BASE_URL}.")
     console.rule()

codeflash/api/cfapi.py

@@ -1,3 +1,4 @@
+import importlib.util
 import logging
 import sys
 from argparse import SUPPRESS, ArgumentParser, Namespace
@@ -96,6 +97,12 @@ def parse_args() -> Namespace:
     )
     parser.add_argument("--no-draft", default=False, action="store_true", help="Skip optimization for draft PRs")
     parser.add_argument("--worktree", default=False, action="store_true", help="Use worktree for optimization")
+    parser.add_argument(
+        "--async",
+        default=False,
+        action="store_true",
+        help="Enable optimization of async functions. By default, async functions are excluded from optimization.",
+    )
 
     args, unknown_args = parser.parse_known_args()
     sys.argv[:] = [sys.argv[0], *unknown_args]
@@ -139,6 +146,14 @@ def process_and_validate_cmd_args(args: Namespace) -> Namespace:
         if env_utils.is_ci():
             args.no_pr = True
 
+    if getattr(args, "async", False) and importlib.util.find_spec("pytest_asyncio") is None:
+        logger.warning(
+            "Warning: The --async flag requires pytest-asyncio to be installed.\n"
+            "Please install it using:\n"
+            '  pip install "codeflash[asyncio]"'
+        )
+        raise SystemExit(1)
+
     return args