import Tabs from '@theme/Tabs'; import TabItem from '@theme/TabItem';

AI/ML Framework Integration with DataHub

Why Integrate Your AI/ML System with DataHub?

As a data practitioner, keeping track of your AI experiments, models, and their relationships can be challenging. DataHub makes this easier by providing a central place to organize and track your AI assets.

This guide will show you how to integrate your AI workflows with DataHub. With integrations for popular ML platforms like MLflow and Amazon SageMaker, DataHub enables you to easily find and share AI models across your organization, track how models evolve over time, and understand how training data connects to each model. Most importantly, it enables seamless collaboration on AI projects by making everything discoverable and connected.

Goals Of This Guide

In this guide, you'll learn how to:

Create your basic AI components (models, experiments, runs)
Connect these components to build a complete AI system
Track relationships between models, data, and experiments

Core AI Concepts

Here's what you need to know about the key components in DataHub:

Experiments are collections of training runs for the same project, like all attempts to build a churn predictor
Training Runs are attempts to train a model within an experiment, capturing parameters and results
Model Groups organize related models together, like all versions of your churn predictor
Models are successful training runs registered for production use

The hierarchy works like this:

Every run belongs to an experiment
Successful runs can be registered as models
Models belong to a model group
Not every run becomes a model

:::note Terminology Mapping Different AI platforms (MLflow, Amazon SageMaker) have their own terminology. To keep things consistent, we'll use DataHub's terms throughout this guide. Here's how DataHub's terminology maps to these platforms:

DataHub	Description	MLflow	SageMaker
ML Model Group	Collection of related models	Model	Model Group
ML Model	Versioned artifact in a model group	Model Version	Model Version
ML Training Run	Single training attempt	Run	Training Job
ML Experiment	Collection of training runs	Experiment	Experiment

:::

For platform-specific details, see our integration guides for MLflow and Amazon SageMaker.

Basic Setup

To follow this tutorial, you'll need DataHub Quickstart deployed locally. For detailed steps, see the DataHub Quickstart Guide.

Next, set up the Python client for DataHub using DatahubAIClientdefined in here.

Create a token in DataHub UI and replace <your_token> with your token:

python

from dh_ai_client import DatahubAIClient

client = DatahubAIClient(token="<your_token>", server_url="http://localhost:9002")

:::note Verifying via GraphQL Throughout this guide, we'll show how to verify changes using GraphQL queries. You can run these queries in the DataHub UI at https://localhost:9002/api/graphiql. :::

Create AI Assets

Let's create the basic building blocks of your ML system. These components will help you organize your AI work and make it discoverable by your team.

Create a Model Group

A model group contains different versions of a similar model. For example, all versions of your "Customer Churn Predictor" would go in one group.

<Tabs> <TabItem value="basic" label="Basic"> Create a basic model group with just an identifier:

python

model_group = MLModelGroup(
    id="airline_forecast_models_group",
    platform="mlflow",
)

client._emit_mcps(model_group.as_mcps())

</TabItem> <TabItem value="advanced" label="Advanced"> Add rich metadata like descriptions, creation timestamps, and team information:

python

model_group = MLModelGroup(
    id="airline_forecast_models_group",
    platform="mlflow",
    name="Airline Forecast Models Group",
    description="Group of models for airline passenger forecasting",
    created=datetime.now(),
    last_modified=datetime.now(),
    owners=[CorpUserUrn("urn:li:corpuser:datahub")],
    external_url="https://www.linkedin.com/in/datahub",
    tags=["urn:li:tag:forecasting", "urn:li:tag:arima"],
    terms=["urn:li:glossaryTerm:forecasting"],
    custom_properties={"team": "forecasting"},
)

client._emit_mcps(model_group.as_mcps())

</TabItem> </Tabs>

Let's verify that our model group was created:

<Tabs> <TabItem value="UI" label="UI"> See your new model group in the DataHub UI: </TabItem> <TabItem value="graphql" label="GraphQL"> Query your model group to check its properties:

graphql

query {
  mlModelGroup(
    urn: "urn:li:mlModelGroup:(urn:li:dataPlatform:mlflow,airline_forecast_models_group,PROD)"
  ) {
    properties {
      name
      description
      created {
        time
      }
    }
  }
}

The response will show your model group's details:

json

{
  "data": {
    "mlModelGroup": {
      "properties": {
        "name": "Airline Forecast Models Group",
        "description": "Group of models for airline passenger forecasting",
        "created": {
          "time": 1744356062485
        }
      }
    }
  },
  "extensions": {}
}

</TabItem> </Tabs>

Create a Model

Next, let's create a specific model version that represents a trained model ready for deployment.

<Tabs> <TabItem value="basic" label="Basic"> Create a model with just the required version:

python

model = MLModel(
    id="arima_model",
    platform="mlflow",
)

client._emit_mcps(model.as_mcps())

</TabItem> <TabItem value="advanced" label="Advanced"> Include metrics, parameters, and metadata for production use:

python

model = MLModel(
    id="arima_model",
    platform="mlflow",
    name="ARIMA Model",
    description="ARIMA model for airline passenger forecasting",
    created=datetime.now(),
    last_modified=datetime.now(),
    owners=[CorpUserUrn("urn:li:corpuser:datahub")],
    external_url="https://www.linkedin.com/in/datahub",
    tags=["urn:li:tag:forecasting", "urn:li:tag:arima"],
    terms=["urn:li:glossaryTerm:forecasting"],
    custom_properties={"team": "forecasting"},
    version="1",
    aliases=["champion"],
    hyper_params={"learning_rate": "0.01"},
    training_metrics={"accuracy": "0.9"},
)

client._emit_mcps(model.as_mcps())

</TabItem> </Tabs>

Let's verify our model:

<Tabs> <TabItem value="UI" label="UI"> Check your model's details in the DataHub UI: </TabItem> <TabItem value="graphql" label="GraphQL"> Query your model's information:

graphql

query {
  mlModel(urn: "urn:li:mlModel:(urn:li:dataPlatform:mlflow,arima_model,PROD)") {
    properties {
      name
      description
    }
    versionProperties {
      version {
        versionTag
      }
    }
  }
}

The response will show your model's details:

json

{
  "data": {
    "mlModel": {
      "properties": {
        "name": "ARIMA Model",
        "description": "ARIMA model for airline passenger forecasting"
      },
      "versionProperties": {
        "version": {
          "versionTag": "1"
        }
      }
    }
  },
  "extensions": {}
}

</TabItem> </Tabs>

Create an Experiment

An experiment helps organize multiple training runs for a specific project.

<Tabs> <TabItem value="basic" label="Basic"> Create a basic experiment:

python

experiment = Container(
    container_key=ContainerKey(
        platform="mlflow",
        name="airline_forecast_experiment"
    ),
    display_name="Airline Forecast Experiment"
)

client._emit_mcps(experiment.as_mcps())

</TabItem> <TabItem value="advanced" label="Advanced"> Add context and metadata:

python

experiment = Container(
    container_key=ContainerKey(
        platform="mlflow",
        name="airline_forecast_experiment"
    ),
    display_name="Airline Forecast Experiment",
    description="Experiment to forecast airline passenger numbers",
    extra_properties={"team": "forecasting"},
    created=datetime(2025, 4, 9, 22, 30),
    last_modified=datetime(2025, 4, 9, 22, 30),
    subtype=MLAssetSubTypes.MLFLOW_EXPERIMENT,
)

client._emit_mcps(experiment.as_mcps())

</TabItem> </Tabs>

Verify your experiment:

<Tabs> <TabItem value="UI" label="UI"> See your experiment's details in the UI: </TabItem> <TabItem value="graphql" label="GraphQL"> Query your experiment's information:

graphql

query {
  container(urn: "urn:li:container:airline_forecast_experiment") {
    properties {
      name
      description
    }
  }
}

Check the response:

json

{
  "data": {
    "container": {
      "properties": {
        "name": "Airline Forecast Experiment",
        "description": "Experiment to forecast airline passenger numbers"
      }
    }
  },
  "extensions": {}
}

</TabItem> </Tabs>

Create a Training Run

A training run captures all details about a specific model training attempt.

<Tabs> <TabItem value="basic" label="Basic"> Create a basic training run:

python

client.create_training_run(
    run_id="simple_training_run",
)

</TabItem> <TabItem value="advanced" label="Advanced"> Include metrics, parameters, and other important metadata:

python

client.create_training_run(
    run_id="simple_training_run",
    properties=DataProcessInstancePropertiesClass(
        name="Simple Training Run",
        created=AuditStampClass(
            time=1628580000000, actor="urn:li:corpuser:datahub"
        ),
        customProperties={"team": "forecasting"},
    ),
    training_run_properties=MLTrainingRunPropertiesClass(
        id="simple_training_run",
        outputUrls=["s3://my-bucket/output"],
        trainingMetrics=[MLMetricClass(name="accuracy", value="0.9")],
        hyperParams=[MLHyperParamClass(name="learning_rate", value="0.01")],
        externalUrl="https:localhost:5000",
    ),
    run_result=RunResultType.FAILURE,
    start_timestamp=1628580000000,
    end_timestamp=1628580001000,
)

</TabItem> </Tabs>

Verify your training run:

<Tabs> <TabItem value="UI" label="UI"> View the run details in the UI: </TabItem> <TabItem value="graphql" label="GraphQL"> Query your training run:

graphql

query {
  dataProcessInstance(urn: "urn:li:dataProcessInstance:simple_training_run") {
    name
    created {
      time
    }
    properties {
      customProperties
    }
  }
}

Check the response:

json

{
  "data": {
    "dataProcessInstance": {
      "name": "Simple Training Run",
      "created": {
        "time": 1628580000000
      },
      "properties": {
        "customProperties": {
          "team": "forecasting"
        }
      }
    }
  }
}

</TabItem> </Tabs>

Create a Dataset

Datasets are crucial components in your ML system, serving as inputs and outputs for your training runs. Creating a dataset in DataHub allows you to track data lineage and understand how data flows through your ML pipeline.

<Tabs> <TabItem value="basic" label="Basic"> Create a basic dataset with minimal information:

python

input_dataset = Dataset(
    platform="snowflake",
    name="iris_input",
)
client._emit_mcps(input_dataset.as_mcps())

</TabItem> <TabItem value="advanced" label="Advanced">

Create a dataset with more detailed information:

python

input_dataset = Dataset(
    platform="snowflake",
    name="iris_input",
    description="Raw Iris dataset used for training ML models",
    schema=[("id", "number"), ("name", "string"), ("species", "string")],
    display_name="Iris Training Input Data",
    tags=["urn:li:tag:ml_data", "urn:li:tag:iris"],
    terms=["urn:li:glossaryTerm:raw_data"],
    owners=[CorpUserUrn("urn:li:corpuser:datahub")],
    custom_properties={
        "data_source": "UCI Repository",
        "records": "150",
        "features": "4",
    },
)
client._emit_mcps(input_dataset.as_mcps())

output_dataset = Dataset(
    platform="snowflake",
    name="iris_output",
    description="Processed Iris dataset with model predictions",
    schema=[("id", "number"), ("name", "string"), ("species", "string")],
    display_name="Iris Model Output Data",
    tags=["urn:li:tag:ml_data", "urn:li:tag:predictions"],
    terms=["urn:li:glossaryTerm:model_output"],
    owners=[CorpUserUrn("urn:li:corpuser:datahub")],
    custom_properties={
        "model_version": "1.0",
        "records": "150",
        "accuracy": "0.95",
    },
)
client._emit_mcps(output_dataset.as_mcps())

</TabItem> </Tabs>

Verify your datasets:

<Tabs> <TabItem value="UI" label="UI"> View dataset details in the DataHub UI: </TabItem> <TabItem value="graphql" label="GraphQL">

Query your dataset's information:

graphql

query {
  dataset(
    urn: "urn:li:dataset:(urn:li:dataPlatform:snowflake,iris_input,PROD)"
  ) {
    name
    properties {
      customProperties
    }
  }
}

Check the response:

graphql

{
  "data": {
    "dataset": {
      "name": "iris_input",
      "properties": {
        "customProperties": {
          "data_source": "UCI Repository",
          "records": "150",
          "features": "4"
        }
      }
    }
  }
}

</TabItem> </Tabs>

Datasets in DataHub can also include schema information, data quality metrics, and lineage details, which are particularly valuable for ML workflows where understanding data characteristics is crucial for model performance.

Define Relationships

Now let's connect these components to create a comprehensive ML system. These connections enable you to track model lineage, monitor model evolution, understand dependencies, and search effectively across your ML assets.

Add Model To Model Group

Connect your model to its group:

python

model.add_group(model_group.urn)

client._emit_mcps(model.as_mcps())

View model versions in the Model Group under the Models section:

Find group information in the Model page under the Group tab:

</TabItem> <TabItem value="graphql" label="GraphQL"> Query the model-group relationship:

graphql

query {
  mlModel(urn: "urn:li:mlModel:(urn:li:dataPlatform:mlflow,arima_model,PROD)") {
    name
    properties {
      groups {
        urn
        properties {
          name
        }
      }
    }
  }
}

Check the response:

json

{
  "data": {
    "mlModel": {
      "name": "ARIMA Model",
      "properties": {
        "groups": [
          {
            "urn": "urn:li:mlModelGroup:(urn:li:dataPlatform:mlflow,airline_forecast_models_group)",
            "properties": {
              "name": "Airline Forecast Models Group"
            }
          }
        ]
      }
    }
  }
}

</TabItem> </Tabs>

Add Run To Experiment

Connect a training run to its experiment:

python

client.add_run_to_experiment(run_urn=run_urn, experiment_urn=experiment_urn)

Find your runs in the Experiment page under the Entities tab:

See the experiment details in the Run page:

</TabItem> <TabItem value="graphql" label="GraphQL"> Query the run-experiment relationship:

graphql

query {
  dataProcessInstance(urn: "urn:li:dataProcessInstance:simple_training_run") {
    name
    parentContainers {
      containers {
        urn
        properties {
          name
        }
      }
    }
  }
}

View the relationship details:

json

{
  "data": {
    "dataProcessInstance": {
      "name": "Simple Training Run",
      "parentContainers": {
        "containers": [
          {
            "urn": "urn:li:container:airline_forecast_experiment",
            "properties": {
              "name": "Airline Forecast Experiment"
            }
          }
        ]
      }
    }
  }
}

</TabItem> </Tabs>

Add Run To Model

Connect a training run to its resulting model:

python

model.add_training_job(DataProcessInstanceUrn(run_id))

client._emit_mcps(model.as_mcps())

This relationship enables you to:

Track which runs produced each model
Understand model provenance
Debug model issues
Monitor model evolution

Find the source run in the Model page under the Summary tab:

See related models in the Run page under the Lineage tab:

</TabItem> <TabItem value="graphql" label="GraphQL"> Query the model's training jobs:

graphql

query {
  mlModel(urn: "urn:li:mlModel:(urn:li:dataPlatform:mlflow,arima_model,PROD)") {
    name
    properties {
      mlModelLineageInfo {
        trainingJobs
      }
    }
  }
}

View the relationship:

json

{
  "data": {
    "mlModel": {
      "name": "ARIMA Model",
      "properties": {
        "mlModelLineageInfo": {
          "trainingJobs": ["urn:li:dataProcessInstance:simple_training_run"]
        }
      }
    }
  }
}

</TabItem> </Tabs>

Add Run To Model Group

Create a direct connection between a run and a model group:

python

model_group.add_training_job(DataProcessInstanceUrn(run_id))

client._emit_mcps(model_group.as_mcps())

This connection lets you:

View model groups in the run's lineage
Query training jobs at the group level
Track training history for model families

See model groups in the Run page under the Lineage tab:

</TabItem> <TabItem value="graphql" label="GraphQL"> Query the model group's training jobs:

graphql

query {
  mlModelGroup(
    urn: "urn:li:mlModelGroup:(urn:li:dataPlatform:mlflow,airline_forecast_models_group)"
  ) {
    name
    properties {
      mlModelLineageInfo {
        trainingJobs
      }
    }
  }
}

Check the relationship:

json

{
  "data": {
    "mlModelGroup": {
      "name": "Airline Forecast Models Group",
      "properties": {
        "mlModelLineageInfo": {
          "trainingJobs": ["urn:li:dataProcessInstance:simple_training_run"]
        }
      }
    }
  }
}

</TabItem> </Tabs>

Add Dataset To Run

Track input and output datasets for your training runs:

python

client.add_input_datasets_to_run(
    run_urn=run_urn,
    dataset_urns=[str(input_dataset_urn)]
)

client.add_output_datasets_to_run(
    run_urn=run_urn,
    dataset_urns=[str(output_dataset_urn)]
)

These connections help you:

Track data lineage
Understand data dependencies
Ensure reproducibility
Monitor data quality impacts

Find dataset relationships in the Lineage tab of either the Dataset or Run page:

Update Properties

Update Model Group Properties

Model groups can be updated with additional information:

python

# Update description
model_group.set_description("Updated description for airline forecast models")

# Add tags and terms
model_group.add_tag(TagUrn("production"))
model_group.add_term(GlossaryTermUrn("time-series"))

# Update custom properties
model_group.set_custom_properties({
    "team": "forecasting",
    "business_unit": "operations",
    "status": "active"
})

# Save the changes
client._emit_mcps(model_group.as_mcps())

These updates allow you to:

Improve documentation with detailed descriptions
Apply consistent business context with tags and terms
Track organizational ownership and status

Update Model Properties

Models can be updated with additional information as they evolve:

python

# Update model version
model.set_version("2")

# Add tags and terms
model.add_tag(TagUrn("marketing"))
model.add_term(GlossaryTermUrn("marketing"))

# Add version alias
model.add_version_alias("challenger")

# Save the changes
client._emit_mcps(model.as_mcps())

These updates allow you to:

Track model iterations through versioning
Apply business context with tags and terms
Manage deployment aliases like "champion" and "challenger"

Update Experiment Properties

Experiments can be updated with additional metadata as your project evolves:

python

# Create a container object for the existing experiment
existing_experiment = Container(
    container_key=ContainerKey(
        platform="mlflow",
        name="airline_forecast_experiment"
    ),
)

# Update properties
existing_experiment.set_description("Updated experiment for forecasting passenger numbers")
existing_experiment.add_tag(TagUrn("time-series"))
existing_experiment.add_term(GlossaryTermUrn("forecasting"))
existing_experiment.set_custom_properties({
    "team": "forecasting",
    "priority": "high",
    "status": "active"
})

# Save the changes
client._emit_mcps(existing_experiment.as_mcps())

These updates help you:

Document evolving experiment objectives
Categorize experiments with consistent tags
Track experiment status and priority

Full Overview

Here's your complete ML system with all components connected:

You now have a complete lineage view of your ML assets, from training data through runs to production models!

You can check out the full code for this tutorial here.

What's Next?

To see these integrations in action:

Watch our Townhall demo showcasing the MLflow integration
Explore our detailed documentation:
- MLflow Integration Guide
- Amazon SageMaker Integration Guide