Should this be combined with the code here https://github.com/apache/spark/blob/master/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala#L618 (or maybe that code can be removed now?)