Merge pull request #2055 from codeflash-ai/perf/defer-cli-imports

perf: defer cli.py imports for 7.7x faster --help

Author: KRRT7

.github/workflows/ci.yaml

@@ -199,6 +199,8 @@ jobs:
         run: |
           uv run ruff check --fix . || true
           uv run ruff format .
+          # uv-dynamic-versioning rewrites version.py on every `uv run` — discard those changes
+          git checkout HEAD -- codeflash/version.py codeflash-benchmark/codeflash_benchmark/version.py 2>/dev/null || true
 
       - name: Commit and push fixes
         run: |

benchmarks/__init__.py

@@ -0,0 +1,72 @@
+"""Benchmark CLI startup latency for codeflash compare --script mode.
+
+Run from a worktree root. Installs deps via uv sync, then times several
+CLI entry points and writes a JSON file mapping command names to median
+wall-clock seconds.
+
+Usage:
+    codeflash compare main codeflash/optimize \
+        --script "python benchmarks/bench_cli_startup.py" \
+        --script-output benchmarks/results.json
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import time
+from pathlib import Path
+
+WARMUP = 3
+RUNS = 30
+OUTPUT = os.environ.get("BENCH_OUTPUT", "benchmarks/results.json")
+
+COMMANDS: dict[str, list[str]] = {
+    "version": ["uv", "run", "codeflash", "--version"],
+    "help": ["uv", "run", "codeflash", "--help"],
+    "auth_status": ["uv", "run", "codeflash", "auth", "status"],
+    "compare_help": ["uv", "run", "codeflash", "compare", "--help"],
+}
+
+
+def measure(cmd: list[str], warmup: int = WARMUP, runs: int = RUNS) -> float:
+    """Return median wall-clock seconds for *cmd* over *runs* iterations."""
+    env = {**os.environ, "CODEFLASH_API_KEY": "bench_dummy_key"}
+    for _ in range(warmup):
+        subprocess.run(cmd, capture_output=True, check=False, env=env)
+
+    times: list[float] = []
+    for _ in range(runs):
+        t0 = time.perf_counter()
+        subprocess.run(cmd, capture_output=True, check=False, env=env)
+        times.append(time.perf_counter() - t0)
+
+    times.sort()
+    mid = len(times) // 2
+    return times[mid] if len(times) % 2 else (times[mid - 1] + times[mid]) / 2
+
+
+def main() -> None:
+    # Ensure deps are installed in the worktree
+    subprocess.run(["uv", "sync"], check=True, capture_output=True)
+
+    results: dict[str, float] = {}
+    for name, cmd in COMMANDS.items():
+        print(f"  {name}: ", end="", flush=True)
+        median = measure(cmd)
+        results[name] = round(median, 4)
+        print(f"{median * 1000:.0f} ms")
+
+    # Total = sum of medians (useful for a single summary number)
+    results["__total__"] = round(sum(results.values()), 4)
+
+    output_path = Path(OUTPUT)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with output_path.open("w") as f:
+        json.dump(results, f, indent=2)
+    print(f"\nResults written to {OUTPUT}")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/bench_cli_startup.py

@@ -4,6 +4,7 @@
 
 import libcst as cst
 
+import codeflash.code_utils._libcst_cache  # noqa: F401
 from codeflash.code_utils.formatter import sort_imports
 
 if TYPE_CHECKING:

codeflash/benchmarking/instrument_codeflash_trace.py

@@ -5,15 +5,6 @@
 from functools import lru_cache
 from pathlib import Path
 
-from codeflash.cli_cmds import logging_config
-from codeflash.cli_cmds.console import apologize_and_exit, logger
-from codeflash.code_utils import env_utils
-from codeflash.code_utils.code_utils import exit_with_message, normalize_ignore_paths
-from codeflash.code_utils.config_parser import parse_config_file
-from codeflash.languages.test_framework import set_current_test_framework
-from codeflash.lsp.helpers import is_LSP_enabled
-from codeflash.version import __version__ as version
-
 
 def parse_args() -> Namespace:
     parser = _build_parser()
@@ -30,12 +21,17 @@ def parse_args() -> Namespace:
 
 
 def process_and_validate_cmd_args(args: Namespace) -> Namespace:
+    from codeflash.cli_cmds import logging_config
+    from codeflash.cli_cmds.console import logger
+    from codeflash.code_utils import env_utils
+    from codeflash.code_utils.code_utils import exit_with_message
     from codeflash.code_utils.git_utils import (
         check_running_in_git_repo,
         confirm_proceeding_with_no_git_repo,
         get_repo_owner_and_name,
     )
     from codeflash.code_utils.github_utils import require_github_app_or_exit
+    from codeflash.version import __version__ as version
 
     if args.server:
         os.environ["CODEFLASH_AIS_SERVER"] = args.server
@@ -85,6 +81,12 @@ def process_and_validate_cmd_args(args: Namespace) -> Namespace:
 
 
 def process_pyproject_config(args: Namespace) -> Namespace:
+    from codeflash.code_utils import env_utils
+    from codeflash.code_utils.code_utils import exit_with_message, normalize_ignore_paths
+    from codeflash.code_utils.config_parser import parse_config_file
+    from codeflash.languages.test_framework import set_current_test_framework
+    from codeflash.lsp.helpers import is_LSP_enabled
+
     try:
         pyproject_config, pyproject_file_path = parse_config_file(args.config_file)
     except ValueError as e:
@@ -222,6 +224,9 @@ def project_root_from_module_root(module_root: Path, pyproject_file_path: Path)
 
 
 def handle_optimize_all_arg_parsing(args: Namespace) -> Namespace:
+    from codeflash.cli_cmds.console import apologize_and_exit, logger
+    from codeflash.code_utils.code_utils import exit_with_message
+
     if hasattr(args, "all") or (hasattr(args, "file") and args.file):
         no_pr = getattr(args, "no_pr", False)
 

codeflash/cli_cmds/cli.py

@@ -2,17 +2,17 @@
 
 import os
 
-import click
-
-from codeflash.cli_cmds.console import console
-from codeflash.cli_cmds.oauth_handler import perform_oauth_signin
-from codeflash.code_utils.env_utils import get_codeflash_api_key
-from codeflash.code_utils.shell_utils import save_api_key_to_rc
-from codeflash.either import is_successful
-
 
 def auth_login() -> None:
     """Perform OAuth login and save the API key."""
+    import click
+
+    from codeflash.cli_cmds.console import console
+    from codeflash.cli_cmds.oauth_handler import perform_oauth_signin
+    from codeflash.code_utils.env_utils import get_codeflash_api_key
+    from codeflash.code_utils.shell_utils import save_api_key_to_rc
+    from codeflash.either import is_successful
+
     try:
         existing_api_key = get_codeflash_api_key()
     except OSError:
@@ -41,6 +41,9 @@ def auth_login() -> None:
 
 def auth_status() -> None:
     """Check and display current authentication status."""
+    from codeflash.cli_cmds.console import console
+    from codeflash.code_utils.env_utils import get_codeflash_api_key
+
     try:
         api_key = get_codeflash_api_key()
     except OSError:

codeflash/cli_cmds/cmd_auth.py

@@ -0,0 +1,64 @@
+"""Cache libcst visitor dispatch table construction.
+
+libcst's ``MatcherDecoratableTransformer`` and
+``MatcherDecoratableVisitor`` rebuild visitor dispatch tables on
+every instantiation by iterating ``dir(self)`` (~600 attributes)
+and calling ``getattr`` + ``inspect.ismethod`` on each.  The
+results depend only on the class, not the instance, so caching
+by ``type(obj)`` is safe.
+
+Import this module before any libcst visitors are instantiated
+to install the cache.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import libcst.matchers._visitors as _mv
+
+_visit_cache: dict[type, Any] = {}
+_leave_cache: dict[type, Any] = {}
+_matchers_cache: dict[type, Any] = {}
+
+_original_visit = _mv._gather_constructed_visit_funcs  # noqa: SLF001
+_original_leave = _mv._gather_constructed_leave_funcs  # noqa: SLF001
+_original_matchers = _mv._gather_matchers  # noqa: SLF001
+
+
+def _cached_visit(obj: object) -> Any:
+    """Return cached visit-function dispatch table for the object's class."""
+    cls = type(obj)
+    try:
+        return _visit_cache[cls]
+    except KeyError:
+        result = _original_visit(obj)
+        _visit_cache[cls] = result
+        return result
+
+
+def _cached_leave(obj: object) -> Any:
+    """Return cached leave-function dispatch table for the object's class."""
+    cls = type(obj)
+    try:
+        return _leave_cache[cls]
+    except KeyError:
+        result = _original_leave(obj)
+        _leave_cache[cls] = result
+        return result
+
+
+def _cached_matchers(obj: object) -> Any:
+    """Return cached matcher dispatch table for the object's class."""
+    cls = type(obj)
+    try:
+        return dict(_matchers_cache[cls])
+    except KeyError:
+        result = _original_matchers(obj)
+        _matchers_cache[cls] = result
+        return dict(result)
+
+
+_mv._gather_constructed_visit_funcs = _cached_visit  # noqa: SLF001
+_mv._gather_constructed_leave_funcs = _cached_leave  # noqa: SLF001
+_mv._gather_matchers = _cached_matchers  # noqa: SLF001

codeflash/code_utils/_libcst_cache.py

@@ -9,17 +9,16 @@
 from pathlib import Path
 from typing import Any, Optional
 
-from codeflash.cli_cmds.console import logger
-from codeflash.code_utils.code_utils import exit_with_message
-from codeflash.code_utils.formatter import format_code
 from codeflash.code_utils.shell_utils import read_api_key_from_shell_config, save_api_key_to_rc
-from codeflash.languages.registry import get_language_support_by_common_formatters
-from codeflash.lsp.helpers import is_LSP_enabled
 
 
 def check_formatter_installed(
     formatter_cmds: list[str], exit_on_failure: bool = True, language: str = "python"
 ) -> bool:
+    from codeflash.cli_cmds.console import logger
+    from codeflash.code_utils.formatter import format_code
+    from codeflash.languages.registry import get_language_support_by_common_formatters
+
     if not formatter_cmds or formatter_cmds[0] == "disabled":
         return True
     first_cmd = formatter_cmds[0]
@@ -69,6 +68,8 @@ def check_formatter_installed(
 
 @lru_cache(maxsize=1)
 def get_codeflash_api_key() -> str:
+    from codeflash.cli_cmds.console import logger
+
     # Check environment variable first
     env_api_key = os.environ.get("CODEFLASH_API_KEY")
     shell_api_key = read_api_key_from_shell_config()
@@ -96,7 +97,8 @@ def get_codeflash_api_key() -> str:
     # Prefer the shell configuration over environment variables for lsp,
     # as the API key may change in the RC file during lsp runtime. Since the LSP client (extension) can restart
     # within the same process, the environment variable could become outdated.
-    api_key = shell_api_key or env_api_key if is_LSP_enabled() else env_api_key or shell_api_key
+    is_lsp = os.getenv("CODEFLASH_LSP", default="false").lower() == "true"
+    api_key = shell_api_key or env_api_key if is_lsp else env_api_key or shell_api_key
 
     api_secret_docs_message = "For more information, refer to the documentation at [https://docs.codeflash.ai/optimizing-with-codeflash/codeflash-github-actions#manual-setup]."  # noqa
     if not api_key:
@@ -106,6 +108,8 @@ def get_codeflash_api_key() -> str:
             f"{api_secret_docs_message}"
         )
         if is_repo_a_fork():
+            from codeflash.code_utils.code_utils import exit_with_message
+
             msg = (
                 "Codeflash API key not detected in your environment. It appears you're running Codeflash from a GitHub fork.\n"
                 "For external contributors, please ensure you've added your own API key to your fork's repository secrets and set it as the CODEFLASH_API_KEY environment variable.\n"
@@ -124,6 +128,8 @@ def get_codeflash_api_key() -> str:
 
 
 def ensure_codeflash_api_key() -> bool:
+    from codeflash.cli_cmds.console import logger
+
     try:
         get_codeflash_api_key()
     except OSError:

codeflash/code_utils/env_utils.py

@@ -7,6 +7,7 @@
 
 import libcst as cst
 
+import codeflash.code_utils._libcst_cache  # noqa: F401
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import get_run_tmp_file, module_name_from_file_path
 from codeflash.code_utils.formatter import sort_imports

codeflash/code_utils/instrument_existing_tests.py

@@ -8,7 +8,6 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
-from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.compat import LF
 from codeflash.either import Failure, Success
 
@@ -41,6 +40,8 @@ def is_powershell() -> bool:
     2. COMSPEC pointing to powershell.exe
     3. TERM_PROGRAM indicating Windows Terminal (often uses PowerShell)
     """
+    from codeflash.cli_cmds.console import logger
+
     if os.name != "nt":
         return False
 
@@ -72,6 +73,8 @@ def is_powershell() -> bool:
 
 def read_api_key_from_shell_config() -> Optional[str]:
     """Read API key from shell configuration file."""
+    from codeflash.cli_cmds.console import logger
+
     shell_rc_path = get_shell_rc_path()
     # Ensure shell_rc_path is a Path object for consistent handling
     if not isinstance(shell_rc_path, Path):
@@ -127,6 +130,8 @@ def get_api_key_export_line(api_key: str) -> str:
 
 def save_api_key_to_rc(api_key: str) -> Result[str, str]:
     """Save API key to the appropriate shell configuration file."""
+    from codeflash.cli_cmds.console import logger
+
     shell_rc_path = get_shell_rc_path()
     # Ensure shell_rc_path is a Path object for consistent handling
     if not isinstance(shell_rc_path, Path):