Skip to content

Commit 67bd717

Browse files
Merge pull request #466 from codeflash-ai/codeflash/optimize-pr384-2025-07-01T22.08.43
⚡️ Speed up method `FunctionRanker._get_function_stats` by 51% in PR #384 (`trace-and-optimize`)
2 parents b7258a9 + 72b51c1 commit 67bd717

File tree

1 file changed

+9
-8
lines changed

1 file changed

+9
-8
lines changed

codeflash/benchmarking/function_ranker.py

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
from __future__ import annotations
22

3+
from pathlib import Path
34
from typing import TYPE_CHECKING
45

56
from codeflash.cli_cmds.console import logger
67
from codeflash.code_utils.config_consts import DEFAULT_IMPORTANCE_THRESHOLD
8+
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
79
from codeflash.tracing.profile_stats import ProfileStats
810

911
if TYPE_CHECKING:
@@ -81,14 +83,13 @@ def load_function_stats(self) -> None:
8183
self._function_stats = {}
8284

8385
def _get_function_stats(self, function_to_optimize: FunctionToOptimize) -> dict | None:
84-
possible_keys = [
85-
f"{function_to_optimize.file_path}:{function_to_optimize.qualified_name}",
86-
f"{function_to_optimize.file_path}:{function_to_optimize.function_name}",
87-
]
88-
for key in possible_keys:
89-
if key in self._function_stats:
90-
return self._function_stats[key]
91-
return None
86+
# First try qualified_name, then function_name, avoid allocating a list
87+
key1 = f"{function_to_optimize.file_path}:{function_to_optimize.qualified_name}"
88+
stats = self._function_stats.get(key1)
89+
if stats is not None:
90+
return stats
91+
key2 = f"{function_to_optimize.file_path}:{function_to_optimize.function_name}"
92+
return self._function_stats.get(key2, None)
9293

9394
def get_function_ttx_score(self, function_to_optimize: FunctionToOptimize) -> float:
9495
stats = self._get_function_stats(function_to_optimize)

0 commit comments

Comments
 (0)