MLEM: A One Stop Shop for Productionizing your ML Models
Despite from what this image might lead you to believe, MLEM isn’t about cats :/
MLEM is an open source tool by Iterative.ai, to help you easily package, deploy and serve your Machine Learning models all in one place. What makes MLEM special to me is its “codification” principle —while pickle/joblib simply serializes/deserializes model objects, MLEM also saves additional model metadata, which saves additional data about the model that could be useful to reload or utilize it. This model metadata is inferred automatically by MLEM, no matter how complex you architecture is. This metadata is saved as a human-readable YAML to make things more convenient for ML Teams.
Now let’s delve into more detail with a quick example!
Getting Started with MLEM
For this example, you need a trained ML model (could be any popular framework — TensorFlow, PyTorch, scikit-learn). I trained a simple Random Forest Regressor from scikit learn to predict movie ratings on the IMDb Dataset (code on GitHub).
Prerequisites
For this tutorial, you would want to install the MLEM API.
pip install mlemFor deploying on FastAPI
pip install fastapi uvicornFor Docker, you can follow setup instructions here.
Saving Models
Once you have your model trained and ready, normally you would save the model to a pkl file using joblib.dump() or pickle.dump(). However, in this scenario we will be using MLEM API’s save function.
from mlem.api import save
PATH = "models/ratings_model"
save(model, PATH, sample_data = df)model — the model you would want to save
PATH — filename to save the model as
sample_data — If the object is a model or function, you can provide input data sample, so MLEM will include it’s schema in the model’s metadata
Once you call the save function, your directory would look like this
> ls model/
ratings_model ratings_model.mlemratings_model is the binary file with the actual model
ratings_model.mlem is the metadata file with model details
Looking into the metadata file
artifacts:
data:
hash: 30ab43b8da9cb06a931a1e10923adfa3
size: 63822733
uri: ratings_model
call_orders:
predict:
- - model
- predict
object_type: model
processors:
model:
methods:
predict:
args:
- name: X
type_:
columns:
- isAdult
- startYear
- runtimeMinutes
dtypes:
- object
- float64
- float64
index_cols: []
type: dataframe
name: predict
returns:
dtype: float64
shape:
- null
type: ndarray
type: sklearn
requirements:
- module: sklearn
package_name: scikit-learn
version: 1.1.2
- module: numpy
version: 1.24.2
- module: pandas
version: 1.5.3ratings_model.mlem is a YAML file containing metadata about the model. It contains useful information about the model, like -
- Functions calls for the model object
- What arguments the functions take
- If you provided
sample_data, the columns/datatypes - All package requirements for this model
This data is extremely useful while coordinating with other developers/within ML teams since it provides enough information about using the model without you having to go through the stress manually documenting instructions.
Note: You can also save pickle files as MLEM Models by simply loading them back in and saving it again
m = joblib.load("model/rfm.pkl")
save(m, "model/rfm2", sample_data = df)Serving Models Locally
This is my personal favorite part of MLEM, how easy it is to serve models! For this tutorial, I’ll be sticking to FastAPI, while you could also use Streamlit or RabbitMQ.
You just need to run the following command
> mlem serve fastapi --model ratings_model
⏳️ Loading model from ratings_model.mlem
Starting fastapi server...
🖇️ Adding route for /predict
Checkout openapi docs at <http://0.0.0.0:8080/docs>
INFO: Started server process [54845]
INFO: Waiting for application startup.
INFO: Application startup complete.
INFO: Uvicorn running on http://0.0.0.0:8080 (Press CTRL+C to quit)This sets up the server, which is now serving your model and listening on http://0.0.0.0:8080 , you can navigate to http://0.0.0.0:8080/docs to navigate through the Swagger UI.
http://0.0.0.0:8080/docs — Providing information about routes to useYou can test calls on the Swagger UI for model prediction, in the image below I filled in the sample request body with a few dummy values.
It not only generates curl calls for you to use, but also provides response code and details. As you can see, the model does successfully return a predicted rating given movie details.
As you can see, it’s pretty easy to serve models locally, providing you with the perfect API to be integrated into any web app which requires ML model predictions. This serves as the perfect solution for all of my course projects where it gets really difficult to deploy models and serve them without any extra effort or money.
Deploying and Other Use Cases
Similar to what we saw in the previous section, its pretty easy to deploy Models in production as well. MLEM works with Docker, Heroku and GitHub Actions for fast and reliable deployment, all with a simple single command!
It’s easy to version models with MLEM as well, one of it’s main features is a Git-Native Model registry, that works with other Iterative tools like DVC. This helps you to track deployment parameters and use it in CI/CD pipelines more easily.
MLEM also supports Kubernetes and SageMaker now!
Final Thoughts
Based off my limited experience with MLEM, I’ve really liked it so far! It makes the whole process of saving, serving and deploying models so much easier and convenient, without the hassle of using multiple tools. I would definitely be using this tool in my future projects, and I hope this article inspires you to do the same!
As for one of the downsides I faced personally, while MLEM Model Objects do support almost all major frameworks, I did face issues with smaller frameworks (in my case it was with scikit-surprise) — while they are slowly rolling out support for most frameworks I would recommend checking compatibility before deciding to use MLEM.
I hope that you enjoyed reading this article and learnt something new today! This article is part of an assignment for my ML in Production course at Carnegie Mellon University, under Profs Christian Kästner and Eunsuk Kang. This was definitely a great way to learn about the MLOps industry and the crazy amount of companies out there that do some amazing work!
