def _getClients(self):
""" returns the clients used in the threads - this is another function that should be extended.
The clients provided here are defaults, and should be adapted
"""
threadTransformationClient = TransformationClient()
threadTaskManager = WorkflowTasks() # this is for wms tasks, replace it with something else if needed
threadTaskManager.pluginLocation = self.pluginLocation
return {"TransformationClient": threadTransformationClient, "TaskManager": threadTaskManager}
def initialize( self ):
""" Sets defaults """
TaskManagerAgentBase.initialize( self )
WorkflowTasks.__init__( self )
self.transType = self.am_getOption( "TransType", ['MCSimulation', 'DataReconstruction', 'DataStripping', 'MCStripping', 'Merge'] )
# This sets the Default Proxy to used as that defined under
# /Operations/Shifter/ProductionManager
# the shifterProxy option in the Configuration can be used to change this default.
self.am_setOption( 'shifterProxy', 'ProductionManager' )
return S_OK()
def _getClients(self):
""" returns the clients used in the threads - this is another function that should be extended.
The clients provided here are defaults, and should be adapted
"""
threadTransformationClient = TransformationClient()
threadTaskManager = WorkflowTasks(
) # this is for wms tasks, replace it with something else if needed
threadTaskManager.pluginLocation = self.pluginLocation
return {
'TransformationClient': threadTransformationClient,
'TaskManager': threadTaskManager
}
def initialize(self):
""" Sets defaults """
TaskManagerAgentBase.initialize(self)
WorkflowTasks.__init__(self)
self.transType = self.am_getOption("TransType", [
'MCSimulation', 'DataReconstruction', 'DataStripping',
'MCStripping', 'Merge'
])
# This sets the Default Proxy to used as that defined under
# /Operations/Shifter/ProductionManager
# the shifterProxy option in the Configuration can be used to change this default.
self.am_setOption('shifterProxy', 'ProductionManager')
return S_OK()
def setUp( self ):
self.mockTransClient = Mock()
self.mockTransClient.setTaskStatusAndWmsID.return_value = {'OK':True}
self.WMSClientMock = Mock()
self.jobMonitoringClient = Mock()
self.mockReqClient = Mock()
self.jobMock = Mock()
self.jobMock2 = Mock()
mockWF = Mock()
mockPar = Mock()
mockWF.findParameter.return_value = mockPar
mockPar.getValue.return_value = 'MySite'
self.jobMock2.workflow = mockWF
self.jobMock2.setDestination.return_value = {'OK':True}
self.jobMock.workflow.return_value = ''
self.jobMock.return_value = self.jobMock2
self.taskBase = TaskBase( transClient = self.mockTransClient )
self.wfTasks = WorkflowTasks( transClient = self.mockTransClient,
submissionClient = self.WMSClientMock,
jobMonitoringClient = self.jobMonitoringClient,
outputDataModule = "mock",
jobClass = self.jobMock )
self.requestTasks = RequestTasks( transClient = self.mockTransClient,
requestClient = self.mockReqClient
)
self.tc = TransformationClient()
self.transformation = Transformation()
self.maxDiff = None
def _getClients(self, ownerDN=None, ownerGroup=None):
"""Returns the clients used in the threads
This is another function that should be extended.
The clients provided here are defaults, and should be adapted
If ownerDN and ownerGroup are not None the clients will delegate to these credentials
:param str ownerDN: DN of the owner of the submitted jobs
:param str ownerGroup: group of the owner of the submitted jobs
:returns: dict of Clients
"""
threadTransformationClient = TransformationClient()
threadTaskManager = WorkflowTasks(ownerDN=ownerDN, ownerGroup=ownerGroup)
threadTaskManager.pluginLocation = self.pluginLocation
return {'TransformationClient': threadTransformationClient,
'TaskManager': threadTaskManager}
def _getClients(self, ownerDN=None, ownerGroup=None):
"""Returns the clients used in the threads
This is another function that should be extended.
The clients provided here are defaults, and should be adapted
If ownerDN and ownerGroup are not None the clients will delegate to these credentials
:param str ownerDN: DN of the owner of the submitted jobs
:param str ownerGroup: group of the owner of the submitted jobs
:returns: dict of Clients
"""
threadTransformationClient = TransformationClient()
threadTaskManager = WorkflowTasks(ownerDN=ownerDN, ownerGroup=ownerGroup)
threadTaskManager.pluginLocation = self.pluginLocation
return {'TransformationClient': threadTransformationClient,
'TaskManager': threadTaskManager}
def __init__(self, *args, **kwargs):
''' c'tor
'''
TaskManagerAgentBase.__init__(self, *args, **kwargs)
self.submissionClient = WMSClient()
self.taskManager = WorkflowTasks(
transClient=self.transClient,
submissionClient=self.submissionClient)
self.shifterProxy = 'ProductionManager'
agentTSTypes = self.am_getOption('TransType', [])
if agentTSTypes:
self.transType = agentTSTypes
else:
self.transType = Operations().getValue(
'Transformations/DataProcessing', ['MCSimulation', 'Merge'])
def setUp(self):
self.mockTransClient = MagicMock()
self.mockTransClient.setTaskStatusAndWmsID.return_value = {'OK': True}
self.WMSClientMock = MagicMock()
self.jobMonitoringClient = MagicMock()
self.mockReqClient = MagicMock()
self.jobMock = MagicMock()
self.jobMock2 = MagicMock()
mockWF = MagicMock()
mockPar = MagicMock()
mockWF.findParameter.return_value = mockPar
mockPar.getValue.return_value = 'MySite'
self.jobMock2.workflow = mockWF
self.jobMock2.setDestination.return_value = {'OK': True}
self.jobMock.workflow.return_value = ''
self.jobMock.return_value = self.jobMock2
self.reqValidatorMock = MagicMock()
self.reqValidatorMock.validate.return_value = {'OK': True}
self.taskBase = TaskBase(transClient=self.mockTransClient)
self.pu = PluginUtilities(transClient=self.mockTransClient)
self.wfTasks = WorkflowTasks(
transClient=self.mockTransClient,
submissionClient=self.WMSClientMock,
jobMonitoringClient=self.jobMonitoringClient,
outputDataModule="mock")
self.requestTasks = RequestTasks(transClient=self.mockTransClient,
requestClient=self.mockReqClient,
requestValidator=reqValFake)
self.tc = TransformationClient()
self.transformation = Transformation()
self.maxDiff = None
gLogger.setLevel('DEBUG')
def setUp(self):
self.mockTransClient = Mock()
self.mockTransClient.setTaskStatusAndWmsID.return_value = {'OK': True}
self.WMSClientMock = Mock()
self.jobMonitoringClient = Mock()
self.mockRequestClient = Mock()
self.jobMock = Mock()
self.jobMock.setDestination.return_value = {'OK': True}
self.taskBase = TaskBase(transClient=self.mockTransClient)
self.wfTasks = WorkflowTasks(
transClient=self.mockTransClient,
submissionClient=self.WMSClientMock,
jobMonitoringClient=self.jobMonitoringClient,
outputDataModule="mock",
jobClass=self.jobMock)
self.requestTasks = RequestTasks(transClient=self.mockTransClient,
requestClient=self.mockRequestClient)
from DIRAC.Interfaces.API.Job import Job
# sut
from DIRAC.TransformationSystem.Client.TaskManager import WorkflowTasks
mockTransClient = MagicMock()
mockTransClient.setTaskStatusAndWmsID.return_value = {'OK': True}
WMSClientMock = MagicMock()
jobMonitoringClient = MagicMock()
wfTasks = WorkflowTasks(transClient=mockTransClient,
submissionClient=WMSClientMock,
jobMonitoringClient=jobMonitoringClient,
outputDataModule="mock")
odm_o = MagicMock()
odm_o.execute.return_value = {'OK': True, 'Value': {}}
wfTasks.outputDataModule_o = odm_o
taskDict = {1: {'TransformationID': 1, 'a1': 'aa1', 'b1': 'bb1', 'Site': 'MySite'},
2: {'TransformationID': 1, 'a2': 'aa2', 'b2': 'bb2', 'InputData': ['a1', 'a2']},
3: {'TransformationID': 2, 'a3': 'aa3', 'b3': 'bb3'}, }
taskDictNoInputs = {1: {'TransformationID': 1, 'a1': 'aa1', 'b1': 'bb1', 'Site': 'MySite'},
2: {'TransformationID': 1, 'a2': 'aa2', 'b2': 'bb2'},
3: {'TransformationID': 2, 'a3': 'aa3', 'b3': 'bb3'}, }
expected = {'OK': True,
'Value': {1: {'a1': 'aa1', 'TaskObject': '', 'TransformationID': 1,
from mock import MagicMock
import pytest
from DIRAC.Interfaces.API.Job import Job
# sut
from DIRAC.TransformationSystem.Client.TaskManager import WorkflowTasks
mockTransClient = MagicMock()
mockTransClient.setTaskStatusAndWmsID.return_value = {'OK': True}
WMSClientMock = MagicMock()
jobMonitoringClient = MagicMock()
wfTasks = WorkflowTasks(transClient=mockTransClient,
submissionClient=WMSClientMock,
jobMonitoringClient=jobMonitoringClient,
outputDataModule="mock")
odm_o = MagicMock()
odm_o.execute.return_value = {'OK': True, 'Value': {}}
wfTasks.outputDataModule_o = odm_o
taskDict = {
1: {
'TransformationID': 1,
'a1': 'aa1',
'b1': 'bb1',
'Site': 'MySite'
},
2: {
'TransformationID': 1,
'a2': 'aa2',
def initialize( self ): self.section = self.am_getOption( "section" ) gMonitor.registerActivity( "SubmittedTasks", "Automatically submitted tasks", "Transformation Monitoring", "Tasks", gMonitor.OP_ACUM ) self.transClient = TransformationClient() self.wfTasks = WorkflowTasks() return S_OK()
class TaskManagerAgentBase( AgentModule ):
#############################################################################
def initialize( self ):
self.section = self.am_getOption( "section" )
gMonitor.registerActivity( "SubmittedTasks", "Automatically submitted tasks", "Transformation Monitoring", "Tasks", gMonitor.OP_ACUM )
self.transClient = TransformationClient()
self.wfTasks = WorkflowTasks()
return S_OK()
#############################################################################
def execute( self ):
""" The TaskManagerBase execution method. """
# Determine whether the task status is to be monitored and updated
enableTaskMonitor = self.am_getOption( 'MonitorTasks', '' )
if not enableTaskMonitor:
gLogger.info( "execute: Monitoring of tasks is disabled." )
gLogger.info( "execute: To enable create the 'MonitorTasks' option" )
else:
res = self.updateTaskStatus()
if not res['OK']:
gLogger.warn( 'execute: Failed to update task states', res['Message'] )
# Determine whether the task files status is to be monitored and updated
enableFileMonitor = self.am_getOption( 'MonitorFiles', '' )
if not enableFileMonitor:
gLogger.info( "execute: Monitoring of files is disabled." )
gLogger.info( "execute: To enable create the 'MonitorFiles' option" )
else:
res = self.updateFileStatus()
if not res['OK']:
gLogger.warn( 'execute: Failed to update file states', res['Message'] )
# Determine whether the checking of reserved tasks is to be performed
enableCheckReserved = self.am_getOption( 'CheckReserved', '' )
if not enableCheckReserved:
gLogger.info( "execute: Checking of reserved tasks is disabled." )
gLogger.info( "execute: To enable create the 'CheckReserved' option" )
else:
res = self.checkReservedTasks()
if not res['OK']:
gLogger.warn( 'execute: Failed to checked reserved tasks', res['Message'] )
# Determine whether the submission of tasks is to be executed
enableSubmission = self.am_getOption( 'SubmitTasks', '' )
if not enableSubmission:
gLogger.info( "execute: Submission of tasks is disabled." )
gLogger.info( "execute: To enable create the 'SubmitTasks' option" )
else:
res = self.submitTasks()
if not res['OK']:
gLogger.warn( 'execute: Failed to submit created tasks', res['Message'] )
return S_OK()
def _selectTransformations( self, transType = [], status = ['Active', 'Completing'], agentType = ['Automatic'] ):
selectCond = {}
if status:
selectCond['Status'] = status
if transType:
selectCond['Type'] = transType
if agentType:
selectCond['AgentType'] = agentType
res = self.transClient.getTransformations( condDict = selectCond )
if not res['OK']:
gLogger.error( "_selectTransformations: Failed to get transformations for selection.", res['Message'] )
elif not res['Value']:
gLogger.info( "_selectTransformations: No transformations found for selection." )
else:
gLogger.info( "_selectTransformations: Obtained %d transformations for selection" % len( res['Value'] ) )
return res
def updateTaskStatus( self ):
gLogger.info( "updateTaskStatus: Updating the Status of tasks" )
# Get the transformations to be updated
status = self.am_getOption( 'UpdateTasksStatus', ['Active', 'Completing', 'Stopped'] )
res = self._selectTransformations( transType = self.transType, status = status, agentType = [] )
if not res['OK']:
return res
for transformation in res['Value']:
transID = transformation['TransformationID']
# Get the tasks which are in a UPDATE state
updateStatus = self.am_getOption( 'TaskUpdateStatus', ['Checking', 'Deleted', 'Killed', 'Staging', 'Stalled', 'Matched', 'Rescheduled', 'Completed', 'Submitted', 'Received', 'Waiting', 'Running'] )
condDict = {"TransformationID":transID, "ExternalStatus":updateStatus}
timeStamp = str( datetime.datetime.utcnow() - datetime.timedelta( minutes = 10 ) )
res = self.transClient.getTransformationTasks( condDict = condDict, older = timeStamp, timeStamp = 'LastUpdateTime' )
if not res['OK']:
gLogger.error( "updateTaskStatus: Failed to get tasks to update for transformation", "%s %s" % ( transID, res['Message'] ) )
continue
if not res['Value']:
gLogger.verbose( "updateTaskStatus: No tasks found to update for transformation %s" % transID )
continue
res = self.getSubmittedTaskStatus( res['Value'] )
if not res['OK']:
gLogger.error( "updateTaskStatus: Failed to get updated task statuses for transformation", "%s %s" % ( transID, res['Message'] ) )
continue
statusDict = res['Value']
for status in sortList( statusDict.keys() ):
taskIDs = statusDict[status]
gLogger.info( "updateTaskStatus: Updating %d task(s) from transformation %d to %s" % ( len( taskIDs ), transID, status ) )
res = self.transClient.setTaskStatus( transID, taskIDs, status )
if not res['OK']:
gLogger.error( "updateTaskStatus: Failed to update task status for transformation", "%s %s" % ( transID, res['Message'] ) )
gLogger.info( "updateTaskStatus: Transformation task status update complete" )
return S_OK()
def updateFileStatus( self ):
gLogger.info( "updateFileStatus: Updating Status of task files" )
#Get the transformations to be updated
status = self.am_getOption( 'UpdateFilesStatus', ['Active', 'Completing', 'Stopped'] )
res = self._selectTransformations( transType = self.transType, status = status, agentType = [] )
if not res['OK']:
return res
for transformation in res['Value']:
transID = transformation['TransformationID']
# Get the files which are in a UPDATE state
updateStatus = self.am_getOption( 'FileUpdateStatus', ['Submitted', 'Received', 'Waiting', 'Running'] )
timeStamp = str( datetime.datetime.utcnow() - datetime.timedelta( minutes = 10 ) )
condDict = {'TransformationID' : transID, 'Status' : ['Assigned']}
res = self.transClient.getTransformationFiles( condDict = condDict, older = timeStamp, timeStamp = 'LastUpdate' )
if not res['OK']:
gLogger.error( "updateFileStatus: Failed to get transformation files to update.", res['Message'] )
continue
if not res['Value']:
gLogger.info( "updateFileStatus: No files to be updated for transformation %s." % transID )
continue
res = self.getSubmittedFileStatus( res['Value'] )
if not res['OK']:
gLogger.error( "updateFileStatus: Failed to get updated file statuses for transformation", "%s %s" % ( transID, res['Message'] ) )
continue
statusDict = res['Value']
if not statusDict:
gLogger.info( "updateFileStatus: No file statuses to be updated for transformation %s." % transID )
continue
fileReport = FileReport( server = self.transClient.getServer() )
for lfn, status in statusDict.items():
fileReport.setFileStatus( int( transID ), lfn, status )
res = fileReport.commit()
if not res['OK']:
gLogger.error( "updateFileStatus: Failed to update file status for transformation", "%s %s" % ( transID, res['Message'] ) )
else:
for status, update in res['Value'].items():
gLogger.info( "updateFileStatus: Updated %s files for %s to %s." % ( update, transID, status ) )
gLogger.info( "updateFileStatus: Transformation file status update complete" )
return S_OK()
def checkReservedTasks( self ):
gLogger.info( "checkReservedTasks: Checking Reserved tasks" )
# Get the transformations which should be checked
status = self.am_getOption( 'CheckReservedStatus', ['Active', 'Completing', 'Stopped'] )
res = self._selectTransformations( transType = self.transType, status = status, agentType = [] )
if not res['OK']:
return res
for transformation in res['Value']:
transID = transformation['TransformationID']
# Select the tasks which have been in Reserved status for more than 1 hour for selected transformations
condDict = {"TransformationID":transID, "ExternalStatus":'Reserved'}
time_stamp_older = str( datetime.datetime.utcnow() - datetime.timedelta( hours = 1 ) )
time_stamp_newer = str( datetime.datetime.utcnow() - datetime.timedelta( days = 7 ) )
res = self.transClient.getTransformationTasks( condDict = condDict, older = time_stamp_older, newer = time_stamp_newer, timeStamp = 'LastUpdateTime' )
if not res['OK']:
gLogger.error( "checkReservedTasks: Failed to get Reserved tasks for transformation", "%s %s" % ( transID, res['Message'] ) )
continue
if not res['Value']:
gLogger.verbose( "checkReservedTasks: No Reserved tasks found for transformation %s" % transID )
continue
res = self.updateTransformationReservedTasks( res['Value'] )
if not res['OK']:
gLogger.info( "checkReservedTasks: No Reserved tasks found for transformation %s" % transID )
continue
noTasks = res['Value']['NoTasks']
taskNameIDs = res['Value']['TaskNameIDs']
# For the tasks with no associated request found re-set the status of the task in the transformationDB
for taskName in noTasks:
transID, taskID = taskName.split( '_' )
gLogger.info( "checkReservedTasks: Resetting status of %s to Created as no associated task found" % ( taskName ) )
res = self.transClient.setTaskStatus( int( transID ), int( taskID ), 'Created' )
if not res['OK']:
gLogger.warn( "checkReservedTasks: Failed to update task status and ID after recovery", "%s %s" % ( taskName, res['Message'] ) )
# For the tasks for which an associated request was found update the task details in the transformationDB
for taskName, extTaskID in taskNameIDs.items():
transID, taskID = taskName.split( '_' )
gLogger.info( "checkReservedTasks: Resetting status of %s to Created with ID %s" % ( taskName, extTaskID ) )
res = self.transClient.setTaskStatusAndWmsID( int( transID ), int( taskID ), 'Submitted', str( extTaskID ) )
if not res['OK']:
gLogger.warn( "checkReservedTasks: Failed to update task status and ID after recovery", "%s %s" % ( taskName, res['Message'] ) )
gLogger.info( "checkReservedTasks: Updating of reserved tasks complete" )
return S_OK()
def submitTasks( self ):
gLogger.info( "submitTasks: Submitting tasks for transformations" )
res = getProxyInfo( False, False )
if not res['OK']:
gLogger.error( "submitTasks: Failed to determine credentials for submission", res['Message'] )
return res
proxyInfo = res['Value']
owner = proxyInfo['username']
ownerGroup = proxyInfo['group']
gLogger.info( "submitTasks: Tasks will be submitted with the credentials %s:%s" % ( owner, ownerGroup ) )
# Get the transformations which should be submitted
tasksPerLoop = self.am_getOption( 'TasksPerLoop', 50 )
status = self.am_getOption( 'SubmitStatus', ['Active', 'Completing'] )
res = self._selectTransformations( transType = self.transType, status = status )
if not res['OK']:
return res
for transformation in res['Value']:
transID = transformation['TransformationID']
transBody = transformation['Body']
res = self.transClient.getTasksToSubmit( transID, tasksPerLoop )
if not res['OK']:
gLogger.error( "submitTasks: Failed to obtain tasks for transformation", "%s %s" % ( transID, res['Message'] ) )
continue
tasks = res['Value']['JobDictionary']
if not tasks:
gLogger.verbose( "submitTasks: No tasks found for submission for transformation %s" % transID )
continue
gLogger.info( "submitTasks: Obtained %d tasks for submission for transformation %s" % ( len( tasks ), transID ) )
res = self.wfTasks.prepareTransformationTasks( transBody, tasks, owner, ownerGroup )
if not res['OK']:
gLogger.error( "submitTasks: Failed to prepare tasks for transformation", "%s %s" % ( transID, res['Message'] ) )
continue
res = self.wfTasks.submitTransformationTasks( res['Value'] )
if not res['OK']:
gLogger.error( "submitTasks: Failed to submit prepared tasks for transformation", "%s %s" % ( transID, res['Message'] ) )
continue
res = self.wfTasks.updateDBAfterTaskSubmission( res['Value'] )
if not res['OK']:
gLogger.error( "submitTasks: Failed to update DB after task submission for transformation", "%s %s" % ( transID, res['Message'] ) )
continue
gLogger.info( "submitTasks: Submission of transformation tasks complete" )
return S_OK()