diff --git a/codeflash/api/aiservice.py b/codeflash/api/aiservice.py index ded1b64ee..62e1df174 100644 --- a/codeflash/api/aiservice.py +++ b/codeflash/api/aiservice.py @@ -10,7 +10,7 @@ from pydantic.json import pydantic_encoder from codeflash.cli_cmds.console import console, logger -from codeflash.code_utils.config_consts import get_n_candidates, get_n_candidates_lp +from codeflash.code_utils.config_consts import N_CANDIDATES_EFFECTIVE, N_CANDIDATES_LP_EFFECTIVE from codeflash.code_utils.env_utils import get_codeflash_api_key from codeflash.code_utils.git_utils import get_last_commit_author_if_pr_exists, get_repo_owner_and_name from codeflash.lsp.helpers import is_LSP_enabled @@ -132,7 +132,7 @@ def optimize_python_code( # noqa: D417 "current_username": get_last_commit_author_if_pr_exists(None), "repo_owner": git_repo_owner, "repo_name": git_repo_name, - "n_candidates": get_n_candidates(), + "n_candidates": N_CANDIDATES_EFFECTIVE, } logger.info("!lsp|Generating optimized candidates…") @@ -194,7 +194,7 @@ def optimize_python_code_line_profiler( # noqa: D417 "experiment_metadata": experiment_metadata, "codeflash_version": codeflash_version, "lsp_mode": is_LSP_enabled(), - "n_candidates_lp": get_n_candidates_lp(), + "n_candidates_lp": N_CANDIDATES_LP_EFFECTIVE, } console.rule() diff --git a/codeflash/code_utils/config_consts.py b/codeflash/code_utils/config_consts.py index d5c4cdc48..eab4f0fd1 100644 --- a/codeflash/code_utils/config_consts.py +++ b/codeflash/code_utils/config_consts.py @@ -19,26 +19,17 @@ TOTAL_LOOPING_TIME_LSP = 10.0 # Kept same timing for LSP mode to avoid in increase in performance reporting N_CANDIDATES_LP_LSP = 3 +MAX_N_CANDIDATES = 5 +MAX_N_CANDIDATES_LP = 6 -def get_n_candidates() -> int: +try: from codeflash.lsp.helpers import is_LSP_enabled - return N_CANDIDATES_LSP if is_LSP_enabled() else N_CANDIDATES + _IS_LSP_ENABLED = is_LSP_enabled() +except ImportError: + _IS_LSP_ENABLED = False - -def get_n_candidates_lp() -> int: - from codeflash.lsp.helpers import is_LSP_enabled - - return N_CANDIDATES_LP_LSP if is_LSP_enabled() else N_CANDIDATES_LP - - -def get_n_tests_to_generate() -> int: - from codeflash.lsp.helpers import is_LSP_enabled - - return N_TESTS_TO_GENERATE_LSP if is_LSP_enabled() else N_TESTS_TO_GENERATE - - -def get_total_looping_time() -> float: - from codeflash.lsp.helpers import is_LSP_enabled - - return TOTAL_LOOPING_TIME_LSP if is_LSP_enabled() else TOTAL_LOOPING_TIME +N_CANDIDATES_EFFECTIVE = min(N_CANDIDATES_LSP if _IS_LSP_ENABLED else N_CANDIDATES, MAX_N_CANDIDATES) +N_CANDIDATES_LP_EFFECTIVE = min(N_CANDIDATES_LP_LSP if _IS_LSP_ENABLED else N_CANDIDATES_LP, MAX_N_CANDIDATES_LP) +N_TESTS_TO_GENERATE_EFFECTIVE = N_TESTS_TO_GENERATE_LSP if _IS_LSP_ENABLED else N_TESTS_TO_GENERATE +TOTAL_LOOPING_TIME_EFFECTIVE = TOTAL_LOOPING_TIME_LSP if _IS_LSP_ENABLED else TOTAL_LOOPING_TIME diff --git a/codeflash/code_utils/git_utils.py b/codeflash/code_utils/git_utils.py index dbff075ba..40a725692 100644 --- a/codeflash/code_utils/git_utils.py +++ b/codeflash/code_utils/git_utils.py @@ -16,7 +16,7 @@ from unidiff import PatchSet from codeflash.cli_cmds.console import logger -from codeflash.code_utils.config_consts import get_n_candidates +from codeflash.code_utils.config_consts import N_CANDIDATES_EFFECTIVE if TYPE_CHECKING: from git import Repo @@ -164,7 +164,7 @@ def create_git_worktrees( ) -> tuple[Path | None, list[Path]]: if git_root and worktree_root_dir: worktree_root = Path(tempfile.mkdtemp(dir=worktree_root_dir)) - worktrees = [Path(tempfile.mkdtemp(dir=worktree_root)) for _ in range(get_n_candidates() + 1)] + worktrees = [Path(tempfile.mkdtemp(dir=worktree_root)) for _ in range(N_CANDIDATES_EFFECTIVE + 1)] for worktree in worktrees: subprocess.run(["git", "worktree", "add", "-d", worktree], cwd=module_root, check=True) else: diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 40c5cd8d4..926adff95 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -44,11 +44,11 @@ from codeflash.code_utils.config_consts import ( COVERAGE_THRESHOLD, INDIVIDUAL_TESTCASE_TIMEOUT, + N_CANDIDATES_EFFECTIVE, + N_CANDIDATES_LP_EFFECTIVE, + N_TESTS_TO_GENERATE_EFFECTIVE, REPEAT_OPTIMIZATION_PROBABILITY, - get_n_candidates, - get_n_candidates_lp, - get_n_tests_to_generate, - get_total_looping_time, + TOTAL_LOOPING_TIME_EFFECTIVE, ) from codeflash.code_utils.deduplicate_code import normalize_code from codeflash.code_utils.edit_generated_tests import ( @@ -237,7 +237,7 @@ def __init__( self.generate_and_instrument_tests_results: ( tuple[GeneratedTestsList, dict[str, set[FunctionCalledInTest]], OptimizationSet] | None ) = None - n_tests = get_n_tests_to_generate() + n_tests = N_TESTS_TO_GENERATE_EFFECTIVE self.executor = concurrent.futures.ThreadPoolExecutor( max_workers=n_tests + 2 if self.experiment_id is None else n_tests + 3 ) @@ -289,7 +289,7 @@ def generate_and_instrument_tests( ] ]: """Generate and instrument tests, returning all necessary data for optimization.""" - n_tests = get_n_tests_to_generate() + n_tests = N_TESTS_TO_GENERATE_EFFECTIVE generated_test_paths = [ get_test_file_path( self.test_cfg.tests_root, self.function_to_optimize.function_name, test_index, test_type="unit" @@ -487,7 +487,7 @@ def determine_best_candidate( dependency_code=code_context.read_only_context_code, trace_id=self.function_trace_id[:-4] + exp_type if self.experiment_id else self.function_trace_id, line_profiler_results=original_code_baseline.line_profile_results["str_out"], - num_candidates=get_n_candidates_lp(), + num_candidates=N_CANDIDATES_LP_EFFECTIVE, experiment_metadata=ExperimentMetadata( id=self.experiment_id, group="control" if exp_type == "EXP0" else "experiment" ) @@ -1061,7 +1061,7 @@ def generate_tests_and_optimizations( generated_perf_test_paths: list[Path], run_experiment: bool = False, # noqa: FBT001, FBT002 ) -> Result[tuple[GeneratedTestsList, dict[str, set[FunctionCalledInTest]], OptimizationSet], str]: - n_tests = get_n_tests_to_generate() + n_tests = N_TESTS_TO_GENERATE_EFFECTIVE assert len(generated_test_paths) == n_tests console.rule() # Submit the test generation task as future @@ -1072,7 +1072,7 @@ def generate_tests_and_optimizations( generated_test_paths, generated_perf_test_paths, ) - n_candidates = get_n_candidates() + n_candidates = N_CANDIDATES_EFFECTIVE future_optimization_candidates = self.executor.submit( self.aiservice_client.optimize_python_code, read_writable_code.markdown, @@ -1482,7 +1482,7 @@ def establish_original_code_baseline( instrument_codeflash_capture( self.function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root ) - total_looping_time = get_total_looping_time() + total_looping_time = TOTAL_LOOPING_TIME_EFFECTIVE behavioral_results, coverage_results = self.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, @@ -1623,7 +1623,7 @@ def run_optimized_candidate( self.function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root ) - total_looping_time = get_total_looping_time() + total_looping_time = TOTAL_LOOPING_TIME_EFFECTIVE candidate_behavior_results, _ = self.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, @@ -1678,7 +1678,7 @@ def run_optimized_candidate( start_time: float = time.time() loop_count = 0 for i in range(100): - if i >= 5 and time.time() - start_time >= get_total_looping_time() * 1.5: + if i >= 5 and time.time() - start_time >= TOTAL_LOOPING_TIME_EFFECTIVE * 1.5: # * 1.5 to give unittest a bit more time to run break test_env["CODEFLASH_LOOP_INDEX"] = str(i + 1) @@ -1687,7 +1687,7 @@ def run_optimized_candidate( test_env=test_env, test_files=self.test_files, optimization_iteration=optimization_candidate_index, - testing_time=get_total_looping_time(), + testing_time=TOTAL_LOOPING_TIME_EFFECTIVE, unittest_loop_index=i + 1, ) loop_count = i + 1 @@ -1726,7 +1726,7 @@ def run_and_parse_tests( test_env: dict[str, str], test_files: TestFiles, optimization_iteration: int, - testing_time: float = get_total_looping_time(), + testing_time: float = TOTAL_LOOPING_TIME_EFFECTIVE, *, enable_coverage: bool = False, pytest_min_loops: int = 5, @@ -1877,7 +1877,7 @@ def line_profiler_step( test_env=test_env, test_files=self.test_files, optimization_iteration=0, - testing_time=get_total_looping_time(), + testing_time=TOTAL_LOOPING_TIME_EFFECTIVE, enable_coverage=False, code_context=code_context, line_profiler_output_file=line_profiler_output_file,