refactor: extracted getJobParameters in utility, using it in agents

Author: fstagni

src/DIRAC/WorkloadManagementSystem/Agent/PushJobAgent.py

@@ -12,16 +12,14 @@
 import hashlib
 import json
 import os
-from pathlib import Path
 import random
 import shutil
 import sys
-from collections import defaultdict
 import time
+from collections import defaultdict
+from pathlib import Path
 
-from diraccfg import CFG
-
-from DIRAC import gConfig, S_OK, S_ERROR
+from DIRAC import S_ERROR, S_OK, gConfig
 from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations
 from DIRAC.ConfigurationSystem.Client.Helpers.Resources import getQueues
 from DIRAC.Core.Utilities import DErrno
@@ -30,6 +28,7 @@
 from DIRAC.Core.Utilities.Version import getVersion
 from DIRAC.FrameworkSystem.Client.ProxyManagerClient import gProxyManager
 from DIRAC.Resources.Computing import ComputingElement
+from DIRAC.WorkloadManagementSystem.Agent.JobAgent import JobAgent
 from DIRAC.WorkloadManagementSystem.Client import JobMinorStatus, JobStatus, PilotStatus
 from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
 from DIRAC.WorkloadManagementSystem.Client.JobReport import JobReport
@@ -41,11 +40,11 @@
     resolveInputData,
     transferInputSandbox,
 )
-from DIRAC.WorkloadManagementSystem.Utilities.QueueUtilities import getQueuesResolved
-from DIRAC.WorkloadManagementSystem.Agent.JobAgent import JobAgent
-from DIRAC.WorkloadManagementSystem.Utilities.Utils import createJobWrapper
 from DIRAC.WorkloadManagementSystem.private.ConfigHelper import findGenericPilotCredentials
 from DIRAC.WorkloadManagementSystem.Utilities.QueueUtilities import getQueuesResolved
+from DIRAC.WorkloadManagementSystem.Utilities.Utils import createJobWrapper
+from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
+
 
 MAX_JOBS_MANAGED = 100
 
@@ -740,7 +739,7 @@ def _checkSubmittedJobWrappers(self, ce: ComputingElement, site: str):
             return S_OK()
 
         # Get their parameters
-        if not (result := self.jobMonitoring.getJobParameters(jobs, ["GridCE", "TaskID", "Stamp"]))["OK"]:
+        if not (result := getJobParameters(jobs, ["GridCE", "TaskID", "Stamp"]))["OK"]:
             self.log.error("Failed to get the list of taskIDs", result["Message"])
             return result
 

src/DIRAC/WorkloadManagementSystem/Agent/StalledJobAgent.py

@@ -21,11 +21,11 @@
 from DIRAC.Core.Utilities.TimeUtilities import fromString, second, toEpoch
 from DIRAC.WorkloadManagementSystem.Client import JobMinorStatus, JobStatus
 from DIRAC.WorkloadManagementSystem.Client.JobManagerClient import JobManagerClient
-from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
 from DIRAC.WorkloadManagementSystem.Client.WMSClient import WMSClient
 from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
 from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB
 from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB
+from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
 
 
 class StalledJobAgent(AgentModule):
@@ -254,11 +254,11 @@ def _failStalledJobs(self, jobID):
 
     def _getJobPilotStatus(self, jobID):
         """Get the job pilot status."""
-        result = JobMonitoringClient().getJobParameter(jobID, "Pilot_Reference")
+        result = getJobParameters(jobID, "Pilot_Reference")
         if not result["OK"]:
             return result
-        pilotReference = result["Value"].get("Pilot_Reference", "Unknown")
-        if pilotReference == "Unknown":
+        pilotReference = result["Value"].get("Pilot_Reference")
+        if not pilotReference:
             # There is no pilot reference, hence its status is unknown
             return S_OK("NoPilot")
 
@@ -389,7 +389,7 @@ def _sendAccounting(self, jobID):
             if lastHeartBeatTime is not None and lastHeartBeatTime > endTime:
                 endTime = lastHeartBeatTime
 
-            result = JobMonitoringClient().getJobParameter(jobID, "CPUNormalizationFactor")
+            result = getJobParameters(jobID, "CPUNormalizationFactor")
             if not result["OK"] or not result["Value"]:
                 self.log.error(
                     "Error getting Job Parameter CPUNormalizationFactor, setting 0",

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

@@ -25,9 +25,9 @@ def sja(mocker):
     )
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.JobDB")
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.JobLoggingDB")
-    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.JobMonitoringClient", return_value=MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.JobManagerClient", return_value=MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.PilotAgentsDB", return_value=MagicMock())
+    mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.getJobParameters", return_value=MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.WMSClient", return_value=MagicMock())
     mocker.patch("DIRAC.WorkloadManagementSystem.Agent.StalledJobAgent.getDNForUsername", return_value=MagicMock())
 

src/DIRAC/WorkloadManagementSystem/Service/JobMonitoringHandler.py

@@ -11,6 +11,7 @@
 from DIRAC.Core.Utilities.JEncode import strToIntDict
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
 from DIRAC.WorkloadManagementSystem.Client.PilotManagerClient import PilotManagerClient
+from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
 
 
 class JobMonitoringHandlerMixin:
@@ -270,48 +271,8 @@ def export_getJobParameters(self, jobIDs, parName=None):
         if not isinstance(jobIDs, list):
             jobIDs = [jobIDs]
         jobIDs = [int(jobID) for jobID in jobIDs]
-        if self.vo:  # a user is connecting, with a proxy
-            res = self.elasticJobParametersDB.getJobParameters(jobIDs, self.vo, parName)
-            if not res["OK"]:
-                return res
-            parameters = res["Value"]
-        else:  # a service is connecting, no proxy, e.g. StalledJobAgent
-            q = f"SELECT JobID, VO FROM Jobs WHERE JobID IN ({','.join([str(jobID) for jobID in jobIDs])})"
-            res = self.jobDB._query(q)
-            if not res["OK"]:
-                return res
-            if not res["Value"]:
-                return S_OK({})
-            # get the VO for each jobID
-            voDict = {}
-            for jobID, vo in res["Value"]:
-                if vo not in voDict:
-                    voDict[vo] = []
-                voDict[vo].append(jobID)
-            # get the parameters for each VO
-            parameters = {}
-            for vo, jobIDs in voDict.items():
-                res = self.elasticJobParametersDB.getJobParameters(jobIDs, vo, parName)
-                if not res["OK"]:
-                    return res
-                parameters.update(res["Value"])
-
-        # Need anyway to get also from JobDB, for those jobs with parameters registered in MySQL or in both backends
-        res = self.jobDB.getJobParameters(jobIDs, parName)
-        if not res["OK"]:
-            return res
-        parametersM = res["Value"]
-
-        # and now combine
-        final = dict(parametersM)
-        # if job in JobDB, update with parameters from ES if any
-        for jobID in final:
-            final[jobID].update(parameters.get(jobID, {}))
-        # if job in ES and not in JobDB, take ES
-        for jobID in parameters:
-            if jobID not in final:
-                final[jobID] = parameters[jobID]
-        return S_OK(final)
+
+        return getJobParameters(jobIDs, parName, self.vo or "")
 
     ##############################################################################
     types_getAtticJobParameters = [int]

src/DIRAC/WorkloadManagementSystem/Utilities/JobParameters.py

@@ -2,8 +2,10 @@
 """
 import multiprocessing
 
-from DIRAC import gConfig, gLogger
+from DIRAC import gConfig, gLogger, S_OK
 from DIRAC.Core.Utilities.List import fromChar
+from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
+from DIRAC.WorkloadManagementSystem.DB.JobParametersDB import JobParametersDB
 
 
 def getMemoryFromProc():
@@ -139,3 +141,61 @@ def getNumberOfGPUs(siteName=None, gridCE=None, queue=None):
     # 3) return 0
     gLogger.info("NumberOfGPUs could not be found in CS")
     return 0
+
+
+def getJobParameters(jobIDs: list[int], parName: str | None, vo: str = "") -> dict:
+    """Utility to get a job parameter for a list of jobIDs pertaining to a VO.
+    If the jobID is not in the JobParametersDB, it will be looked up in the JobDB.
+    
+    Requires direct access to the JobParametersDB and JobDB.
+    :param jobIDs: list of jobIDs
+    :param parName: name of the parameter to be retrieved
+    :param vo: VO of the jobIDs
+    :return: dictionary with jobID as key and the parameter as value
+    :rtype: dict
+    """
+    elasticJobParametersDB = JobParametersDB()
+    jobDB = JobDB()
+
+    if vo:  # a user is connecting, with a proxy
+        res = elasticJobParametersDB.getJobParameters(jobIDs, vo, parName)
+        if not res["OK"]:
+            return res
+        parameters = res["Value"]
+    else:  # a service is connecting, no proxy, e.g. StalledJobAgent
+        q = f"SELECT JobID, VO FROM Jobs WHERE JobID IN ({','.join([str(jobID) for jobID in jobIDs])})"
+        res = jobDB._query(q)
+        if not res["OK"]:
+            return res
+        if not res["Value"]:
+            return S_OK({})
+        # get the VO for each jobID
+        voDict = {}
+        for jobID, vo in res["Value"]:
+            if vo not in voDict:
+                voDict[vo] = []
+            voDict[vo].append(jobID)
+        # get the parameters for each VO
+        parameters = {}
+        for vo, jobIDs in voDict.items():
+            res = elasticJobParametersDB.getJobParameters(jobIDs, vo, parName)
+            if not res["OK"]:
+                return res
+            parameters.update(res["Value"])
+
+    # Need anyway to get also from JobDB, for those jobs with parameters registered in MySQL or in both backends
+    res = jobDB.getJobParameters(jobIDs, parName)
+    if not res["OK"]:
+        return res
+    parametersM = res["Value"]
+
+    # and now combine
+    final = dict(parametersM)
+    # if job in JobDB, update with parameters from ES if any
+    for jobID in final:
+        final[jobID].update(parameters.get(jobID, {}))
+    # if job in ES and not in JobDB, take ES
+    for jobID in parameters:
+        if jobID not in final:
+            final[jobID] = parameters[jobID]
+    return S_OK(final)