[AQUA Telemetry] Update MD Tracking

ads/aqua/constants.py

@@ -54,6 +54,8 @@
 SERVICE_MANAGED_CONTAINER_URI_SCHEME = "dsmc://"
 SUPPORTED_FILE_FORMATS = ["jsonl"]
 MODEL_BY_REFERENCE_OSS_PATH_KEY = "artifact_location"
+DEFAULT_WAIT_TIME = 1200
+DEFAULT_POLL_INTERVAL = 10
 
 CONSOLE_LINK_RESOURCE_TYPE_MAPPING = {
     "datasciencemodel": "models",

ads/aqua/modeldeployment/deployment.py

@@ -8,6 +8,8 @@
 from typing import Dict, List, Optional, Union
 
 from cachetools import TTLCache, cached
+import concurrent
+from ads.common.work_request import DataScienceWorkRequest
 from oci.data_science.models import ModelDeploymentShapeSummary
 from pydantic import ValidationError
 
@@ -43,6 +45,8 @@
     MODEL_BY_REFERENCE_OSS_PATH_KEY,
     MODEL_NAME_DELIMITER,
     UNKNOWN_DICT,
+    DEFAULT_WAIT_TIME, 
+    DEFAULT_POLL_INTERVAL
 )
 from ads.aqua.data import AquaResourceIdentifier
 from ads.aqua.model import AquaModelApp
@@ -80,6 +84,9 @@
 from ads.model.model_metadata import ModelCustomMetadataItem
 from ads.telemetry import telemetry
 
+THREAD_POOL_SIZE = 16
+thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=THREAD_POOL_SIZE)
+
 
 class AquaDeploymentApp(AquaApp):
     """Provides a suite of APIs to interact with Aqua model deployments within the Oracle
@@ -780,11 +787,18 @@ def _create_deployment(
             .with_runtime(container_runtime)
         ).deploy(wait_for_completion=False)
 
-        deployment_id = deployment.id
+
+
+        deployment_id = deployment.id()
         logger.info(
             f"Aqua model deployment {deployment_id} created for model {aqua_model_id}."
         )
 
+        thread_pool.submit(self.get_deployment_status, 
+                            model_deployment_id=deployment_id, 
+                            work_request_id=deployment.dsc_model_deployment.workflow_req_id,
+                            model_type=model_type) 
+
         # we arbitrarily choose last 8 characters of OCID to identify MD in telemetry
         telemetry_kwargs = {"ocid": get_ocid_substring(deployment_id, key_len=8)}
 
@@ -1309,4 +1323,33 @@ def list_shapes(self, **kwargs) -> List[ComputeShapeSummary]:
                 or gpu_specs.shapes.get(oci_shape.name.upper()),
             )
             for oci_shape in oci_shapes
-        ]
+        ]
+
+
+    def get_deployment_status(self,model_deployment_id: str, work_request_id : str, model_type : str) :
+
+        telemetry_kwargs = {"ocid": get_ocid_substring(model_deployment_id, key_len=8)}
+
+        try:
+            DataScienceWorkRequest(work_request_id).wait_work_request(
+                progress_bar_description="Creating model deployment",
+                max_wait_time=DEFAULT_WAIT_TIME, 
+                poll_interval=DEFAULT_POLL_INTERVAL
+            )
+        except Exception as e:
+            logger.error(
+                "Error while trying to create model deployment: " + str(e)
+            )
+            self.telemetry.record_event_async(
+                category=f"aqua/{model_type}/deployment/status",
+                action="FAILED",
+                detail="Error creating model deployment"
+                **telemetry_kwargs
+            )
+
+        self.telemetry.record_event_async(
+                category=f"aqua/{model_type}/deployment/status",
+                action="SUCCEEDED",
+                detail=" Create model deployment successful", 
+                **telemetry_kwargs
+            )

ads/common/work_request.py

@@ -38,6 +38,7 @@ def __init__(
         config: dict = None, 
         signer: Signer = None, 
         client_kwargs: dict = None, 
+        _error_message: str = None,
         **kwargs
     ) -> None:
         """Initializes ADSWorkRequest object.
@@ -65,6 +66,7 @@ def __init__(
         self._description = description
         self._percentage = 0
         self._status = None
+        _error_message = _error_message
         super().__init__(config, signer, client_kwargs, **kwargs)
 
 
@@ -78,6 +80,7 @@ def _sync(self):
         self._percentage= work_request.percent_complete
         self._status = work_request.status
         self._description = work_request_logs[-1].message if work_request_logs else "Processing"
+        if work_request.status == 'FAILED' : self._error_message = self.client.list_work_request_errors
 
     def watch(
         self, 

ads/model/service/oci_datascience_model_deployment.py

@@ -185,9 +185,9 @@ def activate(
                 self.id,
             )
 
+
+            self.workflow_req_id = response.headers.get("opc-work-request-id", None)
             if wait_for_completion:
-                self.workflow_req_id = response.headers.get("opc-work-request-id", None)
-
                 try:
                     DataScienceWorkRequest(self.workflow_req_id).wait_work_request(
                         progress_bar_description="Activating model deployment",

ads/telemetry/client.py

@@ -119,4 +119,4 @@ def record_event_async(
         Thread
             A started thread to send a head request to generate an event record.
         """
-        thread_pool.submit(self.record_event, args=(category, action, detail), kwargs=kwargs)
+        thread_pool.submit(self.record_event, category, action, detail, **kwargs)