Skip to content

Commit 374abd1

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
fix: Use PrivateServiceConnectConfig in service_networking in PrivateEndpoint instead of the wrapper class PrivateServiceConnectConfig
PiperOrigin-RevId: 775512697
1 parent d570fc9 commit 374abd1

File tree

2 files changed

+58
-12
lines changed

2 files changed

+58
-12
lines changed

google/cloud/aiplatform/models.py

Lines changed: 24 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3363,6 +3363,8 @@ def health_http_uri(self) -> Optional[str]:
33633363
return None
33643364
return self._gca_resource.deployed_models[0].private_endpoints.health_http_uri
33653365

3366+
# PrivateServiceConnectConfig is deprecated.
3367+
# Use service_networking.PrivateServiceConnectConfig instead.
33663368
class PrivateServiceConnectConfig:
33673369
"""Represents a Vertex AI PrivateServiceConnectConfig resource."""
33683370

@@ -3399,7 +3401,10 @@ def create(
33993401
credentials: Optional[auth_credentials.Credentials] = None,
34003402
encryption_spec_key_name: Optional[str] = None,
34013403
sync=True,
3402-
private_service_connect_config: Optional[PrivateServiceConnectConfig] = None,
3404+
private_service_connect_config: Union[
3405+
Optional[PrivateServiceConnectConfig],
3406+
Optional[gca_service_networking.PrivateServiceConnectConfig],
3407+
] = None,
34033408
enable_request_response_logging=False,
34043409
request_response_logging_sampling_rate: Optional[float] = None,
34053410
request_response_logging_bq_destination_table: Optional[str] = None,
@@ -3428,15 +3433,17 @@ def create(
34283433
display_name="my_endpoint_name",
34293434
project="my_project_id",
34303435
location="us-central1",
3431-
private_service_connect=aiplatform.PrivateEndpoint.PrivateServiceConnectConfig(
3436+
private_service_connect=aiplatform.compat.types.service_networking.PrivateServiceConnectConfig(
3437+
enable_private_service_connect=True,
34323438
project_allowlist=["test-project"]),
34333439
)
34343440
34353441
or (when project and location are initialized)
34363442
34373443
my_private_endpoint = aiplatform.PrivateEndpoint.create(
34383444
display_name="my_endpoint_name",
3439-
private_service_connect=aiplatform.PrivateEndpoint.PrivateServiceConnectConfig(
3445+
private_service_connect=aiplatform.compat.types.service_networking.PrivateServiceConnectConfig(
3446+
enable_private_service_connect=True,
34403447
project_allowlist=["test-project"]),
34413448
)
34423449
Args:
@@ -3475,12 +3482,11 @@ def create(
34753482
sync (bool): Whether to execute this method synchronously. If False,
34763483
this method will be executed in concurrent Future and any downstream
34773484
object will be immediately returned and synced when the Future has
3478-
completed. private_service_connect_config
3479-
(aiplatform.PrivateEndpoint.PrivateServiceConnectConfig): [Private
3480-
Service
3481-
Connect](https://cloud.google.com/vpc/docs/private-service-connect)
3482-
configuration for the endpoint. Cannot be set when network is
3483-
specified.
3485+
completed.
3486+
private_service_connect_config
3487+
(aiplatform.compat.types.service_networking.PrivateServiceConnectConfig): [Private
3488+
Service Connect Configuration](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/PrivateServiceConnectConfig)
3489+
for the endpoint. Cannot be set when network is specified.
34843490
enable_request_response_logging (bool): Optional. Whether to enable
34853491
request & response logging for this endpoint.
34863492
request_response_logging_sampling_rate (float): Optional. The request
@@ -3527,9 +3533,15 @@ def create(
35273533

35283534
config = None
35293535
if private_service_connect_config:
3530-
config = (
3531-
private_service_connect_config._gapic_private_service_connect_config
3532-
)
3536+
if hasattr(
3537+
private_service_connect_config,
3538+
"_gapic_private_service_connect_config",
3539+
):
3540+
config = (
3541+
private_service_connect_config._gapic_private_service_connect_config
3542+
)
3543+
else:
3544+
config = private_service_connect_config
35333545

35343546
predict_request_response_logging_config = None
35353547
if enable_request_response_logging:

tests/unit/aiplatform/test_endpoints.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3679,6 +3679,40 @@ def test_create_psc(self, create_psc_private_endpoint_mock, sync):
36793679
endpoint_id=None,
36803680
)
36813681

3682+
@pytest.mark.parametrize("sync", [True, False])
3683+
def test_create_psc_with_service_networking(
3684+
self, create_psc_private_endpoint_mock, sync
3685+
):
3686+
test_endpoint = models.PrivateEndpoint.create(
3687+
display_name=_TEST_DISPLAY_NAME,
3688+
project=_TEST_PROJECT,
3689+
location=_TEST_LOCATION,
3690+
private_service_connect_config=gca_service_networking.PrivateServiceConnectConfig(
3691+
enable_private_service_connect=True,
3692+
project_allowlist=_TEST_PROJECT_ALLOWLIST,
3693+
),
3694+
sync=sync,
3695+
)
3696+
3697+
if not sync:
3698+
test_endpoint.wait()
3699+
3700+
expected_endpoint = gca_endpoint.Endpoint(
3701+
display_name=_TEST_DISPLAY_NAME,
3702+
private_service_connect_config=gca_service_networking.PrivateServiceConnectConfig(
3703+
enable_private_service_connect=True,
3704+
project_allowlist=_TEST_PROJECT_ALLOWLIST,
3705+
),
3706+
)
3707+
3708+
create_psc_private_endpoint_mock.assert_called_once_with(
3709+
parent=_TEST_PARENT,
3710+
endpoint=expected_endpoint,
3711+
metadata=(),
3712+
timeout=None,
3713+
endpoint_id=None,
3714+
)
3715+
36823716
@pytest.mark.parametrize("sync", [True, False])
36833717
def test_create_psc_with_timeout(self, create_psc_private_endpoint_mock, sync):
36843718
test_endpoint = models.PrivateEndpoint.create(

0 commit comments

Comments
 (0)