Merge remote-tracking branch 'origin/main' into feat-save-log-line-profiler-results

Author: HeshamHM28

codeflash/cli_cmds/cmd_init.py

@@ -941,6 +941,8 @@ def run_end_to_end_test(args: Namespace, bubble_sort_path: str, bubble_sort_test
     command = ["codeflash", "--file", "bubble_sort.py", "--function", "sorter"]
     if args.no_pr:
         command.append("--no-pr")
+    if args.verbose:
+        command.append("--verbose")
 
     logger.info("Running sample optimization…")
     console.rule()
@@ -953,7 +955,7 @@ def run_end_to_end_test(args: Namespace, bubble_sort_path: str, bubble_sort_test
             if process.stdout:
                 for line in process.stdout:
                     stripped = line.strip()
-                    console.print(stripped)
+                    console.out(stripped)
                     output.append(stripped)
             process.wait()
         console.rule()

codeflash/optimization/function_optimizer.py

@@ -1389,23 +1389,11 @@ def generate_and_instrument_tests(
         ]
 
     def cleanup_generated_files(self) -> None:
-        paths_to_cleanup = (
-            [
-                test_file.instrumented_behavior_file_path
-                for test_type in [
-                TestType.GENERATED_REGRESSION,
-                TestType.EXISTING_UNIT_TEST,
-                TestType.CONCOLIC_COVERAGE_TEST,
-            ]
-                for test_file in self.test_files.get_by_type(test_type).test_files
-            ]
-            + [
-                test_file.benchmarking_file_path
-                for test_type in [TestType.GENERATED_REGRESSION, TestType.EXISTING_UNIT_TEST]
-                for test_file in self.test_files.get_by_type(test_type).test_files
-            ]
-            + [self.test_cfg.concolic_test_root_dir]
-        )
+        paths_to_cleanup = [self.test_cfg.concolic_test_root_dir]
+        for test_file in self.test_files:
+            paths_to_cleanup.append(test_file.instrumented_behavior_file_path)
+            paths_to_cleanup.append(test_file.benchmarking_file_path)
+
         cleanup_paths(paths_to_cleanup)
         if hasattr(get_run_tmp_file, "tmpdir"):
             get_run_tmp_file.tmpdir.cleanup()

codeflash/optimization/optimizer.py

@@ -45,6 +45,7 @@ def __init__(self, args: Namespace) -> None:
         self.local_aiservice_client = LocalAiServiceClient() if self.experiment_id else None
         self.replay_tests_dir = None
         self.functions_checkpoint: CodeflashRunCheckpoint | None = None
+        self.current_function_optimizer: FunctionOptimizer | None = None
 
     def create_function_optimizer(
         self,
@@ -265,7 +266,9 @@ def run(self) -> None:
                             function_to_tests,
                             validated_original_code[original_module_path].source_code,
                         )
-
+                    self.current_function_optimizer = (
+                        function_optimizer  # needed to clean up from the outside of this function
+                    )
                     best_optimization = function_optimizer.optimize_function()
                     if self.functions_checkpoint:
                         self.functions_checkpoint.add_function_to_checkpoint(
@@ -293,6 +296,9 @@ def run(self) -> None:
     def cleanup_temporary_paths(self) -> None:
         from codeflash.code_utils.code_utils import cleanup_paths
 
+        if self.current_function_optimizer:
+            self.current_function_optimizer.cleanup_generated_files()
+
         cleanup_paths([self.test_cfg.concolic_test_root_dir, self.replay_tests_dir])
 
 

codeflash/telemetry/sentry.py

@@ -12,6 +12,7 @@ def init_sentry(enabled: bool = False, exclude_errors: bool = False) -> None:  #
             if exclude_errors
             else logging.ERROR,  # Otherwise, error logs will create sentry events
         )
+
         sentry_sdk.init(
             dsn="https://4b9a1902f9361b48c04376df6483bc96@o4506833230561280.ingest.sentry.io/4506833262477312",
             integrations=[sentry_logging],
@@ -22,4 +23,5 @@ def init_sentry(enabled: bool = False, exclude_errors: bool = False) -> None:  #
             # of sampled transactions.
             # We recommend adjusting this value in production.
             profiles_sample_rate=1.0,
+            ignore_errors=[KeyboardInterrupt],
         )

codeflash/tracer.py

@@ -29,7 +29,6 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, ClassVar
 
-import dill
 import isort
 from rich.align import Align
 from rich.panel import Panel
@@ -41,6 +40,7 @@
 from codeflash.code_utils.code_utils import module_name_from_file_path
 from codeflash.code_utils.config_parser import parse_config_file
 from codeflash.discovery.functions_to_optimize import filter_files_optimized
+from codeflash.picklepatch.pickle_patcher import PicklePatcher
 from codeflash.tracing.replay_test import create_trace_replay_test
 from codeflash.tracing.tracing_utils import FunctionModules
 from codeflash.verification.verification_utils import get_test_file_path
@@ -399,22 +399,12 @@ def tracer_logic(self, frame: FrameType, event: str) -> None:  # noqa: PLR0911
                 arguments_copy = dict(arguments.items())  # Use the local 'arguments' from frame.f_locals
                 if class_name and code.co_name == "__init__" and "self" in arguments_copy:
                     del arguments_copy["self"]
-                local_vars = pickle.dumps(arguments_copy, protocol=pickle.HIGHEST_PROTOCOL)
+                local_vars = PicklePatcher.dumps(arguments_copy, protocol=pickle.HIGHEST_PROTOCOL)
                 sys.setrecursionlimit(original_recursion_limit)
             except Exception:
-                # we retry with dill if pickle fails. It's slower but more comprehensive
-                try:
-                    sys.setrecursionlimit(10000)  # Ensure limit is high for dill too
-                    # arguments_copy should be used here as well if defined above
-                    local_vars = dill.dumps(
-                        arguments_copy if "arguments_copy" in locals() else dict(arguments.items()),
-                        protocol=dill.HIGHEST_PROTOCOL,
-                    )
-                    sys.setrecursionlimit(original_recursion_limit)
-
-                except Exception:
-                    self.function_count[function_qualified_name] -= 1
-                    return
+                self.function_count[function_qualified_name] -= 1
+                sys.setrecursionlimit(original_recursion_limit)
+                return
 
             cur.execute(
                 "INSERT INTO function_calls VALUES(?, ?, ?, ?, ?, ?, ?, ?)",

codeflash/version.py

@@ -1,3 +1,3 @@
 # These version placeholders will be replaced by uv-dynamic-versioning during build.
-__version__ = "0.14.5"
-__version_tuple__ = (0, 14, 5)
+__version__ = "0.14.6"
+__version_tuple__ = (0, 14, 6)