Support warm-up of loader caches

Author: jakob-keller

CHANGES.rst

@@ -1,6 +1,10 @@
 Changes
 -------
 
+3.2.0 (2026-01-03)
+^^^^^^^^^^^^^^^^^^^
+* support `warm_up_loader_caches` in `AioConfig`
+
 3.1.0 (2026-01-02)
 ^^^^^^^^^^^^^^^^^^^
 * support passing `socket_factory` as part of `connector_args` in `AioConfig`

aiobotocore/__init__.py

@@ -1 +1 @@
-__version__ = '3.1.0'
+__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:

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

@@ -1,4 +1,5 @@
 import logging
+from unittest import mock
 
 import pytest
 from _pytest.logging import LogCaptureFixture
@@ -54,3 +55,37 @@ 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):
+    loader = mock.Mock()
+    with mock.patch.object(
+        session, "get_component", return_value=loader
+    ) as get_component:
+        session.warm_up_loader_caches(service_name, api_version)
+
+    get_component.assert_called_once_with("data_loader")
+    loader.mock_calls == [
+        mock.call.load_data_with_path("endpoints"),
+        mock.call.load_data("sdk-default-configuration"),
+        mock.call.list_available_services(type_name="service-2"),
+        mock.call.load_data("partitions"),
+        mock.call.load_service_model(
+            service_name, "service-2", api_version=api_version
+        ),
+        mock.call.load_service_model(
+            service_name, "endpoint-rule-set-1", api_version=api_version
+        ),
+        mock.call.load_data("_retry"),
+        mock.call.load_service_model(
+            service_name, "examples-1", api_version=api_version
+        ),
+    ]