fix: Performance regression in /get-online-features (#4892)

* put blocking opps in get online features in threadpool Signed-off-by: Rob Howley <[email protected]> * chore: leave comment in path operation Signed-off-by: Rob Howley <[email protected]> * use quoted type hint Signed-off-by: Rob Howley <[email protected]> * remove hint Signed-off-by: Rob Howley <[email protected]> --------- Signed-off-by: Rob Howley <[email protected]>
feast-dev · Jan 15, 2025 · 0db56a2 · 0db56a2
1 parent bd9f071
commit 0db56a2
Showing 1 changed file with 39 additions and 29 deletions.
diff --git a/sdk/python/feast/feature_server.py b/sdk/python/feast/feature_server.py
@@ -76,6 +76,38 @@ class GetOnlineFeaturesRequest(BaseModel):
     full_feature_names: bool = False
 
 
+def _get_features(request: GetOnlineFeaturesRequest, store: "feast.FeatureStore"):
+    if request.feature_service:
+        feature_service = 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,
+            )
+        )
+        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]
+            )
+        features = request.features  # type: ignore
+
+    return features
+
+
 def get_app(
     store: "feast.FeatureStore",
     registry_ttl_sec: int = DEFAULT_FEATURE_SERVER_REGISTRY_TTL,
@@ -121,33 +153,7 @@ async def lifespan(app: FastAPI):
     )
     async def get_online_features(request: GetOnlineFeaturesRequest) -> Dict[str, Any]:
         # Initialize parameters for FeatureStore.get_online_features(...) call
-        if request.feature_service:
-            feature_service = 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,
-                )
-            )
-            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]
-                )
-            features = request.features  # type: ignore
+        features = await run_in_threadpool(_get_features, request, store)
 
         read_params = dict(
             features=features,
@@ -163,9 +169,13 @@ async def get_online_features(request: GetOnlineFeaturesRequest) -> Dict[str, An
             )
 
         # Convert the Protobuf object to JSON and return it
-        return MessageToDict(
-            response.proto, preserving_proto_field_name=True, float_precision=18
+        response_dict = await run_in_threadpool(
+            MessageToDict,
+            response.proto,
+            preserving_proto_field_name=True,
+            float_precision=18,
         )
+        return response_dict
 
     @app.post("/push", dependencies=[Depends(inject_user_details)])
     async def push(request: PushFeaturesRequest) -> None: