acryl-datahub 1.0.0.1rc1__py3-none-any.whl → 1.0.0.1rc2__py3-none-any.whl

datahub/ingestion/source/vertexai/vertexai_config.py

@@ -0,0 +1,29 @@
+from typing import Dict, Optional
+
+from pydantic import Field
+
+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"),
+    )
+
+    def get_credentials(self) -> Optional[Dict[str, str]]:
+        if self.credential:
+            return self.credential.to_dict(self.project_id)
+        return None

datahub/ingestion/source/vertexai/vertexai_result_type_utils.py

@@ -0,0 +1,68 @@
+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(state: PipelineState) -> Union[str, RunResultTypeClass]:
+    state_mapping = {
+        PipelineState.PIPELINE_STATE_SUCCEEDED: RunResultTypeClass.SUCCESS,
+        PipelineState.PIPELINE_STATE_FAILED: RunResultTypeClass.FAILURE,
+        PipelineState.PIPELINE_STATE_CANCELLED: "Cancelled",
+        PipelineState.PIPELINE_STATE_PAUSED: "Paused",
+        PipelineState.PIPELINE_STATE_QUEUED: "Queued",
+        PipelineState.PIPELINE_STATE_RUNNING: "Running",
+        PipelineState.PIPELINE_STATE_UNSPECIFIED: "Unspecific",
+    }
+
+    return state_mapping.get(state, "UNKNOWN")
+
+
+def get_custom_job_result_type(state: JobState) -> Union[str, RunResultTypeClass]:
+    state_mapping = {
+        JobState.JOB_STATE_SUCCEEDED: RunResultTypeClass.SUCCESS,
+        JobState.JOB_STATE_FAILED: RunResultTypeClass.FAILURE,
+        JobState.JOB_STATE_CANCELLED: "Cancelled",
+        JobState.JOB_STATE_PAUSED: "Paused",
+        JobState.JOB_STATE_QUEUED: "Queued",
+        JobState.JOB_STATE_RUNNING: "Running",
+        JobState.JOB_STATE_CANCELLING: "Cancelling",
+        JobState.JOB_STATE_EXPIRED: "Expired",
+        JobState.JOB_STATE_UPDATING: "Updating",
+    }
+    return state_mapping.get(state, "UNKNOWN")
+
+
+def get_job_result_status(job: VertexAiResourceNoun) -> Union[str, RunResultTypeClass]:
+    if isinstance(job, _TrainingJob) and isinstance(job.state, PipelineState):
+        return get_automl_job_result_type(job.state)
+    elif isinstance(job, _RunnableJob) and isinstance(job.state, JobState):
+        return get_custom_job_result_type(job.state)
+    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
+    """
+    status_mapping = {
+        0: "STATE_UNSPECIFIED",
+        1: "PENDING",
+        2: "RUNNING",
+        3: RunResultTypeClass.SUCCESS,
+        4: RunResultTypeClass.FAILURE,
+    }
+    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]