Merge branch 'main' into ODSC-34939/explicit_download_artifact

Author: VipulMascarenhas

ads/catalog/model.py

@@ -82,11 +82,11 @@
 _WORK_REQUEST_INTERVAL_IN_SEC = 3
 
 
-class ModelWithActiveDeploymentError(Exception):   # pragma: no cover
+class ModelWithActiveDeploymentError(Exception):  # pragma: no cover
     pass
 
 
-class ModelArtifactSizeError(Exception):   # pragma: no cover
+class ModelArtifactSizeError(Exception):  # pragma: no cover
     def __init__(self, max_artifact_size: str):
         super().__init__(
             f"The model artifacts size is greater than `{max_artifact_size}`. "
@@ -391,7 +391,6 @@ def show_in_notebook(self, display_format: str = "dataframe") -> None:
             Nothing.
         """
         if display_format == "dataframe":
-
             from IPython.core.display import display
 
             display(self.to_dataframe())
@@ -1560,9 +1559,9 @@ def _wait_for_work_request(
                 if work_request_logs:
                     new_work_request_logs = work_request_logs[i:]
 
-                for wr_item in new_work_request_logs:
-                    progress.update(wr_item.message)
-                    i += 1
+                    for wr_item in new_work_request_logs:
+                        progress.update(wr_item.message)
+                        i += 1
 
                 if work_request.data.status in STOP_STATE:
                     if work_request.data.status != WorkRequest.STATUS_SUCCEEDED:

ads/config.py

@@ -45,6 +45,7 @@
     NB_SESSION_OCID or JOB_RUN_OCID or MD_OCID or PIPELINE_RUN_OCID or DATAFLOW_RUN_OCID
 )
 NO_CONTAINER = os.environ.get("NO_CONTAINER")
+TMPDIR = os.environ.get("TMPDIR")
 
 
 def export(

ads/jobs/templates/driver_utils.py

@@ -413,7 +413,7 @@ def run_command(
                     logger.log(level=level, msg=msg)
             # Add a small delay so that
             # outputs from the subsequent code will have different timestamp for oci logging
-            time.sleep(0.05)
+            time.sleep(0.02)
         if check and process.returncode != 0:
             # If there is an error, exit the main process with the same return code.
             sys.exit(process.returncode)

ads/model/artifact_downloader.py

@@ -63,6 +63,7 @@ def download(self):
                     "Set `force_overwrite` to `True` if you wish to overwrite."
                 )
             shutil.rmtree(self.target_dir)
+            os.mkdir(self.target_dir)
         with utils.get_progress_bar(
             ArtifactDownloader.PROGRESS_STEPS_COUNT + self.PROGRESS_STEPS_COUNT
         ) as progress:
@@ -85,15 +86,16 @@ def _download(self):
         """Downloads model artifacts."""
         self.progress.update("Importing model artifacts from catalog")
         zip_content = self.dsc_model.get_model_artifact_content()
-        with tempfile.TemporaryDirectory() as temp_dir:
-            self.progress.update("Copying model artifacts to the artifact directory")
-            zip_file_path = os.path.join(temp_dir, f"{str(uuid.uuid4())}.zip")
-            with open(zip_file_path, "wb") as zip_file:
-                zip_file.write(zip_content)
-            self.progress.update("Extracting model artifacts")
-            with ZipFile(zip_file_path) as zip_file:
-                zip_file.extractall(self.target_dir)
-
+        self.progress.update("Copying model artifacts to the artifact directory")
+        
+        zip_file_path = os.path.join(self.target_dir, f"{str(uuid.uuid4())}.zip")
+        with open(zip_file_path, "wb") as zip_file:
+            zip_file.write(zip_content)
+        self.progress.update("Extracting model artifacts")
+        with ZipFile(zip_file_path) as zip_file:
+            zip_file.extractall(self.target_dir)
+        
+        utils.remove_file(zip_file_path)
 
 class LargeArtifactDownloader(ArtifactDownloader):
     PROGRESS_STEPS_COUNT = 4
@@ -157,17 +159,19 @@ def _download(self):
 
         self.dsc_model.import_model_artifact(bucket_uri=bucket_uri, region=self.region)
         self.progress.update("Copying model artifacts to the artifact directory")
-        with tempfile.TemporaryDirectory() as temp_dir:
-            zip_file_path = os.path.join(temp_dir, f"{str(uuid.uuid4())}.zip")
-            zip_file_path = utils.copy_file(
-                uri_src=bucket_uri,
-                uri_dst=zip_file_path,
-                auth=self.auth,
-                progressbar_description="Copying model artifacts to the artifact directory",
-            )
-            self.progress.update("Extracting model artifacts")
-            with ZipFile(zip_file_path) as zip_file:
-                zip_file.extractall(self.target_dir)
+        
+        zip_file_path = os.path.join(self.target_dir, f"{str(uuid.uuid4())}.zip")
+        zip_file_path = utils.copy_file(
+            uri_src=bucket_uri,
+            uri_dst=zip_file_path,
+            auth=self.auth,
+            progressbar_description="Copying model artifacts to the artifact directory",
+        )
+        self.progress.update("Extracting model artifacts")
+        with ZipFile(zip_file_path) as zip_file:
+            zip_file.extractall(self.target_dir)
+
+        utils.remove_file(zip_file_path)
 
         if self.remove_existing_artifact:
             self.progress.update(

ads/model/deployment/model_deployment.py

@@ -64,7 +64,6 @@
 DELETE_WORKFLOW_STEPS = 2
 DEACTIVATE_WORKFLOW_STEPS = 2
 DEFAULT_RETRYING_REQUEST_ATTEMPTS = 3
-TERMINAL_STATES = [State.ACTIVE, State.FAILED, State.DELETED, State.INACTIVE]
 
 MODEL_DEPLOYMENT_KIND = "deployment"
 MODEL_DEPLOYMENT_TYPE = "modelDeployment"
@@ -720,7 +719,7 @@ def update(
 
     def watch(
         self,
-        log_type: str = ModelDeploymentLogType.ACCESS,
+        log_type: str = None,
         time_start: datetime = None,
         interval: int = LOG_INTERVAL,
         log_filter: str = None,
@@ -731,7 +730,7 @@ def watch(
         ----------
         log_type: str, optional
             The log type. Can be `access`, `predict` or None.
-            Defaults to access.
+            Defaults to None.
         time_start : datetime.datetime, optional
             Starting time for the log query.
             Defaults to None.
@@ -757,7 +756,7 @@ def watch(
             count = self.logs(log_type).stream(
                 source=self.model_deployment_id,
                 interval=interval,
-                stop_condition=self._stop_condition,
+                stop_condition=self._stream_stop_condition,
                 time_start=time_start,
                 log_filter=log_filter,
             )
@@ -773,7 +772,11 @@ def watch(
 
     def _stop_condition(self):
         """Stops the sync once the model deployment is in a terminal state."""
-        return self.state in TERMINAL_STATES
+        return self.state in [State.ACTIVE, State.FAILED, State.DELETED, State.INACTIVE]
+
+    def _stream_stop_condition(self):
+        """Stops the stream sync once the model deployment is in a terminal state."""
+        return self.state in [State.FAILED, State.DELETED, State.INACTIVE]
 
     def _check_and_print_status(self, prev_status) -> str:
         """Check and print the next status.

ads/model/deployment/model_deployment_infrastructure.py

@@ -20,7 +20,6 @@
 MODEL_DEPLOYMENT_INFRASTRUCTURE_KIND = "infrastructure"
 
 DEFAULT_BANDWIDTH_MBPS = 10
-DEFAULT_WEB_CONCURRENCY = 10
 DEFAULT_REPLICA = 1
 DEFAULT_SHAPE_NAME = "VM.Standard.E4.Flex"
 DEFAULT_OCPUS = 1
@@ -219,7 +218,6 @@ def _load_default_properties(self) -> Dict:
             defaults[self.CONST_PROJECT_ID] = PROJECT_OCID
 
         defaults[self.CONST_BANDWIDTH_MBPS] = DEFAULT_BANDWIDTH_MBPS
-        defaults[self.CONST_WEB_CONCURRENCY] = DEFAULT_WEB_CONCURRENCY
         defaults[self.CONST_REPLICA] = DEFAULT_REPLICA
 
         if NB_SESSION_OCID:
@@ -628,7 +626,6 @@ def init(self, **kwargs) -> "ModelDeploymentInfrastructure":
             .with_compartment_id(self.compartment_id or "{Provide a compartment OCID}")
             .with_project_id(self.project_id or "{Provide a project OCID}")
             .with_bandwidth_mbps(self.bandwidth_mbps or DEFAULT_BANDWIDTH_MBPS)
-            .with_web_concurrency(self.web_concurrency or DEFAULT_WEB_CONCURRENCY)
             .with_replica(self.replica or DEFAULT_REPLICA)
             .with_shape_name(self.shape_name or DEFAULT_SHAPE_NAME)
             .with_shape_config_details(

ads/model/generic_model.py

@@ -31,6 +31,7 @@
     NB_SESSION_OCID,
     PIPELINE_RUN_COMPARTMENT_OCID,
     PROJECT_OCID,
+    TMPDIR,
 )
 from ads.evaluations import EvaluatorMixin
 from ads.feature_engineering import ADSImage
@@ -190,7 +191,7 @@ def _prepare_artifact_dir(artifact_dir: str = None) -> str:
     if artifact_dir and isinstance(artifact_dir, str):
         return os.path.abspath(os.path.expanduser(artifact_dir))
 
-    artifact_dir = tempfile.mkdtemp()
+    artifact_dir = TMPDIR or tempfile.mkdtemp()
     logger.info(
         f"The `artifact_dir` was not provided and "
         f"automatically set to: {artifact_dir}"

ads/model/service/oci_datascience_model_deployment.py

@@ -67,11 +67,11 @@ def wrapper(self, *args, **kwargs):
     return decorator
 
 
-class MissingModelDeploymentIdError(Exception):   # pragma: no cover
+class MissingModelDeploymentIdError(Exception):  # pragma: no cover
     pass
 
 
-class MissingModelDeploymentWorkflowIdError(Exception):   # pragma: no cover
+class MissingModelDeploymentWorkflowIdError(Exception):  # pragma: no cover
     pass
 
 
@@ -175,32 +175,23 @@ def activate(
             The `OCIDataScienceModelDeployment` instance (self).
         """
         dsc_model_deployment = OCIDataScienceModelDeployment.from_id(self.id)
-        if (
-            dsc_model_deployment.lifecycle_state
-            == self.LIFECYCLE_STATE_ACTIVE
-        ):
+        if dsc_model_deployment.lifecycle_state == self.LIFECYCLE_STATE_ACTIVE:
             raise Exception(
                 f"Model deployment {dsc_model_deployment.id} is already in active state."
             )
 
-        if (
-            dsc_model_deployment.lifecycle_state
-            == self.LIFECYCLE_STATE_INACTIVE
-        ):
+        if dsc_model_deployment.lifecycle_state == self.LIFECYCLE_STATE_INACTIVE:
             logger.info(f"Activating model deployment `{self.id}`.")
             response = self.client.activate_model_deployment(
                 self.id,
             )
 
             if wait_for_completion:
-
                 self.workflow_req_id = response.headers.get("opc-work-request-id", None)
 
                 try:
                     self.wait_for_progress(
-                        self.workflow_req_id, 
-                        max_wait_time, 
-                        poll_interval
+                        self.workflow_req_id, max_wait_time, poll_interval
                     )
                 except Exception as e:
                     logger.error(
@@ -241,21 +232,17 @@ 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}")
 
         if wait_for_completion:
-
             self.workflow_req_id = response.headers.get("opc-work-request-id", None)
 
             try:
                 self.wait_for_progress(
-                    self.workflow_req_id, 
-                    max_wait_time, 
-                    poll_interval
+                    self.workflow_req_id, max_wait_time, poll_interval
                 )
             except Exception as e:
-                logger.error(
-                    "Error while trying to create model deployment: " + str(e)
-                )
+                logger.error("Error while trying to create model deployment: " + str(e))
 
         return self.sync()
 
@@ -287,32 +274,23 @@ def deactivate(
             The `OCIDataScienceModelDeployment` instance (self).
         """
         dsc_model_deployment = OCIDataScienceModelDeployment.from_id(self.id)
-        if (
-            dsc_model_deployment.lifecycle_state
-            == self.LIFECYCLE_STATE_INACTIVE
-        ):
+        if dsc_model_deployment.lifecycle_state == self.LIFECYCLE_STATE_INACTIVE:
             raise Exception(
                 f"Model deployment {dsc_model_deployment.id} is already in inactive state."
             )
 
-        if (
-            dsc_model_deployment.lifecycle_state
-            == self.LIFECYCLE_STATE_ACTIVE
-        ):
+        if dsc_model_deployment.lifecycle_state == self.LIFECYCLE_STATE_ACTIVE:
             logger.info(f"Deactivating model deployment `{self.id}`.")
             response = self.client.deactivate_model_deployment(
                 self.id,
             )
 
             if wait_for_completion:
-
                 self.workflow_req_id = response.headers.get("opc-work-request-id", None)
 
                 try:
                     self.wait_for_progress(
-                        self.workflow_req_id, 
-                        max_wait_time, 
-                        poll_interval
+                        self.workflow_req_id, max_wait_time, poll_interval
                     )
                 except Exception as e:
                     logger.error(
@@ -355,7 +333,7 @@ def delete(
         dsc_model_deployment = OCIDataScienceModelDeployment.from_id(self.id)
         if dsc_model_deployment.lifecycle_state in [
             self.LIFECYCLE_STATE_DELETED,
-            self.LIFECYCLE_STATE_DELETING
+            self.LIFECYCLE_STATE_DELETING,
         ]:
             raise Exception(
                 f"Model deployment {dsc_model_deployment.id} is either deleted or being deleted."
@@ -374,19 +352,14 @@ def delete(
         )
 
         if wait_for_completion:
-
             self.workflow_req_id = response.headers.get("opc-work-request-id", None)
 
             try:
                 self.wait_for_progress(
-                    self.workflow_req_id, 
-                    max_wait_time, 
-                    poll_interval
+                    self.workflow_req_id, max_wait_time, poll_interval
                 )
             except Exception as e:
-                logger.error(
-                    "Error while trying to delete model deployment: " + str(e)
-                )
+                logger.error("Error while trying to delete model deployment: " + str(e))
 
         return self.sync()
 
@@ -440,9 +413,7 @@ def update(
             )
             self.workflow_req_id = response.headers.get("opc-work-request-id", None)
         except Exception as e:
-            logger.error(
-                "Error while trying to update model deployment: " + str(e)
-            )
+            logger.error("Error while trying to update model deployment: " + str(e))
 
         return self.sync()
 

ads/opctl/operator/common/operator_loader.py

@@ -136,6 +136,9 @@ class OperatorInfo(DataClassSerializable):
     dataflow_default_params (DataFlowDefaultParams)
         The default params for the DataFlow service.
         Will be used when operator run on the DataFlow service.
+    logo: str
+        The logo of the operator.
+        Needs to be attached in the "svg+xml;base64" format.
 
     Properties
     ----------
@@ -157,6 +160,7 @@ class OperatorInfo(DataClassSerializable):
     dataflow_default_params: DataFlowDefaultParams = field(
         default_factory=DataFlowDefaultParams
     )
+    logo: str = ""
 
     @property
     def conda_prefix(self) -> str: