Merge pull request #225 from LSSTDESC/u/jchiang/add_option_to_move_ch…

…eckpoint_files

add option to specify archive directory for checkpoint files

bin.src/imsim.py

@@ -50,6 +50,10 @@
                     "This will be prepended to any existing IMSIM_IMAGE_PATH "
                     "environment variable, for which $CWD is included by "
                     "default.")
+parser.add_argument('--ckpt_archive_dir', type=str, default=None,
+                    help="Archive directory for checkpoint files. "
+                    "If None, then delete them (if the checkpointing.cleanup "
+                    "configuration is True).")
 
 args = parser.parse_args()
 
@@ -89,6 +93,7 @@
                                     apply_sensor_model=apply_sensor_model,
                                     create_centroid_file=args.create_centroid_file,
                                     file_id=args.file_id,
-                                    log_level=args.log_level)
+                                    log_level=args.log_level,
+                                    ckpt_archive_dir=args.ckpt_archive_dir)
 
     image_simulator.run(processes=args.processes)

python/desc/imsim/ImageSimulator.py

@@ -8,6 +8,7 @@
 import warnings
 import gzip
 import shutil
+import tempfile
 import sqlite3
 import numpy as np
 from astropy._erfa import ErfaWarning
@@ -49,7 +50,8 @@ class ImageSimulator:
     """
     def __init__(self, instcat, psf, numRows=None, config=None, seed=267,
                  outdir='fits', sensor_list=None, apply_sensor_model=True,
-                 create_centroid_file=False, file_id=None, log_level='WARN'):
+                 create_centroid_file=False, file_id=None, log_level='WARN',
+                 ckpt_archive_dir=None):
         """
         Parameters
         ----------
@@ -80,6 +82,12 @@ def __init__(self, instcat, psf, numRows=None, config=None, seed=267,
             If None, then no checkpoint file will be used
         log_level: str ['WARN']
             Logging level ('DEBUG', 'INFO', 'WARN', 'ERROR', 'CRITICAL').
+        ckpt_archive_dir: str [None]
+            If this is not None, then after a sensor-visit has written
+            its output FITS file(s), the associated checkpoint file will
+            be moved to this directory instead of deleted.  If set to
+            `None`, then delete the checkpoint file (assuming the
+            `checkpointing.cleanup` config parameter is True).
         """
         self.config = read_config(config)
         self.log_level = log_level
@@ -103,6 +111,10 @@ def __init__(self, instcat, psf, numRows=None, config=None, seed=267,
         self._make_gs_interpreters(seed, sensor_list, file_id)
         self.log_level = log_level
         self.logger = get_logger(self.log_level, name='ImageSimulator')
+        self.ckpt_archive_dir = ckpt_archive_dir
+        if (self.ckpt_archive_dir is not None and
+            not os.path.isdir(self.ckpt_archive_dir)):
+            os.makedirs(self.ckpt_archive_dir)
 
     def _gather_checkpoint_files(self, sensor_list, file_id=None):
         """
@@ -423,13 +435,23 @@ def __call__(self, gs_objects):
         # Write out the centroid files if they were made.
         gs_interpreter.write_centroid_files()
 
-        # The image for the sensor-visit has been drawn, so delete any
-        # existing checkpoint file if the config says to do so.
+        # The image for the sensor-visit has been drawn, so delete or
+        # move to the archive area any existing checkpoint file if the
+        # config says perform the cleanup.
         if (gs_interpreter.checkpoint_file is not None
                 and os.path.isfile(gs_interpreter.checkpoint_file)
                 and IMAGE_SIMULATOR.config['checkpointing']['cleanup']):
-            os.remove(gs_interpreter.checkpoint_file)
-
+            if IMAGE_SIMULATOR.ckpt_archive_dir is None:
+                os.remove(gs_interpreter.checkpoint_file)
+            else:
+                src_file = gs_interpreter.checkpoint_file
+                dest_dir = os.path.abspath(IMAGE_SIMULATOR.ckpt_archive_dir)
+                dest_file = os.path.join(dest_dir, os.path.basename(src_file))
+                with tempfile.NamedTemporaryFile(mode='wb', delete=False,
+                                                 dir=dest_dir) as tmp:
+                    shutil.copy(src_file, tmp.name)
+                    os.rename(tmp.name, dest_file)
+                os.remove(src_file)
         # Remove reference to gs_interpreter in order to recover the
         # memory associated with that object.
         IMAGE_SIMULATOR.gs_interpreters[self.sensor_name] = None