Skip to content

fix: Add DynamoDB connection pre-warming, fix mypy TorchTensor type alias, and fix CI mcp pin mismatch #6710

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
nithin42 wants to merge 1 commit into from
Closed

fix: Add DynamoDB connection pre-warming, fix mypy TorchTensor type alias, and fix CI mcp pin mismatch #6710

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,7 +179,7 @@ ci = [
"moto<5",
"mypy>=1.4.1,<1.11.3",
"urllib3>=2.6.3,<3",
"psutil==5.9.0",
"psutil>=5.9.0",
"pytest-cov",
"Sphinx>4.0.0,<7",
"sqlglot[rs]>=23.4",
Expand Down
11 changes: 10 additions & 1 deletion sdk/python/feast/infra/online_stores/dynamodb.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,9 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel):
consistent_reads: StrictBool = False
"""Whether to read from Dynamodb by forcing consistent reads"""

warmup_connections: StrictBool = False
"""Whether to warm up the connection pool with a lightweight call on initialization"""

tags: Union[Dict[str, str], None] = None
"""AWS resource tags added to each table"""

Expand Down Expand Up @@ -146,7 +149,7 @@ def __init__(self):
async def initialize(self, config: RepoConfig):
online_config = config.online_store

await self._get_aiodynamodb_client(
client = await self._get_aiodynamodb_client(
online_config.region,
online_config.max_pool_connections,
online_config.keepalive_timeout,
Expand All @@ -157,6 +160,12 @@ async def initialize(self, config: RepoConfig):
online_config.endpoint_url,
)

if online_config.warmup_connections:
try:
await client.describe_limits()
except Exception:
logger.warning("Failed to warmup DynamoDB connection pool", exc_info=True)

async def close(self):
await self._aiodynamodb_close()

Expand Down
7 changes: 3 additions & 4 deletions sdk/python/feast/online_response.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
# limitations under the License.

import uuid as uuid_module
from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypeAlias, Union
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union

import pandas as pd
import pyarrow as pa
Expand All @@ -26,10 +26,9 @@

if TYPE_CHECKING:
import torch

TorchTensor: TypeAlias = torch.Tensor
from torch import Tensor as TorchTensor
else:
TorchTensor: TypeAlias = Any
TorchTensor = Any

TIMESTAMP_POSTFIX: str = "__ts"

Expand Down
10 changes: 3 additions & 7 deletions sdk/python/requirements/py3.10-ci-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2796,14 +2796,10 @@ matplotlib-inline==0.2.2 \
# via
# ipykernel
# ipython
mcp==2.0.0 \
--hash=sha256:0f440e735c13ece8bb19bc62cf0b86f4313448432fbb77d35e14034f4e050728 \
--hash=sha256:1cb4c75d2d2c7b8c1d756355e5d82a39f2822cc7f13e22a2051d7ca3592349d6
mcp==1.29.0 \
--hash=sha256:52d01f334de1868cc3bb2d6604931126a67631f99a6c5d3b82ba47290315ec36 \
--hash=sha256:f5a075bb611f23d6f4d080c6a1699fa62772eebc562ba9e66b306ddde1c755f7
# via fastapi-mcp
mcp-types==2.0.0 \
--hash=sha256:6b2de797ca2797f568b79529e1b25948e34de511bcc0bd82fef1039a6d1b8eb0 \
--hash=sha256:d7d939b9285c9961ae8866ba75ef85da34d12bafe276efbf4eb6a131786d8379
# via mcp
mdurl==0.1.2 \
--hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \
--hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba
Expand Down
10 changes: 3 additions & 7 deletions sdk/python/requirements/py3.10-minimal-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1367,14 +1367,10 @@ markupsafe==3.0.3 \
--hash=sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a \
--hash=sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50
# via jinja2
mcp==2.0.0 \
--hash=sha256:0f440e735c13ece8bb19bc62cf0b86f4313448432fbb77d35e14034f4e050728 \
--hash=sha256:1cb4c75d2d2c7b8c1d756355e5d82a39f2822cc7f13e22a2051d7ca3592349d6
mcp==1.29.0 \
--hash=sha256:52d01f334de1868cc3bb2d6604931126a67631f99a6c5d3b82ba47290315ec36 \
--hash=sha256:f5a075bb611f23d6f4d080c6a1699fa62772eebc562ba9e66b306ddde1c755f7
# via fastapi-mcp
mcp-types==2.0.0 \
--hash=sha256:6b2de797ca2797f568b79529e1b25948e34de511bcc0bd82fef1039a6d1b8eb0 \
--hash=sha256:d7d939b9285c9961ae8866ba75ef85da34d12bafe276efbf4eb6a131786d8379
# via mcp
mdurl==0.1.2 \
--hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \
--hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba
Expand Down
10 changes: 3 additions & 7 deletions sdk/python/requirements/py3.10-minimal-sdist-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1544,14 +1544,10 @@ markupsafe==3.0.3 \
--hash=sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a \
--hash=sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50
# via jinja2
mcp==2.0.0 \
--hash=sha256:0f440e735c13ece8bb19bc62cf0b86f4313448432fbb77d35e14034f4e050728 \
--hash=sha256:1cb4c75d2d2c7b8c1d756355e5d82a39f2822cc7f13e22a2051d7ca3592349d6
mcp==1.29.0 \
--hash=sha256:52d01f334de1868cc3bb2d6604931126a67631f99a6c5d3b82ba47290315ec36 \
--hash=sha256:f5a075bb611f23d6f4d080c6a1699fa62772eebc562ba9e66b306ddde1c755f7
# via fastapi-mcp
mcp-types==2.0.0 \
--hash=sha256:6b2de797ca2797f568b79529e1b25948e34de511bcc0bd82fef1039a6d1b8eb0 \
--hash=sha256:d7d939b9285c9961ae8866ba75ef85da34d12bafe276efbf4eb6a131786d8379
# via mcp
mdurl==0.1.2 \
--hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \
--hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba
Expand Down
10 changes: 3 additions & 7 deletions sdk/python/requirements/py3.11-ci-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2896,14 +2896,10 @@ matplotlib-inline==0.2.2 \
# via
# ipykernel
# ipython
mcp==2.0.0 \
--hash=sha256:0f440e735c13ece8bb19bc62cf0b86f4313448432fbb77d35e14034f4e050728 \
--hash=sha256:1cb4c75d2d2c7b8c1d756355e5d82a39f2822cc7f13e22a2051d7ca3592349d6
mcp==1.29.0 \
--hash=sha256:52d01f334de1868cc3bb2d6604931126a67631f99a6c5d3b82ba47290315ec36 \
--hash=sha256:f5a075bb611f23d6f4d080c6a1699fa62772eebc562ba9e66b306ddde1c755f7
# via fastapi-mcp
mcp-types==2.0.0 \
--hash=sha256:6b2de797ca2797f568b79529e1b25948e34de511bcc0bd82fef1039a6d1b8eb0 \
--hash=sha256:d7d939b9285c9961ae8866ba75ef85da34d12bafe276efbf4eb6a131786d8379
# via mcp
mdurl==0.1.2 \
--hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \
--hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba
Expand Down
10 changes: 3 additions & 7 deletions sdk/python/requirements/py3.11-minimal-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1361,14 +1361,10 @@ markupsafe==3.0.3 \
--hash=sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a \
--hash=sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50
# via jinja2
mcp==2.0.0 \
--hash=sha256:0f440e735c13ece8bb19bc62cf0b86f4313448432fbb77d35e14034f4e050728 \
--hash=sha256:1cb4c75d2d2c7b8c1d756355e5d82a39f2822cc7f13e22a2051d7ca3592349d6
mcp==1.29.0 \
--hash=sha256:52d01f334de1868cc3bb2d6604931126a67631f99a6c5d3b82ba47290315ec36 \
--hash=sha256:f5a075bb611f23d6f4d080c6a1699fa62772eebc562ba9e66b306ddde1c755f7
# via fastapi-mcp
mcp-types==2.0.0 \
--hash=sha256:6b2de797ca2797f568b79529e1b25948e34de511bcc0bd82fef1039a6d1b8eb0 \
--hash=sha256:d7d939b9285c9961ae8866ba75ef85da34d12bafe276efbf4eb6a131786d8379
# via mcp
mdurl==0.1.2 \
--hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \
--hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba
Expand Down
10 changes: 3 additions & 7 deletions sdk/python/requirements/py3.11-minimal-sdist-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1536,14 +1536,10 @@ markupsafe==3.0.3 \
--hash=sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a \
--hash=sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50
# via jinja2
mcp==2.0.0 \
--hash=sha256:0f440e735c13ece8bb19bc62cf0b86f4313448432fbb77d35e14034f4e050728 \
--hash=sha256:1cb4c75d2d2c7b8c1d756355e5d82a39f2822cc7f13e22a2051d7ca3592349d6
mcp==1.29.0 \
--hash=sha256:52d01f334de1868cc3bb2d6604931126a67631f99a6c5d3b82ba47290315ec36 \
--hash=sha256:f5a075bb611f23d6f4d080c6a1699fa62772eebc562ba9e66b306ddde1c755f7
# via fastapi-mcp
mcp-types==2.0.0 \
--hash=sha256:6b2de797ca2797f568b79529e1b25948e34de511bcc0bd82fef1039a6d1b8eb0 \
--hash=sha256:d7d939b9285c9961ae8866ba75ef85da34d12bafe276efbf4eb6a131786d8379
# via mcp
mdurl==0.1.2 \
--hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \
--hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba
Expand Down
10 changes: 3 additions & 7 deletions sdk/python/requirements/py3.12-ci-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2905,14 +2905,10 @@ matplotlib-inline==0.2.2 \
# via
# ipykernel
# ipython
mcp==2.0.0 \
--hash=sha256:0f440e735c13ece8bb19bc62cf0b86f4313448432fbb77d35e14034f4e050728 \
--hash=sha256:1cb4c75d2d2c7b8c1d756355e5d82a39f2822cc7f13e22a2051d7ca3592349d6
mcp==1.29.0 \
--hash=sha256:52d01f334de1868cc3bb2d6604931126a67631f99a6c5d3b82ba47290315ec36 \
--hash=sha256:f5a075bb611f23d6f4d080c6a1699fa62772eebc562ba9e66b306ddde1c755f7
# via fastapi-mcp
mcp-types==2.0.0 \
--hash=sha256:6b2de797ca2797f568b79529e1b25948e34de511bcc0bd82fef1039a6d1b8eb0 \
--hash=sha256:d7d939b9285c9961ae8866ba75ef85da34d12bafe276efbf4eb6a131786d8379
# via mcp
mdurl==0.1.2 \
--hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \
--hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba
Expand Down
10 changes: 3 additions & 7 deletions sdk/python/requirements/py3.12-minimal-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1353,14 +1353,10 @@ markupsafe==3.0.3 \
--hash=sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a \
--hash=sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50
# via jinja2
mcp==2.0.0 \
--hash=sha256:0f440e735c13ece8bb19bc62cf0b86f4313448432fbb77d35e14034f4e050728 \
--hash=sha256:1cb4c75d2d2c7b8c1d756355e5d82a39f2822cc7f13e22a2051d7ca3592349d6
mcp==1.29.0 \
--hash=sha256:52d01f334de1868cc3bb2d6604931126a67631f99a6c5d3b82ba47290315ec36 \
--hash=sha256:f5a075bb611f23d6f4d080c6a1699fa62772eebc562ba9e66b306ddde1c755f7
# via fastapi-mcp
mcp-types==2.0.0 \
--hash=sha256:6b2de797ca2797f568b79529e1b25948e34de511bcc0bd82fef1039a6d1b8eb0 \
--hash=sha256:d7d939b9285c9961ae8866ba75ef85da34d12bafe276efbf4eb6a131786d8379
# via mcp
mdurl==0.1.2 \
--hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \
--hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba
Expand Down
10 changes: 3 additions & 7 deletions sdk/python/requirements/py3.12-minimal-sdist-requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1528,14 +1528,10 @@ markupsafe==3.0.3 \
--hash=sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a \
--hash=sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50
# via jinja2
mcp==2.0.0 \
--hash=sha256:0f440e735c13ece8bb19bc62cf0b86f4313448432fbb77d35e14034f4e050728 \
--hash=sha256:1cb4c75d2d2c7b8c1d756355e5d82a39f2822cc7f13e22a2051d7ca3592349d6
mcp==1.29.0 \
--hash=sha256:52d01f334de1868cc3bb2d6604931126a67631f99a6c5d3b82ba47290315ec36 \
--hash=sha256:f5a075bb611f23d6f4d080c6a1699fa62772eebc562ba9e66b306ddde1c755f7
# via fastapi-mcp
mcp-types==2.0.0 \
--hash=sha256:6b2de797ca2797f568b79529e1b25948e34de511bcc0bd82fef1039a6d1b8eb0 \
--hash=sha256:d7d939b9285c9961ae8866ba75ef85da34d12bafe276efbf4eb6a131786d8379
# via mcp
mdurl==0.1.2 \
--hash=sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8 \
--hash=sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba
Expand Down
41 changes: 41 additions & 0 deletions sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ def test_dynamodb_online_store_config_default():
assert dynamodb_store_config.read_timeout == 10
assert dynamodb_store_config.total_max_retry_attempts == 3
assert dynamodb_store_config.retry_mode == "adaptive"
assert dynamodb_store_config.warmup_connections is False


def test_dynamodb_online_store_config_custom_params():
Expand All @@ -88,12 +89,52 @@ def test_dynamodb_online_store_config_custom_params():
batch_size=batch_size,
endpoint_url=endpoint_url,
table_name_template=table_name_template,
warmup_connections=True,
)
assert dynamodb_store_config.type == "dynamodb"
assert dynamodb_store_config.batch_size == batch_size
assert dynamodb_store_config.endpoint_url == endpoint_url
assert dynamodb_store_config.region == aws_region
assert dynamodb_store_config.table_name_template == table_name_template
assert dynamodb_store_config.warmup_connections is True


@pytest.mark.asyncio
async def test_dynamodb_online_store_warmup_connections():
"""Test DynamoDBOnlineStore warmup connections in initialize method."""
from unittest.mock import AsyncMock

online_store = DynamoDBOnlineStore()

# Mock _get_aiodynamodb_client to return a mock client
mock_client = AsyncMock()
mock_client.describe_limits = AsyncMock()
online_store._get_aiodynamodb_client = AsyncMock(return_value=mock_client)

# Test case 1: warmup_connections=True
config_warmup = RepoConfig(
registry=REGISTRY,
project=PROJECT,
provider=PROVIDER,
online_store=DynamoDBOnlineStoreConfig(region=REGION, warmup_connections=True),
offline_store=DaskOfflineStoreConfig(),
entity_key_serialization_version=3,
)
await online_store.initialize(config_warmup)
mock_client.describe_limits.assert_called_once()

# Test case 2: warmup_connections=False
mock_client.describe_limits.reset_mock()
config_no_warmup = RepoConfig(
registry=REGISTRY,
project=PROJECT,
provider=PROVIDER,
online_store=DynamoDBOnlineStoreConfig(region=REGION, warmup_connections=False),
offline_store=DaskOfflineStoreConfig(),
entity_key_serialization_version=3,
)
await online_store.initialize(config_no_warmup)
mock_client.describe_limits.assert_not_called()


def test_dynamodb_online_store_config_dynamodb_client(dynamodb_online_store):
Expand Down