Visitar URL original
perf: Optimize feature serving latency with batched async Redis, cach… · feast-dev/feast@103809a · GitHub
Skip to content

Commit 103809a

Browse files
ntkatholentkathole
committed
perf: Optimize feature serving latency with batched async Redis, cached checks fix
Signed-off-by: ntkathole <nikhilkathole2683@gmail.com>
1 parent 4203eb7 commit 103809a

4 files changed

Lines changed: 230 additions & 92 deletions

File tree

‎sdk/python/feast/feature_server.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,11 @@
5454
from feast.feature_server_utils import convert_response_to_dict
5555
from feast.feature_view_utils import get_feature_view_from_feature_store
5656
from feast.permissions.action import WRITE, AuthzedAction
57-
from feast.permissions.security_manager import assert_permissions
57+
from feast.permissions.security_manager import (
58+
assert_permissions,
59+
get_security_manager,
60+
is_auth_necessary,
61+
)
5862
from feast.permissions.server.rest import inject_user_details
5963
from feast.permissions.server.utils import (
6064
ServerType,
@@ -184,7 +188,7 @@ async def _get_features(
184188
resource=feature_service, actions=[AuthzedAction.READ_ONLINE]
185189
)
186190
features = feature_service # type: ignore
187-
else:
191+
elif is_auth_necessary(get_security_manager()):
188192
all_feature_views, all_on_demand_feature_views = await run_in_threadpool(
189193
utils._get_feature_views_to_use,
190194
store.registry,
@@ -202,6 +206,8 @@ async def _get_features(
202206
resource=od_feature_view, actions=[AuthzedAction.READ_ONLINE]
203207
)
204208
features = request.features # type: ignore
209+
else:
210+
features = request.features # type: ignore
205211
return features
206212

207213

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

Lines changed: 33 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -253,55 +253,14 @@ def get_online_features(
253253
)
254254
return OnlineResponse(online_features_response, feature_types=feature_types)
255255

256+
_versioned_read_supported: Optional[bool] = None
257+
256258
def _check_versioned_read_support(self, grouped_refs):
257259
"""Raise an error if versioned reads are attempted on unsupported stores."""
258-
from feast.infra.online_stores.sqlite import SqliteOnlineStore
259-
260-
supported_types: list[type] = [SqliteOnlineStore]
261-
try:
262-
from feast.infra.online_stores.mysql_online_store.mysql import (
263-
MySQLOnlineStore,
264-
)
265-
266-
supported_types.append(MySQLOnlineStore)
267-
except ImportError:
268-
pass
269-
try:
270-
from feast.infra.online_stores.postgres_online_store.postgres import (
271-
PostgreSQLOnlineStore,
272-
)
273-
274-
supported_types.append(PostgreSQLOnlineStore)
275-
except ImportError:
276-
pass
277-
try:
278-
from feast.infra.online_stores.faiss_online_store import FaissOnlineStore
279-
280-
supported_types.append(FaissOnlineStore)
281-
except ImportError:
282-
pass
283-
try:
284-
from feast.infra.online_stores.redis import RedisOnlineStore
260+
if self._versioned_read_supported is None:
261+
self._versioned_read_supported = self._is_versioned_read_supported()
285262

286-
supported_types.append(RedisOnlineStore)
287-
except Exception:
288-
pass
289-
try:
290-
from feast.infra.online_stores.dynamodb import DynamoDBOnlineStore
291-
292-
supported_types.append(DynamoDBOnlineStore)
293-
except Exception:
294-
pass
295-
try:
296-
from feast.infra.online_stores.milvus_online_store.milvus import (
297-
MilvusOnlineStore,
298-
)
299-
300-
supported_types.append(MilvusOnlineStore)
301-
except ImportError:
302-
pass
303-
304-
if isinstance(self, tuple(supported_types)):
263+
if self._versioned_read_supported:
305264
return
306265
for table, _ in grouped_refs:
307266
version_tag = getattr(table.projection, "version_tag", None)
@@ -310,6 +269,34 @@ def _check_versioned_read_support(self, grouped_refs):
310269
self.__class__.__name__, version_tag
311270
)
312271

272+
def _is_versioned_read_supported(self) -> bool:
273+
"""Check if this store type supports versioned reads (resolved once, cached)."""
274+
from feast.infra.online_stores.sqlite import SqliteOnlineStore
275+
276+
supported_types: list[type] = [SqliteOnlineStore]
277+
for module, cls_name in (
278+
("feast.infra.online_stores.mysql_online_store.mysql", "MySQLOnlineStore"),
279+
(
280+
"feast.infra.online_stores.postgres_online_store.postgres",
281+
"PostgreSQLOnlineStore",
282+
),
283+
("feast.infra.online_stores.faiss_online_store", "FaissOnlineStore"),
284+
("feast.infra.online_stores.redis", "RedisOnlineStore"),
285+
("feast.infra.online_stores.dynamodb", "DynamoDBOnlineStore"),
286+
(
287+
"feast.infra.online_stores.milvus_online_store.milvus",
288+
"MilvusOnlineStore",
289+
),
290+
):
291+
try:
292+
import importlib
293+
294+
mod = importlib.import_module(module)
295+
supported_types.append(getattr(mod, cls_name))
296+
except Exception:
297+
pass
298+
return isinstance(self, tuple(supported_types))
299+
313300
async def get_online_features_async(
314301
self,
315302
config: RepoConfig,

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

Lines changed: 149 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434

3535
from feast import Entity, FeatureView, RepoConfig, utils
3636
from feast.feature_service import FeatureService
37+
from feast.infra.key_encoding_utils import serialize_entity_key
3738
from feast.infra.online_stores.helpers import (
3839
_mmh3,
3940
_redis_key,
@@ -506,15 +507,14 @@ async def online_write_batch_async(
506507
def _generate_redis_keys_for_entities(
507508
self, config: RepoConfig, entity_keys: List[EntityKeyProto]
508509
) -> List[bytes]:
509-
keys = []
510-
for entity_key in entity_keys:
511-
redis_key_bin = _redis_key(
512-
config.project,
513-
entity_key,
514-
entity_key_serialization_version=config.entity_key_serialization_version,
515-
)
516-
keys.append(redis_key_bin)
517-
return keys
510+
project = config.project
511+
version = config.entity_key_serialization_version
512+
project_bytes = project.encode("utf-8")
513+
return [
514+
serialize_entity_key(ek, entity_key_serialization_version=version)
515+
+ project_bytes
516+
for ek in entity_keys
517+
]
518518

519519
def _generate_hset_keys_for_features(
520520
self,
@@ -751,6 +751,146 @@ def get_online_features(
751751
)
752752
return OnlineResponse(online_features_response, feature_types=feature_types)
753753

754+
async def get_online_features_async(
755+
self,
756+
config: RepoConfig,
757+
features: Union[List[str], FeatureService],
758+
entity_rows: Union[
759+
List[Dict[str, Any]],
760+
Mapping[str, Union[Sequence[Any], Sequence[ValueProto], RepeatedValue]],
761+
],
762+
registry: BaseRegistry,
763+
project: str,
764+
full_feature_names: bool = False,
765+
include_feature_view_version_metadata: bool = False,
766+
) -> OnlineResponse:
767+
"""
768+
Async version of get_online_features using a single batched Redis pipeline.
769+
770+
Mirrors the sync override: all HMGET commands across all feature views are
771+
issued in one async pipeline execution (O(1) round trips).
772+
"""
773+
if isinstance(entity_rows, list):
774+
columnar: Dict[str, List[Any]] = {k: [] for k in entity_rows[0].keys()}
775+
for entity_row in entity_rows:
776+
for key, value in entity_row.items():
777+
try:
778+
columnar[key].append(value)
779+
except KeyError as e:
780+
raise ValueError(
781+
"All entity_rows must have the same keys."
782+
) from e
783+
entity_rows = columnar
784+
785+
(
786+
join_key_values,
787+
grouped_refs,
788+
entity_name_to_join_key_map,
789+
requested_on_demand_feature_views,
790+
feature_refs,
791+
requested_result_row_names,
792+
online_features_response,
793+
) = utils._prepare_entities_to_read_from_online_store(
794+
registry=registry,
795+
project=project,
796+
features=features,
797+
entity_values=entity_rows,
798+
full_feature_names=full_feature_names,
799+
native_entity_values=True,
800+
)
801+
802+
self._check_versioned_read_support(grouped_refs)
803+
804+
_track_read = False
805+
try:
806+
from feast.metrics import _config as _metrics_config
807+
808+
_track_read = _metrics_config.online_features
809+
except Exception:
810+
pass
811+
812+
if _track_read:
813+
import time as _time
814+
815+
_read_start = _time.monotonic()
816+
817+
work_items = []
818+
for table, requested_features in grouped_refs:
819+
table_entity_values, idxs, output_len = utils._get_unique_entities(
820+
table,
821+
join_key_values,
822+
entity_name_to_join_key_map,
823+
)
824+
entity_key_protos = utils._get_entity_key_protos(table_entity_values)
825+
fv_name = _versioned_fv_name(table, config)
826+
redis_keys = self._generate_redis_keys_for_entities(
827+
config, entity_key_protos
828+
)
829+
req_features, hset_keys = self._generate_hset_keys_for_features(
830+
table, requested_features, fv_name_override=fv_name
831+
)
832+
work_items.append(
833+
(table, req_features, fv_name, hset_keys, redis_keys, idxs, output_len)
834+
)
835+
836+
if work_items:
837+
client = await self._get_client_async(config.online_store)
838+
async with client.pipeline(transaction=False) as pipe:
839+
for _, _, _, hset_keys, redis_keys, _, _ in work_items:
840+
for redis_key in redis_keys:
841+
pipe.hmget(redis_key, hset_keys)
842+
all_results = await pipe.execute()
843+
844+
offset = 0
845+
for (
846+
table,
847+
req_features,
848+
fv_name,
849+
_,
850+
redis_keys,
851+
idxs,
852+
output_len,
853+
) in work_items:
854+
n = len(redis_keys)
855+
redis_values = all_results[offset : offset + n]
856+
offset += n
857+
858+
read_rows = self._convert_redis_values_to_protobuf(
859+
redis_values, fv_name, req_features
860+
)
861+
862+
utils._populate_response_from_feature_data(
863+
req_features,
864+
read_rows,
865+
idxs,
866+
online_features_response,
867+
full_feature_names,
868+
table,
869+
output_len,
870+
include_feature_view_version_metadata,
871+
)
872+
873+
if _track_read:
874+
from feast.metrics import track_online_store_read
875+
876+
track_online_store_read(_time.monotonic() - _read_start)
877+
878+
feature_types = self._build_feature_types(grouped_refs)
879+
880+
if requested_on_demand_feature_views:
881+
utils._augment_response_with_on_demand_transforms(
882+
online_features_response,
883+
feature_refs,
884+
requested_on_demand_feature_views,
885+
full_feature_names,
886+
feature_types=feature_types,
887+
)
888+
889+
utils._drop_unneeded_columns(
890+
online_features_response, requested_result_row_names
891+
)
892+
return OnlineResponse(online_features_response, feature_types=feature_types)
893+
754894
def _get_features_for_entity(
755895
self,
756896
values: List[ByteString],

0 commit comments

Comments
 (0)