MLflow插件用于Google Cloud Vertex AI

注意：该插件是实验性的，将来可能会更改或删除。

安装

python3 -m pip install google_cloud_mlflow

部署插件使用

命令行

创建部署

mlflow deployments create --target google_cloud --name "deployment name" --model-uri "models:/mymodel/mymodelversion" --config destination_image_uri="gcr.io/<repo>/<path>"

列出部署

mlflow deployments list --target google_cloud

获取部署

mlflow deployments get --target google_cloud --name "deployment name"

删除部署

mlflow deployments delete --target google_cloud --name "deployment name"

更新部署

mlflow deployments update --target google_cloud --name "deployment name" --model-uri "models:/mymodel/mymodelversion" --config destination_image_uri="gcr.io/<repo>/<path>"

预测

mlflow deployments predict --target google_cloud --name "deployment name" --input-path "inputs.json" --output-path "outputs.json

获取帮助

mlflow deployments help --target google_cloud

Python

from mlflow import deployments
client = deployments.get_deploy_client("google_cloud")

# Create deployment
model_uri = "models:/mymodel/mymodelversion"
deployment = client.create_deployment(
    name="deployment name",
    model_uri=model_uri,
    # Config is optional
    config=dict(
        # Deployed model config
        machine_type="n1-standard-2",
        min_replica_count=None,
        max_replica_count=None,
        accelerator_type=None,
        accelerator_count=None,
        service_account=None,
        explanation_metadata=None, # JSON string
        explanation_parameters=None, # JSON string

        # Model container image building config
        destination_image_uri=None,

        # Endpoint config
        endpoint_description=None,
        endpoint_deploy_timeout=None,

        # Vertex AI config
        project=None,
        location=None,
        encryption_spec_key_name=None,
        staging_bucket=None,
    )
)

# List deployments
deployments = client.list_deployments()

# Get deployment
deployments = client.get_deployment(name="deployment name")

# Delete deployment
deployment = client.delete_deployment(name="deployment name")

# Update deployment
deployment = client.create_deployment(
    name="deployment name",
    model_uri=model_uri,
    # Config is optional
    config=dict(...),
)

# Predict
import pandas
df = pandas.DataFrame([
    {"a": 1,"b": 2,"c": 3},
    {"a": 4,"b": 5,"c": 6}
])
predictions = client.predict("deployment name", df)

模型注册插件使用

将MLflow模型注册URI设置为某个Google Cloud Storage桶中的目录，然后使用mlflow.log_model记录模型，就像平常一样。

mlflow.set_registry_uri("gs://<bucket>/models/")