Merge branch 'develop' into ODSC-40043/dataflow_archive

Author: qiuosier

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/deployment/model_deployment.py

@@ -1532,6 +1532,9 @@ def _build_model_deployment_configuration_details(self) -> Dict:
                 infrastructure.CONST_MEMORY_IN_GBS: infrastructure.shape_config_details.get(
                     "memory_in_gbs", None
                 )
+                or infrastructure.shape_config_details.get(
+                    "memoryInGBs", None
+                )
                 or MODEL_DEPLOYMENT_INSTANCE_MEMORY_IN_GBS,
             }