googleapis · nayaknishant · Jun 28, 2022 · Dec 20, 2021 · Dec 20, 2021 · Dec 20, 2021
diff --git a/google/cloud/aiplatform/__init__.py b/google/cloud/aiplatform/__init__.py
@@ -44,6 +44,7 @@
 )
 from google.cloud.aiplatform import metadata
 from google.cloud.aiplatform.models import Endpoint
+from google.cloud.aiplatform.models import PrivateEndpoint
 from google.cloud.aiplatform.models import Model
 from google.cloud.aiplatform.model_evaluation import ModelEvaluation
 from google.cloud.aiplatform.jobs import (
@@ -136,6 +137,7 @@
     "Model",
     "ModelEvaluation",
     "PipelineJob",
+    "PrivateEndpoint",
     "SequenceToSequencePlusForecastingTrainingJob",
     "TabularDataset",
     "Tensorboard",

diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py
@@ -98,7 +98,7 @@ def init(
         Raises:
             ValueError:
                 If experiment_description is provided but experiment is not.
-                If experiment_tensorboard is provided but expeirment is not.
+                If experiment_tensorboard is provided but experiment is not.
         """
 
         if experiment_description and experiment is None:

diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py
diff --git a/setup.py b/setup.py
@@ -52,21 +52,25 @@
     "pandas >= 1.0.0",
     "pyarrow >= 6.0.1",
 ]
-pipelines_extra_requires = [
+pipelines_extra_require = [
     "pyyaml>=5.3,<6",
 ]
 datasets_extra_require = [
     "pyarrow >= 3.0.0, < 8.0dev",
 ]
+private_endpoints_extra_require = [
+    "urllib3 >=1.21.1, <1.27",
+]
 full_extra_require = list(
     set(
         tensorboard_extra_require
         + metadata_extra_require
         + xai_extra_require
         + lit_extra_require
         + featurestore_extra_require
-        + pipelines_extra_requires
+        + pipelines_extra_require
         + datasets_extra_require
+        + private_endpoints_extra_require
     )
 )
 testing_extra_require = (
@@ -118,7 +122,9 @@
         "xai": xai_extra_require,
         "lit": lit_extra_require,
         "cloud_profiler": profiler_extra_require,
-        "pipelines": pipelines_extra_requires,
+        "pipelines": pipelines_extra_require,
+        "datasets": datasets_extra_require,
+        "private_endpoints": private_endpoints_extra_require,
     },
     python_requires=">=3.7",
     classifiers=[

diff --git a/tests/system/aiplatform/e2e_base.py b/tests/system/aiplatform/e2e_base.py
@@ -173,7 +173,6 @@ def tear_down_resources(self, shared_state: Dict[str, Any]):
         yield
 
         # TODO(b/218310362): Add resource deletion system tests
-
         # Bring all Endpoints to the front of the list
         # Ensures Models are undeployed first before we attempt deletion
         shared_state["resources"].sort(

diff --git a/tests/system/aiplatform/test_private_endpoint.py b/tests/system/aiplatform/test_private_endpoint.py
@@ -0,0 +1,69 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import pytest
+
+from google.cloud import aiplatform
+
+from tests.system.aiplatform import e2e_base
+
+# permanent_custom_mnist_model
+_MODEL_ID = "6430031960164270080"
+_PRIVATE_ENDPOINT_NETWORK = "projects/580378083368/global/networks/private-endpoint-vpc"
+
+
+@pytest.mark.usefixtures("tear_down_resources")
+class TestPrivateEndpoint(e2e_base.TestEndToEnd):
+
+    _temp_prefix = "temp_vertex_sdk_e2e"
+
+    def test_create_deploy_delete_private_endpoint(self, shared_state):
+        # Collection of resources generated by this test, to be deleted during teardown
+        shared_state["resources"] = []
+
+        aiplatform.init(
+            project=e2e_base._PROJECT,
+            location=e2e_base._LOCATION,
+        )
+
+        private_endpoint = aiplatform.PrivateEndpoint.create(
+            display_name=self._make_display_name("private_endpoint_test"),
+            network=_PRIVATE_ENDPOINT_NETWORK,
+        )
+        shared_state["resources"].append(private_endpoint)
+
+        # Verify that the retrieved private Endpoint is the same
+        my_private_endpoint = aiplatform.PrivateEndpoint(
+            endpoint_name=private_endpoint.resource_name
+        )
+        assert private_endpoint.resource_name == my_private_endpoint.resource_name
+        assert private_endpoint.display_name == my_private_endpoint.display_name
+
+        # Verify the endpoint is in the private Endpoint list
+        list_private_endpoint = aiplatform.PrivateEndpoint.list()
+        assert private_endpoint.resource_name in [
+            private_endpoint.resource_name for private_endpoint in list_private_endpoint
+        ]
+
+        # Retrieve permanent model, deploy to private Endpoint, then undeploy
+        my_model = aiplatform.Model(model_name=_MODEL_ID)
+
+        my_private_endpoint.deploy(model=my_model)
+        assert my_private_endpoint._gca_resource.deployed_models
+
+        deployed_model_id = my_private_endpoint.list_models()[0].id
+        my_private_endpoint.undeploy(deployed_model_id=deployed_model_id)
+        assert not my_private_endpoint._gca_resource.deployed_models