feat: use temporary table instead of large IN query in PilotAgentsDB

src/DIRAC/WorkloadManagementSystem/Agent/PilotStatusAgent.py

@@ -82,22 +82,15 @@ def handleOldPilots(self, connection):
         """
         pilotsToAccount = {}
         timeLimitToConsider = TimeUtilities.toString(DiracTime.utcnow() - TimeUtilities.day * self.pilotStalledDays)
-        result = self.pilotDB.selectPilots(
+        result = self.pilotDB.getPilotInfoForSelection(
             {"Status": PilotStatus.PILOT_TRANSIENT_STATES}, older=timeLimitToConsider, timeStamp="LastUpdateTime"
         )
         if not result["OK"]:
-            self.log.error("Failed to get the Pilot Agents")
+            self.log.error("Failed to get Info for Pilot Agents", result["Message"])
             return result
-        if not result["Value"]:
-            return S_OK()
-
-        refList = result["Value"]
-        result = self.pilotDB.getPilotInfo(refList)
-        if not result["OK"]:
-            self.log.error("Failed to get Info for Pilot Agents")
-            return result
-
         pilotsDict = result["Value"]
+        if not pilotsDict:
+            return S_OK()
 
         for pRef in pilotsDict:
             if pilotsDict[pRef].get("Jobs") and self._checkJobLastUpdateTime(

src/DIRAC/WorkloadManagementSystem/Agent/SiteDirector.py

@@ -748,8 +748,8 @@ def _monitorPilotsPerQueue(self, queue: str):
         if callable(getattr(ce, "cleanupPilots", None)):
             ce.cleanupPilots()
 
-        # Select pilots in a transient states
-        result = self.pilotAgentsDB.selectPilots(
+        # Get the information of the pilots in a transient state
+        result = self.pilotAgentsDB.getPilotInfoForSelection(
             {
                 "DestinationSite": ceName,
                 "Queue": queueName,
@@ -759,19 +759,13 @@ def _monitorPilotsPerQueue(self, queue: str):
                 "VO": self.vo,
             }
         )
-        if not result["OK"]:
-            self.log.error("Failed to select pilots", f": {result['Message']}")
-            return result
-        pilotRefs = result["Value"]
-        if not pilotRefs:
-            return S_OK()
-
-        # Get their information
-        result = self.pilotAgentsDB.getPilotInfo(pilotRefs)
         if not result["OK"]:
             self.log.error("Failed to get pilots info from DB", result["Message"])
             return result
         pilotDict = result["Value"]
+        if not pilotDict:
+            return S_OK()
+        pilotRefs = list(pilotDict)
 
         # The proxy is used for checking the pilot status and renewals
         # We really need at least a few hours otherwise the renewed
@@ -799,7 +793,7 @@ def _monitorPilotsPerQueue(self, queue: str):
 
         # Check if the accounting is to be sent
         if self.sendAccounting:
-            result = self.pilotAgentsDB.selectPilots(
+            result = self.pilotAgentsDB.getPilotInfoForSelection(
                 {
                     "DestinationSite": ceName,
                     "Queue": queueName,
@@ -809,20 +803,14 @@ def _monitorPilotsPerQueue(self, queue: str):
                     "Status": PilotStatus.PILOT_FINAL_STATES,
                 }
             )
-            if not result["OK"]:
-                self.log.error("Failed to select pilots", result["Message"])
-                return result
-
-            pilotRefs = result["Value"]
-            if not pilotRefs:
-                return S_OK()
-
-            result = self.pilotAgentsDB.getPilotInfo(pilotRefs)
             if not result["OK"]:
                 self.log.error("Failed to get pilots info from DB", result["Message"])
                 return result
 
             pilotDict = result["Value"]
+            if not pilotDict:
+                return S_OK()
+
             result = self._sendPilotAccounting(pilotDict)
             if not result["OK"]:
                 self.log.error("Failed to send pilot agent accounting")

src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py

@@ -29,7 +29,7 @@ def psa(mocker):
         side_effect=lambda x, y=None: y,
         create=True,
     )
-    module_str = "DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB.selectPilots"
+    module_str = "DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB.getPilotInfoForSelection"
     mocker.patch(module_str, side_effect=mockReply)
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.PilotAgentsDB.__init__", side_effect=mockNone)
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.PilotStatusAgent.JobDB.__init__", side_effect=mockNone)
@@ -55,16 +55,8 @@ def psa(mocker):
 @pytest.mark.parametrize(
     "mockReplyInput, expected",
     [
-        ({"OK": True, "Value": False}, {"OK": True, "Value": None}),
-        (
-            {"OK": True, "Value": ["Test"]},
-            {
-                "OK": False,
-                "Message": "No pilots found for PilotJobReference(s): ['Test']",
-            },
-        ),
+        ({"OK": True, "Value": {}}, {"OK": True, "Value": None}),
         ({"OK": False, "Message": "Test"}, {"OK": False, "Message": "Test"}),
-        ({"OK": True, "Value": []}, {"OK": True, "Value": None}),
     ],
 )
 def test_handleOldPilots(psa, mockReplyInput, expected):

src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py

@@ -260,8 +260,52 @@ def clearPilots(self, interval=30, aborted_interval=7):
 
     ##########################################################################################
     def getPilotInfo(self, pilotRef=False, conn=False, paramNames=[], pilotID=False):
-        """Get all the information for the pilot job reference or reference list"""
+        """Get all the information for the pilot job reference or reference list
 
+        To get the information of pilots matching some criteria, prefer :meth:`getPilotInfoForSelection`
+        over :meth:`selectPilots` followed by this method: it avoids sending a potentially huge list of
+        references back to the DB in an ``IN`` clause.
+        """
+        condSQL = []
+        for key, value in [
+            ("PilotJobReference", pilotRef),
+            ("PilotID", pilotID),
+        ]:
+            resList = []
+            for v in value if isinstance(value, list) else [value] if value else []:
+                result = self._escapeString(v)
+                if not result["OK"]:
+                    return result
+                resList.append(result["Value"])
+            if resList:
+                condSQL.append(f"{key} IN ({','.join(resList)})")
+        condition = f"WHERE {' AND '.join(condSQL)}" if condSQL else ""
+
+        result = self._getPilotInfoForCondition(condition, paramNames, conn=conn)
+        if not result["OK"]:
+            return result
+        if not result["Value"]:
+            msg = "No pilots found"
+            if pilotRef:
+                msg += f" for PilotJobReference(s): {pilotRef}"
+            return S_ERROR(DErrno.EWMSNOPILOT, msg)
+        return result
+
+    def getPilotInfoForSelection(
+        self, condDict, older=None, newer=None, timeStamp="SubmissionTime", paramNames=None, conn=False
+    ):
+        """Get the information of the pilots matching the selection, in the format of :meth:`getPilotInfo`
+
+        The selection arguments are the ones of :meth:`selectPilots`.
+
+        :return: S_OK with a dictionary keyed by the first parameter (PilotJobReference by default);
+                 empty if no pilot matches
+        """
+        condition = self.buildCondition(condDict, older, newer, timeStamp)
+        return self._getPilotInfoForCondition(condition, paramNames, conn=conn)
+
+    def _getPilotInfoForCondition(self, condition, paramNames=None, conn=False):
+        """Get the pilot information for the pilots matching a SQL condition (e.g. ``WHERE ...``)"""
         parameters = (
             [
                 "PilotJobReference",
@@ -285,51 +329,28 @@ def getPilotInfo(self, pilotRef=False, conn=False, paramNames=[], pilotID=False)
             if not col in self.ALLOWED_COLUMNS:
                 return S_ERROR(f"Unknown column: {col}")
 
-        cmd = f"SELECT {', '.join(parameters)} FROM PilotAgents"  # nosec
-        condSQL = []
-        for key, value in [
-            ("PilotJobReference", pilotRef),
-            ("PilotID", pilotID),
-        ]:
-            resList = []
-            for v in value if isinstance(value, list) else [value] if value else []:
-                result = self._escapeString(v)
-                if not result["OK"]:
-                    return result
-                resList.append(result["Value"])
-            if resList:
-                condSQL.append(f"{key} IN ({','.join(resList)})")
-        if condSQL:
-            cmd = f"{cmd} WHERE {' AND '.join(condSQL)}"
-
-        result = self._query(cmd, conn=conn)
+        result = self._query(f"SELECT {', '.join(parameters)} FROM PilotAgents {condition}", conn=conn)  # nosec
         if not result["OK"]:
             return result
-        if not result["Value"]:
-            msg = "No pilots found"
-            if pilotRef:
-                msg += f" for PilotJobReference(s): {pilotRef}"
-            return S_ERROR(DErrno.EWMSNOPILOT, msg)
 
-        resDict = {}
-        pilotIDs = []
-        for row in result["Value"]:
-            pilotDict = {}
-            for i, par in enumerate(parameters):
-                pilotDict[par] = row[i]
-                if par == "PilotID":
-                    pilotIDs.append(row[i])
-            resDict[row[0]] = pilotDict
-
-        result = self.getJobsForPilot(pilotIDs)
-        if not result["OK"]:
+        resDict = {row[0]: dict(zip(parameters, row)) for row in result["Value"]}
+        if not resDict or "PilotID" not in parameters:
             return S_OK(resDict)
 
-        jobsDict = result["Value"]
-        for pilotRef, pilotInfo in resDict.items():
-            pilotID = pilotInfo["PilotID"]
-            if pilotID in jobsDict:
-                pilotInfo["Jobs"] = jobsDict[pilotID]
+        # Reuse the condition in a subquery rather than sending the list of PilotIDs back to the DB
+        cmd = (
+            "SELECT PilotID, JobID FROM JobToPilotMapping"
+            f" WHERE PilotID IN (SELECT PilotID FROM PilotAgents {condition})"  # nosec
+        )
+        result = self._query(cmd, conn=conn)
+        if not result["OK"]:
+            return S_OK(resDict)
+        jobsDict = {}
+        for pID, jobID in result["Value"]:
+            jobsDict.setdefault(pID, []).append(jobID)
+        for pilotInfo in resDict.values():
+            if pilotInfo["PilotID"] in jobsDict:
+                pilotInfo["Jobs"] = jobsDict[pilotInfo["PilotID"]]
 
         return S_OK(resDict)
 

tests/Integration/WorkloadManagementSystem/Test_PilotAgentsDB.py

@@ -186,3 +186,27 @@ def test_PivotedPilotSummaryTable():
     assert "Total" in columns
 
     cleanUpPilots(pilotRef)
+
+
+def test_getPilotInfoForSelection():
+    """Selecting pilot info by criteria gives the same result as selectPilots + getPilotInfo"""
+    pilotRef = preparePilots([2, 3], "TestSite", "TestCE", "VO")
+    try:
+        res = paDB.setJobForPilot(123, "pilotRef_0", updateStatus=False)
+        assert res["OK"] is True, res["Message"]
+
+        condDict = {"DestinationSite": "TestCE", "Status": "Submitted"}
+        res = paDB.getPilotInfoForSelection(condDict)
+        assert res["OK"] is True, res["Message"]
+        assert sorted(res["Value"]) == ["pilotRef_0", "pilotRef_1"]
+        assert res["Value"]["pilotRef_0"]["Jobs"] == [123]
+        assert "Jobs" not in res["Value"]["pilotRef_1"]
+
+        refs = paDB.selectPilots(condDict)["Value"]
+        assert paDB.getPilotInfo(refs)["Value"] == res["Value"]
+
+        res = paDB.getPilotInfoForSelection({"DestinationSite": "TestCE", "Status": "Unknown"})
+        assert res["OK"] is True, res["Message"]
+        assert res["Value"] == {}
+    finally:
+        cleanUpPilots(pilotRef)