This is a work in progress document to have some ideas around event based scheduling

Introduction

The current scheduler polls every dagrun for runnable tasks. If a task has not run yet the scheduler will ask the task if it is runnable (ti.is_runnable) which calls ti.are_dependencies_met. For a single task this costs around 0.05-015s. Thus for many tasks in a dag this grows quite quickly and this is the reason why you see increasing scheduling times ("Loop time took xxxx"). The code for are_dependencies_met looks currently like this:

@provide_session
def are_dependencies_met(
        self,
        session=None,
        flag_upstream_failed=False,
        ignore_depends_on_past=False,
        verbose=False):
    """
    Returns a boolean on whether the upstream tasks are in a SUCCESS state
    and considers depends_on_past and the previous run's state.

    :param flag_upstream_failed: This is a hack to generate
        the upstream_failed state creation while checking to see
        whether the task instance is runnable. It was the shortest
        path to add the feature
    :type flag_upstream_failed: boolean
    :param ignore_depends_on_past: if True, ignores depends_on_past
        dependencies. Defaults to False.
    :type ignore_depends_on_past: boolean
    :param verbose: verbose provides more logging in the case where the
        task instance is evaluated as a check right before being executed.
        In the case of the scheduler evaluating the dependencies, this
        logging would be way too verbose.
    :type verbose: boolean
    """
    TI = TaskInstance
    TR = TriggerRule

    task = self.task

    # Checking that the depends_on_past is fulfilled
    if (task.depends_on_past and not ignore_depends_on_past and
            not self.execution_date == task.start_date):
        previous_ti = session.query(TI).filter(
            TI.dag_id == self.dag_id,
            TI.task_id == task.task_id,
            TI.execution_date ==
                self.task.dag.previous_schedule(self.execution_date),
            TI.state.in_({State.SUCCESS, State.SKIPPED}),
        ).first()
        if not previous_ti:
            if verbose:
                logging.warning("depends_on_past not satisfied")
            return False

        # Applying wait_for_downstream
        previous_ti.task = self.task
        if task.wait_for_downstream and not \
                previous_ti.are_dependents_done(session=session):
            if verbose:
                logging.warning("wait_for_downstream not satisfied")
            return False

    # Checking that all upstream dependencies have succeeded
    if not task.upstream_list or task.trigger_rule == TR.DUMMY:
        return True

    # todo: this query becomes quite expensive with dags that have
    # many tasks. It should be refactored to let the task report
    # to the dag run and get the aggregates from there
    qry = (
        session
        .query(
            func.coalesce(func.sum(
                case([(TI.state == State.SUCCESS, 1)], else_=0)), 0),
            func.coalesce(func.sum(
                case([(TI.state == State.SKIPPED, 1)], else_=0)), 0),
            func.coalesce(func.sum(
                case([(TI.state == State.FAILED, 1)], else_=0)), 0),
            func.coalesce(func.sum(
                case([(TI.state == State.UPSTREAM_FAILED, 1)], else_=0)), 0),
            func.count(TI.task_id),
        )
        .filter(
            TI.dag_id == self.dag_id,
            TI.task_id.in_(task.upstream_task_ids),
            TI.execution_date == self.execution_date,
            TI.state.in_([
                State.SUCCESS, State.FAILED,
                State.UPSTREAM_FAILED, State.SKIPPED]),
        )
    )

    successes, skipped, failed, upstream_failed, done = qry.first()

    satisfied = self.evaluate_trigger_rule(
        session=session, successes=successes, skipped=skipped,
        failed=failed, upstream_failed=upstream_failed, done=done,
        flag_upstream_failed=flag_upstream_failed)

    if verbose and not satisfied:
        logging.warning("Trigger rule `{}` not satisfied".format(task.trigger_rule))
    return satisfied


The part that takes most of the time are the aggregation functions for the database. While a slight optimization can be obtained from properly configuring indexes a table scan stays required for every run. This is suboptimal.

The Event Based Scheduler

The structure of a DAG has entry points that are the first tasks that can get executed. These tasks do not have any upstream tasks. Lets call these tasks the kick_starters. After these kick_starters are executed a snowball effect starts that allows the downstream tasks to run. A task does know in what kind of state it is and to what kind of state it is transitioning to. If it reports these state changes downstream it allows the downstream task to determine if it is ready to run.

In pseudo code it should look something like this:

Task A is a kickstarter and has a downstream Task B. Task A is run and reports new state, by calling notify_downstream
 
notify_downstream(state_from, state_to):
  for task_instance in my_downstream_tasks:
     task_instance.update(self, state_from, state_to)
 
The update method looks like this:
 
update(caller, state_from, state_to):
   state_from.counter = state_from.counter - 1
   state_to.counter = state_to.counter + 1
 
   if self.runnable:
      self.state = READY_TO_RUN
 
   # we need to cascade the update to all downstream tasks
   notify_downstream(state_from, state_to)

The scheduler can then just check for READY_TO_RUN tasks and put those in the executor. Obviously, a further optimization would be to let the downstream task schedule itself and let the scheduler handle only kick_starters and exceptions. This however has further implications like communication between the executor and the tasks and some security implications that need further thinking.

Concerns

The downside of an event based approach is that you need to capture all events. If you change the state of a task without notifying the downstream tasks everything will grind to a halt. To make sure that does not happen a state change of a Task should not be done outside the task. Fortunately, python makes this quite simple for us by allowing us to define "@property" so if you will do something like this "ti.state = State.SCHEDULED" if it will then call a setter that should call notify_downstream. Nevertheless code should be carefully reviewed: removing tasks by issuing a delete from the database will leave everything in limbo. Therefore a garbage_collector (or orphan_collector) should run from time to time.

Design Questions

  1. Handling tasks in distributed environments. Lets say task C has two upstream tasks A and B, both of which are currently being executed but on two different machines. {{The update()}} mechanism will need to deal with exchanging state via the database. Otherwise A and B will each update a local version of C but the task will never see both dependencies complete and run. This is a concern even across airflow processes on a single machine.
  2. What happens if an upstream task's state is updated twice (perhaps to the same value)? For example, because of a conflict (two executors encounter the same task and execute it) or for a legitimate reason (immediately after a task runs, a user manually runs it a second time). Some sort of per-task tracking would be useful inside the {{update()}} method to ensure that only the most recent state from each upstream task is being considered. If a task sends a state a second time, then we overwrite the stored state for that task rather than incrementing the counters regardless.
  3. Error handling. Let's say the kick_starter tasks run (I've been thinking of them as "leaves" in contrast to the "roots" at the opposite end of the dag) but the scheduler dies after they complete. When a new scheduler comes online, it needs to know to kick off the new set of "leaf" tasks for which all upstream tasks are complete but which aren't usually the first tasks to run.


Garbage Collector

t.b.d.




  • No labels