Refactor OS plugin to incorporate OCI OS UploadManager for enhanced functionality.

Author: mrDzurb

container-image/environment.yaml

@@ -3,7 +3,7 @@ dependencies:
   - main::pip
   - pip:
     - oracledb
-    - mlflow==2.3.2
+    - mlflow
     - oracle-ads>=2.8.5
     - mysql-connector-python
     - oci-mlflow

docs/source/project.rst

@@ -50,7 +50,7 @@ The ``{work_dir}`` can be used to point out that the YAML template located insid
 Data Science Job Template
 =========================
 
-The template file contains the information about the infrastructure on which a Data Science job should be run, and also the runtime information. More details can be found in the `ADS documentation <https://accelerated-data-science.readthedocs.io/en/latest/user_guide/jobs/data_science_job.html>`__. The template file is divided into two main sections: ``infrastructure`` and ``runtime``.
+The template file contains the information about the infrastructure on which a Data Science job should be run, and also the runtime information. More details can be found in the `ADS documentation <https://accelerated-data-science.readthedocs.io/en/latest/user_guide/jobs/data_science_job.html>`__. The template file is divided into two main sections: ``infrastructure`` and ``runtime``. The template also can be generated using ``ads opctl init`` command. More details can be found in the `ADS documentation <https://accelerated-data-science.readthedocs.io/en/latest/user_guide/cli/opctl/configure.html#generate-starter-yaml>`__.
 
 Data Science Job Infrastructure
 ###############################
@@ -507,7 +507,7 @@ This example demonstrates an MLflow project that trains a logistic regression mo
 
   Copy the ``oci-datascience-config.json`` file to the ``pyspark_ml_autologging`` folder.
 
-- Prepare a ``oci-datascience-template.yaml`` job configuration file.
+- Prepare a ``oci-datascience-template.yaml`` job configuration file. The template can be generated using ``ads opctl init`` command. More details can be found in the `ADS documentation <https://accelerated-data-science.readthedocs.io/en/latest/user_guide/cli/opctl/configure.html#generate-starter-yaml>`__.
 
   .. tabs::
 

oci_mlflow/oci_object_storage.py

@@ -5,40 +5,101 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import os
+from typing import List
 from urllib.parse import urlparse
 
 import fsspec
+from ads.common.auth import AuthType, default_signer, set_auth
+from ads.common.oci_client import OCIClientFactory
 from mlflow.entities import FileInfo
 from mlflow.store.artifact.artifact_repo import ArtifactRepository
 from mlflow.utils.file_utils import relative_path_to_artifact_path
+from oci import object_storage
 from ocifs import OCIFileSystem
 
 from oci_mlflow import logger
 
-OCI_PREFIX = "oci://"
+OCI_SCHEME = "oci"
+OCI_PREFIX = f"{OCI_SCHEME}://"
 
 
-class OCIObjectStorageArtifactRepository(ArtifactRepository):
+def parse_os_uri(uri: str):
+    """
+    Parse an OCI object storage URI, returning tuple (bucket, namespace, path).
+
+    Parameters
+    ----------
+    uri: str
+        The OCI Object Storage URI.
+
+    Returns
+    -------
+    Tuple
+        The (bucket, ns, type)
+
+    Raise
+    -----
+    Exception
+        If provided URI is not an OCI OS bucket URI.
     """
-    MLFlow Plugin implementation for storing artifacts to OCI Object Storage
+    parsed = urlparse(uri)
+    if parsed.scheme.lower() != OCI_SCHEME:
+        raise Exception("Not an OCI object storage URI: %s" % uri)
+    path = parsed.path
+
+    if path.startswith("/"):
+        path = path[1:]
+
+    bucket, ns = parsed.netloc.split("@")
+
+    return bucket, ns, path
+
+
+class ArtifactUploader:
     """
+    The class helper to upload model artifacts.
+
+    Attributes
+    ----------
+    upload_manager: UploadManager
+        The uploadManager simplifies interaction with the Object Storage service.
+    """
+
+    def __init__(self):
+        """Initializes `ArtifactUploader` instance."""
+        auth_type = os.environ.get(
+            "OCI_IAM_TYPE", os.environ.get("OCIFS_IAM_TYPE", AuthType.API_KEY)
+        )
+        logger.debug(f"Using auth {auth_type=}")
+        set_auth(auth_type)
+
+        self.upload_manager = object_storage.UploadManager(
+            OCIClientFactory(**default_signer()).object_storage
+        )
+
+    def upload(self, file_path: str, dst_path: str):
+        """Uploads model artifacts.
+
+        Parameters
+        ----------
+        file_path: str
+            The source file path.
+        dst_path: str
+            The destination path.
+        """
+        bucket_name, namespace_name, object_name = parse_os_uri(dst_path)
+        logger.debug(f"{bucket_name=}, {namespace_name=}, {object_name=}")
+        response = self.upload_manager.upload_file(
+            namespace_name=namespace_name,
+            bucket_name=bucket_name,
+            object_name=object_name,
+            file_path=file_path,
+        )
+        logger.debug(response)
 
-    @staticmethod
-    def parse_os_uri(uri):
-        """Parse an OCI object storage URI, returning (bucket, namespace, path)"""
-        parsed = urlparse(uri)
-        if parsed.scheme != "oci":
-            raise Exception("Not an OCI object storage URI: %s" % uri)
-        path = parsed.path
-        if path.startswith("/"):
-            path = path[1:]
-        bucket, ns = parsed.netloc.split("@")
-        return bucket, ns, path
-
-    def _upload_file(self, local_file, dest_path):
-        with open(local_file, "rb") as data:
-            with fsspec.open(dest_path, "wb") as outfile:
-                outfile.write(data.read())
+
+class OCIObjectStorageArtifactRepository(ArtifactRepository):
+    """MLFlow Plugin implementation for storing artifacts to OCI Object Storage."""
 
     def _download_file(self, remote_file_path, local_path):
         if not remote_file_path.startswith(self.artifact_uri):
@@ -49,43 +110,84 @@ def _download_file(self, remote_file_path, local_path):
         logger.info(f"{full_path}, {remote_file_path}")
         fs.download(full_path, local_path)
 
-    def log_artifact(self, local_file, artifact_path=None):
+    def log_artifact(self, local_file: str, artifact_path: str = None):
+        """
+        Logs a local file as an artifact, optionally taking an ``artifact_path`` to place it in
+        within the run's artifacts. Run artifacts can be organized into directories, so you can
+        place the artifact in a directory this way.
+
+        Parameters
+        ----------
+        local_file:str
+            Path to artifact to log.
+        artifact_path:str
+            Directory within the run's artifact directory in which to log the artifact.
+        """
         if artifact_path:
             dest_path = os.path.join(self.artifact_uri, artifact_path)
         else:
             dest_path = self.artifact_uri
         dest_path = os.path.join(dest_path, os.path.basename(local_file))
-        self._upload_file(local_file, dest_path)
 
-    def log_artifacts(self, local_dir, artifact_path=None):
+        ArtifactUploader().upload(local_file, dest_path)
+
+    def log_artifacts(self, local_dir: str, artifact_path: str = None):
+        """
+        Logs the files in the specified local directory as artifacts, optionally taking
+        an ``artifact_path`` to place them in within the run's artifacts.
+
+        Parameters
+        ----------
+        local_dir:str
+            Directory of local artifacts to log.
+        artifact_path:str
+            Directory within the run's artifact directory in which to log the artifacts.
+        """
+        artifact_uploader = ArtifactUploader()
+
         if artifact_path:
             dest_path = os.path.join(self.artifact_uri, artifact_path)
         else:
             dest_path = artifact_path
         local_dir = os.path.abspath(local_dir)
-        for (root, _, filenames) in os.walk(local_dir):
+
+        for root, _, filenames in os.walk(local_dir):
             upload_path = dest_path
             if root != local_dir:
                 rel_path = os.path.relpath(root, local_dir)
                 rel_path = relative_path_to_artifact_path(rel_path)
                 upload_path = os.path.join(dest_path, rel_path)
             for f in filenames:
-                self._upload_file(
-                    local_file=os.path.join(root, f),
-                    dest_path=os.path.join(upload_path, f),
+                artifact_uploader.upload(
+                    file_path=os.path.join(root, f),
+                    dst_path=os.path.join(upload_path, f),
                 )
 
     def get_fs(self):
         """
-        Get fssepc filesystem based on the uri scheme
+        Gets fssepc filesystem based on the uri scheme.
         """
         self.fs = fsspec.filesystem(
             urlparse(self.artifact_uri).scheme
         )  # FileSystem class corresponding to the URI scheme.
 
         return self.fs
 
-    def list_artifacts(self, path: str = ""):
+    def list_artifacts(self, path: str = "") -> List[FileInfo]:
+        """
+        Return all the artifacts for this run_id directly under path. If path is a file, returns
+        an empty list. Will error if path is neither a file nor directory.
+
+        Parameters
+        ----------
+        path:str
+            Relative source path that contains desired artifacts
+
+        Returns
+        -------
+        List[FileInfo]
+            List of artifacts as FileInfo listed directly under path.
+        """
         result = []
         dest_path = self.artifact_uri
         if path:
@@ -111,7 +213,17 @@ def list_artifacts(self, path: str = ""):
         result.sort(key=lambda f: f.path)
         return result
 
-    def delete_artifacts(self, artifact_path=None):
+    def delete_artifacts(self, artifact_path: str = None):
+        """
+        Delete the artifacts at the specified location.
+        Supports the deletion of a single file or of a directory. Deletion of a directory
+        is recursive.
+
+        Parameters
+        ----------
+        artifact_path: str
+            Path of the artifact to delete.
+        """
         dest_path = self.artifact_uri
         if artifact_path:
             dest_path = os.path.join(self.artifact_uri, artifact_path)

tests/plugins/unitary/test_oci_object_storage.py

@@ -1,17 +1,21 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 
+import os
+import tempfile
+from unittest.mock import MagicMock, Mock, patch
+
 # Copyright (c) 2023 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 import pytest
-from unittest.mock import MagicMock, Mock, patch
-import tempfile
-import os
-
 from mlflow.entities import FileInfo
 
 from oci_mlflow import oci_object_storage
-from oci_mlflow.oci_object_storage import OCIObjectStorageArtifactRepository
+from oci_mlflow.oci_object_storage import (
+    ArtifactUploader,
+    OCIObjectStorageArtifactRepository,
+)
+from oci import object_storage
 
 
 class DataObject:
@@ -41,7 +45,7 @@ def mock_fsspec_open(self):
             yield mock_open
 
     def test_parse_os_uri(self, oci_artifact_repo):
-        bucket, namespace, path = oci_artifact_repo.parse_os_uri(
+        bucket, namespace, path = oci_object_storage.parse_os_uri(
             "oci://my-bucket@my-namespace/my-artifact-path"
         )
         assert bucket == "my-bucket"
@@ -50,20 +54,7 @@ def test_parse_os_uri(self, oci_artifact_repo):
 
     def test_parse_os_uri_with_invalid_scheme(self, oci_artifact_repo):
         with pytest.raises(Exception):
-            oci_artifact_repo.parse_os_uri("s3://my-bucket/my-artifact-path")
-
-    def test_upload_file(self, mock_fsspec_open):
-        local_file = os.path.join(self.curr_dir, "test_files/test.txt")
-        dest_path = "oci://my-bucket@my-namespace/path/to/test.txt"
-        repository = OCIObjectStorageArtifactRepository(artifact_uri=dest_path)
-        mock_outfile = Mock()
-        mock_fsspec_open.return_value.__enter__.return_value = mock_outfile
-
-        repository._upload_file(local_file, dest_path)
-
-        mock_fsspec_open.assert_called_once_with(dest_path, "wb")
-        with open(local_file, "rb") as f:
-            mock_outfile.write.assert_called_once_with(f.read())
+            oci_object_storage.parse_os_uri("s3://my-bucket/my-artifact-path")
 
     def test_download_file(self, oci_artifact_repo):
         mock_fs = MagicMock()
@@ -82,7 +73,7 @@ def test_download_file(self, oci_artifact_repo):
                 local_path,
             )
 
-    @patch.object(OCIObjectStorageArtifactRepository, "_upload_file")
+    @patch.object(ArtifactUploader, "upload")
     def test_log_artifact(self, mock_upload_file, oci_artifact_repo):
         local_file = "test_files/test.txt"
         artifact_path = "logs"
@@ -92,7 +83,7 @@ def test_log_artifact(self, mock_upload_file, oci_artifact_repo):
         )
         mock_upload_file.assert_called_once_with(local_file, expected_dest_path)
 
-    @patch.object(OCIObjectStorageArtifactRepository, "_upload_file")
+    @patch.object(ArtifactUploader, "upload")
     def test_log_artifacts(self, mock_upload_file, oci_artifact_repo):
         local_dir = os.path.join(self.curr_dir, "test_files")
         dest_path = "path/to/dest"
@@ -128,3 +119,29 @@ def test_list_artifacts(self):
             FileInfo("sub_folder", True, 0),
         ]
         assert artifacts == expected_artifacts
+
+
+class TestArtifactUploader:
+    def test_init(self):
+        """Ensures the ArtifactUploader instance can be initialized."""
+        artifact_uploader = ArtifactUploader()
+        assert isinstance(
+            artifact_uploader.upload_manager, object_storage.UploadManager
+        )
+
+    @patch.object(object_storage.UploadManager, "upload_file")
+    def test_upload(self, mock_upload_file):
+        """Tests uploading model artifacts."""
+        artifact_uploader = ArtifactUploader()
+
+        local_file = "test_files/test.txt"
+        dest_path = "oci://my-bucket@my-namespace/my-artifact-path/logs/test.txt"
+        artifact_uploader.upload(local_file, dest_path)
+
+        mock_upload_file.assert_called_with(
+            namespace_name="my-namespace",
+            bucket_name="my-bucket",
+            object_name="my-artifact-path/logs/test.txt",
+            file_path=local_file,
+        )
+