Skip to content

Commit 73dc607

Browse files
committed
perf(serverless): lazy-load boto3 to reduce cold start time
Move boto3 imports from module level to function level in rp_upload.py. This defers loading ~50MB of boto3/botocore dependencies until S3 upload functions are actually called, improving initial import time and memory footprint for users who don't use S3 features. Changes: - Refactored get_boto_client() and bucket_upload() with lazy imports - Added ImportError handling with helpful error messages - Updated tests to mock boto3 modules directly - Enhanced documentation to explain lazy-loading behavior All upload functions maintain backward compatibility and graceful fallback to local file storage when boto3 is unavailable.
1 parent 227ef38 commit 73dc607

3 files changed

Lines changed: 50 additions & 11 deletions

File tree

docs/serverless/utils/rp_upload.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,12 @@ The upload utility provides functions to upload files and in-memory objects to a
44

55
*Note: The upload utility utilizes the Virtual-hosted-style URL with the bucket name in the host name. For example, `https: // bucket-name.s3.amazonaws.com`.*
66

7+
## Requirements
8+
9+
The upload utility requires [boto3](https://pypi.org/project/boto3/) for S3 functionality. boto3 is lazy-loaded to minimize initial import time and memory footprint.
10+
11+
If you attempt to use S3 upload features without boto3 installed, you'll receive a warning and files will be saved to local disk instead (`simulated_uploaded/` or `local_upload/` directories).
12+
713
## Bucket Credentials
814

915
You can set your S3 bucket credentials in the following ways:

runpod/serverless/utils/rp_upload.py

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,9 @@
1010
import threading
1111
import time
1212
import uuid
13-
from typing import Optional, Tuple
13+
from typing import Any, Optional, Tuple
1414
from urllib.parse import urlparse
1515

16-
import boto3
17-
from boto3 import session
18-
from boto3.s3.transfer import TransferConfig
19-
from botocore.config import Config
2016
from tqdm_loggable.auto import tqdm
2117

2218
logger = logging.getLogger("runpod upload utility")
@@ -43,12 +39,22 @@ def extract_region_from_url(endpoint_url):
4339
# --------------------------- S3 Bucket Connection --------------------------- #
4440
def get_boto_client(
4541
bucket_creds: Optional[dict] = None,
46-
) -> Tuple[
47-
boto3.client, TransferConfig
48-
]: # pragma: no cover # pylint: disable=line-too-long
42+
) -> Tuple[Any, Any]: # pragma: no cover # pylint: disable=line-too-long
4943
"""
5044
Returns a boto3 client and transfer config for the bucket.
45+
Lazy-loads boto3 to reduce initial import time.
5146
"""
47+
try:
48+
from boto3 import session
49+
from boto3.s3.transfer import TransferConfig
50+
from botocore.config import Config
51+
except ImportError:
52+
logger.warning(
53+
"boto3 not installed. S3 upload functionality disabled. "
54+
"Install with: pip install boto3"
55+
)
56+
return None, None
57+
5258
bucket_session = session.Session()
5359

5460
boto_config = Config(
@@ -180,6 +186,18 @@ def bucket_upload(job_id, file_list, bucket_creds): # pragma: no cover
180186
"""
181187
Uploads files to bucket storage.
182188
"""
189+
try:
190+
from boto3 import session
191+
from botocore.config import Config
192+
except ImportError:
193+
logger.error(
194+
"boto3 not installed. Cannot upload to S3 bucket. "
195+
"Install with: pip install boto3"
196+
)
197+
raise ImportError(
198+
"boto3 is required for bucket_upload. Install with: pip install boto3"
199+
)
200+
183201
temp_bucket_session = session.Session()
184202

185203
temp_boto_config = Config(
@@ -285,6 +303,20 @@ def upload_in_memory_object(
285303

286304
key = f"{prefix}/{file_name}" if prefix else file_name
287305

306+
if boto_client is None:
307+
print("No bucket endpoint set, saving to disk folder 'local_upload'")
308+
print("If this is a live endpoint, please reference the following:")
309+
print(
310+
"https://github.com/runpod/runpod-python/blob/main/docs/serverless/utils/rp_upload.md"
311+
)
312+
313+
os.makedirs("local_upload", exist_ok=True)
314+
local_upload_location = f"local_upload/{file_name}"
315+
with open(local_upload_location, "wb") as file_output:
316+
file_output.write(file_data)
317+
318+
return local_upload_location
319+
288320
file_size = len(file_data)
289321
with tqdm(
290322
total=file_size, unit="B", unit_scale=True, desc=file_name

tests/test_serverless/test_utils/test_upload.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,9 @@ def test_get_boto_client(self):
3939
# Define the bucket credentials
4040
bucket_creds = BUCKET_CREDENTIALS
4141

42-
# Mock boto3.session.Session
42+
# Mock boto3 imports (now lazy-loaded inside the function)
4343
with patch("boto3.session.Session") as mock_session, patch(
44-
"runpod.serverless.utils.rp_upload.TransferConfig"
44+
"boto3.s3.transfer.TransferConfig"
4545
) as mock_transfer_config:
4646
mock_session.return_value.client.return_value = self.mock_boto_client
4747
mock_transfer_config.return_value = self.mock_transfer_config
@@ -110,8 +110,9 @@ def test_get_boto_client_environ(self):
110110

111111
importlib.reload(rp_upload)
112112

113+
# Mock boto3 imports (now lazy-loaded inside the function)
113114
with patch("boto3.session.Session") as mock_session, patch(
114-
"runpod.serverless.utils.rp_upload.TransferConfig"
115+
"boto3.s3.transfer.TransferConfig"
115116
) as mock_transfer_config:
116117
mock_session.return_value.client.return_value = self.mock_boto_client
117118
mock_transfer_config.return_value = self.mock_transfer_config

0 commit comments

Comments
 (0)