def execute_dag(self, dag, workflow_id, data=None): """ Celery task that runs a single dag on a worker. This celery task starts, manages and monitors the individual tasks of a dag. Args: self (Task): Reference to itself, the celery task object. dag (Dag): Reference to a Dag object that is being used to start, manage and monitor tasks. workflow_id (string): The unique ID of the workflow run that started this dag. data (MultiTaskData): An optional MultiTaskData object that is being passed to the first tasks in the dag. This allows the transfer of data from dag to dag. """ start_time = datetime.now() logger.info('Running DAG <{}>'.format(dag.name)) # send custom celery event that the dag has been started self.send_event(JobEventName.Started, job_type=JobType.Dag, name=dag.name, time=datetime.utcnow(), workflow_id=workflow_id, duration=None) # store job specific meta information wth the job self.update_state(meta={ 'name': dag.name, 'type': JobType.Dag, 'workflow_id': workflow_id }) # run the tasks in the DAG signal = DagSignal( Client(SignalConnection(**self.app.user_options['config'].signal, auto_connect=True), request_key=workflow_id), dag.name) dag.run(config=self.app.user_options['config'], workflow_id=workflow_id, signal=signal, data=data) # send custom celery event that the dag has succeeded event_name = JobEventName.Succeeded if not signal.is_stopped else JobEventName.Aborted self.send_event(event_name, job_type=JobType.Dag, name=dag.name, time=datetime.utcnow(), workflow_id=workflow_id, duration=(datetime.now() - start_time).total_seconds()) logger.info('Finished DAG <{}>'.format(dag.name))
def stop(self, consumer): """ This function is called when the worker received a request to terminate. Upon the termination of the worker, the workflows for all running jobs are stopped gracefully. Args: consumer (Consumer): Reference to the consumer object that handles messages from the broker. """ stopped_workflows = [] for request in [r for r in consumer.controller.state.active_requests]: job = AsyncResult(request.id) workflow_id = job.result['workflow_id'] if workflow_id not in stopped_workflows: client = Client(SignalConnection( **consumer.app.user_options['config'].signal, auto_connect=True), request_key=workflow_id) client.send(Request(action='stop_workflow')) stopped_workflows.append(workflow_id)
def execute_workflow(self, workflow, workflow_id=None): """ Celery task (aka job) that runs a workflow on a worker. This celery task starts, manages and monitors the dags that make up a workflow. Args: self (Task): Reference to itself, the celery task object. workflow (Workflow): Reference to the workflow object that is being used to start, manage and monitor dags. workflow_id (string): If a workflow ID is provided the workflow run will use this ID, if not a new ID will be auto generated. """ start_time = datetime.now() logger.info('Running workflow <{}>'.format(workflow.name)) data_store = DataStore(**self.app.user_options['config'].data_store, auto_connect=True) # create a unique workflow id for this run if data_store.exists(workflow_id): logger.info('Using existing workflow ID: {}'.format(workflow_id)) else: workflow_id = data_store.add(payload={ 'name': workflow.name, 'start_time': datetime.utcnow() }) logger.info('Created workflow ID: {}'.format(workflow_id)) # send custom celery event that the workflow has been started self.send_event(JobEventName.Started, job_type=JobType.Workflow, name=workflow.name, time=datetime.utcnow(), workflow_id=workflow_id, duration=None) # create server for inter-task messaging signal_server = Server( SignalConnection(**self.app.user_options['config'].signal, auto_connect=True), request_key=workflow_id) # store job specific meta information wth the job self.update_state( meta={ 'name': workflow.name, 'type': JobType.Workflow, 'workflow_id': workflow_id, 'arguments': workflow.provided_arguments }) # run the DAGs in the workflow workflow.run(config=self.app.user_options['config'], data_store=data_store, signal_server=signal_server, workflow_id=workflow_id) # send custom celery event that the workflow has succeeded event_name = JobEventName.Succeeded if not workflow.is_stopped \ else JobEventName.Aborted self.send_event(event_name, job_type=JobType.Workflow, name=workflow.name, time=datetime.utcnow(), workflow_id=workflow_id, duration=(datetime.now() - start_time).total_seconds()) logger.info('Finished workflow <{}>'.format(workflow.name))
def execute_task(self, task, workflow_id, data=None): """ Celery task that runs a single task on a worker. Args: self (Task): Reference to itself, the celery task object. task (BaseTask): Reference to the task object that performs the work in its run() method. workflow_id (string): The unique ID of the workflow run that started this task. data (MultiTaskData): An optional MultiTaskData object that contains the data that has been passed down from upstream tasks. """ start_time = datetime.now() store_doc = DataStore(**self.app.user_options['config'].data_store, auto_connect=True).get(workflow_id) def handle_callback(message, event_type, exc=None): msg = '{}: {}'.format(message, str(exc)) if exc is not None else message # set the logging level if event_type == JobEventName.Stopped: logger.warning(msg) elif event_type == JobEventName.Aborted: logger.error(msg) else: logger.info(msg) # store a log into the persistent data store if event_type != JobEventName.Started: duration = (datetime.now() - start_time).total_seconds() store_doc.set(key='log.{}.{}.duration'.format( task.dag_name, task.name), value=duration, section=DataStoreDocumentSection.Meta) store_doc.set(key='log.{}.{}.worker'.format( task.dag_name, task.name), value=self.request.hostname, section=DataStoreDocumentSection.Meta) else: duration = None # send custom celery event self.send_event(event_type, job_type=JobType.Task, name=task.name, time=datetime.utcnow(), workflow_id=workflow_id, duration=duration) # store job specific meta information wth the job self.update_state(meta={ 'name': task.name, 'type': JobType.Task, 'workflow_id': workflow_id }) # send start celery event handle_callback('Start task <{}>'.format(task.name), JobEventName.Started) # run the task and capture the result return task._run( data=data, store=store_doc, signal=TaskSignal( Client(SignalConnection(**self.app.user_options['config'].signal, auto_connect=True), request_key=workflow_id), task.dag_name), context=TaskContext(task.name, task.dag_name, task.workflow_name, workflow_id, self.request.hostname), success_callback=partial(handle_callback, message='Complete task <{}>'.format( task.name), event_type=JobEventName.Succeeded), stop_callback=partial(handle_callback, message='Stop task <{}>'.format(task.name), event_type=JobEventName.Stopped), abort_callback=partial( handle_callback, message='Abort workflow <{}> by task <{}>'.format( task.workflow_name, task.name), event_type=JobEventName.Aborted))