def kill_running_tasks(self):
"""Stop running the specified task instance and downstream tasks.
Obtain task_instance from session according to dag_id, run_id and task_id,
If task_id is not empty, get task_instance with RUNNIN or NONE status from dag_run according to task_id,
and set task_instance status to FAILED.
If task_id is empty, get all task_instances whose status is RUNNIN or NONE from dag_run,
and set the status of these task_instances to FAILED.
args:
dag_id: dag id
run_id: the run id of dag run
task_id: the task id of task instance of dag
"""
logging.info("Executing custom 'kill_running_tasks' function")
dagbag = self.get_dagbag()
dag_id = self.get_argument(request, 'dag_id')
run_id = self.get_argument(request, 'run_id')
task_id = self.get_argument(request, 'task_id')
session = settings.Session()
query = session.query(DagRun)
dag_run = query.filter(
DagRun.dag_id == dag_id,
DagRun.run_id == run_id
).first()
if dag_run is None:
return ApiResponse.not_found("dag run is not found")
if dag_id not in dagbag.dags:
return ApiResponse.bad_request("Dag id {} not found".format(dag_id))
dag = dagbag.get_dag(dag_id)
logging.info('dag: ' + str(dag))
logging.info('dag_subdag: ' + str(dag.subdags))
tis = []
if task_id:
task_instance = DagRun.get_task_instance(dag_run, task_id)
if task_instance is None or task_instance.state not in [State.RUNNING, State.NONE]:
return ApiResponse.not_found("task is not found or state is neither RUNNING nor NONE")
else:
tis.append(task_instance)
else:
tis = DagRun.get_task_instances(dag_run, [State.RUNNING, State.NONE])
logging.info('tis: ' + str(tis))
running_task_count = len(tis)
if running_task_count > 0:
for ti in tis:
ti.state = State.FAILED
ti.end_date = timezone.utcnow()
session.merge(ti)
session.commit()
else:
return ApiResponse.not_found("dagRun don't have running tasks")
session.close()
return ApiResponse.success()
def restart_failed_task(self):
"""Restart the failed task in the specified dag run.
According to dag_id, run_id get dag_run from session,
query task_instances that status is FAILED in dag_run,
restart them and clear status of all task_instance's downstream of them.
args:
dag_id: dag id
run_id: the run id of dag run
"""
logging.info("Executing custom 'restart_failed_task' function")
dagbag = self.get_dagbag()
dag_id = self.get_argument(request, 'dag_id')
run_id = self.get_argument(request, 'run_id')
session = settings.Session()
query = session.query(DagRun)
dag_run = query.filter(
DagRun.dag_id == dag_id,
DagRun.run_id == run_id
).first()
if dag_run is None:
return ApiResponse.not_found("dag run is not found")
if dag_id not in dagbag.dags:
return ApiResponse.bad_request("Dag id {} not found".format(dag_id))
dag = dagbag.get_dag(dag_id)
if dag is None:
return ApiResponse.not_found("dag is not found")
tis = DagRun.get_task_instances(dag_run, State.FAILED)
logging.info('task_instances: ' + str(tis))
failed_task_count = len(tis)
if failed_task_count > 0:
for ti in tis:
dag = DAG.sub_dag(
self=dag,
task_regex=r"^{0}$".format(ti.task_id),
include_downstream=True,
include_upstream=False)
count = DAG.clear(
self=dag,
start_date=dag_run.execution_date,
end_date=dag_run.execution_date,
)
logging.info('count：' + str(count))
else:
return ApiResponse.not_found("dagRun don't have failed tasks")
session.close()
return ApiResponse.success({
'failed_task_count': failed_task_count,
'clear_task_count': count
})