Support for Private Endpoint on Model Deployment and AQUA MD (#979)

Author: lu-ohai

ads/aqua/extension/deployment_handler.py

@@ -102,6 +102,7 @@ def post(self, *args, **kwargs):
         ocpus = input_data.get("ocpus")
         memory_in_gbs = input_data.get("memory_in_gbs")
         model_file = input_data.get("model_file")
+        private_endpoint_id = input_data.get("private_endpoint_id")
 
         self.finish(
             AquaDeploymentApp().create(
@@ -124,6 +125,7 @@ def post(self, *args, **kwargs):
                 ocpus=ocpus,
                 memory_in_gbs=memory_in_gbs,
                 model_file=model_file,
+                private_endpoint_id=private_endpoint_id,
             )
         )
 

ads/aqua/modeldeployment/deployment.py

@@ -106,6 +106,7 @@ def create(
         memory_in_gbs: Optional[float] = None,
         ocpus: Optional[float] = None,
         model_file: Optional[str] = None,
+        private_endpoint_id: Optional[str] = None,
     ) -> "AquaDeployment":
         """
         Creates a new Aqua deployment
@@ -152,6 +153,9 @@ def create(
             The ocpu count for the shape selected.
         model_file: str
             The file used for model deployment.
+        private_endpoint_id: str
+            The private endpoint id of model deployment.
+
         Returns
         -------
         AquaDeployment
@@ -345,6 +349,7 @@ def create(
             .with_bandwidth_mbps(bandwidth_mbps)
             .with_replica(instance_count)
             .with_web_concurrency(web_concurrency)
+            .with_private_endpoint_id(private_endpoint_id)
             .with_access_log(
                 log_group_id=log_group_id,
                 log_id=access_log_id,

ads/aqua/modeldeployment/entities.py

@@ -3,7 +3,7 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 from dataclasses import dataclass, field
-from typing import Union
+from typing import Union, Optional
 
 from oci.data_science.models import (
     ModelDeployment,
@@ -47,9 +47,10 @@ class AquaDeployment(DataClassSerializable):
     created_on: str = None
     created_by: str = None
     endpoint: str = None
+    private_endpoint_id: str = None
     console_link: str = None
     lifecycle_details: str = None
-    shape_info: field(default_factory=ShapeInfo) = None
+    shape_info: Optional[ShapeInfo] = None
     tags: dict = None
     environment_variables: dict = None
 
@@ -98,6 +99,7 @@ def from_oci_model_deployment(
         freeform_tags = oci_model_deployment.freeform_tags or UNKNOWN_DICT
         aqua_service_model_tag = freeform_tags.get(Tags.AQUA_SERVICE_MODEL_TAG, None)
         aqua_model_name = freeform_tags.get(Tags.AQUA_MODEL_NAME_TAG, UNKNOWN)
+        private_endpoint_id = getattr(instance_configuration, "private_endpoint_id", UNKNOWN)
 
         return AquaDeployment(
             id=oci_model_deployment.id,
@@ -113,6 +115,7 @@ def from_oci_model_deployment(
             created_on=str(oci_model_deployment.time_created),
             created_by=oci_model_deployment.created_by,
             endpoint=oci_model_deployment.model_deployment_url,
+            private_endpoint_id=private_endpoint_id,
             console_link=get_console_link(
                 resource="model-deployments",
                 ocid=oci_model_deployment.id,

ads/model/deployment/model_deployment.py

@@ -1648,6 +1648,19 @@ def _build_model_deployment_configuration_details(self) -> Dict:
                 infrastructure.CONST_SUBNET_ID
             ] = infrastructure.subnet_id
 
+        if infrastructure.private_endpoint_id:
+            if not hasattr(
+                oci.data_science.models.InstanceConfiguration, "private_endpoint_id"
+            ):
+                # TODO: add oci version with private endpoint support.
+                raise EnvironmentError(
+                    "Private endpoint is not supported in the current OCI SDK installed."
+                )
+
+            instance_configuration[
+                infrastructure.CONST_PRIVATE_ENDPOINT_ID
+            ] = infrastructure.private_endpoint_id
+
         scaling_policy = {
             infrastructure.CONST_POLICY_TYPE: "FIXED_SIZE",
             infrastructure.CONST_INSTANCE_COUNT: infrastructure.replica

ads/model/deployment/model_deployment_infrastructure.py

@@ -57,6 +57,8 @@ class ModelDeploymentInfrastructure(Builder):
         The web concurrency of model deployment
     subnet_id: str
         The subnet id of model deployment
+    private_endpoint_id: str
+        The private endpoint id of model deployment
 
     Methods
     -------
@@ -84,6 +86,8 @@ class ModelDeploymentInfrastructure(Builder):
         Sets the web concurrency of model deployment
     with_subnet_id(subnet_id)
         Sets the subnet id of model deployment
+    with_private_endpoint_id(private_endpoint)
+        Sets the private endpoint id of model deployment
 
     Example
     -------
@@ -100,6 +104,7 @@ class ModelDeploymentInfrastructure(Builder):
     ...        .with_bandwidth_mbps(10)
     ...        .with_web_concurrency(10)
     ...        .with_subnet_id(<subnet_id>)
+    ...        .with_private_endpoint_id(<private_endpoint_id>)
     ...        .with_access_log(
     ...            log_group_id=<log_group_id>,
     ...            log_id=<log_id>
@@ -143,6 +148,7 @@ class ModelDeploymentInfrastructure(Builder):
     CONST_WEB_CONCURRENCY = "webConcurrency"
     CONST_STREAM_CONFIG_DETAILS = "streamConfigurationDetails"
     CONST_SUBNET_ID = "subnetId"
+    CONST_PRIVATE_ENDPOINT_ID = "privateEndpointId"
 
     attribute_map = {
         CONST_PROJECT_ID: "project_id",
@@ -159,6 +165,7 @@ class ModelDeploymentInfrastructure(Builder):
         CONST_LOG_GROUP_ID: "log_group_id",
         CONST_WEB_CONCURRENCY: "web_concurrency",
         CONST_SUBNET_ID: "subnet_id",
+        CONST_PRIVATE_ENDPOINT_ID: "private_endpoint_id",
     }
 
     shape_config_details_attribute_map = {
@@ -186,6 +193,7 @@ class ModelDeploymentInfrastructure(Builder):
         CONST_SHAPE_NAME: f"{MODEL_CONFIG_DETAILS_PATH}.instance_configuration.instance_shape_name",
         CONST_SHAPE_CONFIG_DETAILS: f"{MODEL_CONFIG_DETAILS_PATH}.instance_configuration.model_deployment_instance_shape_config_details",
         CONST_SUBNET_ID: f"{MODEL_CONFIG_DETAILS_PATH}.instance_configuration.subnet_id",
+        CONST_PRIVATE_ENDPOINT_ID: f"{MODEL_CONFIG_DETAILS_PATH}.instance_configuration.private_endpoint_id",
         CONST_REPLICA: f"{MODEL_CONFIG_DETAILS_PATH}.scaling_policy.instance_count",
         CONST_BANDWIDTH_MBPS: f"{MODEL_CONFIG_DETAILS_PATH}.bandwidth_mbps",
         CONST_ACCESS_LOG: "category_log_details.access",
@@ -613,6 +621,32 @@ def subnet_id(self) -> str:
         """
         return self.get_spec(self.CONST_SUBNET_ID, None)
 
+    def with_private_endpoint_id(self, private_endpoint_id: str) -> "ModelDeploymentInfrastructure":
+        """Sets the private endpoint id of model deployment.
+
+        Parameters
+        ----------
+        private_endpoint_id : str
+            The private endpoint id of model deployment.
+
+        Returns
+        -------
+        ModelDeploymentInfrastructure
+            The ModelDeploymentInfrastructure instance (self).
+        """
+        return self.set_spec(self.CONST_PRIVATE_ENDPOINT_ID, private_endpoint_id)
+
+    @property
+    def private_endpoint_id(self) -> str:
+        """The model deployment private endpoint id.
+
+        Returns
+        -------
+        str
+            The model deployment private endpoint id.
+        """
+        return self.get_spec(self.CONST_PRIVATE_ENDPOINT_ID, None)
+
     def init(self, **kwargs) -> "ModelDeploymentInfrastructure":
         """Initializes a starter specification for the ModelDeploymentInfrastructure.
 

ads/model/generic_model.py

@@ -2262,6 +2262,7 @@ def deploy(
         description: Optional[str] = None,
         deployment_instance_shape: Optional[str] = None,
         deployment_instance_subnet_id: Optional[str] = None,
+        deployment_instance_private_endpoint_id: Optional[str] = None,
         deployment_instance_count: Optional[int] = None,
         deployment_bandwidth_mbps: Optional[int] = None,
         deployment_log_group_id: Optional[str] = None,
@@ -2312,6 +2313,8 @@ def deploy(
             The shape of the instance used for deployment.
         deployment_instance_subnet_id: (str, optional). Default to None.
             The subnet id of the instance used for deployment.
+        deployment_instance_private_endpoint_id: (str, optional). Default to None.
+            The private endpoint id of instance used for deployment.
         deployment_instance_count: (int, optional). Defaults to 1.
             The number of instance used for deployment.
         deployment_bandwidth_mbps: (int, optional). Defaults to 10.
@@ -2432,6 +2435,8 @@ def deploy(
             or self.properties.deployment_image,
             deployment_instance_subnet_id=existing_infrastructure.subnet_id
             or self.properties.deployment_instance_subnet_id,
+            deployment_instance_private_endpoint_id=existing_infrastructure.private_endpoint_id
+            or self.properties.deployment_instance_private_endpoint_id,
         ).to_dict()
 
         property_dict.update(override_properties)
@@ -2465,6 +2470,7 @@ def deploy(
             .with_shape_name(self.properties.deployment_instance_shape)
             .with_replica(self.properties.deployment_instance_count)
             .with_subnet_id(self.properties.deployment_instance_subnet_id)
+            .with_private_endpoint_id(self.properties.deployment_instance_private_endpoint_id)
         )
 
         web_concurrency = (
@@ -2611,6 +2617,7 @@ def prepare_save_deploy(
         deployment_description: Optional[str] = None,
         deployment_instance_shape: Optional[str] = None,
         deployment_instance_subnet_id: Optional[str] = None,
+        deployment_instance_private_endpoint_id: Optional[str] = None,
         deployment_instance_count: Optional[int] = None,
         deployment_bandwidth_mbps: Optional[int] = None,
         deployment_log_group_id: Optional[str] = None,
@@ -2701,6 +2708,8 @@ def prepare_save_deploy(
             The shape of the instance used for deployment.
         deployment_instance_subnet_id: (str, optional). Default to None.
             The subnet id of the instance used for deployment.
+        deployment_instance_private_endpoint_id: (str, optional). Default to None.
+            The private endpoint id of instance used for deployment.
         deployment_instance_count: (int, optional). Defaults to 1.
             The number of instance used for deployment.
         deployment_bandwidth_mbps: (int, optional). Defaults to 10.
@@ -2846,6 +2855,7 @@ def prepare_save_deploy(
             description=deployment_description,
             deployment_instance_shape=self.properties.deployment_instance_shape,
             deployment_instance_subnet_id=self.properties.deployment_instance_subnet_id,
+            deployment_instance_private_endpoint_id=self.properties.deployment_instance_private_endpoint_id,
             deployment_instance_count=self.properties.deployment_instance_count,
             deployment_bandwidth_mbps=self.properties.deployment_bandwidth_mbps,
             deployment_log_group_id=self.properties.deployment_log_group_id,

ads/model/model_properties.py

@@ -29,6 +29,7 @@ class ModelProperties(BaseProperties):
     overwrite_existing_artifact: bool = None
     deployment_instance_shape: str = None
     deployment_instance_subnet_id: str = None
+    deployment_instance_private_endpoint_id: str = None
     deployment_instance_count: int = None
     deployment_bandwidth_mbps: int = None
     deployment_log_group_id: str = None

tests/unitary/with_extras/aqua/test_deployment.py

@@ -181,6 +181,7 @@ class TestDataset:
         "created_on": "2024-01-01T00:00:00.000000+00:00",
         "created_by": "ocid1.user.oc1..<OCID>",
         "endpoint": MODEL_DEPLOYMENT_URL,
+        "private_endpoint_id": "",
         "environment_variables": {
             "BASE_MODEL": "service_models/model-name/artifact",
             "MODEL_DEPLOY_ENABLE_STREAMING": "true",

tests/unitary/with_extras/aqua/test_deployment_handler.py

@@ -129,6 +129,7 @@ def test_post(self, mock_create):
             memory_in_gbs=None,
             ocpus=None,
             model_file=None,
+            private_endpoint_id=None,
         )
 
 

tests/unitary/with_extras/model/test_generic_model.py

@@ -1514,6 +1514,7 @@ def test__to_yaml(self):
                     "description": None,
                     "deployment_instance_shape": None,
                     "deployment_instance_subnet_id": None,
+                    "deployment_instance_private_endpoint_id": None,
                     "deployment_instance_count": None,
                     "deployment_bandwidth_mbps": None,
                     "deployment_log_group_id": None,
@@ -1554,6 +1555,7 @@ def test__to_yaml(self):
                     "deployment_description": None,
                     "deployment_instance_shape": None,
                     "deployment_instance_subnet_id": None,
+                    "deployment_instance_private_endpoint_id": None,
                     "deployment_instance_count": None,
                     "deployment_bandwidth_mbps": None,
                     "deployment_log_group_id": None,
@@ -1606,6 +1608,7 @@ def test__to_yaml(self):
                     "description": None,
                     "deployment_instance_shape": None,
                     "deployment_instance_subnet_id": None,
+                    "deployment_instance_private_endpoint_id": None,
                     "deployment_instance_count": None,
                     "deployment_bandwidth_mbps": None,
                     "deployment_log_group_id": None,
@@ -1646,6 +1649,7 @@ def test__to_yaml(self):
                     "deployment_description": "fake_deployment_description",
                     "deployment_instance_shape": "2.1",
                     "deployment_instance_subnet_id": "ocid1.subnet.oc1.iad.<unique_ocid>",
+                    "deployment_instance_private_endpoint_id": "ocid1.datascienceprivateendpointint.oc1.iad.<unique_ocid>",
                     "deployment_instance_count": 1,
                     "deployment_bandwidth_mbps": 10,
                     "deployment_log_group_id": "ocid1.loggroup.oc1.iad.<unique_ocid>",
@@ -1704,6 +1708,7 @@ def test__to_yaml(self):
                     "deployment_instance_count": 1,
                     "deployment_bandwidth_mbps": 10,
                     "deployment_instance_subnet_id": "ocid1.subnet.oc1.iad.<unique_ocid>",
+                    "deployment_instance_private_endpoint_id": "ocid1.datascienceprivateendpointint.oc1.iad.<unique_ocid>",
                     "deployment_log_group_id": "ocid1.loggroup.oc1.iad.<unique_ocid>",
                     "deployment_access_log_id": "ocid1.log.oc1.iad.<unique_ocid>",
                     "deployment_predict_log_id": "ocid1.log.oc1.iad.<unique_ocid>",
@@ -1746,6 +1751,7 @@ def test__to_yaml(self):
                     "deployment_description": "fake_deployment_description",
                     "deployment_instance_shape": "2.1",
                     "deployment_instance_subnet_id": "ocid1.subnet.oc1.iad.<unique_ocid>",
+                    "deployment_instance_private_endpoint_id": "ocid1.datascienceprivateendpointint.oc1.iad.<unique_ocid>",
                     "deployment_instance_count": 1,
                     "deployment_bandwidth_mbps": 10,
                     "deployment_log_group_id": "ocid",
@@ -1810,6 +1816,7 @@ def test__to_yaml(self):
                     "deployment_instance_count": 1,
                     "deployment_bandwidth_mbps": 10,
                     "deployment_instance_subnet_id": "ocid1.subnet.oc1.iad.<unique_ocid>",
+                    "deployment_instance_private_endpoint_id": "ocid1.datascienceprivateendpointint.oc1.iad.<unique_ocid>",
                     "deployment_log_group_id": "ocid",
                     "deployment_access_log_id": "ocid",
                     "deployment_predict_log_id": "ocid",
@@ -1882,6 +1889,7 @@ def test_prepare_save_deploy_with_default_display_name(
             "description": None,
             "deployment_instance_shape": None,
             "deployment_instance_subnet_id": None,
+            "deployment_instance_private_endpoint_id": None,
             "deployment_instance_count": None,
             "deployment_bandwidth_mbps": None,
             "deployment_memory_in_gbs": None,