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

27 lines
834 B
Python

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)