Save before cleaning up threads

cozy/main.py

@@ -101,7 +101,12 @@ def run():
 
     start = datetime.datetime.now()
 
-    if not args.simple:
+    if args.simple:
+        if args.save:
+            with open(args.save, "wb") as f:
+                pickle.dump(ast, f)
+                print("Saved implementation to file {}".format(args.save))
+    else:
         callback = None
         server = None
 
@@ -146,16 +151,12 @@ def callback(impl):
             ast,
             timeout           = datetime.timedelta(seconds=args.timeout),
             progress_callback = callback,
-            improve_count=improve_count)
+            improve_count=improve_count,
+            dump_synthesized_in_file=args.save)
 
         if server is not None:
             server.join()
 
-    if args.save:
-        with open(args.save, "wb") as f:
-            pickle.dump(ast, f)
-            print("Saved implementation to file {}".format(args.save))
-
     print("Generating IR...")
     code = ast.code
 

cozy/synthesis/high_level_interface.py

@@ -10,6 +10,7 @@
 from typing import Callable, Any
 import sys
 import os
+import pickle
 from queue import Empty
 from multiprocessing import Value
 
@@ -99,7 +100,8 @@ def improve_implementation(
         impl              : Implementation,
         timeout           : datetime.timedelta = datetime.timedelta(seconds=60),
         progress_callback : Callable[[Implementation], Any] = None,
-        improve_count     : Value = None) -> Implementation:
+        improve_count     : Value = None,
+        dump_synthesized_in_file: str = None) -> Implementation:
     """Improve an implementation.
 
     This function tries to synthesize a better version of the given
@@ -108,6 +110,10 @@ def improve_implementation(
     If provided, progress_callback will be called whenever a better
     implementation is found.  It will be given the better implementation, which
     it should not modify or cache.
+
+    If provided, the synthesized implementation will be dumped to dump_synthesized_in_file
+    before cleaning up the running threads when the loop terminates (because of time-outs etc.).
+    This is useful when thread that invokes Z3 is not responsive to cleanup.
     """
 
     start_time = datetime.datetime.now()
@@ -241,6 +247,11 @@ def index_of(l, p):
                 else:
                     print("  (skipped; {} was aleady cleaned up)".format(q.name))
 
+        if dump_synthesized_in_file is not None:
+            with open(dump_synthesized_in_file, "wb") as f:
+                pickle.dump(impl, f)
+                print("Dumped implementation to file {}".format(dump_synthesized_in_file))
+
         # stop jobs
         print("Stopping jobs")
         stop_jobs(list(improvement_jobs))