expose stageInfo in SparkContext

Author: Imran Rashid

core/src/main/scala/spark/SparkContext.scala

@@ -40,13 +40,16 @@ import spark.deploy.LocalSparkCluster
 import spark.partial.ApproximateEvaluator
 import spark.partial.PartialResult
 import rdd.{CheckpointRDD, HadoopRDD, NewHadoopRDD, UnionRDD}
-import scheduler.{ResultTask, ShuffleMapTask, DAGScheduler, TaskScheduler}
+import scheduler._
 import spark.scheduler.local.LocalScheduler
 import spark.scheduler.cluster.{SparkDeploySchedulerBackend, SchedulerBackend, ClusterScheduler}
 import spark.scheduler.mesos.{CoarseMesosSchedulerBackend, MesosSchedulerBackend}
 import storage.BlockManagerUI
+import storage.RDDInfo
+import storage.StorageStatus
 import util.{MetadataCleaner, TimeStampedHashMap}
 import storage.{StorageStatus, StorageUtils, RDDInfo}
+import scala.Some
 
 /**
  * Main entry point for Spark functionality. A SparkContext represents the connection to a Spark
@@ -482,6 +485,10 @@ class SparkContext(
     StorageUtils.rddInfoFromStorageStatus(getSlavesStorageStatus, this)
   }
 
+  def getStageInfo: Map[Stage,StageInfo] = {
+    dagScheduler.stageToInfos
+  }
+
   /**
    * Return information about blocks stored in all of the slaves
    */

core/src/main/scala/spark/scheduler/DAGScheduler.scala

@@ -66,7 +66,7 @@ class DAGScheduler(taskSched: TaskScheduler) extends TaskSchedulerListener with
 
   val shuffleToMapStage = new TimeStampedHashMap[Int, Stage]
 
-  private val stageToInfos = new TimeStampedHashMap[Stage, StageInfo]
+  private[spark] val stageToInfos = new TimeStampedHashMap[Stage, StageInfo]
 
   private val sparkListeners = Traversable[SparkListener]()