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

18 lines
496 B
Python

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}")