Baseten now supports MLflow models via Truss. MLflow is a popular library for model experimentation and model management with over ten million monthly downloads on PyPi. With MLflow, you can train a model in any framework (PyTorch, TensorFlow, XGBoost, etc) and access features for tracking, packaging, and registering your model. And now, deploying to Baseten is a natural extension of MLflow-based workflows.
Deploying an MLflow model looks a bit like this:
import mlflow import baseten model = mlflow.pyfunc.load_model(MODEL_URI) baseten.deploy(model, "MLflow model")
For a complete runnable example, check out this demo on Google Colab.