refactor: moved getJobParameters in JobDBUtils

Author: fstagni

src/DIRAC/WorkloadManagementSystem/Agent/JobCleaningAgent.py

@@ -1,4 +1,4 @@
-""" The Job Cleaning Agent controls removing jobs from the WMS in the end of their life cycle.
+"""The Job Cleaning Agent controls removing jobs from the WMS in the end of their life cycle.
 
     This agent will take care of:
     - removing all jobs that are in status JobStatus.DELETED
@@ -22,6 +22,7 @@
 than 0.
 
 """
+
 import datetime
 import os
 
@@ -37,10 +38,10 @@
 from DIRAC.WorkloadManagementSystem.Client import JobStatus
 from DIRAC.WorkloadManagementSystem.Client.WMSClient import WMSClient
 from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
+from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import getJobParameters
 from DIRAC.WorkloadManagementSystem.DB.SandboxMetadataDB import SandboxMetadataDB
 from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
 from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_DELETE
-from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
 
 
 class JobCleaningAgent(AgentModule):

src/DIRAC/WorkloadManagementSystem/Agent/PushJobAgent.py

@@ -1,4 +1,4 @@
-"""  The Push Job Agent class inherits from Job Agent and aims to support job submission in
+"""The Push Job Agent class inherits from Job Agent and aims to support job submission in
      sites with no external connectivity (e.g. some supercomputers).
 
 .. literalinclude:: ../ConfigTemplate.cfg
@@ -32,6 +32,7 @@
 from DIRAC.WorkloadManagementSystem.Client import JobMinorStatus, JobStatus, PilotStatus
 from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
 from DIRAC.WorkloadManagementSystem.Client.JobReport import JobReport
+from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import getJobParameters
 from DIRAC.WorkloadManagementSystem.JobWrapper.JobWrapper import JobWrapper
 from DIRAC.WorkloadManagementSystem.JobWrapper.JobWrapperUtilities import (
     getJobWrapper,
@@ -41,7 +42,6 @@
     transferInputSandbox,
 )
 from DIRAC.WorkloadManagementSystem.private.ConfigHelper import findGenericPilotCredentials
-from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
 from DIRAC.WorkloadManagementSystem.Utilities.QueueUtilities import getQueuesResolved
 from DIRAC.WorkloadManagementSystem.Utilities.Utils import createJobWrapper
 

src/DIRAC/WorkloadManagementSystem/Agent/StalledJobAgent.py

@@ -8,6 +8,7 @@
   :dedent: 2
   :caption: StalledJobAgent options
 """
+
 import concurrent.futures
 import datetime
 
@@ -20,9 +21,9 @@
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
 from DIRAC.Core.Utilities.TimeUtilities import fromString, second, toEpoch
 from DIRAC.WorkloadManagementSystem.Client import JobMinorStatus, JobStatus
+from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import getJobParameters
 from DIRAC.WorkloadManagementSystem.DB.StatusUtils import kill_delete_jobs
 from DIRAC.WorkloadManagementSystem.Service.JobPolicy import RIGHT_KILL
-from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
 from DIRAC.WorkloadManagementSystem.Utilities.Utils import rescheduleJobs
 
 

src/DIRAC/WorkloadManagementSystem/DB/JobDBUtils.py

@@ -3,19 +3,20 @@
 # Import stateless functions from DIRACCommon for backward compatibility
 from DIRACCommon.WorkloadManagementSystem.DB.JobDBUtils import *
 
+from DIRAC import S_OK
 from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
 from DIRAC.Core.Utilities.ReturnValues import returnValueOrRaise
 from DIRAC.WorkloadManagementSystem.Client.JobState.JobManifest import makeJobManifestConfig
+from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
+from DIRAC.WorkloadManagementSystem.DB.JobParametersDB import JobParametersDB
 
 getDIRACPlatform = returnValueOrRaise(
     ObjectLoader().loadObject("ConfigurationSystem.Client.Helpers.Resources", "getDIRACPlatform")
 )
 
 
-def checkAndPrepareJob(
-    jobID, classAdJob, classAdReq, owner, ownerGroup, jobAttrs, vo
-):  # pylint: disable=function-redefined
+def checkAndPrepareJob(jobID, classAdJob, classAdReq, owner, ownerGroup, jobAttrs, vo):  # pylint: disable=function-redefined
     from DIRACCommon.WorkloadManagementSystem.DB.JobDBUtils import checkAndPrepareJob
 
     config = {
@@ -31,3 +32,63 @@ def checkAndAddOwner(jdl: str, owner: str, ownerGroup: str):  # pylint: disable=
     from DIRACCommon.WorkloadManagementSystem.DB.JobDBUtils import checkAndAddOwner
 
     return checkAndAddOwner(jdl, owner, ownerGroup, job_manifest_config=makeJobManifestConfig(ownerGroup))
+
+
+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)

src/DIRAC/WorkloadManagementSystem/Service/JobMonitoringHandler.py

@@ -1,7 +1,7 @@
-""" JobMonitoringHandler is the implementation of the JobMonitoring service
-    in the DISET framework
+"""JobMonitoringHandler is the implementation of the JobMonitoring service
+in the DISET framework
 
-    The following methods are available in the Service interface
+The following methods are available in the Service interface
 """
 
 from DIRAC import S_ERROR, S_OK
@@ -10,7 +10,7 @@
 from DIRAC.Core.Utilities.DEncode import ignoreEncodeWarning
 from DIRAC.Core.Utilities.JEncode import strToIntDict
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
-from DIRAC.WorkloadManagementSystem.Utilities.JobParameters import getJobParameters
+from DIRAC.WorkloadManagementSystem.DB.JobDBUtils import getJobParameters
 
 
 class JobMonitoringHandlerMixin:

src/DIRAC/WorkloadManagementSystem/Utilities/JobParameters.py

@@ -141,68 +141,6 @@ def getNumberOfGPUs(siteName=None, gridCE=None, queue=None):
     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
-    """
-    from DIRAC.WorkloadManagementSystem.DB.JobParametersDB import JobParametersDB
-    from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
-
-    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)
-
-
 def getAvailableRAM(siteName=None, gridCE=None, queue=None):
     """Gets the available RAM on a certain CE/queue/node (what the pilot administers)