FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

fix: Add DynamoDB connection pre-warming, fix mypy TorchTensor type alias, and fix CI mcp pin mismatch by nithin42 · Pull Request #6710 · feast-dev/feast · GitHub

Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension .py  (3) .toml  (1) .txt  (9) All 3 file types selected
Only manifest files
Viewed files
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Unified
Split
Hide whitespace
Diff view
Unified
Split
Hide whitespace
2 changes: 1 addition & 1 deletion pyproject.toml
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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

Back | FazBrowse Home | New Git URL