acryl-datahub 1.0.0.3rc2__py3-none-any.whl → 1.0.0.3rc5__py3-none-any.whl

datahub/ingestion/source/vertexai/vertexai.py

@@ -1,5 +1,6 @@
 import dataclasses
 import logging
+from datetime import datetime, timedelta
 from typing import Dict, Iterable, List, Optional, Tuple, TypeVar, Union
 
 from google.api_core.exceptions import GoogleAPICallError
@@ -12,15 +13,22 @@ from google.cloud.aiplatform import (
     AutoMLVideoTrainingJob,
     Endpoint,
     ExperimentRun,
+    PipelineJob,
 )
 from google.cloud.aiplatform.base import VertexAiResourceNoun
 from google.cloud.aiplatform.metadata.execution import Execution
 from google.cloud.aiplatform.metadata.experiment_resources import Experiment
 from google.cloud.aiplatform.models import Model, VersionInfo
 from google.cloud.aiplatform.training_jobs import _TrainingJob
+from google.cloud.aiplatform_v1.types import (
+    PipelineJob as PipelineJobType,
+    PipelineTaskDetail,
+)
 from google.oauth2 import service_account
+from google.protobuf import timestamp_pb2
 
 import datahub.emitter.mce_builder as builder
+from datahub.api.entities.datajob import DataFlow, DataJob
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import (
     ExperimentKey,
@@ -43,6 +51,7 @@ from datahub.ingestion.source.vertexai.vertexai_config import VertexAIConfig
 from datahub.ingestion.source.vertexai.vertexai_result_type_utils import (
     get_execution_result_status,
     get_job_result_status,
+    get_pipeline_task_result_status,
     is_status_for_run_event_class,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.dataprocess import (
@@ -76,7 +85,13 @@ from datahub.metadata.schema_classes import (
     VersionPropertiesClass,
     VersionTagClass,
 )
-from datahub.metadata.urns import DataPlatformUrn, MlModelUrn, VersionSetUrn
+from datahub.metadata.urns import (
+    DataFlowUrn,
+    DataJobUrn,
+    DataPlatformUrn,
+    MlModelUrn,
+    VersionSetUrn,
+)
 from datahub.utilities.time import datetime_to_ts_millis
 
 T = TypeVar("T")
@@ -100,6 +115,34 @@ class ModelMetadata:
     endpoints: Optional[List[Endpoint]] = None
 
 
+@dataclasses.dataclass
+class PipelineTaskMetadata:
+    name: str
+    urn: DataJobUrn
+    id: Optional[int] = None
+    type: Optional[str] = None
+    state: Optional[PipelineTaskDetail.State] = None
+    start_time: Optional[timestamp_pb2.Timestamp] = None
+    create_time: Optional[timestamp_pb2.Timestamp] = None
+    end_time: Optional[timestamp_pb2.Timestamp] = None
+    upstreams: Optional[List[DataJobUrn]] = None
+    duration: Optional[int] = None
+
+
+@dataclasses.dataclass
+class PipelineMetadata:
+    name: str
+    resource_name: str
+    tasks: List[PipelineTaskMetadata]
+    urn: DataFlowUrn
+    id: Optional[str] = None
+    labels: Optional[Dict[str, str]] = None
+    create_time: Optional[datetime] = None
+    update_time: Optional[datetime] = None
+    duration: Optional[timedelta] = None
+    region: Optional[str] = None
+
+
 @platform_name("Vertex AI", id="vertexai")
 @config_class(VertexAIConfig)
 @support_status(SupportStatus.TESTING)
@@ -150,6 +193,255 @@ class VertexAISource(Source):
         yield from self._get_experiments_workunits()
         # Fetch and Ingest Experiment Runs
         yield from auto_workunit(self._get_experiment_runs_mcps())
+        # Fetch Pipelines and Tasks
+        yield from auto_workunit(self._get_pipelines_mcps())
+
+    def _get_pipelines_mcps(self) -> Iterable[MetadataChangeProposalWrapper]:
+        """
+        Fetches pipelines from Vertex AI and generates corresponding mcps.
+        """
+
+        pipeline_jobs = self.client.PipelineJob.list()
+
+        for pipeline in pipeline_jobs:
+            logger.info(f"fetching pipeline ({pipeline.name})")
+            pipeline_meta = self._get_pipeline_metadata(pipeline)
+            yield from self._get_pipeline_mcps(pipeline_meta)
+            yield from self._gen_pipeline_task_mcps(pipeline_meta)
+
+    def _get_pipeline_tasks_metadata(
+        self, pipeline: PipelineJob, pipeline_urn: DataFlowUrn
+    ) -> List[PipelineTaskMetadata]:
+        tasks: List[PipelineTaskMetadata] = list()
+        task_map: Dict[str, PipelineTaskDetail] = dict()
+        for task in pipeline.task_details:
+            task_map[task.task_name] = task
+
+        resource = pipeline.gca_resource
+        if isinstance(resource, PipelineJobType):
+            for task_name in resource.pipeline_spec["root"]["dag"]["tasks"]:
+                logger.debug(
+                    f"fetching pipeline task ({task_name}) in pipeline ({pipeline.name})"
+                )
+                task_urn = DataJobUrn.create_from_ids(
+                    data_flow_urn=str(pipeline_urn),
+                    job_id=self._make_vertexai_pipeline_task_id(task_name),
+                )
+                task_meta = PipelineTaskMetadata(name=task_name, urn=task_urn)
+                if (
+                    "dependentTasks"
+                    in resource.pipeline_spec["root"]["dag"]["tasks"][task_name]
+                ):
+                    upstream_tasks = resource.pipeline_spec["root"]["dag"]["tasks"][
+                        task_name
+                    ]["dependentTasks"]
+                    upstream_urls = [
+                        DataJobUrn.create_from_ids(
+                            data_flow_urn=str(pipeline_urn),
+                            job_id=self._make_vertexai_pipeline_task_id(upstream_task),
+                        )
+                        for upstream_task in upstream_tasks
+                    ]
+                    task_meta.upstreams = upstream_urls
+
+                task_detail = task_map.get(task_name)
+                if task_detail:
+                    task_meta.id = task_detail.task_id
+                    task_meta.state = task_detail.state
+                    task_meta.start_time = task_detail.start_time
+                    task_meta.create_time = task_detail.create_time
+                    if task_detail.end_time:
+                        task_meta.end_time = task_detail.end_time
+                        task_meta.duration = int(
+                            (
+                                task_meta.end_time.timestamp()
+                                - task_meta.start_time.timestamp()
+                            )
+                            * 1000
+                        )
+
+                tasks.append(task_meta)
+        return tasks
+
+    def _get_pipeline_metadata(self, pipeline: PipelineJob) -> PipelineMetadata:
+        dataflow_urn = DataFlowUrn.create_from_ids(
+            orchestrator=self.platform,
+            env=self.config.env,
+            flow_id=self._make_vertexai_pipeline_id(pipeline.name),
+            platform_instance=self.platform,
+        )
+        tasks = self._get_pipeline_tasks_metadata(
+            pipeline=pipeline, pipeline_urn=dataflow_urn
+        )
+
+        pipeline_meta = PipelineMetadata(
+            name=pipeline.name,
+            resource_name=pipeline.resource_name,
+            urn=dataflow_urn,
+            tasks=tasks,
+        )
+        pipeline_meta.resource_name = pipeline.resource_name
+        pipeline_meta.labels = pipeline.labels
+        pipeline_meta.create_time = pipeline.create_time
+        pipeline_meta.region = pipeline.location
+        if pipeline.update_time:
+            pipeline_meta.update_time = pipeline.update_time
+            pipeline_meta.duration = timedelta(
+                milliseconds=datetime_to_ts_millis(pipeline.update_time)
+                - datetime_to_ts_millis(pipeline.create_time)
+            )
+        return pipeline_meta
+
+    def _gen_pipeline_task_run_mcps(
+        self, task: PipelineTaskMetadata, datajob: DataJob, pipeline: PipelineMetadata
+    ) -> (Iterable)[MetadataChangeProposalWrapper]:
+        dpi_urn = builder.make_data_process_instance_urn(
+            self._make_vertexai_pipeline_task_run_id(entity_id=task.name)
+        )
+        result_status: Union[str, RunResultTypeClass] = get_pipeline_task_result_status(
+            task.state
+        )
+
+        yield from MetadataChangeProposalWrapper.construct_many(
+            dpi_urn,
+            aspects=[
+                DataProcessInstancePropertiesClass(
+                    name=task.name,
+                    created=AuditStampClass(
+                        time=(
+                            int(task.create_time.timestamp() * 1000)
+                            if task.create_time
+                            else 0
+                        ),
+                        actor="urn:li:corpuser:datahub",
+                    ),
+                    externalUrl=self._make_pipeline_external_url(pipeline.name),
+                    customProperties={},
+                ),
+                SubTypesClass(typeNames=[MLAssetSubTypes.VERTEX_PIPELINE_TASK_RUN]),
+                ContainerClass(container=self._get_project_container().as_urn()),
+                DataPlatformInstanceClass(platform=str(DataPlatformUrn(self.platform))),
+                DataProcessInstanceRelationships(
+                    upstreamInstances=[], parentTemplate=str(datajob.urn)
+                ),
+                (
+                    DataProcessInstanceRunEventClass(
+                        status=DataProcessRunStatusClass.COMPLETE,
+                        timestampMillis=(
+                            int(task.create_time.timestamp() * 1000)
+                            if task.create_time
+                            else 0
+                        ),
+                        result=DataProcessInstanceRunResultClass(
+                            type=result_status,
+                            nativeResultType=self.platform,
+                        ),
+                        durationMillis=task.duration,
+                    )
+                    if is_status_for_run_event_class(result_status) and task.duration
+                    else None
+                ),
+            ],
+        )
+
+    def _gen_pipeline_task_mcps(
+        self, pipeline: PipelineMetadata
+    ) -> Iterable[MetadataChangeProposalWrapper]:
+        dataflow_urn = pipeline.urn
+
+        for task in pipeline.tasks:
+            datajob = DataJob(
+                id=self._make_vertexai_pipeline_task_id(task.name),
+                flow_urn=dataflow_urn,
+                name=task.name,
+                properties={},
+                owners={"urn:li:corpuser:datahub"},
+                upstream_urns=task.upstreams if task.upstreams else [],
+                url=self._make_pipeline_external_url(pipeline.name),
+            )
+            yield from MetadataChangeProposalWrapper.construct_many(
+                entityUrn=str(datajob.urn),
+                aspects=[
+                    ContainerClass(container=self._get_project_container().as_urn()),
+                    SubTypesClass(typeNames=[MLAssetSubTypes.VERTEX_PIPELINE_TASK]),
+                ],
+            )
+            yield from datajob.generate_mcp()
+            yield from self._gen_pipeline_task_run_mcps(task, datajob, pipeline)
+
+    def _format_pipeline_duration(self, td: timedelta) -> str:
+        days = td.days
+        hours, remainder = divmod(td.seconds, 3600)
+        minutes, seconds = divmod(remainder, 60)
+        milliseconds = td.microseconds // 1000
+
+        parts = []
+        if days:
+            parts.append(f"{days}d")
+        if hours:
+            parts.append(f"{hours}h")
+        if minutes:
+            parts.append(f"{minutes}m")
+        if seconds:
+            parts.append(f"{seconds}s")
+        if milliseconds:
+            parts.append(f"{milliseconds}ms")
+        return " ".join(parts) if parts else "0s"
+
+    def _get_pipeline_task_properties(
+        self, task: PipelineTaskMetadata
+    ) -> Dict[str, str]:
+        return {
+            "created_time": (
+                task.create_time.strftime("%Y-%m-%d %H:%M:%S")
+                if task.create_time
+                else ""
+            )
+        }
+
+    def _get_pipeline_properties(self, pipeline: PipelineMetadata) -> Dict[str, str]:
+        return {
+            "resource_name": pipeline.resource_name if pipeline.resource_name else "",
+            "create_time": (
+                pipeline.create_time.isoformat() if pipeline.create_time else ""
+            ),
+            "update_time": (
+                pipeline.update_time.isoformat() if pipeline.update_time else ""
+            ),
+            "duration": (
+                self._format_pipeline_duration(pipeline.duration)
+                if pipeline.duration
+                else ""
+            ),
+            "location": (pipeline.region if pipeline.region else ""),
+            "labels": ",".join([f"{k}:{v}" for k, v in pipeline.labels.items()])
+            if pipeline.labels
+            else "",
+        }
+
+    def _get_pipeline_mcps(
+        self, pipeline: PipelineMetadata
+    ) -> Iterable[MetadataChangeProposalWrapper]:
+        dataflow = DataFlow(
+            orchestrator=self.platform,
+            id=self._make_vertexai_pipeline_id(pipeline.name),
+            env=self.config.env,
+            name=pipeline.name,
+            platform_instance=self.platform,
+            properties=self._get_pipeline_properties(pipeline),
+            owners={"urn:li:corpuser:datahub"},
+            url=self._make_pipeline_external_url(pipeline_name=pipeline.name),
+        )
+
+        yield from dataflow.generate_mcp()
+
+        yield from MetadataChangeProposalWrapper.construct_many(
+            entityUrn=str(dataflow.urn),
+            aspects=[
+                ContainerClass(container=self._get_project_container().as_urn()),
+                SubTypesClass(typeNames=[MLAssetSubTypes.VERTEX_PIPELINE]),
+            ],
+        )
 
     def _get_experiments_workunits(self) -> Iterable[MetadataWorkUnit]:
         # List all experiments
@@ -175,7 +467,7 @@ class VertexAISource(Source):
             parent_container_key=self._get_project_container(),
             container_key=ExperimentKey(
                 platform=self.platform,
-                id=self._make_vertexai_experiment_name(experiment.name),
+                id=self._make_vertexai_experiment_id(experiment.name),
             ),
             name=experiment.name,
             sub_types=[MLAssetSubTypes.VERTEX_EXPERIMENT],
@@ -311,7 +603,7 @@ class VertexAISource(Source):
     ) -> Iterable[MetadataChangeProposalWrapper]:
         experiment_key = ExperimentKey(
             platform=self.platform,
-            id=self._make_vertexai_experiment_name(experiment.name),
+            id=self._make_vertexai_experiment_id(experiment.name),
         )
         run_urn = self._make_experiment_run_urn(experiment, run)
         created_time, duration = self._get_run_timestamps(run)
@@ -968,7 +1260,7 @@ class VertexAISource(Source):
     ) -> str:
         return f"{self.config.project_id}.job.{entity_id}"
 
-    def _make_vertexai_experiment_name(self, entity_id: Optional[str]) -> str:
+    def _make_vertexai_experiment_id(self, entity_id: Optional[str]) -> str:
         return f"{self.config.project_id}.experiment.{entity_id}"
 
     def _make_vertexai_experiment_run_name(self, entity_id: Optional[str]) -> str:
@@ -977,6 +1269,15 @@ class VertexAISource(Source):
     def _make_vertexai_run_execution_name(self, entity_id: Optional[str]) -> str:
         return f"{self.config.project_id}.execution.{entity_id}"
 
+    def _make_vertexai_pipeline_id(self, entity_id: Optional[str]) -> str:
+        return f"{self.config.project_id}.pipeline.{entity_id}"
+
+    def _make_vertexai_pipeline_task_id(self, entity_id: Optional[str]) -> str:
+        return f"{self.config.project_id}.pipeline_task.{entity_id}"
+
+    def _make_vertexai_pipeline_task_run_id(self, entity_id: Optional[str]) -> str:
+        return f"{self.config.project_id}.pipeline_task_run.{entity_id}"
+
     def _make_artifact_external_url(
         self, experiment: Experiment, run: ExperimentRun
     ) -> str:
@@ -1053,3 +1354,14 @@ class VertexAISource(Source):
             f"/runs/{experiment.name}-{run.name}/charts?project={self.config.project_id}"
         )
         return external_url
+
+    def _make_pipeline_external_url(self, pipeline_name: str) -> str:
+        """
+        Pipeline Run external URL in Vertex AI
+        https://console.cloud.google.com/vertex-ai/pipelines/locations/us-west2/runs/pipeline-example-more-tasks-3-20250320210739?project=acryl-poc
+        """
+        external_url: str = (
+            f"{self.config.vertexai_url}/pipelines/locations/{self.config.region}/runs/{pipeline_name}"
+            f"?project={self.config.project_id}"
+        )
+        return external_url

datahub/ingestion/source/vertexai/vertexai_result_type_utils.py

@@ -1,9 +1,9 @@
-from typing import Union
+from typing import Optional, Union
 
 from google.cloud.aiplatform.base import VertexAiResourceNoun
 from google.cloud.aiplatform.jobs import _RunnableJob
 from google.cloud.aiplatform.training_jobs import _TrainingJob
-from google.cloud.aiplatform_v1.types import JobState, PipelineState
+from google.cloud.aiplatform_v1.types import JobState, PipelineState, PipelineTaskDetail
 
 from datahub.metadata.schema_classes import RunResultTypeClass
 
@@ -64,5 +64,26 @@ def get_execution_result_status(status: int) -> Union[str, RunResultTypeClass]:
     return status_mapping.get(status, "UNKNOWN")
 
 
+def get_pipeline_task_result_status(
+    status: Optional[PipelineTaskDetail.State],
+) -> Union[str, RunResultTypeClass]:
+    # TODO: DataProcessInstanceRunResultClass fails with status string except for SUCCESS, FAILURE, SKIPPED,
+    #  which will be fixed in the future
+    status_mapping = {
+        # PipelineTaskDetail.State.STATE_UNSPECIFIED: "STATE_UNSPECIFIED",
+        # PipelineTaskDetail.State.PENDING: "PENDING",
+        # PipelineTaskDetail.State.RUNNING: "RUNNING",
+        # PipelineTaskDetail.State.CANCEL_PENDING: "CANCEL_PENDING",
+        # PipelineTaskDetail.State.CANCELLING: "CANCELLING",
+        # PipelineTaskDetail.State.NOT_TRIGGERED: "NOT_TRIGGERED",
+        PipelineTaskDetail.State.SUCCEEDED: RunResultTypeClass.SUCCESS,
+        PipelineTaskDetail.State.FAILED: RunResultTypeClass.FAILURE,
+        PipelineTaskDetail.State.SKIPPED: RunResultTypeClass.SKIPPED,
+    }
+    if status is None:
+        return "UNKNOWN"
+    return status_mapping.get(status, "UNKNOWN")
+
+
 def is_status_for_run_event_class(status: Union[str, RunResultTypeClass]) -> bool:
     return status in [RunResultTypeClass.SUCCESS, RunResultTypeClass.FAILURE]

datahub/metadata/_schema_classes.py

@@ -15247,7 +15247,7 @@ class DataContractKeyClass(_Aspect):
 
 
     ASPECT_NAME = 'dataContractKey'
-    ASPECT_INFO = {'keyForEntity': 'dataContract', 'entityCategory': 'core', 'entityAspects': ['dataContractProperties', 'dataContractStatus', 'status']}
+    ASPECT_INFO = {'keyForEntity': 'dataContract', 'entityCategory': 'core', 'entityAspects': ['dataContractProperties', 'dataContractStatus', 'status', 'structuredProperties']}
     RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.metadata.key.DataContractKey")
 
     def __init__(self,

datahub/metadata/schema.avsc

@@ -16024,7 +16024,8 @@
       "entityAspects": [
         "dataContractProperties",
         "dataContractStatus",
-        "status"
+        "status",
+        "structuredProperties"
       ]
     },
     "name": "DataContractKey",

datahub/metadata/schemas/DataContractKey.avsc

@@ -7,7 +7,8 @@
     "entityAspects": [
       "dataContractProperties",
       "dataContractStatus",
-      "status"
+      "status",
+      "structuredProperties"
     ]
   },
   "name": "DataContractKey",

datahub/sdk/_all_entities.py

@@ -3,11 +3,15 @@ from typing import Dict, List, Type
 from datahub.sdk.container import Container
 from datahub.sdk.dataset import Dataset
 from datahub.sdk.entity import Entity
+from datahub.sdk.mlmodel import MLModel
+from datahub.sdk.mlmodelgroup import MLModelGroup
 
 # TODO: Is there a better way to declare this?
 ENTITY_CLASSES_LIST: List[Type[Entity]] = [
     Container,
     Dataset,
+    MLModel,
+    MLModelGroup,
 ]
 
 ENTITY_CLASSES: Dict[str, Type[Entity]] = {

datahub/sdk/_shared.py

@@ -5,6 +5,7 @@ from datetime import datetime
 from typing import (
     TYPE_CHECKING,
     Callable,
+    Dict,
     List,
     Optional,
     Sequence,
@@ -14,6 +15,7 @@ from typing import (
 
 from typing_extensions import TypeAlias, assert_never
 
+import datahub.emitter.mce_builder as builder
 import datahub.metadata.schema_classes as models
 from datahub.emitter.mce_builder import (
     make_ts_millis,
@@ -30,12 +32,14 @@ from datahub.metadata.urns import (
     DataJobUrn,
     DataPlatformInstanceUrn,
     DataPlatformUrn,
+    DataProcessInstanceUrn,
     DatasetUrn,
     DomainUrn,
     GlossaryTermUrn,
     OwnershipTypeUrn,
     TagUrn,
     Urn,
+    VersionSetUrn,
 )
 from datahub.sdk._utils import add_list_unique, remove_list_unique
 from datahub.sdk.entity import Entity
@@ -52,6 +56,36 @@ ActorUrn: TypeAlias = Union[CorpUserUrn, CorpGroupUrn]
 
 _DEFAULT_ACTOR_URN = CorpUserUrn("__ingestion").urn()
 
+TrainingMetricsInputType: TypeAlias = Union[
+    List[models.MLMetricClass], Dict[str, Optional[str]]
+]
+HyperParamsInputType: TypeAlias = Union[
+    List[models.MLHyperParamClass], Dict[str, Optional[str]]
+]
+MLTrainingJobInputType: TypeAlias = Union[Sequence[Union[str, DataProcessInstanceUrn]]]
+
+
+def convert_training_metrics(
+    metrics: TrainingMetricsInputType,
+) -> List[models.MLMetricClass]:
+    if isinstance(metrics, dict):
+        return [
+            models.MLMetricClass(name=name, value=str(value))
+            for name, value in metrics.items()
+        ]
+    return metrics
+
+
+def convert_hyper_params(
+    params: HyperParamsInputType,
+) -> List[models.MLHyperParamClass]:
+    if isinstance(params, dict):
+        return [
+            models.MLHyperParamClass(name=name, value=str(value))
+            for name, value in params.items()
+        ]
+    return params
+
 
 def make_time_stamp(ts: Optional[datetime]) -> Optional[models.TimeStampClass]:
     if ts is None:
@@ -578,3 +612,109 @@ class HasInstitutionalMemory(Entity):
             self._link_key,
             self._parse_link_association_class(link),
         )
+
+
+class HasVersion(Entity):
+    """Mixin for entities that have version properties."""
+
+    def _get_version_props(self) -> Optional[models.VersionPropertiesClass]:
+        return self._get_aspect(models.VersionPropertiesClass)
+
+    def _ensure_version_props(self) -> models.VersionPropertiesClass:
+        version_props = self._get_version_props()
+        if version_props is None:
+            guid_dict = {"urn": str(self.urn)}
+            version_set_urn = VersionSetUrn(
+                id=builder.datahub_guid(guid_dict), entity_type=self.urn.ENTITY_TYPE
+            )
+
+            version_props = models.VersionPropertiesClass(
+                versionSet=str(version_set_urn),
+                version=models.VersionTagClass(versionTag="0.1.0"),
+                sortId="0000000.1.0",
+            )
+            self._set_aspect(version_props)
+        return version_props
+
+    @property
+    def version(self) -> Optional[str]:
+        version_props = self._get_version_props()
+        if version_props and version_props.version:
+            return version_props.version.versionTag
+        return None
+
+    def set_version(self, version: str) -> None:
+        """Set the version of the entity."""
+        guid_dict = {"urn": str(self.urn)}
+        version_set_urn = VersionSetUrn(
+            id=builder.datahub_guid(guid_dict), entity_type=self.urn.ENTITY_TYPE
+        )
+
+        version_props = self._get_version_props()
+        if version_props is None:
+            # If no version properties exist, create a new one
+            version_props = models.VersionPropertiesClass(
+                version=models.VersionTagClass(versionTag=version),
+                versionSet=str(version_set_urn),
+                sortId=version.zfill(10),  # Pad with zeros for sorting
+            )
+        else:
+            # Update existing version properties
+            version_props.version = models.VersionTagClass(versionTag=version)
+            version_props.versionSet = str(version_set_urn)
+            version_props.sortId = version.zfill(10)
+
+        self._set_aspect(version_props)
+
+    @property
+    def version_aliases(self) -> List[str]:
+        version_props = self._get_version_props()
+        if version_props and version_props.aliases:
+            return [
+                alias.versionTag
+                for alias in version_props.aliases
+                if alias.versionTag is not None
+            ]
+        return []  # Return empty list instead of None
+
+    def set_version_aliases(self, aliases: List[str]) -> None:
+        version_props = self._get_aspect(models.VersionPropertiesClass)
+        if version_props:
+            version_props.aliases = [
+                models.VersionTagClass(versionTag=alias) for alias in aliases
+            ]
+        else:
+            # If no version properties exist, we need to create one with a default version
+            guid_dict = {"urn": str(self.urn)}
+            version_set_urn = VersionSetUrn(
+                id=builder.datahub_guid(guid_dict), entity_type=self.urn.ENTITY_TYPE
+            )
+            self._set_aspect(
+                models.VersionPropertiesClass(
+                    version=models.VersionTagClass(
+                        versionTag="0.1.0"
+                    ),  # Default version
+                    versionSet=str(version_set_urn),
+                    sortId="0000000.1.0",
+                    aliases=[
+                        models.VersionTagClass(versionTag=alias) for alias in aliases
+                    ],
+                )
+            )
+
+    def add_version_alias(self, alias: str) -> None:
+        if not alias:
+            raise ValueError("Alias cannot be empty")
+        version_props = self._ensure_version_props()
+        if version_props.aliases is None:
+            version_props.aliases = []
+        version_props.aliases.append(models.VersionTagClass(versionTag=alias))
+        self._set_aspect(version_props)
+
+    def remove_version_alias(self, alias: str) -> None:
+        version_props = self._get_version_props()
+        if version_props and version_props.aliases:
+            version_props.aliases = [
+                a for a in version_props.aliases if a.versionTag != alias
+            ]
+            self._set_aspect(version_props)

datahub/sdk/entity_client.py

@@ -11,6 +11,8 @@ from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.urns import (
     ContainerUrn,
     DatasetUrn,
+    MlModelGroupUrn,
+    MlModelUrn,
     Urn,
 )
 from datahub.sdk._all_entities import ENTITY_CLASSES
@@ -18,6 +20,8 @@ from datahub.sdk._shared import UrnOrStr
 from datahub.sdk.container import Container
 from datahub.sdk.dataset import Dataset
 from datahub.sdk.entity import Entity
+from datahub.sdk.mlmodel import MLModel
+from datahub.sdk.mlmodelgroup import MLModelGroup
 
 if TYPE_CHECKING:
     from datahub.sdk.main_client import DataHubClient
@@ -49,6 +53,10 @@ class EntityClient:
     @overload
     def get(self, urn: DatasetUrn) -> Dataset: ...
     @overload
+    def get(self, urn: MlModelUrn) -> MLModel: ...
+    @overload
+    def get(self, urn: MlModelGroupUrn) -> MLModelGroup: ...
+    @overload
     def get(self, urn: Union[Urn, str]) -> Entity: ...
     def get(self, urn: UrnOrStr) -> Entity:
         """Retrieve an entity by its urn.