Skip to content

Commit cc05a5b

Browse files
authored
perf(serverless): lazy-load boto3, fastapi, and pydantic to reduce cold start time (#466)
* 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. * refactor(serverless): address PR review feedback for boto3 lazy-loading Address Copilot review comments from PR #466: 1. Extract boto3 import logic into shared helper function - Created _import_boto3_dependencies() to reduce duplication - Used by both get_boto_client() and bucket_upload() - Consistent error handling across all S3 upload functions 2. Add TransferConfig import to bucket_upload() - Now imports all boto3 dependencies via shared helper - Maintains consistency with get_boto_client() 3. Clarify documentation about fallback directories - Documented that upload_image() uses simulated_uploaded/ - Documented that public API functions use local_upload/ - Added context about when fallback behavior occurs All 358 tests pass with 97% coverage. * refactor(serverless): replace print() with logger in upload fallback paths Address Copilot review comment from PR #466: Replace print() statements with logger.warning() for consistency with the module's logging setup. This allows proper log level control and maintains consistent logging behavior throughout the module. Changes: - upload_image(): Use logger.warning() instead of print() - upload_file_to_bucket(): Use logger.warning() instead of print() - upload_in_memory_object(): Use logger.warning() instead of print() All fallback messages now use structured logging with single-line format for better log parsing and filtering. * refactor(serverless): extract local fallback logic into shared helper Address Copilot review comment from PR #466: Extract duplicated fallback logic from upload_file_to_bucket() and upload_in_memory_object() into a shared _save_to_local_fallback() helper function to reduce code duplication and ensure consistent behavior. Changes: - Created _save_to_local_fallback() helper function - Handles both file-based (source_path) and in-memory (file_data) uploads - Consolidated logging, directory creation, and file saving logic - upload_file_to_bucket() now calls helper with source_path parameter - upload_in_memory_object() now calls helper with file_data parameter Test improvements: - Added test_upload_file_to_bucket_fallback() for file-based fallback - Added test_upload_in_memory_object_fallback() for in-memory fallback - Added test_save_to_local_fallback_invalid_args() for error handling - Added test_import_boto3_dependencies_missing() for ImportError path - Achieved 100% test coverage for rp_upload.py module Benefits: - Reduced code duplication (removed 12 lines of duplicate code) - Single source of truth for fallback behavior - Easier to maintain and test - Consistent error messages and logging - Complete test coverage ensures reliability All 10 upload tests pass with 100% module coverage. * refactor(serverless): consolidate fallback logic across all upload functions Address latest Copilot review comment from PR #466: Eliminate remaining duplication by making upload_image() use the _save_to_local_fallback() helper function. Added a 'directory' parameter to the helper to support different fallback directories. Changes: - Added 'directory' parameter to _save_to_local_fallback() (default: 'local_upload') - Updated upload_image() to use helper with directory='simulated_uploaded' - Removed duplicate warning message and URL from upload_image() - Consolidated all fallback logic into single helper function Benefits: - Complete elimination of code duplication - Single source of truth for all fallback behavior - Consistent warning messages across all upload functions - Easier to maintain and update fallback logic All 362 tests pass with 97% overall coverage, 100% coverage on rp_upload.py. * refactor(serverless): restore type hints for boto3 lazy-loading Use TYPE_CHECKING to import boto3 types only during static type checking, maintaining proper type safety without runtime import cost. Changes: - Import BaseClient and TransferConfig under TYPE_CHECKING guard - Restore get_boto_client() return type from Tuple[Any, Any] to Tuple[Optional[BaseClient], Optional[TransferConfig]] - Remove # pragma: no cover comment as it's no longer needed This addresses PR review feedback about maintaining type safety while preserving the lazy-loading optimization. * perf(serverless): lazy-load FastAPI to reduce cold start time Move FastAPI/Uvicorn/Pydantic imports from module-level to conditional blocks where they're actually used. This stack is only needed when --rp_serve_api flag is set (local dev) or realtime mode is enabled. Performance Impact: - Cold start: 480ms → 280-326ms (32-42% faster) - Modules loaded: 841 → 640 (24% reduction, ~200 fewer) - Production workers: Never load FastAPI/Uvicorn/Pydantic stack - Dev mode: FastAPI loads on-demand when needed Changes: - Remove eager import of rp_fastapi from module level - Add lazy import in start() when rp_serve_api flag is True - Add lazy import in start() when realtime mode is enabled All tests pass. No breaking changes. * test(serverless): fix tests for lazy-loaded FastAPI Update test mocks to use correct import path for lazy-loaded rp_fastapi module. Since FastAPI is now imported on-demand inside start() function rather than at module level, tests need to mock the actual module path. Changes: - Update test_local_api to mock runpod.serverless.modules.rp_fastapi.WorkerAPI - Update test_start_does_not_set_excepthook to mock correct module path All 362 tests pass with 96.76% coverage.
1 parent 227ef38 commit cc05a5b

File tree

6 files changed

+179
-40
lines changed

6 files changed

+179
-40
lines changed

docs/serverless/utils/rp_upload.md

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,14 @@ 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 or S3 credentials are not configured, files will be saved to local disk instead:
12+
- `upload_image()` saves to `simulated_uploaded/` directory
13+
- `upload_file_to_bucket()` and `upload_in_memory_object()` save to `local_upload/` directory
14+
715
## Bucket Credentials
816

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

runpod/serverless/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414

1515
from ..version import __version__ as runpod_version
1616
from . import worker
17-
from .modules import rp_fastapi
1817
from .modules.rp_logger import RunPodLogger
1918
from .modules.rp_progress import progress_update
2019

@@ -155,6 +154,7 @@ def start(config: Dict[str, Any]):
155154

156155
if config["rp_args"]["rp_serve_api"]:
157156
log.info("Starting API server.")
157+
from .modules import rp_fastapi
158158
api_server = rp_fastapi.WorkerAPI(config)
159159

160160
api_server.start_uvicorn(
@@ -166,6 +166,7 @@ def start(config: Dict[str, Any]):
166166

167167
if realtime_port:
168168
log.info(f"Starting API server for realtime on port {realtime_port}.")
169+
from .modules import rp_fastapi
169170
api_server = rp_fastapi.WorkerAPI(config)
170171

171172
api_server.start_uvicorn(

runpod/serverless/modules/rp_scale.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -195,7 +195,7 @@ async def get_jobs(self, session: ClientSession):
195195

196196
except TooManyRequests:
197197
log.debug(
198-
f"JobScaler.get_jobs | Too many requests. Debounce for 5 seconds."
198+
"JobScaler.get_jobs | Too many requests. Debounce for 5 seconds."
199199
)
200200
await asyncio.sleep(5) # debounce for 5 seconds
201201
except asyncio.CancelledError:

runpod/serverless/utils/rp_upload.py

Lines changed: 91 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -10,20 +10,75 @@
1010
import threading
1111
import time
1212
import uuid
13-
from typing import Optional, Tuple
13+
from typing import TYPE_CHECKING, 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

18+
if TYPE_CHECKING:
19+
from boto3.s3.transfer import TransferConfig
20+
from botocore.client import BaseClient
21+
2222
logger = logging.getLogger("runpod upload utility")
2323
FMT = "%(filename)-20s:%(lineno)-4d %(asctime)s %(message)s"
2424
logging.basicConfig(level=logging.INFO, format=FMT, handlers=[logging.StreamHandler()])
2525

2626

27+
def _import_boto3_dependencies():
28+
"""
29+
Lazy-load boto3 dependencies.
30+
Returns tuple of (session, TransferConfig, Config) or raises ImportError.
31+
"""
32+
try:
33+
from boto3 import session
34+
from boto3.s3.transfer import TransferConfig
35+
from botocore.config import Config
36+
return session, TransferConfig, Config
37+
except ImportError as e:
38+
raise ImportError(
39+
"boto3 is required for S3 upload functionality. "
40+
"Install with: pip install boto3"
41+
) from e
42+
43+
44+
def _save_to_local_fallback(
45+
file_name: str,
46+
source_path: Optional[str] = None,
47+
file_data: Optional[bytes] = None,
48+
directory: str = "local_upload"
49+
) -> str:
50+
"""
51+
Save file to local directory as fallback when S3 is unavailable.
52+
53+
Args:
54+
file_name: Name of the file to save
55+
source_path: Path to source file to copy (for file-based uploads)
56+
file_data: Bytes to write (for in-memory uploads)
57+
directory: Local directory to save to (default: 'local_upload')
58+
59+
Returns:
60+
Path to the saved local file
61+
"""
62+
logger.warning(
63+
f"No bucket endpoint set, saving to disk folder '{directory}'. "
64+
"If this is a live endpoint, please reference: "
65+
"https://github.com/runpod/runpod-python/blob/main/docs/serverless/utils/rp_upload.md"
66+
)
67+
68+
os.makedirs(directory, exist_ok=True)
69+
local_upload_location = f"{directory}/{file_name}"
70+
71+
if source_path:
72+
shutil.copyfile(source_path, local_upload_location)
73+
elif file_data is not None:
74+
with open(local_upload_location, "wb") as file_output:
75+
file_output.write(file_data)
76+
else:
77+
raise ValueError("Either source_path or file_data must be provided")
78+
79+
return local_upload_location
80+
81+
2782
def extract_region_from_url(endpoint_url):
2883
"""
2984
Extracts the region from the endpoint URL.
@@ -43,12 +98,20 @@ def extract_region_from_url(endpoint_url):
4398
# --------------------------- S3 Bucket Connection --------------------------- #
4499
def get_boto_client(
45100
bucket_creds: Optional[dict] = None,
46-
) -> Tuple[
47-
boto3.client, TransferConfig
48-
]: # pragma: no cover # pylint: disable=line-too-long
101+
) -> Tuple[Optional["BaseClient"], Optional["TransferConfig"]]:
49102
"""
50103
Returns a boto3 client and transfer config for the bucket.
104+
Lazy-loads boto3 to reduce initial import time.
51105
"""
106+
try:
107+
session, TransferConfig, Config = _import_boto3_dependencies()
108+
except ImportError:
109+
logger.warning(
110+
"boto3 not installed. S3 upload functionality disabled. "
111+
"Install with: pip install boto3"
112+
)
113+
return None, None
114+
52115
bucket_session = session.Session()
53116

54117
boto_config = Config(
@@ -111,18 +174,13 @@ def upload_image(
111174
output = input_file.read()
112175

113176
if boto_client is None:
114-
# Save the output to a file
115-
print("No bucket endpoint set, saving to disk folder 'simulated_uploaded'")
116-
print("If this is a live endpoint, please reference the following:")
117-
print(
118-
"https://github.com/runpod/runpod-python/blob/main/docs/serverless/utils/rp_upload.md"
119-
) # pylint: disable=line-too-long
120-
121-
os.makedirs("simulated_uploaded", exist_ok=True)
122-
sim_upload_location = f"simulated_uploaded/{image_name}{file_extension}"
123-
124-
with open(sim_upload_location, "wb") as file_output:
125-
file_output.write(output)
177+
# Save the output to a file using fallback helper
178+
file_name_with_ext = f"{image_name}{file_extension}"
179+
sim_upload_location = _save_to_local_fallback(
180+
file_name_with_ext,
181+
file_data=output,
182+
directory="simulated_uploaded"
183+
)
126184

127185
if results_list is not None:
128186
results_list[result_index] = sim_upload_location
@@ -180,6 +238,15 @@ def bucket_upload(job_id, file_list, bucket_creds): # pragma: no cover
180238
"""
181239
Uploads files to bucket storage.
182240
"""
241+
try:
242+
session, _, Config = _import_boto3_dependencies()
243+
except ImportError:
244+
logger.error(
245+
"boto3 not installed. Cannot upload to S3 bucket. "
246+
"Install with: pip install boto3"
247+
)
248+
raise
249+
183250
temp_bucket_session = session.Session()
184251

185252
temp_boto_config = Config(
@@ -231,17 +298,7 @@ def upload_file_to_bucket(
231298
key = f"{prefix}/{file_name}" if prefix else file_name
232299

233300
if boto_client is None:
234-
print("No bucket endpoint set, saving to disk folder 'local_upload'")
235-
print("If this is a live endpoint, please reference the following:")
236-
print(
237-
"https://github.com/runpod/runpod-python/blob/main/docs/serverless/utils/rp_upload.md"
238-
) # pylint: disable=line-too-long
239-
240-
os.makedirs("local_upload", exist_ok=True)
241-
local_upload_location = f"local_upload/{file_name}"
242-
shutil.copyfile(file_location, local_upload_location)
243-
244-
return local_upload_location
301+
return _save_to_local_fallback(file_name, source_path=file_location)
245302

246303
file_size = os.path.getsize(file_location)
247304
with tqdm(
@@ -285,6 +342,9 @@ def upload_in_memory_object(
285342

286343
key = f"{prefix}/{file_name}" if prefix else file_name
287344

345+
if boto_client is None:
346+
return _save_to_local_fallback(file_name, file_data=file_data)
347+
288348
file_size = len(file_data)
289349
with tqdm(
290350
total=file_size, unit="B", unit_scale=True, desc=file_name

tests/test_serverless/test_utils/test_upload.py

Lines changed: 73 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,16 +32,25 @@ def setUp(self) -> None:
3232
def tearDown(self):
3333
os.environ = self.original_environ
3434

35+
def test_import_boto3_dependencies_missing(self):
36+
"""
37+
Tests _import_boto3_dependencies when boto3 is not available
38+
"""
39+
with patch("builtins.__import__", side_effect=ImportError("No module named 'boto3'")):
40+
with self.assertRaises(ImportError) as context:
41+
rp_upload._import_boto3_dependencies()
42+
self.assertIn("boto3 is required for S3 upload functionality", str(context.exception))
43+
3544
def test_get_boto_client(self):
3645
"""
3746
Tests get_boto_client
3847
"""
3948
# Define the bucket credentials
4049
bucket_creds = BUCKET_CREDENTIALS
4150

42-
# Mock boto3.session.Session
51+
# Mock boto3 imports (now lazy-loaded inside the function)
4352
with patch("boto3.session.Session") as mock_session, patch(
44-
"runpod.serverless.utils.rp_upload.TransferConfig"
53+
"boto3.s3.transfer.TransferConfig"
4554
) as mock_transfer_config:
4655
mock_session.return_value.client.return_value = self.mock_boto_client
4756
mock_transfer_config.return_value = self.mock_transfer_config
@@ -110,8 +119,9 @@ def test_get_boto_client_environ(self):
110119

111120
importlib.reload(rp_upload)
112121

122+
# Mock boto3 imports (now lazy-loaded inside the function)
113123
with patch("boto3.session.Session") as mock_session, patch(
114-
"runpod.serverless.utils.rp_upload.TransferConfig"
124+
"boto3.s3.transfer.TransferConfig"
115125
) as mock_transfer_config:
116126
mock_session.return_value.client.return_value = self.mock_boto_client
117127
mock_transfer_config.return_value = self.mock_transfer_config
@@ -178,9 +188,46 @@ def test_upload_image_s3(self, mock_open, mock_get_boto_client):
178188
mock_boto_client.generate_presigned_url.assert_called_once()
179189

180190

191+
class TestLocalFallback(unittest.TestCase):
192+
"""Tests for _save_to_local_fallback helper function"""
193+
194+
@patch("os.makedirs")
195+
def test_save_to_local_fallback_invalid_args(self, mock_makedirs):
196+
"""
197+
Tests _save_to_local_fallback raises ValueError when neither source_path nor file_data provided
198+
"""
199+
with self.assertRaises(ValueError) as context:
200+
rp_upload._save_to_local_fallback("test.txt")
201+
self.assertIn("Either source_path or file_data must be provided", str(context.exception))
202+
203+
181204
class TestUploadUtility(unittest.TestCase):
182205
"""Tests for upload utility"""
183206

207+
@patch("runpod.serverless.utils.rp_upload.get_boto_client")
208+
@patch("os.path.exists")
209+
@patch("shutil.copyfile")
210+
@patch("os.makedirs")
211+
def test_upload_file_to_bucket_fallback(
212+
self, mock_makedirs, mock_copyfile, mock_exists, mock_get_boto_client
213+
):
214+
"""
215+
Tests upload_file_to_bucket fallback when boto_client is None
216+
"""
217+
# Mock get_boto_client to return None
218+
mock_get_boto_client.return_value = (None, None)
219+
mock_exists.return_value = True
220+
221+
file_name = "example.txt"
222+
file_location = "/path/to/file.txt"
223+
224+
result = upload_file_to_bucket(file_name, file_location)
225+
226+
# Check fallback behavior
227+
assert result == "local_upload/example.txt"
228+
mock_makedirs.assert_called_once_with("local_upload", exist_ok=True)
229+
mock_copyfile.assert_called_once_with(file_location, "local_upload/example.txt")
230+
184231
@patch("runpod.serverless.utils.rp_upload.get_boto_client")
185232
def test_upload_file_to_bucket(self, mock_get_boto_client):
186233
"""
@@ -220,6 +267,29 @@ def test_upload_file_to_bucket(self, mock_get_boto_client):
220267
ExpiresIn=604800,
221268
)
222269

270+
@patch("runpod.serverless.utils.rp_upload.get_boto_client")
271+
@patch("builtins.open", new_callable=unittest.mock.mock_open)
272+
@patch("os.makedirs")
273+
def test_upload_in_memory_object_fallback(
274+
self, mock_makedirs, mock_open_file, mock_get_boto_client
275+
):
276+
"""
277+
Tests upload_in_memory_object fallback when boto_client is None
278+
"""
279+
# Mock get_boto_client to return None
280+
mock_get_boto_client.return_value = (None, None)
281+
282+
file_name = "example.txt"
283+
file_data = b"This is test data."
284+
285+
result = upload_in_memory_object(file_name, file_data)
286+
287+
# Check fallback behavior
288+
assert result == "local_upload/example.txt"
289+
mock_makedirs.assert_called_once_with("local_upload", exist_ok=True)
290+
mock_open_file.assert_called_once_with("local_upload/example.txt", "wb")
291+
mock_open_file().write.assert_called_once_with(file_data)
292+
223293
@patch("runpod.serverless.utils.rp_upload.get_boto_client")
224294
def test_upload_in_memory_object(self, mock_get_boto_client):
225295
"""

tests/test_serverless/test_worker.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -84,13 +84,13 @@ def test_local_api(self):
8484
with patch(
8585
"argparse.ArgumentParser.parse_known_args"
8686
) as mock_parse_known_args, patch(
87-
"runpod.serverless.rp_fastapi"
88-
) as mock_fastapi:
87+
"runpod.serverless.modules.rp_fastapi.WorkerAPI"
88+
) as mock_worker_api:
8989

9090
mock_parse_known_args.return_value = known_args, []
9191
runpod.serverless.start({"handler": self.mock_handler})
9292

93-
assert mock_fastapi.WorkerAPI.called
93+
assert mock_worker_api.called
9494

9595
@patch("runpod.serverless.log")
9696
@patch("runpod.serverless.sys.exit")
@@ -544,7 +544,7 @@ def test_start_sets_excepthook(self, _, __):
544544
assert sys.excepthook == _handle_uncaught_exception
545545

546546
@patch("runpod.serverless.signal.signal")
547-
@patch("runpod.serverless.rp_fastapi.WorkerAPI.start_uvicorn")
547+
@patch("runpod.serverless.modules.rp_fastapi.WorkerAPI.start_uvicorn")
548548
@patch("runpod.serverless._set_config_args")
549549
def test_start_does_not_set_excepthook(self, mock_set_config_args, _, __):
550550
mock_set_config_args.return_value = self.config

0 commit comments

Comments
 (0)