From 803db4c26c2ca479a3cb45d19470e14898c5e884 Mon Sep 17 00:00:00 2001 From: Nithin Date: Fri, 7 Aug 2026 14:28:22 -0400 Subject: [PATCH] fix: Add DynamoDB connection pre-warming, fix mypy TorchTensor type alias, and fix CI mcp pin mismatch - Add warmup_connections option to DynamoDBOnlineStoreConfig and warm up connections via describe_limits to solve cold-start latency. - Fix TypeAlias annotation in online_response.py fallback branch to stop mypy from rejecting TorchTensor as a type. - Loosen psutil dependency to >=5.9.0 to allow installing binary wheels on Windows Python 3.11. - Downgrade mcp dependency in Python requirements lock files to 1.29.0 and remove mcp-types to satisfy the bound 'mcp>=1.0,<2' declared in pyproject.toml and fix the feature server crash. Fixes #6060, #5563, #6706 Signed-off-by: Nithin --- pyproject.toml | 2 +- .../feast/infra/online_stores/dynamodb.py | 11 ++++- sdk/python/feast/online_response.py | 7 ++-- .../requirements/py3.10-ci-requirements.txt | 10 ++--- .../py3.10-minimal-requirements.txt | 10 ++--- .../py3.10-minimal-sdist-requirements.txt | 10 ++--- .../requirements/py3.11-ci-requirements.txt | 10 ++--- .../py3.11-minimal-requirements.txt | 10 ++--- .../py3.11-minimal-sdist-requirements.txt | 10 ++--- .../requirements/py3.12-ci-requirements.txt | 10 ++--- .../py3.12-minimal-requirements.txt | 10 ++--- .../py3.12-minimal-sdist-requirements.txt | 10 ++--- .../test_dynamodb_online_store.py | 41 +++++++++++++++++++ 13 files changed, 82 insertions(+), 69 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e2441730106..69f2358d901 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py index b67c8a81f6b..17d01d93a74 100644 --- a/sdk/python/feast/infra/online_stores/dynamodb.py +++ b/sdk/python/feast/infra/online_stores/dynamodb.py @@ -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""" @@ -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, @@ -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() diff --git a/sdk/python/feast/online_response.py b/sdk/python/feast/online_response.py index 7b6b4806e4d..0e122ff794a 100644 --- a/sdk/python/feast/online_response.py +++ b/sdk/python/feast/online_response.py @@ -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 @@ -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" diff --git a/sdk/python/requirements/py3.10-ci-requirements.txt b/sdk/python/requirements/py3.10-ci-requirements.txt index d431f142320..ab6f04bfdf4 100644 --- a/sdk/python/requirements/py3.10-ci-requirements.txt +++ b/sdk/python/requirements/py3.10-ci-requirements.txt @@ -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 diff --git a/sdk/python/requirements/py3.10-minimal-requirements.txt b/sdk/python/requirements/py3.10-minimal-requirements.txt index 47ad051f9e7..3412ace21cf 100644 --- a/sdk/python/requirements/py3.10-minimal-requirements.txt +++ b/sdk/python/requirements/py3.10-minimal-requirements.txt @@ -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 diff --git a/sdk/python/requirements/py3.10-minimal-sdist-requirements.txt b/sdk/python/requirements/py3.10-minimal-sdist-requirements.txt index b4a45c99dc2..fa22f822099 100644 --- a/sdk/python/requirements/py3.10-minimal-sdist-requirements.txt +++ b/sdk/python/requirements/py3.10-minimal-sdist-requirements.txt @@ -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 diff --git a/sdk/python/requirements/py3.11-ci-requirements.txt b/sdk/python/requirements/py3.11-ci-requirements.txt index b1243a03d23..6a5a72fa0e7 100644 --- a/sdk/python/requirements/py3.11-ci-requirements.txt +++ b/sdk/python/requirements/py3.11-ci-requirements.txt @@ -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 diff --git a/sdk/python/requirements/py3.11-minimal-requirements.txt b/sdk/python/requirements/py3.11-minimal-requirements.txt index dc313f87fdc..2ba95f90fc2 100644 --- a/sdk/python/requirements/py3.11-minimal-requirements.txt +++ b/sdk/python/requirements/py3.11-minimal-requirements.txt @@ -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 diff --git a/sdk/python/requirements/py3.11-minimal-sdist-requirements.txt b/sdk/python/requirements/py3.11-minimal-sdist-requirements.txt index 5648439b025..27ae941f0c2 100644 --- a/sdk/python/requirements/py3.11-minimal-sdist-requirements.txt +++ b/sdk/python/requirements/py3.11-minimal-sdist-requirements.txt @@ -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 diff --git a/sdk/python/requirements/py3.12-ci-requirements.txt b/sdk/python/requirements/py3.12-ci-requirements.txt index bef84732a3a..5852510afc6 100644 --- a/sdk/python/requirements/py3.12-ci-requirements.txt +++ b/sdk/python/requirements/py3.12-ci-requirements.txt @@ -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 diff --git a/sdk/python/requirements/py3.12-minimal-requirements.txt b/sdk/python/requirements/py3.12-minimal-requirements.txt index 4affd220de9..b7214595a1b 100644 --- a/sdk/python/requirements/py3.12-minimal-requirements.txt +++ b/sdk/python/requirements/py3.12-minimal-requirements.txt @@ -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 diff --git a/sdk/python/requirements/py3.12-minimal-sdist-requirements.txt b/sdk/python/requirements/py3.12-minimal-sdist-requirements.txt index a11b0b0c0b3..ea9f2596de6 100644 --- a/sdk/python/requirements/py3.12-minimal-sdist-requirements.txt +++ b/sdk/python/requirements/py3.12-minimal-sdist-requirements.txt @@ -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 diff --git a/sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py b/sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py index 3b84b1a3547..8d78526ca2e 100644 --- a/sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py +++ b/sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py @@ -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(): @@ -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):