yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-12-06 07:34:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(metadata-ingestion) fix connector test error for Vertex AI (#12963)

Browse Source
This commit is contained in:
ryota-cloud 2025-03-24 13:17:57 -07:00 committed by GitHub
parent 2cc8856c6b
commit 11353d1172
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
metadata-ingestion/src/datahub/ingestion/source/vertexai/vertexai.py
View File

@ -22,7 +22,7 @@ from google.oauth2 import service_account
import datahub.emitter.mce_builder as builder import datahub.emitter.mce_builder as builder
from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.emitter.mcp_builder import ProjectIdKey, gen_containers from datahub.emitter.mcp_builder import ContainerKey, ProjectIdKey, gen_containers
from datahub.ingestion.api.common import PipelineContext from datahub.ingestion.api.common import PipelineContext
from datahub.ingestion.api.decorators import ( from datahub.ingestion.api.decorators import (
SupportStatus, SupportStatus,
@ -35,7 +35,6 @@ from datahub.ingestion.api.source import Source, SourceCapability, SourceReport
from datahub.ingestion.api.source_helpers import auto_workunit from datahub.ingestion.api.source_helpers import auto_workunit
from datahub.ingestion.api.workunit import MetadataWorkUnit from datahub.ingestion.api.workunit import MetadataWorkUnit
from datahub.ingestion.source.common.subtypes import MLAssetSubTypes from datahub.ingestion.source.common.subtypes import MLAssetSubTypes
from datahub.ingestion.source.mlflow import ContainerKeyWithId
from datahub.ingestion.source.vertexai.vertexai_config import VertexAIConfig from datahub.ingestion.source.vertexai.vertexai_config import VertexAIConfig
from datahub.ingestion.source.vertexai.vertexai_result_type_utils import ( from datahub.ingestion.source.vertexai.vertexai_result_type_utils import (
get_execution_result_status, get_execution_result_status,
@ -97,6 +96,10 @@ class ModelMetadata:
endpoints: Optional[List[Endpoint]] = None endpoints: Optional[List[Endpoint]] = None
class ContainerKeyWithId(ContainerKey):
id: str
@platform_name("Vertex AI", id="vertexai") @platform_name("Vertex AI", id="vertexai")
@config_class(VertexAIConfig) @config_class(VertexAIConfig)
@support_status(SupportStatus.TESTING) @support_status(SupportStatus.TESTING)