PR review comments

Author: agrimk

ads/aqua/modeldeployment/deployment.py

@@ -88,6 +88,7 @@
 
 from ads.model.model_metadata import ModelCustomMetadataItem
 from ads.telemetry import telemetry
+from ads.common.decorator.threaded import thread_pool
 
 class AquaDeploymentApp(AquaApp):
     """Provides a suite of APIs to interact with Aqua model deployments within the Oracle
@@ -790,17 +791,15 @@ def _create_deployment(
             .with_runtime(container_runtime)
         ).deploy(wait_for_completion=False)
 
-       
-
         deployment_id = deployment.id
         logger.info(
-            f"Aqua model deployment {deployment_id} created for model {aqua_model_id}."
+            f"Aqua model deployment {deployment_id} created for model {aqua_model_id}. Work request Id is {deployment.dsc_model_deployment.workflow_req_id}"
         )
-
-        self.get_deployment_status( 
-            model_deployment_id=deployment_id, 
-            work_request_id=deployment.dsc_model_deployment.workflow_req_id,
-            model_type=model_type)
+        
+        thread_pool.submit( self.get_deployment_status ,
+            deployment_id, 
+            deployment.dsc_model_deployment.workflow_req_id,
+            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,7 +1324,6 @@ 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.
 

ads/model/service/oci_datascience_model_deployment.py

@@ -233,11 +233,9 @@ def create(
         response = self.client.create_model_deployment(create_model_deployment_details)
         self.update_from_oci_model(response.data)
         logger.info(f"Creating model deployment `{self.id}`.")
-        print(f"Model Deployment OCID: {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="Creating model deployment",
@@ -287,10 +285,8 @@ def deactivate(
             response = self.client.deactivate_model_deployment(
                 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="Deactivating model deployment",
@@ -355,10 +351,9 @@ def delete(
         response = self.client.delete_model_deployment(
             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="Deleting model deployment",

tests/unitary/with_extras/aqua/test_deployment.py

@@ -11,6 +11,8 @@
 from importlib import reload
 from unittest.mock import MagicMock, patch
 
+from ads.aqua.modeldeployment.constants import DEFAULT_POLL_INTERVAL, DEFAULT_WAIT_TIME
+from ads.model.service.oci_datascience_model_deployment import OCIDataScienceModelDeployment
 import oci
 import pytest
 from oci.data_science.models import (
@@ -2282,3 +2284,26 @@ def test_validate_multimodel_deployment_feasibility_positive_single(
             total_gpus,
             "test_data/deployment/aqua_summary_multi_model_single.json",
         )
+
+    def test_get_deployment_status(self) :
+        deployment_id = "fakeid.datasciencemodeldeployment.oc1.iad.xxx"
+        work_request_id = "fakeid.workrequest.oc1.iad.xxx"
+        model_type = "custom"
+        
+        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
+                )
+
+                mock_ds_work_request.assert_called_with("test")
+                mock_wait.assert_called_with(
+                    progress_bar_description='Creating model deployment', 
+                    max_wait_time=DEFAULT_WAIT_TIME, 
+                    poll_interval=DEFAULT_POLL_INTERVAL
+                )