Skip to content

Commit

Permalink
Merge pull request #133 from runpod/main
Browse files Browse the repository at this point in the history
update DevEx branch
  • Loading branch information
justinmerrell authored Sep 19, 2023
2 parents ae5fb51 + d47d9c5 commit 2f63d0a
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 3 deletions.
2 changes: 1 addition & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -22,4 +22,4 @@ fastapi[all] == 0.103.1
urllib3 >= 1.26.6

# Testing Requirements
nest_asyncio == 1.5.7
nest_asyncio == 1.5.8
19 changes: 18 additions & 1 deletion runpod/serverless/utils/rp_upload.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,24 @@
from boto3.s3.transfer import TransferConfig
from botocore.config import Config
from tqdm_loggable.auto import tqdm
from urllib.parse import urlparse

logger = logging.getLogger("runpod upload utility")
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):
parsed_url = urlparse(endpoint_url)
# AWS/backblaze S3-like URL
if '.s3.' in endpoint_url:
return endpoint_url.split('.s3.')[1].split('.')[0]
# DigitalOcean Spaces-like URL
elif parsed_url.netloc.endswith('.digitaloceanspaces.com'):
return endpoint_url.split('.')[1].split('.digitaloceanspaces.com')[0]
else:
# Additional cases can be added here
return None


# --------------------------- S3 Bucket Connection --------------------------- #
def get_boto_client(
Expand Down Expand Up @@ -57,12 +70,16 @@ def get_boto_client(
secret_access_key = os.environ.get('BUCKET_SECRET_ACCESS_KEY', None)

if endpoint_url and access_key_id and secret_access_key:
# Extract region from the endpoint URL
region = extract_region_from_url(endpoint_url)

boto_client = bucket_session.client(
's3',
endpoint_url=endpoint_url,
aws_access_key_id=access_key_id,
aws_secret_access_key=secret_access_key,
config=boto_config
config=boto_config,
region_name=region
)
else:
boto_client = None
Expand Down
3 changes: 2 additions & 1 deletion tests/test_serverless/test_utils/test_upload.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,8 @@ def test_get_boto_client(self):
endpoint_url=bucket_creds['endpointUrl'],
aws_access_key_id=bucket_creds['accessId'],
aws_secret_access_key=bucket_creds['accessSecret'],
config=unittest.mock.ANY
config=unittest.mock.ANY,
region_name=None
)

def test_get_boto_client_environ(self):
Expand Down

0 comments on commit 2f63d0a

Please sign in to comment.