ODSC-34939: Implement optional loading of artifacts in the GenericModel (#488)

Author: VipulMascarenhas

ads/model/generic_model.py

@@ -153,6 +153,13 @@ def __init__(self, state: str):
         super().__init__(msg)
 
 
+class ArtifactsNotAvailableError(Exception):
+    def __init__(
+        self, msg="Model artifacts are either not generated or not available locally."
+    ):
+        super().__init__(msg)
+
+
 class SerializeModelNotImplementedError(NotImplementedError):  # pragma: no cover
     pass
 
@@ -281,6 +288,8 @@ class GenericModel(MetadataMixin, Introspectable, EvaluatorMixin):
         Tests if deployment works in local environment.
     upload_artifact(...)
         Uploads model artifacts to the provided `uri`.
+    download_artifact(...)
+        Downloads model artifacts from the model catalog.
 
 
     Examples
@@ -1249,6 +1258,9 @@ def verify(
         Dict
             A dictionary which contains prediction results.
         """
+        if self.model_artifact is None:
+            raise ArtifactsNotAvailableError
+
         endpoint = f"http://127.0.0.1:8000/predict"
         data = self._handle_input_data(data, auto_serialize_data, **kwargs)
 
@@ -1402,6 +1414,117 @@ def from_model_artifact(
 
         return model
 
+    def download_artifact(
+        self,
+        artifact_dir: Optional[str] = None,
+        auth: Optional[Dict] = None,
+        force_overwrite: Optional[bool] = False,
+        bucket_uri: Optional[str] = None,
+        remove_existing_artifact: Optional[bool] = True,
+        **kwargs,
+    ) -> "GenericModel":
+        """Downloads model artifacts from the model catalog.
+
+         Parameters
+         ----------
+        artifact_dir: (str, optional). Defaults to `None`.
+            The artifact directory to store the files needed for deployment.
+            Will be created if not exists.
+        auth: (Dict, optional). Defaults to None.
+            The default authentication is set using `ads.set_auth` API. If you need to override the
+            default, use the `ads.common.auth.api_keys` or `ads.common.auth.resource_principal` to create appropriate
+            authentication signer and kwargs required to instantiate IdentityClient object.
+        force_overwrite: (bool, optional). Defaults to False.
+            Whether to overwrite existing files or not.
+        bucket_uri: (str, optional). Defaults to None.
+            The OCI Object Storage URI where model artifacts will be copied to.
+            The `bucket_uri` is only necessary for downloading large artifacts with
+            size is greater than 2GB. Example: `oci://<bucket_name>@<namespace>/prefix/`.
+        remove_existing_artifact: (bool, optional). Defaults to `True`.
+            Whether artifacts uploaded to object storage bucket need to be removed or not.
+
+         Returns
+         -------
+         Self
+             An instance of `GenericModel` class.
+
+         Raises
+         ------
+         ValueError
+             If `model_id` is not available in the GenericModel object.
+        """
+        model_id = self.model_id
+        if not model_id:
+            raise ValueError(
+                "`model_id` is not available, load the GenericModel object first."
+            )
+
+        if not artifact_dir:
+            artifact_dir = self.artifact_dir
+        artifact_dir = _prepare_artifact_dir(artifact_dir)
+
+        target_dir = (
+            _prepare_artifact_dir()
+            if ObjectStorageDetails.is_oci_path(artifact_dir)
+            else artifact_dir
+        )
+
+        dsc_model = DataScienceModel.from_id(model_id)
+        dsc_model.download_artifact(
+            target_dir=target_dir,
+            force_overwrite=force_overwrite,
+            bucket_uri=bucket_uri,
+            remove_existing_artifact=remove_existing_artifact,
+            auth=auth,
+            region=kwargs.pop("region", None),
+            timeout=kwargs.pop("timeout", None),
+        )
+        model_artifact = ModelArtifact.from_uri(
+            uri=target_dir,
+            artifact_dir=artifact_dir,
+            model_file_name=self.model_file_name,
+            force_overwrite=force_overwrite,
+            auth=auth,
+            ignore_conda_error=self.ignore_conda_error,
+        )
+        self.dsc_model = dsc_model
+        self.local_copy_dir = model_artifact.local_copy_dir
+        self.model_artifact = model_artifact
+        self.reload_runtime_info()
+
+        self._summary_status.update_status(
+            detail="Generated score.py",
+            status=ModelState.DONE.value,
+        )
+        self._summary_status.update_status(
+            detail="Generated runtime.yaml",
+            status=ModelState.DONE.value,
+        )
+        self._summary_status.update_status(
+            detail="Serialized model", status=ModelState.DONE.value
+        )
+        self._summary_status.update_status(
+            detail="Populated metadata(Custom, Taxonomy and Provenance)",
+            status=ModelState.DONE.value,
+        )
+        self._summary_status.update_status(
+            detail="Local tested .predict from score.py",
+            status=ModelState.AVAILABLE.value,
+        )
+        self._summary_status.update_action(
+            detail="Local tested .predict from score.py",
+            action="",
+        )
+        self._summary_status.update_status(
+            detail="Conducted Introspect Test",
+            status=ModelState.AVAILABLE.value,
+        )
+        self._summary_status.update_status(
+            detail="Uploaded artifact to model catalog",
+            status=ModelState.AVAILABLE.value,
+        )
+        return self
+
     @classmethod
     def from_model_catalog(
         cls: Type[Self],
@@ -1414,6 +1537,7 @@ def from_model_catalog(
         bucket_uri: Optional[str] = None,
         remove_existing_artifact: Optional[bool] = True,
         ignore_conda_error: Optional[bool] = False,
+        download_artifact: Optional[bool] = True,
         **kwargs,
     ) -> Self:
         """Loads model from model catalog.
@@ -1443,6 +1567,8 @@ def from_model_catalog(
             Wether artifacts uploaded to object storage bucket need to be removed or not.
         ignore_conda_error: (bool, optional). Defaults to False.
             Parameter to ignore error when collecting conda information.
+        download_artifact: (bool, optional). Defaults to True.
+            Whether to download the model pickle or checkpoints
         kwargs:
             compartment_id : (str, optional)
                 Compartment OCID. If not specified, the value will be taken from the environment variables.
@@ -1475,14 +1601,60 @@ def from_model_catalog(
         artifact_dir = _prepare_artifact_dir(artifact_dir)
 
         target_dir = (
-            artifact_dir
-            if not ObjectStorageDetails.is_oci_path(artifact_dir)
-            else tempfile.mkdtemp()
+            _prepare_artifact_dir()
+            if ObjectStorageDetails.is_oci_path(artifact_dir)
+            else artifact_dir
         )
         bucket_uri = bucket_uri or (
             artifact_dir if ObjectStorageDetails.is_oci_path(artifact_dir) else None
         )
         dsc_model = DataScienceModel.from_id(model_id)
+
+        if not download_artifact:
+            result_model = cls(
+                artifact_dir=artifact_dir,
+                bucket_uri=bucket_uri,
+                auth=auth,
+                properties=properties,
+                ignore_conda_error=ignore_conda_error,
+                **kwargs,
+            )
+            result_model._summary_status.update_status(
+                detail="Generated score.py",
+                status=ModelState.NOTAPPLICABLE.value,
+            )
+            result_model._summary_status.update_status(
+                detail="Generated runtime.yaml",
+                status=ModelState.NOTAPPLICABLE.value,
+            )
+            result_model._summary_status.update_status(
+                detail="Serialized model", status=ModelState.NOTAPPLICABLE.value
+            )
+            result_model._summary_status.update_status(
+                detail="Populated metadata(Custom, Taxonomy and Provenance)",
+                status=ModelState.NOTAPPLICABLE.value,
+            )
+            result_model._summary_status.update_status(
+                detail="Local tested .predict from score.py",
+                status=ModelState.NOTAPPLICABLE.value,
+            )
+            result_model._summary_status.update_action(
+                detail="Local tested .predict from score.py",
+                action="Local artifact is not available. "
+                "Set load_artifact flag to True while loading the model or "
+                "call .download_artifact().",
+            )
+            result_model._summary_status.update_status(
+                detail="Conducted Introspect Test",
+                status=ModelState.NOTAPPLICABLE.value,
+            )
+            result_model._summary_status.update_status(
+                detail="Uploaded artifact to model catalog",
+                status=ModelState.NOTAPPLICABLE.value,
+            )
+            result_model.dsc_model = dsc_model
+            return result_model
+
         dsc_model.download_artifact(
             target_dir=target_dir,
             force_overwrite=force_overwrite,
@@ -1536,6 +1708,7 @@ def from_model_deployment(
         bucket_uri: Optional[str] = None,
         remove_existing_artifact: Optional[bool] = True,
         ignore_conda_error: Optional[bool] = False,
+        download_artifact: Optional[bool] = True,
         **kwargs,
     ) -> Self:
         """Loads model from model deployment.
@@ -1565,6 +1738,8 @@ def from_model_deployment(
             Wether artifacts uploaded to object storage bucket need to be removed or not.
         ignore_conda_error: (bool, optional). Defaults to False.
             Parameter to ignore error when collecting conda information.
+        download_artifact: (bool, optional). Defaults to True.
+            Whether to download the model pickle or checkpoints
         kwargs:
             compartment_id : (str, optional)
                 Compartment OCID. If not specified, the value will be taken from the environment variables.
@@ -1608,6 +1783,7 @@ def from_model_deployment(
             bucket_uri=bucket_uri,
             remove_existing_artifact=remove_existing_artifact,
             ignore_conda_error=ignore_conda_error,
+            download_artifact=download_artifact,
             **kwargs,
         )
         model._summary_status.update_status(
@@ -1730,6 +1906,7 @@ def from_id(
         bucket_uri: Optional[str] = None,
         remove_existing_artifact: Optional[bool] = True,
         ignore_conda_error: Optional[bool] = False,
+        download_artifact: Optional[bool] = True,
         **kwargs,
     ) -> Self:
         """Loads model from model OCID or model deployment OCID.
@@ -1757,6 +1934,10 @@ def from_id(
             size is greater than 2GB. Example: `oci://<bucket_name>@<namespace>/prefix/`.
         remove_existing_artifact: (bool, optional). Defaults to `True`.
             Wether artifacts uploaded to object storage bucket need to be removed or not.
+        ignore_conda_error: (bool, optional). Defaults to False.
+            Parameter to ignore error when collecting conda information.
+        download_artifact: (bool, optional). Defaults to True.
+            Whether to download the model pickle or checkpoints
         kwargs:
             compartment_id : (str, optional)
                 Compartment OCID. If not specified, the value will be taken from the environment variables.
@@ -1780,6 +1961,7 @@ def from_id(
                 bucket_uri=bucket_uri,
                 remove_existing_artifact=remove_existing_artifact,
                 ignore_conda_error=ignore_conda_error,
+                download_artifact=download_artifact,
                 **kwargs,
             )
         elif DataScienceModelType.MODEL in ocid:
@@ -1793,6 +1975,7 @@ def from_id(
                 bucket_uri=bucket_uri,
                 remove_existing_artifact=remove_existing_artifact,
                 ignore_conda_error=ignore_conda_error,
+                download_artifact=download_artifact,
                 **kwargs,
             )
         else:
@@ -1924,11 +2107,13 @@ def save(
         ...     bucket_uri="oci://my-bucket@my-tenancy/",
         ...     overwrite_existing_artifact=True,
         ...     remove_existing_artifact=True,
-        ...     remove_existing_artifact=True,
         ...     parallel_process_count=9,
         ... )
 
         """
+        if self.model_artifact is None:
+            raise ArtifactsNotAvailableError
+
         # Set default display_name if not specified - randomly generated easy to remember name generated
         if not display_name:
             display_name = self._random_display_name()