oracle-ads 2.13.9rc0__py3-none-any.whl → 2.13.10rc0__py3-none-any.whl

ads/model/artifact_downloader.py

@@ -0,0 +1,254 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2022, 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+import os
+import shutil
+import uuid
+from abc import ABC, abstractmethod
+from typing import Dict, Optional
+from zipfile import ZipFile
+
+from ads.common import utils
+from ads.common.object_storage_details import ObjectStorageDetails
+from ads.common.utils import extract_region
+from ads.model.service.oci_datascience_model import OCIDataScienceModel
+
+
+class ArtifactDownloader(ABC):
+    """The abstract class to download model artifacts."""
+
+    PROGRESS_STEPS_COUNT = 1
+
+    def __init__(
+        self,
+        dsc_model: OCIDataScienceModel,
+        target_dir: str,
+        force_overwrite: Optional[bool] = False,
+    ):
+        """Initializes `ArtifactDownloader` instance.
+
+        Parameters
+        ----------
+        dsc_model: OCIDataScienceModel
+            The data scince model instance.
+        target_dir: str
+            The target location of model after download.
+        force_overwrite: bool
+            Overwrite target_dir if exists.
+        """
+        self.dsc_model = dsc_model
+        self.target_dir = target_dir
+        self.force_overwrite = force_overwrite
+        self.progress = None
+
+    def download(self):
+        """Downloads model artifacts.
+
+        Returns
+        -------
+        None
+
+        Raises
+        ------
+        ValueError
+            If target directory does not exist.
+        """
+        if os.path.exists(self.target_dir) and len(os.listdir(self.target_dir)) > 0:
+            if not self.force_overwrite:
+                raise ValueError(
+                    f"The `{self.target_dir}` directory already exists. "
+                    "Set `force_overwrite` to `True` if you wish to overwrite."
+                )
+            shutil.rmtree(self.target_dir)
+        os.makedirs(self.target_dir, exist_ok=True)
+        with utils.get_progress_bar(
+            ArtifactDownloader.PROGRESS_STEPS_COUNT + self.PROGRESS_STEPS_COUNT
+        ) as progress:
+            self.progress = progress
+            self._download()
+            self.progress.update(
+                "Downloading model artifacts has been successfully completed."
+            )
+            self.progress.update("Done.")
+
+    @abstractmethod
+    def _download(self):
+        """Downloads model artifacts."""
+
+
+class SmallArtifactDownloader(ArtifactDownloader):
+    PROGRESS_STEPS_COUNT = 3
+
+    def _download(self):
+        """Downloads model artifacts."""
+        self.progress.update("Importing model artifacts from catalog")
+
+        artifact_info = self.dsc_model.get_artifact_info()
+        artifact_name = artifact_info["Content-Disposition"].replace(
+            "attachment; filename=", ""
+        )
+        _, file_extension = os.path.splitext(artifact_name)
+        file_extension = file_extension.lower() if file_extension else ".zip"
+
+        file_content = self.dsc_model.get_model_artifact_content()
+        self.progress.update("Copying model artifacts to the artifact directory")
+
+        file_name = (
+            "model_description" if file_extension == ".json" else str(uuid.uuid4())
+        )
+        artifact_file_path = os.path.join(
+            self.target_dir, f"{file_name}{file_extension}"
+        )
+        with open(artifact_file_path, "wb") as _file:
+            _file.write(file_content)
+
+        if file_extension == ".zip":
+            self.progress.update("Extracting model artifacts")
+            with ZipFile(artifact_file_path) as _file:
+                _file.extractall(self.target_dir)
+            utils.remove_file(artifact_file_path)
+
+
+class LargeArtifactDownloader(ArtifactDownloader):
+    PROGRESS_STEPS_COUNT = 4
+
+    def __init__(
+        self,
+        dsc_model: OCIDataScienceModel,
+        target_dir: str,
+        auth: Optional[Dict] = None,
+        force_overwrite: Optional[bool] = False,
+        region: Optional[str] = None,
+        bucket_uri: Optional[str] = None,
+        overwrite_existing_artifact: Optional[bool] = True,
+        remove_existing_artifact: Optional[bool] = True,
+        model_file_description: Optional[dict] = None,
+    ):
+        """Initializes `LargeArtifactDownloader` instance.
+
+        Parameters
+        ----------
+        dsc_model: OCIDataScienceModel
+            The data scince model instance.
+        target_dir: str
+            The target location of model after download.
+        auth: (Dict, optional). Defaults to `None`.
+            The default authetication 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`.
+            Overwrite target_dir if exists.
+        region: (str, optional). Defaults to `None`.
+            The destination Object Storage bucket region.
+            By default the value will be extracted from the `OCI_REGION_METADATA` environment variables.
+        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 uploading large artifacts which
+            size is greater than 2GB. Example: `oci://<bucket_name>@<namespace>/prefix/`.
+        overwrite_existing_artifact: (bool, optional). Defaults to `True`.
+            Overwrite target bucket artifact if exists.
+        remove_existing_artifact: (bool, optional). Defaults to `True`.
+            Wether artifacts uploaded to object storage bucket need to be removed or not.
+        model_file_description: (dict, optional). Defaults to None.
+            Contains object path details for models created by reference.
+
+        """
+        super().__init__(
+            dsc_model=dsc_model, target_dir=target_dir, force_overwrite=force_overwrite
+        )
+        self.auth = auth or dsc_model.auth
+        self.region = region or extract_region(self.auth)
+        self.bucket_uri = bucket_uri
+        self.overwrite_existing_artifact = overwrite_existing_artifact
+        self.remove_existing_artifact = remove_existing_artifact
+        self.model_file_description = model_file_description
+
+    def _download(self):
+        """Downloads model artifacts."""
+        self.progress.update("Importing model artifacts from catalog")
+
+        if (
+            self.dsc_model.is_model_created_by_reference()
+            and self.model_file_description
+        ):
+            self.download_from_model_file_description()
+            self.progress.update()
+            return
+
+        bucket_uri = self.bucket_uri
+
+        if not os.path.basename(bucket_uri):
+            bucket_uri = os.path.join(bucket_uri, f"{self.dsc_model.id}.zip")
+        elif not bucket_uri.lower().endswith(".zip"):
+            bucket_uri = f"{bucket_uri}.zip"
+
+        self.dsc_model.import_model_artifact(bucket_uri=bucket_uri, region=self.region)
+        self.progress.update("Copying model artifacts to the artifact directory")
+        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(
+                "Removing temporary artifacts from the Object Storage bucket"
+            )
+            utils.remove_file(bucket_uri)
+        else:
+            self.progress.update()
+
+    def download_from_model_file_description(self):
+        """Helper function to download the objects using model file description content to the target directory."""
+
+        models = self.model_file_description["models"]
+        total_size = 0
+
+        for model in models:
+            namespace, bucket_name, prefix = (
+                model["namespace"],
+                model["bucketName"],
+                model["prefix"],
+            )
+            bucket_uri = f"oci://{bucket_name}@{namespace}/{prefix}"
+
+            message = f"Copying model artifacts by reference from {bucket_uri} to {self.target_dir}"
+            self.progress.update(message)
+
+            objects = model["objects"]
+            os_details_list = list()
+
+            for obj in objects:
+                name = obj["name"]
+                version = None if obj["version"] == "" else obj["version"]
+                size = obj["sizeInBytes"]
+                if size == 0:
+                    continue
+                total_size += size
+                object_uri = f"oci://{bucket_name}@{namespace}/{name}"
+
+                os_details = ObjectStorageDetails.from_path(object_uri)
+                os_details.version = version
+                os_details_list.append(os_details)
+            try:
+                ObjectStorageDetails.from_path(
+                    bucket_uri
+                ).bulk_download_from_object_storage(
+                    paths=os_details_list,
+                    target_dir=self.target_dir,
+                    progress_bar=self.progress,
+                )
+            except Exception as ex:
+                raise RuntimeError(
+                    f"Failed to download model artifact by reference from the given Object Storage path `{bucket_uri}`."
+                    f"See Exception: {ex}"
+                )

ads/model/artifact_uploader.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python
+# -*- coding: utf-8; -*-
+import logging
+
+# Copyright (c) 2022, 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+import os
+import shutil
+from abc import ABC, abstractmethod
+from typing import Dict, Optional
+
+from ads.common import utils
+from ads.common.object_storage_details import ObjectStorageDetails
+from ads.model.common import utils as model_utils
+from ads.model.service.oci_datascience_model import OCIDataScienceModel
+
+
+class ArtifactUploader(ABC):
+    """The abstract class to upload model artifacts."""
+
+    PROGRESS_STEPS_COUNT = 3
+
+    def __init__(self, dsc_model: OCIDataScienceModel, artifact_path: str):
+        """Initializes `ArtifactUploader` instance.
+
+        Parameters
+        ----------
+        dsc_model: OCIDataScienceModel
+            The data scince model instance.
+        artifact_path: str
+            The model artifact location.
+        """
+        if not (
+            ObjectStorageDetails.is_oci_path(artifact_path)
+            or os.path.exists(artifact_path)
+        ):
+            raise ValueError(f"The `{artifact_path}` does not exist")
+
+        self.dsc_model = dsc_model
+        self.artifact_path = artifact_path
+        self.artifact_file_path = None
+        self.progress = None
+
+    def upload(self):
+        """Uploads model artifacts."""
+        try:
+            with utils.get_progress_bar(
+                ArtifactUploader.PROGRESS_STEPS_COUNT + self.PROGRESS_STEPS_COUNT
+            ) as progress:
+                self.progress = progress
+                self.progress.update("Preparing model artifacts file.")
+                self._prepare_artifact_tmp_file()
+                self.progress.update("Uploading model artifacts.")
+                self._upload()
+                self.progress.update(
+                    "Uploading model artifacts has been successfully completed."
+                )
+                self.progress.update("Done.")
+        except Exception:
+            raise
+        finally:
+            self._remove_artifact_tmp_file()
+
+    def _prepare_artifact_tmp_file(self) -> str:
+        """Prepares model artifacts file.
+
+        Returns
+        -------
+        str
+            Path to the model artifact file.
+        """
+        if ObjectStorageDetails.is_oci_path(self.artifact_path):
+            self.artifact_file_path = self.artifact_path
+        elif os.path.isfile(self.artifact_path) and self.artifact_path.lower().endswith(
+            (".zip", ".json")
+        ):
+            self.artifact_file_path = self.artifact_path
+        else:
+            self.artifact_file_path = model_utils.zip_artifact(
+                artifact_dir=self.artifact_path
+            )
+        return self.artifact_file_path
+
+    def _remove_artifact_tmp_file(self):
+        """Removes temporary created artifact file."""
+        if (
+            self.artifact_file_path
+            and self.artifact_file_path.lower() != self.artifact_path.lower()
+        ):
+            shutil.rmtree(self.artifact_file_path, ignore_errors=True)
+
+    @abstractmethod
+    def _upload(self):
+        """Uploads model artifacts. Needs to be implemented in a child class."""
+
+
+class SmallArtifactUploader(ArtifactUploader):
+    """The class helper to upload small model artifacts."""
+
+    PROGRESS_STEPS_COUNT = 1
+
+    def _upload(self):
+        """Uploads model artifacts to the model catalog."""
+        _, ext = os.path.splitext(self.artifact_file_path)
+        self.progress.update("Uploading model artifacts to the catalog")
+        with open(self.artifact_file_path, "rb") as file_data:
+            self.dsc_model.create_model_artifact(bytes_content=file_data, extension=ext)
+
+
+class LargeArtifactUploader(ArtifactUploader):
+    """
+    The class helper to upload large model artifacts.
+
+    Attributes
+    ----------
+    artifact_path: str
+        The model artifact location. Possible values are:
+            - object storage path to zip archive. Example: `oci://<bucket_name>@<namespace>/prefix/mymodel.zip`.
+            - local path to zip archive. Example: `./mymodel.zip`.
+            - local path to folder with artifacts. Example: `./mymodel`.
+    artifact_file_path: str
+        The uri of the zip of model artifact.
+    auth: dict
+        The default authetication 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.
+    bucket_uri: str
+        The OCI Object Storage URI where model artifacts will be copied to.
+        The `bucket_uri` is only necessary for uploading large artifacts which
+        size is greater than 2GB. Example: `oci://<bucket_name>@<namespace>/prefix/`.
+
+        .. versionadded:: 2.8.10
+
+            If artifact_path is object storage path to a zip archive, bucket_uri will be ignored.
+
+    dsc_model: OCIDataScienceModel
+        The data scince model instance.
+    overwrite_existing_artifact: bool
+        Overwrite target bucket artifact if exists.
+    progress: TqdmProgressBar
+        An instance of the TqdmProgressBar.
+    region: str
+        The destination Object Storage bucket region.
+        By default the value will be extracted from the `OCI_REGION_METADATA` environment variables.
+    remove_existing_artifact: bool
+        Wether artifacts uploaded to object storage bucket need to be removed or not.
+    upload_manager: UploadManager
+        The uploadManager simplifies interaction with the Object Storage service.
+    """
+
+    PROGRESS_STEPS_COUNT = 4
+
+    def __init__(
+        self,
+        dsc_model: OCIDataScienceModel,
+        artifact_path: str,
+        bucket_uri: str = None,
+        auth: Optional[Dict] = None,
+        region: Optional[str] = None,
+        overwrite_existing_artifact: Optional[bool] = True,
+        remove_existing_artifact: Optional[bool] = True,
+        parallel_process_count: int = utils.DEFAULT_PARALLEL_PROCESS_COUNT,
+    ):
+        """Initializes `LargeArtifactUploader` instance.
+
+        Parameters
+        ----------
+        dsc_model: OCIDataScienceModel
+            The data scince model instance.
+        artifact_path: str
+            The model artifact location. Possible values are:
+                - object storage path to zip archive. Example: `oci://<bucket_name>@<namespace>/prefix/mymodel.zip`.
+                - local path to zip archive. Example: `./mymodel.zip`.
+                - local path to folder with artifacts. Example: `./mymodel`.
+        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 uploading large artifacts from local which
+            size is greater than 2GB. Example: `oci://<bucket_name>@<namespace>/prefix/`.
+
+            .. versionadded:: 2.8.10
+
+                If `artifact_path` is object storage path to a zip archive, `bucket_uri` will be ignored.
+
+        auth: (Dict, optional). Defaults to `None`.
+            The default authetication 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.
+        region: (str, optional). Defaults to `None`.
+            The destination Object Storage bucket region.
+            By default the value will be extracted from the `OCI_REGION_METADATA` environment variables.
+        overwrite_existing_artifact: (bool, optional). Defaults to `True`.
+            Overwrite target bucket artifact if exists.
+        remove_existing_artifact: (bool, optional). Defaults to `True`.
+            Whether artifacts uploaded to object storage bucket need to be removed or not.
+        parallel_process_count: (int, optional).
+            The number of worker processes to use in parallel for uploading individual parts of a multipart upload.
+        """
+        self.auth = auth or dsc_model.auth
+        if ObjectStorageDetails.is_oci_path(artifact_path):
+            if not artifact_path.endswith(".zip"):
+                raise ValueError(
+                    f"The `artifact_path={artifact_path}` is invalid."
+                    "The remote path for model artifact should be a zip archive, "
+                    "e.g. `oci://<bucket_name>@<namespace>/prefix/mymodel.zip`."
+                )
+            if not utils.is_path_exists(uri=artifact_path, auth=self.auth):
+                raise ValueError(f"The `{artifact_path}` does not exist.")
+            bucket_uri = artifact_path
+
+        if not bucket_uri:
+            raise ValueError("The `bucket_uri` must be provided.")
+
+        super().__init__(dsc_model=dsc_model, artifact_path=artifact_path)
+        self.region = region or utils.extract_region(self.auth)
+        self.bucket_uri = bucket_uri
+        self.overwrite_existing_artifact = overwrite_existing_artifact
+        self.remove_existing_artifact = remove_existing_artifact
+        self._parallel_process_count = parallel_process_count
+
+    def _upload(self):
+        """Uploads model artifacts to the model catalog."""
+        bucket_uri = self.bucket_uri
+        self.progress.update("Copying model artifact to the Object Storage bucket")
+        if not bucket_uri == self.artifact_file_path:
+            bucket_uri_file_name = os.path.basename(bucket_uri)
+
+            if not bucket_uri_file_name:
+                bucket_uri = os.path.join(bucket_uri, f"{self.dsc_model.id}.zip")
+            elif not bucket_uri.lower().endswith(".zip"):
+                bucket_uri = f"{bucket_uri}.zip"
+
+            if not self.overwrite_existing_artifact and utils.is_path_exists(
+                uri=bucket_uri, auth=self.auth
+            ):
+                raise FileExistsError(
+                    f"The bucket_uri=`{self.bucket_uri}` exists. Please use a new file name or "
+                    "set `overwrite_existing_artifact` to `True` if you wish to overwrite."
+                )
+
+            try:
+                utils.upload_to_os(
+                    src_uri=self.artifact_file_path,
+                    dst_uri=bucket_uri,
+                    auth=self.auth,
+                    parallel_process_count=self._parallel_process_count,
+                    force_overwrite=self.overwrite_existing_artifact,
+                    progressbar_description="Copying model artifact to the Object Storage bucket.",
+                )
+            except Exception as ex:
+                raise RuntimeError(
+                    f"Failed to upload model artifact to the given Object Storage path `{self.bucket_uri}`."
+                    f"See Exception: {ex}"
+                )
+
+        self.progress.update("Exporting model artifact to the model catalog")
+        self.dsc_model.export_model_artifact(bucket_uri=bucket_uri, region=self.region)
+
+        if self.remove_existing_artifact:
+            self.progress.update(
+                "Removing temporary artifacts from the Object Storage bucket"
+            )
+            utils.remove_file(bucket_uri)
+        else:
+            self.progress.update()