feat: expose PSC for CustomModel

PiperOrigin-RevId: 815235726

Author: vertex-sdk-bot

tests/unit/vertexai/model_garden/test_model_garden.py

@@ -1551,6 +1551,45 @@ def test_deploy_custom_model_with_all_config_success(self, deploy_mock):
             )
         )
 
+    def test_deploy_custom_model_with_psc_success(self, deploy_mock):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+        )
+        model = model_garden_preview.CustomModel(gcs_uri=_TEST_GCS_URI)
+        model.deploy(
+            machine_type="n1-standard-4",
+            accelerator_type="NVIDIA_TESLA_T4",
+            accelerator_count=1,
+            enable_private_service_connect=True,
+            psc_project_allow_list=["test-project"],
+        )
+        deploy_mock.assert_called_once_with(
+            types.DeployRequest(
+                destination=f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}",
+                custom_model=types.DeployRequest.CustomModel(
+                    gcs_uri=_TEST_GCS_URI,
+                ),
+                endpoint_config=types.DeployRequest.EndpointConfig(
+                    private_service_connect_config=types.PrivateServiceConnectConfig(
+                        enable_private_service_connect=True,
+                        project_allowlist=["test-project"],
+                    ),
+                ),
+                deploy_config=types.DeployRequest.DeployConfig(
+                    dedicated_resources=types.DedicatedResources(
+                        min_replica_count=1,
+                        max_replica_count=1,
+                        machine_spec=types.MachineSpec(
+                            machine_type="n1-standard-4",
+                            accelerator_type="NVIDIA_TESLA_T4",
+                            accelerator_count=1,
+                        ),
+                    ),
+                ),
+            )
+        )
+
     def test_deploy_custom_model_with_reservation_success(self, deploy_mock):
         aiplatform.init(
             project=_TEST_PROJECT,

vertexai/model_garden/_model_garden.py

@@ -1063,6 +1063,8 @@ def deploy(
         reservation_affinity_values: Optional[List[str]] = None,
         endpoint_display_name: Optional[str] = None,
         model_display_name: Optional[str] = None,
+        enable_private_service_connect: bool = False,
+        psc_project_allow_list: Optional[List[str]] = None,
         deploy_request_timeout: Optional[float] = None,
     ) -> aiplatform.Endpoint:
         """Deploys a Custom Model to an endpoint.
@@ -1105,6 +1107,10 @@ def deploy(
                   'projects/{project_id_or_number}/zones/{zone}/reservations/{reservation_name}'
             endpoint_display_name: The display name of the created endpoint.
             model_display_name: The display name of the custom model.
+            enable_private_service_connect (bool): Whether to enable private service
+              connect.
+            psc_project_allow_list (List[str]): The list of projects that are allowed to
+              access the endpoint over private service connect.
             deploy_request_timeout: The timeout for the deploy request. Default is 2
               hours.
 
@@ -1123,6 +1129,8 @@ def deploy(
             reservation_affinity_values=reservation_affinity_values,
             endpoint_display_name=endpoint_display_name,
             model_display_name=model_display_name,
+            enable_private_service_connect=enable_private_service_connect,
+            psc_project_allow_list=psc_project_allow_list,
             deploy_request_timeout=deploy_request_timeout,
         )
 
@@ -1139,6 +1147,8 @@ def _deploy_gcs_uri(
         max_replica_count: int = 1,
         accelerator_type: Optional[str] = None,
         accelerator_count: Optional[int] = None,
+        enable_private_service_connect: bool = False,
+        psc_project_allow_list: Optional[List[str]] = None,
         reservation_affinity_type: Optional[str] = None,
         reservation_affinity_key: Optional[str] = None,
         reservation_affinity_values: Optional[List[str]] = None,
@@ -1188,6 +1198,10 @@ def _deploy_gcs_uri(
                   'projects/{project_id_or_number}/zones/{zone}/reservations/{reservation_name}'
             endpoint_display_name: The display name of the created endpoint.
             model_display_name: The display name of the custom model.
+            enable_private_service_connect (bool): Whether to enable private service
+              connect.
+            psc_project_allow_list (List[str]): The list of projects that are allowed to
+              access the endpoint over private service connect.
             deploy_request_timeout: The timeout for the deploy request. Default is 2
               hours.
 
@@ -1217,6 +1231,14 @@ def has_all_or_none_values(var1, var2, var3) -> bool:
         if model_display_name:
             request.model_config.model_display_name = model_display_name
 
+        if enable_private_service_connect and psc_project_allow_list:
+            request.endpoint_config.private_service_connect_config = (
+                types.PrivateServiceConnectConfig(
+                    enable_private_service_connect=enable_private_service_connect,
+                    project_allowlist=psc_project_allow_list,
+                )
+            )
+
         if machine_type and accelerator_type and accelerator_count:
             request.deploy_config.dedicated_resources = types.DedicatedResources(
                 machine_spec=types.MachineSpec(