update barrier mode

Signed-off-by: minmingzhu <[email protected]>

mllib-dal/src/main/scala/com/intel/oap/mllib/OneDAL.scala

@@ -409,29 +409,28 @@ object OneDAL {
 
     val spark = SparkSession.active
     import spark.implicits._
-    val labeledPointsRDD = labeledPoints.select(labelCol, featuresCol).toDF().rdd
-      .barrier().mapPartitions(iter => iter)
+    val labeledPointsRDD = labeledPoints.rdd
 
     // Repartition to executorNum if not enough partitions
     val dataForConversion = if (labeledPointsRDD.getNumPartitions < executorNum) {
       logger.info(s"Repartition to executorNum if not enough partitions")
-      val rePartitions = labeledPointsRDD.repartition(executorNum).cache()
+      val rePartitions = labeledPoints.repartition(executorNum).cache()
       rePartitions.count()
       rePartitions
     } else {
-      labeledPointsRDD
+      labeledPoints
     }
 
     // Get dimensions for each partition
-    val partitionDims = Utils.getPartitionDims(dataForConversion.map{ row =>
-      val vector = row.getAs[Vector](1)
+    val partitionDims = Utils.getPartitionDims(dataForConversion.select(featuresCol).rdd.map{ row =>
+      val vector = row.getAs[Vector](0)
       vector
     })
 
     // Filter out empty partitions, if there is no such rdd, coalesce will report an error
     // "No partitions or no locations for partitions found".
     // TODO: ML-312: Improve ExecutorInProcessCoalescePartitioner
-    val nonEmptyPartitions = dataForConversion.mapPartitionsWithIndex {
+    val nonEmptyPartitions = dataForConversion.select(labelCol, featuresCol).toDF().rdd.mapPartitionsWithIndex {
       (index: Int, it: Iterator[Row]) => Iterator(Tuple3(partitionDims(index)._1, index, it))
     }.filter {
       _._1 > 0
@@ -679,14 +678,12 @@ object OneDAL {
     require(executorNum > 0)
 
     logger.info(s"Processing partitions with $executorNum executors")
-    val barrierRDD = vectors.barrier()mapPartitions(iter => iter)
-
 
     // Repartition to executorNum if not enough partitions
     val dataForConversion = if (vectors.getNumPartitions < executorNum) {
-      barrierRDD.repartition(executorNum).setName("Repartitioned for conversion").cache()
+      vectors.repartition(executorNum).setName("Repartitioned for conversion").cache()
     } else {
-      barrierRDD
+      vectors
     }
 
     // Get dimensions for each partition

mllib-dal/src/main/scala/com/intel/oap/mllib/clustering/KMeansDALImpl.scala

@@ -42,11 +42,11 @@ class KMeansDALImpl(var nClusters: Int,
     val useDevice = sparkContext.getConf.get("spark.oap.mllib.device", Utils.DefaultComputeDevice)
     val computeDevice = Common.ComputeDevice.getDeviceByName(useDevice)
     kmeansTimer.record("Preprocessing")
-
+    val barrierRDD = data.barrier().mapPartitions(iter => iter)
     val coalescedTables = if (useDevice == "GPU") {
-      OneDAL.coalesceVectorsToHomogenTables(data, executorNum, computeDevice)
+      OneDAL.coalesceVectorsToHomogenTables(barrierRDD, executorNum, computeDevice)
     } else {
-      OneDAL.coalesceVectorsToNumericTables(data, executorNum)
+      OneDAL.coalesceVectorsToNumericTables(barrierRDD, executorNum)
     }
     kmeansTimer.record("Data Convertion")
 
@@ -95,7 +95,7 @@ class KMeansDALImpl(var nClusters: Int,
         }
       OneCCL.cleanup()
       ret
-    }.barrier().mapPartitions(iter => iter).collect()
+    }.collect()
 
     // Make sure there is only one result from rank 0
     assert(results.length == 1)

mllib-dal/src/main/scala/com/intel/oap/mllib/feature/PCADALImpl.scala

@@ -45,16 +45,18 @@ class PCADALImpl(val k: Int,
   def train(data: RDD[Vector]): PCADALModel = {
     val pcaTimer = new Utils.AlgoTimeMetrics("PCA")
     val normalizedData = normalizeData(data)
+    val barrierRDD = normalizedData.barrier().mapPartitions(iter => iter)
+
     val sparkContext = normalizedData.sparkContext
     val useDevice = sparkContext.getConf.get("spark.oap.mllib.device", Utils.DefaultComputeDevice)
     val computeDevice = Common.ComputeDevice.getDeviceByName(useDevice)
     pcaTimer.record("Preprocessing")
 
     val coalescedTables = if (useDevice == "GPU") {
-      OneDAL.coalesceVectorsToHomogenTables(normalizedData, executorNum,
+      OneDAL.coalesceVectorsToHomogenTables(barrierRDD, executorNum,
         computeDevice)
     } else {
-      OneDAL.coalesceVectorsToNumericTables(normalizedData, executorNum)
+      OneDAL.coalesceVectorsToNumericTables(barrierRDD, executorNum)
     }
     val kvsIPPort = getOneCCLIPPort(coalescedTables)
     pcaTimer.record("Data Convertion")

mllib-dal/src/main/scala/com/intel/oap/mllib/recommendation/ALSDALImpl.scala

@@ -72,8 +72,9 @@ class ALSDALImpl[@specialized(Int, Long) ID: ClassTag]( data: RDD[Rating[ID]],
 
     logInfo(s"ALSDAL fit using $executorNum Executors " +
       s"for $nVectors vectors and $nFeatures features")
+    val barrierRDD = data.barrier().mapPartitions(iter => iter)
 
-    val numericTables = data.repartition(executorNum)
+    val numericTables = barrierRDD.repartition(executorNum)
       .setName("Repartitioned for conversion").cache()
 
     val kvsIPPort = getOneCCLIPPort(numericTables)

mllib-dal/src/main/scala/com/intel/oap/mllib/stat/CorrelationDALImpl.scala

@@ -35,12 +35,13 @@ class CorrelationDALImpl(
     val useDevice = sparkContext.getConf.get("spark.oap.mllib.device", Utils.DefaultComputeDevice)
     val computeDevice = Common.ComputeDevice.getDeviceByName(useDevice)
     corTimer.record("Preprocessing")
+    val barrierRDD = data.barrier().mapPartitions(iter => iter)
 
     val coalescedTables = if (useDevice == "GPU") {
-      OneDAL.coalesceVectorsToHomogenTables(data, executorNum,
+      OneDAL.coalesceVectorsToHomogenTables(barrierRDD, executorNum,
         computeDevice)
     } else {
-      OneDAL.coalesceVectorsToNumericTables(data, executorNum)
+      OneDAL.coalesceVectorsToNumericTables(barrierRDD, executorNum)
     }
     corTimer.record("Data Convertion")
 

mllib-dal/src/main/scala/com/intel/oap/mllib/stat/SummarizerDALImpl.scala

@@ -36,12 +36,13 @@ class SummarizerDALImpl(val executorNum: Int,
     val useDevice = sparkContext.getConf.get("spark.oap.mllib.device", Utils.DefaultComputeDevice)
     val computeDevice = Common.ComputeDevice.getDeviceByName(useDevice)
     sumTimer.record("Preprocessing")
+    val barrierRDD = data.barrier().mapPartitions(iter => iter)
 
     val coalescedTables = if (useDevice == "GPU") {
-      OneDAL.coalesceVectorsToHomogenTables(data, executorNum,
+      OneDAL.coalesceVectorsToHomogenTables(barrierRDD, executorNum,
         computeDevice)
     } else {
-      OneDAL.coalesceVectorsToNumericTables(data, executorNum)
+      OneDAL.coalesceVectorsToNumericTables(barrierRDD, executorNum)
     }
     sumTimer.record("Data Convertion")
 

mllib-dal/src/test/scala/org/apache/spark/ml/clustering/MLlibKMeansSuite.scala

@@ -47,6 +47,7 @@ class MLlibKMeansSuite extends MLTest with DefaultReadWriteTest with PMMLReadWri
   override def sparkConf: SparkConf = {
     val conf = super.sparkConf
     conf.set("spark.oap.mllib.device", TestCommon.getComputeDevice.toString)
+    conf.set("spark.driver.bindAddress", "10.239.34.1");
   }
 
   test("default parameters") {

mllib-dal/src/test/scala/org/apache/spark/ml/feature/MLlibPCASuite.scala

@@ -32,6 +32,8 @@ class MLlibPCASuite extends MLTest with DefaultReadWriteTest {
   override def sparkConf: SparkConf = {
     val conf = super.sparkConf
     conf.set("spark.oap.mllib.device", TestCommon.getComputeDevice.toString)
+    conf.set("spark.driver.bindAddress", "10.239.34.1");
+
   }
 
   test("params") {