Merge pull request #993 from codeflash-ai/don't-cleanup-replays

Don't cleanup replays and traces

Author: KRRT7

codeflash/optimization/optimizer.py

@@ -578,13 +578,16 @@ def cleanup_temporary_paths(self) -> None:
             get_run_tmp_file.tmpdir.cleanup()
             del get_run_tmp_file.tmpdir
 
+        # Always clean up concolic test directory
+        cleanup_paths([self.test_cfg.concolic_test_root_dir])
+
         if self.current_worktree:
             remove_worktree(self.current_worktree)
             return
 
         if self.current_function_optimizer:
             self.current_function_optimizer.cleanup_generated_files()
-        paths_to_cleanup = [self.test_cfg.concolic_test_root_dir, self.replay_tests_dir]
+        paths_to_cleanup = [self.replay_tests_dir]
         if self.trace_file:
             paths_to_cleanup.append(self.trace_file)
         if self.test_cfg.tests_root.exists():

codeflash/picklepatch/pickle_patcher.py

@@ -8,12 +8,16 @@
 
 import contextlib
 import pickle
+import warnings
 from typing import Any, ClassVar
 
 import dill
+from dill import PicklingWarning
 
 from .pickle_placeholder import PicklePlaceholder
 
+warnings.filterwarnings("ignore", category=PicklingWarning)
+
 
 class PicklePatcher:
     """A utility class for safely pickling objects with unpicklable components.

codeflash/tracing/replay_test.py

@@ -44,7 +44,10 @@ def get_function_alias(module: str, function_name: str) -> str:
 
 
 def create_trace_replay_test(trace_file: str, functions: list[FunctionModules], max_run_count: int = 100) -> str:
-    imports = """import dill as pickle
+    imports = """import warnings
+import dill as pickle
+from dill import PicklingWarning
+warnings.filterwarnings("ignore", category=PicklingWarning)
 from codeflash.tracing.replay_test import get_next_arg_and_return
 """
 

codeflash/tracing/tracing_new_process.py

@@ -12,6 +12,7 @@
 import sys
 import threading
 import time
+import warnings
 from collections import defaultdict
 from importlib.util import find_spec
 from pathlib import Path
@@ -26,6 +27,10 @@
 from codeflash.picklepatch.pickle_patcher import PicklePatcher
 from codeflash.tracing.tracing_utils import FunctionModules, filter_files_optimized, module_name_from_file_path
 
+# Suppress dill PicklingWarning
+warnings.filterwarnings("ignore", message="Cannot locate reference to")
+warnings.filterwarnings("ignore", message="Cannot pickle.*recursive self-references")
+
 if TYPE_CHECKING:
     from types import FrameType, TracebackType
 

codeflash/verification/codeflash_capture.py

@@ -7,11 +7,15 @@
 import os
 import sqlite3
 import time
+import warnings
 from enum import Enum
 from pathlib import Path
 from typing import Callable
 
 import dill as pickle
+from dill import PicklingWarning
+
+warnings.filterwarnings("ignore", category=PicklingWarning)
 
 
 class VerificationType(str, Enum):