Merge pull request #445 from LSSTDESC/u/jchiang/checkpoint_fixes

Checkpoint tweaks

Author: jchiang87

config/imsim-config.yaml

@@ -126,11 +126,15 @@ image:
                 # a small numer of photons.
 
     # The objects are processed in batches.
-    # If checkpointing is turned on, then the checkpoint file will be updated after each batch.
+    # If checkpointing is turned on, then the checkpoint file will be updated after every
+    # group of nbatch_per_checkpoint batches.  This factor helps lower the overall checkpoint
+    # rate when running thousands of processes concurrently; otherwise, the contention
+    # from lots of concurrent checkpoint writes can overwhelm the file system.
     # Even if not checkpointing, using batches helps keep the memory down, since a bunch of
     # temporary data can be purged after each batch.
     # The default number of batches is 100, but you can change it if desired.
     nbatch: 100
+    nbatch_per_checkpoint: 1
 
     det_name: "$det_name"
 

imsim/lsst_image.py

@@ -45,7 +45,8 @@ def setup(self, config, base, image_num, obj_num, ignore, logger):
         req = { 'det_name': str }
         opt = { 'size': int , 'xsize': int , 'ysize': int, 'dtype': None,
                  'apply_sky_gradient': bool, 'apply_fringing': bool,
-                 'boresight': galsim.CelestialCoord, 'camera': str, 'nbatch': int}
+                 'boresight': galsim.CelestialCoord, 'camera': str, 'nbatch': int,
+                 'nbatch_per_checkpoint': int}
         params = GetAllParams(config, base, req=req, opt=opt, ignore=ignore+extra_ignore)[0]
 
         # Let the user override the image size
@@ -86,6 +87,7 @@ def setup(self, config, base, image_num, obj_num, ignore, logger):
         try:
             self.checkpoint = galsim.config.GetInputObj('checkpoint', config, base, 'LSST_Image')
             self.nbatch = params.get('nbatch', 100)
+            self.nbatch_per_checkpoint = params.get('nbatch_per_checkpoint', 1)
         except galsim.config.GalSimConfigError:
             self.checkpoint = None
             # Batching is also useful for memory reasons, to limit the number of stamps held
@@ -149,7 +151,8 @@ def buildImage(self, config, base, image_num, obj_num, logger):
             chk_name = 'buildImage_%s'%(self.det_name)
             saved = self.checkpoint.load(chk_name)
             if saved is not None:
-                full_image, all_bounds, all_vars, start_num, extra_builder = saved
+                full_image, all_bounds, all_vars, delta_start_num, extra_builder = saved
+                start_num = delta_start_num + base.get('start_obj_num', 0)
                 if extra_builder is not None:
                     base['extra_builder'] = extra_builder
                 all_stamps = [galsim._Image(np.array([]), b, full_image.wcs) for b in all_bounds]
@@ -223,13 +226,15 @@ def buildImage(self, config, base, image_num, obj_num, logger):
                 # Don't save the full stamps.  All we need for FlattenNoiseVariance is the bounds.
                 # Everything else about the stamps has already been handled above.
                 all_bounds = [stamp.bounds for stamp in all_stamps]
-                data = (full_image, all_bounds, all_vars, end_obj_num,
+                delta_end_obj_num = end_obj_num - base.get('start_obj_num', 0)
+                data = (full_image, all_bounds, all_vars, delta_end_obj_num,
                         base.get('extra_builder',None))
-                self.checkpoint.save(chk_name, data)
-                logger.warning('File %d: Completed batch %d with objects [%d, %d), and wrote '
-                               'checkpoint data to %s',
-                               base.get('file_num', 0), batch+1, start_obj_num, end_obj_num,
-                               self.checkpoint.file_name)
+                logger.warning('File %d: Completed batch %d with objects [%d, %d)',
+                               base.get('file_num', 0), batch+1, start_obj_num, end_obj_num)
+                if (batch % self.nbatch_per_checkpoint == 0
+                    or batch + 1 == nbatch):
+                    self.checkpoint.save(chk_name, data)
+                    logger.warning('Wrote checkpoint data to %s', self.checkpoint.file_name)
 
         # Bring the image so far up to a flat noise variance
         current_var = galsim.config.FlattenNoiseVariance(

tests/test_checkpoint.py

@@ -1,6 +1,7 @@
 import os
 from pathlib import Path
 import numpy as np
+import hashlib
 import logging
 import galsim
 import time
@@ -213,6 +214,62 @@ def test_checkpoint_image():
     np.testing.assert_array_equal(centroid5, centroid2)
 
 
+def test_nbatch_per_checkpoint():
+    """Test that the final checkpoint files written with two different values
+    of nbatch_per_checkpoint, both not factors of nbatch, produce the same
+    checkpoint output."""
+    wcs = galsim.PixelScale(0.2)
+    config = {
+        'input': {
+            'checkpoint': {
+                'dir': 'output',
+                'file_name': '$"checkpoint_%d.hdf"%image_num',
+            },
+        },
+        'gal': {
+            'type': 'Exponential',
+            'half_light_radius': {'type': 'Random', 'min': 0.3, 'max': 1.2},
+            'ellip': {
+                'type': 'EBeta',
+                'e': {'type': 'Random', 'min': 0.0, 'max': 0.3},
+                'beta': {'type': 'Random'},
+            }
+        },
+        'image': {
+            'type': 'LSST_Image',
+            'det_name': 'R22_S11',
+            'xsize': 2048,
+            'ysize': 2048,
+            'wcs': wcs,
+            'random_seed': 12345,
+            'nobjects': 500,
+            'nbatch': 100,
+        },
+    }
+
+    checkpoint_0 = 'output/checkpoint_0.hdf'
+    if os.path.exists(checkpoint_0):
+        os.remove(checkpoint_0)
+    config['image_num'] = 0
+    config['image']['nbatch_per_checkpoint'] = 11
+    galsim.config.ProcessInput(config)
+    image0 = galsim.config.BuildImage(config)
+    with open(checkpoint_0, 'rb') as fobj:
+        md5_0 = hashlib.md5(fobj.read()).hexdigest()
+
+    checkpoint_1 = 'output/checkpoint_1.hdf'
+    if os.path.exists(checkpoint_1):
+        os.remove(checkpoint_1)
+    config['image_num'] = 1
+    config['image']['nbatch_per_checkpoint'] = 13
+    galsim.config.ProcessInput(config)
+    image1 = galsim.config.BuildImage(config)
+    with open(checkpoint_1, 'rb') as fobj:
+        md5_1 = hashlib.md5(fobj.read()).hexdigest()
+
+    assert md5_0 == md5_1
+
+
 def test_checkpoint_flatten():
     """Test the flatten() step of buildImage when using checkpointing
     """