feat(ingestion) Adding vertexAI ingestion source (v2 - experiment and experiment run)

metadata-ingestion/docs/sources/vertexai/vertexai_recipe.yml

@@ -13,4 +13,4 @@ source:
 sink:
   type: "datahub-rest"
   config:
-    server: "http://localhost:8080"
+    server: "http://localhost:8080"

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

@@ -97,3 +97,13 @@ class BIAssetSubTypes(StrEnum):
 class MLAssetSubTypes(StrEnum):
     MLFLOW_TRAINING_RUN = "ML Training Run"
     MLFLOW_EXPERIMENT = "ML Experiment"
+    VERTEX_EXPERIMENT = "Experiment"
+    VERTEX_EXPERIMENT_RUN = "Experiment Run"
+    VERTEX_EXECUTION = "Execution"
+
+    VERTEX_MODEL = "ML Model"
+    VERTEX_MODEL_GROUP = "ML Model Group"
+    VERTEX_TRAINING_JOB = "Training Job"
+    VERTEX_ENDPOINT = "Endpoint"
+    VERTEX_DATASET = "Dataset"
+    VERTEX_PROJECT = "Project"

metadata-ingestion/src/datahub/ingestion/source/vertexai/__init__.py

@@ -0,0 +1,7 @@
+from datahub.ingestion.source.vertexai.config import VertexAIConfig
+from datahub.ingestion.source.vertexai.vertexai import (
+    ContainerKeyWithId,
+    ModelMetadata,
+    TrainingJobMetadata,
+    VertexAISource,
+)

metadata-ingestion/src/datahub/ingestion/source/vertexai/config.py

@@ -0,0 +1,33 @@
+from typing import Any, Optional
+
+from pydantic import Field, PrivateAttr
+
+from datahub.configuration.source_common import EnvConfigMixin
+from datahub.ingestion.source.common.gcp_credentials_config import GCPCredential
+
+
+class VertexAIConfig(EnvConfigMixin):
+    credential: Optional[GCPCredential] = Field(
+        default=None, description="GCP credential information"
+    )
+    project_id: str = Field(description=("Project ID in Google Cloud Platform"))
+    region: str = Field(
+        description=("Region of your project in Google Cloud Platform"),
+    )
+    bucket_uri: Optional[str] = Field(
+        default=None,
+        description=("Bucket URI used in your project"),
+    )
+    vertexai_url: Optional[str] = Field(
+        default="https://console.cloud.google.com/vertex-ai",
+        description=("VertexUI URI"),
+    )
+
+    _credentials_path: Optional[str] = PrivateAttr(None)
+
+    def __init__(self, **data: Any):
+        super().__init__(**data)
+        if self.credential:
+            self._credentials_path = self.credential.create_credential_temp_file(
+                project_id=self.project_id
+            )

metadata-ingestion/src/datahub/ingestion/source/vertexai/result_type_utils.py

@@ -0,0 +1,87 @@
+from typing import 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 datahub.metadata.schema_classes import RunResultTypeClass
+
+
+def get_automl_job_result_type(job: _TrainingJob) -> Union[str, RunResultTypeClass]:
+    state = job.state
+    if state == PipelineState.PIPELINE_STATE_SUCCEEDED:
+        return RunResultTypeClass.SUCCESS
+    elif state == PipelineState.PIPELINE_STATE_FAILED:
+        return RunResultTypeClass.FAILURE
+    elif state == PipelineState.PIPELINE_STATE_CANCELLED:
+        return "Cancelled"
+    elif state == PipelineState.PIPELINE_STATE_PAUSED:
+        return "Paused"
+    elif state == PipelineState.PIPELINE_STATE_QUEUED:
+        return "Queued"
+    elif state == PipelineState.PIPELINE_STATE_RUNNING:
+        return "Running"
+    elif state == PipelineState.PIPELINE_STATE_UNSPECIFIED:
+        return "Unspecific"
+    else:
+        return "UNKNOWN"
+
+
+def get_custom_job_result_type(job: _RunnableJob) -> Union[str, RunResultTypeClass]:
+    state = job.state
+    if state == JobState.JOB_STATE_SUCCEEDED:
+        return RunResultTypeClass.SUCCESS
+    elif state == JobState.JOB_STATE_FAILED:
+        return RunResultTypeClass.FAILURE
+    elif state == JobState.JOB_STATE_CANCELLED:
+        return "Cancelled"
+    elif state == JobState.JOB_STATE_PAUSED:
+        return "Paused"
+    elif state == JobState.JOB_STATE_QUEUED:
+        return "Queued"
+    elif state == JobState.JOB_STATE_RUNNING:
+        return "Running"
+    elif state == JobState.JOB_STATE_CANCELLING:
+        return "Cancelling"
+    elif state == JobState.JOB_STATE_EXPIRED:
+        return "Expired"
+    elif state == JobState.JOB_STATE_UPDATING:
+        return "Updating"
+    else:
+        return "UNKNOWN"
+
+
+def get_job_result_status(job: VertexAiResourceNoun) -> Union[str, RunResultTypeClass]:
+    if isinstance(job, _TrainingJob):
+        return get_automl_job_result_type(job)
+    elif isinstance(job, _RunnableJob):
+        return get_custom_job_result_type(job)
+    return "UNKNOWN"
+
+
+def get_execution_result_status(status: int) -> Union[str, RunResultTypeClass]:
+    """
+    State of the execution.
+    STATE_UNSPECIFIED = 0
+    PENDING = 1
+    RUNNING = 2
+    SUCCEEDED = 3
+    FAILED = 4
+    """
+    if status == 3:
+        return RunResultTypeClass.SUCCESS
+    elif status == 4:
+        return RunResultTypeClass.FAILURE
+    elif status == 2:
+        return "RUNNING"
+    elif status == 1:
+        return "PENDING"
+    elif status == 0:
+        return "STATE_UNSPECIFIED"
+    else:
+        return "UNKNOWN"
+
+
+def is_status_for_run_event_class(status: Union[str, RunResultTypeClass]) -> bool:
+    return status in [RunResultTypeClass.SUCCESS, RunResultTypeClass.FAILURE]