Support for PE on AQUA MD.

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

@@ -107,6 +107,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
@@ -153,6 +154,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
@@ -342,6 +346,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

@@ -47,6 +47,7 @@ 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
@@ -98,6 +99,9 @@ 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(oci_model_deployment.model_deployment_configuration_details.model_configuration_details.instance_configuration, "private_endpoint_id", UNKNOWN)
+        )
 
         return AquaDeployment(
             id=oci_model_deployment.id,
@@ -113,6 +117,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,18 @@ 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"
+            ):
+                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