formatting changes

Author: aseembits93

codeflash/api/aiservice.py

@@ -172,6 +172,10 @@ def optimize_python_code_line_profiler(
 
         logger.info("Generating optimized candidates…")
         console.rule()
+        if line_profiler_results=="":
+            logger.info("No LineProfiler results were provided, Skipping optimization.")
+            console.rule()
+            return []
         try:
             response = self.make_ai_service_request("/optimize-line-profiler", payload=payload, timeout=600)
         except requests.exceptions.RequestException as e:

codeflash/code_utils/line_profile_utils.py

@@ -1,25 +1,25 @@
 """Adapted from line_profiler (https://github.com/pyutils/line_profiler) written by Enthought, Inc. (BSD License)"""
 from collections import defaultdict
+from pathlib import Path
+from typing import Union
 
 import isort
 import libcst as cst
-from pathlib import Path
-from typing import Union, List
-from libcst import ImportFrom, ImportAlias, Name
 
 from codeflash.code_utils.code_utils import get_run_tmp_file
 
 
 class LineProfilerDecoratorAdder(cst.CSTTransformer):
     """Transformer that adds a decorator to a function with a specific qualified name."""
-    #Todo we don't support nested functions yet so they can only be inside classes, dont use qualified names, instead use the structure
+
+    #TODO we don't support nested functions yet so they can only be inside classes, dont use qualified names, instead use the structure
     def __init__(self, qualified_name: str, decorator_name: str):
-        """
-        Initialize the transformer.
+        """Initialize the transformer.
 
         Args:
             qualified_name: The fully qualified name of the function to add the decorator to (e.g., "MyClass.nested_func.target_func").
             decorator_name: The name of the decorator to add.
+
         """
         super().__init__()
         self.qualified_name_parts = qualified_name.split(".")
@@ -45,7 +45,7 @@ def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.Fu
         function_name = original_node.name.value
 
         # Check if the current context path matches our target qualified name
-        if self._matches_qualified_path():
+        if self.context_stack==self.qualified_name_parts:
             # Check if the decorator is already present
             has_decorator = any(
                 self._is_target_decorator(decorator.decorator)
@@ -68,22 +68,11 @@ def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.Fu
         self.context_stack.pop()
         return updated_node
 
-    def _matches_qualified_path(self) -> bool:
-        """Check if the current context stack matches the qualified name."""
-        if len(self.context_stack) != len(self.qualified_name_parts):
-            return False
-
-        for i, name in enumerate(self.qualified_name_parts):
-            if self.context_stack[i] != name:
-                return False
-
-        return True
-
     def _is_target_decorator(self, decorator_node: Union[cst.Name, cst.Attribute, cst.Call]) -> bool:
         """Check if a decorator matches our target decorator name."""
         if isinstance(decorator_node, cst.Name):
             return decorator_node.value == self.decorator_name
-        elif isinstance(decorator_node, cst.Call) and isinstance(decorator_node.func, cst.Name):
+        if isinstance(decorator_node, cst.Call) and isinstance(decorator_node.func, cst.Name):
             return decorator_node.func.value == self.decorator_name
         return False
 
@@ -147,8 +136,7 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
         return updated_node.with_changes(body=new_body)
 
 def add_decorator_to_qualified_function(module, qualified_name, decorator_name):
-    """
-    Add a decorator to a function with the exact qualified name in the source code.
+    """Add a decorator to a function with the exact qualified name in the source code.
 
     Args:
         module: The Python source code as a string.
@@ -157,6 +145,7 @@ def add_decorator_to_qualified_function(module, qualified_name, decorator_name):
 
     Returns:
         The modified source code as a string.
+
     """
     # Parse the source code into a CST
 
@@ -168,7 +157,7 @@ def add_decorator_to_qualified_function(module, qualified_name, decorator_name):
     return modified_module
 
 def add_profile_enable(original_code: str, line_profile_output_file: str) -> str:
-    # todo modify by using a libcst transformer
+    # TODO modify by using a libcst transformer
     module = cst.parse_module(original_code)
     transformer = ProfileEnableTransformer(line_profile_output_file)
     modified_module = module.visit(transformer)
@@ -217,7 +206,7 @@ def add_decorator_imports(function_to_optimize, code_context):
         module_node = cst.parse_module(file_contents)
         for fn_name in fns_present:
             # add decorator
-            module_node = add_decorator_to_qualified_function(module_node, fn_name, 'codeflash_line_profile')
+            module_node = add_decorator_to_qualified_function(module_node, fn_name, "codeflash_line_profile")
         # add imports
         # Create a transformer to add the import
         transformer = ImportAdder("from line_profiler import profile as codeflash_line_profile")
@@ -231,4 +220,4 @@ def add_decorator_imports(function_to_optimize, code_context):
     file_contents = function_to_optimize.file_path.read_text("utf-8")
     modified_code = add_profile_enable(file_contents,str(line_profile_output_file))
     function_to_optimize.file_path.write_text(modified_code,"utf-8")
-    return line_profile_output_file
+    return line_profile_output_file

codeflash/optimization/function_optimizer.py

@@ -232,7 +232,8 @@ def optimize_function(self) -> Result[BestOptimization, str]:
         ):
             cleanup_paths(paths_to_cleanup)
             return Failure("The threshold for test coverage was not met.")
-        # request for new optimizations but don't block execution, check for completion later, only adding to control set right now
+        # request for new optimizations but don't block execution, check for completion later
+        # adding to control and experiment set but with same traceid
         best_optimization = None
 
         for _u, candidates in enumerate([optimizations_set.control, optimizations_set.experiment]):
@@ -359,30 +360,36 @@ def determine_best_candidate(
         )
         console.rule()
         candidates = deque(candidates)
+        # Start a new thread for AI service request, start loop in main thread
+        # check if aiservice request is complete, when it is complete, append result to the candidates list
         with concurrent.futures.ThreadPoolExecutor() as executor:
-            future_line_profile_results = executor.submit(self.aiservice_client.optimize_python_code_line_profiler,
-            source_code=code_context.read_writable_code,
-            dependency_code=code_context.read_only_context_code,
-            trace_id=self.function_trace_id,
-            line_profiler_results=original_code_baseline.line_profile_results['str_out'],
-            num_candidates = 10,
-            experiment_metadata = None)
+            future_line_profile_results = executor.submit(
+                self.aiservice_client.optimize_python_code_line_profiler,
+                source_code=code_context.read_writable_code,
+                dependency_code=code_context.read_only_context_code,
+                trace_id=self.function_trace_id,
+                line_profiler_results=original_code_baseline.line_profile_results["str_out"],
+                num_candidates=10,
+                experiment_metadata=None,
+            )
             try:
                 candidate_index = 0
                 done = False
+                original_len = len(candidates)
                 while candidates:
-                #for candidate_index, candidate in enumerate(candidates, start=1):
+                    # for candidate_index, candidate in enumerate(candidates, start=1):
                     done = True if future_line_profile_results is None else future_line_profile_results.done()
                     if done and (future_line_profile_results is not None):
                         line_profile_results = future_line_profile_results.result()
                         candidates.extend(line_profile_results)
-                        logger.info(f"Added result from line profiler to candidates: {len(line_profile_results)}")
+                        original_len+= len(candidates)
+                        logger.info(f"Added results from line profiler to candidates, total candidates now: {original_len}")
                         future_line_profile_results = None
                     candidate_index += 1
                     candidate = candidates.popleft()
                     get_run_tmp_file(Path(f"test_return_values_{candidate_index}.bin")).unlink(missing_ok=True)
                     get_run_tmp_file(Path(f"test_return_values_{candidate_index}.sqlite")).unlink(missing_ok=True)
-                    logger.info(f"Optimization candidate {candidate_index}/{len(candidates)}:")
+                    logger.info(f"Optimization candidate {candidate_index}/{original_len}:")
                     code_print(candidate.source_code)
                     try:
                         did_update = self.replace_function_and_helpers_with_optimized_code(
@@ -397,7 +404,9 @@ def determine_best_candidate(
                     except (ValueError, SyntaxError, cst.ParserSyntaxError, AttributeError) as e:
                         logger.error(e)
                         self.write_code_and_helpers(
-                            self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
+                            self.function_to_optimize_source_code,
+                            original_helper_code,
+                            self.function_to_optimize.file_path,
                         )
                         continue
 
@@ -781,7 +790,7 @@ def establish_original_code_baseline(
         original_helper_code: dict[Path, str],
         file_path_to_helper_classes: dict[Path, set[str]],
     ) -> Result[tuple[OriginalCodeBaseline, list[str]], str]:
-        line_profile_results = {'timings':{},'unit':0, 'str_out':''}
+        line_profile_results = {"timings": {}, "unit": 0, "str_out": ""}
         # For the original function - run the tests and get the runtime, plus coverage
         with progress_bar(f"Establishing original code baseline for {self.function_to_optimize.function_name}"):
             assert (test_framework := self.args.test_framework) in ["pytest", "unittest"]
@@ -826,8 +835,7 @@ def establish_original_code_baseline(
                 return Failure("The threshold for test coverage was not met.")
             if test_framework == "pytest":
                 try:
-                    line_profiler_output_file = add_decorator_imports(
-                        self.function_to_optimize, code_context)
+                    line_profiler_output_file = add_decorator_imports(self.function_to_optimize, code_context)
                     line_profile_results, _ = self.run_and_parse_tests(
                         testing_type=TestingMode.LINE_PROFILE,
                         test_env=test_env,
@@ -843,7 +851,7 @@ def establish_original_code_baseline(
                     self.write_code_and_helpers(
                         self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
                     )
-                if line_profile_results['str_out']=='':
+                if line_profile_results["str_out"] == "":
                     logger.warning(
                         f"Couldn't run line profiler for original function {self.function_to_optimize.function_name}"
                     )
@@ -1081,7 +1089,7 @@ def run_and_parse_tests(
                     pytest_min_loops=1,
                     pytest_max_loops=1,
                     test_framework=self.test_cfg.test_framework,
-                    line_profiler_output_file=line_profiler_output_file
+                    line_profiler_output_file=line_profiler_output_file,
                 )
             elif testing_type == TestingMode.PERFORMANCE:
                 result_file_path, run_result = run_benchmarking_tests(