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

robhowley · web-flow · commit 0db56a2cb588 · 2025-01-14T20:38:36.000-05:00
* put blocking opps in get online features in threadpool

Signed-off-by: Rob Howley &lt;howley.robert@gmail.com&gt;

* chore: leave comment in path operation

Signed-off-by: Rob Howley &lt;howley.robert@gmail.com&gt;

* use quoted type hint

Signed-off-by: Rob Howley &lt;howley.robert@gmail.com&gt;

* remove hint

Signed-off-by: Rob Howley &lt;howley.robert@gmail.com&gt;

---------

Signed-off-by: Rob Howley &lt;howley.robert@gmail.com&gt;
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:
