airflow源码分析-任务调度器实现分析

Airflow源码分析-任务调度器实现分析

概述

本文介绍Airflow执行器的总体实现流程。通过函数调用的方式说明了Airflow scheduler的实现原理，对整个调度过程的源码进行了分析。
通过本文，可以基本把握住Airflow的调度器的运行原理主线。

启动调度器

可以通过命令来启动调度器：

airflow scheduler

启动airflow的调度器后的总执行流程如下：

1. 执行命令后会调用scheduler_command.py#scheduler(args)函数。
2. scheduler(args)函数会打开一个日志文件，然后调用_run_scheduler_job(args)函数，该函数会：创建一个SchedulerJob对象，把DAG文件的目录传入，启动后就开始解析DAG文件目录中的DAG文件(*.py)文件。
3. 然后开始执行SchedulerJob#run()这个函数，启动job侧。该函数在其父类BaseJob中定义。
4. BaseJob#run()函数调用其实现类的_execute()函数，这里就是调用SchedulerJob#_execute函数来启动调度服务。

调度器的实现分析

SchedulerJob#_execute函数

1. 打印启动调度器的日志：可以在调度器服务的日志文件中看到以下日志：Starting the scheduler。

2. 如果不是独立模式进程，则创建一个DagFileProcessorAgent对象，用于读取DAG文件。并把该对象保存到：SchedulerJob#processor_agent成员变量中。

3. 初始化执行器，并设置执行器的回调函数。创建callback_sink对象，它可能是PipeCallbackSink（DagFileProcessorAgent是非独立模式）或DatabaseCallbackSink（DagFileProcessorAgent是独立模式）；创建完成后，对象保存到BaseExecutor#callback_sink变量中（代码中实现的是SchedulerJob#executor函数的返回值）。

4. 调用： self.executor.start()函数来启动执行器executor：调用SchedulerJob#executor.start()

5. 调用self.register_signals()函数来：注册信号处理程序，以便在需要时可以停止儿子进程。

6. 启动DagFileProcessorAgent服务：SchedulerJob.processor_agent.start()

7. 调用： SchedulerJob#_run_scheduler_loop，进入调度服务循环，根据DAG的调度计划，执行DAG中的任务。

8. DAG处理完成后，如果是非独立模式，则停止DagFileProcessorAgent对象，并检查所有文件是否都已经处x理。如果所有文件都已处理，则停用未被调度器触及的DAG。

9. 最后，执行器结束工作，关闭处理DAG文件的代理对象和回调函数，并移除会话对象。

调度服务主循环: SchedulerJob#_run_scheduler_loop

1. 通过if not self.processor_agent来检查DAG文件处理服务是否已经启动

2. 创建一个定时器调度对象timers： timers = EventScheduler()，该对象会定时执行定义的函数。用于周期性地运行一些任务。

3. 向该timer对象中注册一些回调函数，来定时进行一些任务。这些任务包括检查孤立的任务、检查触发器超时、更新池指标、查找僵尸任务等。

4. 通过一个无限循环来不断地进行调度，直到达到指定的运行次数或达到了DAG解析次数的上限。在每次循环中，它会执行以下步骤：

   a. 进入调度循环，若使用sqlite，则运行一个单独的的DAG文件解析进程：processor_agent.run_single_parsing_loop()。

   b. 调用SchedulerJob#_do_scheduling进入实际的调度实现代码中。并返回排队的任务数量。

   c. 启动executor的心跳处理进程：self.executor.heartbeat()

   d. 启动executor的事件处理器：SchedulerJob#_process_executor_events

   e. 启动DagFileProcessorAgent的心跳服务：self.processor_agent.heartbeat()

   f. 运行定时任务；

5. 如果没有工作需要执行，则等待一段时间。

最后，如果达到了指定的运行次数或达到了DAG解析次数的上限，则退出循环。如果使用DagFileProcessorAgent，则在达到解析次数上限时也会退出循环。

调度实现函数：SchedulerJob#_do_scheduling

1. 调用prohibit_commit的函数来返回一个上下文管理器：CommitProhibitorGuard。该上下文管理器可以防止在其作用域之外通过会话对象提交事务，从而严格控制事务的生命周期，以确保在核心调度器循环中的严格控制。如果在上下文管理器之外通过会话对象提交事务，将引发RuntimeError异常。
2. 调用SchedulerJob#_create_dagruns_for_dags函数来根据DagModel中的next_dagrun_create_after列创建任何必要的DagRun。默认情况下，只选择10个DAG，可以通scheduler.max_dagruns_to_create_per_loop设置进行配置。
3. 调用函数: SchedulerJob#_start_queued_dagruns：在DagRuns集合中对象中查找“下n个最老的”正在运行的DAGRun进行调度（默认n = 20，可以通过“scheduler.max_dagruns_per_loop_to_schedule”进行配置），并尝试进度状态（将TIs调度为SCHEDULED，或将DagRuns调度为SUCCESS / FAILURE等）。
4. 调用SchedulerJob#_get_next_dagruns_to_examine检查dagrun的参数。
5. 调用SchedulerJob#_schedule_all_dag_runs函数来决定所有dagrun的调度决定。该函数会遍历所有的dagrun，对每个dagrun调用SchedulerJob#_schedule_dag_run函数来决定是否调度该dagrun。
6. 通过临界区（锁定Pool模型的行），将任务排队，然后将其发送到执行器中。详见_critical_section_enqueue_task_instances()文档。
7. 返回在此迭代中入队的TIs的数量。

其中，步骤2和步骤3需要注意，因为它们会锁定某些行，并且只有一个调度器可以同时处理这些行，因此它们可能会影响调度器的吞吐量。步骤2中默认选择的20个DAG Run是基于它们最久没有被检查/调度的时间来选择的。步骤3中，通过临界区锁定行的目的是为了防止多个调度器同时修改同一个任务实例，这会导致竞态条件和不一致性。

SchedulerJob#_start_queued_dagruns

该方法用于启动处于排队状态的DagRuns。

1. 该方法调用_get_next_dagruns_to_examine方法，获取处于QUEUED状态的DagRuns。

2. 然后，该方法使用DagRun.active_runs_of_dags方法计算每个DAG当前正在运行的DagRun数量，并将结果存储在active_runs_of_dags字典中。

3. 它遍历了每个处于QUEUED状态的DagRun，检查是否可以将其移动到RUNNING状态。

   a. 对于每个DagRun，该方法首先使用DagBag.get_dag方法获取其对应的DAG对象，并将其赋值给dag_run.dag属性。如果DAG不存在，则记录错误并继续处理下一个DagRun。

   b. 接下来，该方法使用active_runs_of_dags字典获取DAG当前正在运行的DagRun数量，并将结果存储在active_runs变量中。

   c. 该方法检查DAG的max_active_runs属性是否为None，如果不是，则检查DAG当前正在运行的DagRun数量是否超过了该属性的值。如果是，则记录调试日志并不将该DagRun移动到RUNNING状态。否则，该方法会将DagRun的状态设置为RUNNING，并更新其start_date属性。此外，该方法还会将DAG当前正在运行的DagRun数量加1，并调用DagRun.notify_dagrun_state_changed方法通知状态已更改。

   d. 最后，该方法调用了_update_state方法，该方法用于设置DagRun的状态为RUNNING，并计算调度延迟（如果DAG是周期性的）。需要注意的是，_update_state方法是一个内部函数，定义在_start_queued_dagruns方法内部。

SchedulerJob#_critical_section_enqueue_task_instances

该方法用于将TaskInstances添加到执行队列中。

该方法包含以下三个步骤：

1. 使用优先级选择TaskInstances，并确保它们处于预期状态，并且不会超过最大活动运行数或池限制。
   原子地更改上述TaskInstances的状态。

2. 将TaskInstances添加到执行器的队列中。需要注意的是，该方法是一个“关键段”，意味着只有一个执行器进程可以同时执行该方法。为了实现这一点，该方法使用了SELECT…FOR UPDATE语句锁定了池表，以确保只有一个进程可以进行修改。

3. 该方法还包含了两个辅助方法：_executable_task_instances_to_queued和_enqueue_task_instances_with_queued_state。其中，_executable_task_instances_to_queued方法用于选择可执行的TaskInstances，而_enqueue_task_instances_with_queued_state方法用于将TaskInstances添加到执行器的队列中。

4. 最后，该方法返回状态发生了变化的TaskInstance的数量。

注意：该方法使用了一个名为max_tis_per_query的属性，它表示每次查询最多选择的TaskInstance数量。如果该属性的值为0，则选择所有可用的TaskInstances；否则，仅选择最多max_tis_per_query个TaskInstances。此外，该方法还使用了一个名为executor的属性，它表示Airflow的执行器对象，用于将TaskInstances添加到执行队列中。

SchedulerJob#_enqueue_task_instances_with_queued_state

该方法的作用是将状态为"queued"的任务实例添加到执行器（executor）的队列中（queued_tasks字典中）等待执行。

该方法接受两个参数：

• task_instances：待执行的任务实例列表。
• session：数据库会话对象。

1. 该方法首先遍历任务实例列表，对于每个状态为"queued"的任务实例，将其命令添加到执行器的队列中等待执行。如果任务实例所属的DAG运行状态为finished，则将任务实例状态设置为"None"，并跳过该任务实例的执行。

2. 对于每个任务实例，该方法使用任务实例的command_as_list方法获取该任务实例的命令，并设置该任务实例的优先级和队列。然后，该方法调用执行器的queue_command方法将该任务实例的命令添加到执行器的队列中等待执行。

需要注意的是，该方法并不会等待任务实例执行完毕，而是将任务实例的执行交给了执行器处理。而执行器的具体处理逻辑在_process_tasks函数的_executor.execute_async方法中。

该函数的实现代码如下：

def _enqueue_task_instances_with_queued_state(self, task_instances: list[TI], session: Session) -> None:
        """
        Takes task_instances, which should have been set to queued, and enqueues them
        with the executor.

        :param task_instances: TaskInstances to enqueue
        :param session: The session object
        """
        # actually enqueue them
        for ti in task_instances:
            if ti.dag_run.state in State.finished:
                ti.set_state(State.NONE, session=session)
                continue
            command = ti.command_as_list(
                local=True,
                pickle_id=ti.dag_model.pickle_id,
            )

            priority = ti.priority_weight
            queue = ti.queue
            self.log.info("Sending %s to executor with priority %s and queue %s", ti.key, priority, queue)

            self.executor.queue_command(
                ti,
                command,
                priority=priority,
                queue=queue,
            )

task通过queue_command已经放到了执行器的任务执行队列queued_tasks中，该变量其实是一个有序的字典，由OrderedDict类来定义。这样不同类型的执行器就可以消费该队列，执行任务了。

Task放入执行器队列

执行器会调用以下函数来执行task。每个执行器实现的实现逻辑不同，可以进入每个执行器中继续分析其实现。

    def _process_tasks(self, task_tuples: list[TaskTuple]) -> None:
        for key, command, queue, executor_config in task_tuples:
            del self.queued_tasks[key]
            self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config)
            self.running.add(key)

小结

本文分析了airflow的任务执行总体流程。分析了从dag文件处理，到task的调度和执行的整个流程。

通过本文的分析可以说基本把握住了Airflow的运行原理的主线。可以根据这条主线，继续分析每个执行器的执行原理，以及任务优先级，DAG文件处理的细节。