DAGScheduler vs TaskScheduler

Concepts covered: Spark Execution Model

An action hands its job to the driver, and 2 schedulers inside the driver divide the work between them. The DAGScheduler thinks in stages: it cuts the job at every stage boundary, works out which stages read from which, and releases a stage only when every stage it reads from has finished. The TaskScheduler thinks in tasks: it takes a released stage as a set of tasks, 1 per partition, and matches each task to a free core on an executor. Neither does the other's job, and most of this lesson follows from where that line falls.

DAGScheduler

• Cuts the job into stages by walking back from the action's result to each stage boundary
• Submits a stage once every parent stage has finished
• Turns a stage into a task set: 1 task per partition still to compute
• Records which executor holds each finished task's output
• Reruns earlier-stage tasks whose output a later stage can no longer read

TaskScheduler

• Receives task sets and knows nothing of the stages behind them
• Matches each task to a free core, preferring executors near its data
• Retries a failed task until it succeeds or runs out of attempts
• Launches backup copies of slow tasks when speculation is on
• Shares cores among task sets from concurrent jobs, first in first out by default

One job through both schedulers

The program below sums ride fares per city over a table of 9 rides read as 3 partitions. A groupBy needs every row for a city on the same task, so the DAGScheduler cuts the job in 2. Stage 0 runs 3 tasks that read the partitions and prepare their rows for regrouping by city; stage 1 runs 3 tasks that each gather their cities' rows and sum them. Watch the order: no stage 1 task starts until all 3 stage 0 tasks have reported success.

from pyspark.sql import functions as F

totals = rides.groupBy("city").agg(F.sum("fare").alias("total_fare"))
totals.collect()

Count the work: 6 tasks in 2 stages, 1 job, with adaptive query execution off as in the animation. The DAGScheduler built both task sets and decided their order, and it never chose an executor. The TaskScheduler placed all 6 tasks and would have retried any that failed, and it never knew they belonged to a groupBy. When the last task of stage 1 succeeded, the DAGScheduler marked the job done and collect() returned the 3 city totals. Adaptive query execution is on in a default session since Spark 3.2, so the Jobs tab lists stage 0 as a job of its own, then as skipped in the job that returns the totals, and the 3 tiny reduce partitions may run as 1 task.

Which scheduler answers which question

Which scheduler handles a failure depends on how far the damage reaches. A problem that stays inside 1 task, such as an exception or a slow machine, belongs to the TaskScheduler: the DAGScheduler is told as each task ends but acts only when the task set finishes or gives up. A problem that destroys output a later stage needs goes to the DAGScheduler, because only it knows which earlier stage produced that output and can run those tasks again. The next 2 sections follow each path.

Stage 3 has 400 tasks. Task 211 throws an exception on its first attempt while the others keep running. What does Spark do next?

• The DAGScheduler resubmits stage 3 as a new task set of all 400 tasks.
• The TaskScheduler queues task 211 for another attempt; the DAGScheduler is not involved yet.
• The TaskScheduler retries task 211, and the DAGScheduler reruns stage 2 to supply fresh input.
• The job fails at once, since a code exception would repeat on every attempt anyway.

About this lesson

This section is part of the Execution Model: Advanced lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.