Repository: kserve/kserve

Issue:
[AWS] Support IAM Role for Service Account in KFServing
/kind feature

**Describe the solution you'd like**
[A clear and concise description of what you want to happen.]
Currently, it needs a S3 credential to download model. We need more fine grain control ways. [IRSA](https://docs.aws.amazon.com/eks/latest/userguide/iam-roles-for-service-accounts.html) is a better option. This helps us deploy model without credentials
 
**Anything else you would like to add:**
[Miscellaneous information that will assist in solving the issue.]


Discussion:
Issue-Label Bot is automatically applying the labels:

| Label  | Probability |
| ------------- | ------------- |
| area/engprod | 0.67 |

Please mark this comment with :thumbsup: or :thumbsdown: to give our bot feedback! 
Links: [app homepage](https://github.com/marketplace/issue-label-bot), [dashboard](https://github.com/marketplace/issue-label-botdata/kubeflow/kfserving) and [code](https://github.com/hamelsmu/MLapp) for this bot.
/assign @Jeffwan 
Issue Label Bot is not confident enough to auto-label this issue.
                See [dashboard](https://github.com/marketplace/issue-label-botdata/kubeflow/kfserving) for more details.
                
Hey @Jeffwan! Great feature. I'm less familiar with the AWS terminology. It looks like this will enable short lived credentials within pods, similar to GKE's managed pod identity. 

I just want to confirm that these credentials will go through the Kubernetes Service Account API and that no API change is necessary from a KFServing perspective. What is the difference from an implementation perspective? IIUC, we currently read some AWS environment variables to retrieve the credentials -- what's different with IRSA?
@ellistarn 

>  It looks like this will enable short lived credentials within pods, similar to GKE's managed pod identity.

Exactly. This will be same as workload identity. 

> I just want to confirm that these credentials will go through the Kubernetes Service Account API and that no API change is necessary from a KFServing perspective. What is the difference from an implementation perspective? IIUC, we currently read some AWS environment variables to retrieve the credentials -- what's different with IRSA?

KFServing side doesn't need to make any API level change or even calls. User needs to choose service account with IAM role attached for inference service. EKS has the webhook to inject the temporary token from IAM role into the pod. There's two things need to be done. 

1. Upgrade AWS or Minio SDK to the version that support `assume-web-identity-role` 
2. Add `assume-web-identity-role` in the credential resolver chain and pod will load token to exchange temporary credentials. 

This maybe supported out-of-box in kfserving and I have not checked it yet. I just create the issue for tracking purpose. 

There is something else that it would be nice to have support for, which is to be able to chain into another role, once the web-identity-role has been configured. That's because our data-lake in S3 doesn't support the new web-identity roles yet, and we need to chain (assume roles) from the one configured by the EKS webhook.

`pod-role-with-web-identity ---[assume]---> data-lake-role ===> download model`

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
@@ -13,22 +13,24 @@
 # limitations under the License.
 
 import glob
+import gzip
 import logging
-import tempfile
 import mimetypes
 import os
 import re
 import json
 import shutil
 import tarfile
+import tempfile
 import zipfile
-import gzip
 from urllib.parse import urlparse
+
+import boto3
 import requests
 from azure.storage.blob import BlockBlobService
 from google.auth import exceptions
 from google.cloud import storage
-from minio import Minio
+
 from kfserving.kfmodel_repository import MODEL_MOUNT_DIRS
 
 _GCS_PREFIX = "gs://"
@@ -81,25 +83,27 @@ def download(uri: str, out_dir: str = None) -> str:
 
     @staticmethod
     def _download_s3(uri, temp_dir: str):
-        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)
-        count = 0
-        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("/")
-            # fget_object handles directory creation if does not exist
-            if not obj.is_dir:
-                if subdir_object_key == "":
-                    subdir_object_key = obj.object_name
-                client.fget_object(bucket_name, obj.object_name,
-                                   os.path.join(temp_dir, subdir_object_key))
-            count = count + 1
-        if count == 0:
-            raise RuntimeError("Failed to fetch model. \
-The path or model %s does not exist." % (uri))
+        s3 = boto3.resource('s3', endpoint_url=os.getenv("AWS_ENDPOINT_URL", "http://s3.amazonaws.com"))
+        parsed = urlparse(uri, scheme='s3')
+        bucket_name = parsed.netloc
+        bucket_path = parsed.path.lstrip('/')
+
+        bucket = s3.Bucket(bucket_name)
+        for obj in bucket.objects.filter(Prefix=bucket_path):
+            # Skip where boto3 lists the directory as an object
+            if obj.key.endswith("/"):
+                continue
+            # In the case where bucket_path points to a single object, set the target key to bucket_path
+            # Otherwise, remove the bucket_path prefix, strip any extra slashes, then prepend the target_dir
+            target_key = (
+                obj.key
+                if bucket_path == obj.key
+                else obj.key.replace(bucket_path, "", 1).lstrip("/")
+            )
+            target = f"{temp_dir}/{target_key}"
+            if not os.path.exists(os.path.dirname(target)):
+                os.makedirs(os.path.dirname(target), exist_ok=True)
+            bucket.download_file(obj.key, target)
 
     @staticmethod
     def _download_gcs(uri, temp_dir: str):
@@ -274,14 +278,3 @@ def _download_from_uri(uri, out_dir=None):
             os.remove(local_path)
 
         return out_dir
-
-    @staticmethod
-    def _create_minio_client():
-        # Adding prefixing "http" in urlparse is necessary for it to be the netloc
-        url = urlparse(os.getenv("AWS_ENDPOINT_URL", "http://s3.amazonaws.com"))
-        use_ssl = url.scheme == 'https' if url.scheme else bool(os.getenv("S3_USE_HTTPS", "true"))
-        return Minio(url.netloc,
-                     access_key=os.getenv("AWS_ACCESS_KEY_ID", ""),
-                     secret_key=os.getenv("AWS_SECRET_ACCESS_KEY", ""),
-                     region=os.getenv("AWS_REGION", ""),
-                     secure=use_ssl)


Test patch:
diff --git a/python/kfserving/test/test_s3_storage.py b/python/kfserving/test/test_s3_storage.py
--- a/python/kfserving/test/test_s3_storage.py
+++ b/python/kfserving/test/test_s3_storage.py
@@ -13,20 +13,26 @@
 # limitations under the License.
 
 import unittest.mock as mock
+
 import kfserving
 
 
 def create_mock_obj(path):
     mock_obj = mock.MagicMock()
-    mock_obj.object_name = path
+    mock_obj.key = path
     mock_obj.is_dir = False
     return mock_obj
 
 
-def create_mock_minio_client(mock_storage, paths):
-    mock_minio_client = mock_storage.return_value
-    mock_minio_client.list_objects.return_value = [create_mock_obj(p) for p in paths]
-    return mock_minio_client
+def create_mock_boto3_bucket(mock_storage, paths):
+    mock_s3_resource = mock.MagicMock()
+    mock_s3_bucket = mock.MagicMock()
+    mock_s3_bucket.objects.filter.return_value = [create_mock_obj(p) for p in paths]
+
+    mock_s3_resource.Bucket.return_value = mock_s3_bucket
+    mock_storage.resource.return_value = mock_s3_resource
+
+    return mock_s3_bucket
 
 
 def get_call_args(call_args_list):
@@ -37,14 +43,14 @@ def get_call_args(call_args_list):
     return arg_list
 
 
-def expected_call_args_list(bucket_name, parent_key, dest, paths):
-    return [(bucket_name, f'{parent_key}/{p}'.strip('/'), f'{dest}/{p}'.strip('/'))
+def expected_call_args_list(parent_key, dest, paths):
+    return [(f'{parent_key}/{p}'.strip('/'), f'{dest}/{p}'.strip('/'))
             for p in paths]
 
 # pylint: disable=protected-access
 
 
-@mock.patch('kfserving.storage.Minio')
+@mock.patch('kfserving.storage.boto3')
 def test_parent_key(mock_storage):
 
     # given
@@ -53,17 +59,17 @@ def test_parent_key(mock_storage):
     object_paths = ['bar/' + p for p in paths]
 
     # when
-    mock_minio_client = create_mock_minio_client(mock_storage, object_paths)
+    mock_boto3_bucket = create_mock_boto3_bucket(mock_storage, object_paths)
     kfserving.Storage._download_s3(f's3://{bucket_name}/bar', 'dest_path')
 
     # then
-    arg_list = get_call_args(mock_minio_client.fget_object.call_args_list)
-    assert arg_list == expected_call_args_list(bucket_name, 'bar', 'dest_path', paths)
+    arg_list = get_call_args(mock_boto3_bucket.download_file.call_args_list)
+    assert arg_list == expected_call_args_list('bar', 'dest_path', paths)
 
-    mock_minio_client.list_objects.assert_called_with(bucket_name, prefix='bar', recursive=True)
+    mock_boto3_bucket.objects.filter.assert_called_with(Prefix='bar')
 
 
-@mock.patch('kfserving.storage.Minio')
+@mock.patch('kfserving.storage.boto3')
 def test_no_key(mock_storage):
 
     # given
@@ -71,17 +77,17 @@ def test_no_key(mock_storage):
     object_paths = ['models/weights.pt', '0002.h5', 'a/very/long/path/config.json']
 
     # when
-    mock_minio_client = create_mock_minio_client(mock_storage, object_paths)
+    mock_boto3_bucket = create_mock_boto3_bucket(mock_storage, object_paths)
     kfserving.Storage._download_s3(f's3://{bucket_name}/', 'dest_path')
 
     # then
-    arg_list = get_call_args(mock_minio_client.fget_object.call_args_list)
-    assert arg_list == expected_call_args_list(bucket_name, '', 'dest_path', object_paths)
+    arg_list = get_call_args(mock_boto3_bucket.download_file.call_args_list)
+    assert arg_list == expected_call_args_list('', 'dest_path', object_paths)
 
-    mock_minio_client.list_objects.assert_called_with(bucket_name, prefix='', recursive=True)
+    mock_boto3_bucket.objects.filter.assert_called_with(Prefix='')
 
 
-@mock.patch('kfserving.storage.Minio')
+@mock.patch('kfserving.storage.boto3')
 def test_full_name_key(mock_storage):
 
     # given
@@ -89,13 +95,12 @@ def test_full_name_key(mock_storage):
     object_key = 'path/to/model/name.pt'
 
     # when
-    mock_minio_client = create_mock_minio_client(mock_storage, [object_key])
+    mock_boto3_bucket = create_mock_boto3_bucket(mock_storage, [object_key])
     kfserving.Storage._download_s3(f's3://{bucket_name}/{object_key}', 'dest_path')
 
     # then
-    arg_list = get_call_args(mock_minio_client.fget_object.call_args_list)
-    assert arg_list == expected_call_args_list(bucket_name, '', 'dest_path',
+    arg_list = get_call_args(mock_boto3_bucket.download_file.call_args_list)
+    assert arg_list == expected_call_args_list('', 'dest_path',
                                                [object_key])
 
-    mock_minio_client.list_objects.assert_called_with(bucket_name, prefix=object_key,
-                                                      recursive=True)
+    mock_boto3_bucket.objects.filter.assert_called_with(Prefix=object_key)
diff --git a/python/kfserving/test/test_storage.py b/python/kfserving/test/test_storage.py
--- a/python/kfserving/test/test_storage.py
+++ b/python/kfserving/test/test_storage.py
@@ -16,11 +16,11 @@
 import os
 import tempfile
 import binascii
+import unittest.mock as mock
 
-import pytest
+import botocore
 import kfserving
-from minio import Minio, error
-import unittest.mock as mock
+import pytest
 
 STORAGE_MODULE = 'kfserving.storage'
 
@@ -166,25 +166,31 @@ def test_storage_blob_exception():
         kfserving.Storage.download(blob_path)
 
 
-@mock.patch('urllib3.PoolManager')
-@mock.patch(STORAGE_MODULE + '.Minio')
-def test_storage_s3_exception(mock_connection, mock_minio):
-    minio_path = 's3://foo/bar'
-    # Create mock connection
-    mock_server = mock.MagicMock()
-    mock_connection.return_value = mock_server
+@mock.patch(STORAGE_MODULE + '.boto3')
+def test_storage_s3_exception(mock_boto3):
+    path = 's3://foo/bar'
     # Create mock client
-    mock_minio.return_value = Minio("s3.us.cloud-object-storage.appdomain.cloud", secure=True)
+    mock_s3_resource = mock.MagicMock()
+    mock_s3_resource.Bucket.side_effect = Exception()
+    mock_boto3.resource.return_value = mock_s3_resource
+
     with pytest.raises(Exception):
-        kfserving.Storage.download(minio_path)
+        kfserving.Storage.download(path)
 
 
+@mock.patch(STORAGE_MODULE + '.boto3')
 @mock.patch('urllib3.PoolManager')
-@mock.patch(STORAGE_MODULE + '.Minio')
-def test_no_permission_buckets(mock_connection, mock_minio):
+def test_no_permission_buckets(mock_connection, mock_boto3):
     bad_s3_path = "s3://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()
-    with pytest.raises(error.AccessDenied):
+    mock_s3_resource = mock.MagicMock()
+    mock_s3_bucket = mock.MagicMock()
+    mock_s3_bucket.objects.filter.return_value = [mock.MagicMock()]
+    mock_s3_bucket.objects.filter.side_effect = botocore.exceptions.ClientError(
+        {}, "GetObject"
+    )
+    mock_s3_resource.Bucket.return_value = mock_s3_bucket
+    mock_boto3.resource.return_value = mock_s3_resource
+
+    with pytest.raises(botocore.exceptions.ClientError):
         kfserving.Storage.download(bad_s3_path)