added validation logic for model_task and unit test in test_model.py

Author: elizjo

ads/aqua/model/enums.py

@@ -28,8 +28,6 @@ class FineTuningCustomMetadata(ExtendedEnum):
 class MultiModelSupportedTaskType(ExtendedEnum):
     TEXT_GENERATION = "text-generation"
     TEXT_GENERATION_ALT = "text_generation"
-    EMBEDDING_ALT = "text_embedding"
-
-class MultiModelConfigMode(ExtendedEnum):
-    EMBEDDING = "embedding"
-    DEFAULT = "completion"
+    IMAGE_TEXT_TO_TEXT = "image_text_to_text"
+    CODE_SYNTHESIS = "code_synthesis"
+    EMBEDDING = "text_embedding"

ads/aqua/model/model.py

@@ -4,6 +4,7 @@
 import json
 import os
 import pathlib
+import re
 from datetime import datetime, timedelta
 from threading import Lock
 from typing import Any, Dict, List, Optional, Set, Union
@@ -80,7 +81,7 @@
     ImportModelDetails,
     ModelValidationResult,
 )
-from ads.aqua.model.enums import MultiModelSupportedTaskType, MultiModelConfigMode
+from ads.aqua.model.enums import MultiModelSupportedTaskType
 from ads.common.auth import default_signer
 from ads.common.oci_resource import SEARCH_TYPE, OCIResource
 from ads.common.utils import (
@@ -709,14 +710,16 @@ def edit_registered_model(
         else:
             raise AquaRuntimeError("Only registered unverified models can be edited.")
 
-    def _get_task(
-        self,
-        model: AquaMultiModelRef,
-        freeform_task_tag: str
-    ) -> str:
-        """In a Multi Model Deployment, will set model task if freeform task tag from model needs a non-completion endpoint (embedding)"""
-        if freeform_task_tag == MultiModelSupportedTaskType.EMBEDDING_ALT:
-            model.model_task = MultiModelConfigMode.EMBEDDING
+    def _get_task(self, model: AquaMultiModelRef, freeform_task_tag: str) -> str:
+        """In a Multi Model Deployment, will set model_task parameter in AquaMultiModelRef from freeform tags or user"""
+        task_tag = re.sub(r"-", "_", freeform_task_tag)
+
+        if task_tag in MultiModelSupportedTaskType:
+            model.model_task = task_tag
+        else:
+            raise AquaValueError(
+                f"{freeform_task_tag} is not supported. Valid model_task inputs are: {MultiModelSupportedTaskType.values()}."
+            )
 
     def _fetch_metric_from_metadata(
         self,

tests/unitary/with_extras/aqua/test_deployment.py

@@ -45,6 +45,7 @@
     ModelDeploymentConfigSummary,
     ModelParams,
 )
+from ads.aqua.model.enums import MultiModelSupportedTaskType
 from ads.aqua.modeldeployment.utils import MultiModelDeploymentConfigLoader
 from ads.model.datascience_model import DataScienceModel
 from ads.model.deployment.model_deployment import ModelDeployment
@@ -276,7 +277,7 @@ class TestDataset:
                         "environment_configuration_type": "OCIR_CONTAINER",
                         "environment_variables": {
                             "MODEL_DEPLOY_PREDICT_ENDPOINT": "/v1/completions",
-                            "MULTI_MODEL_CONFIG": '{ "models": [{ "params": "--served-model-name model_one --tensor-parallel-size 1 --max-model-len 2096", "model_path": "models/model_one/5be6479/artifact/", "model_task": "embedding"}, {"params": "--served-model-name model_two --tensor-parallel-size 1 --max-model-len 2096", "model_path": "models/model_two/83e9aa1/artifact/"}, {"params": "--served-model-name model_three --tensor-parallel-size 1 --max-model-len 2096", "model_path": "models/model_three/83e9aa1/artifact/"}]}',
+                            "MULTI_MODEL_CONFIG": '{ "models": [{ "params": "--served-model-name model_one --tensor-parallel-size 1 --max-model-len 2096", "model_path": "models/model_one/5be6479/artifact/", "model_task": "text_embedding"}, {"params": "--served-model-name model_two --tensor-parallel-size 1 --max-model-len 2096", "model_path": "models/model_two/83e9aa1/artifact/", "model_task": "image_text_to_text"}, {"params": "--served-model-name model_three --tensor-parallel-size 1 --max-model-len 2096", "model_path": "models/model_three/83e9aa1/artifact/", "model_task": "code_synthesis"}]}',
                         },
                         "health_check_port": 8080,
                         "image": "dsmc://image-name:1.0.0.0",
@@ -486,30 +487,30 @@ class TestDataset:
                 "gpu_count": 2,
                 "model_id": "test_model_id_1",
                 "model_name": "test_model_1",
-                "model_task": "embedding",
+                "model_task": "text_embedding",
                 "artifact_location": "test_location_1",
             },
             {
                 "env_var": {},
                 "gpu_count": 2,
                 "model_id": "test_model_id_2",
                 "model_name": "test_model_2",
-                "model_task": None,
+                "model_task": "image_text_to_text",
                 "artifact_location": "test_location_2",
             },
             {
                 "env_var": {},
                 "gpu_count": 2,
                 "model_id": "test_model_id_3",
                 "model_name": "test_model_3",
-                "model_task": None,
+                "model_task": "code_synthesis",
                 "artifact_location": "test_location_3",
             },
         ],
         "model_id": "ocid1.datasciencemodel.oc1.<region>.<OCID>",
         "environment_variables": {
             "MODEL_DEPLOY_PREDICT_ENDPOINT": "/v1/completions",
-            "MULTI_MODEL_CONFIG": '{ "models": [{ "params": "--served-model-name model_one --tensor-parallel-size 1 --max-model-len 2096", "model_path": "models/model_one/5be6479/artifact/", "model_task": "embedding"}, {"params": "--served-model-name model_two --tensor-parallel-size 1 --max-model-len 2096", "model_path": "models/model_two/83e9aa1/artifact/"}, {"params": "--served-model-name model_three --tensor-parallel-size 1 --max-model-len 2096", "model_path": "models/model_three/83e9aa1/artifact/"}]}',
+            "MULTI_MODEL_CONFIG": '{ "models": [{ "params": "--served-model-name model_one --tensor-parallel-size 1 --max-model-len 2096", "model_path": "models/model_one/5be6479/artifact/", "model_task": "text_embedding"}, {"params": "--served-model-name model_two --tensor-parallel-size 1 --max-model-len 2096", "model_path": "models/model_two/83e9aa1/artifact/", "model_task": "image_text_to_text"}, {"params": "--served-model-name model_three --tensor-parallel-size 1 --max-model-len 2096", "model_path": "models/model_three/83e9aa1/artifact/", "model_task": "code_synthesis"}]}',
         },
         "cmd": [],
         "console_link": "https://cloud.oracle.com/data-science/model-deployments/ocid1.datasciencemodeldeployment.oc1.<region>.<MD_OCID>?region=region-name",
@@ -968,23 +969,23 @@ class TestDataset:
             "gpu_count": 1,
             "model_id": "ocid1.compartment.oc1..<OCID>",
             "model_name": "model_one",
-            "model_task": "embedding",
+            "model_task": "text_embedding",
             "artifact_location": "artifact_location_one",
         },
         {
             "env_var": {"--test_key_two": "test_value_two"},
             "gpu_count": 1,
             "model_id": "ocid1.compartment.oc1..<OCID>",
             "model_name": "model_two",
-            "model_task": None,
+            "model_task": "image_text_to_text",
             "artifact_location": "artifact_location_two",
         },
         {
             "env_var": {"--test_key_three": "test_value_three"},
             "gpu_count": 1,
             "model_id": "ocid1.compartment.oc1..<OCID>",
             "model_name": "model_three",
-            "model_task": None,
+            "model_task": "code_synthesis",
             "artifact_location": "artifact_location_three",
         },
     ]
@@ -1793,23 +1794,23 @@ def test_create_deployment_for_multi_model(
         model_info_1 = AquaMultiModelRef(
             model_id="test_model_id_1",
             model_name="test_model_1",
-            model_task="embedding",
+            model_task="text_embedding",
             gpu_count=2,
             artifact_location="test_location_1",
         )
 
         model_info_2 = AquaMultiModelRef(
             model_id="test_model_id_2",
             model_name="test_model_2",
-            model_task=None,
+            model_task="image_text_to_text",
             gpu_count=2,
             artifact_location="test_location_2",
         )
 
         model_info_3 = AquaMultiModelRef(
             model_id="test_model_id_3",
             model_name="test_model_3",
-            model_task=None,
+            model_task="code_synthesis",
             gpu_count=2,
             artifact_location="test_location_3",
         )

tests/unitary/with_extras/aqua/test_model.py

@@ -5,6 +5,7 @@
 
 import json
 import os
+import re
 import shlex
 import tempfile
 from dataclasses import asdict
@@ -13,17 +14,14 @@
 
 import oci
 import pytest
-
-from ads.aqua.app import AquaApp
-from ads.aqua.config.container_config import AquaContainerConfig
 from huggingface_hub.hf_api import HfApi, ModelInfo
 from parameterized import parameterized
 
 import ads.aqua.model
 import ads.common
 import ads.common.oci_client
 import ads.config
-
+from ads.aqua.app import AquaApp
 from ads.aqua.common.entities import AquaMultiModelRef
 from ads.aqua.common.enums import ModelFormat, Tags
 from ads.aqua.common.errors import (
@@ -32,6 +30,7 @@
     AquaValueError,
 )
 from ads.aqua.common.utils import get_hf_model_info
+from ads.aqua.config.container_config import AquaContainerConfig
 from ads.aqua.constants import HF_METADATA_FOLDER
 from ads.aqua.model import AquaModelApp
 from ads.aqua.model.entities import (
@@ -40,14 +39,14 @@
     ImportModelDetails,
     ModelValidationResult,
 )
+from ads.aqua.model.enums import MultiModelSupportedTaskType
 from ads.common.object_storage_details import ObjectStorageDetails
 from ads.model.datascience_model import DataScienceModel
 from ads.model.model_metadata import (
     ModelCustomMetadata,
     ModelProvenanceMetadata,
     ModelTaxonomyMetadata,
 )
-
 from tests.unitary.with_extras.aqua.utils import ServiceManagedContainers
 
 
@@ -397,12 +396,14 @@ def test_create_multimodel(
         model_info_1 = AquaMultiModelRef(
             model_id="test_model_id_1",
             gpu_count=2,
+            model_task = "text_embedding",
             env_var={"params": "--trust-remote-code --max-model-len 60000"},
         )
 
         model_info_2 = AquaMultiModelRef(
             model_id="test_model_id_2",
             gpu_count=2,
+            model_task = "image_text_to_text",
             env_var={"params": "--trust-remote-code --max-model-len 32000"},
         )
 
@@ -439,6 +440,17 @@ def test_create_multimodel(
         mock_model.custom_metadata_list = custom_metadata_list
         mock_from_id.return_value = mock_model
 
+        mock_model.freeform_tags["task"] = "invalid_task"
+
+        with pytest.raises(AquaValueError):
+            model = self.app.create_multi(
+                models=[model_info_1, model_info_2],
+                project_id="test_project_id",
+                compartment_id="test_compartment_id",
+            )
+
+        mock_model.freeform_tags["task"] = "text-generation"
+
         # will create a multi-model group
         model = self.app.create_multi(
             models=[model_info_1, model_info_2],