refactor how async decorators are applied at source site (#897)

* 5%

* better reporting for throughput

* first pass

* improve reporting for async optimizations

* remove some deduplication

* simplify usage

* refactor add_async_decorator_to_function

* doesn't even work

* fix mypy complants

* handle libcst exception

* fix type checking

* revert config

Author: KRRT7

codeflash/code_utils/instrument_existing_tests.py

@@ -684,27 +684,6 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
                     )
 
 
-def instrument_source_module_with_async_decorators(
-    source_path: Path, function_to_optimize: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR
-) -> tuple[bool, str | None]:
-    if not function_to_optimize.is_async:
-        return False, None
-
-    try:
-        with source_path.open(encoding="utf8") as f:
-            source_code = f.read()
-
-        modified_code, decorator_added = add_async_decorator_to_function(source_code, function_to_optimize, mode)
-
-        if decorator_added:
-            return True, modified_code
-
-    except Exception:
-        return False, None
-    else:
-        return False, None
-
-
 def inject_async_profiling_into_existing_test(
     test_path: Path,
     call_positions: list[CodePosition],
@@ -1288,25 +1267,29 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
 
 
 def add_async_decorator_to_function(
-    source_code: str, function: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR
-) -> tuple[str, bool]:
-    """Add async decorator to an async function definition.
+    source_path: Path, function: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR
+) -> bool:
+    """Add async decorator to an async function definition and write back to file.
 
     Args:
     ----
-        source_code: The source code to modify.
+        source_path: Path to the source file to modify in-place.
         function: The FunctionToOptimize object representing the target async function.
         mode: The testing mode to determine which decorator to apply.
 
     Returns:
     -------
-        Tuple of (modified_source_code, was_decorator_added).
+        Boolean indicating whether the decorator was successfully added.
 
     """
     if not function.is_async:
-        return source_code, False
+        return False
 
     try:
+        # Read source code
+        with source_path.open(encoding="utf8") as f:
+            source_code = f.read()
+
         module = cst.parse_module(source_code)
 
         # Add the decorator to the function
@@ -1318,10 +1301,17 @@ def add_async_decorator_to_function(
             import_transformer = AsyncDecoratorImportAdder(mode)
             module = module.visit(import_transformer)
 
-        return sort_imports(code=module.code, float_to_top=True), decorator_transformer.added_decorator
+        modified_code = sort_imports(code=module.code, float_to_top=True)
     except Exception as e:
         logger.exception(f"Error adding async decorator to function {function.qualified_name}: {e}")
-        return source_code, False
+        return False
+    else:
+        if decorator_transformer.added_decorator:
+            with source_path.open("w", encoding="utf8") as f:
+                f.write(modified_code)
+            logger.debug(f"Applied async {mode.value} instrumentation to {source_path}")
+            return True
+        return False
 
 
 def create_instrumented_source_module_path(source_path: Path, temp_dir: Path) -> Path:

codeflash/context/unused_definition_remover.py

@@ -469,22 +469,32 @@ def remove_unused_definitions_by_function_names(code: str, qualified_function_na
         qualified_function_names: Set of function names to keep. For methods, use format 'classname.methodname'
 
     """
-    module = cst.parse_module(code)
-    # Collect all definitions (top level classes, variables or function)
-    definitions = collect_top_level_definitions(module)
+    try:
+        module = cst.parse_module(code)
+    except Exception as e:
+        logger.debug(f"Failed to parse code with libcst: {type(e).__name__}: {e}")
+        return code
 
-    # Collect dependencies between definitions using the visitor pattern
-    dependency_collector = DependencyCollector(definitions)
-    module.visit(dependency_collector)
+    try:
+        # Collect all definitions (top level classes, variables or function)
+        definitions = collect_top_level_definitions(module)
 
-    # Mark definitions used by specified functions, and their dependencies recursively
-    usage_marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)
-    usage_marker.mark_used_definitions()
+        # Collect dependencies between definitions using the visitor pattern
+        dependency_collector = DependencyCollector(definitions)
+        module.visit(dependency_collector)
 
-    # Apply the recursive removal transformation
-    modified_module, _ = remove_unused_definitions_recursively(module, definitions)
+        # Mark definitions used by specified functions, and their dependencies recursively
+        usage_marker = QualifiedFunctionUsageMarker(definitions, qualified_function_names)
+        usage_marker.mark_used_definitions()
 
-    return modified_module.code if modified_module else ""
+        # Apply the recursive removal transformation
+        modified_module, _ = remove_unused_definitions_recursively(module, definitions)
+
+        return modified_module.code if modified_module else ""  # noqa: TRY300
+    except Exception as e:
+        # If any other error occurs during processing, return the original code
+        logger.debug(f"Error processing code to remove unused definitions: {type(e).__name__}: {e}")
+        return code
 
 
 def print_definitions(definitions: dict[str, UsageInfo]) -> None:

codeflash/github/PrComment.py

@@ -21,15 +21,17 @@ class PrComment:
     winning_behavior_test_results: TestResults
     winning_benchmarking_test_results: TestResults
     benchmark_details: Optional[list[BenchmarkDetail]] = None
+    original_async_throughput: Optional[int] = None
+    best_async_throughput: Optional[int] = None
 
-    def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str, Optional[list[BenchmarkDetail]]]]:
+    def to_json(self) -> dict[str, Union[str, int, dict[str, dict[str, int]], list[BenchmarkDetail], None]]:
         report_table = {
             test_type.to_name(): result
             for test_type, result in self.winning_behavior_test_results.get_test_pass_fail_report_by_type().items()
             if test_type.to_name()
         }
 
-        return {
+        result: dict[str, Union[str, int, dict[str, dict[str, int]], list[BenchmarkDetail], None]] = {
             "optimization_explanation": self.optimization_explanation,
             "best_runtime": humanize_runtime(self.best_runtime),
             "original_runtime": humanize_runtime(self.original_runtime),
@@ -42,6 +44,12 @@ def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str, Option
             "benchmark_details": self.benchmark_details if self.benchmark_details else None,
         }
 
+        if self.original_async_throughput is not None and self.best_async_throughput is not None:
+            result["original_async_throughput"] = str(self.original_async_throughput)
+            result["best_async_throughput"] = str(self.best_async_throughput)
+
+        return result
+
 
 class FileDiffContent(BaseModel):
     oldContent: str  # noqa: N815

codeflash/optimization/function_optimizer.py

@@ -607,26 +607,32 @@ def determine_best_candidate(
                         original_async_throughput=original_code_baseline.async_throughput,
                         best_throughput_until_now=None,
                     ) and quantity_of_tests_critic(candidate_result):
-                        tree.add("This candidate is faster than the original code. 🚀")  # TODO: Change this description
-                        tree.add(f"Original summed runtime: {humanize_runtime(original_code_baseline.runtime)}")
-                        tree.add(
-                            f"Best summed runtime: {humanize_runtime(candidate_result.best_test_runtime)} "
-                            f"(measured over {candidate_result.max_loop_count} "
-                            f"loop{'s' if candidate_result.max_loop_count > 1 else ''})"
-                        )
-                        tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%")
-                        tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X")
-                        if (
+                        # For async functions, prioritize throughput metrics over runtime
+                        is_async = (
                             original_code_baseline.async_throughput is not None
                             and candidate_result.async_throughput is not None
-                        ):
+                        )
+
+                        if is_async:
                             throughput_gain_value = throughput_gain(
                                 original_throughput=original_code_baseline.async_throughput,
                                 optimized_throughput=candidate_result.async_throughput,
                             )
+                            tree.add("This candidate has better async throughput than the original code. 🚀")
                             tree.add(f"Original async throughput: {original_code_baseline.async_throughput} executions")
                             tree.add(f"Optimized async throughput: {candidate_result.async_throughput} executions")
                             tree.add(f"Throughput improvement: {throughput_gain_value * 100:.1f}%")
+                            tree.add(f"Throughput ratio: {throughput_gain_value + 1:.3f}X")
+                        else:
+                            tree.add("This candidate is faster than the original code. 🚀")
+                            tree.add(f"Original summed runtime: {humanize_runtime(original_code_baseline.runtime)}")
+                            tree.add(
+                                f"Best summed runtime: {humanize_runtime(candidate_result.best_test_runtime)} "
+                                f"(measured over {candidate_result.max_loop_count} "
+                                f"loop{'s' if candidate_result.max_loop_count > 1 else ''})"
+                            )
+                            tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%")
+                            tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X")
                         line_profile_test_results = self.line_profiler_step(
                             code_context=code_context,
                             original_helper_code=original_helper_code,
@@ -681,22 +687,31 @@ def determine_best_candidate(
                                 )
                             )
                     else:
-                        tree.add(
-                            f"Summed runtime: {humanize_runtime(best_test_runtime)} "
-                            f"(measured over {candidate_result.max_loop_count} "
-                            f"loop{'s' if candidate_result.max_loop_count > 1 else ''})"
-                        )
-                        tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%")
-                        tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X")
-                        if (
+                        # For async functions, prioritize throughput metrics over runtime even for slow candidates
+                        is_async = (
                             original_code_baseline.async_throughput is not None
                             and candidate_result.async_throughput is not None
-                        ):
+                        )
+
+                        if is_async:
                             throughput_gain_value = throughput_gain(
                                 original_throughput=original_code_baseline.async_throughput,
                                 optimized_throughput=candidate_result.async_throughput,
                             )
-                            tree.add(f"Throughput gain: {throughput_gain_value * 100:.1f}%")
+                            tree.add(f"Async throughput: {candidate_result.async_throughput} executions")
+                            tree.add(f"Throughput change: {throughput_gain_value * 100:.1f}%")
+                            tree.add(
+                                f"(Runtime for reference: {humanize_runtime(best_test_runtime)} over "
+                                f"{candidate_result.max_loop_count} loop{'s' if candidate_result.max_loop_count > 1 else ''})"
+                            )
+                        else:
+                            tree.add(
+                                f"Summed runtime: {humanize_runtime(best_test_runtime)} "
+                                f"(measured over {candidate_result.max_loop_count} "
+                                f"loop{'s' if candidate_result.max_loop_count > 1 else ''})"
+                            )
+                            tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%")
+                            tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X")
 
                     if is_LSP_enabled():
                         lsp_log(LspMarkdownMessage(markdown=tree_to_markdown(tree)))
@@ -1502,16 +1517,21 @@ def process_review(
         raise_pr = not self.args.no_pr
         staging_review = self.args.staging_review
         opt_review_response = ""
-        if raise_pr or staging_review:
+        # Skip optimization review for async functions for now
+        if (raise_pr or staging_review) and not self.function_to_optimize.is_async:
             data["root_dir"] = git_root_dir()
             try:
                 opt_review_response = self.aiservice_client.get_optimization_review(
                     **data, calling_fn_details=function_references
                 )
             except Exception as e:
                 logger.debug(f"optimization review response failed, investigate {e}")
-            data["optimization_review"] = opt_review_response
+        # Always set optimization_review in data (empty string for async functions)
+        data["optimization_review"] = opt_review_response
         if raise_pr and not staging_review and opt_review_response != "low":
+            # Ensure root_dir is set for PR creation (needed for async functions that skip opt_review)
+            if "root_dir" not in data:
+                data["root_dir"] = git_root_dir()
             data["git_remote"] = self.args.git_remote
             check_create_pr(**data)
         elif staging_review:
@@ -1579,15 +1599,11 @@ def establish_original_code_baseline(
         test_env = self.get_test_env(codeflash_loop_index=0, codeflash_test_iteration=0, codeflash_tracer_disable=1)
 
         if self.function_to_optimize.is_async:
-            from codeflash.code_utils.instrument_existing_tests import instrument_source_module_with_async_decorators
+            from codeflash.code_utils.instrument_existing_tests import add_async_decorator_to_function
 
-            success, instrumented_source = instrument_source_module_with_async_decorators(
+            success = add_async_decorator_to_function(
                 self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.BEHAVIOR
             )
-            if success and instrumented_source:
-                with self.function_to_optimize.file_path.open("w", encoding="utf8") as f:
-                    f.write(instrumented_source)
-                logger.debug(f"Applied async instrumentation to {self.function_to_optimize.file_path}")
 
         # Instrument codeflash capture
         with progress_bar("Running tests to establish original code behavior..."):
@@ -1632,19 +1648,11 @@ def establish_original_code_baseline(
             console.rule()
             with progress_bar("Running performance benchmarks..."):
                 if self.function_to_optimize.is_async:
-                    from codeflash.code_utils.instrument_existing_tests import (
-                        instrument_source_module_with_async_decorators,
-                    )
+                    from codeflash.code_utils.instrument_existing_tests import add_async_decorator_to_function
 
-                    success, instrumented_source = instrument_source_module_with_async_decorators(
+                    add_async_decorator_to_function(
                         self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.PERFORMANCE
                     )
-                    if success and instrumented_source:
-                        with self.function_to_optimize.file_path.open("w", encoding="utf8") as f:
-                            f.write(instrumented_source)
-                        logger.debug(
-                            f"Applied async performance instrumentation to {self.function_to_optimize.file_path}"
-                        )
 
                 try:
                     benchmarking_results, _ = self.run_and_parse_tests(
@@ -1767,19 +1775,11 @@ def run_optimized_candidate(
             for module_abspath in original_helper_code:
                 candidate_helper_code[module_abspath] = Path(module_abspath).read_text("utf-8")
             if self.function_to_optimize.is_async:
-                from codeflash.code_utils.instrument_existing_tests import (
-                    instrument_source_module_with_async_decorators,
-                )
+                from codeflash.code_utils.instrument_existing_tests import add_async_decorator_to_function
 
-                success, instrumented_source = instrument_source_module_with_async_decorators(
+                add_async_decorator_to_function(
                     self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.BEHAVIOR
                 )
-                if success and instrumented_source:
-                    with self.function_to_optimize.file_path.open("w", encoding="utf8") as f:
-                        f.write(instrumented_source)
-                    logger.debug(
-                        f"Applied async behavioral instrumentation to {self.function_to_optimize.file_path} for candidate {optimization_candidate_index}"
-                    )
 
             try:
                 instrument_codeflash_capture(
@@ -1820,19 +1820,11 @@ def run_optimized_candidate(
             if test_framework == "pytest":
                 # For async functions, instrument at definition site for performance benchmarking
                 if self.function_to_optimize.is_async:
-                    from codeflash.code_utils.instrument_existing_tests import (
-                        instrument_source_module_with_async_decorators,
-                    )
+                    from codeflash.code_utils.instrument_existing_tests import add_async_decorator_to_function
 
-                    success, instrumented_source = instrument_source_module_with_async_decorators(
+                    add_async_decorator_to_function(
                         self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.PERFORMANCE
                     )
-                    if success and instrumented_source:
-                        with self.function_to_optimize.file_path.open("w", encoding="utf8") as f:
-                            f.write(instrumented_source)
-                        logger.debug(
-                            f"Applied async performance instrumentation to {self.function_to_optimize.file_path} for candidate {optimization_candidate_index}"
-                        )
 
                 try:
                     candidate_benchmarking_results, _ = self.run_and_parse_tests(

codeflash/result/create_pr.py

@@ -220,6 +220,8 @@ def check_create_pr(
                 winning_behavior_test_results=explanation.winning_behavior_test_results,
                 winning_benchmarking_test_results=explanation.winning_benchmarking_test_results,
                 benchmark_details=explanation.benchmark_details,
+                original_async_throughput=explanation.original_async_throughput,
+                best_async_throughput=explanation.best_async_throughput,
             ),
             existing_tests=existing_tests_source,
             generated_tests=generated_original_test_source,
@@ -270,6 +272,8 @@ def check_create_pr(
                 winning_behavior_test_results=explanation.winning_behavior_test_results,
                 winning_benchmarking_test_results=explanation.winning_benchmarking_test_results,
                 benchmark_details=explanation.benchmark_details,
+                original_async_throughput=explanation.original_async_throughput,
+                best_async_throughput=explanation.best_async_throughput,
             ),
             existing_tests=existing_tests_source,
             generated_tests=generated_original_test_source,