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
- Schema
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 |
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.
Is this page helpful?