From 81dac4e4623e0bea6e6ce1cdd6c213dc6e1d8989 Mon Sep 17 00:00:00 2001
From: Dmitrii Cherkasov <dmitrii.cherkasov@oracle.com>
Date: Mon, 10 Nov 2025 20:38:49 -0800
Subject: [PATCH] Adds supporting custom networking in AQUA MD

---
 ads/aqua/modeldeployment/deployment.py            | 2 ++
 ads/aqua/modeldeployment/entities.py              | 8 ++++++++
 ads/aqua/server/aqua_spec.yml                     | 7 +++++++
 tests/unitary/with_extras/aqua/test_deployment.py | 8 +++-----
 4 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/ads/aqua/modeldeployment/deployment.py b/ads/aqua/modeldeployment/deployment.py
index 214d98d1f..e5d4d1fb5 100644
--- a/ads/aqua/modeldeployment/deployment.py
+++ b/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
diff --git a/ads/aqua/modeldeployment/entities.py b/ads/aqua/modeldeployment/entities.py
index a80bc8142..206caeb03 100644
--- a/ads/aqua/modeldeployment/entities.py
+++ b/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
diff --git a/ads/aqua/server/aqua_spec.yml b/ads/aqua/server/aqua_spec.yml
index 672d84507..3523086d6 100644
--- a/ads/aqua/server/aqua_spec.yml
+++ b/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>"
diff --git a/tests/unitary/with_extras/aqua/test_deployment.py b/tests/unitary/with_extras/aqua/test_deployment.py
index 0d5b1016e..2a22a7f42 100644
--- a/tests/unitary/with_extras/aqua/test_deployment.py
+++ b/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": {},