register and deployment changes for byoc

Author: VipulMascarenhas

ads/aqua/common/enums.py

@@ -52,6 +52,7 @@ class InferenceContainerTypeFamily(str, metaclass=ExtendedEnumMeta):
     AQUA_VLLM_CONTAINER_FAMILY = "odsc-vllm-serving"
     AQUA_TGI_CONTAINER_FAMILY = "odsc-tgi-serving"
     AQUA_LLAMA_CPP_CONTAINER_FAMILY = "odsc-llama-cpp-serving"
+    AQUA_TEI_CONTAINER_FAMILY = "odsc-tei-serving"
 
 
 class InferenceContainerParamType(str, metaclass=ExtendedEnumMeta):
@@ -80,3 +81,11 @@ class RqsAdditionalDetails(str, metaclass=ExtendedEnumMeta):
     MODEL_VERSION_SET_NAME = "modelVersionSetName"
     PROJECT_ID = "projectId"
     VERSION_LABEL = "versionLabel"
+
+
+class TextEmbeddingInferenceContainerParams(str, metaclass=ExtendedEnumMeta):
+    """Contains a subset of params that are required for enabling model deployment in OCI Data Science. More options
+    are available at https://huggingface.co/docs/text-embeddings-inference/en/cli_arguments"""
+
+    MODEL_ID = "model-id"
+    PORT = "port"

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")
@@ -1079,3 +1086,27 @@ 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_vars(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
+    -------
+        List of command line arguments
+    """
+
+    cmd_prefix = "--"
+    cmd_vars = [
+        cmd_prefix + TextEmbeddingInferenceContainerParams.MODEL_ID,
+        f"{AQUA_MODEL_DEPLOYMENT_FOLDER}{ObjectStorageDetails.from_path(os_path.rstrip('/')).filepath}/",
+        cmd_prefix + TextEmbeddingInferenceContainerParams.PORT,
+        TEI_CONTAINER_DEFAULT_HOST,
+    ]
+
+    return cmd_vars

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_vars,
     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,
     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_vars(os_path)
+                metadata.add(
+                    key=AQUA_DEPLOYMENT_CONTAINER_CMD_VAR,
+                    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(
@@ -706,12 +741,6 @@ def _create_model_catalog_entry(
                     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 +1041,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,6 +1225,7 @@ 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

ads/aqua/modeldeployment/deployment.py

@@ -46,6 +46,7 @@
 from ads.config import (
     AQUA_CONFIG_FOLDER,
     AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME,
+    AQUA_DEPLOYMENT_CONTAINER_URI_METADATA_NAME,
     AQUA_MODEL_DEPLOYMENT_CONFIG,
     AQUA_MODEL_DEPLOYMENT_CONFIG_DEFAULTS,
     COMPARTMENT_OCID,
@@ -106,6 +107,7 @@ def create(
         memory_in_gbs: Optional[float] = None,
         ocpus: Optional[float] = None,
         model_file: Optional[str] = None,
+        container_image_uri: Optional[None] = None,
         cmd_var: List[str] = None,
     ) -> "AquaDeployment":
         """
@@ -153,6 +155,9 @@ def create(
             The ocpu count for the shape selected.
         model_file: str
             The file used for 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
@@ -198,9 +203,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(
@@ -236,11 +243,37 @@ def create(
             model=aqua_model, container_family=container_family
         )
 
-        # fetch image name from config
-        container_image = get_container_image(container_type=container_type_key)
+        # todo: revisit this when TEI is added to SMC list. Currently, container_image_uri is ignored if container
+        #   family is SMC.
+        if container_type_key == InferenceContainerTypeFamily.AQUA_TEI_CONTAINER_FAMILY:
+            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
+
+            try:
+                cmd_var_string = aqua_model.custom_metadata_list.get(
+                    AQUA_DEPLOYMENT_CONTAINER_URI_METADATA_NAME
+                ).value
+                cmd_var.append(cmd_var_string.split(","))
+            except ValueError as err:
+                raise AquaValueError(
+                    f"{AQUA_DEPLOYMENT_CONTAINER_URI_METADATA_NAME} key is not available in the custom metadata "
+                    f"field. Please check if the model was registered with {container_type_key} inference container."
+                ) from err
+        else:
+            # fetch image name from config
+            container_image_uri = get_container_image(container_type=container_type_key)
 
         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}"
         )
 
         model_formats_str = aqua_model.freeform_tags.get(
@@ -310,7 +343,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."
                 )
 
@@ -364,7 +397,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)

ads/config.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-# -*- coding: utf-8; -*-
 
 # Copyright (c) 2020, 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
@@ -66,6 +65,8 @@
 AQUA_DEPLOYMENT_CONTAINER_METADATA_NAME = "deployment-container"
 AQUA_FINETUNING_CONTAINER_METADATA_NAME = "finetune-container"
 AQUA_EVALUATION_CONTAINER_METADATA_NAME = "evaluation-container"
+AQUA_DEPLOYMENT_CONTAINER_CMD_VAR = "container_cmd_var"
+AQUA_DEPLOYMENT_CONTAINER_URI_METADATA_NAME = "deployment-container-uri"
 AQUA_DEPLOYMENT_CONTAINER_OVERRIDE_FLAG_METADATA_NAME = "deployment-container-custom"
 AQUA_FINETUNING_CONTAINER_OVERRIDE_FLAG_METADATA_NAME = "finetune-container-custom"
 AQUA_MODEL_DEPLOYMENT_FOLDER = "/opt/ds/model/deployed_model/"
@@ -212,7 +213,7 @@ def open(
         frame.f_globals.pop("config", None)
 
         # Restores original globals
-        for key in defined_globals.keys():
+        for key in defined_globals:
             frame.f_globals[key] = defined_globals[key]
 
         # Saving config if it necessary