feat: concurrency-aware async benchmarking

codeflash/api/aiservice.py

@@ -403,6 +403,10 @@ def get_new_explanation(  # noqa: D417
         optimized_throughput: str | None = None,
         throughput_improvement: str | None = None,
         function_references: str | None = None,
+        acceptance_reason: str | None = None,
+        original_concurrency_ratio: str | None = None,
+        optimized_concurrency_ratio: str | None = None,
+        concurrency_improvement: str | None = None,
         codeflash_version: str = codeflash_version,
     ) -> str:
         """Optimize the given python code for performance by making a request to the Django endpoint.
@@ -423,8 +427,12 @@ def get_new_explanation(  # noqa: D417
         - 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
-        - current codeflash version
         - function_references: str | None - where the function is called in the codebase
+        - acceptance_reason: str | None - why the optimization was accepted (runtime, throughput, or concurrency)
+        - original_concurrency_ratio: str | None - concurrency ratio for the baseline code
+        - optimized_concurrency_ratio: str | None - concurrency ratio for the optimized code
+        - concurrency_improvement: str | None - concurrency improvement percentage
+        - codeflash_version: str - current codeflash version
 
         Returns
         -------
@@ -448,6 +456,10 @@ def get_new_explanation(  # noqa: D417
             "optimized_throughput": optimized_throughput,
             "throughput_improvement": throughput_improvement,
             "function_references": function_references,
+            "acceptance_reason": acceptance_reason,
+            "original_concurrency_ratio": original_concurrency_ratio,
+            "optimized_concurrency_ratio": optimized_concurrency_ratio,
+            "concurrency_improvement": concurrency_improvement,
             "codeflash_version": codeflash_version,
             "call_sequence": self.get_next_sequence(),
         }

codeflash/code_utils/codeflash_wrap_decorator.py

@@ -4,6 +4,7 @@
 import gc
 import os
 import sqlite3
+import time
 from enum import Enum
 from functools import wraps
 from pathlib import Path
@@ -165,3 +166,45 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
         return return_value
 
     return async_wrapper
+
+
+def codeflash_concurrency_async(func: F) -> F:
+    """Measures concurrent vs sequential execution performance for async functions."""
+
+    @wraps(func)
+    async def async_wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+        function_name = func.__name__
+        concurrency_factor = int(os.environ.get("CODEFLASH_CONCURRENCY_FACTOR", "10"))
+
+        test_module_name = os.environ.get("CODEFLASH_TEST_MODULE", "")
+        test_class_name = os.environ.get("CODEFLASH_TEST_CLASS", "")
+        test_function = os.environ.get("CODEFLASH_TEST_FUNCTION", "")
+        loop_index = os.environ.get("CODEFLASH_LOOP_INDEX", "0")
+
+        # Phase 1: Sequential execution timing
+        gc.disable()
+        try:
+            seq_start = time.perf_counter_ns()
+            for _ in range(concurrency_factor):
+                result = await func(*args, **kwargs)
+            sequential_time = time.perf_counter_ns() - seq_start
+        finally:
+            gc.enable()
+
+        # Phase 2: Concurrent execution timing
+        gc.disable()
+        try:
+            conc_start = time.perf_counter_ns()
+            tasks = [func(*args, **kwargs) for _ in range(concurrency_factor)]
+            await asyncio.gather(*tasks)
+            concurrent_time = time.perf_counter_ns() - conc_start
+        finally:
+            gc.enable()
+
+        # Output parseable metrics
+        tag = f"{test_module_name}:{test_class_name}:{test_function}:{function_name}:{loop_index}"
+        print(f"!@######CONC:{tag}:{sequential_time}:{concurrent_time}:{concurrency_factor}######@!")
+
+        return result
+
+    return async_wrapper

codeflash/code_utils/config_consts.py

@@ -4,6 +4,8 @@
 N_CANDIDATES = 5
 MIN_IMPROVEMENT_THRESHOLD = 0.05
 MIN_THROUGHPUT_IMPROVEMENT_THRESHOLD = 0.10  # 10% minimum improvement for async throughput
+MIN_CONCURRENCY_IMPROVEMENT_THRESHOLD = 0.20  # 20% concurrency ratio improvement required
+CONCURRENCY_FACTOR = 10  # Number of concurrent executions for concurrency benchmark
 MAX_TEST_FUNCTION_RUNS = 50
 MAX_CUMULATIVE_TEST_RUNTIME_NANOSECONDS = 100e6  # 100ms
 N_TESTS_TO_GENERATE = 2

codeflash/code_utils/instrument_existing_tests.py

@@ -1109,9 +1109,12 @@ def __init__(self, function: FunctionToOptimize, mode: TestingMode = TestingMode
         self.added_decorator = False
 
         # Choose decorator based on mode
-        self.decorator_name = (
-            "codeflash_behavior_async" if mode == TestingMode.BEHAVIOR else "codeflash_performance_async"
-        )
+        if mode == TestingMode.BEHAVIOR:
+            self.decorator_name = "codeflash_behavior_async"
+        elif mode == TestingMode.CONCURRENCY:
+            self.decorator_name = "codeflash_concurrency_async"
+        else:
+            self.decorator_name = "codeflash_performance_async"
 
     def visit_ClassDef(self, node: cst.ClassDef) -> None:
         # Track when we enter a class
@@ -1154,12 +1157,14 @@ def _is_target_decorator(self, decorator_node: cst.Name | cst.Attribute | cst.Ca
                 "codeflash_trace_async",
                 "codeflash_behavior_async",
                 "codeflash_performance_async",
+                "codeflash_concurrency_async",
             }
         if isinstance(decorator_node, cst.Call) and isinstance(decorator_node.func, cst.Name):
             return decorator_node.func.value in {
                 "codeflash_trace_async",
                 "codeflash_behavior_async",
                 "codeflash_performance_async",
+                "codeflash_concurrency_async",
             }
         return False
 
@@ -1171,6 +1176,14 @@ def __init__(self, mode: TestingMode = TestingMode.BEHAVIOR) -> None:
         self.mode = mode
         self.has_import = False
 
+    def _get_decorator_name(self) -> str:
+        """Get the decorator name based on the testing mode."""
+        if self.mode == TestingMode.BEHAVIOR:
+            return "codeflash_behavior_async"
+        if self.mode == TestingMode.CONCURRENCY:
+            return "codeflash_concurrency_async"
+        return "codeflash_performance_async"
+
     def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
         # Check if the async decorator import is already present
         if (
@@ -1182,9 +1195,7 @@ def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
             and node.module.attr.value == "codeflash_wrap_decorator"
             and not isinstance(node.names, cst.ImportStar)
         ):
-            decorator_name = (
-                "codeflash_behavior_async" if self.mode == TestingMode.BEHAVIOR else "codeflash_performance_async"
-            )
+            decorator_name = self._get_decorator_name()
             for import_alias in node.names:
                 if import_alias.name.value == decorator_name:
                     self.has_import = True
@@ -1195,9 +1206,7 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
             return updated_node
 
         # Choose import based on mode
-        decorator_name = (
-            "codeflash_behavior_async" if self.mode == TestingMode.BEHAVIOR else "codeflash_performance_async"
-        )
+        decorator_name = self._get_decorator_name()
 
         # Parse the import statement into a CST node
         import_node = cst.parse_statement(f"from codeflash.code_utils.codeflash_wrap_decorator import {decorator_name}")

codeflash/models/models.py

@@ -160,6 +160,7 @@ class BestOptimization(BaseModel):
     winning_replay_benchmarking_test_results: Optional[TestResults] = None
     line_profiler_test_results: dict
     async_throughput: Optional[int] = None
+    concurrency_metrics: Optional[ConcurrencyMetrics] = None
 
 
 @dataclass(frozen=True)
@@ -171,6 +172,14 @@ def __str__(self) -> str:
         return f"{self.module_path}::{self.function_name}"
 
 
+@dataclass
+class ConcurrencyMetrics:
+    sequential_time_ns: int
+    concurrent_time_ns: int
+    concurrency_factor: int
+    concurrency_ratio: float  # sequential_time / concurrent_time
+
+
 @dataclass
 class BenchmarkDetail:
     benchmark_name: str
@@ -335,6 +344,7 @@ class OptimizedCandidateResult(BaseModel):
     optimization_candidate_index: int
     total_candidate_timing: int
     async_throughput: Optional[int] = None
+    concurrency_metrics: Optional[ConcurrencyMetrics] = None
 
 
 class GeneratedTests(BaseModel):
@@ -526,6 +536,7 @@ class OriginalCodeBaseline(BaseModel):
     runtime: int
     coverage_results: Optional[CoverageData]
     async_throughput: Optional[int] = None
+    concurrency_metrics: Optional[ConcurrencyMetrics] = None
 
 
 class CoverageStatus(Enum):
@@ -617,6 +628,7 @@ class TestingMode(enum.Enum):
     BEHAVIOR = "behavior"
     PERFORMANCE = "performance"
     LINE_PROFILE = "line_profile"
+    CONCURRENCY = "concurrency"
 
 
 # TODO this class is duplicated in codeflash_capture