Add serializers/desirializers for models

Author: blublinsky

src/main/scala/com/lightbend/modelServer/ModelServingFlatJob.scala

@@ -4,6 +4,7 @@ import java.util.Properties
 
 import com.lightbend.kafka.ModelServingConfiguration
 import com.lightbend.model.winerecord.WineRecord
+import com.lightbend.modelServer.model.{PMMLModelSerializerKryo, TensorFlowModel, TensorFlowModelSerializerKryo}
 import com.lightbend.modelServer.typeschema.ByteArraySchema
 import org.apache.flink.api.scala._
 import org.apache.flink.configuration.{ConfigConstants, Configuration, QueryableStateOptions}
@@ -80,6 +81,10 @@ object ModelServingFlatJob {
   def buildGraph(env : StreamExecutionEnvironment) : Unit = {
     env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime)
     env.enableCheckpointing(5000)
+    // Add custom serializers
+    env.getConfig.addDefaultKryoSerializer(TensorFlowModel.getClass, classOf[TensorFlowModelSerializerKryo])
+    env.getConfig.addDefaultKryoSerializer(TensorFlowModel.getClass, classOf[PMMLModelSerializerKryo])
+
 
     // configure Kafka consumer
     // Data

src/main/scala/com/lightbend/modelServer/ModelServingKeyedJob.scala

@@ -4,12 +4,12 @@ import java.util.Properties
 
 import com.lightbend.kafka.ModelServingConfiguration
 import com.lightbend.model.winerecord.WineRecord
+import com.lightbend.modelServer.model.{PMMLModelSerializerKryo, TensorFlowModel, TensorFlowModelSerializerKryo}
 import org.apache.flink.api.scala._
 import com.lightbend.modelServer.typeschema.ByteArraySchema
 import org.apache.flink.streaming.api.TimeCharacteristic
 import org.apache.flink.streaming.api.scala.StreamExecutionEnvironment
 import org.apache.flink.streaming.connectors.kafka.FlinkKafkaConsumer010
-
 import org.apache.flink.configuration.Configuration
 import org.apache.flink.configuration.ConfigConstants
 import org.apache.flink.configuration.QueryableStateOptions
@@ -84,6 +84,9 @@ object ModelServingKeyedJob {
   def buildGraph(env : StreamExecutionEnvironment) : Unit = {
     env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime)
     env.enableCheckpointing(5000)
+    // Add custom serializer
+    env.getConfig.addDefaultKryoSerializer(TensorFlowModel.getClass, classOf[TensorFlowModelSerializerKryo])
+    env.getConfig.addDefaultKryoSerializer(TensorFlowModel.getClass, classOf[PMMLModelSerializerKryo])
 
     // configure Kafka consumer
     // Data

src/main/scala/com/lightbend/modelServer/model/PMMLModel.scala

@@ -9,7 +9,7 @@ package com.lightbend.modelServer.model
 import org.jpmml.evaluator.{FieldValue, ModelEvaluatorFactory, TargetField}
 import org.jpmml.evaluator.visitors._
 import org.jpmml.model.PMMLUtil
-import java.io.{ByteArrayInputStream, InputStream}
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream, InputStream}
 
 import org.jpmml.evaluator.Computable
 import com.lightbend.model.winerecord.WineRecord
@@ -67,6 +67,12 @@ class PMMLModel(inputStream: Array[Byte]) extends Model {
     case _ => .0
   }
 
+  def toBytes : Array[Byte] = {
+    var stream = new ByteArrayOutputStream()
+    PMMLUtil.marshal(pmml, stream)
+    stream.toByteArray
+  }
+
 }
 
 object PMMLModel{

src/main/scala/com/lightbend/modelServer/model/PMMLModelSerializerKryo.scala

@@ -0,0 +1,41 @@
+package com.lightbend.modelServer.model
+
+/**
+  * Created by boris on 6/2/17.
+  */
+import com.esotericsoftware.kryo.io.{Input, Output}
+import com.esotericsoftware.kryo.{Kryo, Serializer}
+
+
+class PMMLModelSerializerKryo extends Serializer[PMMLModel]{
+  
+  super.setAcceptsNull(false)
+  super.setImmutable(true)
+
+  /** Reads bytes and returns a new object of the specified concrete type.
+    * <p>
+    * Before Kryo can be used to read child objects, {@link Kryo#reference(Object)} must be called with the parent object to
+    * ensure it can be referenced by the child objects. Any serializer that uses {@link Kryo} to read a child object may need to
+    * be reentrant.
+    * <p>
+    * This method should not be called directly, instead this serializer can be passed to {@link Kryo} read methods that accept a
+    * serialier.
+    *
+    * @return May be null if { @link #getAcceptsNull()} is true. */
+  
+  override def read(kryo: Kryo, input: Input, `type`: Class[PMMLModel]): PMMLModel = {
+    val bytes = Stream.continually(input.readByte()).takeWhile(_ != -1).toArray
+    PMMLModel(bytes).get
+  }
+
+  /** Writes the bytes for the object to the output.
+    * <p>
+    * This method should not be called directly, instead this serializer can be passed to {@link Kryo} write methods that accept a
+    * serialier.
+    *
+    * @param value May be null if { @link #getAcceptsNull()} is true. */
+  
+  override def write(kryo: Kryo, output: Output, value: PMMLModel): Unit = {
+    output.write(value.toBytes)
+  }
+}

src/main/scala/com/lightbend/modelServer/model/TensorFlowModel.scala

@@ -1,6 +1,5 @@
 package com.lightbend.modelServer.model
 
-import java.io.InputStream
 
 import com.lightbend.model.winerecord.WineRecord
 import org.tensorflow.{Graph, Session, Tensor}
@@ -9,11 +8,12 @@ import org.tensorflow.{Graph, Session, Tensor}
   * Created by boris on 5/26/17.
   * Implementation of tensorflow model
   */
-class TensorFlowModel(inputStream: Array[Byte]) extends Model  {
+
+class TensorFlowModel(inputStream : Array[Byte]) extends Model  {
 
   val graph = new Graph
   graph.importGraphDef(inputStream)
-  val session = new Session (graph)
+  val session = new Session(graph)
 
   override def score(input: AnyVal): AnyVal = {
 
@@ -45,8 +45,16 @@ class TensorFlowModel(inputStream: Array[Byte]) extends Model  {
   }
 
   override def cleanup(): Unit = {
-    session.close
-    graph.close
+    try{
+      session.close
+    }catch {
+      case t: Throwable =>    // Swallow
+    }
+    try{
+      graph.close
+    }catch {
+      case t: Throwable =>    // Swallow
+    }
   }
 }
 

src/main/scala/com/lightbend/modelServer/model/TensorFlowModelSerializerKryo.scala

@@ -0,0 +1,41 @@
+package com.lightbend.modelServer.model
+
+/**
+  * Created by boris on 6/2/17.
+  */
+import com.esotericsoftware.kryo.{Kryo, Serializer}
+import com.esotericsoftware.kryo.io.{Input, Output}
+
+
+class TensorFlowModelSerializerKryo extends Serializer[TensorFlowModel]{
+  
+  super.setAcceptsNull(false)
+  super.setImmutable(true)
+
+  /** Reads bytes and returns a new object of the specified concrete type.
+    * <p>
+    * Before Kryo can be used to read child objects, {@link Kryo#reference(Object)} must be called with the parent object to
+    * ensure it can be referenced by the child objects. Any serializer that uses {@link Kryo} to read a child object may need to
+    * be reentrant.
+    * <p>
+    * This method should not be called directly, instead this serializer can be passed to {@link Kryo} read methods that accept a
+    * serialier.
+    *
+    * @return May be null if { @link #getAcceptsNull()} is true. */
+  
+  override def read(kryo: Kryo, input: Input, `type`: Class[TensorFlowModel]): TensorFlowModel = {
+    val bytes = Stream.continually(input.readByte()).takeWhile(_ != -1).toArray
+    TensorFlowModel(bytes).get
+  }
+
+  /** Writes the bytes for the object to the output.
+    * <p>
+    * This method should not be called directly, instead this serializer can be passed to {@link Kryo} write methods that accept a
+    * serialier.
+    *
+    * @param value May be null if { @link #getAcceptsNull()} is true. */
+  
+  override def write(kryo: Kryo, output: Output, value: TensorFlowModel): Unit = {
+    output.write(value.graph.toGraphDef)
+  }
+}