Skip to content

Commit 3d98b1f

Browse files
committed
Merge branch 'vertex_src_temp' into vertex_exp_3
2 parents 6b3a761 + dedab91 commit 3d98b1f

File tree

1 file changed

+5
-5
lines changed
  • metadata-ingestion/src/datahub/ingestion/source

1 file changed

+5
-5
lines changed

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

+5-5
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import dataclasses
22
import logging
33
from datetime import datetime
4-
from typing import Any, Iterable, List, Optional, TypeVar, Union
4+
from typing import Any, Dict, Iterable, List, Optional, TypeVar, Union
55

66
from google.api_core.exceptions import GoogleAPICallError
77
from google.cloud import aiplatform
@@ -157,8 +157,8 @@ def __init__(self, ctx: PipelineContext, config: VertexAIConfig):
157157
project=config.project_id, location=config.region, credentials=credentials
158158
)
159159
self.client = aiplatform
160-
self.endpoints: Optional[dict[str, List[Endpoint]]] = None
161-
self.datasets: Optional[dict[str, VertexAiResourceNoun]] = None
160+
self.endpoints: Optional[Dict[str, List[Endpoint]]] = None
161+
self.datasets: Optional[Dict[str, VertexAiResourceNoun]] = None
162162
self.experiments: Optional[List[Experiment]] = None
163163

164164
def get_report(self) -> SourceReport:
@@ -560,7 +560,7 @@ def _search_dataset(self, dataset_id: str) -> Optional[VertexAiResourceNoun]:
560560
]
561561

562562
if self.datasets is None:
563-
self.datasets = dict[str, VertexAiResourceNoun]()
563+
self.datasets = {}
564564

565565
for dtype in dataset_types:
566566
dataset_class = getattr(self.client.datasets, dtype)
@@ -799,7 +799,7 @@ def _search_endpoint(self, model: Model) -> List[Endpoint]:
799799
Search for an endpoint associated with the model.
800800
"""
801801
if self.endpoints is None:
802-
endpoint_dict = dict[str, List[Endpoint]]()
802+
endpoint_dict: Dict[str, List[Endpoint]] = {}
803803
for endpoint in self.client.Endpoint.list():
804804
for resource in endpoint.list_models():
805805
if resource.model not in endpoint_dict:

0 commit comments

Comments
 (0)