diff --git a/src/DIRAC/WorkloadManagementSystem/Agent/PilotStatusAgent.py b/src/DIRAC/WorkloadManagementSystem/Agent/PilotStatusAgent.py index a32f60a2ae8..1ebb7669e94 100644 --- a/src/DIRAC/WorkloadManagementSystem/Agent/PilotStatusAgent.py +++ b/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( diff --git a/src/DIRAC/WorkloadManagementSystem/Agent/SiteDirector.py b/src/DIRAC/WorkloadManagementSystem/Agent/SiteDirector.py index ef4a14426f4..05a71ca5734 100644 --- a/src/DIRAC/WorkloadManagementSystem/Agent/SiteDirector.py +++ b/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") diff --git a/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py b/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py index 8273ca12da4..feb0cc74c66 100644 --- a/src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_PilotStatusAgent.py +++ b/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): diff --git a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py index 2bd9f79e292..5c4a598cf37 100755 --- a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py +++ b/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) diff --git a/tests/Integration/WorkloadManagementSystem/Test_PilotAgentsDB.py b/tests/Integration/WorkloadManagementSystem/Test_PilotAgentsDB.py index 3243ae9ad54..eed90f22ee5 100644 --- a/tests/Integration/WorkloadManagementSystem/Test_PilotAgentsDB.py +++ b/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)