Skip to content

Commit 8ce548c

Browse files
CyberRikshuchu
CyberRik
and
shuchu
authored
fix: Warm up the whole DynamoDB connection pool (#6985)
* fix: Warm up the whole DynamoDB connection pool warmup_connections made a single describe_limits call, so only one connection of the pool was opened. Issue max_pool_connections calls concurrently, so that each one needs its own connection. A failed call is logged and does not stop the others. Signed-off-by: Ritankar Mondal <[email protected]> * fix: Use describe_endpoints for DynamoDB connection warmup DescribeLimits is throttled when called more than once a minute, so warming the pool with max_pool_connections concurrent calls would be throttled, and with adaptive retries it would also slow the client down for the first reads. DescribeEndpoints is not tied to a table and has no such limit. Signed-off-by: Ritankar Mondal <[email protected]> --------- Signed-off-by: Ritankar Mondal <[email protected]> Co-authored-by: Shuchu Han <[email protected]>
1 parent 81770b9 commit 8ce548c

4 files changed

Lines changed: 56 additions & 21 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -344,7 +344,7 @@ Key knobs:
344344
- **`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.
345345
- **`max_read_workers`**: Controls parallelism for batch reads. With 10 workers, those 5 batches run concurrently (~10 ms) instead of sequentially (~50 ms).
346346
- **`consistent_reads: false`**: Eventually consistent reads are faster and cheaper. Use `true` only if you need read-after-write consistency.
347-
- **`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.
347+
- **`warmup_connections: true`**: Pre-warms the DynamoDB connection pool on server startup by making `max_pool_connections` concurrent lightweight calls (`describe_endpoints`). This avoids a cold-start latency penalty (~20ms) on the first feature requests, including concurrent ones.
348348
- **`max_pool_connections`**: Increase for high-throughput deployments to improve HTTP connection reuse to the DynamoDB endpoint.
349349
- **`keepalive_timeout`**: Longer keep-alive reduces TLS handshake overhead on reused connections.
350350
- **`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: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ online_store:
5050
| `batch_size` | int | `100` | Number of items per BatchGetItem/BatchWriteItem request (max 100) |
5151
| `max_read_workers` | int | `10` | Maximum parallel threads for batch read operations. Higher values improve throughput for large batch reads but increase resource usage |
5252
| `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`) |
53+
| `warmup_connections` | bool | `false` | Whether to pre-warm the async connection pool on startup with `max_pool_connections` concurrent lightweight calls (`describe_endpoints`) |
5454
| `tags` | dict | `null` | AWS resource tags added to each table |
5555
| `session_based_auth` | bool | `false` | Use AWS session-based client authentication |
5656

@@ -65,7 +65,7 @@ For high-throughput workloads with large entity counts, increase `max_read_worke
6565

6666
**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.
6767

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.
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` opens up to `max_pool_connections` connections during server startup, one per concurrent `describe_endpoints` call. When DynamoDB is reached through a VPC endpoint, these calls need the `dynamodb:DescribeEndpoints` permission; without it the connections are still opened, but a warning is logged.
6969

7070
## Permissions
7171

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

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,8 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel):
7777
"""Whether to read from Dynamodb by forcing consistent reads"""
7878

7979
warmup_connections: StrictBool = False
80-
"""Whether to warm up the connection pool with a lightweight call on initialization"""
80+
"""Whether to warm up the connection pool on initialization, with
81+
max_pool_connections concurrent lightweight calls"""
8182

8283
tags: Union[Dict[str, str], None] = None
8384
"""AWS resource tags added to each table"""
@@ -161,11 +162,24 @@ async def initialize(self, config: RepoConfig):
161162
)
162163

163164
if online_config.warmup_connections:
164-
try:
165-
await client.describe_limits()
166-
except Exception:
165+
# Concurrent calls, so that each one needs its own connection; a
166+
# sequential call would reuse the connection of the previous one.
167+
# describe_endpoints is not tied to a table and, unlike
168+
# describe_limits, has no documented once-a-minute call limit.
169+
results = await asyncio.gather(
170+
*(
171+
client.describe_endpoints()
172+
for _ in range(online_config.max_pool_connections)
173+
),
174+
return_exceptions=True,
175+
)
176+
errors = [r for r in results if isinstance(r, BaseException)]
177+
if errors:
167178
logger.warning(
168-
"Failed to warmup DynamoDB connection pool", exc_info=True
179+
"Failed to warmup %d of %d DynamoDB connections",
180+
len(errors),
181+
len(results),
182+
exc_info=errors[0],
169183
)
170184

171185
async def close(self):

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

Lines changed: 34 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import asyncio
12
from copy import deepcopy
23
from dataclasses import dataclass
34
from datetime import datetime
@@ -100,31 +101,47 @@ def test_dynamodb_online_store_config_custom_params():
100101

101102

102103
@pytest.mark.asyncio
103-
async def test_dynamodb_online_store_warmup_connections():
104+
async def test_dynamodb_online_store_warmup_connections(caplog):
104105
"""Test DynamoDBOnlineStore warmup connections in initialize method."""
105106
from unittest.mock import AsyncMock
106107

107108
online_store = DynamoDBOnlineStore()
108109

109110
# Mock _get_aiodynamodb_client to return a mock client
110111
mock_client = AsyncMock()
111-
mock_client.describe_limits = AsyncMock()
112+
mock_client.describe_endpoints = AsyncMock()
112113
online_store._get_aiodynamodb_client = AsyncMock(return_value=mock_client)
113114

114-
# Test case 1: warmup_connections=True
115+
# Test case 1: warmup_connections=True warms the whole pool. The calls
116+
# must be in flight together, or they would all reuse one connection.
117+
pool_size = 8
118+
in_flight = 0
119+
max_in_flight = 0
120+
121+
async def describe_endpoints():
122+
nonlocal in_flight, max_in_flight
123+
in_flight += 1
124+
max_in_flight = max(max_in_flight, in_flight)
125+
await asyncio.sleep(0)
126+
in_flight -= 1
127+
128+
mock_client.describe_endpoints.side_effect = describe_endpoints
115129
config_warmup = RepoConfig(
116130
registry=REGISTRY,
117131
project=PROJECT,
118132
provider=PROVIDER,
119-
online_store=DynamoDBOnlineStoreConfig(region=REGION, warmup_connections=True),
133+
online_store=DynamoDBOnlineStoreConfig(
134+
region=REGION, warmup_connections=True, max_pool_connections=pool_size
135+
),
120136
offline_store=DaskOfflineStoreConfig(),
121137
entity_key_serialization_version=3,
122138
)
123139
await online_store.initialize(config_warmup)
124-
mock_client.describe_limits.assert_called_once()
140+
assert mock_client.describe_endpoints.call_count == pool_size
141+
assert max_in_flight == pool_size
125142

126143
# Test case 2: warmup_connections=False
127-
mock_client.describe_limits.reset_mock()
144+
mock_client.describe_endpoints.reset_mock()
128145
config_no_warmup = RepoConfig(
129146
registry=REGISTRY,
130147
project=PROJECT,
@@ -134,13 +151,17 @@ async def test_dynamodb_online_store_warmup_connections():
134151
entity_key_serialization_version=3,
135152
)
136153
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()
154+
mock_client.describe_endpoints.assert_not_called()
155+
156+
# Test case 3: a failed warmup call is logged, and the other calls still run
157+
mock_client.describe_endpoints.reset_mock()
158+
mock_client.describe_endpoints.side_effect = [Exception("Connection failed")] + [
159+
None
160+
] * (pool_size - 1)
161+
with caplog.at_level("WARNING"):
162+
await online_store.initialize(config_warmup)
163+
assert mock_client.describe_endpoints.call_count == pool_size
164+
assert f"Failed to warmup 1 of {pool_size} DynamoDB connections" in caplog.text
144165

145166

146167
def test_dynamodb_online_store_config_dynamodb_client(dynamodb_online_store):

0 commit comments

Comments
 (0)