[ODSC-60107] Return container spec with list containers API (#906)

Author: VipulMascarenhas

ads/aqua/common/entities.py

@@ -0,0 +1,16 @@
+#!/usr/bin/env python
+# Copyright (c) 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+
+class ContainerSpec:
+    """
+    Class to hold to hold keys within the container spec.
+    """
+
+    CONTAINER_SPEC = "containerSpec"
+    CLI_PARM = "cliParam"
+    SERVER_PORT = "serverPort"
+    HEALTH_CHECK_PORT = "healthCheckPort"
+    ENV_VARS = "envVars"
+    RESTRICTED_PARAMS = "restrictedParams"

ads/aqua/common/enums.py

@@ -44,6 +44,7 @@ class Tags(str, metaclass=ExtendedEnumMeta):
 class InferenceContainerType(str, metaclass=ExtendedEnumMeta):
     CONTAINER_TYPE_VLLM = "vllm"
     CONTAINER_TYPE_TGI = "tgi"
+    CONTAINER_TYPE_LLAMA_CPP = "llama-cpp"
 
 
 class InferenceContainerTypeFamily(str, metaclass=ExtendedEnumMeta):
@@ -55,6 +56,7 @@ class InferenceContainerTypeFamily(str, metaclass=ExtendedEnumMeta):
 class InferenceContainerParamType(str, metaclass=ExtendedEnumMeta):
     PARAM_TYPE_VLLM = "VLLM_PARAMS"
     PARAM_TYPE_TGI = "TGI_PARAMS"
+    PARAM_TYPE_LLAMA_CPP = "LLAMA_CPP_PARAMS"
 
 
 class HuggingFaceTags(str, metaclass=ExtendedEnumMeta):

ads/aqua/common/utils.py

@@ -17,9 +17,11 @@
 from typing import List, Union
 
 import fsspec
+import oci
 from cachetools import TTLCache, cached
+from oci.data_science.models import JobRun, Model
+from oci.object_storage.models import ObjectSummary
 
-import oci
 from ads.aqua.common.enums import (
     InferenceContainerParamType,
     InferenceContainerType,
@@ -52,8 +54,6 @@
 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.model import DataScienceModel, ModelVersionSet
-from oci.data_science.models import JobRun, Model
-from oci.object_storage.models import ObjectSummary
 
 logger = logging.getLogger("ads.aqua")
 
@@ -909,6 +909,8 @@ def get_container_params_type(container_type_name: str) -> str:
         return InferenceContainerParamType.PARAM_TYPE_VLLM
     elif InferenceContainerType.CONTAINER_TYPE_TGI in container_type_name.lower():
         return InferenceContainerParamType.PARAM_TYPE_TGI
+    elif InferenceContainerType.CONTAINER_TYPE_LLAMA_CPP in container_type_name.lower():
+        return InferenceContainerParamType.PARAM_TYPE_LLAMA_CPP
     else:
         return UNKNOWN
 

ads/aqua/constants.py

@@ -74,3 +74,7 @@
     "--sharded",
     "--trust-remote-code",
 }
+LLAMA_CPP_INFERENCE_RESTRICTED_PARAMS = {
+    "--port",
+    "--host",
+}

ads/aqua/modeldeployment/deployment.py

@@ -6,6 +6,7 @@
 from typing import Dict, List, Optional, Union
 
 from ads.aqua.app import AquaApp, logger
+from ads.aqua.common.entities import ContainerSpec
 from ads.aqua.common.enums import (
     InferenceContainerTypeFamily,
     Tags,
@@ -38,7 +39,6 @@
 from ads.aqua.modeldeployment.entities import (
     AquaDeployment,
     AquaDeploymentDetail,
-    ContainerSpec,
 )
 from ads.aqua.ui import ModelFormat
 from ads.common.object_storage_details import ObjectStorageDetails
@@ -281,8 +281,10 @@ def create(
             f"Aqua Image used for deploying {aqua_model.id} : {container_image}"
         )
 
+        # todo: use AquaContainerConfig.from_container_index_json instead.
         # Fetch the startup cli command for the container
-        # container_index.json will have "containerSpec" section which will provide the cli params for a given container family
+        # container_index.json will have "containerSpec" section which will provide the cli params for
+        # a given container family
         container_config = get_container_config()
         container_spec = container_config.get(ContainerSpec.CONTAINER_SPEC, {}).get(
             container_type_key, {}
@@ -308,6 +310,18 @@ def create(
         # validate user provided params
         user_params = env_var.get("PARAMS", UNKNOWN)
         if user_params:
+            # todo: remove this check in the future version, logic to be moved to container_index
+            if (
+                container_type_key.lower()
+                == InferenceContainerTypeFamily.AQUA_LLAMA_CPP_CONTAINER_FAMILY
+            ):
+                # 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"with deployment without parameter overrides."
+                )
+
             restricted_params = self._find_restricted_params(
                 params, user_params, container_type_key
             )

ads/aqua/modeldeployment/entities.py

@@ -1,12 +1,14 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*-
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 from dataclasses import dataclass, field
 from typing import Union
 
-from oci.data_science.models import ModelDeployment, ModelDeploymentSummary
+from oci.data_science.models import (
+    ModelDeployment,
+    ModelDeploymentSummary,
+)
 
 from ads.aqua.common.enums import Tags
 from ads.aqua.constants import UNKNOWN, UNKNOWN_DICT
@@ -24,18 +26,6 @@ class ModelParams:
     model: str = None
 
 
-class ContainerSpec:
-    """
-    Class to hold to hold keys within the container spec.
-    """
-
-    CONTAINER_SPEC = "containerSpec"
-    CLI_PARM = "cliParam"
-    SERVER_PORT = "serverPort"
-    HEALTH_CHECK_PORT = "healthCheckPort"
-    ENV_VARS = "envVars"
-
-
 @dataclass
 class ShapeInfo:
     instance_shape: str = None
@@ -61,6 +51,7 @@ class AquaDeployment(DataClassSerializable):
     lifecycle_details: str = None
     shape_info: field(default_factory=ShapeInfo) = None
     tags: dict = None
+    environment_variables: dict = None
 
     @classmethod
     def from_oci_model_deployment(
@@ -83,15 +74,12 @@ def from_oci_model_deployment(
         AquaDeployment:
             The instance of the Aqua model deployment.
         """
-        instance_configuration = (
-            oci_model_deployment.model_deployment_configuration_details.model_configuration_details.instance_configuration
-        )
+        instance_configuration = oci_model_deployment.model_deployment_configuration_details.model_configuration_details.instance_configuration
         instance_shape_config_details = (
             instance_configuration.model_deployment_instance_shape_config_details
         )
-        instance_count = (
-            oci_model_deployment.model_deployment_configuration_details.model_configuration_details.scaling_policy.instance_count
-        )
+        instance_count = oci_model_deployment.model_deployment_configuration_details.model_configuration_details.scaling_policy.instance_count
+        environment_variables = oci_model_deployment.model_deployment_configuration_details.environment_configuration_details.environment_variables
         shape_info = ShapeInfo(
             instance_shape=instance_configuration.instance_shape_name,
             instance_count=instance_count,
@@ -131,6 +119,7 @@ def from_oci_model_deployment(
                 region=region,
             ),
             tags=freeform_tags,
+            environment_variables=environment_variables,
         )
 
 

ads/aqua/ui.py

@@ -9,9 +9,12 @@
 from typing import Dict, List, Optional
 
 from cachetools import TTLCache
+from oci.exceptions import ServiceError
+from oci.identity.models import Compartment
 
 from ads.aqua import logger
 from ads.aqua.app import AquaApp
+from ads.aqua.common.entities import ContainerSpec
 from ads.aqua.common.enums import Tags
 from ads.aqua.common.errors import AquaResourceAccessError, AquaValueError
 from ads.aqua.common.utils import get_container_config, load_config, sanitize_response
@@ -27,8 +30,6 @@
     TENANCY_OCID,
 )
 from ads.telemetry import telemetry
-from oci.exceptions import ServiceError
-from oci.identity.models import Compartment
 
 
 class ModelFormat(Enum):
@@ -40,6 +41,19 @@ def to_dict(self):
         return self.value
 
 
+# todo: the container config spec information is shared across ui and deployment modules, move them
+#   within ads.aqua.common.entities. In that case, check for circular imports due to usage of get_container_config.
+
+
+@dataclass(repr=False)
+class AquaContainerConfigSpec(DataClassSerializable):
+    cli_param: str = None
+    server_port: str = None
+    health_check_port: str = None
+    env_vars: List[dict] = None
+    restricted_params: List[str] = None
+
+
 @dataclass(repr=False)
 class AquaContainerConfigItem(DataClassSerializable):
     """Represents an item of the AQUA container configuration."""
@@ -60,6 +74,7 @@ def __repr__(self):
     family: str = None
     platforms: List[Platform] = None
     model_formats: List[ModelFormat] = None
+    spec: AquaContainerConfigSpec = field(default_factory=AquaContainerConfigSpec)
 
 
 @dataclass(repr=False)
@@ -81,7 +96,9 @@ def to_dict(self):
 
     @classmethod
     def from_container_index_json(
-        cls, config: Optional[Dict] = None
+        cls,
+        config: Optional[Dict] = None,
+        enable_spec: Optional[bool] = False,
     ) -> "AquaContainerConfig":
         """
         Create an AquaContainerConfig instance from a container index JSON.
@@ -90,6 +107,8 @@ def from_container_index_json(
         ----------
         config : Dict
             The container index JSON.
+        enable_spec: bool
+            flag to check if container specification details should be fetched.
 
         Returns
         -------
@@ -114,6 +133,13 @@ def from_container_index_json(
                         ModelFormat[model_format]
                         for model_format in container.get("modelFormats", [])
                     ]
+                    container_spec = (
+                        config.get(ContainerSpec.CONTAINER_SPEC, {}).get(
+                            container_type, {}
+                        )
+                        if enable_spec
+                        else None
+                    )
                     container_item = AquaContainerConfigItem(
                         name=container.get("name", ""),
                         version=container.get("version", ""),
@@ -123,6 +149,21 @@ def from_container_index_json(
                         family=container_type,
                         platforms=platforms,
                         model_formats=model_formats,
+                        spec=AquaContainerConfigSpec(
+                            cli_param=container_spec.get(ContainerSpec.CLI_PARM, ""),
+                            server_port=container_spec.get(
+                                ContainerSpec.SERVER_PORT, ""
+                            ),
+                            health_check_port=container_spec.get(
+                                ContainerSpec.HEALTH_CHECK_PORT, ""
+                            ),
+                            env_vars=container_spec.get(ContainerSpec.ENV_VARS, []),
+                            restricted_params=container_spec.get(
+                                ContainerSpec.RESTRICTED_PARAMS, []
+                            ),
+                        )
+                        if container_spec
+                        else None,
                     )
                     if container.get("type") == "inference":
                         inference_items[container_type] = container_item
@@ -571,5 +612,6 @@ def list_containers(self) -> AquaContainerConfig:
             The AQUA containers configurations.
         """
         return AquaContainerConfig.from_container_index_json(
-            config=get_container_config()
+            config=get_container_config(),
+            enable_spec=True,
         )

tests/unitary/with_extras/aqua/test_data/ui/container_index.json

@@ -20,6 +20,7 @@
         }
       ],
       "healthCheckPort": "8080",
+      "restrictedParams": [],
       "serverPort": "8080"
     },
     "odsc-tgi-serving": {
@@ -39,10 +40,17 @@
         }
       ],
       "healthCheckPort": "8080",
+      "restrictedParams": [
+        "--port",
+        "--hostname",
+        "--num-shard",
+        "--sharded",
+        "--trust-remote-code"
+      ],
       "serverPort": "8080"
     },
     "odsc-vllm-serving": {
-      "cliParam": "--served-model-name $(python -c 'import os; print(os.environ.get(\"ODSC_SERVED_MODEL_NAME\",\"odsc-llm\"))') --seed 42 ",
+      "cliParam": "--served-model-name odsc-llm --seed 42 ",
       "envVars": [
         {
           "MODEL_DEPLOY_PREDICT_ENDPOINT": "/v1/completions"
@@ -58,6 +66,12 @@
         }
       ],
       "healthCheckPort": "8080",
+      "restrictedParams": [
+        "--port",
+        "--host",
+        "--served-model-name",
+        "--seed"
+      ],
       "serverPort": "8080"
     }
   },

tests/unitary/with_extras/aqua/test_deployment.py

@@ -181,6 +181,12 @@ class TestDataset:
         "created_on": "2024-01-01T00:00:00.000000+00:00",
         "created_by": "ocid1.user.oc1..<OCID>",
         "endpoint": MODEL_DEPLOYMENT_URL,
+        "environment_variables": {
+            "BASE_MODEL": "service_models/model-name/artifact",
+            "MODEL_DEPLOY_ENABLE_STREAMING": "true",
+            "MODEL_DEPLOY_PREDICT_ENDPOINT": "/v1/completions",
+            "PARAMS": "--served-model-name odsc-llm --seed 42",
+        },
         "console_link": "https://cloud.oracle.com/data-science/model-deployments/ocid1.datasciencemodeldeployment.oc1.<region>.<MD_OCID>?region=region-name",
         "lifecycle_details": "",
         "shape_info": {
@@ -192,6 +198,14 @@ class TestDataset:
         "tags": {"OCI_AQUA": "active", "aqua_model_name": "model-name"},
     }
 
+    aqua_deployment_gguf_env_vars = {
+        "BASE_MODEL": "service_models/model-name/artifact",
+        "BASE_MODEL_FILE": "model-name.gguf",
+        "MODEL_DEPLOY_ENABLE_STREAMING": "true",
+        "MODEL_DEPLOY_HEALTH_ENDPOINT": "/v1/models",
+        "MODEL_DEPLOY_PREDICT_ENDPOINT": "/v1/completions",
+    }
+
     aqua_deployment_gguf_shape_info = {
         "instance_shape": DEPLOYMENT_SHAPE_NAME_CPU,
         "instance_count": 1,
@@ -544,6 +558,9 @@ def test_create_deployment_for_gguf_model(
         expected_result = copy.deepcopy(TestDataset.aqua_deployment_object)
         expected_result["state"] = "CREATING"
         expected_result["shape_info"] = TestDataset.aqua_deployment_gguf_shape_info
+        expected_result["environment_variables"] = (
+            TestDataset.aqua_deployment_gguf_env_vars
+        )
         assert actual_attributes == expected_result
 
     @parameterized.expand(