MLflow

Important Capabilities

Capability	Status	Notes
Descriptions	✅	Extract descriptions for MLflow Registered Models and Model Versions
Extract Tags	✅	Extract tags for MLflow Registered Model Stages

Concept Mapping

This ingestion source maps the following MLflow Concepts to DataHub Concepts:

Source Concept	DataHub Concept	Notes
Registered Model	MlModelGroup	The name of a Model Group is the same as a Registered Model's name (e.g. my_mlflow_model)
Model Version	MlModel	The name of a Model is {registered_model_name}{model_name_separator}{model_version} (e.g. my_mlflow_model_1 for Registered Model named my_mlflow_model and Version 1, my_mlflow_model_2, etc.)
Model Stage	Tag	The mapping between Model Stages and generated Tags is the following: - Production: mlflow_production - Staging: mlflow_staging - Archived: mlflow_archived - None: mlflow_none

CLI based Ingestion

Install the Plugin

pip install 'acryl-datahub[mlflow]'

Starter Recipe

Check out the following recipe to get started with ingestion! See below for full configuration options.

For general pointers on writing and running a recipe, see our main recipe guide.

source:
  type: mlflow
  config:
    # Coordinates
    tracking_uri: tracking_uri

sink:
  # sink configs

Config Details

Options

Note that a . is used to denote nested fields in the YAML recipe.

Field	Description
model_name_separator string	A string which separates model name from its version (e.g. model_1 or model-1) Default: _
registry_uri string	Registry server URI. If not set, an MLflow default registry_uri is used (value of tracking_uri or MLFLOW_REGISTRY_URI environment variable)
tracking_uri string	Tracking server URI. If not set, an MLflow default tracking_uri is used (local mlruns/ directory or MLFLOW_TRACKING_URI environment variable)
env string	The environment that all assets produced by this connector belong to Default: PROD

Schema

The JSONSchema for this configuration is inlined below.

{
  "title": "MLflowConfig",
  "description": "Any source that produces dataset urns in a single environment should inherit this class",
  "type": "object",
  "properties": {
    "env": {
      "title": "Env",
      "description": "The environment that all assets produced by this connector belong to",
      "default": "PROD",
      "type": "string"
    },
    "tracking_uri": {
      "title": "Tracking Uri",
      "description": "Tracking server URI. If not set, an MLflow default tracking_uri is used (local `mlruns/` directory or `MLFLOW_TRACKING_URI` environment variable)",
      "type": "string"
    },
    "registry_uri": {
      "title": "Registry Uri",
      "description": "Registry server URI. If not set, an MLflow default registry_uri is used (value of tracking_uri or `MLFLOW_REGISTRY_URI` environment variable)",
      "type": "string"
    },
    "model_name_separator": {
      "title": "Model Name Separator",
      "description": "A string which separates model name from its version (e.g. model_1 or model-1)",
      "default": "_",
      "type": "string"
    }
  },
  "additionalProperties": false
}

Code Coordinates

• Class Name: datahub.ingestion.source.mlflow.MLflowSource
• Browse on GitHub

Questions

If you've got any questions on configuring ingestion for MLflow, feel free to ping us on our Slack.