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.
Package | Version |
---|---|
fastai | ~=2.7.10 |
gluon | ~=1.1.0 |
keras | ~=2.10.0 |
lightgbm | ~=3.3.3 |
mlflow | ~=1.29.0 |
onnx | ~=1.12.0 |
scikit-learn | ~=1.1.3 |
statsmodels | ~=0.13.5 |
tensorflow | ~=2.10.1 |
torch | ~=1.12.1 |
torchvision | ~=0.13.1 |
xgboost | ~=1.7.1 |
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)