Skip to content

Commit

Permalink
feat: introduce addPilotReferences() method
Browse files Browse the repository at this point in the history
  • Loading branch information
aldbr committed Dec 8, 2023
1 parent 096bb08 commit 94b5c88
Show file tree
Hide file tree
Showing 5 changed files with 48 additions and 1 deletion.
7 changes: 6 additions & 1 deletion src/DIRAC/Interfaces/API/DiracAdmin.py
Original file line number Diff line number Diff line change
Expand Up @@ -553,6 +553,9 @@ def getPilotInfo(self, gridReference):
if not isinstance(gridReference, str):
return self._errorReport("Expected string for pilot reference")

# TODO: to remove from v9.0
gLogger.notice("Notice: 'TaskQueueID' will be removed from the output in v9.0.")

result = PilotManagerClient().getPilotInfo(gridReference)
return result

Expand Down Expand Up @@ -599,14 +602,16 @@ def getJobPilots(self, jobID):
:param job: JobID
:type job: integer or string
:return: S_OK,S_ERROR
"""
if isinstance(jobID, str):
try:
jobID = int(jobID)
except ValueError as x:
return self._errorReport(str(x), "Expected integer or string for existing jobID")

# TODO: remove this comment from v9.0
gLogger.notice("Notice: 'TaskQueueID' will be removed from the output in v9.0.")

result = PilotManagerClient().getPilots(jobID)
if result["OK"]:
gLogger.notice(self.pPrint.pformat(result["Value"]))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,10 @@ def execute(self):
# PilotsHistory to Monitoring
if "Monitoring" in self.pilotMonitoringOption:
self.log.info("Committing PilotsHistory to Monitoring")

# TODO: remove this comment from v9.0
self.log.notice("Notice: 'TaskQueueID' will be removed from the pilotAgentsDB in v9.0.")

result = PilotAgentsDB().getSummarySnapshot()
now = datetime.datetime.utcnow()
if not result["OK"]:
Expand Down
26 changes: 26 additions & 0 deletions src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
Available methods are:
addPilotReferences()
addPilotTQReference()
setPilotStatus()
deletePilot()
Expand All @@ -20,6 +21,7 @@
import datetime
import decimal
import threading
from DIRAC.Core.Utilities.Decorators import deprecated

import DIRAC.Core.Utilities.TimeUtilities as TimeUtilities
from DIRAC import S_ERROR, S_OK
Expand All @@ -38,6 +40,30 @@ def __init__(self, parentLogger=None):
self.lock = threading.Lock()

##########################################################################################
def addPilotReferences(self, pilotRef, ownerGroup, gridType="DIRAC", pilotStampDict={}):
"""Add a new pilot job reference"""
for ref in pilotRef:
stamp = ""
if ref in pilotStampDict:
stamp = pilotStampDict[ref]

req = (
"INSERT INTO PilotAgents "
+ "(PilotJobReference, OwnerGroup, GridType, SubmissionTime, LastUpdateTime, Status, PilotStamp) "
+ "VALUES ('%s','%s','%s',UTC_TIMESTAMP(),UTC_TIMESTAMP(),'Submitted','%s')"
% (ref, ownerGroup, gridType, stamp)
)

result = self._update(req)
if not result["OK"]:
return result

if "lastRowId" not in result:
return S_ERROR("PilotAgentsDB.addPilotReferences: Failed to retrieve a new Id.")

return S_OK()

@deprecated("Use addPilotReferences instead")
def addPilotTQReference(
self, pilotRef, taskQueueID, ownerDN, ownerGroup, broker="Unknown", gridType="DIRAC", pilotStampDict={}
):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import datetime

from DIRAC import S_OK, S_ERROR
from DIRAC.Core.Utilities.Decorators import deprecated
import DIRAC.Core.Utilities.TimeUtilities as TimeUtilities

from DIRAC.Core.DISET.RequestHandler import RequestHandler
Expand Down Expand Up @@ -80,8 +81,16 @@ def export_getCurrentPilotCounters(cls, attrDict={}):
return S_OK(resultDict)

##########################################################################################
types_addPilotReferences = [list, str]

@classmethod
def export_addPilotReferences(cls, pilotRef, ownerGroup, gridType="DIRAC", pilotStampDict={}):
"""Add a new pilot job reference"""
return cls.pilotAgentsDB.addPilotReferences(pilotRef, ownerGroup, gridType, pilotStampDict)

types_addPilotTQReference = [list, int, str, str]

@deprecated("Use addPilotReferences instead")
@classmethod
def export_addPilotTQReference(
cls, pilotRef, taskQueueID, ownerDN, ownerGroup, broker="Unknown", gridType="DIRAC", pilotStampDict={}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,9 @@ def setTaskQueueID(self, value):
:return: S_OK()/S_ERROR()
"""
# TODO: remove this comment from v9.0
gLogger.notice("Notice: 'TaskQueueID' will be removed from the pilotAgentsDB in v9.0.")

try:
self.taskQueueID = int(value)
except ValueError:
Expand Down

0 comments on commit 94b5c88

Please sign in to comment.