accept container family as input

Author: VipulMascarenhas

ads/aqua/extension/deployment_handler.py

@@ -99,6 +99,7 @@ def post(self, *args, **kwargs):
         server_port = input_data.get("server_port")
         health_check_port = input_data.get("health_check_port")
         env_var = input_data.get("env_var")
+        container_family = input_data.get("container_family")
 
         self.finish(
             AquaDeploymentApp().create(
@@ -117,6 +118,7 @@ def post(self, *args, **kwargs):
                 server_port=server_port,
                 health_check_port=health_check_port,
                 env_var=env_var,
+                container_family=container_family,
             )
         )
 
@@ -245,10 +247,12 @@ def post(self, *args, **kwargs):
             raise HTTPError(400, Errors.MISSING_REQUIRED_PARAMETER.format("model_id"))
 
         params = input_data.get("params")
+        container_family = input_data.get("container_family")
         return self.finish(
             AquaDeploymentApp().validate_deployment_params(
                 model_id=model_id,
                 params=params,
+                container_family=container_family,
             )
         )
 

ads/aqua/model/model.py

@@ -38,6 +38,7 @@
 )
 from ads.aqua.model.constants import *
 from ads.aqua.model.entities import *
+from ads.aqua.modeldeployment.enums import InferenceContainerTypeKey
 from ads.common.auth import default_signer
 from ads.common.oci_resource import SEARCH_TYPE, OCIResource
 from ads.common.utils import get_console_link
@@ -621,7 +622,15 @@ def _create_model_catalog_entry(
                 )
             else:
                 logger.warn(
-                    f"Require Inference container information. Model: {model_name} does not have associated inference container defaults. Check docs for more information on how to pass inference container. Proceeding with model registration without the fine-tuning container information. This model will not be available for fine tuning."
+                    f"Proceeding with model registration without the fine-tuning container information. "
+                    f"This model will not be available for fine tuning."
+                )
+
+            if not inference_container:
+                inference_container = InferenceContainerTypeKey.AQUA_TGI_CONTAINER_KEY
+                logger.info(
+                    f"Model: {model_name} does not have associated inference container defaults. "
+                    f"{inference_container} will be used instead."
                 )
             metadata.add(
                 key=AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME,

ads/aqua/modeldeployment/constants.py

@@ -10,4 +10,5 @@
 This module contains constants used in Aqua Model Deployment.
 """
 
-VLLMInferenceRestrictedParams = {"tensor-parallel-size"}
+VLLMInferenceRestrictedParams = {"--tensor-parallel-size"}
+TGIInferenceRestrictedParams = {"--port"}

ads/aqua/modeldeployment/deployment.py

@@ -38,7 +38,10 @@
     AquaDeploymentDetail,
     ContainerSpec,
 )
-from ads.aqua.modeldeployment.constants import VLLMInferenceRestrictedParams
+from ads.aqua.modeldeployment.constants import (
+    VLLMInferenceRestrictedParams,
+    TGIInferenceRestrictedParams,
+)
 from ads.aqua.modeldeployment.enums import (
     InferenceContainerParamType,
     InferenceContainerType,
@@ -106,6 +109,7 @@ def create(
         server_port: int = None,
         health_check_port: int = None,
         env_var: Dict = None,
+        container_family: str = None,
     ) -> "AquaDeployment":
         """
         Creates a new Aqua deployment
@@ -144,6 +148,8 @@ def create(
             The health check port for docker container image.
         env_var : dict, optional
             Environment variable for the deployment, by default None.
+        container_family: str
+            The image family of model deployment container runtime. Required for unverified Aqua models.
         Returns
         -------
         AquaDeployment
@@ -227,9 +233,17 @@ def create(
                 AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME
             ).value
         except ValueError:
-            raise AquaValueError(
-                f"{AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME} key is not available in the custom metadata field for model {aqua_model.id}"
+            message = (
+                f"{AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME} key is not available in the custom metadata field "
+                f"for model {aqua_model.id}."
             )
+            logger.debug(message)
+            if not container_family:
+                raise AquaValueError(
+                    f"{message}. For unverified Aqua models, container_family parameter should be "
+                    f"set and value can be one of {', '.join(InferenceContainerTypeKey.values())}."
+                )
+            container_type_key = container_family
         try:
             # Check if the container override flag is set. If set, then the user has chosen custom image
             if aqua_model.custom_metadata_list.get(
@@ -275,13 +289,12 @@ def create(
             .get(InferenceContainerParamType.PARAM_TYPE_VLLM, UNKNOWN)
         )
 
-        # todo: add support for tgi once parameters are added to configs. _find_restricted_params can take in
-        #   additional parameter container_type_key and should validate against TGIInferenceRestrictedParams set for
-        #   restricted params.
         # validate user provided params
         user_params = env_var.get("PARAMS", UNKNOWN)
         if user_params:
-            restricted_params = self._find_restricted_params(params, user_params)
+            restricted_params = self._find_restricted_params(
+                params, user_params, container_type_key
+            )
             if restricted_params:
                 raise AquaValueError(
                     f"Parameters {restricted_params} are set by Aqua "
@@ -587,7 +600,10 @@ def get_deployment_default_params(
         return default_params
 
     def validate_deployment_params(
-        self, model_id: str, params: List[str] = None
+        self,
+        model_id: str,
+        params: List[str] = None,
+        container_family: str = None,
     ) -> Dict:
         """Validate if the deployment parameters passed by the user can be overridden. Parameter values are not
         validated, only param keys are validated.
@@ -596,9 +612,10 @@ def validate_deployment_params(
         ----------
         model_id: str
             The OCID of the Aqua model.
-
         params : List[str], optional
             Params passed by the user.
+        container_family: str
+            The image family of model deployment container runtime. Required for unverified Aqua models.
 
         Returns
         -------
@@ -613,18 +630,28 @@ def validate_deployment_params(
                     AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME
                 ).value
             except ValueError:
-                container_type_key = UNKNOWN
-                logger.debug(
-                    f"{AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME} key is not available in the custom metadata field for model {model_id}."
+                message = (
+                    f"{AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME} key is not available in the custom metadata field "
+                    f"for model {model_id}."
                 )
-            if container_type_key:
-                container_config = get_container_config()
-                container_spec = container_config.get(
-                    ContainerSpec.CONTAINER_SPEC, {}
-                ).get(container_type_key, {})
-                cli_params = container_spec.get(ContainerSpec.CLI_PARM, "")
+                logger.debug(message)
 
-                restricted_params = self._find_restricted_params(cli_params, params)
+                if not container_family:
+                    raise AquaValueError(
+                        f"{message}. For unverified Aqua models, container_family parameter should be "
+                        f"set and value can be one of {', '.join(InferenceContainerTypeKey.values())}."
+                    )
+                container_type_key = container_family
+
+            container_config = get_container_config()
+            container_spec = container_config.get(ContainerSpec.CONTAINER_SPEC, {}).get(
+                container_type_key, {}
+            )
+            cli_params = container_spec.get(ContainerSpec.CLI_PARM, "")
+
+            restricted_params = self._find_restricted_params(
+                cli_params, params, container_type_key
+            )
 
         if restricted_params:
             raise AquaValueError(
@@ -635,7 +662,9 @@ def validate_deployment_params(
 
     @staticmethod
     def _find_restricted_params(
-        default_params: Union[str, List[str]], user_params: Union[str, List[str]]
+        default_params: Union[str, List[str]],
+        user_params: Union[str, List[str]],
+        container_family: str,
     ) -> List[str]:
         """Returns a list of restricted params that user chooses to override when creating an Aqua deployment.
         The default parameters coming from the container index json file cannot be overridden. In addition to this,
@@ -647,6 +676,8 @@ def _find_restricted_params(
             Inference container parameter string with default values.
         user_params:
             Inference container parameter string with user provided values.
+        container_family: str
+            The image family of model deployment container runtime.
 
         Returns
         -------
@@ -659,7 +690,17 @@ def _find_restricted_params(
             user_params_dict = get_params_dict(user_params)
 
             for key, items in user_params_dict.items():
-                if key in default_params_dict or key in VLLMInferenceRestrictedParams:
+                if (
+                    key in default_params_dict
+                    or (
+                        InferenceContainerType.CONTAINER_TYPE_VLLM in container_family
+                        and key in VLLMInferenceRestrictedParams
+                    )
+                    or (
+                        InferenceContainerType.CONTAINER_TYPE_TGI in container_family
+                        and key in TGIInferenceRestrictedParams
+                    )
+                ):
                     restricted_params.append(key.lstrip("--"))
 
         return restricted_params

tests/unitary/with_extras/aqua/test_deployment.py

@@ -494,6 +494,14 @@ def test_get_deployment_default_params(
                 "custom-container-key",
                 ["--max-model-len 4096", "--seed 42", "--trust-remote-code"],
             ),
+            (
+                "odsc-vllm-serving",
+                ["--tensor-parallel-size 2"],
+            ),
+            (
+                "odsc-tgi-serving",
+                ["--port 8080"],
+            ),
         ]
     )
     @patch("ads.model.datascience_model.DataScienceModel.from_id")
@@ -530,6 +538,55 @@ def test_validate_deployment_params(
             )
             assert result["valid"] is True
 
+    @parameterized.expand(
+        [
+            (
+                "odsc-vllm-serving",
+                ["--max-model-len 4096"],
+            ),
+            (
+                "odsc-tgi-serving",
+                ["--max_stop_sequences 5"],
+            ),
+            (
+                "",
+                ["--some_random_key some_random_value"],
+            ),
+        ]
+    )
+    @patch("ads.model.datascience_model.DataScienceModel.from_id")
+    @patch("ads.aqua.modeldeployment.deployment.get_container_config")
+    def test_validate_deployment_params_for_unverified_models(
+        self, container_type_key, params, mock_get_container_config, mock_from_id
+    ):
+        """Test to check if container family is used when metadata does not have image information
+        for unverified models."""
+        mock_model = MagicMock()
+        mock_model.custom_metadata_list = ModelCustomMetadata()
+        mock_from_id.return_value = mock_model
+
+        container_index_json = os.path.join(
+            self.curr_dir, "test_data/ui/container_index.json"
+        )
+        with open(container_index_json, "r") as _file:
+            container_index_config = json.load(_file)
+        mock_get_container_config.return_value = container_index_config
+
+        if container_type_key in {"odsc-vllm-serving", "odsc-tgi-serving"} and params:
+            result = self.app.validate_deployment_params(
+                model_id="mock-model-id",
+                params=params,
+                container_family=container_type_key,
+            )
+            assert result["valid"] is True
+        else:
+            with pytest.raises(AquaValueError):
+                self.app.validate_deployment_params(
+                    model_id="mock-model-id",
+                    params=params,
+                    container_family=container_type_key,
+                )
+
 
 class TestMDInferenceResponse(unittest.TestCase):
     def setUp(self):

tests/unitary/with_extras/aqua/test_deployment_handler.py

@@ -8,6 +8,7 @@
 import unittest
 from importlib import reload
 from unittest.mock import MagicMock, patch
+from parameterized import parameterized
 
 from notebook.base.handlers import IPythonHandler
 
@@ -124,6 +125,7 @@ def test_post(self, mock_create):
             server_port=None,
             health_check_port=None,
             env_var=None,
+            container_family=None,
         )
 
 
@@ -156,21 +158,33 @@ def test_get_deployment_default_params(
             model_id="test_model_id", instance_shape=TestDataset.INSTANCE_SHAPE
         )
 
+    @parameterized.expand(
+        [
+            None,
+            "container-family-name",
+        ]
+    )
     @patch("notebook.base.handlers.APIHandler.finish")
     @patch("ads.aqua.modeldeployment.AquaDeploymentApp.validate_deployment_params")
     def test_validate_deployment_params(
-        self, mock_validate_deployment_params, mock_finish
+        self, container_family_value, mock_validate_deployment_params, mock_finish
     ):
         mock_validate_deployment_params.return_value = dict(valid=True)
         mock_finish.side_effect = lambda x: x
 
         self.test_instance.get_json_body = MagicMock(
-            return_value=dict(model_id="test-model-id", params=self.default_params)
+            return_value=dict(
+                model_id="test-model-id",
+                params=self.default_params,
+                container_family=container_family_value,
+            )
         )
         result = self.test_instance.post()
         assert result["valid"] is True
         mock_validate_deployment_params.assert_called_with(
-            model_id="test-model-id", params=self.default_params
+            model_id="test-model-id",
+            params=self.default_params,
+            container_family=container_family_value,
         )