cleanup properly

add verbose mode to --verify-setup

Author: misrasaurabh1

codeflash/cli_cmds/cmd_init.py

@@ -34,7 +34,7 @@
     from argparse import Namespace
 
 CODEFLASH_LOGO: str = (
-    f"{LF}"  # noqa: ISC003
+    f"{LF}"
     r"                   _          ___  _               _     " + f"{LF}"
     r"                  | |        / __)| |             | |    " + f"{LF}"
     r"  ____   ___    _ | |  ____ | |__ | |  ____   ___ | | _  " + f"{LF}"
@@ -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

@@ -908,7 +908,7 @@ def generate_tests_and_optimizations(
             function_to_concolic_tests, concolic_test_str = future_concolic_tests.result()
             logger.info(f"Generated {len(tests)} tests for {self.function_to_optimize.function_name}")
             console.rule()
-            generated_tests = GeneratedTestsList(generated_tests=tests)
+            generated_tests = GeneratedTestsList(generated_tests=[])
 
         return Success(
             (
@@ -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])