def kill_run(self):
_is_killed.set()
try:
self.context.kill_api_client.kill_run(str(self.run_uid))
except Exception as e:
raise DatabandFailFastError(
"Could not send request to kill databand run!", e)
def kill_run(self):
_is_killed.set()
try:
return kill_run(str(self.run_uid), ctx=self.context)
except Exception as e:
raise DatabandFailFastError(
"Could not send request to kill databand run!", e)
def kill_run(self, message=None):
_is_killed.set()
# When initiating kill_run, the api's kill_run sends a signal to all running runs,
# to change their state to shutdown, which in the end sets it to cancelled.
# the task which initiated the killing, the current task run, should have a state of Failed, and not Canceled.
# It is important to set it with an error, to allow the passing of error message, to be displayed in the UI
# as the error_message for the whole run.
tr = current_task_run()
if tr.run == self.run:
task_run_error = TaskRunError.build_from_message(
task_run=tr,
msg=message or DEFAULT_TASK_CANCELED_ERR_MSG,
help_msg="task with task_run_uid:%s initiated kill_run" %
(tr.task_run_uid),
ex_class=DbndCanceledRunError,
)
tr.set_task_run_state(TaskRunState.FAILED,
track=True,
error=task_run_error)
try:
kill_run(str(self.run.run_uid), ctx=self.run.context)
except Exception as e:
raise DatabandFailFastError(
"Could not send request to kill databand run!", e)
if tr.run == self.run:
raise DatabandError(message or DEFAULT_TASK_CANCELED_ERR_MSG)
def _process_dag_task_instances(self, ti_status, executor, pickle_id, session=None):
"""
Process a set of task instances from a set of dag runs. Special handling is done
to account for different task instance states that could be present when running
them in a backfill process.
:param ti_status: the internal status of the job
:type ti_status: DagRunJob._DagRunTaskStatus
:param executor: the executor to run the task instances
:type executor: BaseExecutor
:param pickle_id: the pickle_id if dag is pickled, None otherwise
:type pickle_id: int
:param start_date: the start date of the backfill job
:type start_date: datetime
:param session: the current session object
:type session: Session
:return: the list of execution_dates for the finished dag runs
:rtype: list
"""
executed_run_dates = []
# values() returns a view so we copy to maintain a full list of the TIs to run
all_ti = list(ti_status.to_run.values())
waiting_for_executor_result = {}
while (len(ti_status.to_run) > 0 or len(ti_status.running) > 0) and len(
ti_status.deadlocked
) == 0:
if current.is_killed():
raise friendly_error.task_execution.databand_context_killed(
"SingleDagRunJob scheduling main loop"
)
self.log.debug("*** Clearing out not_ready list ***")
ti_status.not_ready.clear()
self.ti_state_manager.refresh_task_instances_state(
all_ti, self.dag.dag_id, self.execution_date, session=session
)
# we need to execute the tasks bottom to top
# or leaf to root, as otherwise tasks might be
# determined deadlocked while they are actually
# waiting for their upstream to finish
for task in self.dag.topological_sort():
# TODO: too complicated mechanism,
# it's not possible that we have multiple tasks with the same id in to run
for key, ti in list(ti_status.to_run.items()):
if task.task_id != ti.task_id:
continue
if not self._optimize:
ti.refresh_from_db()
task = self.dag.get_task(ti.task_id)
ti.task = task
# TODO : do we need that?
# ignore_depends_on_past = (
# self.ignore_first_depends_on_past and
# ti.execution_date == (start_date or ti.start_date))
ignore_depends_on_past = False
self.log.debug("Task instance to run %s state %s", ti, ti.state)
# guard against externally modified tasks instances or
# in case max concurrency has been reached at task runtime
if ti.state == State.NONE:
self.log.warning(
"FIXME: task instance {} state was set to None "
"externally. This should not happen"
)
ti.set_state(State.SCHEDULED, session=session)
# The task was already marked successful or skipped by a
# different Job. Don't rerun it.
if ti.state == State.SUCCESS:
ti_status.succeeded.add(key)
self.log.debug("Task instance %s succeeded. Don't rerun.", ti)
ti_status.to_run.pop(key)
if key in ti_status.running:
ti_status.running.pop(key)
continue
elif ti.state == State.SKIPPED:
ti_status.skipped.add(key)
self.log.debug("Task instance %s skipped. Don't rerun.", ti)
ti_status.to_run.pop(key)
if key in ti_status.running:
ti_status.running.pop(key)
continue
elif ti.state == State.FAILED:
self.log.error("Task instance %s failed", ti)
ti_status.failed.add(key)
ti_status.to_run.pop(key)
if key in ti_status.running:
ti_status.running.pop(key)
continue
elif ti.state == State.UPSTREAM_FAILED:
self.log.error("Task instance %s upstream failed", ti)
ti_status.failed.add(key)
ti_status.to_run.pop(key)
if key in ti_status.running:
ti_status.running.pop(key)
continue
runtime_deps = []
if self.airflow_config.disable_dag_concurrency_rules:
# RUN Deps validate dag and task concurrency
# It's less relevant when we run in stand along mode with SingleDagRunJob
# from airflow.ti_deps.deps.runnable_exec_date_dep import RunnableExecDateDep
from airflow.ti_deps.deps.valid_state_dep import ValidStateDep
# from airflow.ti_deps.deps.dag_ti_slots_available_dep import DagTISlotsAvailableDep
# from airflow.ti_deps.deps.task_concurrency_dep import TaskConcurrencyDep
# from airflow.ti_deps.deps.pool_slots_available_dep import PoolSlotsAvailableDep
runtime_deps = {
# RunnableExecDateDep(),
ValidStateDep(SCHEDUALED_OR_RUNNABLE),
# DagTISlotsAvailableDep(),
# TaskConcurrencyDep(),
# PoolSlotsAvailableDep(),
}
else:
runtime_deps = RUNNING_DEPS
dagrun_dep_context = DepContext(
deps=runtime_deps,
ignore_depends_on_past=ignore_depends_on_past,
ignore_task_deps=self.ignore_task_deps,
flag_upstream_failed=True,
)
# Is the task runnable? -- then run it
# the dependency checker can change states of tis
if ti.are_dependencies_met(
dep_context=dagrun_dep_context,
session=session,
verbose=self.verbose,
):
ti.refresh_from_db(lock_for_update=True, session=session)
if (
ti.state == State.SCHEDULED
or ti.state == State.UP_FOR_RETRY
):
if executor.has_task(ti):
self.log.debug(
"Task Instance %s already in executor "
"waiting for queue to clear",
ti,
)
else:
self.log.debug("Sending %s to executor", ti)
# if ti.state == State.UP_FOR_RETRY:
# ti._try_number += 1
# Skip scheduled state, we are executing immediately
ti.state = State.QUEUED
session.merge(ti)
cfg_path = None
if executor.__class__ in (
executors.LocalExecutor,
executors.SequentialExecutor,
):
cfg_path = tmp_configuration_copy()
executor.queue_task_instance(
ti,
mark_success=self.mark_success,
pickle_id=pickle_id,
ignore_task_deps=self.ignore_task_deps,
ignore_depends_on_past=ignore_depends_on_past,
pool=self.pool,
cfg_path=cfg_path,
)
ti_status.to_run.pop(key)
ti_status.running[key] = ti
waiting_for_executor_result[key] = ti
session.commit()
continue
if ti.state == State.UPSTREAM_FAILED:
self.log.error("Task instance %s upstream failed", ti)
ti_status.failed.add(key)
ti_status.to_run.pop(key)
if key in ti_status.running:
ti_status.running.pop(key)
continue
# special case
if ti.state == State.UP_FOR_RETRY:
self.log.debug(
"Task instance %s retry period not " "expired yet", ti
)
if key in ti_status.running:
ti_status.running.pop(key)
ti_status.to_run[key] = ti
continue
# all remaining tasks
self.log.debug("Adding %s to not_ready", ti)
ti_status.not_ready.add(key)
# execute the tasks in the queue
self.heartbeat()
executor.heartbeat()
# If the set of tasks that aren't ready ever equals the set of
# tasks to run and there are no running tasks then the backfill
# is deadlocked
if (
ti_status.not_ready
and ti_status.not_ready == set(ti_status.to_run)
and len(ti_status.running) == 0
):
self.log.warning(
"scheduler: Deadlock discovered for ti_status.to_run=%s",
ti_status.to_run.values(),
)
ti_status.deadlocked.update(ti_status.to_run.values())
ti_status.to_run.clear()
self.ti_state_manager.refresh_task_instances_state(
all_ti, self.dag.dag_id, self.execution_date, session=session
)
# check executor state
self._manage_executor_state(ti_status.running, waiting_for_executor_result)
if self._zombie_cleaner:
# this code exists in airflow original scheduler
# clean zombies ( we don't need multiple runs here actually
self._zombie_cleaner.find_and_clean_dag_zombies(
dag=self.dag, execution_date=self.execution_date
)
# update the task counters
self._update_counters(ti_status, waiting_for_executor_result)
# update dag run state
_dag_runs = ti_status.active_runs[:]
for run in _dag_runs:
run.update_state(session=session)
self._update_databand_task_run_states(run)
if run.state in State.finished():
ti_status.finished_runs += 1
ti_status.active_runs.remove(run)
executed_run_dates.append(run.execution_date)
self._log_progress(ti_status)
if self.fail_fast and ti_status.failed:
msg = ",".join([t[1] for t in ti_status.failed])
logger.error(
"scheduler: Terminating executor because a task failed and fail_fast mode is enabled %s",
msg,
)
raise DatabandFailFastError(
"Failing whole pipeline as it has failed/canceled tasks %s" % msg,
)
# return updated status
return executed_run_dates