diff --git a/requirements.txt b/requirements.txt index 9d21a3db..313955e7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,14 +7,13 @@ click >= 8.1.7 colorama >= 0.4.6 fastapi[all] >= 0.94.0 paramiko >= 3.3.1 -pillow >= 9.5.0 -prettytable >= 3.8.0 +prettytable >= 3.9.0 py-cpuinfo >= 9.0.0 python-dotenv >= 1.0.0 inquirerpy == 0.3.4 requests >= 2.31.0 tomli >= 2.0.1 -tomlkit >= 0.12.1 +tomlkit >= 0.12.2 tqdm-loggable >= 0.1.4 urllib3 >= 1.26.6 setuptools_scm >= 8.0.4 diff --git a/runpod/serverless/utils/rp_upload.py b/runpod/serverless/utils/rp_upload.py index 5887443b..e5f34447 100644 --- a/runpod/serverless/utils/rp_upload.py +++ b/runpod/serverless/utils/rp_upload.py @@ -9,12 +9,10 @@ import logging import threading import multiprocessing -from io import BytesIO from urllib.parse import urlparse from typing import Optional, Tuple import boto3 -from PIL import Image, UnidentifiedImageError from boto3 import session from boto3.s3.transfer import TransferConfig from botocore.config import Config @@ -25,6 +23,7 @@ FMT = "%(filename)-20s:%(lineno)-4d %(asctime)s %(message)s" logging.basicConfig(level=logging.INFO, format=FMT, handlers=[logging.StreamHandler()]) + def extract_region_from_url(endpoint_url): """ Extracts the region from the endpoint URL. @@ -43,7 +42,7 @@ def extract_region_from_url(endpoint_url): # --------------------------- S3 Bucket Connection --------------------------- # def get_boto_client( - bucket_creds: Optional[dict] = None) -> Tuple[boto3.client, TransferConfig]: # pragma: no cover + bucket_creds: Optional[dict] = None) -> Tuple[boto3.client, TransferConfig]: # pragma: no cover # pylint: disable=line-too-long ''' Returns a boto3 client and transfer config for the bucket. ''' @@ -94,49 +93,35 @@ def get_boto_client( # ---------------------------------------------------------------------------- # # Upload Image # # ---------------------------------------------------------------------------- # -def upload_image(job_id, image_location, result_index=0, results_list=None): # pragma: no cover +def upload_image(job_id, image_location, result_index=0, results_list=None): # pragma: no cover ''' Upload a single file to bucket storage. ''' image_name = str(uuid.uuid4())[:8] boto_client, _ = get_boto_client() file_extension = os.path.splitext(image_location)[1] + content_type = "image/" + file_extension.lstrip(".") + + with open(image_location, "rb") as input_file: + output = input_file.read() if boto_client is None: # Save the output to a file print("No bucket endpoint set, saving to disk folder 'simulated_uploaded'") print("If this is a live endpoint, please reference the following:") - print("https://github.com/runpod/runpod-python/blob/main/docs/serverless/utils/rp_upload.md") # pylint: disable=line-too-long + print("https://github.com/runpod/runpod-python/blob/main/docs/serverless/utils/rp_upload.md") # pylint: disable=line-too-long os.makedirs("simulated_uploaded", exist_ok=True) sim_upload_location = f"simulated_uploaded/{image_name}{file_extension}" - try: - with Image.open(image_location) as img, open(sim_upload_location, "wb") as file_output: - img.save(file_output, format=img.format) - except UnidentifiedImageError: - # If the file is not an image, save it directly - shutil.copy(image_location, sim_upload_location) + with open(sim_upload_location, "wb") as file_output: + file_output.write(output) if results_list is not None: results_list[result_index] = sim_upload_location return sim_upload_location - try: - with Image.open(image_location) as img: - output = BytesIO() - img.save(output, format=img.format) - output.seek(0) - content_type = "image/" + file_extension.lstrip(".") - - except UnidentifiedImageError: - # If the file is not an image, read it directly - with open(image_location, "rb") as f: - output = f.read() - content_type = "application/octet-stream" - - bucket = time.strftime('%m-%y') boto_client.put_object( Bucket=f'{bucket}', @@ -161,7 +146,7 @@ def upload_image(job_id, image_location, result_index=0, results_list=None): # p # ---------------------------------------------------------------------------- # # Files To Upload # # ---------------------------------------------------------------------------- # -def files(job_id, file_list): # pragma: no cover +def files(job_id, file_list): # pragma: no cover ''' Uploads a list of files in parallel. Once all files are uploaded, the function returns the presigned URLs list. @@ -186,7 +171,7 @@ def files(job_id, file_list): # pragma: no cover # --------------------------- Custom Bucket Upload --------------------------- # -def bucket_upload(job_id, file_list, bucket_creds): # pragma: no cover +def bucket_upload(job_id, file_list, bucket_creds): # pragma: no cover ''' Uploads files to bucket storage. ''' @@ -231,7 +216,7 @@ def upload_file_to_bucket( bucket_name: Optional[str] = None, prefix: Optional[str] = None, extra_args: Optional[dict] = None -) -> str: # pragma: no cover +) -> str: # pragma: no cover ''' Uploads a single file to bucket storage and returns a presigned URL. ''' @@ -283,7 +268,7 @@ def upload_in_memory_object( file_name: str, file_data: bytes, bucket_creds: Optional[dict] = None, bucket_name: Optional[str] = None, - prefix: Optional[str] = None) -> str: # pragma: no cover + prefix: Optional[str] = None) -> str: # pragma: no cover ''' Uploads an in-memory object (bytes) to bucket storage and returns a presigned URL. ''' diff --git a/tests/test_serverless/test_utils/test_upload.py b/tests/test_serverless/test_utils/test_upload.py index 619be95f..43c2d3c4 100644 --- a/tests/test_serverless/test_utils/test_upload.py +++ b/tests/test_serverless/test_utils/test_upload.py @@ -38,7 +38,7 @@ def test_get_boto_client(self): # Mock boto3.session.Session with patch('boto3.session.Session') as mock_session, \ - patch('runpod.serverless.utils.rp_upload.TransferConfig') as mock_transfer_config: + patch('runpod.serverless.utils.rp_upload.TransferConfig') as mock_transfer_config: mock_session.return_value.client.return_value = self.mock_boto_client mock_transfer_config.return_value = self.mock_transfer_config @@ -90,8 +90,6 @@ def test_get_boto_client(self): region_name="region-code" ) - - def test_get_boto_client_environ(self): ''' Tests get_boto_client with environment variables @@ -105,7 +103,7 @@ def test_get_boto_client_environ(self): importlib.reload(rp_upload) with patch('boto3.session.Session') as mock_session, \ - patch('runpod.serverless.utils.rp_upload.TransferConfig') as mock_transfer_config: + patch('runpod.serverless.utils.rp_upload.TransferConfig') as mock_transfer_config: mock_session.return_value.client.return_value = self.mock_boto_client mock_transfer_config.return_value = self.mock_transfer_config @@ -117,13 +115,15 @@ def test_get_boto_client_environ(self): # ---------------------------------------------------------------------------- # # Upload Image # # ---------------------------------------------------------------------------- # + + class TestUploadImage(unittest.TestCase): ''' Tests for upload_image ''' @patch("runpod.serverless.utils.rp_upload.get_boto_client") - @patch("runpod.serverless.utils.rp_upload.Image.open") + @patch("builtins.open") @patch("runpod.serverless.utils.rp_upload.os.makedirs") - def test_upload_image_local(self, mock_makedirs, mock_img_open, mock_get_boto_client): + def test_upload_image_local(self, mock_makedirs, mock_open, mock_get_boto_client): ''' Test upload_image function when there is no boto client ''' @@ -131,25 +131,19 @@ def test_upload_image_local(self, mock_makedirs, mock_img_open, mock_get_boto_cl mock_get_boto_client.return_value = (None, None) # Mocking the context manager of Image.open - mock_image = Mock() - mock_image.format = "PNG" - mock_img_open.return_value.__enter__.return_value = mock_image + mock_file = mock_open.return_value.__enter__.return_value + mock_file.read.return_value = b"simulated_uploaded" + mock_file.__exit__.return_value = False - with patch("builtins.open") as mock_open: - mock_open.return_value = io.BytesIO(b"simulated_uploaded") - result = rp_upload.upload_image("job_id", "image_location") + result = rp_upload.upload_image("job_id", "image_location") # Assert that image is saved locally assert "simulated_uploaded" in result mock_makedirs.assert_called_once() - mock_img_open.assert_called_once() - mock_open.assert_called_once() - mock_image.save.assert_called_once() @patch("runpod.serverless.utils.rp_upload.get_boto_client") - @patch("runpod.serverless.utils.rp_upload.Image.open") - @patch("runpod.serverless.utils.rp_upload.BytesIO") - def test_upload_image_s3(self, mock_bytes_io, mock_open, mock_get_boto_client): + @patch("builtins.open") + def test_upload_image_s3(self, mock_open, mock_get_boto_client): ''' Test upload_image function when there is a boto client ''' @@ -165,22 +159,15 @@ def test_upload_image_s3(self, mock_bytes_io, mock_open, mock_get_boto_client): mock_image.format = "PNG" mock_open.return_value.__enter__.return_value = mock_image - # Mocking BytesIO - mock_bytes_io_instance = Mock() - mock_bytes_io_instance.getvalue = Mock(return_value="image_bytes") - mock_bytes_io.return_value = mock_bytes_io_instance - result = rp_upload.upload_image("job_id", "image_location") # Assert the image is uploaded to S3 assert result == "presigned_url" - mock_open.assert_called_once_with("image_location") + mock_open.assert_called_once_with("image_location", "rb") mock_boto_client.put_object.assert_called_once() mock_boto_client.generate_presigned_url.assert_called_once() - - class TestUploadUtility(unittest.TestCase): ''' Tests for upload utility '''