Links

MLflow

MLflow is a supported framework on Truss. To package an MLflow model, follow the steps below or run this Google Colab notebook.

Install packages

If you're using a Jupyter notebook, add a line to install the mlflow and truss packages. Otherwise, ensure the packages are installed in your Python environment.
!pip install --upgrade pip
!pip install --upgrade mlflow truss
Truss officially supports mlflow version 1.30.0 or higher. Especially if you're using an online notebook environment like Google Colab or a bundle of packages like Anaconda, ensure that the version you are using is supported. If it's not, use the --upgrade flag and pip will install the most recent version.

Create an in-memory model

This is the part you want to replace with your own code. We are creating a super simple logistic regression model, but you can package any MLflow model as a Truss.
import mlflow
from sklearn.linear_model import LogisticRegression
import numpy as np
with mlflow.start_run():
X = np.array([-2, -1, 0, 1, 2, 1]).reshape(-1, 1)
y = np.array([0, 0, 1, 1, 1, 0])
lr = LogisticRegression()
lr.fit(X, y)
model_info = mlflow.sklearn.log_model(sk_model=lr, artifact_path="model")
MODEL_URI = model_info.model_uri

Create a Truss

Truss uses MLflow's pyfunc module in the packaging process. Once you have loaded the model, use the mk_truss command to package your model into a Truss.
import os
import truss
model = mlflow.pyfunc.load_model(MODEL_URI)
tr = truss.mk_truss(model, target_directory="./mlflow_truss")
Check the target directory to see your new Truss!

Serve the model

To get a prediction from the Truss, try running:
data = np.array([-4, 1, 0, 10, -2, 1]).reshape(-1, 1)
predictions = tr.server_predict({"inputs": data})
print(predictions)
For more on running the Truss locally, see local development.