refactored _get_task

Author: elizjo

ads/aqua/model/model.py

@@ -317,10 +317,7 @@ def create_multi(
 
             display_name_list.append(display_name)
 
-            model_task = source_model.freeform_tags.get(Tags.TASK, UNKNOWN)
-
-            if model_task != UNKNOWN:
-                self._get_task(model, model_task)
+            self._get_task(model, source_model)
 
             # Retrieve model artifact
             model_artifact_path = source_model.artifact
@@ -710,15 +707,23 @@ 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:
+    def _get_task(
+        self,
+        model: AquaMultiModelRef,
+        source_model: DataScienceModel,
+    ) -> 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)
+        # user does not supply model task, we extract from model metadata
+        if not model.model_task:
+            model.model_task = source_model.freeform_tags.get(Tags.TASK, UNKNOWN)
+
+        task_tag = re.sub(r"-", "_", model.model_task)
 
         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()}."
+                f"{task_tag} is not supported. Valid model_task inputs are: {MultiModelSupportedTaskType.values()}."
             )
 
     def _fetch_metric_from_metadata(

tests/unitary/with_extras/aqua/test_model.py

@@ -440,7 +440,8 @@ 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"
+        # testing _get_task when a user passes an invalid task to AquaMultiModelRef
+        model_info_1.model_task = "invalid_task"
 
         with pytest.raises(AquaValueError):
             model = self.app.create_multi(
@@ -449,7 +450,18 @@ def test_create_multimodel(
                 compartment_id="test_compartment_id",
             )
 
+        # testing if a user tries to invoke a model with a task mode that is not yet supported
+        model_info_1.model_task = None
+        mock_model.freeform_tags["task"] = "unsupported_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"
+        model_info_1.model_task = "text_embedding"
 
         # will create a multi-model group
         model = self.app.create_multi(