Merge branch 'main' into trace-and-optimize

Author: KRRT7

code_to_optimize/code_directories/circular_deps/api_client.py

@@ -0,0 +1,25 @@
+from os import getenv
+
+from attrs import define, evolve
+
+from constants import DEFAULT_API_URL, DEFAULT_APP_URL
+
+
+@define
+class ApiClient():
+    api_key_header_name: str = "API-Key"
+    client_type_header_name: str = "client-type"
+    client_type_header_value: str = "sdk-python"
+
+    @staticmethod
+    def get_console_url() -> str:
+        console_url = getenv("CONSOLE_URL", DEFAULT_API_URL)
+        if DEFAULT_API_URL == console_url:
+            return DEFAULT_APP_URL
+
+        return console_url
+
+    def with_api_key(self, api_key: str) -> "ApiClient": # ---> here is the problem with circular dependency, this makes libcst thinks that ApiClient needs an import despite it's already in the same file.
+        """Get a new client matching this one with a new API key"""
+        return evolve(self, api_key=api_key)
+

code_to_optimize/code_directories/circular_deps/constants.py

@@ -0,0 +1,8 @@
+DEFAULT_API_URL = "https://api.galileo.ai/"
+DEFAULT_APP_URL = "https://app.galileo.ai/"
+
+
+# function_names: GalileoApiClient.get_console_url
+# module_abs_path : /home/mohammed/Work/galileo-python/src/galileo/api_client.py
+# preexisting_objects: {('GalileoApiClient', ()), ('_set_destination', ()), ('get_console_url', (FunctionParent(name='GalileoApiClient', type='ClassDef'),))}
+# project_root_path: /home/mohammed/Work/galileo-python/src

code_to_optimize/code_directories/circular_deps/optimized.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import urllib.parse
+from os import getenv
+
+from attrs import define
+from api_client import ApiClient
+from constants import DEFAULT_API_URL, DEFAULT_APP_URL
+
+
+@define
+class ApiClient():
+
+    @staticmethod
+    def get_console_url() -> str:
+        # Cache env lookup for speed
+        console_url = getenv("CONSOLE_URL")
+        if not console_url or console_url == DEFAULT_API_URL:
+            return DEFAULT_APP_URL
+        return console_url
+
+# Pre-parse netlocs that are checked frequently to avoid parsing repeatedly
+_DEFAULT_APP_URL_NETLOC = urllib.parse.urlparse(DEFAULT_APP_URL).netloc
+_DEFAULT_API_URL_NETLOC = urllib.parse.urlparse(DEFAULT_API_URL).netloc
+
+def get_dest_url(url: str) -> str:
+    destination = url if url else ApiClient.get_console_url()
+    # Replace only if 'console.' is at the beginning to avoid partial matches
+    if destination.startswith("console."):
+        destination = "api." + destination[len("console."):]
+    else:
+        destination = destination.replace("console.", "api.", 1)
+
+    parsed_url = urllib.parse.urlparse(destination)
+    if parsed_url.netloc == _DEFAULT_APP_URL_NETLOC or parsed_url.netloc == _DEFAULT_API_URL_NETLOC:
+        return f"{DEFAULT_APP_URL}api/traces"
+    return f"{parsed_url.scheme}://{parsed_url.netloc}/traces"

code_to_optimize/code_directories/circular_deps/pyproject.toml

@@ -0,0 +1,7 @@
+[tool.codeflash]
+# All paths are relative to this pyproject.toml's directory.
+module-root = "."
+tests-root = "tests"
+test-framework = "pytest"
+ignore-paths = []
+formatter-cmds = ["black $file"]

codeflash/code_utils/code_extractor.py

@@ -331,7 +331,9 @@ def add_needed_imports_from_module(
             RemoveImportsVisitor.remove_unused_import(dst_context, mod)
         for mod, obj_seq in gatherer.object_mapping.items():
             for obj in obj_seq:
-                if f"{mod}.{obj}" in helper_functions_fqn:
+                if (
+                    f"{mod}.{obj}" in helper_functions_fqn or dst_context.full_module_name == mod  # avoid circular deps
+                ):
                     continue  # Skip adding imports for helper functions already in the context
                 AddImportsVisitor.add_needed_import(dst_context, mod, obj)
                 RemoveImportsVisitor.remove_unused_import(dst_context, mod, obj)

codeflash/code_utils/code_replacer.py

@@ -415,12 +415,16 @@ def replace_function_definitions_in_module(
 ) -> bool:
     source_code: str = module_abspath.read_text(encoding="utf8")
     new_code: str = replace_functions_and_add_imports(
-        source_code, function_names, optimized_code, module_abspath, preexisting_objects, project_root_path
+        add_global_assignments(optimized_code, source_code),
+        function_names,
+        optimized_code,
+        module_abspath,
+        preexisting_objects,
+        project_root_path,
     )
     if is_zero_diff(source_code, new_code):
         return False
-    code_with_global_assignments = add_global_assignments(optimized_code, new_code)
-    module_abspath.write_text(code_with_global_assignments, encoding="utf8")
+    module_abspath.write_text(new_code, encoding="utf8")
     return True
 
 

codeflash/lsp/beta.py

@@ -36,27 +36,71 @@ def get_optimizable_functions(
     server: CodeflashLanguageServer, params: OptimizableFunctionsParams
 ) -> dict[str, list[str]]:
     file_path = Path(uris.to_fs_path(params.textDocument.uri))
-    server.optimizer.args.file = file_path
-    server.optimizer.args.previous_checkpoint_functions = False
-    optimizable_funcs, _ = server.optimizer.get_optimizable_functions()
-    path_to_qualified_names = {}
-    for path, functions in optimizable_funcs.items():
-        path_to_qualified_names[path.as_posix()] = [func.qualified_name for func in functions]
-    return path_to_qualified_names
+    server.show_message_log(f"Getting optimizable functions for: {file_path}", "Info")
+
+    # Save original args to restore later
+    original_file = getattr(server.optimizer.args, "file", None)
+    original_function = getattr(server.optimizer.args, "function", None)
+    original_checkpoint = getattr(server.optimizer.args, "previous_checkpoint_functions", None)
+
+    server.show_message_log(f"Original args - file: {original_file}, function: {original_function}", "Info")
+
+    try:
+        # Set temporary args for this request only
+        server.optimizer.args.file = file_path
+        server.optimizer.args.function = None  # Always get ALL functions, not just one
+        server.optimizer.args.previous_checkpoint_functions = False
+
+        server.show_message_log("Calling get_optimizable_functions...", "Info")
+        optimizable_funcs, _ = server.optimizer.get_optimizable_functions()
+
+        path_to_qualified_names = {}
+        for path, functions in optimizable_funcs.items():
+            path_to_qualified_names[path.as_posix()] = [func.qualified_name for func in functions]
+
+        server.show_message_log(
+            f"Found {len(path_to_qualified_names)} files with functions: {path_to_qualified_names}", "Info"
+        )
+        return path_to_qualified_names
+    finally:
+        # Restore original args to prevent state corruption
+        if original_file is not None:
+            server.optimizer.args.file = original_file
+        if original_function is not None:
+            server.optimizer.args.function = original_function
+        else:
+            server.optimizer.args.function = None
+        if original_checkpoint is not None:
+            server.optimizer.args.previous_checkpoint_functions = original_checkpoint
+
+        server.show_message_log(
+            f"Restored args - file: {server.optimizer.args.file}, function: {server.optimizer.args.function}", "Info"
+        )
 
 
 @server.feature("initializeFunctionOptimization")
 def initialize_function_optimization(
     server: CodeflashLanguageServer, params: FunctionOptimizationParams
 ) -> dict[str, str]:
     file_path = Path(uris.to_fs_path(params.textDocument.uri))
+    server.show_message_log(f"Initializing optimization for function: {params.functionName} in {file_path}", "Info")
+
+    # IMPORTANT: Store the specific function for optimization, but don't corrupt global state
     server.optimizer.args.function = params.functionName
     server.optimizer.args.file = file_path
+
+    server.show_message_log(
+        f"Args set - function: {server.optimizer.args.function}, file: {server.optimizer.args.file}", "Info"
+    )
+
     optimizable_funcs, _ = server.optimizer.get_optimizable_functions()
     if not optimizable_funcs:
+        server.show_message_log(f"No optimizable functions found for {params.functionName}", "Warning")
         return {"functionName": params.functionName, "status": "not found", "args": None}
+
     fto = optimizable_funcs.popitem()[1][0]
     server.optimizer.current_function_being_optimized = fto
+    server.show_message_log(f"Successfully initialized optimization for {params.functionName}", "Info")
     return {"functionName": params.functionName, "status": "success"}
 
 
@@ -136,11 +180,20 @@ def generate_tests(server: CodeflashLanguageServer, params: FunctionOptimization
 
 
 @server.feature("performFunctionOptimization")
-def perform_function_optimization(
+def perform_function_optimization(  # noqa: PLR0911
     server: CodeflashLanguageServer, params: FunctionOptimizationParams
 ) -> dict[str, str]:
+    server.show_message_log(f"Starting optimization for function: {params.functionName}", "Info")
     current_function = server.optimizer.current_function_being_optimized
 
+    if not current_function:
+        server.show_message_log(f"No current function being optimized for {params.functionName}", "Error")
+        return {
+            "functionName": params.functionName,
+            "status": "error",
+            "message": "No function currently being optimized",
+        }
+
     module_prep_result = server.optimizer.prepare_module_for_optimization(current_function.file_path)
 
     validated_original_code, original_module_ast = module_prep_result
@@ -214,19 +267,29 @@ def perform_function_optimization(
     )
 
     if not best_optimization:
+        server.show_message_log(
+            f"No best optimizations found for function {function_to_optimize_qualified_name}", "Warning"
+        )
         return {
             "functionName": params.functionName,
             "status": "error",
             "message": f"No best optimizations found for function {function_to_optimize_qualified_name}",
         }
 
     optimized_source = best_optimization.candidate.source_code
+    speedup = original_code_baseline.runtime / best_optimization.runtime
+
+    server.show_message_log(f"Optimization completed for {params.functionName} with {speedup:.2f}x speedup", "Info")
+
+    # CRITICAL: Clear the function filter after optimization to prevent state corruption
+    server.optimizer.args.function = None
+    server.show_message_log("Cleared function filter to prevent state corruption", "Info")
 
     return {
         "functionName": params.functionName,
         "status": "success",
         "message": "Optimization completed successfully",
-        "extra": f"Speedup: {original_code_baseline.runtime / best_optimization.runtime:.2f}x faster",
+        "extra": f"Speedup: {speedup:.2f}x faster",
         "optimization": optimized_source,
     }
 

codeflash/optimization/optimizer.py

@@ -12,6 +12,7 @@
 from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient
 from codeflash.cli_cmds.console import console, logger, progress_bar
 from codeflash.code_utils import env_utils
+from codeflash.code_utils.code_utils import cleanup_paths
 from codeflash.code_utils.env_utils import get_pr_number
 from codeflash.either import is_successful
 from codeflash.models.models import ValidCode
@@ -248,10 +249,10 @@ def run(self) -> None:
             return
         if not env_utils.check_formatter_installed(self.args.formatter_cmds):
             return
-
         if self.args.no_draft and is_pr_draft():
             logger.warning("PR is in draft mode, skipping optimization")
             return
+        cleanup_paths(Optimizer.find_leftover_instrumented_test_files(self.test_cfg.tests_root))
 
         function_optimizer = None
         file_to_funcs_to_optimize, num_optimizable_functions = self.get_optimizable_functions()
@@ -326,9 +327,27 @@ def run(self) -> None:
 
             self.cleanup_temporary_paths()
 
-    def cleanup_temporary_paths(self) -> None:
-        from codeflash.code_utils.code_utils import cleanup_paths
+    @staticmethod
+    def find_leftover_instrumented_test_files(test_root: Path) -> list[Path]:
+        """Search for all paths within the test_root that match the following patterns.
 
+        - 'test.*__perf_test_{0,1}.py'
+        - 'test_.*__unit_test_{0,1}.py'
+        - 'test_.*__perfinstrumented.py'
+        - 'test_.*__perfonlyinstrumented.py'
+        Returns a list of matching file paths.
+        """
+        import re
+
+        pattern = re.compile(
+            r"(?:test.*__perf_test_\d?\.py|test_.*__unit_test_\d?\.py|test_.*__perfinstrumented\.py|test_.*__perfonlyinstrumented\.py)$"
+        )
+
+        return [
+            file_path for file_path in test_root.rglob("*") if file_path.is_file() and pattern.match(file_path.name)
+        ]
+
+    def cleanup_temporary_paths(self) -> None:
         if self.current_function_optimizer:
             self.current_function_optimizer.cleanup_generated_files()
 

tests/test_code_context_extractor.py

@@ -11,6 +11,8 @@
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import FunctionParent
 from codeflash.optimization.optimizer import Optimizer
+from codeflash.code_utils.code_replacer import replace_functions_and_add_imports
+from codeflash.code_utils.code_extractor import add_global_assignments
 
 
 class HelperClass:
@@ -2434,3 +2436,22 @@ def simple_method(self):
         assert "class SimpleClass:" in code_content
         assert "def simple_method(self):" in code_content
         assert "return 42" in code_content
+
+
+
+def test_replace_functions_and_add_imports():
+    path_to_root = Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "circular_deps"
+    file_abs_path = path_to_root / "api_client.py"
+    optimized_code = Path(path_to_root / "optimized.py").read_text(encoding="utf-8")
+    content = Path(file_abs_path).read_text(encoding="utf-8")
+    new_code = replace_functions_and_add_imports(
+        source_code= add_global_assignments(optimized_code, content),
+        function_names= ["ApiClient.get_console_url"],
+        optimized_code= optimized_code,
+        module_abspath= Path(file_abs_path),
+        preexisting_objects= {('ApiClient', ()), ('get_console_url', (FunctionParent(name='ApiClient', type='ClassDef'),))},
+        project_root_path= Path(path_to_root),
+    )
+    assert "import ApiClient" not in new_code, "Error: Circular dependency found"
+    
+    assert "import urllib.parse" in new_code, "Make sure imports for optimization global assignments exist" 

tests/test_code_replacement.py

@@ -1693,8 +1693,8 @@ def new_function2(value):
 print("Hello world")
     """
     expected_code = """import numpy as np
-print("Hello world")
 
+print("Hello world")
 a=2
 print("Hello world")
 def some_fn():
@@ -1712,8 +1712,7 @@ def __init__(self, name):
     def __call__(self, value):
         return "I am still old"
     def new_function2(value):
-        return cst.ensure_type(value, str)
-"""
+        return cst.ensure_type(value, str)"""
     code_path = (Path(__file__).parent.resolve() / "../code_to_optimize/global_var_original.py").resolve()
     code_path.write_text(original_code, encoding="utf-8")
     tests_root = Path("/Users/codeflash/Downloads/codeflash-dev/codeflash/code_to_optimize/tests/pytest/")
@@ -1769,8 +1768,8 @@ def new_function2(value):
 print("Hello world")
         """
     expected_code = """import numpy as np
-print("Hello world")
 
+print("Hello world")
 print("Hello world")
 def some_fn():
     a=np.zeros(10)
@@ -1846,8 +1845,8 @@ def new_function2(value):
 print("Hello world")
     """
     expected_code = """import numpy as np
-print("Hello world")
 
+print("Hello world")
 a=3
 print("Hello world")
 def some_fn():
@@ -1922,8 +1921,8 @@ def new_function2(value):
 print("Hello world")
     """
     expected_code = """import numpy as np
-print("Hello world")
 
+print("Hello world")
 a=2
 print("Hello world")
 def some_fn():
@@ -1999,8 +1998,8 @@ def new_function2(value):
 print("Hello world")
     """
     expected_code = """import numpy as np
-print("Hello world")
 
+print("Hello world")
 a=3
 print("Hello world")
 def some_fn():
@@ -2082,8 +2081,8 @@ def new_function2(value):
 print("Hello world")
 """
     expected_code = """import numpy as np
-print("Hello world")
 
+print("Hello world")
 if 2<3:
     a=4
 else: