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 create command to package your model into a Truss.
import os
import truss
​
model = mlflow.pyfunc.load_model(MODEL_URI)
tr = truss.create(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)
tr.predict(data)
For more on running the Truss locally, see local development.

Last modified 22d ago