chore: Improve async code for feature server

Signed-off-by: Francisco Javier Arceo <[email protected]>

Author: franciscojavierarceo

sdk/python/feast/feature_server.py

@@ -120,27 +120,26 @@ class SaveDocumentRequest(BaseModel):
     data: dict
 
 
-def _get_features(
+async def _get_features(
     request: Union[GetOnlineFeaturesRequest, GetOnlineDocumentsRequest],
     store: "feast.FeatureStore",
 ):
     if request.feature_service:
-        feature_service = store.get_feature_service(
-            request.feature_service, allow_cache=True
+        feature_service = await run_in_threadpool(
+            store.get_feature_service, request.feature_service, allow_cache=True
         )
         assert_permissions(
             resource=feature_service, actions=[AuthzedAction.READ_ONLINE]
         )
         features = feature_service  # type: ignore
     else:
-        all_feature_views, all_on_demand_feature_views = (
-            utils._get_feature_views_to_use(
-                store.registry,
-                store.project,
-                request.features,
-                allow_cache=True,
-                hide_dummy_entity=False,
-            )
+        all_feature_views, all_on_demand_feature_views = await run_in_threadpool(
+            utils._get_feature_views_to_use,
+            store.registry,
+            store.project,
+            request.features,
+            allow_cache=True,
+            hide_dummy_entity=False,
         )
         for feature_view in all_feature_views:
             assert_permissions(
@@ -230,7 +229,7 @@ async def lifespan(app: FastAPI):
     )
     async def get_online_features(request: GetOnlineFeaturesRequest) -> Dict[str, Any]:
         # Initialize parameters for FeatureStore.get_online_features(...) call
-        features = await run_in_threadpool(_get_features, request, store)
+        features = await _get_features(request, store)
 
         read_params = dict(
             features=features,
@@ -265,7 +264,7 @@ async def retrieve_online_documents(
             "This endpoint is in alpha and will be moved to /get-online-features when stable."
         )
         # Initialize parameters for FeatureStore.retrieve_online_documents_v2(...) call
-        features = await run_in_threadpool(_get_features, request, store)
+        features = await _get_features(request, store)
 
         read_params = dict(features=features, query=request.query, top_k=request.top_k)
         if request.api_version == 2 and request.query_string is not None:
@@ -342,26 +341,31 @@ async def push(request: PushFeaturesRequest) -> None:
         else:
             store.push(**push_params)
 
-    def _get_feast_object(
+    async def _get_feast_object(
         feature_view_name: str, allow_registry_cache: bool
     ) -> FeastObject:
         try:
-            return store.get_stream_feature_view(  # type: ignore
-                feature_view_name, allow_registry_cache=allow_registry_cache
+            return await run_in_threadpool(
+                store.get_stream_feature_view,
+                feature_view_name,
+                allow_registry_cache=allow_registry_cache,
             )
         except FeatureViewNotFoundException:
-            return store.get_feature_view(  # type: ignore
-                feature_view_name, allow_registry_cache=allow_registry_cache
+            return await run_in_threadpool(
+                store.get_feature_view,
+                feature_view_name,
+                allow_registry_cache=allow_registry_cache,
             )
 
     @app.post("/write-to-online-store", dependencies=[Depends(inject_user_details)])
-    def write_to_online_store(request: WriteToFeatureStoreRequest) -> None:
+    async def write_to_online_store(request: WriteToFeatureStoreRequest) -> None:
         df = pd.DataFrame(request.df)
         feature_view_name = request.feature_view_name
         allow_registry_cache = request.allow_registry_cache
-        resource = _get_feast_object(feature_view_name, allow_registry_cache)
+        resource = await _get_feast_object(feature_view_name, allow_registry_cache)
         assert_permissions(resource=resource, actions=[AuthzedAction.WRITE_ONLINE])
-        store.write_to_online_store(
+        await run_in_threadpool(
+            store.write_to_online_store,
             feature_view_name=feature_view_name,
             df=df,
             allow_registry_cache=allow_registry_cache,
@@ -428,10 +432,11 @@ async def chat_ui():
         return Response(content=content, media_type="text/html")
 
     @app.post("/materialize", dependencies=[Depends(inject_user_details)])
-    def materialize(request: MaterializeRequest) -> None:
+    async def materialize(request: MaterializeRequest) -> None:
         for feature_view in request.feature_views or []:
+            resource = await _get_feast_object(feature_view, True)
             assert_permissions(
-                resource=_get_feast_object(feature_view, True),
+                resource=resource,
                 actions=[AuthzedAction.WRITE_ONLINE],
             )
 
@@ -450,22 +455,26 @@ def materialize(request: MaterializeRequest) -> None:
             start_date = utils.make_tzaware(parser.parse(request.start_ts))
             end_date = utils.make_tzaware(parser.parse(request.end_ts))
 
-        store.materialize(
+        await run_in_threadpool(
+            store.materialize,
             start_date,
             end_date,
             request.feature_views,
             disable_event_timestamp=request.disable_event_timestamp,
         )
 
     @app.post("/materialize-incremental", dependencies=[Depends(inject_user_details)])
-    def materialize_incremental(request: MaterializeIncrementalRequest) -> None:
+    async def materialize_incremental(request: MaterializeIncrementalRequest) -> None:
         for feature_view in request.feature_views or []:
+            resource = await _get_feast_object(feature_view, True)
             assert_permissions(
-                resource=_get_feast_object(feature_view, True),
+                resource=resource,
                 actions=[AuthzedAction.WRITE_ONLINE],
             )
-        store.materialize_incremental(
-            utils.make_tzaware(parser.parse(request.end_ts)), request.feature_views
+        await run_in_threadpool(
+            store.materialize_incremental,
+            utils.make_tzaware(parser.parse(request.end_ts)),
+            request.feature_views,
         )
 
     @app.exception_handler(Exception)

sdk/python/tests/unit/test_feature_server_async.py

@@ -0,0 +1,28 @@
+from unittest.mock import AsyncMock, MagicMock
+
+from fastapi.testclient import TestClient
+
+from feast.feature_server import get_app
+from feast.online_response import OnlineResponse
+from feast.protos.feast.serving.ServingService_pb2 import GetOnlineFeaturesResponse
+
+
+def test_async_get_online_features():
+    """Test that async get_online_features endpoint works correctly"""
+    fs = MagicMock()
+    fs._get_provider.return_value.async_supported.online.read = True
+    fs.get_online_features_async = AsyncMock(
+        return_value=OnlineResponse(GetOnlineFeaturesResponse())
+    )
+    fs.get_feature_service = MagicMock()
+    fs.initialize = AsyncMock()
+    fs.close = AsyncMock()
+
+    client = TestClient(get_app(fs))
+    response = client.post(
+        "/get-online-features",
+        json={"features": ["test:feature"], "entities": {"entity_id": [123]}},
+    )
+
+    assert response.status_code == 200
+    assert fs.get_online_features_async.await_count == 1