Updated pr.

Author: lu-ohai

ads/common/oci_mixin.py

@@ -11,6 +11,7 @@
 import logging
 import os
 import re
+import time
 import traceback
 from datetime import date, datetime
 from typing import Callable, Optional, Union
@@ -20,7 +21,7 @@
 import yaml
 from ads.common import auth
 from ads.common.decorator.utils import class_or_instance_method
-from ads.common.utils import camel_to_snake
+from ads.common.utils import camel_to_snake, get_progress_bar
 from ads.config import COMPARTMENT_OCID
 from dateutil import tz
 from dateutil.parser import parse
@@ -29,6 +30,10 @@
 logger = logging.getLogger(__name__)
 
 LIFECYCLE_STOP_STATE = ("SUCCEEDED", "FAILED", "CANCELED", "DELETED")
+WORK_REQUEST_STOP_STATE = ("SUCCEEDED", "FAILED", "CANCELED")
+DEFAULT_WAIT_TIME = 1200
+DEFAULT_POLL_INTERVAL = 10
+DEFAULT_WORKFLOW_STEPS = 2
 
 
 class MergeStrategy(Enum):
@@ -931,6 +936,76 @@ def get_work_request_response(
             )
         return work_request_response
 
+    def wait_for_progress(
+        self, 
+        work_request_id: str, 
+        num_steps: int = DEFAULT_WORKFLOW_STEPS, 
+        max_wait_time: int = DEFAULT_WAIT_TIME, 
+        poll_interval: int = DEFAULT_POLL_INTERVAL
+    ):
+        """Waits for the work request progress bar to be completed.
+
+        Parameters
+        ----------
+        work_request_id: str
+            Work Request OCID.
+        num_steps: (int, optional). Defaults to 2.
+            Number of steps for the progress indicator.
+        max_wait_time: int
+            Maximum amount of time to wait in seconds (Defaults to 1200).
+            Negative implies infinite wait time.
+        poll_interval: int
+            Poll interval in seconds (Defaults to 10).
+
+        Returns
+        -------
+        None
+        """
+        work_request_logs = []
+
+        i = 0
+        start_time = time.time()
+        with get_progress_bar(num_steps) as progress:
+            seconds_since = time.time() - start_time
+            exceed_max_time = max_wait_time > 0 and seconds_since >= max_wait_time
+            if exceed_max_time:
+                logger.error(
+                    f"Max wait time ({max_wait_time} seconds) exceeded."
+                )
+            while not exceed_max_time and (not work_request_logs or len(work_request_logs) < num_steps):
+                time.sleep(poll_interval)
+                new_work_request_logs = []
+
+                try:
+                    work_request = self.client.get_work_request(work_request_id).data
+                    work_request_logs = self.client.list_work_request_logs(
+                        work_request_id
+                    ).data
+                except Exception as ex:
+                    logger.warn(ex)
+
+                new_work_request_logs = (
+                    work_request_logs[i:] if work_request_logs else []
+                )
+
+                for wr_item in new_work_request_logs:
+                    progress.update(wr_item.message)
+                    i += 1
+
+                if work_request and work_request.status in WORK_REQUEST_STOP_STATE:
+                    if work_request.status != "SUCCEEDED":
+                        if new_work_request_logs:
+                            raise Exception(new_work_request_logs[-1].message)
+                        else:
+                            raise Exception(
+                                "Error occurred in attempt to perform the operation. "
+                                "Check the service logs to get more details. "
+                                f"{work_request}"
+                            )
+                    else:
+                        break
+            progress.update("Done")
+
 
 class OCIModelWithNameMixin:
     """Mixin class to operate OCI model which contains name property."""

ads/model/service/oci_datascience_model_deployment.py

@@ -5,20 +5,17 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 from functools import wraps
-import time
 import logging
 from typing import Callable, List
 from ads.common.oci_datascience import OCIDataScienceMixin
-from ads.common import utils as progress_bar_utils
+from ads.common.oci_mixin import OCIWorkRequestMixin
 from ads.config import PROJECT_OCID
-from ads.model.deployment.common import utils
 from ads.model.deployment.common.utils import OCIClientManager, State
 import oci
 
 from oci.data_science.models import (
     CreateModelDeploymentDetails,
     UpdateModelDeploymentDetails,
-    WorkRequest
 )
 
 DEFAULT_WAIT_TIME = 1200
@@ -84,6 +81,7 @@ class MissingModelDeploymentWorkflowIdError(Exception):
 
 class OCIDataScienceModelDeployment(
     OCIDataScienceMixin,
+    OCIWorkRequestMixin,
     oci.data_science.models.ModelDeployment,
 ):
     """Represents an OCI Data Science Model Deployment.
@@ -190,7 +188,7 @@ def activate(
             self.workflow_req_id = response.headers.get("opc-work-request-id", None)
 
             try:
-                self._wait_for_work_request(
+                self.wait_for_progress(
                     self.workflow_req_id, 
                     ACTIVATE_WORKFLOW_STEPS, 
                     max_wait_time, 
@@ -238,7 +236,7 @@ def create(
             self.workflow_req_id = response.headers.get("opc-work-request-id", None)
 
             try:
-                self._wait_for_work_request(
+                self.wait_for_progress(
                     self.workflow_req_id, 
                     CREATE_WORKFLOW_STEPS, 
                     max_wait_time, 
@@ -289,7 +287,7 @@ def deactivate(
             self.workflow_req_id = response.headers.get("opc-work-request-id", None)
 
             try:
-                self._wait_for_work_request(
+                self.wait_for_progress(
                     self.workflow_req_id, 
                     DEACTIVATE_WORKFLOW_STEPS, 
                     max_wait_time, 
@@ -340,7 +338,7 @@ def delete(
             self.workflow_req_id = response.headers.get("opc-work-request-id", None)
 
             try:
-                self._wait_for_work_request(
+                self.wait_for_progress(
                     self.workflow_req_id, 
                     DELETE_WORKFLOW_STEPS, 
                     max_wait_time, 
@@ -483,77 +481,3 @@ def from_id(cls, model_deployment_id: str) -> "OCIDataScienceModelDeployment":
             An instance of `OCIDataScienceModelDeployment`.
         """
         return super().from_ocid(model_deployment_id)
-
-    def _wait_for_work_request(
-        self, 
-        work_request_id: str, 
-        num_steps: int = DELETE_WORKFLOW_STEPS, 
-        max_wait_time: int = DEFAULT_WAIT_TIME, 
-        poll_interval: int = DEFAULT_POLL_INTERVAL
-    ) -> None:
-        """Waits for the work request to be completed.
-
-        Parameters
-        ----------
-        work_request_id: str
-            Work Request OCID.
-        num_steps: (int, optional). Defaults to 6.
-            Number of steps for the progress indicator.
-        max_wait_time: int
-            Maximum amount of time to wait in seconds (Defaults to 1200).
-            Negative implies infinite wait time.
-        poll_interval: int
-            Poll interval in seconds (Defaults to 10).
-
-        Returns
-        -------
-        None
-        """
-        STOP_STATE = (
-            WorkRequest.STATUS_SUCCEEDED,
-            WorkRequest.STATUS_CANCELED,
-            WorkRequest.STATUS_FAILED,
-        )
-        work_request_logs = []
-
-        i = 0
-        start_time = time.time()
-        with progress_bar_utils.get_progress_bar(num_steps) as progress:
-            exceed_max_time = max_wait_time > 0 and utils.seconds_since(start_time) >= max_wait_time
-            if exceed_max_time:
-                logger.error(
-                    f"Max wait time ({max_wait_time} seconds) exceeded."
-                )
-            while not exceed_max_time and (not work_request_logs or len(work_request_logs) < num_steps):
-                time.sleep(poll_interval)
-                new_work_request_logs = []
-
-                try:
-                    work_request = self.client.get_work_request(work_request_id).data
-                    work_request_logs = self.client.list_work_request_logs(
-                        work_request_id
-                    ).data
-                except Exception as ex:
-                    logger.warn(ex)
-
-                new_work_request_logs = (
-                    work_request_logs[i:] if work_request_logs else []
-                )
-
-                for wr_item in new_work_request_logs:
-                    progress.update(wr_item.message)
-                    i += 1
-
-                if work_request and work_request.status in STOP_STATE:
-                    if work_request.status != WorkRequest.STATUS_SUCCEEDED:
-                        if new_work_request_logs:
-                            raise Exception(new_work_request_logs[-1].message)
-                        else:
-                            raise Exception(
-                                "Error occurred in attempt to perform the operation. "
-                                "Check the service logs to get more details. "
-                                f"{work_request}"
-                            )
-                    else:
-                        break
-            progress.update("Done")

tests/unitary/default_setup/model_deployment/test_oci_datascience_model_deployment.py

@@ -10,7 +10,7 @@
     ModelDeployment,
 )
 from ads.common.oci_datascience import OCIDataScienceMixin
-from ads.common.oci_mixin import OCIModelMixin
+from ads.common.oci_mixin import OCIModelMixin, OCIWorkRequestMixin
 
 from ads.model.service.oci_datascience_model_deployment import (
     ACTIVATE_WORKFLOW_STEPS,
@@ -116,7 +116,7 @@ def test_activate_with_waiting(self):
                 response.data = data
                 mock_get.return_value = response
                 with patch.object(
-                    OCIDataScienceModelDeployment, "_wait_for_work_request"
+                    OCIWorkRequestMixin, "wait_for_progress"
                 ) as mock_wait:
                     with patch.object(
                         OCIDataScienceModelDeployment, "sync"
@@ -169,7 +169,7 @@ def test_deactivate_with_waiting(self):
                 response.data = data
                 mock_get.return_value = response
                 with patch.object(
-                    OCIDataScienceModelDeployment, "_wait_for_work_request"
+                    OCIWorkRequestMixin, "wait_for_progress"
                 ) as mock_wait:
                     with patch.object(
                         OCIDataScienceModelDeployment, "sync"
@@ -243,7 +243,7 @@ def test_create_with_waiting(self):
                     )
                     mock_to_oci_mode.return_value = oci_model_deployment
                     with patch.object(
-                        OCIDataScienceModelDeployment, "_wait_for_work_request"
+                        OCIWorkRequestMixin, "wait_for_progress"
                     ) as mock_wait:
                         with patch("json.loads") as mock_json_load:
                             create_model_deployment_details = MagicMock()
@@ -354,7 +354,7 @@ def test_delete_with_waiting(self):
             }
             mock_delete.return_value = response
             with patch.object(
-                OCIDataScienceModelDeployment, "_wait_for_work_request"
+                OCIWorkRequestMixin, "wait_for_progress"
             ) as mock_wait:
                 with patch.object(
                     oci.data_science.DataScienceClient, "get_model_deployment"