Merge pull request #7 from lsst-ts/tickets/SITCOM-1761

* Adding butler 'get fits image', tests, and config updates

python/lsst/rubintv/analysis/service/commands/butler.py

@@ -21,6 +21,8 @@
 
 from __future__ import annotations
 
+import base64
+import logging
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
@@ -31,6 +33,8 @@
 
     from ..data import DataCenter
 
+logger = logging.getLogger("lsst.rubintv.analysis.service.commands.butler")
+
 
 def get_camera(instrument_name: str) -> Camera:
     """Load a camera based on the instrument name
@@ -124,3 +128,63 @@ def build_contents(self, data_center: DataCenter) -> dict:
         return {
             "image": image.array,
         }
+
+
+@dataclass(kw_only=True)
+class GetFitsImageCommand(BaseCommand):
+    """Load an image from the Butler.
+
+    Attributes
+    ----------
+    collection : str
+        The name of the collection to load the image from.
+    image_name : str
+        The dataset_type of the image to get.
+    data_id : dict
+        The data ID of the image. Depending on the type of image this could
+        include things like "band" or "visit" or "detector".
+    compress : bool
+        Compress the image before sending.
+    """
+
+    repo: str
+    image_name: str
+    collection: dict
+    data_id: dict
+    compress: bool
+    response_type: str = "fits"
+
+    def build_contents(self, data_center: DataCenter) -> dict:
+        from lsst.afw.fits import MemFileManager
+
+        # Load the image from the Butler
+        assert data_center.butlers is not None
+        logger.info("Querying butler...")
+        exposure = data_center.butlers[self.repo].get(
+            self.image_name, collections=[self.collection], **self.data_id
+        )
+        logger.info("Received exposure.")
+
+        manager = MemFileManager()
+        opts = dict()
+        if self.compress:
+            from lsst.afw.fits import ImageCompressionOptions, ImageScalingOptions, ImageWriteOptions
+
+            logger.info("Configuring compression.")
+            quantize = 10.0
+            compression = ImageCompressionOptions(ImageCompressionOptions.RICE, True, 0.0)
+            scaling = ImageScalingOptions(ImageScalingOptions.STDEV_BOTH, 32, quantizeLevel=quantize)
+            opts["imageOptions"] = ImageWriteOptions(compression, scaling)
+            opts["maskOptions"] = ImageWriteOptions(compression)
+            opts["varianceOptions"] = opts["imageOptions"]
+        exposure.writeFits(manager, **opts)
+
+        logger.info("Encoding exposure...")
+        content = base64.b64encode(manager.getData()).decode("ascii")
+
+        return {
+            "fits": content,
+        }
+
+
+GetFitsImageCommand.register("get fits image")

scripts/config.yaml

@@ -10,5 +10,7 @@ schemas:
 repos:
   - /repo/main
   - embargo_or4
+butlers:
+  - embargo
 #efd:
 # url: connection info here

scripts/mock_server.py

@@ -51,6 +51,7 @@ class WebSocketHandler(tornado.websocket.WebSocketHandler):
     workers: dict[str, WorkerPod] = dict()  # Keep track of connected worker pods
     clients: dict[str, WebSocketHandler] = dict()  # Keep track of connected clients
     queue: list[QueueItem] = list()  # Queue of messages to be processed
+    max_message_size: int = 100 * 1024 * 1024  # Maximum message size
 
     @classmethod
     def urls(cls) -> list[tuple[str, type[tornado.web.RequestHandler], dict[str, str]]]:

scripts/rubintv_worker.py

@@ -159,6 +159,7 @@ def main():
     butlers: dict[str, Butler] | None = None
     if "butlers" in config:
         logger.info("Connecting to Butlers")
+        butlers = {}
         for repo in config["butlers"]:
             butlers[repo] = Butler(repo)  # type: ignore
 

tests/test_command.py

@@ -24,6 +24,7 @@
 
 import astropy.table
 import lsst.rubintv.analysis.service as lras
+import pytest
 import utils
 
 
@@ -186,3 +187,23 @@ def test_errors(self):
             content,
             "execution error",
         )
+
+
+# Only runs if butler instantiated
+@pytest.mark.skip(reason="Needs butler access")
+class TestSendFitsImageCommand(TestCommand):
+    def test_send_fits_image_command(self):
+        command = {
+            "name": "get fits image",
+            "parameters": {
+                "repo": "embargo",
+                "collection": "u/kadrlica/binCalexp4",
+                "image_name": "calexpBinned8",
+                "data_id": {"instrument": "LSSTComCam", "detector": 3, "visit": 2024110900185},
+            },
+        }
+
+        print(lras.command.BaseCommand.command_registry)
+        content = self.execute_command(command, "get fits image")
+        length = len(content["fits"])
+        self.assertEqual(length, 4608000)