support tei deployment

Author: VipulMascarenhas

ads/aqua/common/utils.py

@@ -35,6 +35,7 @@
     InferenceContainerParamType,
     InferenceContainerType,
     RqsAdditionalDetails,
+    TextEmbeddingInferenceContainerParams,
 )
 from ads.aqua.common.errors import (
     AquaFileNotFoundError,
@@ -51,6 +52,7 @@
     MODEL_BY_REFERENCE_OSS_PATH_KEY,
     SERVICE_MANAGED_CONTAINER_URI_SCHEME,
     SUPPORTED_FILE_FORMATS,
+    TEI_CONTAINER_DEFAULT_HOST,
     TGI_INFERENCE_RESTRICTED_PARAMS,
     UNKNOWN,
     UNKNOWN_JSON_STR,
@@ -63,7 +65,12 @@
 from ads.common.object_storage_details import ObjectStorageDetails
 from ads.common.oci_resource import SEARCH_TYPE, OCIResource
 from ads.common.utils import copy_file, get_console_link, upload_to_os
-from ads.config import AQUA_SERVICE_MODELS_BUCKET, CONDA_BUCKET_NS, TENANCY_OCID
+from ads.config import (
+    AQUA_MODEL_DEPLOYMENT_FOLDER,
+    AQUA_SERVICE_MODELS_BUCKET,
+    CONDA_BUCKET_NS,
+    TENANCY_OCID,
+)
 from ads.model import DataScienceModel, ModelVersionSet
 
 logger = logging.getLogger("ads.aqua")
@@ -569,15 +576,13 @@ def get_container_image(
         A dict of allowed configs.
     """
 
+    container_image = UNKNOWN
     config = config_file_name or get_container_config()
     config_file_name = service_config_path()
 
     if container_type not in config:
-        raise AquaValueError(
-            f"{config_file_name} does not have config details for model: {container_type}"
-        )
+        return UNKNOWN
 
-    container_image = None
     mapping = config[container_type]
     versions = [obj["version"] for obj in mapping]
     # assumes numbered versions, update if `latest` is used
@@ -1078,3 +1083,71 @@ def list_hf_models(query: str) -> List[str]:
         return [model.id for model in models if model.disabled is None]
     except HfHubHTTPError as err:
         raise format_hf_custom_error_message(err) from err
+
+
+def generate_tei_cmd_var(os_path: str) -> List[str]:
+    """This utility functions generates CMD params for Text Embedding Inference container. Only the
+    essential parameters for OCI model deployment are added, defaults are used for the rest.
+    Parameters
+    ----------
+    os_path: str
+        OCI bucket path where the model artifacts are uploaded - oci://bucket@namespace/prefix
+    Returns
+    -------
+        cmd_var:
+            List of command line arguments
+    """
+
+    cmd_prefix = "--"
+    cmd_var = [
+        f"{cmd_prefix}{TextEmbeddingInferenceContainerParams.MODEL_ID}",
+        f"{AQUA_MODEL_DEPLOYMENT_FOLDER}{ObjectStorageDetails.from_path(os_path.rstrip('/')).filepath}/",
+        f"{cmd_prefix}{TextEmbeddingInferenceContainerParams.PORT}",
+        TEI_CONTAINER_DEFAULT_HOST,
+    ]
+
+    return cmd_var
+
+
+def parse_cmd_var(cmd_list: List[str]) -> dict:
+    """Helper functions that parses a list into a key-value dictionary. The list contains keys separated by the prefix
+    '--' and the value of the key is the subsequent element.
+    """
+    it = iter(cmd_list)
+    return {
+        key: next(it, None) if not key.startswith("--") else next(it, None)
+        for key in it
+        if key.startswith("--")
+    }
+
+
+def validate_cmd_var(cmd_var: List[str], overrides: List[str]) -> List[str]:
+    """This function accepts two lists of parameters and combines them. If the second list shares the common parameter
+    names/keys, then it raises an error.
+    Parameters
+    ----------
+    cmd_var: List[str]
+        Default list of parameters
+    overrides: List[str]
+        List of parameters to override
+    Returns
+    -------
+        List[str] of combined parameters
+    """
+    cmd_var = [str(x) for x in cmd_var]
+    if not overrides:
+        return cmd_var
+    overrides = [str(x) for x in overrides]
+
+    cmd_dict = parse_cmd_var(cmd_var)
+    overrides_dict = parse_cmd_var(overrides)
+
+    # check for conflicts
+    common_keys = set(cmd_dict.keys()) & set(overrides_dict.keys())
+    if common_keys:
+        raise AquaValueError(
+            f"The following CMD input cannot be overridden for model deployment: {', '.join(common_keys)}"
+        )
+
+    combined_cmd_var = cmd_var + overrides
+    return combined_cmd_var

ads/aqua/constants.py

@@ -79,3 +79,4 @@
     "--port",
     "--host",
 }
+TEI_CONTAINER_DEFAULT_HOST = "8080"

ads/aqua/model/entities.py

@@ -287,6 +287,7 @@ class ImportModelDetails(CLIBuilderMixin):
     compartment_id: Optional[str] = None
     project_id: Optional[str] = None
     model_file: Optional[str] = None
+    inference_container_uri: Optional[str] = None
 
     def __post_init__(self):
         self._command = "model register"

ads/aqua/model/model.py

@@ -14,13 +14,14 @@
 
 from ads.aqua import ODSC_MODEL_COMPARTMENT_OCID, logger
 from ads.aqua.app import AquaApp
-from ads.aqua.common.enums import Tags
+from ads.aqua.common.enums import InferenceContainerTypeFamily, Tags
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
 from ads.aqua.common.utils import (
     LifecycleStatus,
     _build_resource_identifier,
     copy_model_config,
     create_word_icon,
+    generate_tei_cmd_var,
     get_artifact_path,
     get_hf_model_info,
     list_os_files_with_extension,
@@ -67,7 +68,9 @@
 from ads.common.oci_resource import SEARCH_TYPE, OCIResource
 from ads.common.utils import get_console_link
 from ads.config import (
+    AQUA_DEPLOYMENT_CONTAINER_CMD_VAR_METADATA_NAME,
     AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME,
+    AQUA_DEPLOYMENT_CONTAINER_URI_METADATA_NAME,
     AQUA_EVALUATION_CONTAINER_METADATA_NAME,
     AQUA_FINETUNING_CONTAINER_METADATA_NAME,
     COMPARTMENT_OCID,
@@ -629,6 +632,7 @@ def _create_model_catalog_entry(
         validation_result: ModelValidationResult,
         compartment_id: Optional[str],
         project_id: Optional[str],
+        inference_container_uri: Optional[str],
     ) -> DataScienceModel:
         """Create model by reference from the object storage path
 
@@ -640,6 +644,7 @@ def _create_model_catalog_entry(
             verified_model (DataScienceModel): If set, then copies all the tags and custom metadata information from the service verified model
             compartment_id (Optional[str]): Compartment Id of the compartment where the model has to be created
             project_id (Optional[str]): Project id of the project where the model has to be created
+            inference_container_uri (Optional[str]): Inference container uri for BYOC
 
         Returns:
             DataScienceModel: Returns Datascience model instance.
@@ -685,6 +690,36 @@ def _create_model_catalog_entry(
                 raise AquaRuntimeError(
                     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."
                 )
+            metadata.add(
+                key=AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME,
+                value=inference_container,
+                description=f"Inference container mapping for {model_name}",
+                category="Other",
+            )
+            if (
+                inference_container
+                == InferenceContainerTypeFamily.AQUA_TEI_CONTAINER_FAMILY
+            ):
+                if not inference_container_uri:
+                    logger.warn(
+                        f"Proceeding with model registration without the inference container URI for "
+                        f"{inference_container}. You can still add this configuration during model deployment."
+                    )
+                else:
+                    metadata.add(
+                        key=AQUA_DEPLOYMENT_CONTAINER_URI_METADATA_NAME,
+                        value=inference_container_uri,
+                        description=f"Inference container URI for {model_name}",
+                        category="Other",
+                    )
+
+                cmd_vars = generate_tei_cmd_var(os_path)
+                metadata.add(
+                    key=AQUA_DEPLOYMENT_CONTAINER_CMD_VAR_METADATA_NAME,
+                    value=",".join(cmd_vars),
+                    description=f"Inference container cmd vars for {model_name}",
+                    category="Other",
+                )
             if finetuning_container:
                 tags[Tags.READY_TO_FINE_TUNE] = "true"
                 metadata.add(
@@ -705,13 +740,6 @@ def _create_model_catalog_entry(
                     description=f"The model file for {model_name}",
                     category="Other",
                 )
-
-            metadata.add(
-                key=AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME,
-                value=inference_container,
-                description=f"Inference container mapping for {model_name}",
-                category="Other",
-            )
             metadata.add(
                 key=AQUA_EVALUATION_CONTAINER_METADATA_NAME,
                 value="odsc-llm-evaluate",
@@ -1012,7 +1040,10 @@ def _validate_model(
                                 AQUA_MODEL_TYPE_CUSTOM
                             )
             elif model_format == ModelFormat.GGUF and len(gguf_model_files) > 0:
-                if import_model_details.finetuning_container and not safetensors_model_files:
+                if (
+                    import_model_details.finetuning_container
+                    and not safetensors_model_files
+                ):
                     raise AquaValueError(
                         "Fine-tuning is currently not supported with GGUF model format."
                     )
@@ -1193,15 +1224,19 @@ def register(
             validation_result=validation_result,
             compartment_id=import_model_details.compartment_id,
             project_id=import_model_details.project_id,
+            inference_container_uri=import_model_details.inference_container_uri,
         )
         # registered model will always have inference and evaluation container, but
         # fine-tuning container may be not set
         inference_container = ds_model.custom_metadata_list.get(
             ModelCustomMetadataFields.DEPLOYMENT_CONTAINER
         ).value
-        evaluation_container = ds_model.custom_metadata_list.get(
-            ModelCustomMetadataFields.EVALUATION_CONTAINER,
-        ).value
+        try:
+            evaluation_container = ds_model.custom_metadata_list.get(
+                ModelCustomMetadataFields.EVALUATION_CONTAINER,
+            ).value
+        except Exception:
+            evaluation_container = None
         try:
             finetuning_container = ds_model.custom_metadata_list.get(
                 ModelCustomMetadataFields.FINETUNE_CONTAINER,

ads/aqua/modeldeployment/deployment.py

@@ -24,6 +24,7 @@
     get_resource_name,
     get_restricted_params_by_container,
     load_config,
+    validate_cmd_var,
 )
 from ads.aqua.constants import (
     AQUA_MODEL_ARTIFACT_FILE,
@@ -45,7 +46,9 @@
 from ads.common.utils import get_log_links
 from ads.config import (
     AQUA_CONFIG_FOLDER,
+    AQUA_DEPLOYMENT_CONTAINER_CMD_VAR_METADATA_NAME,
     AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME,
+    AQUA_DEPLOYMENT_CONTAINER_URI_METADATA_NAME,
     AQUA_MODEL_DEPLOYMENT_CONFIG,
     AQUA_MODEL_DEPLOYMENT_CONFIG_DEFAULTS,
     COMPARTMENT_OCID,
@@ -107,6 +110,8 @@ def create(
         ocpus: Optional[float] = None,
         model_file: Optional[str] = None,
         private_endpoint_id: Optional[str] = None,
+        container_image_uri: Optional[None] = None,
+        cmd_var: List[str] = None,
     ) -> "AquaDeployment":
         """
         Creates a new Aqua deployment
@@ -155,6 +160,11 @@ def create(
             The file used for model deployment.
         private_endpoint_id: str
             The private endpoint id of model deployment.
+        container_image_uri: str
+            The image of model deployment container runtime, ignored for service managed containers.
+            Required parameter for BYOC based deployments if this parameter was not set during model registration.
+        cmd_var: List[str]
+            The cmd of model deployment container runtime.
 
         Returns
         -------
@@ -200,9 +210,11 @@ def create(
                     f"from custom metadata for the model {config_source_id}"
                 ) from err
 
-        # set up env vars
+        # set up env and cmd var
         if not env_var:
             env_var = {}
+        if not cmd_var:
+            cmd_var = []
 
         try:
             model_path_prefix = aqua_model.custom_metadata_list.get(
@@ -238,13 +250,44 @@ def create(
             model=aqua_model, container_family=container_family
         )
 
-        # fetch image name from config
-        container_image = get_container_image(container_type=container_type_key)
-
+        # fetch image name
+        container_image_uri = container_image_uri or get_container_image(
+            container_type=container_type_key
+        )
+        if not container_image_uri:
+            try:
+                container_image_uri = aqua_model.custom_metadata_list.get(
+                    AQUA_DEPLOYMENT_CONTAINER_URI_METADATA_NAME
+                ).value
+            except ValueError as err:
+                raise AquaValueError(
+                    f"{AQUA_DEPLOYMENT_CONTAINER_URI_METADATA_NAME} key is not available in the custom metadata "
+                    f"field. Either re-register the model with custom container URI, or set container_image_uri "
+                    f"parameter when creating this deployment."
+                ) from err
         logging.info(
-            f"Aqua Image used for deploying {aqua_model.id} : {container_image}"
+            f"Aqua Image used for deploying {aqua_model.id} : {container_image_uri}"
         )
 
+        # set up cmd var if applicable
+        try:
+            cmd_var_string = aqua_model.custom_metadata_list.get(
+                AQUA_DEPLOYMENT_CONTAINER_CMD_VAR_METADATA_NAME
+            ).value
+            default_cmd_var = cmd_var_string.split(",")
+            if default_cmd_var:
+                cmd_var = validate_cmd_var(default_cmd_var, cmd_var)
+            logging.info(f"CMD used for deploying {aqua_model.id} :{cmd_var}")
+        except ValueError:
+            logging.debug(
+                f"CMD will be ignored for this deployment as {AQUA_DEPLOYMENT_CONTAINER_CMD_VAR_METADATA_NAME} "
+                f"key is not available in the custom metadata field for this model."
+            )
+        except Exception as e:
+            logging.error(
+                f"There was an issue processing CMD arguments. Error: {str(e)}"
+            )
+
         model_formats_str = aqua_model.freeform_tags.get(
             Tags.MODEL_FORMAT, ModelFormat.SAFETENSORS.value
         ).upper()
@@ -312,7 +355,7 @@ def create(
                 # AQUA_LLAMA_CPP_CONTAINER_FAMILY container uses uvicorn that required model/server params
                 # to be set as env vars
                 raise AquaValueError(
-                    f"Currently, parameters cannot be overridden for the container: {container_image}. Please proceed "
+                    f"Currently, parameters cannot be overridden for the container: {container_image_uri}. Please proceed "
                     f"with deployment without parameter overrides."
                 )
 
@@ -367,7 +410,7 @@ def create(
         # configure model deployment runtime
         container_runtime = (
             ModelDeploymentContainerRuntime()
-            .with_image(container_image)
+            .with_image(container_image_uri)
             .with_server_port(server_port)
             .with_health_check_port(health_check_port)
             .with_env(env_var)
@@ -377,6 +420,8 @@ def create(
             .with_overwrite_existing_artifact(True)
             .with_remove_existing_artifact(True)
         )
+        if cmd_var:
+            container_runtime.with_cmd(cmd_var)
 
         # configure model deployment and deploy model on container runtime
         deployment = (