mlflow_provider.operators.deployment

Module Contents

Classes

CreateDeploymentOperator

Deploy MLflow models

PredictOperator

Get predictions from an MLflow deployment
class mlflow_provider.operators.deployment.CreateDeploymentOperator(*, mlflow_conn_id='mlflow_default', name, model_uri, target_uri, target_conn_id=None, flavor=None, config=None, endpoint=None, **kwargs)

Bases: airflow.models.BaseOperator

Deploy MLflow models

Parameters:

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

  • model_uri (str) – URI of MLflow model

  • target_uri (str) – URI of location to deploy the model (ie ‘sagemaker’)

  • target_conn_id (str) – Airflow connection id for target system

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

  • config (dict) – Target-specific configuration for the deployment

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

template_fields = ['name', 'model_uri', 'endpoint', 'target_uri', 'flavor', 'config', 'endpoint']
template_fields_renderers
template_ext = ()
ui_color = '#f4a460'
execute(context)
class mlflow_provider.operators.deployment.PredictOperator(*, mlflow_conn_id='mlflow_default', deployment_name, inputs=None, endpoint=None, target_uri, target_conn_id=None, **kwargs)

Bases: airflow.models.BaseOperator

Get predictions from an MLflow deployment

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 (str) – Endpoint to predict against. May not be supported by all targets

  • target_uri (str) – URI of location to deploy the model (ie ‘sagemaker’)

  • target_conn_id (str) – Connection id for target system

template_fields = ['deployment_name', 'endpoint', 'target_uri', 'target_conn_id']
template_fields_renderers: Dict[str, str]
template_ext = ()
ui_color = '#f4a460'
execute(context)