Merge pull request #375 from codeflash-ai/cleanup-properly

cleanup properly

Author: misrasaurabh1

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

@@ -1321,23 +1321,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],
         )