Skip to main content
Version: Next

ML Model

The ML Model entity represents trained machine learning models across various ML platforms and frameworks. ML Models can be trained using different algorithms and frameworks (TensorFlow, PyTorch, Scikit-learn, etc.) and deployed to various platforms (MLflow, SageMaker, Vertex AI, etc.).

Identity​

ML Models are identified by three pieces of information:

  • The platform where the model is registered or deployed: this is the specific ML platform that hosts or manages this model. Examples are mlflow, sagemaker, vertexai, databricks, etc. See dataplatform for more details.
  • The name of the model: this is the unique identifier for the model within the platform. The naming convention varies by platform:
    • MLflow: typically uses the registered model name (e.g., recommendation-model)
    • SageMaker: uses the model name or model package group name (e.g., product-recommendation-v1)
    • Vertex AI: uses the model resource name (e.g., projects/123/locations/us-central1/models/456)
  • The environment or origin where the model was trained: this is similar to the fabric concept for datasets, allowing you to distinguish between models in different environments (PROD, DEV, QA, etc.). The full list of supported environments is available in FabricType.pdl.

An example of an ML Model identifier is urn:li:mlModel:(urn:li:dataPlatform:mlflow,my-recommendation-model,PROD).

Important Capabilities​

Basic Model Information​

The core information about an ML Model is captured in the mlModelProperties aspect. This includes:

  • Name and Description: Human-readable name and description of what the model does
  • Model Type: The algorithm or architecture used (e.g., "Convolutional Neural Network", "Random Forest", "BERT")
  • Version: Version information using the versionProperties aspect
  • Timestamps: Created and last modified timestamps
  • Custom Properties: Flexible key-value pairs for platform-specific metadata (e.g., framework version, model format)

The following code snippet shows you how to create a basic ML Model:

Python SDK: Create an ML Model
# Inlined from /metadata-ingestion/examples/library/mlmodel_create.py
from datahub.metadata.urns import MlModelGroupUrn
from datahub.sdk import DataHubClient
from datahub.sdk.mlmodel import MLModel

client = DataHubClient.from_env()

mlmodel = MLModel(
    id="customer-churn-predictor",
    name="Customer Churn Prediction Model",
    platform="mlflow",
    description="A gradient boosting model that predicts customer churn based on usage patterns and engagement metrics",
    custom_properties={
        "framework": "xgboost",
        "framework_version": "1.7.0",
        "model_format": "pickle",
    },
    model_group=MlModelGroupUrn(platform="mlflow", name="customer-churn-models"),
)

client.entities.upsert(mlmodel)

Hyperparameters and Metrics​

ML Models can capture both the hyperparameters used during training and various metrics from training and production:

  • Hyperparameters: Configuration values that control the training process (learning rate, batch size, number of epochs, etc.)
  • Training Metrics: Performance metrics from the training process (accuracy, loss, F1 score, etc.)
  • Online Metrics: Performance metrics from production deployment (latency, throughput, drift, etc.)

These are stored in the mlModelProperties aspect as structured lists of parameters and metrics.

Python SDK: Add hyperparameters and metrics to an ML Model
# Inlined from /metadata-ingestion/examples/library/mlmodel_add_metadata.py
from datahub.metadata.urns import CorpUserUrn, DomainUrn, MlModelUrn, TagUrn
from datahub.sdk import DataHubClient

client = DataHubClient.from_env()

mlmodel = client.entities.get(
    MlModelUrn(platform="mlflow", name="customer-churn-predictor")
)

mlmodel.set_hyper_params(
    {
        "learning_rate": "0.1",
        "max_depth": "6",
        "n_estimators": "100",
        "subsample": "0.8",
        "colsample_bytree": "0.8",
    }
)

mlmodel.set_training_metrics(
    {
        "accuracy": "0.87",
        "precision": "0.84",
        "recall": "0.82",
        "f1_score": "0.83",
        "auc_roc": "0.91",
    }
)

mlmodel.add_owner(CorpUserUrn("data_science_team"))

mlmodel.add_tag(TagUrn("production"))
mlmodel.add_tag(TagUrn("classification"))

mlmodel.set_domain(DomainUrn("urn:li:domain:customer-analytics"))

client.entities.update(mlmodel)

Intended Use and Ethical Considerations​

DataHub supports comprehensive model documentation following ML model card best practices. These aspects help stakeholders understand the appropriate use cases and ethical implications of using the model:

  • Intended Use (intendedUse aspect): Documents primary use cases, intended users, and out-of-scope applications
  • Ethical Considerations (mlModelEthicalConsiderations aspect): Documents use of sensitive data, risks and harms, mitigation strategies
  • Caveats and Recommendations (mlModelCaveatsAndRecommendations aspect): Additional considerations, ideal dataset characteristics, and usage recommendations

These aspects align with responsible AI practices and help ensure models are used appropriately.

Training and Evaluation Data​

ML Models can document their training and evaluation datasets in two complementary ways:

Direct Dataset References​

  • Training Data (mlModelTrainingData aspect): Datasets used to train the model, including preprocessing information and motivation for dataset selection
  • Evaluation Data (mlModelEvaluationData aspect): Datasets used for model evaluation and testing

Each dataset reference includes the dataset URN, motivation for using that dataset, and any preprocessing steps applied. This creates direct lineage relationships between models and their training data.

Lineage via Training Runs​

Training runs (dataProcessInstance entities) provide an alternative and often more detailed way to capture training lineage:

  • Training runs declare their input datasets via dataProcessInstanceInput aspect
  • Training runs declare their output datasets via dataProcessInstanceOutput aspect
  • Models reference training runs via the trainingJobs field

This creates indirect lineage: Dataset → Training Run → Model

When to use each approach:

  • Use direct dataset references for simple documentation of what data was used
  • Use training runs for complete lineage tracking including:
    • Multiple training/validation/test datasets
    • Metrics and hyperparameters from the training process
    • Temporal tracking (when the training occurred)
    • Connection to experiments for comparing multiple training attempts

Most production ML systems should use training runs for comprehensive lineage tracking.

Factor Prompts and Quantitative Analysis​

For detailed model analysis and performance reporting:

  • Factor Prompts (mlModelFactorPrompts aspect): Factors that may affect model performance (demographic groups, environmental conditions, etc.)
  • Quantitative Analyses (mlModelQuantitativeAnalyses aspect): Links to dashboards or reports showing disaggregated performance metrics across different factors
  • Metrics (mlModelMetrics aspect): Detailed metrics with descriptions beyond simple training/online metrics

Source Code and Cost​

  • Source Code (sourceCode aspect): Links to model training code, notebooks, or repositories (GitHub, GitLab, etc.)
  • Cost (cost aspect): Cost attribution information for tracking model training and inference expenses

Training Runs and Experiments​

ML Models in DataHub can be linked to their training runs and experiments, providing complete lineage from raw data through training to deployed models.

Training Runs​

Training runs represent specific executions of model training jobs. In DataHub, training runs are modeled as dataProcessInstance entities with a specialized subtype:

  • Entity Type: dataProcessInstance
  • Subtype: MLAssetSubTypes.MLFLOW_TRAINING_RUN
  • Key Aspects:
    • dataProcessInstanceProperties: Basic properties like name, timestamps, and custom properties
    • mlTrainingRunProperties: ML-specific properties including:
      • Training metrics (accuracy, loss, F1 score, etc.)
      • Hyperparameters (learning rate, batch size, epochs, etc.)
      • Output URLs (model artifacts, checkpoints)
      • External URLs (links to training dashboards)
    • dataProcessInstanceInput: Input datasets used for training
    • dataProcessInstanceOutput: Output datasets (predictions, feature importance, etc.)
    • dataProcessInstanceRunEvent: Start, completion, and failure events

Training runs create lineage relationships showing:

  • Upstream: Which datasets were used for training
  • Downstream: Which models were produced by the training run

Models reference their training runs through the trainingJobs field in mlModelProperties, and model groups can also reference training runs to track all training activity for a model family.

Experiments​

Experiments organize related training runs into logical groups, typically representing a series of attempts to optimize a model or compare different approaches. In DataHub, experiments are modeled as container entities:

  • Entity Type: container
  • Subtype: MLAssetSubTypes.MLFLOW_EXPERIMENT
  • Purpose: Group related training runs for organization and comparison

Training runs belong to experiments through the container aspect, creating a hierarchy:

Experiment: "Customer Churn Prediction"
├── Training Run 1: baseline model
├── Training Run 2: with feature engineering
├── Training Run 3: hyperparameter tuning
└── Training Run 4: final production model

This structure mirrors common ML platform patterns (like MLflow's experiment/run hierarchy) and enables:

  • Comparing metrics across multiple training attempts
  • Tracking the evolution of a model through iterations
  • Understanding which approaches were tried and their results
  • Organizing training work by project or objective
Python SDK: Create training runs and experiments
# Inlined from /metadata-ingestion/examples/ai/dh_ai_docs_demo.py
import argparse
from datetime import datetime

from dh_ai_client import DatahubAIClient

from datahub.emitter.mcp_builder import (
    ContainerKey,
)
from datahub.ingestion.source.common.subtypes import MLAssetSubTypes
from datahub.metadata.com.linkedin.pegasus2avro.dataprocess import RunResultType
from datahub.metadata.schema_classes import (
    AuditStampClass,
    DataProcessInstancePropertiesClass,
    MLHyperParamClass,
    MLMetricClass,
    MLTrainingRunPropertiesClass,
)
from datahub.metadata.urns import (
    CorpUserUrn,
    DataProcessInstanceUrn,
    GlossaryTermUrn,
    TagUrn,
)
from datahub.sdk.container import Container
from datahub.sdk.dataset import Dataset
from datahub.sdk.mlmodel import MLModel
from datahub.sdk.mlmodelgroup import MLModelGroup

parser = argparse.ArgumentParser()
parser.add_argument("--token", required=False, help="DataHub access token")
parser.add_argument(
    "--server_url",
    required=False,
    default="http://localhost:8080",
    help="DataHub server URL (defaults to http://localhost:8080)",
)
args = parser.parse_args()

# Initialize client
client = DatahubAIClient(token=args.token, server_url=args.server_url)

# Use a unique prefix for all IDs to avoid conflicts
prefix = "test"

# Define all entity IDs upfront
# Basic entity IDs
basic_model_group_id = f"{prefix}_basic_group"
basic_model_id = f"{prefix}_basic_model"
basic_experiment_id = f"{prefix}_basic_experiment"
basic_run_id = f"{prefix}_basic_run"
basic_dataset_id = f"{prefix}_basic_dataset"

# Advanced entity IDs
advanced_model_group_id = f"{prefix}_airline_forecast_models_group"
advanced_model_id = f"{prefix}_arima_model"
advanced_experiment_id = f"{prefix}_airline_forecast_experiment"
advanced_run_id = f"{prefix}_simple_training_run"
advanced_input_dataset_id = f"{prefix}_iris_input"
advanced_output_dataset_id = f"{prefix}_iris_output"

# Display names with prefix
basic_model_group_name = f"{prefix} Basic Group"
basic_model_name = f"{prefix} Basic Model"
basic_experiment_name = f"{prefix} Basic Experiment"
basic_run_name = f"{prefix} Basic Run"
basic_dataset_name = f"{prefix} Basic Dataset"

advanced_model_group_name = f"{prefix} Airline Forecast Models Group"
advanced_model_name = f"{prefix} ARIMA Model"
advanced_experiment_name = f"{prefix} Airline Forecast Experiment"
advanced_run_name = f"{prefix} Simple Training Run"
advanced_input_dataset_name = f"{prefix} Iris Training Input Data"
advanced_output_dataset_name = f"{prefix} Iris Model Output Data"


def create_basic_model_group():
    """Create a basic model group."""
    print("Creating basic model group...")
    basic_model_group = MLModelGroup(
        id=basic_model_group_id,
        platform="mlflow",
        name=basic_model_group_name,
    )
    client._emit_mcps(basic_model_group.as_mcps())
    return basic_model_group


def create_advanced_model_group():
    """Create an advanced model group."""
    print("Creating advanced model group...")
    advanced_model_group = MLModelGroup(
        id=advanced_model_group_id,
        platform="mlflow",
        name=advanced_model_group_name,
        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(advanced_model_group.as_mcps())
    return advanced_model_group


def create_basic_model():
    """Create a basic model."""
    print("Creating basic model...")
    basic_model = MLModel(
        id=basic_model_id,
        platform="mlflow",
        name=basic_model_name,
    )
    client._emit_mcps(basic_model.as_mcps())
    return basic_model


def create_advanced_model():
    """Create an advanced model."""
    print("Creating advanced model...")
    advanced_model = MLModel(
        id=advanced_model_id,
        platform="mlflow",
        name=advanced_model_name,
        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(advanced_model.as_mcps())
    return advanced_model


def create_basic_experiment():
    """Create a basic experiment."""
    print("Creating basic experiment...")
    basic_experiment = Container(
        container_key=ContainerKey(platform="mlflow", name=basic_experiment_id),
        display_name=basic_experiment_name,
    )
    client._emit_mcps(basic_experiment.as_mcps())
    return basic_experiment


def create_advanced_experiment():
    """Create an advanced experiment."""
    print("Creating advanced experiment...")
    advanced_experiment = Container(
        container_key=ContainerKey(platform="mlflow", name=advanced_experiment_id),
        display_name=advanced_experiment_name,
        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(advanced_experiment.as_mcps())
    return advanced_experiment


def create_basic_training_run():
    """Create a basic training run."""
    print("Creating basic training run...")
    basic_run_urn = client.create_training_run(
        run_id=basic_run_id,
        run_name=basic_run_name,
    )
    return basic_run_urn


def create_advanced_training_run():
    """Create an advanced training run."""
    print("Creating advanced training run...")
    advanced_run_urn = client.create_training_run(
        run_id=advanced_run_id,
        properties=DataProcessInstancePropertiesClass(
            name=advanced_run_name,
            created=AuditStampClass(
                time=1628580000000, actor="urn:li:corpuser:datahub"
            ),
            customProperties={"team": "forecasting"},
        ),
        training_run_properties=MLTrainingRunPropertiesClass(
            id=advanced_run_id,
            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,
    )
    return advanced_run_urn


def create_basic_dataset():
    """Create a basic dataset."""
    print("Creating basic dataset...")
    basic_input_dataset = Dataset(
        platform="snowflake",
        name=basic_dataset_id,
        display_name=basic_dataset_name,
    )
    client._emit_mcps(basic_input_dataset.as_mcps())
    return basic_input_dataset


def create_advanced_datasets():
    """Create advanced datasets."""
    print("Creating advanced datasets...")
    advanced_input_dataset = Dataset(
        platform="snowflake",
        name=advanced_input_dataset_id,
        description="Raw Iris dataset used for training ML models",
        schema=[("id", "number"), ("name", "string"), ("species", "string")],
        display_name=advanced_input_dataset_name,
        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(advanced_input_dataset.as_mcps())

    advanced_output_dataset = Dataset(
        platform="snowflake",
        name=advanced_output_dataset_id,
        description="Processed Iris dataset with model predictions",
        schema=[("id", "number"), ("name", "string"), ("species", "string")],
        display_name=advanced_output_dataset_name,
        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(advanced_output_dataset.as_mcps())
    return advanced_input_dataset, advanced_output_dataset


# Split relationship functions into individual top-level functions
def add_model_to_model_group(model, model_group):
    """Add model to model group relationship."""
    print("Adding model to model group...")
    model.set_model_group(model_group.urn)
    client._emit_mcps(model.as_mcps())


def add_run_to_experiment(run_urn, experiment):
    """Add run to experiment relationship."""
    print("Adding run to experiment...")
    client.add_run_to_experiment(run_urn=run_urn, experiment_urn=str(experiment.urn))


def add_run_to_model(model, run_id):
    """Add run to model relationship."""
    print("Adding run to model...")
    model.add_training_job(DataProcessInstanceUrn(run_id))
    client._emit_mcps(model.as_mcps())


def add_run_to_model_group(model_group, run_id):
    """Add run to model group relationship."""
    print("Adding run to model group...")
    model_group.add_training_job(DataProcessInstanceUrn(run_id))
    client._emit_mcps(model_group.as_mcps())


def add_input_dataset_to_run(run_urn, input_dataset):
    """Add input dataset to run relationship."""
    print("Adding input dataset to run...")
    client.add_input_datasets_to_run(
        run_urn=run_urn, dataset_urns=[str(input_dataset.urn)]
    )


def add_output_dataset_to_run(run_urn, output_dataset):
    """Add output dataset to run relationship."""
    print("Adding output dataset to run...")
    client.add_output_datasets_to_run(
        run_urn=run_urn, dataset_urns=[str(output_dataset.urn)]
    )


def update_model_properties(model):
    """Update model properties."""
    print("Updating model properties...")

    # 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())


def update_model_group_properties(model_group):
    """Update model group properties."""
    print("Updating model group properties...")

    # 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())


def update_experiment_properties():
    """Update experiment properties."""
    print("Updating experiment properties...")

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

    # 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())


def main():
    # Parse arguments
    print("Creating AI assets...")

    # Comment in/out the functions you want to run
    # Create basic entities
    create_basic_model_group()
    create_basic_model()
    create_basic_experiment()
    create_basic_training_run()
    create_basic_dataset()

    # Create advanced entities
    advanced_model_group = create_advanced_model_group()
    advanced_model = create_advanced_model()
    advanced_experiment = create_advanced_experiment()
    advanced_run_urn = create_advanced_training_run()
    advanced_input_dataset, advanced_output_dataset = create_advanced_datasets()

    # # Create relationships - each can be commented out independently
    add_model_to_model_group(advanced_model, advanced_model_group)
    add_run_to_experiment(advanced_run_urn, advanced_experiment)
    add_run_to_model(advanced_model, advanced_run_id)
    add_run_to_model_group(advanced_model_group, advanced_run_id)
    add_input_dataset_to_run(advanced_run_urn, advanced_input_dataset)
    add_output_dataset_to_run(advanced_run_urn, advanced_output_dataset)

    # # Update properties - each can be commented out independently
    update_model_properties(advanced_model)
    update_model_group_properties(advanced_model_group)
    update_experiment_properties()

    print("All done! AI entities created successfully.")


if __name__ == "__main__":
    main()

Relationships and Lineage​

ML Models support rich relationship modeling through various aspects and fields:

Core Relationships​

  • Model Groups (via groups field in mlModelProperties): Models can belong to mlModelGroup entities, creating a MemberOf relationship. This organizes related models into logical families or collections.

  • Training Runs (via trainingJobs field in mlModelProperties): Models reference dataProcessInstance entities with MLFLOW_TRAINING_RUN subtype that produced them. This creates upstream lineage showing:

    • Which training run created this model
    • What datasets were used for training (via the training run's input datasets)
    • What hyperparameters and metrics were recorded
    • Which experiment the training run belonged to

  • Features (via mlFeatures field in mlModelProperties): Models can consume mlFeature entities, creating a Consumes relationship. This documents:

    • Which features are required for model inference
    • The complete feature set used during training
    • Dependencies on feature stores or feature tables

  • Deployments (via deployments field in mlModelProperties): Models can be deployed to mlModelDeployment entities, representing running model endpoints in various environments (production, staging, etc.)

  • Training Datasets (via mlModelTrainingData aspect): Direct references to datasets used for training, including preprocessing information and motivation for dataset selection

  • Evaluation Datasets (via mlModelEvaluationData aspect): References to datasets used for model evaluation and testing

Lineage Graph Structure​

These relationships create a comprehensive lineage graph:

Training Datasets → Training Run → ML Model → ML Model Deployment
                         ↓
                    Experiment

Feature Tables → ML Features → ML Model

ML Model Group ← ML Model

This enables powerful queries such as:

  • "Show me all datasets that influenced this model's predictions"
  • "Which models will be affected if this dataset schema changes?"
  • "What's the full history of training runs that created versions of this model?"
  • "Which production endpoints are serving this model?"
Python SDK: Update model-specific aspects
# Inlined from /metadata-ingestion/examples/library/mlmodel_update_aspects.py
import datahub.metadata.schema_classes as models
from datahub.metadata.urns import DatasetUrn, MlModelUrn
from datahub.sdk import DataHubClient

client = DataHubClient.from_env()

model_urn = MlModelUrn(platform="mlflow", name="customer-churn-predictor")

mlmodel = client.entities.get(model_urn)

intended_use = models.IntendedUseClass(
    primaryUses=[
        "Predict customer churn to enable proactive retention campaigns",
        "Identify high-risk customers for targeted interventions",
    ],
    primaryUsers=[models.IntendedUserTypeClass.ENTERPRISE],
    outOfScopeUses=[
        "Not suitable for real-time predictions (batch inference only)",
        "Not trained on international markets outside North America",
    ],
)

mlmodel._set_aspect(intended_use)

training_data = models.TrainingDataClass(
    trainingData=[
        models.BaseDataClass(
            dataset=str(
                DatasetUrn(
                    platform="snowflake", name="prod.analytics.customer_features"
                )
            ),
            motivation="Historical customer data with confirmed churn labels",
            preProcessing=[
                "Removed customers with less than 30 days of history",
                "Standardized numerical features using StandardScaler",
                "One-hot encoded categorical variables",
            ],
        )
    ]
)

mlmodel._set_aspect(training_data)

source_code = models.SourceCodeClass(
    sourceCode=[
        models.SourceCodeUrlClass(
            type=models.SourceCodeUrlTypeClass.ML_MODEL_SOURCE_CODE,
            sourceCodeUrl="https://github.com/example/ml-models/tree/main/churn-predictor",
        )
    ]
)

mlmodel._set_aspect(source_code)

ethical_considerations = models.EthicalConsiderationsClass(
    data=["Model uses demographic data (age, location) which may be sensitive"],
    risksAndHarms=[
        "Predictions may disproportionately affect certain customer segments",
        "False positives could lead to unnecessary retention spending",
    ],
    mitigations=[
        "Regular bias audits conducted quarterly",
        "Human review required for high-value customer interventions",
    ],
)

mlmodel._set_aspect(ethical_considerations)

client.entities.update(mlmodel)

print(f"Updated aspects for model: {model_urn}")

Tags, Terms, and Ownership​

Like other DataHub entities, ML Models support:

  • Tags (globalTags aspect): Flexible categorization (e.g., "pii-model", "production-ready", "experimental")
  • Glossary Terms (glossaryTerms aspect): Business concepts (e.g., "Customer Churn", "Fraud Detection")
  • Ownership (ownership aspect): Individuals or teams responsible for the model (data scientists, ML engineers, etc.)
  • Domains (domains aspect): Organizational grouping (e.g., "Recommendations", "Risk Management")

Complete ML Workflow Example​

The following example demonstrates a complete ML model lifecycle in DataHub, showing how all the pieces work together:

1. Create Model Group
   ↓
2. Create Experiment (Container)
   ↓
3. Create Training Run (DataProcessInstance)
   ├── Link input datasets
   ├── Link output datasets
   └── Add metrics and hyperparameters
   ↓
4. Create Model
   ├── Set version and aliases
   ├── Link to model group
   ├── Link to training run
   ├── Add hyperparameters and metrics
   └── Add ownership and tags
   ↓
5. Link Training Run to Experiment
   ↓
6. Update Model properties as needed
   ├── Change version aliases (champion → challenger)
   ├── Add additional tags/terms
   └── Update metrics from production

This workflow creates rich lineage showing:

  • Which datasets trained the model
  • What experiments and training runs were involved
  • How the model evolved through versions
  • Which version is deployed (via aliases)
  • Who owns and maintains the model
Complete Python Example: Full ML Workflow

See the comprehensive example in /metadata-ingestion/examples/ai/dh_ai_docs_demo.py which demonstrates:

  • Creating model groups with metadata
  • Creating experiments to organize training runs
  • Creating training runs with metrics, hyperparameters, and dataset lineage
  • Creating models with versions and aliases
  • Linking all entities together to form complete lineage
  • Updating properties and managing the model lifecycle

The example shows both basic patterns for getting started and advanced patterns for production ML systems.

Code Examples​

Querying ML Model Information​

The standard REST APIs can be used to retrieve ML Model entities and their aspects:

Python: Query an ML Model via REST API
# Inlined from /metadata-ingestion/examples/library/mlmodel_query_rest_api.py
import urllib.parse

import requests

gms_server = "http://localhost:8080"

model_urn = "urn:li:mlModel:(urn:li:dataPlatform:mlflow,customer-churn-predictor,PROD)"
encoded_urn = urllib.parse.quote(model_urn, safe="")

response = requests.get(f"{gms_server}/entities/{encoded_urn}")

if response.status_code == 200:
    entity = response.json()

    print(f"Entity URN: {entity['urn']}")
    print("\nAspects:")

    if "mlModelProperties" in entity["aspects"]:
        props = entity["aspects"]["mlModelProperties"]
        print(f"  Name: {props.get('name')}")
        print(f"  Description: {props.get('description')}")
        print(f"  Type: {props.get('type')}")

        if props.get("hyperParams"):
            print("\n  Hyperparameters:")
            for param in props["hyperParams"]:
                print(f"    - {param['name']}: {param['value']}")

        if props.get("trainingMetrics"):
            print("\n  Training Metrics:")
            for metric in props["trainingMetrics"]:
                print(f"    - {metric['name']}: {metric['value']}")

    if "globalTags" in entity["aspects"]:
        tags = entity["aspects"]["globalTags"]["tags"]
        print(f"\n  Tags: {[tag['tag'] for tag in tags]}")

    if "ownership" in entity["aspects"]:
        owners = entity["aspects"]["ownership"]["owners"]
        print(f"\n  Owners: {[owner['owner'] for owner in owners]}")

    if "intendedUse" in entity["aspects"]:
        intended = entity["aspects"]["intendedUse"]
        print(f"\n  Primary Uses: {intended.get('primaryUses')}")
        print(f"  Out of Scope Uses: {intended.get('outOfScopeUses')}")

else:
    print(f"Failed to fetch entity: {response.status_code}")
    print(response.text)

Integration Points​

ML Models integrate with several other entities in the DataHub metadata model:

  • mlModelGroup: Logical grouping of related model versions (e.g., all versions of a recommendation model)
  • mlModelDeployment: Running instances of deployed models with status, endpoint URLs, and deployment metadata
  • mlFeature: Individual features consumed by the model for inference
  • mlFeatureTable: Collections of features, often from feature stores
  • dataset: Training and evaluation datasets used by the model
  • dataProcessInstance (with MLFLOW_TRAINING_RUN subtype): Specific training runs that created model versions, including metrics, hyperparameters, and lineage to input/output datasets
  • container (with MLFLOW_EXPERIMENT subtype): Experiments that organize related training runs for a model or project
  • versionSet: Groups all versions of a model together for version management

GraphQL Resolvers​

The GraphQL API provides rich querying capabilities for ML Models through resolvers in datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/. These resolvers support:

  • Fetching model details with all aspects
  • Navigating relationships to features, groups, and deployments
  • Searching and filtering models by tags, terms, platform, etc.

Ingestion Sources​

Several ingestion sources automatically extract ML Model metadata:

  • MLflow: Extracts registered models, versions, metrics, parameters, and lineage from MLflow tracking servers
  • SageMaker: Ingests models, model packages, and endpoints from AWS SageMaker
  • Vertex AI: Extracts models and endpoints from Google Cloud Vertex AI
  • Databricks: Ingests MLflow models from Databricks workspaces
  • Unity Catalog: Extracts ML models registered in Unity Catalog

These sources are located in /metadata-ingestion/src/datahub/ingestion/source/ and automatically populate model properties, relationships, and lineage.

Notable Exceptions​

Model Versioning​

ML Model versioning in DataHub uses the versionProperties aspect, which provides a robust framework for tracking model versions across their lifecycle. This is the standard approach demonstrated in production ML platforms.

Version Properties Aspect​

Every ML Model should use the versionProperties aspect, which includes:

  • version: A VersionTagClass containing the version identifier (e.g., "1", "2", "v1.0.0")
  • versionSet: A URN that groups all versions of a model together (e.g., urn:li:versionSet:(mlModel,mlmodel_my-model_versions))
  • sortId: A string used for ordering versions (typically the version number zero-padded)
  • aliases: Optional array of VersionTagClass objects for named version references

Version Aliases for A/B Testing​

Version aliases enable flexible model lifecycle management and A/B testing workflows. Common aliases include:

  • "champion": The currently deployed production model
  • "challenger": A candidate model being tested or evaluated
  • "baseline": A reference model for performance comparison
  • "latest": The most recently trained version

These aliases allow you to reference models by their role rather than specific version numbers, enabling smooth model promotion workflows:

Model v1 (alias: "champion")     # Currently in production
Model v2 (alias: "challenger")   # Being tested in canary deployment
Model v3 (alias: "latest")       # Just completed training

When v2 proves superior, you can update aliases without changing infrastructure:

Model v1 (no alias)              # Retired
Model v2 (alias: "champion")     # Promoted to production
Model v3 (alias: "challenger")   # Now being tested

Model Groups and Versioning​

Model groups (mlModelGroup entities) serve as logical containers for organizing related models. While model groups can contain multiple versions of the same model, versioning is handled through the versionProperties aspect on individual models, not through the group structure itself. Model groups are used for:

  • Organizing all versions of a model family
  • Grouping experimental variants or different architectures solving the same problem
  • Managing lineage and metadata common across multiple related models

The relationship between models and model groups is through the groups field in mlModelProperties, creating a MemberOf relationship.

Platform-Specific Naming​

Different ML platforms have different naming conventions:

  • MLflow: Uses a two-level hierarchy (registered model name + version number). In DataHub, each version can be a separate entity, or versions can be tracked in a single entity.
  • SageMaker: Has multiple model concepts (model, model package, model package group). DataHub can model these as separate entities or consolidate them.
  • Vertex AI: Uses fully qualified resource names. These should be simplified to human-readable names when possible.

When ingesting from these platforms, connectors handle platform-specific naming and convert it to appropriate DataHub URNs.

Model Cards​

The various aspects (intendedUse, mlModelFactorPrompts, mlModelEthicalConsiderations, etc.) follow the Model Cards for Model Reporting framework (Mitchell et al., 2019). While these aspects are optional, they are strongly recommended for production models to ensure responsible AI practices and transparent model documentation.

Technical Reference Guide​

The sections above provide an overview of how to use this entity. The following sections provide detailed technical information about how metadata is stored and represented in DataHub.

Aspects are the individual pieces of metadata that can be attached to an entity. Each aspect contains specific information (like ownership, tags, or properties) and is stored as a separate record, allowing for flexible and incremental metadata updates.

Relationships show how this entity connects to other entities in the metadata graph. These connections are derived from the fields within each aspect and form the foundation of DataHub's knowledge graph.

Reading the Field Tables​

Each aspect's field table includes an Annotations column that provides additional metadata about how fields are used:

  • ⚠️ Deprecated: This field is deprecated and may be removed in a future version. Check the description for the recommended alternative
  • Searchable: This field is indexed and can be searched in DataHub's search interface
  • Searchable (fieldname): When the field name in parentheses is shown, it indicates the field is indexed under a different name in the search index. For example, dashboardTool is indexed as tool
  • → RelationshipName: This field creates a relationship to another entity. The arrow indicates this field contains a reference (URN) to another entity, and the name indicates the type of relationship (e.g., → Contains, → OwnedBy)

Fields with complex types (like Edge, AuditStamp) link to their definitions in the Common Types section below.

Aspects​

mlModelKey​

Key for an ML model

FieldTypeRequiredDescriptionAnnotations
platformstring✓Standardized platform urn for the model
namestring✓Name of the MLModelSearchable
originFabricType✓Fabric type where model belongs to or where it was generatedSearchable

ownership​

Ownership information of an entity.

FieldTypeRequiredDescriptionAnnotations
ownersOwner[]✓List of owners of the entity.
ownerTypesmapOwnership type to Owners map, populated via mutation hook.Searchable
lastModifiedAuditStamp✓Audit stamp containing who last modified the record and when. A value of 0 in the time field indi...

mlModelProperties​

Properties associated with a ML Model

FieldTypeRequiredDescriptionAnnotations
customPropertiesmap✓Custom property bag.Searchable
externalUrlstringURL where the reference existSearchable
trainingJobsstring[]List of jobs or process instances (if any) used to train the model or group. Visible in Lineage. ...→ TrainedBy
downstreamJobsstring[]List of jobs or process instances (if any) that use the model or group.→ UsedBy
namestringDisplay name of the MLModelSearchable
descriptionstringDocumentation of the MLModelSearchable
datelongDate when the MLModel was developed⚠️ Deprecated
createdTimeStampAudit stamp containing who created this and when
lastModifiedTimeStampDate when the MLModel was last modified
versionVersionTagVersion of the MLModel
typestringType of Algorithm or MLModel such as whether it is a Naive Bayes classifier, Convolutional Neural...Searchable
hyperParametersmapHyper Parameters of the MLModel NOTE: these are deprecated in favor of hyperParams
hyperParamsMLHyperParam[]Hyperparameters of the MLModel
trainingMetricsMLMetric[]Metrics of the MLModel used in training
onlineMetricsMLMetric[]Metrics of the MLModel used in production
mlFeaturesstring[]List of features used for MLModel training→ Consumes
tagsstring[]✓Tags for the MLModel
deploymentsstring[]Deployments for the MLModel→ DeployedTo
groupsstring[]Groups the model belongs to→ MemberOf

intendedUse​

Intended Use for the ML Model

FieldTypeRequiredDescriptionAnnotations
primaryUsesstring[]Primary Use cases for the MLModel.
primaryUsersIntendedUserType[]Primary Intended Users - For example, was the MLModel developed for entertainment purposes, for h...
outOfScopeUsesstring[]Highlight technology that the MLModel might easily be confused with, or related contexts that use...

mlModelFactorPrompts​

Prompts which affect the performance of the MLModel

FieldTypeRequiredDescriptionAnnotations
relevantFactorsMLModelFactors[]What are foreseeable salient factors for which MLModel performance may vary, and how were these d...
evaluationFactorsMLModelFactors[]Which factors are being reported, and why were these chosen?

mlModelMetrics​

Metrics to be featured for the MLModel.

FieldTypeRequiredDescriptionAnnotations
performanceMeasuresstring[]Measures of MLModel performance
decisionThresholdstring[]Decision Thresholds used (if any)?

mlModelEvaluationData​

All referenced datasets would ideally point to any set of documents that provide visibility into the source and composition of the dataset.

FieldTypeRequiredDescriptionAnnotations
evaluationDataBaseData[]✓Details on the dataset(s) used for the quantitative analyses in the MLModel

mlModelTrainingData​

Ideally, the MLModel card would contain as much information about the training data as the evaluation data. However, there might be cases where it is not feasible to provide this level of detailed information about the training data. For example, the data may be proprietary, or require a non-disclosure agreement. In these cases, we advocate for basic details about the distributions over groups in the data, as well as any other details that could inform stakeholders on the kinds of biases the model may have encoded.

FieldTypeRequiredDescriptionAnnotations
trainingDataBaseData[]✓Details on the dataset(s) used for training the MLModel

mlModelQuantitativeAnalyses​

Quantitative analyses should be disaggregated, that is, broken down by the chosen factors. Quantitative analyses should provide the results of evaluating the MLModel according to the chosen metrics, providing confidence interval values when possible.

FieldTypeRequiredDescriptionAnnotations
unitaryResultsstringLink to a dashboard with results showing how the MLModel performed with respect to each factor
intersectionalResultsstringLink to a dashboard with results showing how the MLModel performed with respect to the intersecti...

mlModelEthicalConsiderations​

This section is intended to demonstrate the ethical considerations that went into MLModel development, surfacing ethical challenges and solutions to stakeholders.

FieldTypeRequiredDescriptionAnnotations
datastring[]Does the MLModel use any sensitive data (e.g., protected classes)?
humanLifestring[]Is the MLModel intended to inform decisions about matters central to human life or flourishing - ...
mitigationsstring[]What risk mitigation strategies were used during MLModel development?
risksAndHarmsstring[]What risks may be present in MLModel usage? Try to identify the potential recipients, likelihood,...
useCasesstring[]Are there any known MLModel use cases that are especially fraught? This may connect directly to t...

mlModelCaveatsAndRecommendations​

This section should list additional concerns that were not covered in the previous sections. For example, did the results suggest any further testing? Were there any relevant groups that were not represented in the evaluation dataset? Are there additional recommendations for model use?

FieldTypeRequiredDescriptionAnnotations
caveatsCaveatDetailsThis section should list additional concerns that were not covered in the previous sections. For ...
recommendationsstringRecommendations on where this MLModel should be used.
idealDatasetCharacteristicsstring[]Ideal characteristics of an evaluation dataset for this MLModel

institutionalMemory​

Institutional memory of an entity. This is a way to link to relevant documentation and provide description of the documentation. Institutional or tribal knowledge is very important for users to leverage the entity.

FieldTypeRequiredDescriptionAnnotations
elementsInstitutionalMemoryMetadata[]✓List of records that represent institutional memory of an entity. Each record consists of a link,...

sourceCode​

Source Code

FieldTypeRequiredDescriptionAnnotations
sourceCodeSourceCodeUrl[]✓Source Code along with types

status​

The lifecycle status metadata of an entity, e.g. dataset, metric, feature, etc. This aspect is used to represent soft deletes conventionally.

FieldTypeRequiredDescriptionAnnotations
removedboolean✓Whether the entity has been removed (soft-deleted).Searchable

cost​

None

FieldTypeRequiredDescriptionAnnotations
costTypeCostType✓
costCostCost✓

deprecation​

Deprecation status of an entity

FieldTypeRequiredDescriptionAnnotations
deprecatedboolean✓Whether the entity is deprecated.Searchable
decommissionTimelongThe time user plan to decommission this entity.
notestring✓Additional information about the entity deprecation plan, such as the wiki, doc, RB.
actorstring✓The user URN which will be credited for modifying this deprecation content.
replacementstring

browsePaths​

Shared aspect containing Browse Paths to be indexed for an entity.

FieldTypeRequiredDescriptionAnnotations
pathsstring[]✓A list of valid browse paths for the entity. Browse paths are expected to be forward slash-separ...Searchable

globalTags​

Tag aspect used for applying tags to an entity

FieldTypeRequiredDescriptionAnnotations
tagsTagAssociation[]✓Tags associated with a given entitySearchable, → TaggedWith

dataPlatformInstance​

The specific instance of the data platform that this entity belongs to

FieldTypeRequiredDescriptionAnnotations
platformstring✓Data PlatformSearchable
instancestringInstance of the data platform (e.g. db instance)Searchable (platformInstance)

browsePathsV2​

Shared aspect containing a Browse Path to be indexed for an entity.

FieldTypeRequiredDescriptionAnnotations
pathBrowsePathEntry[]✓A valid browse path for the entity. This field is provided by DataHub by default. This aspect is ...Searchable

glossaryTerms​

Related business terms information

FieldTypeRequiredDescriptionAnnotations
termsGlossaryTermAssociation[]✓The related business terms
auditStampAuditStamp✓Audit stamp containing who reported the related business term

editableMlModelProperties​

Properties associated with a ML Model editable from the UI

FieldTypeRequiredDescriptionAnnotations
descriptionstringDocumentation of the ml modelSearchable (editedDescription)

domains​

Links from an Asset to its Domains

FieldTypeRequiredDescriptionAnnotations
domainsstring[]✓The Domains attached to an AssetSearchable, → AssociatedWith

applications​

Links from an Asset to its Applications

FieldTypeRequiredDescriptionAnnotations
applicationsstring[]✓The Applications attached to an AssetSearchable, → AssociatedWith

structuredProperties​

Properties about an entity governed by StructuredPropertyDefinition

FieldTypeRequiredDescriptionAnnotations
propertiesStructuredPropertyValueAssignment[]✓Custom property bag.

forms​

Forms that are assigned to this entity to be filled out

FieldTypeRequiredDescriptionAnnotations
incompleteFormsFormAssociation[]✓All incomplete forms assigned to the entity.Searchable
completedFormsFormAssociation[]✓All complete forms assigned to the entity.Searchable
verificationsFormVerificationAssociation[]✓Verifications that have been applied to the entity via completed forms.Searchable

testResults​

Information about a Test Result

FieldTypeRequiredDescriptionAnnotations
failingTestResult[]✓Results that are failingSearchable, → IsFailing
passingTestResult[]✓Results that are passingSearchable, → IsPassing

versionProperties​

Properties about a versioned asset i.e. dataset, ML Model, etc.

FieldTypeRequiredDescriptionAnnotations
versionSetstring✓The linked Version Set entity that ties multiple versioned assets togetherSearchable, → VersionOf
versionVersionTag✓Label for this versioned asset, is unique within a version setSearchable
aliasesVersionTag[]✓Associated aliases for this versioned assetSearchable
commentstringComment documenting what this version was created for, changes, or represents
sortIdstring✓Sort identifier that determines where a version lives in the order of the Version Set. What this ...Searchable (versionSortId)
versioningSchemeVersioningScheme✓What versioning scheme sortId belongs to. Defaults to a plain string that is lexicographically ...
sourceCreatedTimestampAuditStampTimestamp reflecting when this asset version was created in the source system.
metadataCreatedTimestampAuditStampTimestamp reflecting when the metadata for this version was created in DataHub
isLatestbooleanMarks whether this version is currently the latest. Set by a side effect and should not be modifi...Searchable

subTypes​

Sub Types. Use this aspect to specialize a generic Entity e.g. Making a Dataset also be a View or also be a LookerExplore

FieldTypeRequiredDescriptionAnnotations
typeNamesstring[]✓The names of the specific types.Searchable

container​

Link from an asset to its parent container

FieldTypeRequiredDescriptionAnnotations
containerstring✓The parent container of an assetSearchable, → IsPartOf

Common Types​

These types are used across multiple aspects in this entity.

AuditStamp​

Data captured on a resource/association/sub-resource level giving insight into when that resource/association/sub-resource moved into a particular lifecycle stage, and who acted to move it into that specific lifecycle stage.

Fields:

  • time (long): When did the resource/association/sub-resource move into the specific lifecyc...
  • actor (string): The entity (e.g. a member URN) which will be credited for moving the resource...
  • impersonator (string?): The entity (e.g. a service URN) which performs the change on behalf of the Ac...
  • message (string?): Additional context around how DataHub was informed of the particular change. ...

BaseData​

BaseData record

Fields:

  • dataset (string): What dataset were used in the MLModel?
  • motivation (string?): Why was this dataset chosen?
  • preProcessing (string[]?): How was the data preprocessed (e.g., tokenization of sentences, cropping of i...

FormAssociation​

Properties of an applied form.

Fields:

  • urn (string): Urn of the applied form
  • incompletePrompts (FormPromptAssociation[]): A list of prompts that are not yet complete for this form.
  • completedPrompts (FormPromptAssociation[]): A list of prompts that have been completed for this form.

MLMetric​

Properties associated with an ML Metric

Fields:

  • name (string): Name of the mlMetric
  • description (string?): Documentation of the mlMetric
  • value (string?): The value of the mlMetric
  • createdAt (long?): Date when the mlMetric was developed

MLModelFactors​

Factors affecting the performance of the MLModel.

Fields:

  • groups (string[]?): Groups refers to distinct categories with similar characteristics that are pr...
  • instrumentation (string[]?): The performance of a MLModel can vary depending on what instruments were used...
  • environment (string[]?): A further factor affecting MLModel performance is the environment in which it...

TestResult​

Information about a Test Result

Fields:

  • test (string): The urn of the test
  • type (TestResultType): The type of the result
  • testDefinitionMd5 (string?): The md5 of the test definition that was used to compute this result. See Test...
  • lastComputed (AuditStamp?): The audit stamp of when the result was computed, including the actor who comp...

TimeStamp​

A standard event timestamp

Fields:

  • time (long): When did the event occur
  • actor (string?): Optional: The actor urn involved in the event.

VersionTag​

A resource-defined string representing the resource state for the purpose of concurrency control

Fields:

  • versionTag (string?):
  • metadataAttribution (MetadataAttribution?):

Relationships​

Outgoing​

These are the relationships stored in this entity's aspects

  • OwnedBy

    • Corpuser via ownership.owners.owner
    • CorpGroup via ownership.owners.owner

  • ownershipType

    • OwnershipType via ownership.owners.typeUrn

  • TrainedBy

    • DataJob via mlModelProperties.trainingJobs
    • DataProcessInstance via mlModelProperties.trainingJobs

  • UsedBy

    • DataJob via mlModelProperties.downstreamJobs
    • DataProcessInstance via mlModelProperties.downstreamJobs

  • Consumes

    • MlFeature via mlModelProperties.mlFeatures

  • DeployedTo

    • MlModelDeployment via mlModelProperties.deployments

  • MemberOf

    • MlModelGroup via mlModelProperties.groups

  • TaggedWith

    • Tag via globalTags.tags

  • TermedWith

    • GlossaryTerm via glossaryTerms.terms.urn

  • AssociatedWith

    • Domain via domains.domains
    • Application via applications.applications

  • IsFailing

    • Test via testResults.failing

  • IsPassing

    • Test via testResults.passing

  • VersionOf

    • VersionSet via versionProperties.versionSet

  • IsPartOf

    • Container via container.container

Incoming​

These are the relationships stored in other entity's aspects

  • Consumes

    • DataProcessInstance via dataProcessInstanceInput.inputs

  • DataProcessInstanceConsumes

    • DataProcessInstance via dataProcessInstanceInput.inputEdges

  • Produces

    • DataProcessInstance via dataProcessInstanceOutput.outputs

  • DataProcessInstanceProduces

    • DataProcessInstance via dataProcessInstanceOutput.outputEdges

Global Metadata Model​

Global Graph