[ODSC-59694] Register model changes for 1.0.3 (#903)

Author: kumar-shivam-ranjan

ads/aqua/common/enums.py

@@ -1,5 +1,4 @@
 #!/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/
 
@@ -8,6 +7,7 @@
 ~~~~~~~~~~~~~~
 This module contains the set of enums used in AQUA.
 """
+
 from ads.common.extended_enum import ExtendedEnumMeta
 
 
@@ -28,6 +28,7 @@ class Tags(str, metaclass=ExtendedEnumMeta):
     TASK = "task"
     LICENSE = "license"
     ORGANIZATION = "organization"
+    PLATFORM = "platform"
     AQUA_TAG = "OCI_AQUA"
     AQUA_SERVICE_MODEL_TAG = "aqua_service_model"
     AQUA_FINE_TUNED_MODEL_TAG = "aqua_fine_tuned_model"
@@ -48,6 +49,7 @@ class InferenceContainerType(str, metaclass=ExtendedEnumMeta):
 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"
 
 
 class InferenceContainerParamType(str, metaclass=ExtendedEnumMeta):

ads/aqua/constants.py

@@ -21,7 +21,8 @@
 DEFAULT_FT_REPLICA = 1
 DEFAULT_FT_BATCH_SIZE = 1
 DEFAULT_FT_VALIDATION_SET_SIZE = 0.1
-
+ARM_CPU="arm_cpu"
+NVIDIA_GPU="nvidia_gpu"
 MAXIMUM_ALLOWED_DATASET_IN_BYTE = 52428800  # 1024 x 1024 x 50 = 50MB
 JOB_INFRASTRUCTURE_TYPE_DEFAULT_NETWORKING = "ME_STANDALONE"
 NB_SESSION_IDENTIFIER = "NB_SESSION_OCID"

ads/aqua/model/entities.py

@@ -76,6 +76,7 @@ class AquaModelSummary(DataClassSerializable):
     ready_to_deploy: bool = True
     ready_to_finetune: bool = False
     ready_to_import: bool = False
+    platform: List[str] = field(default_factory=lambda: ["nvidia_gpu"])
 
 
 @dataclass(repr=False)

ads/aqua/model/model.py

@@ -1,5 +1,4 @@
 #!/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/
 import os
@@ -12,7 +11,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, InferenceContainerTypeFamily
 from ads.aqua.common.errors import AquaRuntimeError
 from ads.aqua.common.utils import (
     create_word_icon,
@@ -36,7 +35,7 @@
     AQUA_MODEL_ARTIFACT_CONFIG,
     AQUA_MODEL_ARTIFACT_CONFIG_MODEL_NAME,
     AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE,
-    AQUA_MODEL_TYPE_CUSTOM,
+    AQUA_MODEL_TYPE_CUSTOM, ARM_CPU, NVIDIA_GPU,
 )
 from ads.aqua.model.constants import *
 from ads.aqua.model.entities import *
@@ -268,8 +267,7 @@ def get(self, model_id: str, load_model_card: Optional[bool] = True) -> "AquaMod
 
             job_run_status = (
                 jobrun.lifecycle_state
-                if jobrun
-                and not jobrun.lifecycle_state == JobRun.LIFECYCLE_STATE_DELETED
+                if jobrun and jobrun.lifecycle_state != JobRun.LIFECYCLE_STATE_DELETED
                 else (
                     JobRun.LIFECYCLE_STATE_SUCCEEDED
                     if self.if_artifact_exist(ds_model.id)
@@ -540,7 +538,7 @@ def clear_model_list_cache(
             dict with the key used, and True if cache has the key that needs to be deleted.
         """
         res = {}
-        logger.info(f"Clearing _service_models_cache")
+        logger.info("Clearing _service_models_cache")
         with self._cache_lock:
             if ODSC_MODEL_COMPARTMENT_OCID in self._service_models_cache.keys():
                 self._service_models_cache.pop(key=ODSC_MODEL_COMPARTMENT_OCID)
@@ -561,6 +559,7 @@ def _create_model_catalog_entry(
         verified_model: DataScienceModel,
         compartment_id: Optional[str],
         project_id: Optional[str],
+        is_gguf_model: bool,
     ) -> DataScienceModel:
         """Create model by reference from the object storage path
 
@@ -583,9 +582,13 @@ def _create_model_catalog_entry(
                 Tags.AQUA_SERVICE_MODEL_TAG: verified_model.id,
             }
             if verified_model
-            else {Tags.AQUA_TAG: "active", Tags.BASE_MODEL_CUSTOM: "true"}
+            else {
+                Tags.AQUA_TAG: "active",
+                Tags.BASE_MODEL_CUSTOM: "true",
+            }
         )
         tags.update({Tags.BASE_MODEL_CUSTOM: "true"})
+        tags.update({Tags.PLATFORM: ARM_CPU if is_gguf_model else NVIDIA_GPU})
 
         # Remove `ready_to_import` tag that might get copied from service model.
         tags.pop(Tags.READY_TO_IMPORT, None)
@@ -615,8 +618,8 @@ def _create_model_catalog_entry(
                 )
             else:
                 logger.warn(
-                    f"Proceeding with model registration without the fine-tuning container information. "
-                    f"This model will not be available for fine tuning."
+                    "Proceeding with model registration without the fine-tuning container information. "
+                    "This model will not be available for fine tuning."
                 )
 
             metadata.add(
@@ -693,24 +696,26 @@ def register(
                 The registered model as a AquaModel object.
         """
         verified_model_details: DataScienceModel = None
-
+        model_config = None
         if not import_model_details:
             import_model_details = ImportModelDetails(**kwargs)
-
-        try:
-            model_config = load_config(
-                file_path=import_model_details.os_path,
-                config_file_name=AQUA_MODEL_ARTIFACT_CONFIG,
-            )
-        except Exception as ex:
-            logger.error(
-                f"Exception occurred while loading config file from {import_model_details.os_path}"
-                f"Exception message: {ex}"
-            )
-            raise AquaRuntimeError(
-                f"The model path {import_model_details.os_path} does not contain the file config.json. "
-                f"Please check if the path is correct or the model artifacts are available at this location."
-            )
+        is_gguf_model = import_model_details.inference_container == InferenceContainerTypeFamily.AQUA_LLAMA_CPP_CONTAINER_FAMILY
+        platform = ARM_CPU if is_gguf_model else NVIDIA_GPU
+        if not is_gguf_model:
+            try:
+                model_config = load_config(
+                    file_path=import_model_details.os_path,
+                    config_file_name=AQUA_MODEL_ARTIFACT_CONFIG,
+                )
+            except Exception as ex:
+                logger.error(
+                    f"Exception occurred while loading config file from {import_model_details.os_path}"
+                    f"Exception message: {ex}"
+                )
+                raise AquaRuntimeError(
+                    f"The model path {import_model_details.os_path} does not contain the file config.json. "
+                    f"Please check if the path is correct or the model artifacts are available at this location."
+                )
 
         model_service_id = None
         # If OCID of a model is passed, we need to copy the defaults for Tags and metadata from the service model.
@@ -770,6 +775,7 @@ def register(
             verified_model=verified_model_details,
             compartment_id=import_model_details.compartment_id,
             project_id=import_model_details.project_id,
+            is_gguf_model=is_gguf_model,
         )
         # registered model will always have inference and evaluation container, but
         # fine-tuning container may be not set
@@ -798,17 +804,23 @@ def register(
             inference_container=inference_container,
             finetuning_container=finetuning_container,
             evaluation_container=evaluation_container,
+            platform=platform,
         )
 
         if verified_model_details:
             telemetry_model_name = model_name
+        elif (
+            model_config is not None
+            and AQUA_MODEL_ARTIFACT_CONFIG_MODEL_NAME in model_config
+        ):
+            telemetry_model_name = f"{AQUA_MODEL_TYPE_CUSTOM}_{model_config[AQUA_MODEL_ARTIFACT_CONFIG_MODEL_NAME]}"
+        elif (
+            model_config is not None
+            and AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE in model_config
+        ):
+            telemetry_model_name = f"{AQUA_MODEL_TYPE_CUSTOM}_{model_config[AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE]}"
         else:
-            if AQUA_MODEL_ARTIFACT_CONFIG_MODEL_NAME in model_config:
-                telemetry_model_name = f"{AQUA_MODEL_TYPE_CUSTOM}_{model_config[AQUA_MODEL_ARTIFACT_CONFIG_MODEL_NAME]}"
-            elif AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE in model_config:
-                telemetry_model_name = f"{AQUA_MODEL_TYPE_CUSTOM}_{model_config[AQUA_MODEL_ARTIFACT_CONFIG_MODEL_TYPE]}"
-            else:
-                telemetry_model_name = AQUA_MODEL_TYPE_CUSTOM
+            telemetry_model_name = AQUA_MODEL_TYPE_CUSTOM
 
         self.telemetry.record_event_async(
             category="aqua/model",