Merge branch 'main' into feature/model_group

Author: mrDzurb

ads/aqua/constants.py

@@ -55,6 +55,8 @@
 SUPPORTED_FILE_FORMATS = ["jsonl"]
 MODEL_BY_REFERENCE_OSS_PATH_KEY = "artifact_location"
 
+AQUA_CHAT_TEMPLATE_METADATA_KEY = "chat_template"
+
 CONSOLE_LINK_RESOURCE_TYPE_MAPPING = {
     "datasciencemodel": "models",
     "datasciencemodeldeployment": "model-deployments",

ads/aqua/extension/model_handler.py

@@ -11,12 +11,15 @@
 from ads.aqua.common.enums import CustomInferenceContainerTypeFamily
 from ads.aqua.common.errors import AquaRuntimeError
 from ads.aqua.common.utils import get_hf_model_info, is_valid_ocid, list_hf_models
+from ads.aqua.constants import AQUA_CHAT_TEMPLATE_METADATA_KEY
 from ads.aqua.extension.base_handler import AquaAPIhandler
 from ads.aqua.extension.errors import Errors
 from ads.aqua.model import AquaModelApp
 from ads.aqua.model.entities import AquaModelSummary, HFModelSummary
 from ads.config import SERVICE
+from ads.model import DataScienceModel
 from ads.model.common.utils import MetadataArtifactPathType
+from ads.model.service.oci_datascience_model import OCIDataScienceModel
 
 
 class AquaModelHandler(AquaAPIhandler):
@@ -320,26 +323,65 @@ def post(self, *args, **kwargs):  # noqa: ARG002
         )
 
 
-class AquaModelTokenizerConfigHandler(AquaAPIhandler):
+class AquaModelChatTemplateHandler(AquaAPIhandler):
     def get(self, model_id):
         """
-        Handles requests for retrieving the Hugging Face tokenizer configuration of a specified model.
-        Expected request format: GET /aqua/models/<model-ocid>/tokenizer
+        Handles requests for retrieving the chat template from custom metadata of a specified model.
+        Expected request format: GET /aqua/models/<model-ocid>/chat-template
 
         """
 
         path_list = urlparse(self.request.path).path.strip("/").split("/")
-        # Path should be /aqua/models/ocid1.iad.ahdxxx/tokenizer
-        # path_list=['aqua','models','<model-ocid>','tokenizer']
+        # Path should be /aqua/models/ocid1.iad.ahdxxx/chat-template
+        # path_list=['aqua','models','<model-ocid>','chat-template']
         if (
             len(path_list) == 4
             and is_valid_ocid(path_list[2])
-            and path_list[3] == "tokenizer"
+            and path_list[3] == "chat-template"
         ):
-            return self.finish(AquaModelApp().get_hf_tokenizer_config(model_id))
+            try:
+                oci_data_science_model = OCIDataScienceModel.from_id(model_id)
+            except Exception as e:
+                raise HTTPError(404, f"Model not found for id: {model_id}. Details: {str(e)}")
+            return self.finish(oci_data_science_model.get_custom_metadata_artifact("chat_template"))
 
         raise HTTPError(400, f"The request {self.request.path} is invalid.")
 
+    @handle_exceptions
+    def post(self, model_id: str):
+        """
+        Handles POST requests to add a custom chat_template metadata artifact to a model.
+
+        Expected request format:
+        POST /aqua/models/<model-ocid>/chat-template
+        Body: { "chat_template": "<your_template_string>" }
+
+        """
+        try:
+            input_body = self.get_json_body()
+        except Exception as e:
+            raise HTTPError(400, f"Invalid JSON body: {str(e)}")
+
+        chat_template = input_body.get("chat_template")
+        if not chat_template:
+            raise HTTPError(400, "Missing required field: 'chat_template'")
+
+        try:
+            data_science_model = DataScienceModel.from_id(model_id)
+        except Exception as e:
+            raise HTTPError(404, f"Model not found for id: {model_id}. Details: {str(e)}")
+
+        try:
+            result = data_science_model.create_custom_metadata_artifact(
+                metadata_key_name=AQUA_CHAT_TEMPLATE_METADATA_KEY,
+                path_type=MetadataArtifactPathType.CONTENT,
+                artifact_path_or_content=chat_template.encode()
+            )
+        except Exception as e:
+            raise HTTPError(500, f"Failed to create metadata artifact: {str(e)}")
+
+        return self.finish(result)
+
 
 class AquaModelDefinedMetadataArtifactHandler(AquaAPIhandler):
     """
@@ -381,7 +423,7 @@ def post(self, model_id: str, metadata_key: str):
     ("model/?([^/]*)", AquaModelHandler),
     ("model/?([^/]*)/license", AquaModelLicenseHandler),
     ("model/?([^/]*)/readme", AquaModelReadmeHandler),
-    ("model/?([^/]*)/tokenizer", AquaModelTokenizerConfigHandler),
+    ("model/?([^/]*)/chat-template", AquaModelChatTemplateHandler),
     ("model/hf/search/?([^/]*)", AquaHuggingFaceHandler),
     (
         "model/?([^/]*)/definedMetadata/?([^/]*)",

ads/aqua/model/constants.py

@@ -26,6 +26,7 @@ class ModelTask(ExtendedEnum):
     TEXT_GENERATION = "text-generation"
     IMAGE_TEXT_TO_TEXT = "image-text-to-text"
     IMAGE_TO_TEXT = "image-to-text"
+    TIME_SERIES_FORECASTING = "time-series-forecasting"
 
 
 class FineTuningMetricCategories(ExtendedEnum):

ads/aqua/modeldeployment/deployment.py

@@ -218,6 +218,14 @@ def create(
                 freeform_tags=freeform_tags,
                 defined_tags=defined_tags,
             )
+            task_tag = aqua_model.freeform_tags.get(Tags.TASK, UNKNOWN)
+            if (
+                task_tag == ModelTask.TIME_SERIES_FORECASTING
+                or task_tag == ModelTask.TIME_SERIES_FORECASTING.replace("-", "_")
+            ):
+                create_deployment_details.env_var.update(
+                    {Tags.TASK.upper(): ModelTask.TIME_SERIES_FORECASTING}
+                )
             return self._create(
                 aqua_model=aqua_model,
                 create_deployment_details=create_deployment_details,
@@ -1077,14 +1085,16 @@ def _create_deployment(
         ).deploy(wait_for_completion=False)
 
         deployment_id = deployment.id
+
         logger.info(
             f"Aqua model deployment {deployment_id} created for model {aqua_model_id}. Work request Id is {deployment.dsc_model_deployment.workflow_req_id}"
         )
+        status_list = []
 
         progress_thread = threading.Thread(
             target=self.get_deployment_status,
             args=(
-                deployment_id,
+                deployment,
                 deployment.dsc_model_deployment.workflow_req_id,
                 model_type,
                 model_name,
@@ -1604,7 +1614,7 @@ def list_shapes(self, **kwargs) -> List[ComputeShapeSummary]:
 
     def get_deployment_status(
         self,
-        model_deployment_id: str,
+        deployment: ModelDeployment,
         work_request_id: str,
         model_type: str,
         model_name: str,
@@ -1626,37 +1636,60 @@ def get_deployment_status(
         AquaDeployment
             An Aqua deployment instance.
         """
-        ocid = get_ocid_substring(model_deployment_id, key_len=8)
-        telemetry_kwargs = {"ocid": ocid}
-
+        ocid = get_ocid_substring(deployment.id, key_len=8)
         data_science_work_request: DataScienceWorkRequest = DataScienceWorkRequest(
             work_request_id
         )
-
         try:
             data_science_work_request.wait_work_request(
                 progress_bar_description="Creating model deployment",
                 max_wait_time=DEFAULT_WAIT_TIME,
                 poll_interval=DEFAULT_POLL_INTERVAL,
             )
         except Exception:
+            status = ""
+            logs = deployment.show_logs().sort_values(by="time", ascending=False)
+
+            if logs and len(logs) > 0:
+                status = logs.iloc[0]["message"]
+
+            status = re.sub(r"[^a-zA-Z0-9]", " ", status)
+
             if data_science_work_request._error_message:
                 error_str = ""
                 for error in data_science_work_request._error_message:
                     error_str = error_str + " " + error.message
 
-            self.telemetry.record_event(
-                category=f"aqua/{model_type}/deployment/status",
-                action="FAILED",
-                detail=error_str,
-                value=model_name,
-                **telemetry_kwargs,
-            )
+                error_str = re.sub(r"[^a-zA-Z0-9]", " ", error_str)
+                telemetry_kwargs = {
+                    "ocid": ocid,
+                    "model_name": model_name,
+                    "work_request_error": error_str,
+                    "status": status,
+                }
+
+                self.telemetry.record_event(
+                    category=f"aqua/{model_type}/deployment/status",
+                    action="FAILED",
+                    **telemetry_kwargs,
+                )
+            else:
+                telemetry_kwargs = {
+                    "ocid": ocid,
+                    "model_name": model_name,
+                    "status": status,
+                }
+
+                self.telemetry.record_event(
+                    category=f"aqua/{model_type}/deployment/status",
+                    action="FAILED",
+                    **telemetry_kwargs,
+                )
 
         else:
-            self.telemetry.record_event_async(
+            telemetry_kwargs = {"ocid": ocid, "model_name": model_name}
+            self.telemetry.record_event(
                 category=f"aqua/{model_type}/deployment/status",
                 action="SUCCEEDED",
-                value=model_name,
                 **telemetry_kwargs,
             )

ads/common/oci_logging.py

@@ -1,22 +1,21 @@
 #!/usr/bin/env python
-# -*- coding: utf-8; -*-
 
 # Copyright (c) 2021, 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 datetime
 import logging
 import time
-from typing import Dict, Union, List
+from typing import Dict, List, Union
 
+import oci.exceptions
 import oci.logging
 import oci.loggingsearch
-import oci.exceptions
+
 from ads.common.decorator.utils import class_or_instance_method
 from ads.common.oci_mixin import OCIModelMixin, OCIWorkRequestMixin
 from ads.common.oci_resource import OCIResource, ResourceNotFoundError
 
-
 logger = logging.getLogger(__name__)
 
 # Maximum number of log records to be returned by default.
@@ -862,9 +861,7 @@ def tail(
             time_start=time_start,
             log_filter=log_filter,
         )
-        self._print(
-            sorted(tail_logs, key=lambda log: log["time"])
-        )
+        self._print(sorted(tail_logs, key=lambda log: log["time"]))
 
     def head(
         self,

ads/model/deployment/model_deployment.py

@@ -757,6 +757,8 @@ def watch(
         log_filter : str, optional
             Expression for filtering the logs. This will be the WHERE clause of the query.
             Defaults to None.
+        status_list : List[str], optional
+            List of status of model deployment. This is used to store list of status from logs.
 
         Returns
         -------

tests/unitary/with_extras/aqua/test_deployment.py

@@ -2376,29 +2376,84 @@ def test_validate_multimodel_deployment_feasibility_positive_single(
             "test_data/deployment/aqua_summary_multi_model_single.json",
         )
 
-    def test_get_deployment_status(self):
+    def test_get_deployment_status_success(self):
+        model_deployment = copy.deepcopy(TestDataset.model_deployment_object[0])
         deployment_id = "fakeid.datasciencemodeldeployment.oc1.iad.xxx"
         work_request_id = "fakeid.workrequest.oc1.iad.xxx"
         model_type = "custom"
         model_name = "model_name"
 
         with patch(
-            "ads.model.service.oci_datascience_model_deployment.DataScienceWorkRequest.__init__"
-        ) as mock_ds_work_request:
-            mock_ds_work_request.return_value = None
-            with patch(
-                "ads.model.service.oci_datascience_model_deployment.DataScienceWorkRequest.wait_work_request"
-            ) as mock_wait:
-                self.app.get_deployment_status(
-                    deployment_id, work_request_id, model_type, model_name
-                )
+            "ads.model.service.oci_datascience_model_deployment.DataScienceWorkRequest.__init__",
+            return_value=None,
+        ) as mock_ds_work_request, patch(
+            "ads.model.service.oci_datascience_model_deployment.DataScienceWorkRequest.wait_work_request"
+        ) as mock_wait:
+            self.app.get_deployment_status(
+                oci.data_science.models.ModelDeploymentSummary(**model_deployment),
+                work_request_id,
+                model_type,
+                model_name,
+            )
 
-                mock_ds_work_request.assert_called_with(work_request_id)
-                mock_wait.assert_called_with(
-                    progress_bar_description="Creating model deployment",
-                    max_wait_time=DEFAULT_WAIT_TIME,
-                    poll_interval=DEFAULT_POLL_INTERVAL,
-                )
+            mock_ds_work_request.assert_called_once_with(work_request_id)
+            mock_wait.assert_called_once_with(
+                progress_bar_description="Creating model deployment",
+                max_wait_time=DEFAULT_WAIT_TIME,
+                poll_interval=DEFAULT_POLL_INTERVAL,
+            )
+
+    def raise_exception(*args, **kwargs):
+        raise Exception("Work request failed")
+
+    def test_get_deployment_status_failed(self):
+        model_deployment = copy.deepcopy(TestDataset.model_deployment_object[0])
+        deployment_id = "fakeid.datasciencemodeldeployment.oc1.iad.xxx"
+        work_request_id = "fakeid.workrequest.oc1.iad.xxx"
+        model_type = "custom"
+        model_name = "model_name"
+        with patch(
+            "ads.telemetry.client.TelemetryClient.record_event"
+        ) as mock_record_event, patch(
+            "ads.aqua.modeldeployment.deployment.DataScienceWorkRequest"
+        ) as mock_ds_work_request_class, patch(
+            "ads.model.deployment.model_deployment.ModelDeployment.show_logs"
+        ) as mock_show_log:
+            mock_ds_work_request_instance = MagicMock()
+            mock_ds_work_request_class.return_value = mock_ds_work_request_instance
+
+            mock_ds_work_request_instance._error_message = [
+                MagicMock(message="Some error occurred")
+            ]
+
+            mock_ds_work_request_instance.wait_work_request.side_effect = (
+                self.raise_exception
+            )
+
+            logs_df = MagicMock()
+            logs_df.sort_values.return_value = logs_df
+            logs_df.empty = False
+            logs_df.iloc.__getitem__.return_value = {
+                "message": "Error: deployment failed!"
+            }
+            mock_show_log.return_value = logs_df
+
+            self.app.get_deployment_status(
+                ModelDeployment(),
+                work_request_id,
+                model_type,
+                model_name,
+            )
+            mock_record_event.assert_called_once()
+            args, kwargs = mock_record_event.call_args
+            self.assertEqual(kwargs["category"], f"aqua/{model_type}/deployment/status")
+            self.assertEqual(kwargs["action"], "FAILED")
+            self.assertIn("work_request_error", kwargs)
+            self.assertIn("status", kwargs)
+            self.assertIn("ocid", kwargs)
+            self.assertIn("model_name", kwargs)
+
+            mock_ds_work_request_class.assert_called_once_with(work_request_id)
 
 
 class TestBaseModelSpec: