calculate in own file time

remove unittests remnants

Author: KRRT7

codeflash/benchmarking/function_ranker.py

@@ -100,7 +100,11 @@ def rank_functions(self, functions_to_optimize: list[FunctionToOptimize]) -> lis
         """Ranks and filters functions based on their ttX score and importance.
 
         Filters out functions whose own_time is less than DEFAULT_IMPORTANCE_THRESHOLD
-        of total runtime, then ranks the remaining functions by ttX score.
+        of file-relative runtime, then ranks the remaining functions by ttX score.
+
+        Importance is calculated relative to functions in the same file(s) rather than
+        total program time. This avoids filtering out functions due to test infrastructure
+        overhead.
 
         The ttX score prioritizes functions that are computationally heavy themselves
         or that make expensive calls to other functions.
@@ -116,9 +120,24 @@ def rank_functions(self, functions_to_optimize: list[FunctionToOptimize]) -> lis
             logger.warning("No function stats available to rank functions.")
             return []
 
-        total_program_time = sum(
-            s["own_time_ns"] for s in self._function_stats.values() if s.get("own_time_ns", 0) > 0
-        )
+        # Calculate total time from functions in the same file(s) as functions to optimize
+        if functions_to_optimize:
+            # Get unique files from functions to optimize
+            target_files = {func.file_path.name for func in functions_to_optimize}
+            # Calculate total time only from functions in these files
+            total_program_time = sum(
+                s["own_time_ns"]
+                for s in self._function_stats.values()
+                if s.get("own_time_ns", 0) > 0 and any(target_file in s["filename"] for target_file in target_files)
+            )
+            logger.debug(
+                f"Using file-relative importance for {len(target_files)} file(s): {target_files}. "
+                f"Total file time: {total_program_time:,} ns"
+            )
+        else:
+            total_program_time = sum(
+                s["own_time_ns"] for s in self._function_stats.values() if s.get("own_time_ns", 0) > 0
+            )
 
         if total_program_time == 0:
             logger.warning("Total program time is zero, cannot determine function importance.")

codeflash/benchmarking/replay_test.py

@@ -66,28 +66,23 @@ def get_unique_test_name(module: str, function_name: str, benchmark_name: str, c
 def create_trace_replay_test_code(
     trace_file: str,
     functions_data: list[dict[str, Any]],
-    test_framework: str = "pytest",
-    max_run_count=256,  # noqa: ANN001
+    max_run_count: int = 256,
 ) -> str:
     """Create a replay test for functions based on trace data.
 
     Args:
     ----
         trace_file: Path to the SQLite database file
         functions_data: List of dictionaries with function info extracted from DB
-        test_framework: 'pytest' or 'unittest'
         max_run_count: Maximum number of runs to include in the test
 
     Returns:
     -------
         A string containing the test code
 
     """
-    assert test_framework in ["pytest", "unittest"]
-
     # Create Imports
-    imports = f"""from codeflash.picklepatch.pickle_patcher import PicklePatcher as pickle
-{"import unittest" if test_framework == "unittest" else ""}
+    imports = """from codeflash.picklepatch.pickle_patcher import PicklePatcher as pickle
 from codeflash.benchmarking.replay_test import get_next_arg_and_return
 """
 
@@ -158,13 +153,7 @@ def create_trace_replay_test_code(
     )
 
     # Create main body
-
-    if test_framework == "unittest":
-        self = "self"
-        test_template = "\nclass TestTracedFunctions(unittest.TestCase):\n"
-    else:
-        test_template = ""
-        self = ""
+    test_template = ""
 
     for func in functions_data:
         module_name = func.get("module_name")
@@ -223,30 +212,28 @@ def create_trace_replay_test_code(
                     filter_variables=filter_variables,
                 )
 
-        formatted_test_body = textwrap.indent(test_body, "        " if test_framework == "unittest" else "    ")
+        formatted_test_body = textwrap.indent(test_body, "    ")
 
-        test_template += "    " if test_framework == "unittest" else ""
         unique_test_name = get_unique_test_name(module_name, function_name, benchmark_function_name, class_name)
-        test_template += f"def test_{unique_test_name}({self}):\n{formatted_test_body}\n"
+        test_template += f"def test_{unique_test_name}():\n{formatted_test_body}\n"
 
     return imports + "\n" + metadata + "\n" + test_template
 
 
 def generate_replay_test(
-    trace_file_path: Path, output_dir: Path, test_framework: str = "pytest", max_run_count: int = 100
+    trace_file_path: Path, output_dir: Path, max_run_count: int = 100
 ) -> int:
     """Generate multiple replay tests from the traced function calls, grouped by benchmark.
 
     Args:
     ----
         trace_file_path: Path to the SQLite database file
         output_dir: Directory to write the generated tests (if None, only returns the code)
-        test_framework: 'pytest' or 'unittest'
         max_run_count: Maximum number of runs to include per function
 
     Returns:
     -------
-        Dictionary mapping benchmark names to generated test code
+        The number of replay tests generated
 
     """
     count = 0
@@ -295,7 +282,6 @@ def generate_replay_test(
             test_code = create_trace_replay_test_code(
                 trace_file=trace_file_path.as_posix(),
                 functions_data=functions_data,
-                test_framework=test_framework,
                 max_run_count=max_run_count,
             )
             test_code = sort_imports(code=test_code)

codeflash/tracing/replay_test.py

@@ -46,13 +46,9 @@ def get_function_alias(module: str, function_name: str) -> str:
 def create_trace_replay_test(
     trace_file: str,
     functions: list[FunctionModules],
-    test_framework: str = "pytest",
-    max_run_count=100,  # noqa: ANN001
+    max_run_count: int = 100,
 ) -> str:
-    assert test_framework in {"pytest", "unittest"}
-
-    imports = f"""import dill as pickle
-{"import unittest" if test_framework == "unittest" else ""}
+    imports = """import dill as pickle
 from codeflash.tracing.replay_test import get_next_arg_and_return
 """
 
@@ -112,12 +108,7 @@ def create_trace_replay_test(
             ret = {class_name_alias}{method_name}(**args)
             """
     )
-    if test_framework == "unittest":
-        self = "self"
-        test_template = "\nclass TestTracedFunctions(unittest.TestCase):\n"
-    else:
-        test_template = ""
-        self = ""
+    test_template = ""
     for func, func_property in zip(functions, function_properties):
         if func_property is None:
             continue
@@ -167,9 +158,8 @@ def create_trace_replay_test(
                 max_run_count=max_run_count,
                 filter_variables=filter_variables,
             )
-        formatted_test_body = textwrap.indent(test_body, "        " if test_framework == "unittest" else "    ")
+        formatted_test_body = textwrap.indent(test_body, "    ")
 
-        test_template += "    " if test_framework == "unittest" else ""
-        test_template += f"def test_{alias}({self}):\n{formatted_test_body}\n"
+        test_template += f"def test_{alias}():\n{formatted_test_body}\n"
 
     return imports + "\n" + metadata + "\n" + test_template

codeflash/tracing/tracing_new_process.py

@@ -110,7 +110,6 @@ def __init__(
         self._db_lock = threading.Lock()
 
         self.con = None
-        self.output_file = Path(output).resolve()
         self.functions = functions
         self.function_modules: list[FunctionModules] = []
         self.function_count = defaultdict(int)
@@ -126,6 +125,14 @@ def __init__(
         self.ignored_functions = {"<listcomp>", "<genexpr>", "<dictcomp>", "<setcomp>", "<lambda>", "<module>"}
 
         self.sanitized_filename = self.sanitize_to_filename(command)
+        # Place trace file next to replay tests in the tests directory
+        from codeflash.verification.verification_utils import get_test_file_path
+        function_path = "_".join(functions) if functions else self.sanitized_filename
+        test_file_path = get_test_file_path(
+            test_dir=Path(config["tests_root"]), function_name=function_path, test_type="replay"
+        )
+        trace_filename = test_file_path.stem + ".trace"
+        self.output_file = test_file_path.parent / trace_filename
         self.result_pickle_file_path = result_pickle_file_path
 
         assert timeout is None or timeout > 0, "Timeout should be greater than 0"
@@ -142,7 +149,6 @@ def __init__(
         self.timer = time.process_time_ns
         self.total_tt = 0
         self.simulate_call("profiler")
-        assert "test_framework" in self.config, "Please specify 'test-framework' in pyproject.toml config file"
         self.t = self.timer()
 
         # Store command information for metadata table
@@ -275,7 +281,6 @@ def __exit__(
         replay_test = create_trace_replay_test(
             trace_file=self.output_file,
             functions=self.function_modules,
-            test_framework=self.config["test_framework"],
             max_run_count=self.max_function_count,
         )
         function_path = "_".join(self.functions) if self.functions else self.sanitized_filename
@@ -770,11 +775,11 @@ def make_pstats_compatible(self) -> None:
         self.files = []
         self.top_level = []
         new_stats = {}
-        for func, (cc, ns, tt, ct, callers) in self.stats.items():
+        for func, (cc, ns, tt, ct, callers) in list(self.stats.items()):
             new_callers = {(k[0], k[1], k[2]): v for k, v in callers.items()}
             new_stats[(func[0], func[1], func[2])] = (cc, ns, tt, ct, new_callers)
         new_timings = {}
-        for func, (cc, ns, tt, ct, callers) in self.timings.items():
+        for func, (cc, ns, tt, ct, callers) in list(self.timings.items()):
             new_callers = {(k[0], k[1], k[2]): v for k, v in callers.items()}
             new_timings[(func[0], func[1], func[2])] = (cc, ns, tt, ct, new_callers)
         self.stats = new_stats