rebaseline_server: add "prefetch" directive that just warms the cache without awaiting results

gm/rebaseline_server/compare_rendered_pictures.py

 #!/usr/bin/python
 
 """
 Copyright 2014 Google Inc.
 
 Use of this source code is governed by a BSD-style license that can be
 found in the LICENSE file.
 
 Compare results of two render_pictures runs.
 
@@ skipping 46 matching lines @@
 
 
 class RenderedPicturesComparisons(results.BaseComparisons):
   """Loads results from multiple render_pictures runs into an ImagePairSet.
   """
 
   def __init__(self, setA_dirs, setB_dirs, image_diff_db,
                image_base_gs_url=DEFAULT_IMAGE_BASE_GS_URL,
                diff_base_url=None, setA_label='setA',
                setB_label='setB', gs=None,
-               truncate_results=False):
-    """
+               truncate_results=False, prefetch_only=False,
+               download_all_images=False):
+    """Constructor: downloads images and generates diffs.
+
+    Once the object has been created (which may take a while), you can call its
+    get_packaged_results_of_type() method to quickly retrieve the results...
+    unless you have set prefetch_only to True, in which case we will
+    asynchronously warm up the ImageDiffDB cache but not fill in self._results.
+
     Args:
       setA_dirs: list of root directories to copy all JSON summaries from,
           and to use as setA within the comparisons
       setB_dirs: list of root directories to copy all JSON summaries from,
           and to use as setB within the comparisons
       image_diff_db: ImageDiffDB instance
       image_base_gs_url: "gs://" URL pointing at the Google Storage bucket/dir
           under which all render_pictures result images can
           be found; this will be used to read images for comparison within
           this code, and included in the ImagePairSet (as an HTTP URL) so its
           consumers know where to download the images from
       diff_base_url: base URL within which the client should look for diff
           images; if not specified, defaults to a "file:///" URL representation
           of image_diff_db's storage_root
       setA_label: description to use for results in setA
       setB_label: description to use for results in setB
       gs: instance of GSUtils object we can use to download summary files
       truncate_results: FOR MANUAL TESTING: if True, truncate the set of images
           we process, to speed up testing.
+      prefetch_only: if True, return the new object as quickly as possible
+          with empty self._results (just queue up all the files to process,
+          don't wait around for them to be processed and recorded); otherwise,
+          block until the results have been assembled and recorded in
+          self._results.
+      download_all_images: if True, download all images, even if we don't
+          need them to generate diffs.  This will take much longer to complete,
+          but is useful for warming up the bitmap cache on local disk.
     """
     super(RenderedPicturesComparisons, self).__init__()
     self._image_diff_db = image_diff_db
     self._image_base_gs_url = image_base_gs_url
     self._diff_base_url = (
         diff_base_url or
         url_utils.create_filepath_url(image_diff_db.storage_root))
     self._setA_label = setA_label
     self._setB_label = setB_label
     self._gs = gs
     self.truncate_results = truncate_results
+    self._prefetch_only = prefetch_only
+    self._download_all_images = download_all_images
 
     tempdir = tempfile.mkdtemp()
     try:
       setA_root = os.path.join(tempdir, 'setA')
       setB_root = os.path.join(tempdir, 'setB')
       for source_dir in setA_dirs:
         self._copy_dir_contents(source_dir=source_dir, dest_dir=setA_root)
       for source_dir in setB_dirs:
         self._copy_dir_contents(source_dir=source_dir, dest_dir=setB_root)
 
       time_start = int(time.time())
       # TODO(epoger): For now, this assumes that we are always comparing two
       # sets of actual results, not actuals vs expectations.  Allow the user
       # to control this.
       self._results = self._load_result_pairs(
           setA_root=setA_root, setA_section=gm_json.JSONKEY_ACTUALRESULTS,
           setB_root=setB_root, setB_section=gm_json.JSONKEY_ACTUALRESULTS)
-      self._timestamp = int(time.time())
-      logging.info('Number of download file collisions: %s' %
-                   imagediffdb.global_file_collisions)
-      logging.info('Results complete; took %d seconds.' %
-                   (self._timestamp - time_start))
+      if self._results:
+        self._timestamp = int(time.time())
+        logging.info('Number of download file collisions: %s' %
+                     imagediffdb.global_file_collisions)
+        logging.info('Results complete; took %d seconds.' %
+                     (self._timestamp - time_start))
     finally:
       shutil.rmtree(tempdir)
 
   def _load_result_pairs(self, setA_root, setA_section, setB_root,
                          setB_section):
     """Loads all JSON image summaries from 2 directory trees and compares them.
 
     Args:
       setA_root: root directory containing JSON summaries of rendering results
       setA_section: which section (gm_json.JSONKEY_ACTUALRESULTS or
           gm_json.JSONKEY_EXPECTEDRESULTS) to load from the summaries in setA
       setB_root: root directory containing JSON summaries of rendering results
       setB_section: which section (gm_json.JSONKEY_ACTUALRESULTS or
           gm_json.JSONKEY_EXPECTEDRESULTS) to load from the summaries in setB
 
-    Returns the summary of all image diff results.
+    Returns the summary of all image diff results (or None, depending on
+    self._prefetch_only).
     """
     logging.info('Reading JSON image summaries from dirs %s and %s...' % (
         setA_root, setB_root))
     setA_dicts = self._read_dicts_from_root(setA_root)
     setB_dicts = self._read_dicts_from_root(setB_root)
     logging.info('Comparing summary dicts...')
 
     all_image_pairs = imagepairset.ImagePairSet(
         descriptions=(self._setA_label, self._setB_label),
         diff_base_url=self._diff_base_url)
@@ skipping 20 matching lines @@
         column_id=COLUMN__RESULT_TYPE, values=[
             results.KEY__RESULT_TYPE__FAILED,
             results.KEY__RESULT_TYPE__NOCOMPARISON,
         ])
 
     union_dict_paths = sorted(set(setA_dicts.keys() + setB_dicts.keys()))
     num_union_dict_paths = len(union_dict_paths)
     dict_num = 0
     for dict_path in union_dict_paths:
       dict_num += 1
-      logging.info('Generating pixel diffs for dict #%d of %d, "%s"...' %
-                   (dict_num, num_union_dict_paths, dict_path))
+      logging.info(
+          'Asynchronously requesting pixel diffs for dict #%d of %d, "%s"...' %
+          (dict_num, num_union_dict_paths, dict_path))
 
       dictA = self.get_default(setA_dicts, None, dict_path)
       self._validate_dict_version(dictA)
       dictA_results = self.get_default(dictA, {}, setA_section)
 
       dictB = self.get_default(setB_dicts, None, dict_path)
       self._validate_dict_version(dictB)
       dictB_results = self.get_default(dictB, {}, setB_section)
 
       skp_names = sorted(set(dictA_results.keys() + dictB_results.keys()))
@@ skipping 31 matching lines @@
                 source_skp_name=skp_name, tilenum=tile_num))
 
         for one_imagepair in imagepairs_for_this_skp:
           if one_imagepair:
             all_image_pairs.add_image_pair(one_imagepair)
             result_type = one_imagepair.extra_columns_dict\
                 [COLUMN__RESULT_TYPE]
             if result_type != results.KEY__RESULT_TYPE__SUCCEEDED:
               failing_image_pairs.add_image_pair(one_imagepair)
 
-    return {
-      results.KEY__HEADER__RESULTS_ALL: all_image_pairs.as_dict(
-          column_ids_in_order=ORDERED_COLUMN_IDS),
-      results.KEY__HEADER__RESULTS_FAILURES: failing_image_pairs.as_dict(
-          column_ids_in_order=ORDERED_COLUMN_IDS),
-    }
+    if self._prefetch_only:
+      return None
+    else:
+      return {
+          results.KEY__HEADER__RESULTS_ALL: all_image_pairs.as_dict(
+              column_ids_in_order=ORDERED_COLUMN_IDS),
+          results.KEY__HEADER__RESULTS_FAILURES: failing_image_pairs.as_dict(
+              column_ids_in_order=ORDERED_COLUMN_IDS),
+      }
 
   def _validate_dict_version(self, result_dict):
     """Raises Exception if the dict is not the type/version we know how to read.
 
     Args:
       result_dict: dictionary holding output of render_pictures; if None,
           this method will return without raising an Exception
     """
     expected_header_type = 'ChecksummedImages'
     expected_header_revision = 1
@@ skipping 59 matching lines @@
     else:
       extra_columns_dict[COLUMN__TILED_OR_WHOLE] = 'tiled'
       extra_columns_dict[COLUMN__TILENUM] = str(tilenum)
 
     try:
       return imagepair.ImagePair(
           image_diff_db=self._image_diff_db,
           base_url=self._image_base_gs_url,
           imageA_relative_url=imageA_relative_url,
           imageB_relative_url=imageB_relative_url,
-          extra_columns=extra_columns_dict)
+          extra_columns=extra_columns_dict,
+          download_all_images=self._download_all_images)
     except (KeyError, TypeError):
       logging.exception(
           'got exception while creating ImagePair for'
           ' urlPair=("%s","%s"), source_skp_name="%s", tilenum="%s"' % (
               imageA_relative_url, imageB_relative_url, source_skp_name,
               tilenum))
       return None
 
   def _copy_dir_contents(self, source_dir, dest_dir):
     """Copy all contents of source_dir into dest_dir, recursing into subdirs.
 
     Args:
       source_dir: path to source dir (GS URL or local filepath)
       dest_dir: path to destination dir (local filepath)
 
     The copy operates as a "merge with overwrite": any files in source_dir will
     be "overlaid" on top of the existing content in dest_dir.  Existing files
     with the same names will be overwritten.
     """
     if gs_utils.GSUtils.is_gs_url(source_dir):
       (bucket, path) = gs_utils.GSUtils.split_gs_url(source_dir)
       self._gs.download_dir_contents(source_bucket=bucket, source_dir=path,
                                      dest_dir=dest_dir)
     else:
       shutil.copytree(source_dir, dest_dir)