Accept container family as input for model creation and param validation (#862)

Author: VipulMascarenhas

ads/aqua/common/enums.py

@@ -40,6 +40,25 @@ class Tags(str, metaclass=ExtendedEnumMeta):
     AQUA_EVALUATION_MODEL_ID = "evaluation_model_id"
 
 
+class InferenceContainerType(str, metaclass=ExtendedEnumMeta):
+    CONTAINER_TYPE_VLLM = "vllm"
+    CONTAINER_TYPE_TGI = "tgi"
+
+
+class InferenceContainerTypeFamily(str, metaclass=ExtendedEnumMeta):
+    AQUA_VLLM_CONTAINER_FAMILY = "odsc-vllm-serving"
+    AQUA_TGI_CONTAINER_FAMILY = "odsc-tgi-serving"
+
+
+class InferenceContainerParamType(str, metaclass=ExtendedEnumMeta):
+    PARAM_TYPE_VLLM = "VLLM_PARAMS"
+    PARAM_TYPE_TGI = "TGI_PARAMS"
+
+
+class HuggingFaceTags(str, metaclass=ExtendedEnumMeta):
+    TEXT_GENERATION_INFERENCE = "text-generation-inference"
+
+
 class RqsAdditionalDetails(str, metaclass=ExtendedEnumMeta):
     METADATA = "metadata"
     CREATED_BY = "createdBy"

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

@@ -13,7 +13,7 @@
 
 from ads.aqua import ODSC_MODEL_COMPARTMENT_OCID
 from ads.aqua.app import AquaApp
-from ads.aqua.common.enums import Tags
+from ads.aqua.common.enums import Tags, HuggingFaceTags, InferenceContainerTypeFamily
 from ads.aqua.common.errors import AquaRuntimeError
 from ads.aqua.common.utils import (
     create_word_icon,
@@ -621,7 +621,21 @@ 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 = (
+                    InferenceContainerTypeFamily.AQUA_TGI_CONTAINER_FAMILY
+                    if model_info
+                    and model_info.tags
+                    and HuggingFaceTags.TEXT_GENERATION_INFERENCE in model_info.tags
+                    else InferenceContainerTypeFamily.AQUA_VLLM_CONTAINER_FAMILY
+                )
+                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,
@@ -774,7 +788,7 @@ def register(
                 break
         if i == retry:
             raise Exception(
-                "Could not download the model {model_name} from https://huggingface.co with message {huggingface_download}"
+                f"Could not download the model {model_name} from https://huggingface.co with message {huggingface_download_err_message}"
             )
         os.makedirs(local_dir, exist_ok=True)
         # Copy the model from the cache to destination

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

@@ -10,7 +10,12 @@
 from oci.data_science.models import ModelDeployment
 
 from ads.aqua.app import AquaApp, logger
-from ads.aqua.common.enums import Tags
+from ads.aqua.common.enums import (
+    Tags,
+    InferenceContainerParamType,
+    InferenceContainerType,
+    InferenceContainerTypeFamily,
+)
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
 from ads.aqua.common.utils import (
     get_container_config,
@@ -38,11 +43,9 @@
     AquaDeploymentDetail,
     ContainerSpec,
 )
-from ads.aqua.modeldeployment.constants import VLLMInferenceRestrictedParams
-from ads.aqua.modeldeployment.enums import (
-    InferenceContainerParamType,
-    InferenceContainerType,
-    InferenceContainerTypeKey,
+from ads.aqua.modeldeployment.constants import (
+    VLLMInferenceRestrictedParams,
+    TGIInferenceRestrictedParams,
 )
 from ads.common.object_storage_details import ObjectStorageDetails
 from ads.common.utils import get_log_links
@@ -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(InferenceContainerTypeFamily.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 "
@@ -559,7 +572,7 @@ def get_deployment_default_params(
 
         if container_type_key:
             container_type_key = container_type_key.lower()
-            if container_type_key in InferenceContainerTypeKey.values():
+            if container_type_key in InferenceContainerTypeFamily.values():
                 deployment_config = self.get_deployment_config(model_id)
                 config_parameters = (
                     deployment_config.get("configuration", UNKNOWN_DICT)
@@ -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(InferenceContainerTypeFamily.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

ads/aqua/modeldeployment/enums.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):