mlflow_provider.hooks.deployment

Module Contents

Classes

MLflowDeploymentHook

Hook that interacts with the mlflow.deployments module in the MLflow library.
class mlflow_provider.hooks.deployment.MLflowDeploymentHook(mlflow_conn_id, target_uri, target_conn_id=None)

Bases: mlflow_provider.hooks.base.MLflowBaseHook

Hook that interacts with the mlflow.deployments module in the MLflow library. https://www.mlflow.org/docs/latest/python_api/mlflow.deployments.html

Parameters:

  • aws_conn_id (str) – AWS connection to use with hook

  • target_uri (str) – target system URI to deploy model to. (ie ‘sagemaker’)

hook_name = 'MLflow Deployment'
aws_conn_dict()
get_conn()

Returns MLflow deployment Client.

create_deployment(name, model_uri, flavor=None, config=None, endpoint=None)

Creates a deployment in the target system. https://www.mlflow.org/docs/latest/python_api/mlflow.deployments.html#mlflow.deployments.BaseDeploymentClient.create_deployment

Parameters:

  • name (str) – Unique name to use for deployment.

  • model_uri (str) – URI of model to deploy

  • flavor (str) – (optional) Model flavor to deploy. If unspecified, a default flavor will be chosen.

  • config (dict) – (optional) Dict containing updated target-specific configuration for the deployment

  • endpoint (str) – (optional) Endpoint to create the deployment under. May not be supported by all targets

Returns:

Dict corresponding to created deployment, which must contain the ‘name’ key.

Return type:

dict

predict(deployment_name, inputs, endpoint=None)

Makes a prediction request to the specified deployment. https://www.mlflow.org/docs/latest/python_api/mlflow.deployments.html#mlflow.deployments.BaseDeploymentClient.predict

Parameters:

  • deployment_name (str) – Name of deployment to predict against

  • inputs (Any) – Input data (or arguments) to pass to the deployment or model endpoint for inference

  • endpoint – Endpoint to predict against. May not be supported by all targets

Returns:

A mlflow.deployments.PredictionsResponse instance representing the predictions and associated Model Server response metadata as a JSON.

Return type:

dict