Skip to content

Commit 89240fa

Browse files
nithin42ntkathole
nithin42
authored and
ntkathole
committed
fix: Add connection pre-warming for DynamoDB async client
Add warmup_connections parameter to DynamoDBOnlineStoreConfig and warm up connections via describe_limits inside initialize() method to solve cold-start latency. Also document warmup_connections in the performance tuning guide and reference docs. Fixes #6060 Signed-off-by: Nithin <[email protected]>
1 parent 98e5bca commit 89240fa

4 files changed

Lines changed: 65 additions & 1 deletion

File tree

‎docs/how-to-guides/online-server-performance-tuning.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,7 @@ online_store:
324324
batch_size: 100
325325
max_read_workers: 10
326326
consistent_reads: false
327+
warmup_connections: true
327328
max_pool_connections: 100
328329
keepalive_timeout: 30.0
329330
connect_timeout: 3
@@ -337,6 +338,7 @@ Key knobs:
337338
- **`batch_size`**: DynamoDB's `BatchGetItem` accepts up to 100 items per request. For 500 entities, this means 5 batches. Keep at 100 unless hitting the 16 MB response limit.
338339
- **`max_read_workers`**: Controls parallelism for batch reads. With 10 workers, those 5 batches run concurrently (~10 ms) instead of sequentially (~50 ms).
339340
- **`consistent_reads: false`**: Eventually consistent reads are faster and cheaper. Use `true` only if you need read-after-write consistency.
341+
- **`warmup_connections: true`**: Pre-warms the DynamoDB connection pool on server startup by making a lightweight call (`describe_limits`). This avoids a cold-start latency penalty (~20ms) on the very first feature request.
340342
- **`max_pool_connections`**: Increase for high-throughput deployments to improve HTTP connection reuse to the DynamoDB endpoint.
341343
- **`keepalive_timeout`**: Longer keep-alive reduces TLS handshake overhead on reused connections.
342344
- **`connect_timeout` / `read_timeout`**: Lower values fail fast, improving p99. Set aggressively if your retry strategy covers transient failures.

‎docs/reference/online-stores/dynamodb.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ online_store:
3737
batch_size: 100
3838
max_read_workers: 10
3939
consistent_reads: false
40+
warmup_connections: true
4041
```
4142
{% endcode %}
4243
@@ -49,6 +50,7 @@ online_store:
4950
| `batch_size` | int | `100` | Number of items per BatchGetItem/BatchWriteItem request (max 100) |
5051
| `max_read_workers` | int | `10` | Maximum parallel threads for batch read operations. Higher values improve throughput for large batch reads but increase resource usage |
5152
| `consistent_reads` | bool | `false` | Whether to use strongly consistent reads (higher latency, guaranteed latest data) |
53+
| `warmup_connections` | bool | `false` | Whether to pre-warm the async connection pool on startup with a lightweight call (`describe_limits`) |
5254
| `tags` | dict | `null` | AWS resource tags added to each table |
5355
| `session_based_auth` | bool | `false` | Use AWS session-based client authentication |
5456

@@ -63,6 +65,8 @@ For high-throughput workloads with large entity counts, increase `max_read_worke
6365

6466
**Batch Size**: Increase `batch_size` up to 100 to reduce the number of API calls. However, larger batches may hit DynamoDB's 16MB response limit for tables with large feature values.
6567

68+
**Connection Warmup**: The DynamoDB async client does not establish actual TCP/TLS connections to the AWS endpoint on initialization. The very first feature retrieval request is penalized with a cold-start overhead (~20ms). Setting `warmup_connections: true` establishes the TCP connection pool during server startup.
69+
6670
## Permissions
6771

6872
Feast requires the following permissions in order to execute commands for DynamoDB online store:

‎sdk/python/feast/infra/online_stores/dynamodb.py‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,9 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel):
7676
consistent_reads: StrictBool = False
7777
"""Whether to read from Dynamodb by forcing consistent reads"""
7878

79+
warmup_connections: StrictBool = False
80+
"""Whether to warm up the connection pool with a lightweight call on initialization"""
81+
7982
tags: Union[Dict[str, str], None] = None
8083
"""AWS resource tags added to each table"""
8184

@@ -146,7 +149,7 @@ def __init__(self):
146149
async def initialize(self, config: RepoConfig):
147150
online_config = config.online_store
148151

149-
await self._get_aiodynamodb_client(
152+
client = await self._get_aiodynamodb_client(
150153
online_config.region,
151154
online_config.max_pool_connections,
152155
online_config.keepalive_timeout,
@@ -157,6 +160,14 @@ async def initialize(self, config: RepoConfig):
157160
online_config.endpoint_url,
158161
)
159162

163+
if online_config.warmup_connections:
164+
try:
165+
await client.describe_limits()
166+
except Exception:
167+
logger.warning(
168+
"Failed to warmup DynamoDB connection pool", exc_info=True
169+
)
170+
160171
async def close(self):
161172
await self._aiodynamodb_close()
162173

‎sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ def test_dynamodb_online_store_config_default():
7575
assert dynamodb_store_config.read_timeout == 10
7676
assert dynamodb_store_config.total_max_retry_attempts == 3
7777
assert dynamodb_store_config.retry_mode == "adaptive"
78+
assert dynamodb_store_config.warmup_connections is False
7879

7980

8081
def test_dynamodb_online_store_config_custom_params():
@@ -88,12 +89,58 @@ def test_dynamodb_online_store_config_custom_params():
8889
batch_size=batch_size,
8990
endpoint_url=endpoint_url,
9091
table_name_template=table_name_template,
92+
warmup_connections=True,
9193
)
9294
assert dynamodb_store_config.type == "dynamodb"
9395
assert dynamodb_store_config.batch_size == batch_size
9496
assert dynamodb_store_config.endpoint_url == endpoint_url
9597
assert dynamodb_store_config.region == aws_region
9698
assert dynamodb_store_config.table_name_template == table_name_template
99+
assert dynamodb_store_config.warmup_connections is True
100+
101+
102+
@pytest.mark.asyncio
103+
async def test_dynamodb_online_store_warmup_connections():
104+
"""Test DynamoDBOnlineStore warmup connections in initialize method."""
105+
from unittest.mock import AsyncMock
106+
107+
online_store = DynamoDBOnlineStore()
108+
109+
# Mock _get_aiodynamodb_client to return a mock client
110+
mock_client = AsyncMock()
111+
mock_client.describe_limits = AsyncMock()
112+
online_store._get_aiodynamodb_client = AsyncMock(return_value=mock_client)
113+
114+
# Test case 1: warmup_connections=True
115+
config_warmup = RepoConfig(
116+
registry=REGISTRY,
117+
project=PROJECT,
118+
provider=PROVIDER,
119+
online_store=DynamoDBOnlineStoreConfig(region=REGION, warmup_connections=True),
120+
offline_store=DaskOfflineStoreConfig(),
121+
entity_key_serialization_version=3,
122+
)
123+
await online_store.initialize(config_warmup)
124+
mock_client.describe_limits.assert_called_once()
125+
126+
# Test case 2: warmup_connections=False
127+
mock_client.describe_limits.reset_mock()
128+
config_no_warmup = RepoConfig(
129+
registry=REGISTRY,
130+
project=PROJECT,
131+
provider=PROVIDER,
132+
online_store=DynamoDBOnlineStoreConfig(region=REGION, warmup_connections=False),
133+
offline_store=DaskOfflineStoreConfig(),
134+
entity_key_serialization_version=3,
135+
)
136+
await online_store.initialize(config_no_warmup)
137+
mock_client.describe_limits.assert_not_called()
138+
139+
# Test case 3: warmup_connections=True and describe_limits raises an exception (should catch and log warning)
140+
mock_client.describe_limits.reset_mock()
141+
mock_client.describe_limits.side_effect = Exception("Connection failed")
142+
await online_store.initialize(config_warmup)
143+
mock_client.describe_limits.assert_called_once()
97144

98145

99146
def test_dynamodb_online_store_config_dynamodb_client(dynamodb_online_store):

0 commit comments

Comments
 (0)