[DRCC][Ready to Merge] AQUA MS Enhancements (#1058)

Author: kumar-shivam-ranjan

ads/aqua/app.py

@@ -7,11 +7,16 @@
 import traceback
 from dataclasses import fields
 from datetime import datetime, timedelta
-from typing import Any, Dict, Optional, Union
+from itertools import chain
+from typing import Any, Dict, List, Optional, Union
 
 import oci
 from cachetools import TTLCache, cached
-from oci.data_science.models import UpdateModelDetails, UpdateModelProvenanceDetails
+from oci.data_science.models import (
+    ContainerSummary,
+    UpdateModelDetails,
+    UpdateModelProvenanceDetails,
+)
 
 from ads import set_auth
 from ads.aqua import logger
@@ -24,6 +29,11 @@
     is_valid_ocid,
     load_config,
 )
+from ads.aqua.config.container_config import (
+    AquaContainerConfig,
+    AquaContainerConfigItem,
+)
+from ads.aqua.constants import SERVICE_MANAGED_CONTAINER_URI_SCHEME
 from ads.common import oci_client as oc
 from ads.common.auth import default_signer
 from ads.common.utils import UNKNOWN, extract_region, is_path_exists
@@ -240,7 +250,9 @@ def create_model_catalog(
             .with_custom_metadata_list(model_custom_metadata)
             .with_defined_metadata_list(model_taxonomy_metadata)
             .with_provenance_metadata(ModelProvenanceMetadata(training_id=UNKNOWN))
-            .with_defined_tags(**(defined_tags or {})) # Create defined tags when a model is created.
+            .with_defined_tags(
+                **(defined_tags or {})
+            )  # Create defined tags when a model is created.
             .create(
                 **kwargs,
             )
@@ -271,6 +283,43 @@ def if_artifact_exist(self, model_id: str, **kwargs) -> bool:
                 logger.info(f"Artifact not found in model {model_id}.")
                 return False
 
+    def get_config_from_metadata(
+        self, model_id: str, metadata_key: str
+    ) -> ModelConfigResult:
+        """Gets the config for the given Aqua model from model catalog metadata content.
+
+        Parameters
+        ----------
+        model_id: str
+            The OCID of the Aqua model.
+        metadata_key: str
+            The metadata key name where artifact content is stored
+        Returns
+        -------
+        ModelConfigResult
+            A Pydantic model containing the model_details (extracted from OCI) and the config dictionary.
+        """
+        config = {}
+        oci_model = self.ds_client.get_model(model_id).data
+        try:
+            config = self.ds_client.get_model_defined_metadatum_artifact_content(
+                model_id, metadata_key
+            ).data.content.decode("utf-8")
+            return ModelConfigResult(config=json.loads(config), model_details=oci_model)
+        except UnicodeDecodeError as ex:
+            logger.error(
+                f"Failed to decode content for '{metadata_key}' in defined metadata for model '{model_id}' : {ex}"
+            )
+        except json.JSONDecodeError as ex:
+            logger.error(
+                f"Invalid JSON format for '{metadata_key}' in defined metadata for model '{model_id}' : {ex}"
+            )
+        except Exception as ex:
+            logger.error(
+                f"Failed to retrieve defined metadata key '{metadata_key}' for model '{model_id}': {ex}"
+            )
+        return ModelConfigResult(config=config, model_details=oci_model)
+
     @cached(cache=TTLCache(maxsize=1, ttl=timedelta(minutes=1), timer=datetime.now))
     def get_config(
         self,
@@ -310,22 +359,7 @@ def get_config(
             raise AquaRuntimeError(f"Target model {oci_model.id} is not an Aqua model.")
 
         config: Dict[str, Any] = {}
-
-        # if the current model has a service model tag, then
-        if Tags.AQUA_SERVICE_MODEL_TAG in oci_model.freeform_tags:
-            base_model_ocid = oci_model.freeform_tags[Tags.AQUA_SERVICE_MODEL_TAG]
-            logger.info(
-                f"Base model found for the model: {oci_model.id}. "
-                f"Loading {config_file_name} for base model {base_model_ocid}."
-            )
-            if config_folder == ConfigFolder.ARTIFACT:
-                artifact_path = get_artifact_path(oci_model.custom_metadata_list)
-            else:
-                base_model = self.ds_client.get_model(base_model_ocid).data
-                artifact_path = get_artifact_path(base_model.custom_metadata_list)
-        else:
-            logger.info(f"Loading {config_file_name} for model {oci_model.id}...")
-            artifact_path = get_artifact_path(oci_model.custom_metadata_list)
+        artifact_path = get_artifact_path(oci_model.custom_metadata_list)
         if not artifact_path:
             logger.debug(
                 f"Failed to get artifact path from custom metadata for the model: {model_id}"
@@ -340,7 +374,7 @@ def get_config(
         config_file_path = os.path.join(config_path, config_file_name)
         if is_path_exists(config_file_path):
             try:
-                logger.debug(
+                logger.info(
                     f"Loading config: `{config_file_name}` from `{config_path}`"
                 )
                 config = load_config(
@@ -361,6 +395,85 @@ def get_config(
 
         return ModelConfigResult(config=config, model_details=oci_model)
 
+    def get_container_image(self, container_type: str = None) -> str:
+        """
+        Gets the latest smc container complete image name from the given container type.
+
+        Parameters
+        ----------
+        container_type: str
+            type of container, can be either odsc-vllm-serving, odsc-llm-fine-tuning, odsc-llm-evaluate
+
+        Returns
+        -------
+        str:
+            A complete container name along with version. ex: dsmc://odsc-vllm-serving:0.7.4.1
+        """
+
+        containers = self.list_service_containers()
+        container = next(
+            (c for c in containers if c.is_latest and c.family_name == container_type),
+            None,
+        )
+        if not container:
+            raise AquaValueError(f"Invalid container type : {container_type}")
+        container_image = (
+            SERVICE_MANAGED_CONTAINER_URI_SCHEME
+            + container.container_name
+            + ":"
+            + container.tag
+        )
+        return container_image
+
+    @cached(cache=TTLCache(maxsize=20, ttl=timedelta(minutes=30), timer=datetime.now))
+    def list_service_containers(self) -> List[ContainerSummary]:
+        """
+        List containers from containers.conf in OCI Datascience control plane
+        """
+        containers = self.ds_client.list_containers().data
+        return containers
+
+    def get_container_config(self) -> AquaContainerConfig:
+        """
+        Fetches latest containers from containers.conf in OCI Datascience control plane
+
+        Returns
+        -------
+        AquaContainerConfig
+            An Object that contains latest container info for the given container family
+
+        """
+        return AquaContainerConfig.from_service_config(
+            service_containers=self.list_service_containers()
+        )
+
+    def get_container_config_item(
+        self, container_family: str
+    ) -> AquaContainerConfigItem:
+        """
+        Fetches latest container for given container_family_name from containers.conf in OCI Datascience control plane
+
+        Returns
+        -------
+        AquaContainerConfigItem
+            An Object that contains latest container info for the given container family
+
+        """
+
+        aqua_container_config = self.get_container_config()
+        inference_config = aqua_container_config.inference.values()
+        ft_config = aqua_container_config.finetune.values()
+        eval_config = aqua_container_config.evaluate.values()
+        container = next(
+            (
+                container
+                for container in chain(inference_config, ft_config, eval_config)
+                if container.family.lower() == container_family.lower()
+            ),
+            None,
+        )
+        return container
+
     @property
     def telemetry(self):
         if not self._telemetry:

ads/aqua/common/utils.py

@@ -53,9 +53,13 @@
     COMPARTMENT_MAPPING_KEY,
     CONSOLE_LINK_RESOURCE_TYPE_MAPPING,
     CONTAINER_INDEX,
+    DEPLOYMENT_CONFIG,
+    FINE_TUNING_CONFIG,
     HF_LOGIN_DEFAULT_TIMEOUT,
+    LICENSE,
     MAXIMUM_ALLOWED_DATASET_IN_BYTE,
     MODEL_BY_REFERENCE_OSS_PATH_KEY,
+    README,
     SERVICE_MANAGED_CONTAINER_URI_SCHEME,
     SUPPORTED_FILE_FORMATS,
     TEI_CONTAINER_DEFAULT_HOST,
@@ -89,6 +93,14 @@
 logger = logging.getLogger("ads.aqua")
 
 
+DEFINED_METADATA_TO_FILE_MAP = {
+    "readme": README,
+    "license": LICENSE,
+    "finetuneconfiguration": FINE_TUNING_CONFIG,
+    "deploymentconfiguration": DEPLOYMENT_CONFIG,
+}
+
+
 class LifecycleStatus(ExtendedEnum):
     UNKNOWN = ""
 
@@ -553,57 +565,6 @@ def service_config_path():
     return f"oci://{AQUA_SERVICE_MODELS_BUCKET}@{CONDA_BUCKET_NS}/service_models/config"
 
 
-@cached(cache=TTLCache(maxsize=1, ttl=timedelta(minutes=10), timer=datetime.now))
-def get_container_config():
-    config = load_config(
-        file_path=service_config_path(),
-        config_file_name=CONTAINER_INDEX,
-    )
-
-    return config
-
-
-def get_container_image(
-    config_file_name: str = None, container_type: str = None
-) -> str:
-    """Gets the image name from the given model and container type.
-    Parameters
-    ----------
-    config_file_name: str
-        name of the config file
-    container_type: str
-        type of container, can be either deployment-container, finetune-container, evaluation-container
-
-    Returns
-    -------
-    Dict:
-        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:
-        return UNKNOWN
-
-    mapping = config[container_type]
-    versions = [obj["version"] for obj in mapping]
-    # assumes numbered versions, update if `latest` is used
-    latest = get_max_version(versions)
-    for obj in mapping:
-        if obj["version"] == str(latest):
-            container_image = f"{obj['name']}:{obj['version']}"
-            break
-
-    if not container_image:
-        raise AquaValueError(
-            f"{config_file_name} is missing name and/or version details."
-        )
-
-    return container_image
-
-
 def fetch_service_compartment() -> Union[str, None]:
     """
     Loads the compartment mapping json from service bucket.