Optimize the aperture photometry operation

datalab/datalab_session/analysis/source_catalog.py

@@ -1,8 +1,7 @@
-import numpy as np
 from django.contrib.auth.models import User
 
 from datalab.datalab_session.exceptions import ClientAlertException
-from datalab.datalab_session.utils.file_utils import get_hdu, scale_points
+from datalab.datalab_session.utils.file_utils import get_fits_dimensions, get_hdu, scale_points
 from datalab.datalab_session.utils.filecache import FileCache
 from datalab.datalab_session.utils.flux_to_mag import flux_to_mag
 
@@ -33,7 +32,6 @@ def source_catalog(input: dict, user: User):
     raise ClientAlertException(f"Download of {input['basename']} timed out")
 
   cat_hdu = get_hdu(file_path, 'CAT')
-  sci_hdu = get_hdu(file_path, 'SCI')
 
   DECIMALS_OF_PRECISION = 6
   MAX_SOURCE_CATALOG_SIZE = min(len(cat_hdu.data["x"]), 1000)
@@ -65,7 +63,8 @@ def source_catalog(input: dict, user: User):
     dec = None
 
   # scale the x_points and y_points from the fits pixel coords to the jpg coords
-  fits_height, fits_width = np.shape(sci_hdu.data)
+  # Shape comes from the SCI header: touching sci_hdu.data would decompress the whole image
+  fits_height, fits_width = get_fits_dimensions(file_path)
   x_points, y_points = scale_points(fits_height, fits_width, input['width'], input['height'], x_points=x_points, y_points=y_points)
   x, y = scale_points(fits_height, fits_width, input['width'], input['height'], x_points=x, y_points=y)
 

datalab/datalab_session/analysis/wcs.py

@@ -4,7 +4,7 @@
 
 
 from datalab.datalab_session.exceptions import ClientAlertException
-from datalab.datalab_session.utils.file_utils import get_hdu
+from datalab.datalab_session.utils.file_utils import get_fits_header
 from datalab.datalab_session.utils.filecache import FileCache
 
 
@@ -18,8 +18,8 @@ def wcs(input: dict, user: User):
   """
   try:
     file_path = FileCache().get_fits(input['basename'], input['source'], user)
-    sci_hdu = get_hdu(file_path, 'SCI')
-    fits_dimensions = [sci_hdu.header.get('NAXIS1'), sci_hdu.header.get('NAXIS2')]
+    sci_header = get_fits_header(file_path, 'SCI')
+    fits_dimensions = [sci_header.get('NAXIS1'), sci_header.get('NAXIS2')]
   except TimeoutError as e:
     raise ClientAlertException(f"Download of {input['basename']} FITs timed out")
   except TypeError as e:
@@ -28,7 +28,7 @@ def wcs(input: dict, user: User):
     raise ClientAlertException(f"No FITs file found for this image")
 
   try:
-    wcs = WCS(sci_hdu.header)
+    wcs = WCS(sci_header)
     wcs_solution = wcs.wcs
     wcs_cd = wcs_solution.cd
     output = {

datalab/datalab_session/data_operations/aperture_photometry.py

@@ -3,10 +3,12 @@
 
 from django.contrib.auth.models import User
 
-from datalab.datalab_session.data_operations.data_operation import BaseDataOperation
-from datalab.datalab_session.data_operations.input_data_handler import InputDataHandler
+from datalab.datalab_session.data_operations.data_operation import BaseDataOperation, ProgressStep
 from datalab.datalab_session.exceptions import ClientAlertException
+from datalab.datalab_session.utils.filecache import FileCache
+from datalab.datalab_session.utils.file_utils import temp_file_manager
 from datalab.datalab_session.utils.format import Format
+from datalab.datalab_session.utils.s3_utils import save_files_to_s3
 from datalab.datalab_session.utils.aperture_light_curve import (
     DEFAULT_ANNULUS_INNER_RADIUS,
     DEFAULT_ANNULUS_OUTER_RADIUS,
@@ -31,9 +33,13 @@ class AperturePhotometry(BaseDataOperation):
     MINIMUM_NUMBER_OF_INPUTS = 1
     MAXIMUM_NUMBER_OF_INPUTS = 999
     PROGRESS_STEPS = {
-        'INPUT_PROCESSING_PERCENTAGE_COMPLETION': 0.2,
-        'APERTURE_PHOTOMETRY_PERCENTAGE_COMPLETION': 0.9,
-        'OUTPUT_PERCENTAGE_COMPLETION': 1.0
+        'downloading': ProgressStep('Downloading input frames', 0.25),
+        'validate': ProgressStep('Validating input frames', 0.3),
+        'catalog': ProgressStep('Building comparison star catalog', 0.45),
+        'measure': ProgressStep('Measuring source and comparison stars', 0.6),
+        'select': ProgressStep('Selecting comparison stars', 0.75),
+        'render': ProgressStep('Creating diagnostic images', 0.9),
+        'save': ProgressStep('Saving output images', 1.0)
     }
 
     @staticmethod
@@ -127,28 +133,32 @@ def operate(self, submitter: User):
             annulus_outer_radius = float(self.input_data['annulus_outer_radius'])
             min_comparisons = int(self.input_data.get('min_comparisons', DEFAULT_MIN_COMPARISONS))
             max_comparisons = int(self.input_data.get('max_comparisons', DEFAULT_MAX_COMPARISONS))
-            self.set_operation_progress(AperturePhotometry.PROGRESS_STEPS['INPUT_PROCESSING_PERCENTAGE_COMPLETION'])
-            input_handlers = [
-                InputDataHandler(submitter, input_file['basename'], input_file.get('source'))
-                for input_file in input_files
-            ]
+            # Resolve inputs to local file-cache paths only. Pixel data is loaded (and released)
+            # frame by frame inside generate_light_curve, never held for all inputs at once.
+            file_cache = FileCache()
+            fits_paths = []
+            for index, input_file in enumerate(input_files, start=1):
+                fits_paths.append(file_cache.get_fits(input_file['basename'], input_file.get('source'), submitter))
+                self._report_pipeline_progress('downloading', index / len(input_files))
+
             result = generate_light_curve(
-                input_handlers=input_handlers,
+                fits_paths=fits_paths,
                 target_ra_deg=target_ra,
                 target_dec_deg=target_dec,
                 aperture_radius=aperture_radius,
                 annulus_inner_radius=annulus_inner_radius,
                 annulus_outer_radius=annulus_outer_radius,
                 min_comparisons=min_comparisons,
                 max_comparisons=max_comparisons,
+                progress_callback=self._report_pipeline_progress,
             )
         except LightCurveError as exc:
             log.warning(f"Aperture Photometry failed: {exc}")
             raise ClientAlertException(str(exc)) from exc
         except (KeyError, TypeError, ValueError) as exc:
             raise ClientAlertException(f'Operation {self.name()} received invalid input.') from exc
 
-        self.set_operation_progress(AperturePhotometry.PROGRESS_STEPS['APERTURE_PHOTOMETRY_PERCENTAGE_COMPLETION'])
+        diagnostic_image_urls = self._save_diagnostic_images_to_s3(result.diagnostic_image_jpegs_by_fits_basename)
         filter_value = input_files[0].get('filter', input_files[0].get('primary_optical_element', 'None'))
         output = {
             'output_data': [
@@ -163,16 +173,47 @@ def operate(self, submitter: User):
                         asdict(star) for star in result.selected_comparison_stars
                     ],
                     'diagnostics': result.diagnostics_by_fits_basename,
-                    'diagnostic_images': result.diagnostic_images_by_fits_basename,
+                    'diagnostic_images': diagnostic_image_urls,
                 }
             ]
         }
         self.set_output(output, is_raw=True)
-        self.set_operation_progress(AperturePhotometry.PROGRESS_STEPS['OUTPUT_PERCENTAGE_COMPLETION'])
+        self.set_operation_progress(1.0)
+        self.set_message("")
         self.set_status('COMPLETED')
         log.info(
             "Aperture Photometry output: "
             f"filter={filter_value}, light_curve_rows={len(result.light_curve_rows)}, "
             f"selected_comparison_stars={len(result.selected_comparison_stars)}, "
-            f"diagnostic_images={len(result.diagnostic_images_by_fits_basename)}"
+            f"diagnostic_images={len(diagnostic_image_urls)}"
         )
+
+    def _report_pipeline_progress(self, phase: str, fraction: float):
+        """
+            Advances this operation's overall progress and status message through the PROGRESS_STEPS
+            band identified by phase. The band runs from the previous step's progress to this step's,
+            filling as fraction goes 0 -> 1.
+        """
+        steps = list(AperturePhotometry.PROGRESS_STEPS.values())
+        band = list(AperturePhotometry.PROGRESS_STEPS).index(phase)
+        band_start = steps[band - 1].progress if band > 0 else 0.0
+        step = steps[band]
+        self.set_operation_progress(band_start + (step.progress - band_start) * fraction)
+        self.set_message(f"{step.message}: {fraction * 100:.0f}%")
+
+    def _save_diagnostic_images_to_s3(self, diagnostic_image_jpegs_by_fits_basename: dict) -> dict:
+        """
+            Uploads each frame's diagnostic overlay JPEG to the operation bucket.
+
+            Returns a dict mapping FITS basename to the presigned bucket url for its overlay.
+        """
+        diagnostic_image_urls = {}
+        total = len(diagnostic_image_jpegs_by_fits_basename)
+        for index, (fits_basename, jpeg_bytes) in enumerate(diagnostic_image_jpegs_by_fits_basename.items(), start=1):
+            with temp_file_manager(f'{self.cache_key}-{index}-diagnostic.jpg', dir=self.temp) as jpeg_path:
+                with open(jpeg_path, 'wb') as jpeg_file:
+                    jpeg_file.write(jpeg_bytes)
+                s3_output = save_files_to_s3(self.cache_key, Format.IMAGE, {'diagnostic_jpg_path': jpeg_path}, index=index)
+            diagnostic_image_urls[fits_basename] = s3_output['diagnostic_url']
+            self._report_pipeline_progress('save', index / total)
+        return diagnostic_image_urls

datalab/datalab_session/data_operations/data_operation.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from collections import namedtuple
 import hashlib
 import json
 import os
@@ -17,6 +18,10 @@
 log = logging.getLogger()
 log.setLevel(logging.INFO)
 
+# ProgressStep will be used to tie a human readable message to the amount of time the step takes for a step
+ProgressStep = namedtuple('ProgressStep', ['message', 'progress'])
+
+
 class BaseDataOperation(ABC):
 
     def __init__(self, input_data: dict = None):