DAGScheduler — Stage-Oriented Scheduler

DAGScheduler is the scheduling layer of Apache Spark that implements stage-oriented scheduling. It transforms a logical execution plan (i.e. RDD lineage of dependencies built using RDD transformations) to a physical execution plan (using stages).

After an action has been called, SparkContext hands over a logical plan to DAGScheduler that it in turn translates to a set of stages that are submitted as TaskSets for execution (see Execution Model).

The fundamental concepts of DAGScheduler are jobs and stages (refer to Jobs and Stages respectively) that it tracks through internal registries and counters.

DAGScheduler works solely on the driver and is created as part of SparkContext’s initialization (right after TaskScheduler and SchedulerBackend are ready).

DAGScheduler does three things in Spark (thorough explanations follow):

DAGScheduler computes a directed acyclic graph (DAG) of stages for each job, keeps track of which RDDs and stage outputs are materialized, and finds a minimal schedule to run jobs. It then submits stages to TaskScheduler.

In addition to coming up with the execution DAG, DAGScheduler also determines the preferred locations to run each task on, based on the current cache status, and passes the information to TaskScheduler.

DAGScheduler tracks which RDDs are cached (or persisted) to avoid “recomputing” them, i.e. redoing the map side of a shuffle. DAGScheduler remembers what ShuffleMapStages have already produced output files (that are stored in BlockManagers).

DAGScheduler is only interested in cache location coordinates, i.e. host and executor id, per partition of a RDD.

Furthermore, it handles failures due to shuffle output files being lost, in which case old stages may need to be resubmitted. Failures within a stage that are not caused by shuffle file loss are handled by the TaskScheduler itself, which will retry each task a small number of times before cancelling the whole stage.

DAGScheduler uses an event queue architecture in which a thread can post DAGSchedulerEvent events, e.g. a new job or stage being submitted, that DAGScheduler reads and executes sequentially. See the section Internal Event Loop – dag-scheduler-event-loop.

DAGScheduler runs stages in topological order.

DAGScheduler uses SparkContext, TaskScheduler, LiveListenerBus, MapOutputTracker and BlockManager for its services. However, at the very minimum, DAGScheduler takes a SparkContext only (and requests SparkContext for the other services).

DAGScheduler reports metrics about its execution (refer to the section Metrics).

When DAGScheduler schedules a job as a result of executing an action on a RDD or calling SparkContext.runJob() method directly, it spawns parallel tasks to compute (partial) results per partition.

DAGScheduler takes the following when created:

DAGScheduler initializes the internal registries and counters.

DAGScheduler sets itself in the given TaskScheduler and in the end starts DAGScheduler Event Bus.

listenerBus is a LiveListenerBus to post scheduling events and is passed in when DAGScheduler is created.

executorHeartbeatReceived posts a SparkListenerExecutorMetricsUpdate (to listenerBus) and informs BlockManagerMaster that blockManagerId block manager is alive (by posting BlockManagerHeartbeat).

cleanupStateForJobAndIndependentStages cleans up the state for job and any stages that are not part of any other job.

cleanupStateForJobAndIndependentStages looks the job up in the internal jobIdToStageIds registry.

If no stages are found, the following ERROR is printed out to the logs:

Oterwise, cleanupStateForJobAndIndependentStages uses stageIdToStage registry to find the stages (the real objects not ids!).

For each stage, cleanupStateForJobAndIndependentStages reads the jobs the stage belongs to.

If the job does not belong to the jobs of the stage, the following ERROR is printed out to the logs:

If the job was the only job for the stage, the stage (and the stage id) gets cleaned up from the registries, i.e. runningStages, shuffleIdToMapStage, waitingStages, failedStages and stageIdToStage.

While removing from runningStages, you should see the following DEBUG message in the logs:

While removing from waitingStages, you should see the following DEBUG message in the logs:

While removing from failedStages, you should see the following DEBUG message in the logs:

After all cleaning (using stageIdToStage as the source registry), if the stage belonged to the one and only job, you should see the following DEBUG message in the logs:

The job is removed from jobIdToStageIds, jobIdToActiveJob, activeJobs registries.

The final stage of the job is removed, i.e. ResultStage or ShuffleMapStage.

markMapStageJobAsFinished marks the active job finished and notifies Spark listeners.

Internally, markMapStageJobAsFinished marks the zeroth partition finished and increases the number of tasks finished in job.

The job listener is notified about the 0th task succeeded.

The state of the job and independent stages are cleaned up.

Ultimately, SparkListenerJobEnd is posted to LiveListenerBus (as listenerBus) for the job, the current time (in millis) and JobSucceeded job result.

submitJob creates a JobWaiter and posts a JobSubmitted event.

Internally, submitJob does the following:

You may see a IllegalArgumentException thrown when the input partitions references partitions not in the input rdd:

submitMapStage creates a JobWaiter (that it eventually returns) and posts a MapStageSubmitted event to DAGScheduler Event Bus).

Internally, submitMapStage increments nextJobId internal counter to get the job id.

submitMapStage then creates a JobWaiter (with the job id and with one artificial task that will however get completed only when the entire stage finishes).

submitMapStage announces the map stage submission application-wide (by posting a MapStageSubmitted to LiveListenerBus).

submitMapStage returns the JobWaiter.

If the number of partition to compute is 0, submitMapStage throws a SparkException:

cancelJobGroup merely posts a StageCancelled event to the DAGScheduler Event Bus.

cancelJobGroup prints the following INFO message to the logs followed by posting a JobGroupCancelled event to the DAGScheduler Event Bus.

cancelAllJobs merely posts a AllJobsCancelled event to the DAGScheduler Event Bus.

taskStarted merely posts a BeginEvent event to the DAGScheduler Event Bus.

taskGettingResult merely posts a GettingResultEvent event to the DAGScheduler Event Bus.

taskEnded simply posts a CompletionEvent event to the DAGScheduler Event Bus.

taskSetFailed simply posts a TaskSetFailed to DAGScheduler Event Bus.

executorLost simply posts a ExecutorLost event to DAGScheduler Event Bus.

executorAdded simply posts a ExecutorAdded event to DAGScheduler Event Bus.

cancelJob prints the following INFO message and posts a JobCancelled to DAGScheduler Event Bus.

getOrCreateParentStages finds all direct parent ShuffleDependencies of the input rdd and then finds ShuffleMapStage stages for each ShuffleDependency.

runJob submits an action job to the DAGScheduler and waits for a result.

Internally, runJob executes submitJob and then waits until a result comes using JobWaiter.

When the job succeeds, you should see the following INFO message in the logs:

When the job fails, you should see the following INFO message in the logs and the exception (that led to the failure) is thrown.

getOrCreateShuffleMapStage finds or creates the ShuffleMapStage for the input ShuffleDependency.

Internally, getOrCreateShuffleMapStage finds the ShuffleDependency in shuffleIdToMapStage internal registry and returns one when found.

If no ShuffleDependency was available, getOrCreateShuffleMapStage finds all the missing shuffle dependencies and creates corresponding ShuffleMapStage stages (including one for the input shuffleDep).

createShuffleMapStage creates a ShuffleMapStage for the input ShuffleDependency and jobId (of a ActiveJob) possibly copying shuffle map output locations from previous jobs to avoid recomputing records.

Internally, createShuffleMapStage first finds or creates missing parent ShuffleMapStage stages of the associated RDD.

createShuffleMapStage creates a ShuffleMapStage (with the stage id from nextStageId internal counter).

createShuffleMapStage registers the ShuffleMapStage in stageIdToStage and shuffleIdToMapStage internal registries.

createShuffleMapStage calls updateJobIdStageIdMaps.

If MapOutputTrackerMaster tracks the input ShuffleDependency (because other jobs have already computed it), createShuffleMapStage requests the serialized ShuffleMapStage outputs, deserializes them and registers with the new ShuffleMapStage.

If however MapOutputTrackerMaster does not track the input ShuffleDependency, you should see the following INFO message in the logs and createShuffleMapStage registers the ShuffleDependency with MapOutputTrackerMaster.

createShuffleMapStage returns the new ShuffleMapStage.

clearCacheLocs clears the internal registry of the partition locations per RDD.

getMissingAncestorShuffleDependencies finds all missing shuffle dependencies for the given RDD traversing its dependency chain (aka RDD lineage).

Internally, getMissingAncestorShuffleDependencies finds direct parent shuffle dependencies of the input RDD and collects the ones that are not registered in shuffleIdToMapStage internal registry. It repeats the process for the RDDs of the parent shuffle dependencies.

getShuffleDependencies finds direct parent shuffle dependencies for the given RDD.

Internally, getShuffleDependencies takes the direct shuffle dependencies of the input RDD and direct shuffle dependencies of all the parent non-ShuffleDependencies in the dependency chain (aka RDD lineage).

The internal failJobAndIndependentStages method fails the input job and all the stages that are only used by the job.

Internally, failJobAndIndependentStages uses jobIdToStageIds internal registry to look up the stages registered for the job.

If no stages could be found, you should see the following ERROR message in the logs:

Otherwise, for every stage, failJobAndIndependentStages finds the job ids the stage belongs to.

If no stages could be found or the job is not referenced by the stages, you should see the following ERROR message in the logs:

Only when there is exactly one job registered for the stage and the stage is in RUNNING state (in runningStages internal registry), TaskScheduler is requested to cancel the stage’s tasks and marks the stage finished.

abortStage is an internal method that finds all the active jobs that depend on the failedStage stage and fails them.

Internally, abortStage looks the failedStage stage up in the internal stageIdToStage registry and exits if there the stage was not registered earlier.

If it was, abortStage finds all the active jobs (in the internal activeJobs registry) with the final stage depending on the failedStage stage.

At this time, the completionTime property (of the failed stage’s StageInfo) is assigned to the current time (millis).

All the active jobs that depend on the failed stage (as calculated above) and the stages that do not belong to other jobs (aka independent stages) are failed (with the failure reason being “Job aborted due to stage failure: [reason]” and the input exception).

If there are no jobs depending on the failed stage, you should see the following INFO message in the logs:

stageDependsOn compares two stages and returns whether the stage depends on target stage (i.e. true) or not (i.e. false).

Internally, stageDependsOn walks through the graph of RDDs of the input stage. For every RDD in the RDD’s dependencies (using RDD.dependencies) stageDependsOn adds the RDD of a NarrowDependency to a stack of RDDs to visit while for a ShuffleDependency it finds ShuffleMapStage stages for a ShuffleDependency for the dependency and the stage‘s first job id that it later adds to a stack of RDDs to visit if the map stage is ready, i.e. all the partitions have shuffle outputs.

After all the RDDs of the input stage are visited, stageDependsOn checks if the target‘s RDD is among the RDDs of the stage, i.e. whether the stage depends on target stage.

eventProcessLoop is DAGScheduler’s event bus to which Spark (by submitJob) posts jobs to schedule their execution. Later on, TaskSetManager talks back to DAGScheduler to inform about the status of the tasks using the same “communication channel”.

It allows Spark to release the current thread when posting happens and let the event loop handle events on a separate thread – asynchronously.

…​IMAGE…​FIXME

submitWaitingChildStages submits for execution all waiting stages for which the input parent Stage is the direct parent.

When executed, you should see the following TRACE messages in the logs:

submitWaitingChildStages finds child stages of the input parent stage, removes them from waitingStages internal registry, and submits one by one sorted by their job ids.

submitStage is an internal method that DAGScheduler uses to submit the input stage or its missing parents (if there any stages not computed yet before the input stage could).

submitStage recursively submits any missing parents of the stage.

Internally, submitStage first finds the earliest-created job id that needs the stage.

The following steps depend on whether there is a job or not.

If there are no jobs that require the stage, submitStage aborts it with the reason:

If however there is a job for the stage, you should see the following DEBUG message in the logs:

submitStage checks the status of the stage and continues when it was not recorded in waiting, running or failed internal registries. It simply exits otherwise.

With the stage ready for submission, submitStage calculates the list of missing parent stages of the stage (sorted by their job ids). You should see the following DEBUG message in the logs:

When the stage has no parent stages missing, you should see the following INFO message in the logs:

submitStage submits the stage (with the earliest-created job id) and finishes.

If however there are missing parent stages for the stage, submitStage submits all the parent stages, and the stage is recorded in the internal waitingStages registry.

A single stage can be re-executed in multiple attempts due to fault recovery. The number of attempts is configured (FIXME).

If TaskScheduler reports that a task failed because a map output file from a previous stage was lost, the DAGScheduler resubmits the lost stage. This is detected through a CompletionEvent with FetchFailed, or an ExecutorLost event. DAGScheduler will wait a small amount of time to see whether other nodes or tasks fail, then resubmit TaskSets for any lost stage(s) that compute the missing tasks.

Please note that tasks from the old attempts of a stage could still be running.

A stage object tracks multiple StageInfo objects to pass to Spark listeners or the web UI.

The latest StageInfo for the most recent attempt for a stage is accessible through latestInfo.

DAGScheduler computes where to run each task in a stage based on the preferred locations of its underlying RDDs, or the location of cached or shuffle data.

See SPARK-9850 Adaptive execution in Spark for the design document. The work is currently in progress.

DAGScheduler.submitMapStage method is used for adaptive query planning, to run map stages and look at statistics about their outputs before submitting downstream stages.

DAGScheduler uses the following ScheduledThreadPoolExecutors (with the policy of removing cancelled tasks from a work queue at time of cancellation):

They are created using ThreadUtils.newDaemonSingleThreadScheduledExecutor method that uses Guava DSL to instantiate a ThreadFactory.

getMissingParentStages finds missing parent ShuffleMapStages in the dependency graph of the input stage (using the breadth-first search algorithm).

Internally, getMissingParentStages starts with the stage‘s RDD and walks up the tree of all parent RDDs to find uncached partitions.

getMissingParentStages traverses the parent dependencies of the RDD and acts according to their type, i.e. ShuffleDependency or NarrowDependency.

For each NarrowDependency, getMissingParentStages simply marks the corresponding RDD to visit and moves on to a next dependency of a RDD or works on another unvisited parent RDD.

getMissingParentStages focuses on ShuffleDependency dependencies.

For each ShuffleDependency, getMissingParentStages finds ShuffleMapStage stages. If the ShuffleMapStage is not available, it is added to the set of missing (map) stages.

submitMissingTasks…​FIXME

When executed, you should see the following DEBUG message in the logs:

The input stage‘s pendingPartitions internal field is cleared (it is later filled out with the partitions to run tasks for).

submitMissingTasks requests the stage for missing partitions, i.e. the indices of the partitions to compute.

submitMissingTasks marks the stage as running (i.e. adds it to runningStages internal registry).

submitMissingTasks notifies OutputCommitCoordinator that the stage is started.

For the missing partitions, submitMissingTasks computes their task locality preferences, i.e. pairs of missing partition ids and their task locality information.
HERE
NOTE: The locality information of a RDD is called preferred locations.

In case of non-fatal exceptions at this time (while getting the locality information), submitMissingTasks creates a new stage attempt.

Despite the failure to submit any tasks, submitMissingTasks does announce that at least there was an attempt on LiveListenerBus by posting a SparkListenerStageSubmitted message.

submitMissingTasks then aborts the stage (with the reason being “Task creation failed” followed by the exception).

The stage is removed from the internal runningStages collection of stages and submitMissingTasks exits.

When no exception was thrown (while computing the locality information for tasks), submitMissingTasks creates a new stage attempt and announces it on LiveListenerBus by posting a SparkListenerStageSubmitted message.

At that time, submitMissingTasks serializes the RDD (of the stage for which tasks are submitted for) and, depending on the type of the stage, the ShuffleDependency (for ShuffleMapStage) or the function (for ResultStage).

The serialized so-called task binary bytes are “wrapped” as a broadcast variable (to make it available for executors to execute later on).

Any NotSerializableException exceptions lead to aborting the stage (with the reason being “Task not serializable: [exception]”) and removing the stage from the internal runningStages collection of stages. submitMissingTasks exits.

Any non-fatal exceptions lead to aborting the stage (with the reason being “Task serialization failed” followed by the exception) and removing the stage from the internal runningStages collection of stages. submitMissingTasks exits.

With no exceptions along the way, submitMissingTasks computes a collection of tasks to execute for the missing partitions (of the stage).

submitMissingTasks creates a ShuffleMapTask or ResultTask for every missing partition of the stage being ShuffleMapStage or ResultStage, respectively. submitMissingTasks uses the preferred locations (computed earlier) per partition.

Any non-fatal exceptions lead to aborting the stage (with the reason being “Task creation failed” followed by the exception) and removing the stage from the internal runningStages collection of stages. submitMissingTasks exits.

If there are tasks to submit for execution (i.e. there are missing partitions in the stage), you should see the following INFO message in the logs:

submitMissingTasks records the partitions (of the tasks) in the stage‘s pendingPartitions property.

You should see the following DEBUG message in the logs:

submitMissingTasks submits the tasks to TaskScheduler for execution (with the id of the stage, attempt id, the input jobId, and the properties of the ActiveJob with jobId).

submitMissingTasks records the submission time in the stage’s StageInfo and exits.

If however there are no tasks to submit for execution, submitMissingTasks marks the stage as finished (with no errorMessage).

You should see a DEBUG message that varies per the type of the input stage which are:

or

for ShuffleMapStage and ResultStage, respectively.

In the end, with no tasks to submit for execution, submitMissingTasks submits waiting child stages for execution and exits.

getPreferredLocs is simply an alias for the internal (recursive) getPreferredLocsInternal.

getCacheLocs gives TaskLocations (block locations) for the partitions of the input rdd. getCacheLocs caches lookup results in cacheLocs internal registry.

Internally, getCacheLocs finds rdd in the cacheLocs internal registry (of partition locations per RDD).

If rdd is not in cacheLocs internal registry, getCacheLocs branches per its storage level.

For NONE storage level (i.e. no caching), the result is an empty locations (i.e. no location preference).

For other non-NONE storage levels, getCacheLocs requests BlockManagerMaster for block locations that are then mapped to TaskLocations with the hostname of the owning BlockManager for a block (of a partition) and the executor id.

getCacheLocs records the computed block locations per partition (as TaskLocation) in cacheLocs internal registry.

getPreferredLocsInternal first finds the TaskLocations for the partition of the rdd (using cacheLocs internal cache) and returns them.

Otherwise, if not found, getPreferredLocsInternal requests rdd for the preferred locations of partition and returns them.

Otherwise, if not found, getPreferredLocsInternal finds the first parent NarrowDependency and (recursively) finds TaskLocations.

If all the attempts fail to yield any non-empty result, getPreferredLocsInternal returns an empty collection of TaskLocations.

stop stops the internal dag-scheduler-message thread pool, dag-scheduler-event-loop, and TaskScheduler.

The private updateAccumulators method merges the partial values of accumulators from a completed task into their “source” accumulators on the driver.

For each AccumulableInfo in the CompletionEvent, a partial value from a task is obtained (from AccumulableInfo.update) and added to the driver’s accumulator (using Accumulable.++= method).

For named accumulators with the update value being a non-zero value, i.e. not Accumulable.zero: