def stop_job_execution( self, job_name: Text, workflow_execution_id: Text) -> Optional[JobExecutionInfo]: with create_session() as session: dagrun = session.query(DagRun).filter( DagRun.id == int(workflow_execution_id)).first() if dagrun is None: return None task = dagrun.get_task_instance(job_name, session) if task is None: return None if task.state in State.finished: raise Exception('job:{} state: {} can not stop!'.format( job_name, task.state)) else: self.airflow_client.schedule_task( dag_id=dagrun.dag_id, task_id=job_name, action=SchedulingAction.STOP, context=ExecutionContext(dagrun_id=dagrun.run_id)) project_name, workflow_name = self.dag_id_to_namespace_workflow( dagrun.dag_id) return JobExecutionInfo( job_name=job_name, status=self.airflow_state_to_status(task.state), workflow_execution=WorkflowExecutionInfo( workflow_info=WorkflowInfo(namespace=project_name, workflow_name=workflow_name), workflow_execution_id=workflow_execution_id, status=self.airflow_state_to_status(dagrun.state)))
def start_job_execution(self, job_name: Text, execution_id: Text) -> JobExecutionInfo: job = self.workflow.jobs[job_name] plugins = self.workflow.properties.get( WorkflowPropertyKeys.JOB_PLUGINS) module, name = plugins.get(job.job_config.job_type) class_object = import_string('{}.{}'.format(module, name)) self.job_controller: JobController = class_object() job_execution_info = JobExecutionInfo( job_execution_id='1', job_name=job_name, workflow_execution=WorkflowExecutionInfo( workflow_execution_id='1', workflow_info=WorkflowInfo( workflow_name=self.workflow.workflow_name))) self.job_runtime_env: JobRuntimeEnv = prepare_job_runtime_env( root_working_dir=self.project_context.project_path + '/temp', workflow_snapshot_id=self.workflow.workflow_snapshot_id, workflow_name=self.workflow.workflow_name, job_execution_info=job_execution_info, project_context=self.project_context) self.job_handler = self.job_controller.submit_job( job=job, job_runtime_env=self.job_runtime_env) return job_execution_info
def list_workflow_executions( self, project_name: Text, workflow_name: Text) -> List[WorkflowExecutionInfo]: dag_id = self.airflow_dag_id(project_name, workflow_name) with create_session() as session: dagrun_list = session.query(DagRun).filter( DagRun.dag_id == dag_id).all() if dagrun_list is None: return [] else: result = [] for dagrun in dagrun_list: status_ = self.airflow_state_to_status(dagrun.state) result.append( WorkflowExecutionInfo( workflow_info=WorkflowInfo( namespace=project_name, workflow_name=workflow_name), workflow_execution_id=str(dagrun.id), status=status_, start_date=str(datetime_to_int64( dagrun.start_date)), end_date=str(datetime_to_int64(dagrun.end_date)), )) return result
def submit_workflow(self, workflow: Workflow, project_context: ProjectContext) -> WorkflowInfo: dag_id = self.airflow_dag_id(project_context.project_name, workflow.workflow_name) code_text = self.dag_generator.generate( workflow=workflow, project_name=project_context.project_name) deploy_path = self.config.get('airflow_deploy_path') if deploy_path is None: raise Exception("airflow_deploy_path config not set!") if not os.path.exists(deploy_path): os.makedirs(deploy_path) airflow_file_path = os.path.join(deploy_path, dag_id + '.py') if os.path.exists(airflow_file_path): os.remove(airflow_file_path) with NamedTemporaryFile(mode='w+t', prefix=dag_id, suffix='.py', dir='/tmp', delete=False) as f: f.write(code_text) shutil.move(f.name, airflow_file_path) self.airflow_client.trigger_parse_dag(airflow_file_path) return WorkflowInfo(namespace=project_context.project_name, workflow_name=workflow.workflow_name, properties={'dag_file': airflow_file_path})
def restart_job_execution( self, job_name: Text, workflow_execution_id: Text) -> Optional[JobExecutionInfo]: with create_session() as session: dagrun = session.query(DagRun).filter( DagRun.id == int(workflow_execution_id)).first() if dagrun is None: return None if dagrun.state != State.RUNNING: raise Exception( 'execution: {} state: {} can not trigger job.'.format( workflow_execution_id, dagrun.state)) task = dagrun.get_task_instance(job_name, session) if task is None: return None self.airflow_client.schedule_task( dag_id=dagrun.dag_id, task_id=job_name, action=SchedulingAction.RESTART, context=ExecutionContext(dagrun_id=dagrun.run_id)) project_name, workflow_name = self.dag_id_to_namespace_workflow( dagrun.dag_id) return JobExecutionInfo( job_name=job_name, status=self.airflow_state_to_status(task.state), workflow_execution=WorkflowExecutionInfo( workflow_info=WorkflowInfo(namespace=project_name, workflow_name=workflow_name), workflow_execution_id=workflow_execution_id, status=self.airflow_state_to_status(dagrun.state)))
def submit_workflow(self, workflow: Workflow, project_context: ProjectContext, args: Dict = None) -> WorkflowInfo: code_text = self.dag_generator.generate( workflow=workflow, project_name=project_context.project_name) return WorkflowInfo(workflow_name=workflow.workflow_name, properties={'code': code_text})
def stop_workflow_execution( self, execution_id: Text) -> Optional[WorkflowExecutionInfo]: workflow_info = WorkflowInfo(workflow_name='workflow_name', namespace='project_name') workflow_execution_info = WorkflowExecutionInfo( workflow_execution_id='1', workflow_info=workflow_info, status=Status.KILLED, start_date=str(int(time.time() * 1000))) return workflow_execution_info
def test_resume_workflow(self): with mock.patch(SCHEDULER_CLASS) as mockScheduler: instance = mockScheduler.return_value self.server.scheduler_service._scheduler = instance instance.resume_workflow_scheduling.return_value = WorkflowInfo( workflow_name='test_workflow') client = SchedulerClient("localhost:{}".format(_PORT)) workflow = client.resume_workflow_scheduling( namespace='namespace', workflow_name='test_workflow') self.assertTrue('test_workflow', workflow.name)
def start_new_workflow_execution( self, project_name: Text, workflow_name: Text) -> Optional[WorkflowExecutionInfo]: workflow_info = WorkflowInfo(workflow_name=workflow_name, namespace=project_name) workflow_execution_info = WorkflowExecutionInfo( workflow_execution_id='1', workflow_info=workflow_info, status=Status.RUNNING, start_date=str(int(time.time() * 1000))) return workflow_execution_info
def stop_all_workflow_execution( self, project_name: Text, workflow_name: Text) -> List[WorkflowExecutionInfo]: workflow_info = WorkflowInfo(workflow_name=workflow_name, namespace=project_name) workflow_execution_info = WorkflowExecutionInfo( workflow_execution_id='1', workflow_info=workflow_info, status=Status.KILLED, start_date=str(int(time.time() * 1000))) return [workflow_execution_info]
def test_submit_workflow(self): with mock.patch(SCHEDULER_CLASS) as mockScheduler: instance = mockScheduler.return_value instance.submit_workflow.return_value = WorkflowInfo( workflow_name='test_workflow') client = SchedulerClient("localhost:{}".format(_PORT)) with self.assertRaises(Exception) as context: workflow = client.submit_workflow_to_scheduler( namespace='namespace', workflow_name='test_workflow', workflow_json='') self.assertTrue('workflow json is empty' in str(context.exception))
def delete_workflow(self, project_name: Text, workflow_name: Text) -> WorkflowInfo: dag_id = self.airflow_dag_id(project_name, workflow_name) deploy_path = self.config.get('airflow_deploy_path') if deploy_path is None: raise Exception("airflow_deploy_path config not set!") airflow_file_path = os.path.join(deploy_path, dag_id + '.py') if os.path.exists(airflow_file_path): os.remove(airflow_file_path) return WorkflowInfo(namespace=project_name, workflow_name=workflow_name, properties={'dag_file': airflow_file_path}) else: return None
def get_workflow(workflow_name: Text = None) -> Optional[WorkflowInfo]: """ Gets the :class:`~ai_flow.plugin_interface.scheduler_interface.WorkflowInfo` with the given name of workflow. :param workflow_name: The name of the workflow. :return: The :class:`~ai_flow.plugin_interface.scheduler_interface.WorkflowInfo` which contains the information about the workflow. """ workflow_meta: WorkflowMeta = get_ai_flow_client().get_workflow_by_name(project_name=current_project_config() .get_project_name(), workflow_name=workflow_name) if workflow_meta is None: return None project_meta: ProjectMeta = get_ai_flow_client().get_project_by_id(workflow_meta.project_id) return WorkflowInfo(workflow_name=workflow_meta.name, namespace=project_meta.name)
def list_job_executions(self, execution_id: Text) -> List[JobExecutionInfo]: workflow_info = WorkflowInfo(workflow_name='workflow_name', namespace='project_name') workflow_execution_info = WorkflowExecutionInfo( workflow_execution_id='1', workflow_info=workflow_info, status=Status.KILLED, start_date=str(int(time.time() * 1000))) job_execution_info: JobExecutionInfo = JobExecutionInfo( job_name='job_name', status=Status.RUNNING, workflow_execution=workflow_execution_info, start_date=str(int(time.time() * 1000))) return [job_execution_info]
def stop_job_execution(self, job_name: Text, execution_id: Text) -> JobExecutionInfo: workflow_info = WorkflowInfo(workflow_name='workflow_name', namespace='project_name') workflow_execution_info = WorkflowExecutionInfo( workflow_execution_id='1', workflow_info=workflow_info, status=Status.KILLED, start_date=str(int(time.time() * 1000))) job_execution_info: JobExecutionInfo = JobExecutionInfo( job_name=job_name, status=Status.KILLED, workflow_execution=workflow_execution_info, start_date=str(int(time.time() * 1000))) return job_execution_info
def test_init_job_runtime_context(self): working_dir = os.path.dirname(__file__) job_runtime_env = JobRuntimeEnv( working_dir=working_dir, job_execution_info=JobExecutionInfo( job_name='task_1', workflow_execution=WorkflowExecutionInfo( workflow_execution_id='1', workflow_info=WorkflowInfo(workflow_name='workflow_1')))) init_job_runtime_context(job_runtime_env) self.assertEqual('workflow_1', current_workflow_config().workflow_name) self.assertEqual( 'task_1', current_workflow_config().job_configs[current_job_name()].job_name) self.assertEqual('test_project', current_project_config().get_project_name())
def build_job_execution_info_list(self, dagrun, task_list): project_name, workflow_name = self.dag_id_to_namespace_workflow( dagrun.dag_id) result = [] for task in task_list: job = JobExecutionInfo( job_name=task.task_id, status=self.airflow_state_to_status(task.state), start_date=str(datetime_to_int64(task.start_date)), end_date=str(datetime_to_int64(task.end_date)), workflow_execution=WorkflowExecutionInfo( workflow_info=WorkflowInfo(namespace=project_name, workflow_name=workflow_name), workflow_execution_id=str(dagrun.id), status=self.airflow_state_to_status(dagrun.state))) result.append(job) return result
def list_workflows(page_size: int, offset: int) -> Optional[List[WorkflowInfo]]: """ Lists the :class:`~ai_flow.plugin_interface.scheduler_interface.WorkflowInfo` with the given pagination. :param page_size: The limitation of the listed workflow pagination. :param offset: The offset of the listed workflow pagination. :return: Pagination list of the :class:`~ai_flow.plugin_interface.scheduler_interface.WorkflowInfo`. """ project_name = current_project_config().get_project_name() workflow_list = get_ai_flow_client().list_workflows(project_name=project_name, page_size=page_size, offset=offset) if workflow_list is None: return None workflow_info_list = [] for workflow_meta in workflow_list: workflow_info_list.append(WorkflowInfo(namespace=project_name, workflow_name=workflow_meta.name)) return workflow_info_list
def get_workflow_execution( self, workflow_execution_id: Text) -> Optional[WorkflowExecutionInfo]: with create_session() as session: dagrun = session.query(DagRun).filter( DagRun.id == int(workflow_execution_id)).first() if dagrun is None: return None else: status_ = self.airflow_state_to_status(dagrun.state) project_name, workflow_name = self.dag_id_to_namespace_workflow( dagrun.dag_id) return WorkflowExecutionInfo( workflow_info=WorkflowInfo(namespace=project_name, workflow_name=workflow_name), workflow_execution_id=workflow_execution_id, status=status_, start_date=str(datetime_to_int64(dagrun.start_date)), end_date=str(datetime_to_int64(dagrun.end_date)))
def stop_workflow_execution( self, workflow_execution_id: Text) -> Optional[WorkflowExecutionInfo]: with create_session() as session: dagrun = session.query(DagRun).filter( DagRun.id == int(workflow_execution_id)).first() if dagrun is None: return None project_name, workflow_name = self.dag_id_to_namespace_workflow( dagrun.dag_id) context: ExecutionContext = ExecutionContext( dagrun_id=dagrun.run_id) current_context = self.airflow_client.stop_dag_run( dagrun.dag_id, context) return WorkflowExecutionInfo( workflow_info=WorkflowInfo(namespace=project_name, workflow_name=workflow_name), workflow_execution_id=workflow_execution_id, status=status.Status.KILLED)
def context_to_job_info(self, project_name: Text, context: Any) -> JobExecutionInfo: """ The function of this function is to turn the context of airflow into execution information of a job. """ wi = WorkflowInfo(namespace=project_name, workflow_name=self.workflow.workflow_name) we = WorkflowExecutionInfo( workflow_execution_id=str(context.get('dag_run').id), workflow_info=wi, start_date=str(datetime_to_int64( context.get('dag_run').start_date)), end_date=str(datetime_to_int64(context.get('dag_run').end_date)), status=context.get('dag_run').get_state()) je = JobExecutionInfo( job_name=self.job.job_name, job_execution_id=str(context.get('ti').try_number), status=context.get('ti').state, start_date=str(datetime_to_int64(context.get('ti').start_date)), end_date=str(datetime_to_int64(context.get('ti').end_date)), workflow_execution=we) return je
def start_new_workflow_execution( self, project_name: Text, workflow_name: Text) -> Optional[WorkflowExecutionInfo]: dag_id = self.airflow_dag_id(project_name, workflow_name) deploy_path = self.config.get('airflow_deploy_path') if deploy_path is None: raise Exception("airflow_deploy_path config not set!") if not self.dag_exist(dag_id): return None context: ExecutionContext = self.airflow_client.schedule_dag(dag_id) with create_session() as session: dagrun = DagRun.get_run_by_id(session=session, dag_id=dag_id, run_id=context.dagrun_id) if dagrun is None: return None else: return WorkflowExecutionInfo( workflow_info=WorkflowInfo(namespace=project_name, workflow_name=workflow_name), workflow_execution_id=str(dagrun.id), status=status.Status.INIT)
def resume_workflow_scheduling(self, project_name: Text, workflow_name: Text) -> WorkflowInfo: return WorkflowInfo(namespace=project_name, workflow_name=workflow_name)
def submit_workflow(self, workflow: Workflow, project_context: ProjectContext) -> WorkflowInfo: return WorkflowInfo(namespace=project_context.project_name, workflow_name=workflow.workflow_name)
def delete_workflow(self, project_name: Text, workflow_name: Text) -> WorkflowInfo: return WorkflowInfo(namespace=project_name, workflow_name=workflow_name)
def proto_to_workflow(proto: WorkflowProto) -> WorkflowInfo: if proto is None: return None else: return WorkflowInfo(namespace=proto.namespace, workflow_name=proto.name, properties=dict(proto.properties))
def submit_workflow(self, workflow: Workflow, project_context: ProjectContext) -> WorkflowInfo: self.workflow = workflow self.project_context = project_context return WorkflowInfo(workflow_name=workflow.workflow_name)
def resume_workflow_scheduling(self, project_name: Text, workflow_name: Text) -> WorkflowInfo: dag_id = self.airflow_dag_id(project_name, workflow_name) DagModel.get_dagmodel(dag_id=dag_id).set_is_paused(is_paused=False) return WorkflowInfo(namespace=project_name, workflow_name=workflow_name)