throughput

Author: KRRT7

codeflash/optimization/function_optimizer.py

@@ -90,7 +90,7 @@
 from codeflash.verification.equivalence import compare_test_results
 from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture
 from codeflash.verification.parse_line_profile_test_output import parse_line_profile_results
-from codeflash.verification.parse_test_output import calculate_function_throughput_from_stdout, parse_test_results
+from codeflash.verification.parse_test_output import calculate_function_throughput_from_test_results, parse_test_results
 from codeflash.verification.test_runner import run_behavioral_tests, run_benchmarking_tests, run_line_profile_tests
 from codeflash.verification.verification_utils import get_test_file_path
 from codeflash.verification.verifier import generate_tests
@@ -1528,14 +1528,9 @@ def establish_original_code_baseline(
 
             async_throughput = None
             if self.function_to_optimize.is_async:
-                all_stdout = ""
-                for result in benchmarking_results.test_results:
-                    if result.stdout:
-                        all_stdout += result.stdout
                 logger.info("Calculating async function throughput from test output...")
-                logger.info(f"All stdout for async throughput calculation:\n{all_stdout}")
-                async_throughput = calculate_function_throughput_from_stdout(
-                    all_stdout, self.function_to_optimize.function_name
+                async_throughput = calculate_function_throughput_from_test_results(
+                    benchmarking_results, self.function_to_optimize.function_name
                 )
                 logger.info(f"Original async function throughput: {async_throughput} calls/second")
 
@@ -1700,13 +1695,8 @@ def run_optimized_candidate(
 
             candidate_async_throughput = None
             if self.function_to_optimize.is_async and candidate_benchmarking_results:
-                all_stdout = ""
-                for result in candidate_benchmarking_results.test_results:
-                    if result.stdout:
-                        all_stdout += result.stdout
-
-                candidate_async_throughput = calculate_function_throughput_from_stdout(
-                    all_stdout, self.function_to_optimize.function_name
+                candidate_async_throughput = calculate_function_throughput_from_test_results(
+                    candidate_benchmarking_results, self.function_to_optimize.function_name
                 )
 
             if self.args.benchmark:

codeflash/verification/parse_test_output.py

@@ -44,26 +44,27 @@ def parse_func(file_path: Path) -> XMLParser:
 end_pattern = re.compile(r"!######([^:]*):([^:]*):([^:]*):([^:]*):([^:]+)######!")
 
 
-def calculate_function_throughput_from_stdout(stdout: str, function_name: str) -> int:
-    """Calculate function throughput from stdout. A completed execution is defined as having both a start tag and matching end tag.
+def calculate_function_throughput_from_test_results(test_results: TestResults, function_name: str) -> int:
+    """Calculate function throughput from TestResults by extracting stdout.
 
+    A completed execution is defined as having both a start tag and matching end tag.
     Start: !$######test_module:test_function:function_name:loop_index:iteration_id######$!
     End:   !######test_module:test_function:function_name:loop_index:iteration_id######!
     """
-    start_matches = start_pattern.findall(stdout)
-    end_matches = end_pattern.findall(stdout)
+    all_stdout = ""
+    for result in test_results.test_results:
+        if result.stdout:
+            all_stdout += result.stdout
+
+    start_matches = start_pattern.findall(all_stdout)
+    end_matches = end_pattern.findall(all_stdout)
     end_matches_set = set(end_matches)
 
-    # Count completed executions for the specific function only
     function_throughput = 0
     logger.info(f"Total start matches: {len(start_matches)}, Total end matches: {len(end_matches)}")
     for start_match in start_matches:
-        # Check if this execution is for the function we're interested in and has a matching end tag
-        # function_name is at index 2 in the match tuple
-        logger.info(f"Start match: {start_match}")
-        logger.info(f"End matches: {end_matches_set}")
-        logger.info(f"Function name: {function_name}")
         if start_match in end_matches_set and len(start_match) > 2 and start_match[2] == function_name:
+            logger.info(f"Matched start-end pair for function '{function_name}': {start_match}")
             function_throughput += 1
 
     return function_throughput