generalize MLAssetSubtypes

Author: ryota-cloud

metadata-ingestion/src/datahub/ingestion/source/common/subtypes.py

@@ -95,5 +95,11 @@ class BIAssetSubTypes(StrEnum):
 
 
 class MLAssetSubTypes(StrEnum):
-    MLFLOW_TRAINING_RUN = "ML Training Run"
-    MLFLOW_EXPERIMENT = "ML Experiment"
+    TRAINING_RUN = "ML Training Run"
+    EXPERIMENT = "ML Experiment"
+    MODEL = "ML Model"
+    TRAINING_JOB = "Training Job"
+    MODEL_GROUP = "ML Model Group"
+    ENDPOINT = "Endpoint"
+    DATASET = "Dataset"
+    PROJECT = "Project"

metadata-ingestion/src/datahub/ingestion/source/mlflow.py

@@ -228,7 +228,7 @@ def _get_experiment_container_workunit(
                 platform=str(DataPlatformUrn(platform_name=self.platform)),
                 id=experiment.name,
             ),
-            subtype=MLAssetSubTypes.MLFLOW_EXPERIMENT,
+            subtype=MLAssetSubTypes.EXPERIMENT,
             display_name=experiment.name,
             description=experiment.tags.get("mlflow.note.content"),
             extra_properties=self._get_experiment_custom_properties(experiment),
@@ -358,7 +358,7 @@ def _get_run_workunits(
 
         yield MetadataChangeProposalWrapper(
             entityUrn=str(data_process_instance.urn),
-            aspect=SubTypesClass(typeNames=[MLAssetSubTypes.MLFLOW_TRAINING_RUN]),
+            aspect=SubTypesClass(typeNames=[MLAssetSubTypes.TRAINING_RUN]),
         ).as_workunit()
 
     def _get_mlflow_registered_models(self) -> Iterable[RegisteredModel]: