Empowering Spark with MLflow. This post aims to cover our initial... | by Albert Franzi | Towards Data Science

This post aims to cover our initial experience using MLflow.
We will start discovering MLflow with its own tracking server by logging all the exploratory iterations. Then, we will show our experience linking Spark with MLflow using UDFs.
At Alpha Health we leverage the power of Machine Learning and AI to empower people in taking control of their health and wellbeing. Machine Learning models are hence at the core of the data products we are developing, and this is why MLFLow, an open-source platform covering all aspects of the ML lifecycle drew our attention.
The main goal of MLflow is to provide an extra layer on top of ML allowing Data Scientists to work with virtually any Machine Learning library (h2o, keras, mleap, pytorch, sklearn & tensorflow) meanwhile, it brings their job to another level.
MLflow provides three components:
Tracking — Recording and querying experiments: code, data, config and results. Very useful to keep track of your modelling progress.
Projects — Packaging format for reproducible runs on any platform (i.e Sagemaker).
Models — General format for sending models to diverse deployment tools.
MLflow (currently in alpha) is an open source platform to manage the ML lifecycle, including experimentation, reproducibility and deployment.
In order to use MLflow, we first need to set up all the Python Environment to use MLflow, we would use PyEnv (to setup → Python in Mac). This will provide a virtual environment where we can install all the libraries required to run it.
pyenv install 3.7.0
pyenv global 3.7.0 # Use Python 3.7
mkvirtualenv mlflow # Create a Virtual Env with Python 3.7
workon mlflow

Install required libraries.
pip install mlflow==0.7.0 \
            Cython==0.29 \ 
            numpy==1.14.5 \
            pandas==0.23.4 \
            pyarrow==0.11.0

Note: We are using PyArrow to launch models as UDFs. PyArrow and Numpy versions required fixing since the latest ones had some conflicts between them.
MLflow Tracking allows us to log and query experiments using both Python and REST APIs. Besides, it’s possible to define where are we going to store the model artifacts (Localhost, Amazon S3, Azure Blob Storage, Google Cloud Storage or SFTP server). Since we are using AWS in Alpha Health, we will be experimenting with S3 as an artifact storage.
MLflow recommends using a persistent file-store. The file-store is where the server stores run and experiment metadata. So, when running a server, make sure that this points to a persistent file system location. Here, we are just using /tmp for the experiment.
Keep in mind that if we want to use the mlflow server to run old experiments, they must be present in the file-store. However, without them, we would still be able to use them in UDFs since only the model path would be required.
Note: Keep in mind the Tracking UI and the model client has to have access to the artifact location. That means, regardless of the Tracking UI being on an EC2 instance, if we run MLflow locally, our machine should have direct access to S3 to write the artifact models.
Once the Tracking Server is running, we can start training our models.
As an example we’ll use a modification of the wine example provided in the MLflow Sklearn example.
As mentioned before, MLflow allows logging the params, metrics and artifacts of our models, so we can keep track of how these evolve over our different iterations. This feature is quite useful, since we will be able to reproduce our best model by checking the Tracking Server or validate which code was executing the desired iteration since it logs (for free) the git hash commit.
The MLflow tracking server ,launched using “mlflow server”, also hosts REST APIs for tracking runs and writing data to the local filesystem. You can specify a tracking server URI with the “MLFLOW_TRACKING_URI” environment variable and MLflow tracking APIs automatically communicate with the tracking server at that URI to create/get run information, log metrics, and so on.
Ref: Docs// Running a tracking server
In order to serve a model, all we need is a Tracking Server running (see Launching UI) and a model run ID.
To serve models with the MLflow serve feature we need access to the Tracking UI, so it can retrieve the model information by just specifying the--run_id .
Once the model is being served by the Tracking Server, we can query our new model endpoint.
Although, it’s quite powerful to be able to serve models in real time by just training the model and using the serve feature (ref: mlflow // docs // models # local), applying models using Spark (batch or streaming) is even more powerful since it joins the distribution force.
Imagine having an offline traning and then applying the output model to all your data in a easier way. This is where Spark and MLflow shine to the perfection together.
Ref: Get started PySpark — Jupyter
To show how we are applying the MLflow models to our Spark Dataframes. We need to setup Jupyter notebooks with PySpark.
Start by downloading the latest stable Apache Spark (current 2.4.3).
cd ~/Downloads/
tar -xzf spark-2.4.3-bin-hadoop2.7.tgz
mv ~/Downloads/spark-2.4.3-bin-hadoop2.7 ~/
ln -s ~/spark-2.4.3-bin-hadoop2.7 ~/spark̀

Install PySpark and Jupyter in our virtualEnv
pip install pyspark jupyter

Setup Environmnet variables
export SPARK_HOME=~/spark
export PATH=$SPARK_HOME/bin:$PATH
export PYSPARK_DRIVER_PYTHON=jupyter
export PYSPARK_DRIVER_PYTHON_OPTS="notebook --notebook-dir=${HOME}/Projects/notebooks"

By defining notebook-dir, we will be able to store & persist our notebooks in the desired folder.
Since we configured Jupyter as the PySpark driver, now we can launch Jupyter with a PySpark context attached to our notebooks.
As we have shown above, MLflow provides a feature to log our model artifacts to S3. So, once we have the chosen one model, we can import it as a UDF by using the mlflow.pyfunc module.
Up to this point, we have shown how to use PySpark with MLflow by running a wine quality prediction across all our wine dataset. But what happens when you need to use the Python MLflow modules from Scala Spark?
We also managed to test this by sharing the Spark Context between Scala and Python. That means we registered the MLflow UDF in Python, and then used it from Scala (yeap, not a nice solution, but at least it’s something 🍭).
For this example, we will add the Toree Kernel to our existing Jupyter.
As you can see in the attached notebook, the UDF is shared between Spark and PySpark. We hope this last part is helpful for those teams that love Scala and have to put ML models into production.
Although MLflow is currently in Alpha (🥁), it looks quite promising. Just with the power of being able to run multiple machine learning frameworks and using them from the same endpoint it brings all the recommendation systems to the next level.
Plus, MLflow brings closer Data Engineers and Data Scientists by having a common layer between them.
After this research about MLflow, we are sure we are going to go further with it and start using it in our Spark pipelines and in our recommendation systems.
It would be nice to have the file-store synchronized with a DB instead of using a FS. This should allow having multiple endpoints using the same file-store. Like having multiple Presto and Athena instances using the same Glue metastore.
Just to wrap up, thanks to all the community behind MLflow to make it possible and making our data lives more interesting.
If you are playing with MLflow, feel free to reach us and provide some feedback about how you are using it! Even more, if you are using MLflow in production.