Added custom egress for model deployment (#149)

Author: lu-ohai

ads/model/deployment/model_deployment.py

@@ -1538,6 +1538,9 @@ def _build_model_deployment_configuration_details(self) -> Dict:
                 or MODEL_DEPLOYMENT_INSTANCE_MEMORY_IN_GBS,
             }
 
+        if infrastructure.subnet_id:
+            instance_configuration[infrastructure.CONST_SUBNET_ID] = infrastructure.subnet_id
+
         scaling_policy = {
             infrastructure.CONST_POLICY_TYPE: "FIXED_SIZE",
             infrastructure.CONST_INSTANCE_COUNT: infrastructure.replica

ads/model/deployment/model_deployment_infrastructure.py

@@ -39,6 +39,8 @@ class ModelDeploymentInfrastructure(Builder):
         The access and predict log id of model deployment
     web_concurrency: int
         The web concurrency of model deployment
+    subnet_id: str
+        The subnet id of model deployment
 
     Methods
     -------
@@ -64,6 +66,8 @@ class ModelDeploymentInfrastructure(Builder):
         Sets the access and predict log id of model deployment
     with_web_concurrency(web_concurrency)
         Sets the web concurrency of model deployment
+    with_subnet_id(subnet_id) 
+        Sets the subnet id of model deployment
 
     Example
     -------
@@ -79,6 +83,7 @@ class ModelDeploymentInfrastructure(Builder):
     ...        .with_replica(1)
     ...        .with_bandwidth_mbps(10)
     ...        .with_web_concurrency(10)
+    ...        .with_subnet_id(<subnet_id>)
     ...        .with_access_log(
     ...            log_group_id=<log_group_id>,
     ...            log_id=<log_id>
@@ -121,6 +126,7 @@ class ModelDeploymentInfrastructure(Builder):
     CONST_LOG_GROUP_ID = "logGroupId"
     CONST_WEB_CONCURRENCY = "webConcurrency"
     CONST_STREAM_CONFIG_DETAILS = "streamConfigurationDetails"
+    CONST_SUBNET_ID = "subnetId"
 
     attribute_map = {
         CONST_PROJECT_ID: "project_id",
@@ -136,6 +142,7 @@ class ModelDeploymentInfrastructure(Builder):
         CONST_LOG_ID: "log_id",
         CONST_LOG_GROUP_ID: "log_group_id",
         CONST_WEB_CONCURRENCY: "web_concurrency",
+        CONST_SUBNET_ID: "subnet_id"
     }
 
     shape_config_details_attribute_map = {
@@ -162,6 +169,7 @@ class ModelDeploymentInfrastructure(Builder):
         CONST_COMPARTMENT_ID: "compartment_id",
         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_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",
@@ -518,3 +526,29 @@ def with_web_concurrency(
             The ModelDeploymentInfrastructure instance (self).
         """
         return self.set_spec(self.CONST_WEB_CONCURRENCY, web_concurrency)
+
+    def with_subnet_id(self, subnet_id: str) -> "ModelDeploymentInfrastructure":
+        """Sets the subnet id of model deployment.
+
+        Parameters
+        ----------
+        subnet_id : str
+            The subnet id of model deployment.
+
+        Returns
+        -------
+        ModelDeploymentInfrastructure
+            The ModelDeploymentInfrastructure instance (self).
+        """
+        return self.set_spec(self.CONST_SUBNET_ID, subnet_id)
+
+    @property
+    def subnet_id(self) -> str:
+        """The model deployment subnet id.
+        
+        Returns
+        -------
+        str
+            The model deployment subnet id.
+        """
+        return self.get_spec(self.CONST_SUBNET_ID, None)

ads/model/generic_model.py

@@ -1902,6 +1902,7 @@ def deploy(
         display_name: Optional[str] = None,
         description: Optional[str] = None,
         deployment_instance_shape: Optional[str] = None,
+        deployment_instance_subnet_id: Optional[str] = None,
         deployment_instance_count: Optional[int] = None,
         deployment_bandwidth_mbps: Optional[int] = None,
         deployment_log_group_id: Optional[str] = None,
@@ -1950,6 +1951,8 @@ def deploy(
             The description of the model.
         deployment_instance_shape: (str, optional). Default to `VM.Standard2.1`.
             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_count: (int, optional). Defaults to 1.
             The number of instance used for deployment.
         deployment_bandwidth_mbps: (int, optional). Defaults to 10.
@@ -2084,6 +2087,10 @@ def deploy(
                 self.properties.deployment_instance_shape
                 or existing_infrastructure.shape_name
             )
+            .with_subnet_id(
+                self.properties.deployment_instance_subnet_id
+                or existing_infrastructure.subnet_id
+            )
             .with_replica(
                 self.properties.deployment_instance_count
                 or existing_infrastructure.replica
@@ -2253,6 +2260,7 @@ def prepare_save_deploy(
         deployment_display_name: Optional[str] = None,
         deployment_description: Optional[str] = None,
         deployment_instance_shape: Optional[str] = None,
+        deployment_instance_subnet_id: Optional[str] = None,
         deployment_instance_count: Optional[int] = None,
         deployment_bandwidth_mbps: Optional[int] = None,
         deployment_log_group_id: Optional[str] = None,
@@ -2340,6 +2348,8 @@ def prepare_save_deploy(
             The description of the model.
         deployment_instance_shape: (str, optional). Default to `VM.Standard2.1`.
             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_count: (int, optional). Defaults to 1.
             The number of instance used for deployment.
         deployment_bandwidth_mbps: (int, optional). Defaults to 10.
@@ -2481,6 +2491,7 @@ def prepare_save_deploy(
             display_name=deployment_display_name,
             description=deployment_description,
             deployment_instance_shape=self.properties.deployment_instance_shape,
+            deployment_instance_subnet_id=self.properties.deployment_instance_subnet_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

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