Skip to content

Commit 162b2f2

Browse files
Ark-kuncopybara-github
authored andcommitted
fix: Resource created by _construct_sdk_resource_from_gapic should use the project from the resource name instead of the default project.
This fixes the following failing assertion: ``` assert aiplatform.Endpoint._construct_sdk_resource_from_gapic( aiplatform_models.gca_endpoint_compat.Endpoint(name='projects/my-project/locations/us-central1/publishers/google/models/text-bison@001') ).project == "my-project" ``` PiperOrigin-RevId: 536839787
1 parent 50e0898 commit 162b2f2

File tree

4 files changed

+36
-23
lines changed

4 files changed

+36
-23
lines changed

google/cloud/aiplatform/base.py

+6-1
Original file line numberDiff line numberDiff line change
@@ -1024,8 +1024,13 @@ def _construct_sdk_resource_from_gapic(
10241024
VertexAiResourceNoun:
10251025
An initialized SDK object that represents GAPIC type.
10261026
"""
1027+
resource_name_parts = utils.extract_project_and_location_from_parent(
1028+
gapic_resource.name
1029+
)
10271030
sdk_resource = cls._empty_constructor(
1028-
project=project, location=location, credentials=credentials
1031+
project=resource_name_parts.get("project") or project,
1032+
location=resource_name_parts.get("location") or location,
1033+
credentials=credentials,
10291034
)
10301035
sdk_resource._gca_resource = gapic_resource
10311036
return sdk_resource

google/cloud/aiplatform/models.py

+10-8
Original file line numberDiff line numberDiff line change
@@ -560,12 +560,13 @@ def _construct_sdk_resource_from_gapic(
560560
Endpoint (aiplatform.Endpoint):
561561
An initialized Endpoint resource.
562562
"""
563-
endpoint = cls._empty_constructor(
564-
project=project, location=location, credentials=credentials
563+
endpoint = super()._construct_sdk_resource_from_gapic(
564+
gapic_resource=gapic_resource,
565+
project=project,
566+
location=location,
567+
credentials=credentials,
565568
)
566569

567-
endpoint._gca_resource = gapic_resource
568-
569570
endpoint._prediction_client = cls._instantiate_prediction_client(
570571
location=endpoint.location,
571572
credentials=credentials,
@@ -2021,12 +2022,13 @@ def _construct_sdk_resource_from_gapic(
20212022
"Cannot import the urllib3 HTTP client. Please install google-cloud-aiplatform[private_endpoints]."
20222023
)
20232024

2024-
endpoint = cls._empty_constructor(
2025-
project=project, location=location, credentials=credentials
2025+
endpoint = super()._construct_sdk_resource_from_gapic(
2026+
gapic_resource=gapic_resource,
2027+
project=project,
2028+
location=location,
2029+
credentials=credentials,
20262030
)
20272031

2028-
endpoint._gca_resource = gapic_resource
2029-
20302032
endpoint._http_client = urllib3.PoolManager()
20312033

20322034
return endpoint

google/cloud/aiplatform/preview/models.py

-14
Original file line numberDiff line numberDiff line change
@@ -431,20 +431,6 @@ def list(
431431
credentials=credentials,
432432
)
433433

434-
@classmethod
435-
def _construct_sdk_resource_from_gapic(
436-
cls,
437-
gapic_resource: proto.Message,
438-
project: Optional[str] = None,
439-
location: Optional[str] = None,
440-
credentials: Optional[auth_credentials.Credentials] = None,
441-
) -> "models.DeploymentResourcePool":
442-
drp = cls._empty_constructor(
443-
project=project, location=location, credentials=credentials
444-
)
445-
drp._gca_resource = gapic_resource
446-
return drp
447-
448434

449435
class Endpoint(aiplatform.Endpoint):
450436
@staticmethod

tests/unit/aiplatform/test_endpoints.py

+20
Original file line numberDiff line numberDiff line change
@@ -2202,3 +2202,23 @@ def test_delete_with_force(self, sdk_undeploy_mock, delete_endpoint_mock, sync):
22022202
def test_list(self):
22032203
ep_list = aiplatform.PrivateEndpoint.list()
22042204
assert ep_list # Ensure list is not empty
2205+
2206+
def test_construct_sdk_resource_from_gapic_uses_resource_project(self):
2207+
PROJECT = "my-project"
2208+
LOCATION = "me-west1"
2209+
endpoint_name = f"projects/{PROJECT}/locations/{LOCATION}/endpoints/123"
2210+
endpoint = aiplatform.Endpoint._construct_sdk_resource_from_gapic(
2211+
models.gca_endpoint_compat.Endpoint(name=endpoint_name)
2212+
)
2213+
assert endpoint.project == PROJECT
2214+
assert endpoint.location == LOCATION
2215+
assert endpoint.project != _TEST_PROJECT
2216+
assert endpoint.location != _TEST_LOCATION
2217+
2218+
endpoint2 = aiplatform.Endpoint._construct_sdk_resource_from_gapic(
2219+
models.gca_endpoint_compat.Endpoint(name=endpoint_name),
2220+
project=_TEST_PROJECT,
2221+
location=_TEST_LOCATION,
2222+
)
2223+
assert endpoint2.project != _TEST_PROJECT
2224+
assert endpoint2.location != _TEST_LOCATION

0 commit comments

Comments
 (0)