include value in model tag ingestion

datahub-project · Mar 5, 2025 · ddc8161 · ddc8161
1 parent eb88a71
commit ddc8161
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/mlflow.py b/metadata-ingestion/src/datahub/ingestion/source/mlflow.py
@@ -584,6 +584,7 @@ def _get_ml_model_properties_workunit(
             if model_version.user_id
             else None
         )
+        model_version_tags = [f"{k}:{v}" for k, v in model_version.tags.items()]
 
         ml_model_properties = MLModelPropertiesClass(
             customProperties=model_version.tags,
@@ -599,8 +600,7 @@ def _get_ml_model_properties_workunit(
             ),
             hyperParams=hyperparams,
             trainingMetrics=training_metrics,
-            # mlflow tags are dicts, but datahub tags are lists. currently use only keys from mlflow tags
-            tags=list(model_version.tags.keys()),
+            tags=model_version_tags,
             groups=[ml_model_group_urn],
             trainingJobs=training_jobs,
         )

diff --git a/metadata-ingestion/tests/integration/mlflow/mlflow_mcps_golden.json b/metadata-ingestion/tests/integration/mlflow/mlflow_mcps_golden.json
@@ -362,7 +362,7 @@
                 }
             ],
             "tags": [
-                "model_version_id"
+                "model_version_id:1"
             ],
             "groups": [
                 "urn:li:mlModelGroup:(urn:li:dataPlatform:mlflow,test-model,PROD)"