diff --git a/ads/aqua/extension/deployment_handler.py b/ads/aqua/extension/deployment_handler.py index 3e3a6c277..e8652e46f 100644 --- a/ads/aqua/extension/deployment_handler.py +++ b/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, ) ) diff --git a/ads/aqua/modeldeployment/deployment.py b/ads/aqua/modeldeployment/deployment.py index 6a534ac4a..8d3a07e8d 100644 --- a/ads/aqua/modeldeployment/deployment.py +++ b/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, diff --git a/ads/aqua/modeldeployment/entities.py b/ads/aqua/modeldeployment/entities.py index 411c68529..dd2d72e10 100644 --- a/ads/aqua/modeldeployment/entities.py +++ b/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, diff --git a/ads/model/deployment/model_deployment.py b/ads/model/deployment/model_deployment.py index df6b681a2..56a70c112 100644 --- a/ads/model/deployment/model_deployment.py +++ b/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 diff --git a/ads/model/deployment/model_deployment_infrastructure.py b/ads/model/deployment/model_deployment_infrastructure.py index 3610db919..fb81b8920 100644 --- a/ads/model/deployment/model_deployment_infrastructure.py +++ b/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() + ... .with_private_endpoint_id() ... .with_access_log( ... log_group_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. diff --git a/ads/model/generic_model.py b/ads/model/generic_model.py index 7a44b82de..842ae94a8 100644 --- a/ads/model/generic_model.py +++ b/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, diff --git a/ads/model/model_properties.py b/ads/model/model_properties.py index 6689dc754..4be249e93 100644 --- a/ads/model/model_properties.py +++ b/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 diff --git a/tests/unitary/with_extras/aqua/test_deployment.py b/tests/unitary/with_extras/aqua/test_deployment.py index 7ef73c25f..6a4840ea9 100644 --- a/tests/unitary/with_extras/aqua/test_deployment.py +++ b/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..", "endpoint": MODEL_DEPLOYMENT_URL, + "private_endpoint_id": "", "environment_variables": { "BASE_MODEL": "service_models/model-name/artifact", "MODEL_DEPLOY_ENABLE_STREAMING": "true", diff --git a/tests/unitary/with_extras/aqua/test_deployment_handler.py b/tests/unitary/with_extras/aqua/test_deployment_handler.py index 54756545a..d4def6bb8 100644 --- a/tests/unitary/with_extras/aqua/test_deployment_handler.py +++ b/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, ) diff --git a/tests/unitary/with_extras/model/test_generic_model.py b/tests/unitary/with_extras/model/test_generic_model.py index c8e3ea75f..c070c606e 100644 --- a/tests/unitary/with_extras/model/test_generic_model.py +++ b/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.", + "deployment_instance_private_endpoint_id": "ocid1.datascienceprivateendpointint.oc1.iad.", "deployment_instance_count": 1, "deployment_bandwidth_mbps": 10, "deployment_log_group_id": "ocid1.loggroup.oc1.iad.", @@ -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.", + "deployment_instance_private_endpoint_id": "ocid1.datascienceprivateendpointint.oc1.iad.", "deployment_log_group_id": "ocid1.loggroup.oc1.iad.", "deployment_access_log_id": "ocid1.log.oc1.iad.", "deployment_predict_log_id": "ocid1.log.oc1.iad.", @@ -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.", + "deployment_instance_private_endpoint_id": "ocid1.datascienceprivateendpointint.oc1.iad.", "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.", + "deployment_instance_private_endpoint_id": "ocid1.datascienceprivateendpointint.oc1.iad.", "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,