Part of [SPARK-2456] Removed some HashMaps from DAGScheduler by storing information in Stage. #1561
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -85,12 +85,9 @@ class DAGScheduler( | |
| private val nextStageId = new AtomicInteger(0) | ||
|
|
||
| private[scheduler] val jobIdToStageIds = new HashMap[Int, HashSet[Int]] | ||
| private[scheduler] val stageIdToStage = new HashMap[Int, Stage] | ||
| private[scheduler] val shuffleToMapStage = new HashMap[Int, Stage] | ||
| private[scheduler] val jobIdToActiveJob = new HashMap[Int, ActiveJob] | ||
|
|
||
| // Stages we need to run whose parents aren't done | ||
| private[scheduler] val waitingStages = new HashSet[Stage] | ||
|
|
@@ -101,9 +98,6 @@ class DAGScheduler( | |
| // Stages that must be resubmitted due to fetch failures | ||
| private[scheduler] val failedStages = new HashSet[Stage] | ||
|
|
||
| private[scheduler] val activeJobs = new HashSet[ActiveJob] | ||
|
|
||
| // Contains the locations that each RDD's partitions are cached on | ||
|
|
@@ -223,7 +217,6 @@ class DAGScheduler( | |
| new Stage(id, rdd, numTasks, shuffleDep, getParentStages(rdd, jobId), jobId, callSite) | ||
| stageIdToStage(id) = stage | ||
| updateJobIdStageIdMaps(jobId, stage) | ||
| stage | ||
| } | ||
|
|
||
|
|
@@ -315,13 +308,12 @@ class DAGScheduler( | |
| */ | ||
| private def updateJobIdStageIdMaps(jobId: Int, stage: Stage) { | ||
| def updateJobIdStageIdMapsList(stages: List[Stage]) { | ||
| if (stages.nonEmpty) { | ||
| val s = stages.head | ||
| s.jobIds += jobId | ||
| jobIdToStageIds.getOrElseUpdate(jobId, new HashSet[Int]()) += s.id | ||
| val parents: List[Stage] = getParentStages(s.rdd, jobId) | ||
| val parentsWithoutThisJobId = parents.filter { ! _.jobIds.contains(jobId) } | ||
| updateJobIdStageIdMapsList(parentsWithoutThisJobId ++ stages.tail) | ||
| } | ||
| } | ||
|
|
@@ -333,16 +325,15 @@ class DAGScheduler( | |
| * handle cancelling tasks or notifying the SparkListener about finished jobs/stages/tasks. | ||
| * | ||
| * @param job The job whose state to cleanup. | ||
| */ | ||
| private def cleanupStateForJobAndIndependentStages(job: ActiveJob) { | ||
| val registeredStages = jobIdToStageIds.get(job.jobId) | ||
| if (registeredStages.isEmpty || registeredStages.get.isEmpty) { | ||
| logError("No stages registered for job " + job.jobId) | ||
| } else { | ||
| stageIdToStage.filterKeys(stageId => registeredStages.get.contains(stageId)).foreach { | ||
| case (stageId, stage) => | ||
| val jobSet = stage.jobIds | ||
| if (!jobSet.contains(job.jobId)) { | ||
| logError( | ||
| "Job %d not registered for stage %d even though that stage was registered for the job" | ||
|
|
@@ -355,14 +346,9 @@ class DAGScheduler( | |
| logDebug("Removing running stage %d".format(stageId)) | ||
| runningStages -= stage | ||
| } | ||
| for ((k, v) <- shuffleToMapStage.find(_._2 == stage)) { | ||
| shuffleToMapStage.remove(k) | ||
| } | ||
| if (waitingStages.contains(stage)) { | ||
| logDebug("Removing stage %d from waiting set.".format(stageId)) | ||
| waitingStages -= stage | ||
|
|
@@ -374,7 +360,6 @@ class DAGScheduler( | |
| } | ||
| // data structures based on StageId | ||
| stageIdToStage -= stageId | ||
|
|
||
| ShuffleMapTask.removeStage(stageId) | ||
| ResultTask.removeStage(stageId) | ||
|
|
@@ -393,19 +378,7 @@ class DAGScheduler( | |
| jobIdToStageIds -= job.jobId | ||
| jobIdToActiveJob -= job.jobId | ||
| activeJobs -= job | ||
| job.finalStage.resultOfJob = None | ||
| } | ||
|
|
||
| /** | ||
|
|
@@ -591,9 +564,10 @@ class DAGScheduler( | |
| job.listener.jobFailed(exception) | ||
| } finally { | ||
| val s = job.finalStage | ||
| // clean up data structures that were populated for a local job, | ||
| // but that won't get cleaned up via the normal paths through | ||
| // completion events or stage abort | ||
| stageIdToStage -= s.id | ||
| jobIdToStageIds -= job.jobId | ||
| listenerBus.post(SparkListenerJobEnd(job.jobId, jobResult)) | ||
| } | ||
|
|
@@ -605,12 +579,8 @@ class DAGScheduler( | |
| // That should take care of at least part of the priority inversion problem with | ||
| // cross-job dependencies. | ||
| private def activeJobForStage(stage: Stage): Option[Int] = { | ||
| val jobsThatUseStage: Array[Int] = stage.jobIds.toArray.sorted | ||
| jobsThatUseStage.find(jobIdToActiveJob.contains) | ||
| } | ||
|
|
||
| private[scheduler] def handleJobGroupCancelled(groupId: String) { | ||
|
|
@@ -642,9 +612,8 @@ class DAGScheduler( | |
| // is in the process of getting stopped. | ||
| val stageFailedMessage = "Stage cancelled because SparkContext was shut down" | ||
| runningStages.foreach { stage => | ||
| stage.info.stageFailed(stageFailedMessage) | ||
| listenerBus.post(SparkListenerStageCompleted(stage.info)) | ||
| } | ||
| listenerBus.post(SparkListenerJobEnd(job.jobId, JobFailed(error))) | ||
| } | ||
|
|
@@ -690,7 +659,7 @@ class DAGScheduler( | |
| } else { | ||
| jobIdToActiveJob(jobId) = job | ||
| activeJobs += job | ||
| finalStage.resultOfJob = Some(job) | ||
| listenerBus.post(SparkListenerJobStart(job.jobId, jobIdToStageIds(jobId).toArray, | ||
| properties)) | ||
| submitStage(finalStage) | ||
|
|
@@ -727,8 +696,7 @@ class DAGScheduler( | |
| private def submitMissingTasks(stage: Stage, jobId: Int) { | ||
| logDebug("submitMissingTasks(" + stage + ")") | ||
| // Get our pending tasks and remember them in our pendingTasks entry | ||
| stage.pendingTasks.clear() | ||
| var tasks = ArrayBuffer[Task[_]]() | ||
| if (stage.isShuffleMap) { | ||
| for (p <- 0 until stage.numPartitions if stage.outputLocs(p) == Nil) { | ||
|
|
@@ -737,7 +705,7 @@ class DAGScheduler( | |
| } | ||
| } else { | ||
| // This is a final stage; figure out its job's missing partitions | ||
| val job = stage.resultOfJob.get | ||
| for (id <- 0 until job.numPartitions if !job.finished(id)) { | ||
| val partition = job.partitions(id) | ||
| val locs = getPreferredLocs(stage.rdd, partition) | ||
|
|
@@ -758,7 +726,7 @@ class DAGScheduler( | |
| // serializable. If tasks are not serializable, a SparkListenerStageCompleted event | ||
| // will be posted, which should always come after a corresponding SparkListenerStageSubmitted | ||
| // event. | ||
| listenerBus.post(SparkListenerStageSubmitted(stage.info, properties)) | ||
|
|
||
|
Contributor
Author
There was a problem hiding this comment. do i need to call
Contributor
There was a problem hiding this comment. Actually you should clear it here, I think the issue is if a stage gets resubmitted. |
||
| // Preemptively serialize a task to make sure it can be serialized. We are catching this | ||
| // exception here because it would be fairly hard to catch the non-serializable exception | ||
|
|
@@ -778,11 +746,11 @@ class DAGScheduler( | |
| } | ||
|
|
||
| logInfo("Submitting " + tasks.size + " missing tasks from " + stage + " (" + stage.rdd + ")") | ||
| stage.pendingTasks ++= tasks | ||
| logDebug("New pending tasks: " + stage.pendingTasks) | ||
| taskScheduler.submitTasks( | ||
| new TaskSet(tasks.toArray, stage.id, stage.newAttemptId(), stage.jobId, properties)) | ||
| stage.info.submissionTime = Some(clock.getTime()) | ||
| } else { | ||
| logDebug("Stage " + stage + " is actually done; %b %d %d".format( | ||
| stage.isAvailable, stage.numAvailableOutputs, stage.numPartitions)) | ||
|
|
@@ -807,13 +775,13 @@ class DAGScheduler( | |
| val stage = stageIdToStage(task.stageId) | ||
|
|
||
| def markStageAsFinished(stage: Stage) = { | ||
| val serviceTime = stage.info.submissionTime match { | ||
| case Some(t) => "%.03f".format((clock.getTime() - t) / 1000.0) | ||
| case _ => "Unknown" | ||
| } | ||
| logInfo("%s (%s) finished in %s s".format(stage, stage.name, serviceTime)) | ||
| stage.info.completionTime = Some(clock.getTime()) | ||
| listenerBus.post(SparkListenerStageCompleted(stage.info)) | ||
| runningStages -= stage | ||
| } | ||
| event.reason match { | ||
|
|
@@ -822,18 +790,18 @@ class DAGScheduler( | |
| // TODO: fail the stage if the accumulator update fails... | ||
| Accumulators.add(event.accumUpdates) // TODO: do this only if task wasn't resubmitted | ||
| } | ||
| stage.pendingTasks -= task | ||
| task match { | ||
| case rt: ResultTask[_, _] => | ||
| stage.resultOfJob match { | ||
| case Some(job) => | ||
| if (!job.finished(rt.outputId)) { | ||
| job.finished(rt.outputId) = true | ||
| job.numFinished += 1 | ||
| // If the whole job has finished, remove it | ||
| if (job.numFinished == job.numPartitions) { | ||
| markStageAsFinished(stage) | ||
| cleanupStateForJobAndIndependentStages(job) | ||
| listenerBus.post(SparkListenerJobEnd(job.jobId, JobSucceeded)) | ||
| } | ||
|
|
||
|
|
@@ -860,7 +828,7 @@ class DAGScheduler( | |
| } else { | ||
| stage.addOutputLoc(smt.partitionId, status) | ||
| } | ||
| if (runningStages.contains(stage) && stage.pendingTasks.isEmpty) { | ||
| markStageAsFinished(stage) | ||
| logInfo("looking for newly runnable stages") | ||
| logInfo("running: " + runningStages) | ||
|
|
@@ -909,7 +877,7 @@ class DAGScheduler( | |
|
|
||
| case Resubmitted => | ||
| logInfo("Resubmitted " + task + ", so marking it as still running") | ||
| stage.pendingTasks += task | ||
|
|
||
| case FetchFailed(bmAddress, shuffleId, mapId, reduceId) => | ||
| // Mark the stage that the reducer was in as unrunnable | ||
|
|
@@ -994,13 +962,14 @@ class DAGScheduler( | |
| } | ||
|
|
||
| private[scheduler] def handleStageCancellation(stageId: Int) { | ||
| stageIdToStage.get(stageId) match { | ||
| case Some(stage) => | ||
| val jobsThatUseStage: Array[Int] = stage.jobIds.toArray | ||
| jobsThatUseStage.foreach { jobId => | ||
| handleJobCancellation(jobId, s"because Stage $stageId was cancelled") | ||
| } | ||
| case None => | ||
| logInfo("No active jobs to kill for Stage " + stageId) | ||
| } | ||
| submitWaitingStages() | ||
| } | ||
|
|
@@ -1009,8 +978,8 @@ class DAGScheduler( | |
| if (!jobIdToStageIds.contains(jobId)) { | ||
| logDebug("Trying to cancel unregistered job " + jobId) | ||
| } else { | ||
| failJobAndIndependentStages( | ||
| jobIdToActiveJob(jobId), "Job %d cancelled %s".format(jobId, reason)) | ||
| } | ||
| submitWaitingStages() | ||
| } | ||
|
|
@@ -1024,26 +993,21 @@ class DAGScheduler( | |
| // Skip all the actions if the stage has been removed. | ||
| return | ||
| } | ||
| val dependentJobs: Seq[ActiveJob] = | ||
| activeJobs.filter(job => stageDependsOn(job.finalStage, failedStage)).toSeq | ||
| failedStage.info.completionTime = Some(clock.getTime()) | ||
| for (job <- dependentJobs) { | ||
| failJobAndIndependentStages(job, s"Job aborted due to stage failure: $reason") | ||
| } | ||
| if (dependentJobs.isEmpty) { | ||
| logInfo("Ignoring failure of " + failedStage + " because all jobs depending on it are done") | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Fails a job and all stages that are only used by that job, and cleans up relevant state. | ||
| */ | ||
| private def failJobAndIndependentStages(job: ActiveJob, failureReason: String) { | ||
| val error = new SparkException(failureReason) | ||
| var ableToCancelStages = true | ||
|
|
||
|
|
@@ -1057,7 +1021,7 @@ class DAGScheduler( | |
| logError("No stages registered for job " + job.jobId) | ||
| } | ||
| stages.foreach { stageId => | ||
| val jobsForStage: Option[HashSet[Int]] = stageIdToStage.get(stageId).map(_.jobIds) | ||
| if (jobsForStage.isEmpty || !jobsForStage.get.contains(job.jobId)) { | ||
| logError( | ||
| "Job %d not registered for stage %d even though that stage was registered for the job" | ||
|
|
@@ -1071,9 +1035,8 @@ class DAGScheduler( | |
| if (runningStages.contains(stage)) { | ||
| try { // cancelTasks will fail if a SchedulerBackend does not implement killTask | ||
| taskScheduler.cancelTasks(stageId, shouldInterruptThread) | ||
| stage.info.stageFailed(failureReason) | ||
| listenerBus.post(SparkListenerStageCompleted(stage.info)) | ||
| } catch { | ||
| case e: UnsupportedOperationException => | ||
| logInfo(s"Could not cancel tasks for stage $stageId", e) | ||
|
|
@@ -1086,7 +1049,7 @@ class DAGScheduler( | |
|
|
||
| if (ableToCancelStages) { | ||
| job.listener.jobFailed(error) | ||
| cleanupStateForJobAndIndependentStages(job) | ||
| listenerBus.post(SparkListenerJobEnd(job.jobId, JobFailed(error))) | ||
| } | ||
| } | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why the type annotation?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It is not obvious to me what getParentStages return just staring at the code. Sometimes when we refer to stages we use a set of integers, sometimes just Stage objects.