Repository: kserve/kserve

Issue:
Python Model download for GCS and S3
Downloading from GCS and S3 needs to be completed.

https://github.com/kubeflow/kfserving/blob/2f8d33d1a9773c5694a22ba749192163251fe287/python/kfserving/kfserving/storage.py#L27-L33


Discussion:

Resolving patch:
diff --git a/python/kfserving/kfserving/storage.py b/python/kfserving/kfserving/storage.py
--- a/python/kfserving/kfserving/storage.py
+++ b/python/kfserving/kfserving/storage.py
@@ -15,6 +15,10 @@
 import logging
 import tempfile
 import os
+import re
+from minio import Minio
+from google.cloud import storage
+from google.auth import exceptions
 
 _GCS_PREFIX = "gs://"
 _S3_PREFIX = "s3://"
@@ -43,11 +47,36 @@ def download(uri: str) -> str:
 
     @staticmethod
     def _download_s3(uri, temp_dir: str):
-        raise NotImplementedError
+        client = Storage._create_minio_client()
+        bucket_args = uri.replace(_S3_PREFIX, "", 1).split("/", 1)
+        bucket_name = bucket_args[0]
+        bucket_path = bucket_args[1] if len(bucket_args) > 1 else ""
+        objects = client.list_objects(bucket_name, prefix=bucket_path, recursive=True)
+        for obj in objects:
+            # Replace any prefix from the object key with temp_dir
+            subdir_object_key = obj.object_name.replace(bucket_path, "", 1).strip("/")
+            client.fget_object(bucket_name, obj.object_name, os.path.join(temp_dir, subdir_object_key))
 
     @staticmethod
     def _download_gcs(uri, temp_dir: str):
-        raise NotImplementedError
+        try:
+            storage_client = storage.Client()
+        except exceptions.DefaultCredentialsError as e:
+            storage_client = storage.Client.create_anonymous_client()
+        bucket_args = uri.replace(_GCS_PREFIX, "", 1).split("/", 1)
+        bucket_name = bucket_args[0]
+        bucket_path = bucket_args[1] if len(bucket_args) > 1 else ""
+        bucket = storage_client.bucket(bucket_name)
+        blobs = bucket.list_blobs(prefix=bucket_path)
+        for blob in blobs:
+            # Replace any prefix from the object key with temp_dir
+            subdir_object_key = blob.name.replace(bucket_path, "", 1).strip("/")
+            # Create necessary subdirectory to store the object locally
+            if "/" in subdir_object_key:
+                local_object_dir = os.path.join(temp_dir, subdir_object_key.rsplit("/", 1)[0])
+                if not os.path.isdir(local_object_dir):
+                    os.makedirs(local_object_dir, exist_ok=True)
+            blob.download_to_filename(os.path.join(temp_dir, subdir_object_key))
 
     @staticmethod
     def _download_local(uri):
@@ -55,3 +84,13 @@ def _download_local(uri):
         if not os.path.exists(local_path):
             raise Exception("Local path %s does not exist." % (uri))
         return local_path
+
+    @staticmethod
+    def _create_minio_client():
+        # Remove possible http scheme for Minio
+        url = re.compile(r"https?://")
+        minioClient = Minio(url.sub("", os.getenv("S3_ENDPOINT", "")),
+                            access_key=os.getenv("AWS_ACCESS_KEY_ID", ""),
+                            secret_key=os.getenv("AWS_SECRET_ACCESS_KEY", ""),
+                            secure=True)
+        return minioClient
diff --git a/python/kfserving/setup.py b/python/kfserving/setup.py
--- a/python/kfserving/setup.py
+++ b/python/kfserving/setup.py
@@ -33,6 +33,8 @@
     install_requires=[
         "tornado >= 1.4.1",
         "argparse >= 1.4.0",
+        "minio >= 4.0.9",
+        "google-cloud-storage >= 1.16.0",
         "numpy"
     ],
     tests_require=tests_require,


Test patch:
diff --git a/python/kfserving/kfserving/test_storage.py b/python/kfserving/kfserving/test_storage.py
--- a/python/kfserving/kfserving/test_storage.py
+++ b/python/kfserving/kfserving/test_storage.py
@@ -14,6 +14,12 @@
 
 import pytest
 import kfserving
+import os
+from minio import Minio, error
+from google.cloud import exceptions
+import unittest.mock as mock
+
+STORAGE_MODULE = 'kfserving.storage'
 
 
 def test_storage_local_path():
@@ -34,3 +40,42 @@ def test_no_prefix_local_path():
     relative_path = '.'
     assert kfserving.Storage.download(abs_path) == abs_path
     assert kfserving.Storage.download(relative_path) == relative_path
+
+
+@mock.patch(STORAGE_MODULE + '.storage')
+def test_mock_gcs(mock_storage):
+    gcs_path = 'gs://foo/bar'
+    mock_obj = mock.MagicMock()
+    mock_obj.name = 'mock.object'
+    mock_storage.Client().bucket().list_blobs().__iter__.return_value = [mock_obj]
+    assert kfserving.Storage.download(gcs_path)
+
+
+@mock.patch('urllib3.PoolManager')
+@mock.patch(STORAGE_MODULE + '.Minio')
+def test_mock_minio(mock_connection, mock_minio):
+    minio_path = 's3://foo/bar'
+    # Create mock connection
+    mock_server = mock.MagicMock()
+    mock_connection.return_value = mock_server
+    # Create mock client
+    mock_minio.return_value = Minio("s3.us.cloud-object-storage.appdomain.cloud", secure=True)
+    mock_obj = mock.MagicMock()
+    mock_obj.object_name = 'mock.object'
+    mock_minio.list_objects().__iter__.return_value = [mock_obj]
+    assert kfserving.Storage.download(minio_path)
+
+
+@mock.patch('urllib3.PoolManager')
+@mock.patch(STORAGE_MODULE + '.Minio')
+def test_no_permission_buckets(mock_connection, mock_minio):
+    bad_s3_path = "s3://random/path"
+    bad_gcs_path = "gs://random/path"
+    # Access private buckets without credentials
+    mock_minio.return_value = Minio("s3.us.cloud-object-storage.appdomain.cloud", secure=True)
+    mock_connection.side_effect = error.AccessDenied(None)
+    with pytest.raises(error.AccessDenied):
+        kfserving.Storage.download(bad_s3_path)
+    mock_connection.side_effect = exceptions.Forbidden(None)
+    with pytest.raises(exceptions.Forbidden):
+        kfserving.Storage.download(bad_gcs_path)