How to Containerize Models Trained in Spark: MLLib, ONNX and more | by Facundo Santiago | Towards Data Science

If we shall name two trends in modern applications, one should mention containers and machine learning. More and more applications are taking advantage of containerized microservices architectures in order to enable improved elasticity, fault-tolerance, and scalability — whether or on-premise or in the cloud. At the same time, achine Learning have increasingly become a basic requirement in any application.
However, the two things do not always get on each other, and this is the case of Spark. Spark is now-a-days the industry defacto standard for big data analytics, either by using Cloudera, Azure Databricks, AWS EMR, or whatever. But what about if you want to use the model you trained in Spark somewhere else outside the cluster?
There are a couple of options for this:
If you trained your model using TensorFlow, PyTorch, Scikit-Learn, or so, then you can package your model in different portable formats. Since these frameworks don’t depend directly on Spark, you are good to go. I won’t talk about this option in this post.If you trained your model using MLLib (like in the namespace pyspark.ml.*), then you can export your model to a portable format, like ONNX, and then use the ONNX runtime to run the model. This has some limitations since not all the models in MLLib support ONNX currently.If you trained your model using MLLib, then you can persist your model and load it from inside the container by creating a cluster-less Spark context object.(Updated Feb-2022) Use MLFlow to persist you model and package it using the MLModel specification. MLFlow currently supports Spark and it is able to package your model using the MLModel specification. You can use MLFlow to deploy you model wherever you want. For more details about this strategy see my post: Effortless deployments using MLFlow.
If you trained your model using TensorFlow, PyTorch, Scikit-Learn, or so, then you can package your model in different portable formats. Since these frameworks don’t depend directly on Spark, you are good to go. I won’t talk about this option in this post.
If you trained your model using MLLib (like in the namespace pyspark.ml.*), then you can export your model to a portable format, like ONNX, and then use the ONNX runtime to run the model. This has some limitations since not all the models in MLLib support ONNX currently.
If you trained your model using MLLib, then you can persist your model and load it from inside the container by creating a cluster-less Spark context object.
(Updated Feb-2022) Use MLFlow to persist you model and package it using the MLModel specification. MLFlow currently supports Spark and it is able to package your model using the MLModel specification. You can use MLFlow to deploy you model wherever you want. For more details about this strategy see my post: Effortless deployments using MLFlow.
ONNX is an open-standard for serialization and specification of a machine learning model. Since the format describes the computation graph (input, output and operations), it is self-contained. It is focused on deep-learning and it is championed mainly by Microsoft and Facebook. Supported in many frameworks like TensorFlow and PyTorch.
In Spark, an ML model is represented by aTransformer which transforms an inputDataFrame with features into another DataFrame with predictions. A Transformer can also output another set of features. Since in Machine Learning you usually have to perform several steps to produce the desired prediction, Spark has aPipelineModelobject which is a sequence of steps or Transformers(note that the pipeline object is a bit more complex, but for the sake of simplicity of this post, we will assume it as that). It chains multiple Transformers together to specify a ML workflow.
Consider the following example from spark.apache.org where we have a PipelineModelwith 3 stages: A tokenizer (which takes text and return words), a Hashing function (which takes words and returns vectors) and a Logistic Regression Model (which takes features and return predictions).
You get a PipelineModel by training a Pipeline using the method fit(). Here you have an example:
tokenizer = Tokenizer(inputCol="text", outputCol="words")hashingTF = HashingTF(inputCol=tokenizer.getOutputCol(), outputCol="features")lr = LogisticRegression(maxIter=10, regParam=0.01)pipeline = Pipeline(stages=[tokenizer, hashingTF, lr])model = pipeline.fit(training)
Now the question is, how to run this PipelineModel object outside Spark?
ONNX is designed for deep-learning models, however, it supports in some extends more “traditional” machine learning techniques. Spark is commonly used for those more traditional approaches. In Spark this includes:
Vectorizers and encodings (String indexing, OneHotEncoding, Word2Vec)
Scalers (Scalers, Inputer, Binarizer, Bucketizer)
Models: Linear models, Random Forest, Gradient Boosted, Naive Bayes, SVM, PCA
To export your model to an ONNX format, you need to install first onnxmltools which is available right now only for PySpark. Depending on the platform you are using the way you will install the library. In my case, I’m using Azure Databricks, so I will install this library using the Install library option in the cluster.
Once the library is installed, then you can export your pipeline model into ONNX by:
from onnxmltools import convert_sparkmlfrom onnxmltools.convert.sparkml.utils import buildInitialTypesSimpleinitial_types = buildInitialTypesSimple(test_df.drop("label"))onnx_model = convert_sparkml(model, 'Pyspark model', initial_types, spark_session = spark)
Where,
the buildInitialTypesSimple function creates a list of all the expected inputs from the model (features). It takes a sample DataFrame as parameter. In my case, I’m using the test_df dataframe and dropping the label column to retain only all the features.
model is the name of the fitted Pipeline (PipelineModel)
‘Pyspark model’ is the description of the model
initial_types are the expected input feature names an types. This can be supplied, as I said, by using the buildInitialTypesSimplefunction, or by constructing by hand, like [ (‘education’, StringTensorType([1, 1]))]
spark_session=sparkpass the SparkSession context to the method.
Note: This last parameter isn’t documented in the GitHub repository but after some debuging I needed to specify it to avoid a null spark context being used. In my case I’m using Azure Databricks, so I don’t know if this is specific to Databricks or not.
Once the model is converted, you can persist it into a file by:
with open(os.path.join("/tmp/", "model.onnx"), "wb") as f:    f.write(onnx_model.SerializeToString())
You can load this model easily using ONNX runtime in Python like:
import onnxruntime session = onnxruntime.InferenceSession(model_file_path, None)output = session.get_outputs()[0]inputs = session.get_inputs()input_data= {i.name: v for i, v in zip(inputs, input_sample.values.reshape(len(inputs),1,1).astype(np.float32))}}results = session.run([output.name], input_data)
Note: input_data is a dictionary with keys as feature names, and Tensors (1,1) as values. The reshape function transform the input to an array of Tensor with shape (feature_count,1,1), which is expected. It is also important to cast the values as float32.
You then can ship this model in a docker container with Python and the following libraries installed:
onnxruntime
Your scoring file will load the model using the onnxruntime.InferenceSession() method. You usually perform this once. Your scoring routine will call session.run() on the other hand.
I have a sample scoring file in the following GitHub link.
By the time writing this post, the following features are missing in the Spark to ONNX converter:
Missing exporters for Feature Hashing, TFIDF, RFormula, NGram, SQLTransformer and some models (clustering, FP, ALS, etc)
Support only for PySpark
Some issues: Tokenizer is not supported in the ONNX specification
Another way to run a PipelineModel inside of a container is to export the model and create a Spark context inside of the container even when there is not cluster available.
If you want to persist this PipelineModel, the object implements the interface MLWritable, which exponses the method save(path) and write().overwrite().save(path). This method saves the model in a folder structure like the following:
In this structure, the PipelineModel is persisted inside a folder structure where all the steps in the pipeline are detailed. In this case, this model I created has 2 stages: a VectorAssembler and a GradientBoostedRegressor. Opposite as what you may be used to, all the structure is required to load and restore the trained model. If you are using a model registry, like MLFlow or Azure Machine Learning Services, I suggest you zip the directory in an archive. The following line helps to do so:
import shutilmodel.write().overwrite().save(model_path)path_drv = shutil.make_archive(model_name, format='zip', base_dir=model_path)
Please notice that shutil.make_archive will create the file in the driver node in its local file system.
Since you will be loading the Spark model directly, you will need to install pyspark Python library in the container image. Then in your scoring script you will create a spark session, unpack the archive in a folder and load the PipelineModel object.
import pysparkfrom pyspark.ml import PipelineModelspark = pyspark.sql.SparkSession               .builder.appName("pyspark_runtime").getOrCreate()model_unpacked = "./" + model_nameshutil.unpack_archive(model_path, model_unpacked)trainedModel = PipelineModel.load(model_unpacked)
The variables spark and trainedModel have to be available in all the routines. Personally, I declare those two as global variables.
Then, you run your model by:
input_df = spark.createDataFrame(pandas_df)predictions = trainedModel.transform(input_df).collect()preds = [x['prediction'] for x in predictioprint('[INFO] Results was ' + json.dumps(preds))
I have a sample scoring file in the following GitHub link.
We reviewed in this post two different ways to port models trained inside a Spark cluster to run them inside a container. Using the open-standard ONNX or loading the persisted model in an Spark context. You can check the GitHub repository if you want to see the complete code for the examples. There is also a Databricks notebook to generate both the ONNX file and the MLLib zip file. Hope it helps.