Use Spark job grouping to distinguish steps of the machine learning flow

core/src/main/scala/com/salesforce/op/OpWorkflow.scala

@@ -38,6 +38,7 @@ import com.salesforce.op.stages.impl.feature.TimePeriod
 import com.salesforce.op.stages.impl.preparators.CorrelationType
 import com.salesforce.op.stages.impl.selector.ModelSelector
 import com.salesforce.op.utils.reflection.ReflectionUtils
+import com.salesforce.op.utils.spark.{JobGroupUtil, OpStep}
 import com.salesforce.op.utils.spark.RichDataset._
 import com.salesforce.op.utils.stages.FitStagesUtil
 import com.salesforce.op.utils.stages.FitStagesUtil.{CutDAG, FittedDAG, Layer, StagesDAG}
@@ -232,28 +233,30 @@ class OpWorkflow(val uid: String = UID[OpWorkflow]) extends OpWorkflowCore {
    * @return Dataframe with all the features generated + persisted
    */
   protected def generateRawData()(implicit spark: SparkSession): DataFrame = {
-    (reader, rawFeatureFilter) match {
-      case (None, None) => throw new IllegalArgumentException(
-        "Data reader must be set either directly on the workflow or through the RawFeatureFilter")
-      case (Some(r), None) =>
-        checkReadersAndFeatures()
-        r.generateDataFrame(rawFeatures, parameters).persist()
-      case (rd, Some(rf)) =>
-        rd match {
-          case None => setReader(rf.trainingReader)
-          case Some(r) => if (r != rf.trainingReader) log.warn(
-            "Workflow data reader and RawFeatureFilter training reader do not match! " +
-              "The RawFeatureFilter training reader will be used to generate the data for training")
-        }
-        checkReadersAndFeatures()
+    JobGroupUtil.withJobGroup(OpStep.DataReadingAndFiltering) {
+      (reader, rawFeatureFilter) match {
+        case (None, None) => throw new IllegalArgumentException(
+          "Data reader must be set either directly on the workflow or through the RawFeatureFilter")
+        case (Some(r), None) =>
+          checkReadersAndFeatures()
+          r.generateDataFrame(rawFeatures, parameters).persist()
+        case (rd, Some(rf)) =>
+          rd match {
+            case None => setReader(rf.trainingReader)
+            case Some(r) => if (r != rf.trainingReader) log.warn(
+              "Workflow data reader and RawFeatureFilter training reader do not match! " +
+                "The RawFeatureFilter training reader will be used to generate the data for training")
+          }
+          checkReadersAndFeatures()
 
-        val FilteredRawData(cleanedData, featuresToDrop, mapKeysToDrop, rawFeatureFilterResults) =
-          rf.generateFilteredRaw(rawFeatures, parameters)
+          val FilteredRawData(cleanedData, featuresToDrop, mapKeysToDrop, rawFeatureFilterResults) =
+            rf.generateFilteredRaw(rawFeatures, parameters)
 
-        setRawFeatureFilterResults(rawFeatureFilterResults)
-        setBlacklist(featuresToDrop, rawFeatureFilterResults.rawFeatureDistributions)
-        setBlacklistMapKeys(mapKeysToDrop)
-        cleanedData
+          setRawFeatureFilterResults(rawFeatureFilterResults)
+          setBlacklist(featuresToDrop, rawFeatureFilterResults.rawFeatureDistributions)
+          setBlacklistMapKeys(mapKeysToDrop)
+          cleanedData
+      }
     }
   }
 
@@ -400,28 +403,34 @@ class OpWorkflow(val uid: String = UID[OpWorkflow]) extends OpWorkflowCore {
 
     // doing regular workflow fit without workflow level CV
     if (!isWorkflowCV) {
-      FitStagesUtil.fitAndTransformDAG(
-        dag = dag,
-        train = train,
-        test = test,
-        hasTest = hasTest,
-        indexOfLastEstimator = indexOfLastEstimator,
-        persistEveryKStages = persistEveryKStages
-      ).transformers
+      // The cross-validation job group is handled in the appropriate Estimator
+      JobGroupUtil.withJobGroup(OpStep.FeatureEngineering) {
+        FitStagesUtil.fitAndTransformDAG(
+          dag = dag,
+          train = train,
+          test = test,
+          hasTest = hasTest,
+          indexOfLastEstimator = indexOfLastEstimator,
+          persistEveryKStages = persistEveryKStages
+        ).transformers
+      }
     } else {
       // doing workflow level CV/TS
       // Extract Model Selector and Split the DAG into
       val CutDAG(modelSelectorOpt, before, during, after) = FitStagesUtil.cutDAG(dag)
 
       log.info("Applying initial DAG before CV/TS. Stages: {}", before.flatMap(_.map(_._1.stageName)).mkString(", "))
-      val FittedDAG(beforeTrain, beforeTest, beforeTransformers) = FitStagesUtil.fitAndTransformDAG(
-        dag = before,
-        train = train,
-        test = test,
-        hasTest = hasTest,
-        indexOfLastEstimator = indexOfLastEstimator,
-        persistEveryKStages = persistEveryKStages
-      )
+      val FittedDAG(beforeTrain, beforeTest, beforeTransformers) =
+        JobGroupUtil.withJobGroup(OpStep.FeatureEngineering) {
+          FitStagesUtil.fitAndTransformDAG(
+            dag = before,
+            train = train,
+            test = test,
+            hasTest = hasTest,
+            indexOfLastEstimator = indexOfLastEstimator,
+            persistEveryKStages = persistEveryKStages
+          )
+        }
 
       // Break up catalyst (cause it chokes) by converting into rdd, persisting it and then back to dataframe
       val (trainRDD, testRDD) = (beforeTrain.rdd.persist(), beforeTest.rdd.persist())
@@ -437,19 +446,23 @@ class OpWorkflow(val uid: String = UID[OpWorkflow]) extends OpWorkflowCore {
           log.info("Estimate best Model with CV/TS. Stages included in CV are: {}, {}",
             during.flatMap(_.map(_._1.stageName)).mkString(", "), modelSelector.uid: Any
           )
-          modelSelector.findBestEstimator(trainFixed, Option(during))
+          JobGroupUtil.withJobGroup(OpStep.CrossValidation) {
+            modelSelector.findBestEstimator(trainFixed, Option(during))
+          }
           val remainingDAG: StagesDAG = (during :+ (Array(modelSelector -> distance): Layer)) ++ after
 
           log.info("Applying DAG after CV/TS. Stages: {}", remainingDAG.flatMap(_.map(_._1.stageName)).mkString(", "))
-          val fitted = FitStagesUtil.fitAndTransformDAG(
-            dag = remainingDAG,
-            train = trainFixed,
-            test = testFixed,
-            hasTest = hasTest,
-            indexOfLastEstimator = indexOfLastEstimator,
-            persistEveryKStages = persistEveryKStages,
-            fittedTransformers = beforeTransformers
-          ).transformers
+          val fitted = JobGroupUtil.withJobGroup(OpStep.FeatureEngineering) {
+            FitStagesUtil.fitAndTransformDAG(
+              dag = remainingDAG,
+              train = trainFixed,
+              test = testFixed,
+              hasTest = hasTest,
+              indexOfLastEstimator = indexOfLastEstimator,
+              persistEveryKStages = persistEveryKStages,
+              fittedTransformers = beforeTransformers
+            ).transformers
+          }
           trainRDD.unpersist()
           testRDD.unpersist()
           fitted
@@ -480,7 +493,11 @@ class OpWorkflow(val uid: String = UID[OpWorkflow]) extends OpWorkflowCore {
    * @param path to the trained workflow model
    * @return workflow model
    */
-  def loadModel(path: String): OpWorkflowModel = new OpWorkflowModelReader(Some(this)).load(path)
+  def loadModel(path: String)(implicit spark: SparkSession): OpWorkflowModel = {
+    JobGroupUtil.withJobGroup(OpStep.ModelIO) {
+      new OpWorkflowModelReader(Some(this)).load(path)
+    }
+  }
 
   /**
    * Returns a dataframe containing all the columns generated up to and including the feature input

core/src/main/scala/com/salesforce/op/OpWorkflowCore.scala

@@ -37,6 +37,7 @@ import com.salesforce.op.features.types.FeatureType
 import com.salesforce.op.filters.{FeatureDistribution, RawFeatureFilterResults}
 import com.salesforce.op.readers.{CustomReader, Reader, ReaderKey}
 import com.salesforce.op.stages.{FeatureGeneratorStage, OPStage, OpTransformer}
+import com.salesforce.op.utils.spark.{JobGroupUtil, OpStep}
 import com.salesforce.op.utils.spark.RichDataset._
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.ml._
@@ -108,7 +109,7 @@ private[op] trait OpWorkflowCore {
 
   /**
    * Whether the cross-validation/train-validation-split will be done at workflow level
-   *g c
+   *
    * @return true if the cross-validation  will be done at workflow level, false otherwise
    */
   final def isWorkflowCV: Boolean = isWorkflowCVEnabled
@@ -306,7 +307,9 @@ private[op] trait OpWorkflowCore {
   def computeDataUpTo(feature: OPFeature, path: String)
     (implicit spark: SparkSession): Unit = {
     val df = computeDataUpTo(feature)
-    df.saveAvro(path)
+    JobGroupUtil.withJobGroup(OpStep.ResultsSaving) {
+      df.saveAvro(path)
+    }
   }
 
   /**

core/src/main/scala/com/salesforce/op/OpWorkflowModel.scala

@@ -35,6 +35,7 @@ import com.salesforce.op.features.types.FeatureType
 import com.salesforce.op.features.{Feature, FeatureLike, OPFeature}
 import com.salesforce.op.readers.DataFrameFieldNames._
 import com.salesforce.op.stages.{OPStage, OpPipelineStage, OpTransformer}
+import com.salesforce.op.utils.spark.{JobGroupUtil, OpStep}
 import com.salesforce.op.utils.spark.RichDataset._
 import com.salesforce.op.utils.spark.RichMetadata._
 import com.salesforce.op.utils.stages.FitStagesUtil
@@ -91,9 +92,11 @@ class OpWorkflowModel(val uid: String = UID[OpWorkflowModel], val trainingParams
    * @return Dataframe with all the features generated + persisted
    */
   protected def generateRawData()(implicit spark: SparkSession): DataFrame = {
-    require(reader.nonEmpty, "Data reader must be set")
-    checkReadersAndFeatures()
-    reader.get.generateDataFrame(rawFeatures, parameters).persist() // don't want to redo this
+    JobGroupUtil.withJobGroup(OpStep.DataReadingAndFiltering) {
+      require(reader.nonEmpty, "Data reader must be set")
+      checkReadersAndFeatures()
+      reader.get.generateDataFrame(rawFeatures, parameters).persist() // don't want to redo this
+    }
   }
 
   /**
@@ -217,8 +220,11 @@ class OpWorkflowModel(val uid: String = UID[OpWorkflowModel], val trainingParams
    * @param path      path to save the model
    * @param overwrite should overwrite if the path exists
    */
-  def save(path: String, overwrite: Boolean = true): Unit =
-    OpWorkflowModelWriter.save(this, path = path, overwrite = overwrite)
+  def save(path: String, overwrite: Boolean = true)(implicit spark: SparkSession): Unit = {
+    JobGroupUtil.withJobGroup(OpStep.ModelIO) {
+      OpWorkflowModelWriter.save(this, path = path, overwrite = overwrite)
+    }
+  }
 
   /**
    * Gets the fitted stage that generates the input feature
@@ -416,7 +422,9 @@ class OpWorkflowModel(val uid: String = UID[OpWorkflowModel], val trainingParams
     if (persistScores) scores.persist()
 
     // Save the scores if a path was provided
-    path.foreach(scores.saveAvro(_))
+    JobGroupUtil.withJobGroup(OpStep.ResultsSaving) {
+      path.foreach(scores.saveAvro(_))
+    }
 
     scores -> metrics
   }

core/src/main/scala/com/salesforce/op/stages/impl/selector/ModelSelector.scala

@@ -40,6 +40,7 @@ import com.salesforce.op.stages.impl.selector.ModelSelectorNames.ModelType
 import com.salesforce.op.stages.impl.tuning.{BestEstimator, _}
 import com.salesforce.op.stages.sparkwrappers.generic.SparkWrapperParams
 import com.salesforce.op.stages.sparkwrappers.specific.{OpPredictorWrapperModel, SparkModelConverter}
+import com.salesforce.op.utils.spark.{JobGroupUtil, OpStep}
 import com.salesforce.op.utils.spark.RichMetadata._
 import com.salesforce.op.utils.spark.RichParamMap._
 import com.salesforce.op.utils.stages.FitStagesUtil._
@@ -143,13 +144,14 @@ E <: Estimator[_] with OpPipelineStage2[RealNN, OPVector, Prediction]]
   final override def fit(dataset: Dataset[_]): SelectedModel = {
 
     implicit val spark = dataset.sparkSession
+    JobGroupUtil.setJobGroup(OpStep.CrossValidation)
     setInputSchema(dataset.schema).transformSchema(dataset.schema)
     require(!dataset.isEmpty, "Dataset cannot be empty")
     val data = dataset.select(labelColName, in2.name)
-    val (BestEstimator(name, estimator, summary), splitterSummary, datasetWithID) = bestEstimator.map{ e =>
+    val (BestEstimator(name, estimator, summary), splitterSummary, datasetWithID) = bestEstimator.map { e =>
       val PrevalidationVal(summary, dataOpt) = prepareForValidation(data, labelColName)
       (e, summary, dataOpt.getOrElse(data))
-    }.getOrElse{ findBestEstimator(data.toDF()) }
+    }.getOrElse { findBestEstimator(data.toDF()) }
 
     val preparedData = splitter.map(_.validationPrepare(datasetWithID)).getOrElse(datasetWithID)
     val bestModel = estimator.fit(preparedData).asInstanceOf[M]
@@ -158,7 +160,7 @@ E <: Estimator[_] with OpPipelineStage2[RealNN, OPVector, Prediction]]
     log.info(s"With parameters : ${bestEst.extractParamMap()}")
 
     // set input and output params
-    outputsColNamesMap.foreach{ case (pname, pvalue) => bestModel.set(bestModel.getParam(pname), pvalue) }
+    outputsColNamesMap.foreach { case (pname, pvalue) => bestModel.set(bestModel.getParam(pname), pvalue) }
 
     // get eval results for metadata
     val trainingEval = evaluate(bestModel.transform(preparedData))
@@ -183,12 +185,20 @@ E <: Estimator[_] with OpPipelineStage2[RealNN, OPVector, Prediction]]
     val meta = metadataSummary.toMetadata(skipUnsupported = true)
     setMetadata(meta.toSummaryMetadata())
 
-    new SelectedModel(bestModel.asInstanceOf[ModelType], outputsColNamesMap, uid = uid, operationName = operationName)
+    val selectedModel = new SelectedModel(
+      bestModel.asInstanceOf[ModelType],
+      outputsColNamesMap,
+      uid = uid,
+      operationName = operationName
+    )
       .setInput(in1.asFeatureLike[RealNN], in2.asFeatureLike[OPVector])
       .setParent(this)
       .setMetadata(getMetadata())
       .setOutputFeatureName(getOutputFeatureName)
       .setEvaluators(evaluators)
+    // Reset the job group to feature engineering.
+    JobGroupUtil.setJobGroup(OpStep.FeatureEngineering)
+    selectedModel
   }
 
 }

core/src/main/scala/com/salesforce/op/utils/spark/JobGroupUtil.scala

@@ -0,0 +1,60 @@
+/*
+ * Copyright (c) 2017, Salesforce.com, Inc.
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * * Redistributions of source code must retain the above copyright notice, this
+ *   list of conditions and the following disclaimer.
+ *
+ * * Redistributions in binary form must reproduce the above copyright notice,
+ *   this list of conditions and the following disclaimer in the documentation
+ *   and/or other materials provided with the distribution.
+ *
+ * * Neither the name of the copyright holder nor the names of its
+ *   contributors may be used to endorse or promote products derived from
+ *   this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package com.salesforce.op.utils.spark
+
+import org.apache.spark.sql.SparkSession
+
+/**
+ * Convenience methods for working with Spark's job groups.
+ */
+object JobGroupUtil {
+  /**
+   * Sets the Spark job group name and description for a wrapped code block.
+   * The job group is cleared afterwards.
+   *
+   * @param step The OpStep with which to the mark the Spark job group
+   */
+  def withJobGroup[R](step: OpStep)(block: => R)(implicit spark: SparkSession): R = {
+    spark.sparkContext.setJobGroup(step.toString, step.entryDescription)
+    val result = block
+    spark.sparkContext.clearJobGroup()
+    result
+  }
+
+  /**
+   * Indefinitely sets the Spark job group name and description.
+   *
+   * @param step The OpStep with which to the mark the Spark job group
+   */
+  def setJobGroup(step: OpStep)(implicit spark: SparkSession): Unit = {
+    spark.sparkContext.setJobGroup(step.toString, step.entryDescription)
+  }
+}