Fixing UTs

Author: kumar-shivam-ranjan

ads/aqua/app.py

@@ -461,9 +461,9 @@ def get_container_config_item(
         """
 
         aqua_container_config = self.get_container_config()
-        inference_config = aqua_container_config.to_dict().get("inference")
-        ft_config = aqua_container_config.to_dict().get("finetune")
-        eval_config = aqua_container_config.to_dict().get("evaluate")
+        inference_config = aqua_container_config.inference.values()
+        ft_config = aqua_container_config.finetune.values()
+        eval_config = aqua_container_config.evaluate.values()
         container = next(
             (
                 container

ads/aqua/config/container_config.py

@@ -166,7 +166,8 @@ def from_service_config(
                 spec=None,
             )
             container_type = container.family_name
-            if container.usages[0].lower() in "inference":
+            usages = [x.upper() for x in container.usages]
+            if "INFERENCE" in usages or "MULTI_MODEL" in usages:
                 container_item.platforms.append(
                     container.workload_configuration_details_list[
                         0
@@ -228,11 +229,11 @@ def from_service_config(
                     ),
                 )
                 container_item.spec = container_spec
-            if "INFERENCE" in (x.upper() for x in container.usages):
+            if "INFERENCE" in usages or "MULTI_MODEL" in usages:
                 inference_items[container_type] = container_item
-            if "FINE_TUNE" in (x.upper() for x in container.usages):
+            if "FINE_TUNE" in usages:
                 finetune_items[container_type] = container_item
-            if "EVALUATION" in (x.upper() for x in container.usages):
+            if "EVALUATION" in usages:
                 evaluate_items[container_type] = container_item
         return cls(
             inference=inference_items, finetune=finetune_items, evaluate=evaluate_items

ads/aqua/model/model.py

@@ -273,7 +273,7 @@ def create_multi(
         supported_container_families = [
             container_config_item.family
             for container_config_item in service_inference_containers
-            if Usage.MULTI_MODEL in container_config_item.usages
+            if Usage.MULTI_MODEL.upper() in container_config_item.usages
         ]
 
         if not supported_container_families:

ads/aqua/modeldeployment/deployment.py

@@ -34,7 +34,7 @@
     load_gpu_shapes_index,
     validate_cmd_var,
 )
-from ads.aqua.config.container_config import AquaContainerConfig, Usage
+from ads.aqua.config.container_config import Usage
 from ads.aqua.constants import (
     AQUA_MODEL_ARTIFACT_FILE,
     AQUA_MODEL_TYPE_CUSTOM,
@@ -225,16 +225,12 @@ def create(
             except ConfigValidationError as err:
                 raise AquaValueError(f"{err}") from err
 
-            service_inference_containers = (
-                AquaContainerConfig.from_container_index_json(
-                    config=container_config
-                ).inference.values()
-            )
+            service_inference_containers = container_config.inference.values()
 
             supported_container_families = [
                 container_config_item.family
                 for container_config_item in service_inference_containers
-                if Usage.MULTI_MODEL in container_config_item.usages
+                if Usage.MULTI_MODEL.upper() in container_config_item.usages
             ]
 
             if not supported_container_families: