robhowley · Oct 11, 2024 · Oct 11, 2024 · Oct 11, 2024
diff --git a/sdk/python/feast/cli.py b/sdk/python/feast/cli.py
@@ -918,6 +918,13 @@ def init_command(project_directory, minimal: bool, template: str):
     show_default=True,
     help="Enable the Metrics Server",
 )
+@click.option(
+    "--async-read",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="Use async retrieval of features from the online datastore",
+)
 @click.pass_context
 def serve_command(
     ctx: click.Context,
@@ -929,6 +936,7 @@ def serve_command(
     metrics: bool,
     keep_alive_timeout: int,
     registry_ttl_sec: int = 5,
+    async_read: bool = False,
 ):
     """Start a feature server locally on a given port."""
     store = create_feature_store(ctx)
@@ -942,6 +950,7 @@ def serve_command(
         metrics=metrics,
         keep_alive_timeout=keep_alive_timeout,
         registry_ttl_sec=registry_ttl_sec,
+        use_async_read=async_read,
     )
 
 

diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py
@@ -4,7 +4,7 @@
 import time
 import traceback
 from contextlib import asynccontextmanager
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import pandas as pd
 import psutil
@@ -68,9 +68,50 @@ class MaterializeIncrementalRequest(BaseModel):
     feature_views: Optional[List[str]] = None
 
 
+def _prepare_get_online_features_request(store, body) -> dict[str, Any]:
+    body = json.loads(body)
+    full_feature_names = body.get("full_feature_names", False)
+    entity_rows = body["entities"]
+    # Initialize parameters for FeatureStore.get_online_features(...) call
+    if "feature_service" in body:
+        feature_service = store.get_feature_service(
+            body["feature_service"], allow_cache=True
+        )
+        assert_permissions(
+            resource=feature_service, actions=[AuthzedAction.READ_ONLINE]
+        )
+        features = feature_service
+    else:
+        features = body["features"]
+        all_feature_views, all_on_demand_feature_views = (
+            utils._get_feature_views_to_use(
+                store.registry,
+                store.project,
+                features,
+                allow_cache=True,
+                hide_dummy_entity=False,
+            )
+        )
+        for feature_view in all_feature_views:
+            assert_permissions(
+                resource=feature_view, actions=[AuthzedAction.READ_ONLINE]
+            )
+        for od_feature_view in all_on_demand_feature_views:
+            assert_permissions(
+                resource=od_feature_view, actions=[AuthzedAction.READ_ONLINE]
+            )
+
+    return dict(
+        features=features,
+        entity_rows=entity_rows,
+        full_feature_names=full_feature_names,
+    )
+
+
 def get_app(
     store: "feast.FeatureStore",
     registry_ttl_sec: int = DEFAULT_FEATURE_SERVER_REGISTRY_TTL,
+    use_async_read: bool = False,
 ):
     proto_json.patch()
     # Asynchronously refresh registry, notifying shutdown and canceling the active timer if the app is shutting down
@@ -107,53 +148,36 @@ async def lifespan(app: FastAPI):
     async def get_body(request: Request):
         return await request.body()
 
-    @app.post(
+    _post_get_online_features = app.post(
         "/get-online-features",
         dependencies=[Depends(inject_user_details)],
     )
-    def get_online_features(body=Depends(get_body)):
-        body = json.loads(body)
-        full_feature_names = body.get("full_feature_names", False)
-        entity_rows = body["entities"]
-        # Initialize parameters for FeatureStore.get_online_features(...) call
-        if "feature_service" in body:
-            feature_service = store.get_feature_service(
-                body["feature_service"], allow_cache=True
+
+    if use_async_read:
+
+        @_post_get_online_features
+        async def get_online_features(body=Depends(get_body)):
+            response = await store.get_online_features_async(
+                **_prepare_get_online_features_request(store, body)
             )
-            assert_permissions(
-                resource=feature_service, actions=[AuthzedAction.READ_ONLINE]
+
+            # Convert the Protobuf object to JSON and return it
+            return MessageToDict(
+                response.proto, preserving_proto_field_name=True, float_precision=18
             )
-            features = feature_service
-        else:
-            features = body["features"]
-            all_feature_views, all_on_demand_feature_views = (
-                utils._get_feature_views_to_use(
-                    store.registry,
-                    store.project,
-                    features,
-                    allow_cache=True,
-                    hide_dummy_entity=False,
-                )
+
+    else:
+
+        @_post_get_online_features
+        def get_online_features(body=Depends(get_body)):
+            response_proto = store.get_online_features(
+                **_prepare_get_online_features_request(store, body)
+            ).proto
+
+            # Convert the Protobuf object to JSON and return it
+            return MessageToDict(
+                response_proto, preserving_proto_field_name=True, float_precision=18
             )
-            for feature_view in all_feature_views:
-                assert_permissions(
-                    resource=feature_view, actions=[AuthzedAction.READ_ONLINE]
-                )
-            for od_feature_view in all_on_demand_feature_views:
-                assert_permissions(
-                    resource=od_feature_view, actions=[AuthzedAction.READ_ONLINE]
-                )
-
-        response_proto = store.get_online_features(
-            features=features,
-            entity_rows=entity_rows,
-            full_feature_names=full_feature_names,
-        ).proto
-
-        # Convert the Protobuf object to JSON and return it
-        return MessageToDict(
-            response_proto, preserving_proto_field_name=True, float_precision=18
-        )
 
     @app.post("/push", dependencies=[Depends(inject_user_details)])
     def push(body=Depends(get_body)):
@@ -276,6 +300,7 @@ def __init__(self, store: "feast.FeatureStore", **options):
             self._app = get_app(
                 store=store,
                 registry_ttl_sec=options["registry_ttl_sec"],
+                use_async_read=options["use_async_read"],
             )
             self._options = options
             super().__init__()
@@ -316,6 +341,7 @@ def start_server(
     keep_alive_timeout: int,
     registry_ttl_sec: int,
     metrics: bool,
+    use_async_read: bool = False,
 ):
     if metrics:
         logger.info("Starting Prometheus Server")
@@ -347,9 +373,10 @@ def start_server(
             workers=workers,
             keepalive=keep_alive_timeout,
             registry_ttl_sec=registry_ttl_sec,
+            use_async_read=use_async_read,
         ).run()
     else:
         import uvicorn
 
-        app = get_app(store, registry_ttl_sec)
+        app = get_app(store, registry_ttl_sec, use_async_read=use_async_read)
         uvicorn.run(app, host=host, port=port, access_log=(not no_access_log))
diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
@@ -1799,6 +1799,7 @@ def serve(
         metrics: bool = False,
         keep_alive_timeout: int = 30,
         registry_ttl_sec: int = 2,
+        use_async_read: bool = False,
     ) -> None:
         """Start the feature consumption server locally on a given port."""
         type_ = type_.lower()
@@ -1816,6 +1817,7 @@ def serve(
             metrics=metrics,
             keep_alive_timeout=keep_alive_timeout,
             registry_ttl_sec=registry_ttl_sec,
+            use_async_read=use_async_read,
         )
 
     def get_feature_server_endpoint(self) -> Optional[str]:

diff --git a/sdk/python/tests/integration/online_store/test_python_feature_server.py b/sdk/python/tests/integration/online_store/test_python_feature_server.py
@@ -20,7 +20,9 @@
 
 @pytest.mark.integration
 @pytest.mark.universal_online_stores
-def test_get_online_features(python_fs_client):
+@pytest.mark.parametrize("use_async", [(False,), (True,)])
+def test_get_online_features(use_async, python_fs_client_factory):
+    python_fs_client = python_fs_client_factory(use_async_read=use_async)
     request_data_dict = {
         "features": [
             "driver_stats:conv_rate",
@@ -87,7 +89,8 @@ def test_push(python_fs_client):
 
 @pytest.mark.integration
 @pytest.mark.universal_online_stores
-def test_push_source_does_not_exist(python_fs_client):
+def test_push_source_does_not_exist(python_fs_client_factory):
+    python_fs_client = python_fs_client_factory()
     initial_temp = _get_temperatures_from_feature_server(
         python_fs_client, location_ids=[1]
     )[0]
@@ -128,14 +131,16 @@ def _get_temperatures_from_feature_server(client, location_ids: List[int]):
 
 
 @pytest.fixture
-def python_fs_client(environment, universal_data_sources, request):
-    fs = environment.feature_store
-    entities, datasets, data_sources = universal_data_sources
-    feature_views = construct_universal_feature_views(data_sources)
-    feast_objects: List[FeastObject] = []
-    feast_objects.extend(feature_views.values())
-    feast_objects.extend([driver(), customer(), location()])
-    fs.apply(feast_objects)
-    fs.materialize(environment.start_date, environment.end_date)
-    client = TestClient(get_app(fs))
-    yield client
+def python_fs_client_factory(environment, universal_data_sources):
+    def _(**app_kwargs):
+        fs = environment.feature_store
+        entities, datasets, data_sources = universal_data_sources
+        feature_views = construct_universal_feature_views(data_sources)
+        feast_objects: List[FeastObject] = []
+        feast_objects.extend(feature_views.values())
+        feast_objects.extend([driver(), customer(), location()])
+        fs.apply(feast_objects)
+        fs.materialize(environment.start_date, environment.end_date)
+        return TestClient(get_app(fs, **app_kwargs))
+
+    return _