Add support for custom networking in AQUA Model Deployments (#1298)

Author: mrDzurb

ads/aqua/modeldeployment/deployment.py

@@ -189,6 +189,7 @@ def create(
                 freeform_tags (Optional[Dict]): Freeform tags for model deployment.
                 defined_tags (Optional[Dict]): Defined tags for model deployment.
                 deployment_type (Optional[str]): The type of model deployment.
+                subnet_id (Optional[str]): The custom egress for model deployment.
 
         Returns
         -------
@@ -1150,6 +1151,7 @@ def _create_deployment(
                 log_group_id=create_deployment_details.log_group_id,
                 log_id=create_deployment_details.predict_log_id,
             )
+            .with_subnet_id(create_deployment_details.subnet_id)
         )
         if (
             create_deployment_details.memory_in_gbs

ads/aqua/modeldeployment/entities.py

@@ -132,6 +132,9 @@ class AquaDeployment(Serializable):
     cmd: Optional[List[str]] = Field(
         default_factory=list, description="The cmd of the model deployment."
     )
+    subnet_id: Optional[str] = Field(
+        None, description="The custom egress for model deployment."
+    )
 
     @classmethod
     def from_oci_model_deployment(
@@ -210,6 +213,7 @@ def from_oci_model_deployment(
         private_endpoint_id = getattr(
             instance_configuration, "private_endpoint_id", UNKNOWN
         )
+        subnet_id = getattr(instance_configuration, "subnet_id", UNKNOWN)
 
         return AquaDeployment(
             id=oci_model_deployment.id,
@@ -235,6 +239,7 @@ def from_oci_model_deployment(
             tags=tags,
             environment_variables=environment_variables,
             cmd=cmd,
+            subnet_id=subnet_id,
         )
 
     class Config:
@@ -803,6 +808,9 @@ class CreateModelDeploymentDetails(ModelDeploymentDetails):
     deployment_type: Optional[str] = Field(
         None, description="The type of model deployment."
     )
+    subnet_id: Optional[str] = Field(
+        None, description="The custom egress for model deployment."
+    )
 
     @model_validator(mode="before")
     @classmethod

ads/aqua/server/aqua_spec.yml

@@ -1092,6 +1092,10 @@ components:
           type: string
           nullable: true
           example: "private-endpoint-ocid-001"
+        subnet_id:
+          type: string
+          nullable: true
+          example: "subnet-ocid-001"
         console_link:
           type: string
           example: "https://console.example.com/model-deployments/deployment-ocid-001"
@@ -1180,6 +1184,9 @@ components:
         private_endpoint_id:
           type: string
           example: "ocid1.datasciencemodeldeploymentprivatendpoint.oc1..<UNIQUEID>"
+        subnet_id:
+          type: string
+          example: "ocid1.subnet.oc1.iad..<UNIQUEID>"
         cmd_var:
           type: string
           example: "ocid1.compartment.oc1..<UNIQUEID>"

tests/unitary/with_extras/aqua/test_deployment.py

@@ -12,11 +12,6 @@
 from importlib import reload
 from unittest.mock import MagicMock, patch
 
-from ads.aqua.modeldeployment.constants import DEFAULT_POLL_INTERVAL, DEFAULT_WAIT_TIME
-from ads.model.datascience_model_group import DataScienceModelGroup
-from ads.model.service.oci_datascience_model_deployment import (
-    OCIDataScienceModelDeployment,
-)
 import oci
 import pytest
 from oci.data_science.models import (
@@ -56,6 +51,7 @@
 )
 from ads.aqua.modeldeployment.model_group_config import BaseModelSpec, ModelGroupConfig
 from ads.model.datascience_model import DataScienceModel
+from ads.model.datascience_model_group import DataScienceModelGroup
 from ads.model.deployment.model_deployment import ModelDeployment
 from ads.model.model_metadata import ModelCustomMetadata
 from ads.model.service.oci_datascience_model_deployment import (
@@ -525,6 +521,7 @@ class TestDataset:
         "created_by": "ocid1.user.oc1..<OCID>",
         "endpoint": MODEL_DEPLOYMENT_URL,
         "private_endpoint_id": None,
+        "subnet_id": None,
         "models": [],
         "model_id": "ocid1.datasciencemodel.oc1.<region>.<OCID>",
         "environment_variables": {
@@ -555,6 +552,7 @@ class TestDataset:
         "created_by": "ocid1.user.oc1..<OCID>",
         "endpoint": MODEL_DEPLOYMENT_URL,
         "private_endpoint_id": None,
+        "subnet_id": None,
         "models": [
             {
                 "env_var": {},