DM-52067: Make Sasquatch dispatch robust to an unresponsive or broken proxy

python/lsst/analysis/tools/atools/calibQuantityProfile.py

@@ -33,6 +33,8 @@
 
 from typing import cast
 
+import numpy as np
+
 from lsst.pex.config import Field
 from lsst.pex.config.configurableActions import ConfigurableActionField
 
@@ -89,21 +91,20 @@ class SingleValueRepacker(KeyedDataAction):
     )
 
     def __call__(self, data: KeyedData, **kwargs) -> KeyedData:
-        repackedData = {}
-        uniquePanelKeys = list(set(data[self.panelKey]))
-
-        # Loop over data vector to repack information as it is expected.
-        for i in range(len(uniquePanelKeys)):
-            repackedData[f"{uniquePanelKeys[i]}_x"] = []
-            repackedData[f"{uniquePanelKeys[i]}"] = []
+        repackedData: dict[str, list] = {}
 
         panelVec = cast(Vector, data[self.panelKey])
         dataVec = cast(Vector, data[self.dataKey])
         quantityVec = cast(Vector, data[self.quantityKey])
+        # cp_verify's results tables also hold per-detector and per-exposure
+        # rows, which have no amplifier and so belong in no panel.
+        panelMask = np.ma.getmaskarray(panelVec)
 
         for i in range(len(panelVec)):
-            repackedData[f"{panelVec[i]}_x"].append(dataVec[i])
-            repackedData[f"{panelVec[i]}"].append(quantityVec[i])
+            if panelMask[i]:
+                continue
+            repackedData.setdefault(f"{panelVec[i]}_x", []).append(dataVec[i])
+            repackedData.setdefault(f"{panelVec[i]}", []).append(quantityVec[i])
 
         return repackedData
 

python/lsst/analysis/tools/interfaces/datastore/_dispatcher.py

@@ -54,6 +54,9 @@
 PARTITIONS = 1
 REPLICATION_FACTOR = 3
 
+DEFAULT_TIMEOUT: float = 30.0
+"""Default timeout in seconds applied to each HTTP request to the proxy."""
+
 IDENTIFIER_KEYS = [
     "detector",
     "patch",
@@ -179,6 +182,10 @@ class SasquatchDispatcher:
     namespace: str = "lsst.dm"
     """The namespace in Sasquatch in which to write the uploaded metrics"""
 
+    timeout: float | None = DEFAULT_TIMEOUT
+    """Timeout in seconds for each HTTP request to the proxy, or `None` for
+    no timeout."""
+
     def __post_init__(self) -> None:
         match ResourcePath(self.url).scheme:
             case "http" | "https":
@@ -208,10 +215,10 @@ def _populateClusterId(self) -> None:
 
         try:
             with http_client() as session:
-                r = session.get(f"{self.url}/v3/clusters", headers=headers)
+                r = session.get(f"{self.url}/v3/clusters", headers=headers, timeout=self.timeout)
+                r.raise_for_status()
                 cluster_id = r.json()["data"][0]["cluster_id"]
                 self._cluster_id = str(cluster_id)
-            r.raise_for_status()
         except requests.RequestException as e:
             raise SasquatchDispatchFailure("Could not retrieve the cluster id for the specified url") from e
 
@@ -242,7 +249,10 @@ def _create_topic(self, topic_name: str) -> bool:
         try:
             with http_client() as session:
                 r = session.post(
-                    f"{self.url}/v3/clusters/{self.clusterId}/topics", json=topic_config, headers=headers
+                    f"{self.url}/v3/clusters/{self.clusterId}/topics",
+                    json=topic_config,
+                    headers=headers,
+                    timeout=self.timeout,
                 )
             r.raise_for_status()
             log.debug("Created topic %s.%s", self.namespace, topic_name)
@@ -716,7 +726,10 @@ def dispatch(
 
                 try:
                     r = session.post(
-                        f"{self.url}/topics/{self.namespace}.{metric}", json=data, headers=headers
+                        f"{self.url}/topics/{self.namespace}.{metric}",
+                        json=data,
+                        headers=headers,
+                        timeout=self.timeout,
                     )
                     r.raise_for_status()
                     log.debug("Succesfully sent data for metric %s", metric)

python/lsst/analysis/tools/interfaces/datastore/_sasquatchDatastore.py

@@ -36,6 +36,7 @@
 from lsst.resources import ResourcePath, ResourcePathExpression
 
 from . import SasquatchDispatcher, SasquatchDispatchFailure, SasquatchDispatchPartialFailure
+from ._dispatcher import DEFAULT_TIMEOUT
 
 if TYPE_CHECKING:
     from lsst.daf.butler import Config, DatasetProvenance, DatasetType, LookupKey
@@ -99,6 +100,14 @@ class SasquatchDatastore(GenericBaseDatastore):
     "k1=v1;k2=v2".
     """
 
+    timeout: float | None
+    """Timeout in seconds for each HTTP request made to the rest proxy, or
+    `None` for no timeout.
+
+    Read from the ``"timeout"`` entry of the datastore config, defaulting to
+    `DEFAULT_TIMEOUT`.
+    """
+
     def __init__(
         self,
         config: DatastoreConfig,
@@ -133,7 +142,11 @@ def __init__(
                 extra_fields[k] = v
         self.extra_fields = extra_fields if extra_fields else None
 
-        self._dispatcher = SasquatchDispatcher(self.restProxyUrl, self.accessToken, self.namespace)
+        self.timeout = self.config.get("timeout", DEFAULT_TIMEOUT)
+
+        self._dispatcher = SasquatchDispatcher(
+            self.restProxyUrl, self.accessToken, self.namespace, timeout=self.timeout
+        )
 
     @classmethod
     def _create_from_config(

tests/test_calibQuantityProfile.py

@@ -0,0 +1,111 @@
+# This file is part of analysis_tools.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (https://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""Tests for the calibQuantityProfile tools on cp_verify results tables.
+
+These catch rows with no panel key (the per-detector and per-exposure rows,
+masked in the amplifier column) breaking the repack, and with it
+analyzeFlatDetCore's flatTestsByDate plot.
+"""
+
+import unittest
+
+import matplotlib
+import matplotlib.pyplot as plt
+from astropy.table import Table, vstack
+from matplotlib.figure import Figure
+
+import lsst.utils.tests
+from lsst.analysis.tools.atools.calibQuantityProfile import CalibAmpScatterTool, SingleValueRepacker
+
+matplotlib.use("Agg")
+
+AMP_NAMES = [f"C{i}{j}" for i in (0, 1) for j in range(8)]
+DETECTORS = (90, 91)
+MJD = 61224.4
+
+
+def makeResultsTable() -> Table:
+    """Make a table shaped like cp_verify's ``verifyFlatResults``.
+
+    Returns
+    -------
+    table : `astropy.table.Table`
+        One row per amplifier per detector, plus a per-detector row for each
+        detector and a per-exposure row. Those last have no amplifier, so are
+        masked in the ``amplifier`` column.
+    """
+    nAmpRows = len(DETECTORS) * len(AMP_NAMES)
+    ampRows = Table(
+        {
+            "detector": [detector for detector in DETECTORS for _ in AMP_NAMES],
+            "amplifier": AMP_NAMES * len(DETECTORS),
+            "mjd": [MJD] * nAmpRows,
+            "FLAT_VERIFY_NOISE": [True] * nAmpRows,
+        }
+    )
+    detectorRows = Table(
+        {"detector": list(DETECTORS), "mjd": [MJD] * len(DETECTORS), "FLAT_DET_SCATTER": [0.01, 0.01]}
+    )
+    exposureRow = Table({"detector": [DETECTORS[0]], "flat_SCATTER": [0.001]})
+    return vstack([ampRows, detectorRows, exposureRow])
+
+
+class SingleValueRepackerTestCase(lsst.utils.tests.TestCase):
+    def testSkipsRowsWithoutPanelKey(self) -> None:
+        # The rows with no amplifier used to raise
+        # "TypeError: unhashable type: 'MaskedConstant'", so
+        # analyzeFlatDetCore never wrote its plot.
+        table = makeResultsTable()
+        # Without masked rows in the fixture this test pins nothing.
+        self.assertTrue(table["amplifier"].mask.any())
+        repacker = SingleValueRepacker(panelKey="amplifier", dataKey="mjd", quantityKey="FLAT_VERIFY_NOISE")
+        repacked = repacker(table)
+        self.assertEqual(set(repacked), {f"{amp}{suffix}" for amp in AMP_NAMES for suffix in ("", "_x")})
+        for amp in AMP_NAMES:
+            self.assertEqual(repacked[f"{amp}_x"], [MJD] * len(DETECTORS))
+            self.assertEqual(repacked[amp], [True] * len(DETECTORS))
+
+
+class CalibAmpScatterToolTestCase(lsst.utils.tests.TestCase):
+    def testPlotsResultsTable(self) -> None:
+        # The whole tool, configured as analyzeFlatDetCore's flatTestsByDate
+        # in cpCore.yaml: a masked amplifier anywhere along the way must not
+        # stop the plot being made.
+        tool = CalibAmpScatterTool()
+        tool.prep.quantityKey = "FLAT_VERIFY_NOISE"
+        tool.finalize()
+        results = tool(makeResultsTable())
+        self.assertIsInstance(results["GridPlot"], Figure)
+        plt.close(results["GridPlot"])
+
+
+class MemoryTester(lsst.utils.tests.MemoryTestCase):
+    pass
+
+
+def setup_module(module):
+    lsst.utils.tests.init()
+
+
+if __name__ == "__main__":
+    lsst.utils.tests.init()
+    unittest.main()

tests/test_sasquatchDatastore.py

@@ -19,23 +19,85 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
+"""Tests for SasquatchDatastore and SasquatchDispatcher.
+
+These catch (1) a put that no longer hands the bundle to the dispatcher,
+(2) a broken or unresponsive Sasquatch proxy failing or hanging the put,
+and with it the pipeline task, and (3) explicit timestamp versions being
+misparsed, or accepted when malformed.
+"""
+
 import datetime
 import os
 import unittest
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import astropy.units as u
+import requests
 
 import lsst.daf.butler.tests as butlerTests
 from lsst.analysis.tools.interfaces import MetricMeasurementBundle
-from lsst.analysis.tools.interfaces.datastore import SasquatchDispatcher
+from lsst.analysis.tools.interfaces.datastore import SasquatchDispatcher, SasquatchDispatchFailure
+from lsst.analysis.tools.interfaces.datastore._dispatcher import DEFAULT_TIMEOUT
 from lsst.daf.butler import CollectionType, Config
 from lsst.daf.butler.tests.utils import makeTestTempDir, removeTestTempDir
 from lsst.verify import Measurement
 
 TESTDIR = os.path.abspath(os.path.dirname(__file__))
 CONFIG_FILE = os.path.join(TESTDIR, "config", "butler-sasquatch.yaml")
 
+# Where the dispatcher looks up its HTTP session factory, for patching.
+HTTP_CLIENT = "lsst.analysis.tools.interfaces.datastore._dispatcher.http_client"
+
+
+def _fakeResponse(status: int = 200, json: object = None, jsonError: Exception | None = None) -> MagicMock:
+    """Make a stand-in for a `requests.Response`.
+
+    Parameters
+    ----------
+    status : `int`, optional
+        HTTP status code; ``raise_for_status`` raises for 4xx and 5xx.
+    json : `object`, optional
+        Value returned by ``json()``.
+    jsonError : `Exception`, optional
+        Exception raised by ``json()`` in place of returning ``json``.
+
+    Returns
+    -------
+    response : `unittest.mock.MagicMock`
+        The fake response.
+    """
+    response = MagicMock()
+    response.status_code = status
+    if status >= 400:
+        response.raise_for_status.side_effect = requests.HTTPError(f"HTTP {status}", response=response)
+    if jsonError is not None:
+        response.json.side_effect = jsonError
+    else:
+        response.json.return_value = json
+    return response
+
+
+def _fakeSession(getResponse: MagicMock | None = None, postResponse: MagicMock | None = None) -> MagicMock:
+    """Make a stand-in for a `requests.Session` returning canned responses.
+
+    Parameters
+    ----------
+    getResponse : `unittest.mock.MagicMock`, optional
+        Response returned by ``get``; a bare 200 if not given.
+    postResponse : `unittest.mock.MagicMock`, optional
+        Response returned by ``post``; a bare 200 if not given.
+
+    Returns
+    -------
+    session : `unittest.mock.MagicMock`
+        The fake session.
+    """
+    session = MagicMock()
+    session.get.return_value = getResponse if getResponse is not None else _fakeResponse()
+    session.post.return_value = postResponse if postResponse is not None else _fakeResponse()
+    return session
+
 
 class SasquatchDatastoreTest(unittest.TestCase):
     def setUp(self):
@@ -44,6 +106,7 @@ def setUp(self):
         config = Config()
         config["datastore", "cls"] = "lsst.analysis.tools.interfaces.datastore.SasquatchDatastore"
         config["datastore", "restProxyUrl"] = "https://example.com/sasquatch-rest-proxy"
+        config["datastore", "timeout"] = 5.5
 
         dataIds = {
             "instrument": ["DummyCam"],
@@ -73,6 +136,44 @@ def test_put(self):
         mock_method.assert_called()
         self.assertIs(mock_method.call_args[0][0], bundle)
 
+    def testTimeoutConfiguration(self) -> None:
+        # The configured timeout must reach the dispatcher, and a dispatcher
+        # made without one must still get DEFAULT_TIMEOUT rather than none.
+        self.assertEqual(self.butler._datastore.timeout, 5.5)
+        self.assertEqual(self.butler._datastore._dispatcher.timeout, 5.5)
+        self.assertEqual(SasquatchDispatcher("http://test.local", "na").timeout, DEFAULT_TIMEOUT)
+
+    def testRequestsUseTimeout(self) -> None:
+        # An HTTP call made without the timeout lets a proxy that accepts
+        # connections but never answers hang the put indefinitely.
+        dispatcher = SasquatchDispatcher("http://test.local", "na", timeout=12.5)
+        session = _fakeSession(getResponse=_fakeResponse(json={"data": [{"cluster_id": "abc"}]}))
+
+        with patch(HTTP_CLIENT) as mockHttpClient:
+            mockHttpClient.return_value.__enter__.return_value = session
+            self.assertEqual(dispatcher.clusterId, "abc")
+            self.assertTrue(dispatcher._create_topic("some.metric"))
+
+        self.assertEqual(session.get.call_args.kwargs["timeout"], 12.5)
+        self.assertEqual(session.post.call_args.kwargs["timeout"], 12.5)
+
+    def testClusterIdFailuresRaiseDispatchFailure(self) -> None:
+        # Each bad response must surface as SasquatchDispatchFailure, the
+        # failure the datastore expects. Parsing the body before checking
+        # the status lets the error document through as a bare KeyError.
+        dispatcher = SasquatchDispatcher("http://test.local", "na")
+        jsonError = requests.exceptions.JSONDecodeError("Expecting value", "<html>", 0)
+        badResponses = {
+            "5xx with an error document": _fakeResponse(503, json={"error_code": 50301}),
+            "2xx with a non-JSON body": _fakeResponse(200, jsonError=jsonError),
+        }
+        for description, response in badResponses.items():
+            with self.subTest(description):
+                with patch(HTTP_CLIENT) as mockHttpClient:
+                    mockHttpClient.return_value.__enter__.return_value = _fakeSession(getResponse=response)
+                    with self.assertRaises(SasquatchDispatchFailure):
+                        _ = dispatcher.clusterId
+
     def test_explicit_timestamp_version(self):
         dispatcher = SasquatchDispatcher("http://test.local", "na")
         bundle = MetricMeasurementBundle()