Merge pull request #687 from codeflash-ai/granular-async-instrumentation

Granular async instrumentation

Author: KRRT7

.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/pre-commit.yaml

@@ -1,8 +1,9 @@
 repos:
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.12.7
+  rev: v0.13.1
   hooks:
     # Run the linter.
     - id: ruff-check
+      args: [ --config=pyproject.toml ]
     # Run the formatter.
     - id: ruff-format

.pre-commit-config.yaml

@@ -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/async_bubble_sort.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/main.py

@@ -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/pyproject.toml

@@ -70,7 +70,11 @@
                 "request": "launch",
                 "program": "${workspaceFolder:codeflash}/codeflash/main.py",
                 "args": [
-                    "--all",
+                    "--file",
+                    "src/async_examples/concurrency.py",
+                    "--function",
+                    "task",
+                    "--verbose"
                 ],
                 "cwd": "${input:chooseCwd}",
                 "console": "integratedTerminal",

code_to_optimize/code_directories/async_e2e/tests/__init__.py

@@ -101,6 +101,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.
 
@@ -131,6 +133,7 @@ def optimize_python_code(  # noqa: D417
             "current_username": get_last_commit_author_if_pr_exists(None),
             "repo_owner": git_repo_owner,
             "repo_name": git_repo_name,
+            "is_async": is_async,
         }
 
         logger.info("Generating optimized candidates…")
@@ -295,6 +298,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.
 
@@ -311,6 +317,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
         -------
@@ -330,6 +339,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("Generating explanation")
         console.rule()
@@ -439,7 +451,9 @@ 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)
         except requests.exceptions.RequestException as e:

codeflash.code-workspace

@@ -272,6 +272,8 @@ def _collect_imports_from_block(self, block: cst.IndentedBlock) -> None:
                         if child.module is None:
                             continue
                         module = self.get_full_dotted_name(child.module)
+                        if isinstance(child.names, cst.ImportStar):
+                            continue
                         for alias in child.names:
                             if isinstance(alias, cst.ImportAlias):
                                 name = alias.name.value
@@ -403,6 +405,73 @@ def add_global_assignments(src_module_code: str, dst_module_code: str) -> str:
     return transformed_module.code
 
 
+def resolve_star_import(module_name: str, project_root: Path) -> set[str]:
+    try:
+        module_path = module_name.replace(".", "/")
+        possible_paths = [project_root / f"{module_path}.py", project_root / f"{module_path}/__init__.py"]
+
+        module_file = None
+        for path in possible_paths:
+            if path.exists():
+                module_file = path
+                break
+
+        if module_file is None:
+            logger.warning(f"Could not find module file for {module_name}, skipping star import resolution")
+            return set()
+
+        with module_file.open(encoding="utf8") as f:
+            module_code = f.read()
+
+        tree = ast.parse(module_code)
+
+        all_names = None
+        for node in ast.walk(tree):
+            if (
+                isinstance(node, ast.Assign)
+                and len(node.targets) == 1
+                and isinstance(node.targets[0], ast.Name)
+                and node.targets[0].id == "__all__"
+            ):
+                if isinstance(node.value, (ast.List, ast.Tuple)):
+                    all_names = []
+                    for elt in node.value.elts:
+                        if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
+                            all_names.append(elt.value)
+                        elif isinstance(elt, ast.Str):  # Python < 3.8 compatibility
+                            all_names.append(elt.s)
+                break
+
+        if all_names is not None:
+            return set(all_names)
+
+        public_names = set()
+        for node in tree.body:
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+                if not node.name.startswith("_"):
+                    public_names.add(node.name)
+            elif isinstance(node, ast.Assign):
+                for target in node.targets:
+                    if isinstance(target, ast.Name) and not target.id.startswith("_"):
+                        public_names.add(target.id)
+            elif isinstance(node, ast.AnnAssign):
+                if isinstance(node.target, ast.Name) and not node.target.id.startswith("_"):
+                    public_names.add(node.target.id)
+            elif isinstance(node, ast.Import) or (
+                isinstance(node, ast.ImportFrom) and not any(alias.name == "*" for alias in node.names)
+            ):
+                for alias in node.names:
+                    name = alias.asname or alias.name
+                    if not name.startswith("_"):
+                        public_names.add(name)
+
+        return public_names  # noqa: TRY300
+
+    except Exception as e:
+        logger.warning(f"Error resolving star import for {module_name}: {e}")
+        return set()
+
+
 def add_needed_imports_from_module(
     src_module_code: str,
     dst_module_code: str,
@@ -457,9 +526,23 @@ def add_needed_imports_from_module(
                     f"{mod}.{obj}" in helper_functions_fqn or dst_context.full_module_name == mod  # avoid circular deps
                 ):
                     continue  # Skip adding imports for helper functions already in the context
-                if f"{mod}.{obj}" not in dotted_import_collector.imports:
-                    AddImportsVisitor.add_needed_import(dst_context, mod, obj)
-                RemoveImportsVisitor.remove_unused_import(dst_context, mod, obj)
+
+                # Handle star imports by resolving them to actual symbol names
+                if obj == "*":
+                    resolved_symbols = resolve_star_import(mod, project_root)
+                    logger.debug(f"Resolved star import from {mod}: {resolved_symbols}")
+
+                    for symbol in resolved_symbols:
+                        if (
+                            f"{mod}.{symbol}" not in helper_functions_fqn
+                            and f"{mod}.{symbol}" not in dotted_import_collector.imports
+                        ):
+                            AddImportsVisitor.add_needed_import(dst_context, mod, symbol)
+                        RemoveImportsVisitor.remove_unused_import(dst_context, mod, symbol)
+                else:
+                    if f"{mod}.{obj}" not in dotted_import_collector.imports:
+                        AddImportsVisitor.add_needed_import(dst_context, mod, obj)
+                    RemoveImportsVisitor.remove_unused_import(dst_context, mod, obj)
     except Exception as e:
         logger.exception(f"Error adding imports to destination module code: {e}")
         return dst_module_code