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)