Skip to main content
Version: Next

ML Model Group

ML Model Groups represent collections of related machine learning models within an organization's ML infrastructure. They serve as logical containers for organizing model versions, experimental variants, or families of models that share common characteristics. Model groups are essential for managing the lifecycle of ML models, tracking model evolution over time, and organizing models by purpose, architecture, or business function.

Identity​

ML Model Groups are identified by three pieces of information:

  • Platform: The ML platform or tool where the model group exists. This represents the specific ML technology that hosts the model group. Examples include mlflow, sagemaker, databricks, kubeflow, tensorflow, pytorch, etc. The platform is represented as a URN like urn:li:dataPlatform:mlflow.

  • Name: The unique name of the model group within the specific platform. This is typically a human-readable identifier that describes the purpose or family of models. Names should be meaningful and follow your organization's naming conventions. Examples include recommendation-models, fraud-detection-v2, or customer-churn-prediction.

  • Origin (Fabric): The environment or fabric where the model group belongs or was generated. This qualifier helps distinguish between models in different environments such as Production (PROD), Staging (QA), Development (DEV), or Testing environments. The full list of supported environments is available in FabricType.pdl.

An example of an ML Model Group identifier is urn:li:mlModelGroup:(urn:li:dataPlatform:mlflow,recommendation-models,PROD).

Important Capabilities​

ML Model Group Properties​

Model group properties are stored in the mlModelGroupProperties aspect and contain the core metadata about a model group:

  • Name: The display name of the model group, which can be more descriptive than the identifier
  • Description: Detailed documentation about what the model group represents, its purpose, and any important context
  • Version: An optional version tag for the entire model group (distinct from individual model versions)
  • Created Timestamp: Audit information about when and who created the model group
  • Last Modified Timestamp: When the model group was last updated
  • Custom Properties: Extensible key-value pairs for storing additional metadata specific to your organization or platform
  • External References: Links to external systems or documentation (e.g., model registry URLs, experiment tracking systems)

Here is an example of creating an ML Model Group with properties:

Python SDK: Create an ML Model Group
# Inlined from /metadata-ingestion/examples/library/mlmodel_group_create.py
from datahub.sdk import DataHubClient
from datahub.sdk.mlmodelgroup import MLModelGroup

client = DataHubClient.from_env()

mlmodel_group = MLModelGroup(
    id="my-recommendations-model-group",
    name="My Recommendations Model Group",
    platform="mlflow",
    description="Grouping of ml model related to home page recommendations",
    custom_properties={
        "framework": "pytorch",
    },
)

client.entities.upsert(mlmodel_group)
print(f"Created ML model group: {mlmodel_group.urn}")

Lineage and Training Information​

ML Model Groups inherit lineage capabilities from the MLModelLineageInfo record, which captures important information about how models in the group are created and used:

  • Training Jobs: References to data jobs or process instances that were used to train models in this group. This creates lineage relationships showing where models come from.
  • Downstream Jobs: References to data jobs or process instances that consume or use models from this group. This tracks how models are deployed and utilized.

These lineage relationships are visible in DataHub's lineage graph and help track the full lifecycle of ML models from training data through deployment.

Python SDK: Add training lineage to a model group
# Inlined from /metadata-ingestion/examples/library/mlmodel_group_add_lineage.py
from datahub.emitter import mce_builder
from datahub.metadata.urns import MlModelGroupUrn
from datahub.sdk import DataHubClient

client = DataHubClient.from_env()

group_urn = MlModelGroupUrn(
    platform="mlflow",
    name="recommendation-models",
    env="PROD",
)

training_job_urn = mce_builder.make_data_job_urn(
    orchestrator="airflow",
    flow_id="train_recommendation_model",
    job_id="training_task",
)

group = client.entities.get(group_urn)

group.add_training_job(training_job_urn)

client.entities.update(group)

print(f"Added training job {training_job_urn} to ML model group {group_urn}")

Ownership​

Like other entities in DataHub, ML Model Groups can have owners assigned using the ownership aspect. Model group owners are typically responsible for:

  • Managing which models belong to the group
  • Maintaining model group metadata and documentation
  • Overseeing model quality and governance standards
  • Serving as points of contact for model-related questions
  • Coordinating model deployment and monitoring

Ownership types for model groups follow the same patterns as other entities, including TECHNICAL_OWNER, BUSINESS_OWNER, DATA_STEWARD, DATAOWNER, PRODUCER, DEVELOPER, etc.

Python SDK: Add an owner to a model group
# Inlined from /metadata-ingestion/examples/library/mlmodel_group_add_owner.py
from datahub.sdk import CorpUserUrn, DataHubClient, MlModelGroupUrn

client = DataHubClient.from_env()

group = client.entities.get(
    MlModelGroupUrn(platform="mlflow", name="recommendation-models", env="PROD")
)

group.add_owner(CorpUserUrn("data_science_team"))

client.entities.update(group)

Tags and Glossary Terms​

ML Model Groups support both tags and glossary terms for categorization and discovery:

  • Tags (via globalTags aspect): Informal labels for quick categorization. Use tags for properties like experimental, production-ready, high-priority, computer-vision, nlp, etc.
  • Glossary Terms (via glossaryTerms aspect): Formal business vocabulary terms that provide standardized definitions. Use terms to link model groups to business concepts.
Python SDK: Add tags to a model group
# Inlined from /metadata-ingestion/examples/library/mlmodel_group_add_tags.py
from datahub.metadata.urns import MlModelGroupUrn, TagUrn
from datahub.sdk import DataHubClient

client = DataHubClient.from_env()

group = client.entities.get(
    MlModelGroupUrn(platform="mlflow", name="recommendation-models", env="PROD")
)
group.add_tag(TagUrn("production-ready"))

client.entities.update(group)

Python SDK: Add glossary terms to a model group
# Inlined from /metadata-ingestion/examples/library/mlmodel_group_add_terms.py
from datahub.metadata.urns import MlModelGroupUrn
from datahub.sdk import DataHubClient, GlossaryTermUrn

client = DataHubClient.from_env()

group_urn = MlModelGroupUrn(platform="mlflow", name="recommendation-models", env="PROD")

mlmodel_group = client.entities.get(group_urn)

mlmodel_group.add_term(GlossaryTermUrn("Recommendation"))

client.entities.update(mlmodel_group)

print(f"Added term {GlossaryTermUrn('Recommendation')} to ML model group {group_urn}")

Model groups support documentation through multiple aspects:

  • Description in mlModelGroupProperties: Primary documentation field for describing the model group
  • Institutional Memory (via institutionalMemory aspect): Links to external resources such as:
    • Confluence pages describing the model group's purpose and architecture
    • Model cards and documentation
    • Experiment tracking dashboards (MLflow, Weights & Biases, etc.)
    • Research papers or technical specifications
    • Team wikis or runbooks
Python SDK: Add documentation links to a model group
# Inlined from /metadata-ingestion/examples/library/mlmodel_group_add_documentation.py
from datahub.sdk import DataHubClient
from datahub.sdk.mlmodelgroup import MLModelGroup

client = DataHubClient.from_env()

mlmodel_group = client.entities.get(
    MLModelGroup.get_urn_type()(
        platform="mlflow", name="recommendation-models", env="PROD"
    )
)

doc_url = "https://wiki.example.com/ml/recommendation-models"
doc_description = "Model architecture and training documentation"

mlmodel_group.add_link((doc_url, doc_description))

client.entities.update(mlmodel_group)

Domains​

ML Model Groups can be assigned to domains using the domains aspect. This allows organizing model groups by business unit, department, or functional area. A model group can belong to only one domain at a time.

Python SDK: Assign a model group to a domain
# Inlined from /metadata-ingestion/examples/library/mlmodel_group_add_domain.py
from datahub.metadata.urns import DomainUrn, MlModelGroupUrn
from datahub.sdk import DataHubClient

client = DataHubClient.from_env()

mlmodel_group = client.entities.get(
    MlModelGroupUrn(platform="mlflow", name="recommendation-models", env="PROD")
)

# If you don't know the domain urn, you can look it up:
# domain_urn = client.resolve.domain(name="marketing")

# NOTE: This will overwrite the existing domain
mlmodel_group.set_domain(DomainUrn(id="marketing"))

client.entities.update(mlmodel_group)

Deprecation​

Model groups can be marked as deprecated using the deprecation aspect when they are no longer actively maintained or should be replaced. This helps users understand which model families are still supported.

Python SDK: Deprecate a model group
# Inlined from /metadata-ingestion/examples/library/mlmodel_group_deprecate.py
from datetime import datetime

import datahub.metadata.schema_classes as models
from datahub.sdk import DataHubClient, MlModelGroupUrn

client = DataHubClient.from_env()

group_urn = MlModelGroupUrn(
    platform="mlflow",
    name="legacy-recommendation-models",
    env="PROD",
)

mlmodel_group = client.entities.get(group_urn)

deprecation_aspect = models.DeprecationClass(
    deprecated=True,
    note="This model group has been replaced by the new transformer-based recommendation models",
    decommissionTime=int(datetime.now().timestamp() * 1000),
    actor="urn:li:corpuser:datahub",
)

mlmodel_group._set_aspect(deprecation_aspect)

client.entities.update(mlmodel_group)
print(f"Deprecated ML model group: {group_urn}")

Structured Properties​

Model groups support structured properties for storing typed, schema-validated metadata that goes beyond simple key-value pairs. This is useful for enforcing organizational standards around model metadata.

Integration Points​

Relationship to ML Models​

The primary relationship for ML Model Groups is with ML Models themselves. Models can be associated with a model group through the groups field in the mlModelProperties aspect. This creates a MemberOf relationship from the model to the model group.

Common patterns for organizing models in groups include:

  1. Version-based grouping: All versions of a model (v1, v2, v3) belong to the same group. Note that versioning is handled through the versionProperties aspect on individual models (which includes version numbers, versionSet URNs, and aliases like "champion" or "challenger"), while the model group serves as the organizational container.

  2. Experiment-based grouping: Different experimental variants of a model belong to the same group

  3. Architecture-based grouping: Models sharing the same architecture or approach

  4. Purpose-based grouping: Models serving the same business purpose or use case

Python SDK: Add a model to a model group
# Inlined from /metadata-ingestion/examples/library/mlgroup_add_to_mlmodel.py
from datahub.metadata.urns import MlModelGroupUrn
from datahub.sdk import DataHubClient
from datahub.sdk.mlmodel import MLModel

client = DataHubClient.from_env()

model = MLModel(
    id="my-recommendations-model",
    platform="mlflow",
)

model.set_model_group(
    MlModelGroupUrn(
        platform="mlflow",
        name="my-recommendations-model-group",
    )
)

client.entities.upsert(model)

Relationship to Data Jobs​

Through the MLModelLineageInfo fields, model groups can be connected to:

  • Training Jobs: Data jobs that produce models in the group (upstream lineage)
  • Downstream Jobs: Data jobs that consume models from the group (downstream lineage)

These relationships enable end-to-end lineage tracking from training data through model deployment.

Relationship to ML Features​

While ML Model Groups don't directly reference ML Features, the individual models within the group often consume ML Features. The model group serves as an organizational layer above the model-to-feature relationships.

Relationship to Containers​

Model groups can be organized within containers using the container aspect. This is useful for representing hierarchical structures like:

  • ML Workspace > Project > Model Group
  • Registry > Namespace > Model Group

Querying Model Groups​

You can query model groups and their associated models using both the REST API and the Python SDK.

Fetching Model Group Information via REST API​

REST API: Get model group by URN
curl 'http://localhost:8080/entities/urn%3Ali%3AmlModelGroup%3A(urn%3Ali%3AdataPlatform%3Amlflow,recommendation-models,PROD)' \
  -H 'Authorization: Bearer <token>'

This will return the model group entity with all its aspects, including:

  • mlModelGroupKey: The unique identifier
  • mlModelGroupProperties: Name, description, version, timestamps
  • ownership: Owners of the model group
  • globalTags: Tags attached to the group
  • glossaryTerms: Business terms associated with the group
  • domains: Domain assignment
  • institutionalMemory: Links and documentation
Python SDK: Read a model group
# Inlined from /metadata-ingestion/examples/library/mlmodel_group_read.py
from datahub.metadata.urns import MlModelGroupUrn
from datahub.sdk import DataHubClient

client = DataHubClient.from_env()

# Or get this from the UI (share -> copy urn) and use MlModelGroupUrn.from_string(...)
mlmodel_group_urn = MlModelGroupUrn(
    platform="mlflow", name="my-recommendations-model-group"
)

mlmodel_group_entity = client.entities.get(mlmodel_group_urn)
print("Model Group Name: ", mlmodel_group_entity.name)
print("Model Group Description: ", mlmodel_group_entity.description)
print("Model Group Custom Properties: ", mlmodel_group_entity.custom_properties)

Finding Models in a Model Group​

To find all models that belong to a specific model group, you can query the relationships.

REST API: Find all models in a model group
curl 'http://localhost:8080/relationships?direction=INCOMING&urn=urn%3Ali%3AmlModelGroup%3A(urn%3Ali%3AdataPlatform%3Amlflow,recommendation-models,PROD)&types=MemberOf' \
  -H 'Authorization: Bearer <token>'

This returns all ML Model entities that are members of the specified model group.

Integration with ML Platforms​

ML Model Groups work seamlessly with popular ML platforms:

MLflow​

MLflow's registered models naturally map to DataHub model groups, with model versions becoming individual MLModel entities within the group. The MLflow ingestion connector automatically creates these relationships.

SageMaker​

Amazon SageMaker model packages and model package groups can be represented as model groups in DataHub, providing a unified view across AWS environments.

Databricks​

Databricks ML models registered in Unity Catalog can be organized into model groups for better organization and governance.

Kubeflow​

Kubeflow model registries can leverage model groups to organize models by pipeline or serving configuration.

Custom Platforms​

Any custom ML platform can use model groups by specifying an appropriate platform identifier in the URN.

Notable Exceptions​

Model Group vs Individual Models​

It's important to understand when to use a model group versus tracking individual models:

  • Use Model Groups when: You have multiple related versions or variants of models that should be organized together. For example, all versions of a "customer churn prediction" model.
  • Use Individual Models when: You have standalone models that don't have multiple versions or aren't part of a logical family.

Lineage Inheritance​

Lineage information stored at the model group level (via trainingJobs and downstreamJobs) represents common lineage across all models in the group. Individual models can also have their own specific lineage information in their mlModelProperties aspect. The two levels of lineage are complementary:

  • Model group lineage: Shared training pipelines or common downstream consumers
  • Individual model lineage: Specific training runs or deployment-specific consumers

Naming Considerations​

When creating model groups, consider your naming strategy:

  • Names should be stable over time as they are part of the identifier
  • Avoid including version numbers in the group name (use individual model names for versioning)
  • Use clear, descriptive names that indicate the model family or purpose
  • Follow consistent naming conventions across your organization

Platform Instance Support​

Model groups support platform instances via the dataPlatformInstance aspect. This is useful when you have multiple instances of the same platform (e.g., multiple MLflow registries) and need to distinguish model groups across them.

Search and Discovery​

Model groups are searchable in DataHub by:

  • Name (with autocomplete support)
  • Description (full-text search)
  • Platform
  • Origin/Fabric
  • Tags and glossary terms
  • Domain
  • Owners

This makes it easy to discover relevant model groups through the DataHub UI or search API.

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​

mlModelGroupKey​

Key for an ML model group

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

mlModelGroupProperties​

Properties associated with an ML Model Group

FieldTypeRequiredDescriptionAnnotations
customPropertiesmap✓Custom property bag.Searchable
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
externalUrlstringURL where the reference existSearchable
namestringDisplay name of the MLModelGroupSearchable
descriptionstringDocumentation of the MLModelGroupSearchable
createdAtlongDate when the MLModelGroup was developed⚠️ Deprecated
createdTimeStampTime and Actor who created the MLModelGroup
lastModifiedTimeStampDate when the MLModelGroup was last modified
versionVersionTagVersion of the MLModelGroup

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...

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

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

editableMlModelGroupProperties​

Properties associated with an ML Model Group editable from the UI

FieldTypeRequiredDescriptionAnnotations
descriptionstringDocumentation of the ml model groupSearchable (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

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

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,...

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. ...

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.

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.

Relationships​

Outgoing​

These are the relationships stored in this entity's aspects

  • TrainedBy

    • DataJob via mlModelGroupProperties.trainingJobs
    • DataProcessInstance via mlModelGroupProperties.trainingJobs

  • UsedBy

    • DataJob via mlModelGroupProperties.downstreamJobs
    • DataProcessInstance via mlModelGroupProperties.downstreamJobs

  • OwnedBy

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

  • ownershipType

    • OwnershipType via ownership.owners.typeUrn

  • 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

  • IsPartOf

    • Container via container.container

Incoming​

These are the relationships stored in other entity's aspects

  • MemberOf

    • MlModel via mlModelProperties.groups

Global Metadata Model​

Global Graph