Skip to content

Commit c60df2e

Browse files
jakob-kellerjakob-keller
committed
Support warm-up of loader caches
1 parent 41e938e commit c60df2e

6 files changed

Lines changed: 213 additions & 1 deletion

File tree

‎CHANGES.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
Changes
22
-------
33

4+
3.2.0 (2026-01-03)
5+
^^^^^^^^^^^^^^^^^^^
6+
* support `warm_up_loader_caches` in `AioConfig`
7+
48
3.1.0 (2026-01-02)
59
^^^^^^^^^^^^^^^^^^^
610
* support passing `socket_factory` as part of `connector_args` in `AioConfig`

‎aiobotocore/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
__version__ = '3.1.0'
1+
__version__ = '3.2.0'

‎aiobotocore/config.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ def __init__(
4444
self,
4545
connector_args: Optional[_ConnectorArgs] = None,
4646
http_session_cls: type[_HttpSessionType] = DEFAULT_HTTP_SESSION_CLS,
47+
warm_up_loader_caches: bool = False,
4748
**kwargs,
4849
):
4950
super().__init__(**kwargs)
@@ -52,6 +53,7 @@ def __init__(
5253
copy.copy(connector_args) if connector_args else {}
5354
)
5455
self.http_session_cls: type[_HttpSessionType] = http_session_cls
56+
self.warm_up_loader_caches: bool = warm_up_loader_caches
5557
self._validate_connector_args(
5658
self.connector_args, self.http_session_cls
5759
)

‎aiobotocore/session.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
import asyncio
2+
from typing import Optional
3+
14
from botocore import UNSIGNED, translate
25
from botocore import __version__ as botocore_version
36
from botocore.context import get_context
@@ -126,6 +129,36 @@ async def get_service_data(self, service_name, api_version=None):
126129
)
127130
return service_data
128131

132+
def warm_up_loader_caches(
133+
self,
134+
service_name: str,
135+
api_version: Optional[str] = None,
136+
):
137+
loader = self.get_component('data_loader')
138+
139+
# from session.py
140+
loader.load_data_with_path('endpoints')
141+
loader.load_data('sdk-default-configuration')
142+
loader.load_service_model(service_name, 'waiters-2', api_version)
143+
loader.load_service_model(service_name, 'paginators-1', api_version)
144+
loader.load_service_model(
145+
service_name, type_name='service-2', api_version=api_version
146+
)
147+
loader.list_available_services(type_name='service-2')
148+
149+
# from client.py
150+
loader.load_data('partitions')
151+
loader.load_service_model(
152+
service_name, 'service-2', api_version=api_version
153+
)
154+
loader.load_service_model(
155+
service_name, 'endpoint-rule-set-1', api_version=api_version
156+
)
157+
loader.load_data('_retry')
158+
159+
# from docs/service.py
160+
loader.load_service_model(service_name, 'examples-1', api_version)
161+
129162
def create_client(self, *args, **kwargs):
130163
return ClientCreatorContext(self._create_client(*args, **kwargs))
131164

@@ -167,6 +200,12 @@ async def _create_client(
167200
)
168201

169202
loader = self.get_component('data_loader')
203+
204+
if getattr(config, 'warm_up_loader_caches', False):
205+
await asyncio.to_thread(
206+
self.warm_up_loader_caches, service_name, api_version
207+
)
208+
170209
event_emitter = self.get_component('event_emitter')
171210
response_parser_factory = self.get_component('response_parser_factory')
172211
if config is not None and config.signature_version is UNSIGNED:

‎tests/test_config.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,3 +185,20 @@ async def send(self, request):
185185
):
186186
with pytest.raises(SuccessExc):
187187
await s3_client.get_object(Bucket='foo', Key='bar')
188+
189+
190+
@pytest.mark.parametrize(
191+
"warm_up_loader_caches, expected",
192+
[
193+
(None, False),
194+
(False, False),
195+
(True, True),
196+
],
197+
)
198+
def test_config_warm_up_loader_caches(warm_up_loader_caches, expected):
199+
if warm_up_loader_caches is None:
200+
config = AioConfig()
201+
else:
202+
config = AioConfig(warm_up_loader_caches=warm_up_loader_caches)
203+
204+
assert config.warm_up_loader_caches is expected

‎tests/test_session.py‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import logging
2+
from unittest import mock
23

34
import pytest
45
from _pytest.logging import LogCaptureFixture
@@ -54,3 +55,152 @@ async def test_set_user_agent_for_session(session: AioSession):
5455
assert session.user_agent_name == "aiobotocore"
5556
assert session.user_agent_version == __version__
5657
assert session.user_agent_extra.startswith("botocore/")
58+
59+
60+
@pytest.mark.parametrize(
61+
"service_name, api_version",
62+
[
63+
("s3", None),
64+
("s3", "2006-03-01"),
65+
("ec2", "2016-11-16"),
66+
],
67+
)
68+
def test_warm_up_loader_caches(session: AioSession, service_name, api_version):
69+
loader = mock.Mock()
70+
with mock.patch.object(
71+
session, "get_component", return_value=loader
72+
) as get_component:
73+
session.warm_up_loader_caches(service_name, api_version)
74+
75+
get_component.assert_called_once_with("data_loader")
76+
loader.mock_calls == [
77+
mock.call.load_data_with_path("endpoints"),
78+
mock.call.load_data("sdk-default-configuration"),
79+
mock.call.load_service_model(service_name, "waiters-2", api_version),
80+
mock.call.load_service_model(
81+
service_name, "paginators-1", api_version
82+
),
83+
mock.call.load_service_model(
84+
service_name, type_name="service-2", api_version=api_version
85+
),
86+
mock.call.list_available_services(type_name="service-2"),
87+
mock.call.load_data("partitions"),
88+
mock.call.load_service_model(
89+
service_name, "service-2", api_version=api_version
90+
),
91+
mock.call.load_service_model(
92+
service_name, "endpoint-rule-set-1", api_version=api_version
93+
),
94+
mock.call.load_data("_retry"),
95+
mock.call.load_service_model(service_name, "examples-1", api_version),
96+
]
97+
98+
99+
@pytest.mark.parametrize(
100+
"warm_up_loader_caches",
101+
[False, True],
102+
)
103+
async def test_warm_up_loader_caches_config(
104+
session: AioSession,
105+
warm_up_loader_caches: bool,
106+
):
107+
config = AioConfig(warm_up_loader_caches=warm_up_loader_caches)
108+
109+
with mock.patch.object(
110+
session,
111+
"warm_up_loader_caches",
112+
wraps=session.warm_up_loader_caches,
113+
):
114+
async with session.create_client(
115+
"s3",
116+
config=config,
117+
aws_secret_access_key="xxx",
118+
aws_access_key_id="xxx",
119+
):
120+
pass
121+
122+
if warm_up_loader_caches:
123+
session.warm_up_loader_caches.assert_called_once_with("s3", None)
124+
else:
125+
session.warm_up_loader_caches.assert_not_called()
126+
127+
128+
@pytest.mark.parametrize(
129+
"warm_up_loader_caches",
130+
[False, True],
131+
)
132+
async def test_non_blocking_create_client(
133+
session: AioSession,
134+
warm_up_loader_caches: bool,
135+
):
136+
config = AioConfig(warm_up_loader_caches=warm_up_loader_caches)
137+
138+
loader = session.get_component("data_loader")
139+
with mock.patch.object(
140+
loader,
141+
"file_loader",
142+
wraps=loader.file_loader,
143+
) as file_loader:
144+
# perform implicit warm-up, while avoiding any other file I/O by stubbing relevant codepathes
145+
session._internal_components.lazy_register_component(
146+
'endpoint_resolver', lambda: None
147+
)
148+
with (
149+
mock.patch.object(
150+
session, "_resolve_defaults_mode", return_value="legacy"
151+
),
152+
mock.patch("aiobotocore.session.AioClientCreator", autospec=True),
153+
):
154+
async with session.create_client(
155+
"s3",
156+
config=config,
157+
aws_secret_access_key="xxx",
158+
aws_access_key_id="xxx",
159+
):
160+
pass
161+
162+
if warm_up_loader_caches:
163+
# warm-up triggered file I/O (non-blocking)
164+
file_loader.exists.assert_called()
165+
file_loader.load_file.assert_called()
166+
else:
167+
# no file I/O
168+
file_loader.exists.assert_not_called()
169+
file_loader.load_file.assert_not_called()
170+
171+
session._register_endpoint_resolver()
172+
173+
file_loader.reset_mock()
174+
175+
# regular client creation #1
176+
async with session.create_client(
177+
"s3",
178+
config=config,
179+
aws_secret_access_key="xxx",
180+
aws_access_key_id="xxx",
181+
):
182+
pass
183+
184+
if warm_up_loader_caches:
185+
# no file I/O
186+
file_loader.exists.assert_not_called()
187+
file_loader.load_file.assert_not_called()
188+
else:
189+
# file I/O (blocking)
190+
file_loader.exists.assert_called()
191+
file_loader.load_file.assert_called()
192+
193+
file_loader.reset_mock()
194+
195+
# regular client creation #2
196+
async with session.create_client(
197+
"s3",
198+
config=config,
199+
aws_secret_access_key="xxx",
200+
aws_access_key_id="xxx",
201+
):
202+
pass
203+
204+
# no file I/O
205+
file_loader.exists.assert_not_called()
206+
file_loader.load_file.assert_not_called()

0 commit comments

Comments
 (0)