PR review comments

Author: agrimk

ads/aqua/constants.py

@@ -54,8 +54,6 @@
 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/constants.py

@@ -8,3 +8,5 @@
 
 This module contains constants used in Aqua Model Deployment.
 """
+DEFAULT_WAIT_TIME = 1200
+DEFAULT_POLL_INTERVAL = 10

ads/aqua/modeldeployment/deployment.py

@@ -8,7 +8,8 @@
 from typing import Dict, List, Optional, Union
 
 from cachetools import TTLCache, cached
-import concurrent
+from ads.aqua.modeldeployment.constants import DEFAULT_POLL_INTERVAL, DEFAULT_WAIT_TIME
+from ads.common.decorator.threaded import threaded
 from ads.common.work_request import DataScienceWorkRequest
 from oci.data_science.models import ModelDeploymentShapeSummary
 from pydantic import ValidationError
@@ -44,10 +45,9 @@
     AQUA_MULTI_MODEL_CONFIG,
     MODEL_BY_REFERENCE_OSS_PATH_KEY,
     MODEL_NAME_DELIMITER,
-    UNKNOWN_DICT,
-    DEFAULT_WAIT_TIME, 
-    DEFAULT_POLL_INTERVAL
+    UNKNOWN_DICT
 )
+
 from ads.aqua.data import AquaResourceIdentifier
 from ads.aqua.model import AquaModelApp
 from ads.aqua.model.constants import AquaModelMetadataKeys, ModelCustomMetadataFields
@@ -81,13 +81,10 @@
     ModelDeploymentInfrastructure,
     ModelDeploymentMode,
 )
+
 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
     Cloud Infrastructure Data Science service, serving as an interface for deploying
@@ -789,15 +786,15 @@ def _create_deployment(
 
 
 
-        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) 
+        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)}
@@ -1325,13 +1322,31 @@ def list_shapes(self, **kwargs) -> List[ComputeShapeSummary]:
             for oci_shape in oci_shapes
         ]
 
+    @threaded()
+    def get_deployment_status(self,model_deployment_id: str, work_request_id : str, model_type : str) -> None:
+        """Waits for the data science  model deployment to be completed and log its status in telemetry.
+
+        Parameters
+        ----------
         
-    def get_deployment_status(self,model_deployment_id: str, work_request_id : str, model_type : str) :
+        model_deployment_id: str
+            The id of the deployed aqua model.
+        work_request_id: str
+            The work request Id of the model deployment.
+        model_type: str
+            The type of aqua model to be deployed. Allowed values are: `custom`, `service` and `multi_model`.
 
+        Returns
+        -------
+        AquaDeployment
+            An Aqua deployment instance.
+        """
         telemetry_kwargs = {"ocid": get_ocid_substring(model_deployment_id, key_len=8)}
+
+        data_science_work_request:DataScienceWorkRequest = DataScienceWorkRequest(work_request_id) 
 
         try:
-            DataScienceWorkRequest(work_request_id).wait_work_request(
+            data_science_work_request.wait_work_request(
                 progress_bar_description="Creating model deployment",
                 max_wait_time=DEFAULT_WAIT_TIME, 
                 poll_interval=DEFAULT_POLL_INTERVAL
@@ -1343,13 +1358,13 @@ def get_deployment_status(self,model_deployment_id: str, work_request_id : str,
             self.telemetry.record_event_async(
                 category=f"aqua/{model_type}/deployment/status",
                 action="FAILED",
-                detail="Error creating model deployment"
+                detail=data_science_work_request._error_message
                 **telemetry_kwargs
             )
+        else :
+            self.telemetry.record_event_async(
+                    category=f"aqua/{model_type}/deployment/status",
+                    action="SUCCEEDED",
+                    **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

@@ -37,8 +37,7 @@ def __init__(
         description: str = "Processing",
         config: dict = None, 
         signer: Signer = None, 
-        client_kwargs: dict = None, 
-        _error_message: str = None,
+        client_kwargs: dict = None,
         **kwargs
     ) -> None:
         """Initializes ADSWorkRequest object.
@@ -66,7 +65,7 @@ def __init__(
         self._description = description
         self._percentage = 0
         self._status = None
-        _error_message = _error_message
+        _error_message = None
         super().__init__(config, signer, client_kwargs, **kwargs)