Support warm-up of loader caches

Author: jakob-keller

CHANGES.rst

@@ -1,6 +1,10 @@
 Changes
 -------
 
+3.2.0 (2026-02-06)
+^^^^^^^^^^^^^^^^^^^
+* support `warm_up_loader_caches` in `AioConfig`
+
 3.1.2 (2026-02-05)
 ^^^^^^^^^^^^^^^^^^^
 * relax botocore dependency specification to support ``"botocore >= 1.41.0, < 1.42.43"``

aiobotocore/__init__.py

@@ -1 +1 @@
-__version__ = '3.1.2'
+__version__ = '3.2.0'

aiobotocore/config.py

@@ -44,6 +44,7 @@ def __init__(
         self,
         connector_args: Optional[_ConnectorArgs] = None,
         http_session_cls: type[_HttpSessionType] = DEFAULT_HTTP_SESSION_CLS,
+        warm_up_loader_caches: bool = False,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -52,6 +53,7 @@ def __init__(
             copy.copy(connector_args) if connector_args else {}
         )
         self.http_session_cls: type[_HttpSessionType] = http_session_cls
+        self.warm_up_loader_caches: bool = warm_up_loader_caches
         self._validate_connector_args(
             self.connector_args, self.http_session_cls
         )

aiobotocore/session.py

@@ -1,3 +1,6 @@
+import asyncio
+from typing import Optional
+
 from botocore import UNSIGNED, translate
 from botocore import __version__ as botocore_version
 from botocore.context import get_context
@@ -126,6 +129,36 @@ async def get_service_data(self, service_name, api_version=None):
         )
         return service_data
 
+    def warm_up_loader_caches(
+        self,
+        service_name: str,
+        api_version: Optional[str] = None,
+    ):
+        loader = self.get_component('data_loader')
+
+        # from session.py
+        loader.load_data_with_path('endpoints')
+        loader.load_data('sdk-default-configuration')
+        loader.load_service_model(service_name, 'waiters-2', api_version)
+        loader.load_service_model(service_name, 'paginators-1', api_version)
+        loader.load_service_model(
+            service_name, type_name='service-2', api_version=api_version
+        )
+        loader.list_available_services(type_name='service-2')
+
+        # from client.py
+        loader.load_data('partitions')
+        loader.load_service_model(
+            service_name, 'service-2', api_version=api_version
+        )
+        loader.load_service_model(
+            service_name, 'endpoint-rule-set-1', api_version=api_version
+        )
+        loader.load_data('_retry')
+
+        # from docs/service.py
+        loader.load_service_model(service_name, 'examples-1', api_version)
+
     def create_client(self, *args, **kwargs):
         return ClientCreatorContext(self._create_client(*args, **kwargs))
 
@@ -167,6 +200,12 @@ async def _create_client(
             )
 
         loader = self.get_component('data_loader')
+
+        if getattr(config, 'warm_up_loader_caches', False):
+            await asyncio.to_thread(
+                self.warm_up_loader_caches, service_name, api_version
+            )
+
         event_emitter = self.get_component('event_emitter')
         response_parser_factory = self.get_component('response_parser_factory')
         if config is not None and config.signature_version is UNSIGNED:

pyproject.toml

@@ -74,6 +74,7 @@ dev = [
     "packaging >= 24.1, < 27", # Used in test_version.py
     "pip >= 24.3.1, < 27", # Used in test_version.py
     "pre-commit >= 3.5.0, < 5",
+    "pytest-mock >= 3.14.1, < 4",  # Used in test_session.py
     "pytest-rerunfailures >= 16.0.1, < 17", # Used in test_lambda.py
     "time-machine >= 2.15.0, < 4", # Used in test_signers.py
     "tomli; python_version<'3.11'", # Used in test_version.py

tests/test_config.py

@@ -185,3 +185,20 @@ async def send(self, request):
     ):
         with pytest.raises(SuccessExc):
             await s3_client.get_object(Bucket='foo', Key='bar')
+
+
+@pytest.mark.parametrize(
+    "warm_up_loader_caches, expected",
+    [
+        (None, False),
+        (False, False),
+        (True, True),
+    ],
+)
+def test_config_warm_up_loader_caches(warm_up_loader_caches, expected):
+    if warm_up_loader_caches is None:
+        config = AioConfig()
+    else:
+        config = AioConfig(warm_up_loader_caches=warm_up_loader_caches)
+
+    assert config.warm_up_loader_caches is expected

tests/test_session.py

@@ -54,3 +54,155 @@ async def test_set_user_agent_for_session(session: AioSession):
     assert session.user_agent_name == "aiobotocore"
     assert session.user_agent_version == __version__
     assert session.user_agent_extra.startswith("botocore/")
+
+
+@pytest.mark.parametrize(
+    "service_name, api_version",
+    [
+        ("s3", None),
+        ("s3", "2006-03-01"),
+        ("ec2", "2016-11-16"),
+    ],
+)
+def test_warm_up_loader_caches(
+    session: AioSession, service_name, api_version, mocker
+):
+    loader = mocker.Mock()
+    get_component = mocker.patch.object(
+        session, "get_component", return_value=loader
+    )
+
+    session.warm_up_loader_caches(service_name, api_version)
+
+    get_component.assert_called_once_with("data_loader")
+    assert loader.mock_calls == [
+        mocker.call.load_data_with_path("endpoints"),
+        mocker.call.load_data("sdk-default-configuration"),
+        mocker.call.load_service_model(service_name, "waiters-2", api_version),
+        mocker.call.load_service_model(
+            service_name, "paginators-1", api_version
+        ),
+        mocker.call.load_service_model(
+            service_name, type_name="service-2", api_version=api_version
+        ),
+        mocker.call.list_available_services(type_name="service-2"),
+        mocker.call.load_data("partitions"),
+        mocker.call.load_service_model(
+            service_name, "service-2", api_version=api_version
+        ),
+        mocker.call.load_service_model(
+            service_name, "endpoint-rule-set-1", api_version=api_version
+        ),
+        mocker.call.load_data("_retry"),
+        mocker.call.load_service_model(
+            service_name, "examples-1", api_version
+        ),
+    ]
+
+
+@pytest.mark.parametrize(
+    "warm_up_loader_caches",
+    [False, True],
+)
+async def test_warm_up_loader_caches_config(
+    session: AioSession,
+    warm_up_loader_caches: bool,
+    mocker,
+):
+    config = AioConfig(warm_up_loader_caches=warm_up_loader_caches)
+    mocker.patch.object(
+        session, "warm_up_loader_caches", wraps=session.warm_up_loader_caches
+    )
+
+    async with session.create_client(
+        "s3",
+        config=config,
+        aws_secret_access_key="xxx",
+        aws_access_key_id="xxx",
+    ):
+        pass
+
+    if warm_up_loader_caches:
+        session.warm_up_loader_caches.assert_called_once_with("s3", None)
+    else:
+        session.warm_up_loader_caches.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "warm_up_loader_caches",
+    [False, True],
+)
+async def test_non_blocking_create_client(
+    session: AioSession,
+    warm_up_loader_caches: bool,
+    mocker,
+):
+    config = AioConfig(warm_up_loader_caches=warm_up_loader_caches)
+    loader = session.get_component("data_loader")
+    file_loader = mocker.patch.object(
+        loader, "file_loader", wraps=loader.file_loader
+    )
+    # perform implicit warm-up, while avoiding any other file I/O by stubbing relevant codepathes
+    session._internal_components.lazy_register_component(
+        'endpoint_resolver', lambda: None
+    )
+    mocker.patch.object(
+        session, "_resolve_defaults_mode", return_value="legacy"
+    )
+    client_creator_cls_mock = mocker.patch(
+        "aiobotocore.session.AioClientCreator", autospec=True
+    )
+
+    async with session.create_client(
+        "s3",
+        config=config,
+        aws_secret_access_key="xxx",
+        aws_access_key_id="xxx",
+    ):
+        pass
+
+    if warm_up_loader_caches:
+        # warm-up triggered file I/O (non-blocking)
+        file_loader.exists.assert_called()
+        file_loader.load_file.assert_called()
+    else:
+        # no file I/O
+        file_loader.exists.assert_not_called()
+        file_loader.load_file.assert_not_called()
+
+    mocker.stop(client_creator_cls_mock)
+    session._register_endpoint_resolver()
+    file_loader.reset_mock()
+
+    # regular client creation #1
+    async with session.create_client(
+        "s3",
+        config=config,
+        aws_secret_access_key="xxx",
+        aws_access_key_id="xxx",
+    ):
+        pass
+
+    if warm_up_loader_caches:
+        # no file I/O
+        file_loader.exists.assert_not_called()
+        file_loader.load_file.assert_not_called()
+    else:
+        # file I/O (blocking)
+        file_loader.exists.assert_called()
+        file_loader.load_file.assert_called()
+
+    file_loader.reset_mock()
+
+    # regular client creation #2
+    async with session.create_client(
+        "s3",
+        config=config,
+        aws_secret_access_key="xxx",
+        aws_access_key_id="xxx",
+    ):
+        pass
+
+    # no file I/O
+    file_loader.exists.assert_not_called()
+    file_loader.load_file.assert_not_called()