yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-12-04 22:52:54 +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
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.decorators import (
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.workunit import MetadataWorkUnit
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_result_type_utils import (
get_execution_result_status,
@ -97,6 +96,10 @@ class ModelMetadata:
endpoints: Optional[List[Endpoint]] = None
class ContainerKeyWithId(ContainerKey):
id: str
@platform_name("Vertex AI", id="vertexai")
@config_class(VertexAIConfig)
@support_status(SupportStatus.TESTING)