update CI

Author: KRRT7

codeflash/benchmarking/function_ranker.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 from codeflash.cli_cmds.console import console, logger
@@ -83,13 +82,17 @@ def load_function_stats(self) -> None:
             self._function_stats = {}
 
     def _get_function_stats(self, function_to_optimize: FunctionToOptimize) -> dict | None:
-        # First try qualified_name, then function_name, avoid allocating a list
-        key1 = f"{function_to_optimize.file_path}:{function_to_optimize.qualified_name}"
-        stats = self._function_stats.get(key1)
-        if stats is not None:
-            return stats
-        key2 = f"{function_to_optimize.file_path}:{function_to_optimize.function_name}"
-        return self._function_stats.get(key2, None)
+        target_filename = function_to_optimize.file_path.name
+        for key, stats in self._function_stats.items():
+            if stats.get("function_name") == function_to_optimize.function_name and (
+                key.endswith(f"/{target_filename}") or target_filename in key
+            ):
+                return stats
+
+        logger.debug(
+            f"Could not find stats for function {function_to_optimize.function_name} in file {target_filename}"
+        )
+        return None
 
     def get_function_ttx_score(self, function_to_optimize: FunctionToOptimize) -> float:
         stats = self._get_function_stats(function_to_optimize)

tests/scripts/end_to_end_test_tracer_replay.py

@@ -8,7 +8,7 @@ def run_test(expected_improvement_pct: int) -> bool:
     config = TestConfig(
         trace_mode=True,
         min_improvement_x=0.1,
-        expected_unit_tests=7,
+        expected_unit_tests=8,
         coverage_expectations=[
             CoverageExpectation(function_name="funcA", expected_coverage=100.0, expected_lines=[5, 6, 7, 8, 10, 13])
         ],

tests/scripts/end_to_end_test_utilities.py

@@ -185,10 +185,9 @@ def validate_stdout_in_candidate(stdout: str, expected_in_stdout: list[str]) ->
 
 
 def run_trace_test(cwd: pathlib.Path, config: TestConfig, expected_improvement_pct: int) -> bool:
-    # First command: Run the tracer
     test_root = cwd / "tests" / (config.test_framework or "")
     clear_directory(test_root)
-    command = ["python", "-m", "codeflash.tracer", "-o", "codeflash.trace", "workload.py"]
+    command = ["python", "-m", "codeflash.main", "optimize", "workload.py"]
     process = subprocess.Popen(
         command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, cwd=str(cwd), env=os.environ.copy()
     )
@@ -202,33 +201,20 @@ def run_trace_test(cwd: pathlib.Path, config: TestConfig, expected_improvement_p
     stdout = "".join(output)
 
     if return_code != 0:
-        logging.error(f"Tracer command returned exit code {return_code}")
+        logging.error(f"Tracer with optimization command returned exit code {return_code}")
         return False
 
-    functions_traced = re.search(r"Traced (\d+) function calls successfully and replay test created at - (.*)$", stdout)
-    if not functions_traced or int(functions_traced.group(1)) != 13:
-        logging.error("Expected 13 traced functions")
+    functions_traced = re.search(r"Traced (\d+) function calls successfully", stdout)
+    logging.info(functions_traced.groups() if functions_traced else "No functions traced")
+    if not functions_traced:
+        logging.error("Failed to find traced functions in output")
         return False
-
-    replay_test_path = pathlib.Path(functions_traced.group(2))
-    if not replay_test_path.exists():
-        logging.error(f"Replay test file missing at {replay_test_path}")
+    if int(functions_traced.group(1)) != 13:
+        logging.error(functions_traced.groups())
+        logging.error("Expected 13 traced functions")
         return False
 
-    # Second command: Run optimization
-    command = ["python", "../../../codeflash/main.py", "--replay-test", str(replay_test_path), "--no-pr"]
-    process = subprocess.Popen(
-        command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, cwd=str(cwd), env=os.environ.copy()
-    )
-
-    output = []
-    for line in process.stdout:
-        logging.info(line.strip())
-        output.append(line)
-
-    return_code = process.wait()
-    stdout = "".join(output)
-
+    # Validate optimization results (from optimization phase)
     return validate_output(stdout, return_code, expected_improvement_pct, config)