mlflow_provider.operators.registry

Module Contents

Classes

CreateRegisteredModelOperator

Creates a new registered model in MLflow.

GetRegisteredModelOperator

Gets a registered model from MLflow based on name.

DeleteRegisteredModelOperator

Deletes a registered model from MLflow based on name.

GetLatestModelVersionsOperator

Gets the latest model versions from MLflow Registry based on name.

CreateModelVersionOperator

Create a model version in MLflow Registry.

GetModelVersionOperator

Get specific model version from MLflow Registry.

DeleteModelVersionOperator

Delete specific model version from MLflow Registry.

TransitionModelVersionStageOperator

Transition model version to new stage.
class mlflow_provider.operators.registry.CreateRegisteredModelOperator(*, mlflow_conn_id='mlflow_default', name, tags=None, description=None, **kwargs)

Bases: airflow.models.BaseOperator

Creates a new registered model in MLflow.

Parameters:

  • mlflow_conn_id (str) – connection to run the operator with

  • name (str) – name of the registered model to be created

  • tags (list[Dict[str, str]]) – tags to add to the registered model

  • description (str) – description of model

template_fields = ['name', 'tags', 'description']
template_fields_renderers: Dict[str, str]
template_ext = ()
ui_color = '#f4a460'
execute(context)
class mlflow_provider.operators.registry.GetRegisteredModelOperator(*, mlflow_conn_id='mlflow_default', name, **kwargs)

Bases: airflow.models.BaseOperator

Gets a registered model from MLflow based on name.

Parameters:

  • mlflow_conn_id (str) – connection to run the operator with

  • name (str) – name of the registered model to get

template_fields = ['name']
template_fields_renderers: Dict[str, str]
template_ext = ()
ui_color = '#f4a460'
execute(context)
class mlflow_provider.operators.registry.DeleteRegisteredModelOperator(*, mlflow_conn_id='mlflow_default', name, **kwargs)

Bases: airflow.models.BaseOperator

Deletes a registered model from MLflow based on name.

Parameters:

  • mlflow_conn_id (str) – connection to run the operator with

  • name (str) – name of the registered model to delete

template_fields = ['name']
template_fields_renderers: Dict[str, str]
template_ext = ()
ui_color = '#f4a460'
execute(context)
class mlflow_provider.operators.registry.GetLatestModelVersionsOperator(*, mlflow_conn_id='mlflow_default', name, stages=None, **kwargs)

Bases: airflow.models.BaseOperator

Gets the latest model versions from MLflow Registry based on name.

Parameters:

  • mlflow_conn_id (str) – connection to run the operator with

  • name (str) – name of the registered model to get versions for

  • stages (list) – List of stages to get

template_fields = ['name']
template_fields_renderers: Dict[str, str]
template_ext = ()
ui_color = '#f4a460'
execute(context)
class mlflow_provider.operators.registry.CreateModelVersionOperator(*, mlflow_conn_id='mlflow_default', name, source, run_id=None, tags=None, run_link=None, description=None, **kwargs)

Bases: airflow.models.BaseOperator

Create a model version in MLflow Registry.

Parameters:

  • mlflow_conn_id (str) – connection to run the operator with

  • name (str) – name of the registered model

  • source (str) – URI indicating the location of the model artifacts

  • run_id (str) – MLflow run ID for correlation, if source was generated by an experiment run in MLflow tracking server

  • tags (list) – Additional metadata for model version

  • run_link (str) – MLflow run link - this is the exact link of the run that generated this model version, potentially hosted at another instance of MLflow.

  • description (str) – description for model version

template_fields = ['name', 'source', 'run_id', 'tags', 'run_link', 'description']
template_fields_renderers
template_ext = ()
ui_color = '#f4a460'
execute(context)
class mlflow_provider.operators.registry.GetModelVersionOperator(*, mlflow_conn_id='mlflow_default', name, version, **kwargs)

Bases: airflow.models.BaseOperator

Get specific model version from MLflow Registry.

Parameters:

  • mlflow_conn_id (str) – connection to run the operator with

  • name (str) – name of the registered model

  • version (str) – Model version number

template_fields = ['name', 'version']
template_fields_renderers: Dict[str, str]
template_ext = ()
ui_color = '#f4a460'
execute(context)
class mlflow_provider.operators.registry.DeleteModelVersionOperator(*, mlflow_conn_id='mlflow_default', name, version, **kwargs)

Bases: airflow.models.BaseOperator

Delete specific model version from MLflow Registry.

Parameters:

  • mlflow_conn_id (str) – connection to run the operator with

  • name (str) – name of the registered model

  • version (str) – Model version number

template_fields = ['name', 'version']
template_fields_renderers: Dict[str, str]
template_ext = ()
ui_color = '#f4a460'
execute(context)
class mlflow_provider.operators.registry.TransitionModelVersionStageOperator(*, mlflow_conn_id='mlflow_default', name, version, stage, archive_existing_versions=False, **kwargs)

Bases: airflow.models.BaseOperator

Transition model version to new stage.

Parameters:

  • mlflow_conn_id (str) – connection to run the operator with

  • name (str) – name of the registered model

  • version (str) – Model version number

  • stage (str) – Transition model_version to new stage

  • archive_existing_versions (bool) – When transitioning a model version to a particular stage, this flag dictates whether all existing model versions in that stage should be atomically moved to the “archived” stage.

template_fields = ['name', 'version']
template_fields_renderers: Dict[str, str]
template_ext = ()
ui_color = '#f4a460'
execute(context)