Skip to main content
Version: v1.0.0

MLflow model support

H2O MLOps lets you upload and deploy MLflow models. The following sections describe this feature.

Supported third-party models​

The following is a list of tested and supported third-party Python models.

PackageVersion
fastai>=2.7.15,<2.9.0
keras~=3.9.1
lightgbm~=4.6.0
mlflow>=2.14.2
onnx~=1.17.0
scikit-learn>=1.1.3,<1.4.0
statsmodels~=0.14.1
tensorflow~=2.16.1
torch<2.6.0
torchvision<0.21.0
xgboost~=1.7.1
opencv-python-headless~=4.11.0
click<9.0.0,>=8.1.0

Create MLflow artifacts for third-party frameworks​

The following is an example of how to create MLflow artifacts for third-party frameworks.

import shutil
import mlflow.sklearn
import sklearn.datasets
import sklearn.ensemble
from mlflow.models import signature
from mlflow.types import Schema, ColSpec, DataType

# Train sklearn model
X_train, y_train = sklearn.datasets.load_wine(return_X_y=True, as_frame=True)
y_train = (y_train >= 7).astype(int)

sklearn_model = sklearn.ensemble.RandomForestClassifier(n_estimators=10)
sklearn_model.fit(X_train, y_train)

# Infer and set model signature
model_signature = signature.infer_signature(X_train, sklearn_model.predict(X_train))
model_signature.outputs = Schema(
    [ColSpec(name="quality", type=DataType.float)]
)

# Define the path to store the model in the current directory
model_dir_path = "wine_model"

# Save the trained sklearn model with MLflow
mlflow.sklearn.save_model(
    sklearn_model, model_dir_path, signature=model_signature
)

# Create a zip archive of the saved model
shutil.make_archive("artifact", "zip", model_dir_path)
Feedback