Adding params for optimization level and device type; adding support …

…for distributing loads on multiple GPUs

deep-learning/src/main/scala/com/microsoft/ml/spark/GPUUtils.scala

@@ -0,0 +1,23 @@
+package com.microsoft.ml.spark
+
+import org.apache.spark.TaskContext
+
+import scala.util.Random
+
+object GPUUtils {
+  @transient
+  private val Rand = new Random()
+
+  /**
+   * Select a random GPU device on an executor node.
+   * @return None if there is no GPU devices found; a randomly selected device id otherwise.
+   */
+  def selectRandomGpu: Option[Int] = {
+    TaskContext.get().resources().get("gpu").map{
+      resource =>
+        val addresses = resource.addresses.map(_.toInt)
+        val i = Rand.nextInt() % addresses.length
+        addresses(i)
+    }
+  }
+}

deep-learning/src/main/scala/com/microsoft/ml/spark/SharedParams.scala

@@ -3,7 +3,7 @@
 
 package com.microsoft.ml.spark
 
-import org.apache.spark.ml.param.{MapParam, Params}
+import org.apache.spark.ml.param.{MapParam, Param, Params}
 import spray.json.DefaultJsonProtocol._
 
 trait HasFeedFetchDicts extends Params {

deep-learning/src/main/scala/com/microsoft/ml/spark/onnx/ONNXModel.scala

@@ -8,14 +8,14 @@ import ai.onnxruntime.OrtSession.SessionOptions
 import ai.onnxruntime.OrtSession.SessionOptions.OptLevel
 import ai.onnxruntime._
 import breeze.linalg.{argmax, softmax, DenseVector => BDV}
-import com.microsoft.ml.spark.HasFeedFetchDicts
 import com.microsoft.ml.spark.codegen.Wrappable
 import com.microsoft.ml.spark.core.env.StreamUtilities.using
 import com.microsoft.ml.spark.core.schema.DatasetExtensions
 import com.microsoft.ml.spark.core.utils.BreezeUtils._
 import com.microsoft.ml.spark.logging.BasicLogging
 import com.microsoft.ml.spark.stages._
-import org.apache.spark.{SparkContext, TaskContext}
+import com.microsoft.ml.spark.{GPUUtils, HasFeedFetchDicts}
+import org.apache.spark.SparkContext
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.injections.UDFUtils
 import org.apache.spark.internal.Logging
@@ -24,10 +24,10 @@ import org.apache.spark.ml.linalg.SQLDataTypes._
 import org.apache.spark.ml.linalg.{SQLDataTypes, Vector}
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.encoders.RowEncoder
 import org.apache.spark.sql.functions.col
 import org.apache.spark.sql.types._
-import org.apache.spark.sql._
 import spray.json.DefaultJsonProtocol._
 
 import java.nio._
@@ -73,7 +73,34 @@ trait ONNXModelParams extends Params with HasMiniBatcher with HasFeedFetchDicts
 
   def getArgMaxDict: Map[String, String] = get(argMaxDict).getOrElse(Map.empty)
 
+  val deviceType: Param[String] = new Param[String](
+    this,
+    "deviceType",
+    "Specify a device type the model inference runs on. Supported types are: CPU or CUDA." +
+      "If not specified, auto detection will be used.",
+    ParamValidators.inArray(Array("CPU", "CUDA"))
+  )
+
+  def getDeviceType: String = $(deviceType)
+
+  def setDeviceType(value: String): this.type = set(deviceType, value)
+
+  //noinspection ScalaStyle
+  val optimizationLevel: Param[String] = new Param[String](
+    this,
+    "optimizationLevel",
+    "Specify the optimization level for the ONNX graph optimizations. Details at " +
+      "https://onnxruntime.ai/docs/resources/graph-optimizations.html#graph-optimization-levels. " +
+      "Supported values are: NO_OPT; BASIC_OPT; EXTENDED_OPT; ALL_OPT. Default: ALL_OPT.",
+    ParamValidators.inArray(Array("NO_OPT", "BASIC_OPT", "EXTENDED_OPT", "ALL_OPT"))
+  )
+
+  def getOptimizationLevel: String = $(optimizationLevel)
+
+  def setOptimizationLevel(value: String): this.type = set(optimizationLevel, value)
+
   setDefault(
+    optimizationLevel -> "ALL_OPT",
     miniBatcher -> new FixedMiniBatchTransformer().setBatchSize(10) //scalastyle:ignore magic.number
   )
 }
@@ -142,7 +169,10 @@ private class ClosableIterator[+T](delegate: Iterator[T], cleanup: => Unit) exte
  * MapInfo defines keyType, valueType and size. It is usually used inside SequenceInfo.
  */
 object ONNXModel extends ComplexParamsReadable[ONNXModel] with Logging {
-  private[onnx] def initializeOrt(modelContent: Array[Byte], ortEnv: OrtEnvironment, gpuDeviceId: Option[Int] = None)
+  private[onnx] def initializeOrt(modelContent: Array[Byte],
+                                  ortEnv: OrtEnvironment,
+                                  optLevel: OptLevel = OptLevel.ALL_OPT,
+                                  gpuDeviceId: Option[Int] = None)
   : OrtSession = {
     val options = new SessionOptions()
 
@@ -157,7 +187,7 @@ object ONNXModel extends ComplexParamsReadable[ONNXModel] with Logging {
         logError(err)
     }
 
-    options.setOptimizationLevel(OptLevel.BASIC_OPT)
+    options.setOptimizationLevel(optLevel)
     ortEnv.createSession(modelContent, options)
   }
 
@@ -243,17 +273,21 @@ object ONNXModel extends ComplexParamsReadable[ONNXModel] with Logging {
     }
   }
 
-  //noinspection ScalaStyle
-  private[onnx] def applyModel(modelPayloadBC: Broadcast[Array[Byte]],
+  private[onnx] def selectGpuDevice(deviceType: Option[String]): Option[Int] = {
+    deviceType match {
+      case None | Some("CUDA") =>
+        GPUUtils.selectRandomGpu
+      case Some("CPU") =>
+        None
+    }
+  }
+
+  private[onnx] def applyModel(session: OrtSession,
+                               env: OrtEnvironment,
                                feedMap: Map[String, String],
                                fetchMap: Map[String, String],
                                inputSchema: StructType
                               )(rows: Iterator[Row]): Iterator[Row] = {
-
-    val payload = modelPayloadBC.value
-    val env = OrtEnvironment.getEnvironment
-    val session = initializeOrt(payload, env, GpuDeviceId)
-
     val results = rows.map {
       row =>
         // Each row contains a batch
@@ -352,9 +386,6 @@ object ONNXModel extends ComplexParamsReadable[ONNXModel] with Logging {
       case (fromDataType, toDataType) => fromDataType == toDataType
     }
   }
-
-  @transient
-  lazy val GpuDeviceId: Option[Int] = TaskContext.get().resources().get("gpu").map(_.addresses.head.toInt)
 }
 
 class ONNXModel(override val uid: String)
@@ -424,10 +455,17 @@ class ONNXModel(override val uid: String)
     val batchedDF = getMiniBatcher.transform(dataset)
     val batchedCache = if (batchedDF.isStreaming) batchedDF else batchedDF.cache().unpersist()
     val (coerced, feedDict) = coerceBatchedDf(batchedCache)
-    val fetchDict = getFetchDict
     val modelBc = broadcastedModelPayload.getOrElse(rebroadcastModelPayload(dataset.sparkSession))
-    val modelFunc = applyModel(modelBc, feedDict, fetchDict, inputSchema) _
-    val outputDf = coerced.mapPartitions(modelFunc)
+    val (fetchDicts, devType, optLevel) = (getFetchDict, get(deviceType), OptLevel.valueOf(getOptimizationLevel))
+    val outputDf = coerced.mapPartitions {
+      rows =>
+        val payload = modelBc.value
+        val gpuDeviceId = selectGpuDevice(devType)
+        val env = OrtEnvironment.getEnvironment
+        logInfo(s"DeviceType=$devType;DeviceId=$gpuDeviceId;OptimizationLevel=$optLevel")
+        val session = initializeOrt(payload, env, optLevel, gpuDeviceId)
+        applyModel(session, env, feedDict, fetchDicts, inputSchema)(rows)
+    }
 
     // The cache call is a workaround for GH issue 1075:
     //  https://github.com/Azure/mmlspark/issues/1075