Updated pr.

lu-ohai · lu-ohai · commit 8e24e15b5023 · 2023-12-14T15:33:28.000-05:00
diff --git a/ads/common/oci_mixin.py b/ads/common/oci_mixin.py
@@ -1040,96 +1040,3 @@ def from_name(cls, name: str, compartment_id: Optional[str] = None):
         if not res:
             raise OCIModelNotExists()
         return cls.from_oci_model(res[0])
-
-
-class ADSWorkRequest(OCIClientMixin):
-
-    def __init__(self, id: str, description: str = "Processing"):
-        self.id = id
-        self._description = description
-        self._percentage = 0
-        self._status = None
-
-    def _sync(self):
-        work_request = self.client.get_work_request(self.id).data
-        work_request_logs = self.client.list_work_request_logs(
-            self.id
-        ).data
-
-        self._percentage= work_request.percent_complete
-        self._status = work_request.status
-        self._description = work_request_logs[:-1]
-
-    def watch(
-        self, 
-        progress_callback: Callable,
-        max_wait_time: int,
-        poll_interval: int,
-    ):
-        previous_percent_complete = 0
-        previous_log = None
-
-        start_time = time.time()
-        while self._percentage < 100:
-
-            seconds_since = time.time() - start_time
-            if max_wait_time > 0 and seconds_since >= max_wait_time:
-                logger.error(f"Max wait time ({max_wait_time} seconds) exceeded.")
-                return
-
-            time.sleep(poll_interval)
-
-            try:
-                self._sync()
-            except Exception as ex:
-                logger.warn(ex)
-                continue
-
-            percent_change = self._percentage - previous_percent_complete
-            previous_percent_complete = self._percentage
-            description = self._description if previous_log != self._description else ""
-            progress_callback(percent_change, description)
-            previous_log = self._description
-
-            if self._status in WORK_REQUEST_STOP_STATE:
-                if self._status != oci.work_requests.models.WorkRequest.STATUS_SUCCEEDED:
-                    if self._description:
-                        raise Exception(self._description)
-                    else:
-                        raise Exception(
-                            "Error occurred in attempt to perform the operation. "
-                            "Check the service logs to get more details. "
-                        )
-                else:
-                    break
-
-        progress_callback(0, "Done")
-
-
-def wait_work_request(
-    id: str, 
-    progress_bar_description: str, 
-    max_wait_time: int=DEFAULT_WAIT_TIME,
-    poll_interval: int=DEFAULT_POLL_INTERVAL
-):
-    ads_work_request = ADSWorkRequest(id)
-
-    with tqdm(
-        leave=False,
-        mininterval=0,
-        file=sys.stdout,
-        desc=progress_bar_description,
-    ) as pbar:
-
-        def progress_callback(percent_change, description):
-            if percent_change != 0:
-                pbar.update(percent_change)
-            if description:
-                pbar.set_description(description)
-
-        ads_work_request.watch(
-            progress_callback,
-            max_wait_time,
-            poll_interval
-        )
- 
diff --git a/ads/common/work_request.py b/ads/common/work_request.py
@@ -0,0 +1,123 @@
+
+import logging
+import sys
+import time
+from typing import Callable
+
+import oci
+from oci import Signer
+from tqdm.auto import tqdm
+from ads.common.oci_datascience import OCIDataScienceMixin
+
+logger = logging.getLogger(__name__)
+
+WORK_REQUEST_STOP_STATE = ("SUCCEEDED", "FAILED", "CANCELED")
+DEFAULT_WAIT_TIME = 1200
+DEFAULT_POLL_INTERVAL = 10
+WORK_REQUEST_PERCENTAGE = 100
+# default tqdm progress bar format: 
+# {l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, ' '{rate_fmt}{postfix}]
+# customize the bar format to remove the {n_fmt}/{total_fmt} from the right side
+DEFAULT_BAR_FORMAT = '{l_bar}{bar}| [{elapsed}<{remaining}, ' '{rate_fmt}{postfix}]'
+
+
+class ADSWorkRequest(OCIDataScienceMixin):
+
+    def __init__(
+        self, 
+        id: str, 
+        description: str = "Processing",
+        config: dict = None, 
+        signer: Signer = None, 
+        client_kwargs: dict = None, 
+        **kwargs
+    ) -> None:
+        self.id = id
+        self._description = description
+        self._percentage = 0
+        self._status = None
+        super().__init__(config, signer, client_kwargs, **kwargs)
+        
+
+    def _sync(self):
+        work_request = self.client.get_work_request(self.id).data
+        work_request_logs = self.client.list_work_request_logs(
+            self.id
+        ).data
+
+        self._percentage= work_request.percent_complete
+        self._status = work_request.status
+        self._description = work_request_logs[-1].message if work_request_logs else "Processing"
+
+    def watch(
+        self, 
+        progress_callback: Callable,
+        max_wait_time: int,
+        poll_interval: int,
+    ):
+        previous_percent_complete = 0
+
+        start_time = time.time()
+        while self._percentage < 100:
+
+            seconds_since = time.time() - start_time
+            if max_wait_time > 0 and seconds_since >= max_wait_time:
+                logger.error(f"Max wait time ({max_wait_time} seconds) exceeded.")
+                return
+
+            time.sleep(poll_interval)
+
+            try:
+                self._sync()
+            except Exception as ex:
+                logger.warn(ex)
+                continue
+
+            percent_change = self._percentage - previous_percent_complete
+            previous_percent_complete = self._percentage
+            progress_callback(percent_change, self._description)
+
+            if self._status in WORK_REQUEST_STOP_STATE:
+                if self._status != oci.work_requests.models.WorkRequest.STATUS_SUCCEEDED:
+                    if self._description:
+                        raise Exception(self._description)
+                    else:
+                        raise Exception(
+                            "Error occurred in attempt to perform the operation. "
+                            "Check the service logs to get more details. "
+                        )
+                else:
+                    break
+
+        progress_callback(0, "Done")
+
+
+def wait_work_request(
+    id: str, 
+    progress_bar_description: str, 
+    max_wait_time: int=DEFAULT_WAIT_TIME,
+    poll_interval: int=DEFAULT_POLL_INTERVAL
+):
+    ads_work_request = ADSWorkRequest(id)
+
+    with tqdm(
+        total=WORK_REQUEST_PERCENTAGE,
+        leave=False,
+        mininterval=0,
+        file=sys.stdout,
+        desc=progress_bar_description,
+        bar_format=DEFAULT_BAR_FORMAT
+    ) as pbar:
+
+        def progress_callback(percent_change, description):
+            if percent_change != 0:
+                pbar.update(percent_change)
+            if description:
+                pbar.set_description(description)
+
+        ads_work_request.watch(
+            progress_callback,
+            max_wait_time,
+            poll_interval
+        )
+ 
diff --git a/ads/model/service/oci_datascience_model.py b/ads/model/service/oci_datascience_model.py
@@ -17,6 +17,7 @@
 from ads.common.oci_mixin import OCIWorkRequestMixin
 from ads.common.oci_resource import SEARCH_TYPE, OCIResource
 from ads.common.utils import extract_region
+from ads.common.work_request import wait_work_request
 from ads.model.deployment import ModelDeployment
 from oci.data_science.models import (
     ArtifactExportDetailsObjectStorage,
@@ -361,10 +362,7 @@ def import_model_artifact(self, bucket_uri: str, region: str = None) -> None:
             ).headers["opc-work-request-id"]
 
             # Show progress of importing artifacts
-            self._wait_for_work_request(
-                work_request_id=work_request_id,
-                num_steps=2,
-            )
+            wait_work_request(work_request_id)
         except ServiceError as ex:
             if ex.status == 404:
                 raise ModelArtifactNotFoundError()
@@ -408,10 +406,7 @@ def export_model_artifact(self, bucket_uri: str, region: str = None):
         ).headers["opc-work-request-id"]
 
         # Show progress of exporting model artifacts
-        self._wait_for_work_request(
-            work_request_id=work_request_id,
-            num_steps=2,
-        )
+        wait_work_request(work_request_id)
 
     @check_for_model_id(
         msg="Model needs to be saved to the Model Catalog before it can be updated."
@@ -540,63 +535,3 @@ def from_id(cls, ocid: str) -> "OCIDataScienceModel":
         if not ocid:
             raise ValueError("Model OCID not provided.")
         return super().from_ocid(ocid)
-
-    def _wait_for_work_request(self, work_request_id: str, num_steps: int = 3) -> None:
-        """Waits for the work request to be completed.
-
-        Parameters
-        ----------
-        work_request_id: str
-            Work Request OCID.
-        num_steps: (int, optional). Defaults to 3.
-            Number of steps for the progress indicator.
-
-        Returns
-        -------
-        None
-        """
-        STOP_STATE = (
-            WorkRequest.STATUS_SUCCEEDED,
-            WorkRequest.STATUS_CANCELED,
-            WorkRequest.STATUS_FAILED,
-        )
-        work_request_logs = []
-
-        i = 0
-        with utils.get_progress_bar(num_steps) as progress:
-            while not work_request_logs or len(work_request_logs) < num_steps:
-                time.sleep(_REQUEST_INTERVAL_IN_SEC)
-                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
-
-            while i < num_steps:
-                progress.update()
-                i += 1
diff --git a/ads/model/service/oci_datascience_model_deployment.py b/ads/model/service/oci_datascience_model_deployment.py
@@ -8,7 +8,7 @@
 import logging
 from typing import Callable, List
 from ads.common.oci_datascience import OCIDataScienceMixin
-from ads.common.oci_mixin import OCIWorkRequestMixin, wait_work_request
+from ads.common.work_request import wait_work_request
 from ads.config import PROJECT_OCID
 from ads.model.deployment.common.utils import OCIClientManager, State
 import oci
@@ -77,7 +77,6 @@ class MissingModelDeploymentWorkflowIdError(Exception):  # pragma: no cover
 
 class OCIDataScienceModelDeployment(
     OCIDataScienceMixin,
-    OCIWorkRequestMixin,
     oci.data_science.models.ModelDeployment,
 ):
     """Represents an OCI Data Science Model Deployment.
@@ -192,7 +191,7 @@ def activate(
                 try:
                     wait_work_request(
                         self.workflow_req_id,
-                        f"Activating model deployment: {self.display_name}",
+                        "Activating model deployment",
                         max_wait_time, 
                         poll_interval
                     )
@@ -243,7 +242,7 @@ def create(
             try:
                 wait_work_request(
                     self.workflow_req_id,
-                    f"Creating model deployment: {self.display_name}",
+                    "Creating model deployment",
                     max_wait_time, 
                     poll_interval
                 )
@@ -297,7 +296,7 @@ def deactivate(
                 try:
                     wait_work_request(
                         self.workflow_req_id,
-                        f"Deactivating model deployment: {self.display_name}",
+                        "Deactivating model deployment",
                         max_wait_time, 
                         poll_interval
                     )
@@ -366,7 +365,7 @@ def delete(
             try:
                 wait_work_request(
                     self.workflow_req_id,
-                    f"Deleting model deployment: {self.display_name}",
+                    "Deleting model deployment",
                     max_wait_time, 
                     poll_interval
                 )
