datahub/metadata-ingestion/examples/library/read_mlmodel_group.py

from datahub.metadata.urns import MlModelGroupUrn
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",
    description="A group for recommendations models",
    platform="mlflow",
    custom_properties={
        "owner": "John Doe",
        "team": "recommendations",
        "domain": "marketing",
    },
)

client.entities.upsert(mlmodel_group)
mlmodel_group = client.entities.get(
    MlModelGroupUrn(platform="mlflow", name="my-recommendations-model-group")
)

print("Model Group Name: ", mlmodel_group.name)
print("Model Group Description: ", mlmodel_group.description)
print("Model Group Custom Properties: ", mlmodel_group.custom_properties)