add explicit async support tags, use in server

robhowley · robhowley · commit f68ebbfd21ee · 2024-10-17T21:20:53.000-04:00
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
@@ -10,6 +10,7 @@
 import psutil
 from dateutil import parser
 from fastapi import Depends, FastAPI, Request, Response, status
+from fastapi.concurrency import run_in_threadpool
 from fastapi.logger import logger
 from fastapi.responses import JSONResponse
 from google.protobuf.json_format import MessageToDict
@@ -100,10 +101,8 @@ def async_refresh():
     @asynccontextmanager
     async def lifespan(app: FastAPI):
         async_refresh()
-        await store.initialize()
         yield
         stop_refresh()
-        await store.close()
 
     app = FastAPI(lifespan=lifespan)
 
@@ -114,7 +113,7 @@ async def get_body(request: Request):
         "/get-online-features",
         dependencies=[Depends(inject_user_details)],
     )
-    def get_online_features(body=Depends(get_body)):
+    async 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"]
@@ -147,15 +146,22 @@ def get_online_features(body=Depends(get_body)):
                     resource=od_feature_view, actions=[AuthzedAction.READ_ONLINE]
                 )
 
-        response_proto = store.get_online_features(
+        read_params = dict(
             features=features,
             entity_rows=entity_rows,
             full_feature_names=full_feature_names,
-        ).proto
+        )
+
+        if store._get_provider().async_supported.online.read:
+            response = await store.get_online_features_async(**read_params)
+        else:
+            response = await run_in_threadpool(
+                lambda: store.get_online_features(**read_params)
+            )
 
         # Convert the Protobuf object to JSON and return it
         return MessageToDict(
-            response_proto, preserving_proto_field_name=True, float_precision=18
+            response.proto, preserving_proto_field_name=True, float_precision=18
         )
 
     @app.post("/push", dependencies=[Depends(inject_user_details)])
diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
@@ -2078,14 +2078,6 @@ def list_saved_datasets(
             self.project, allow_cache=allow_cache, tags=tags
         )
 
-    async def initialize(self) -> None:
-        """Initialize long-lived clients and/or resources needed for accessing datastores"""
-        await self._get_provider().initialize(self.config)
-
-    async def close(self) -> None:
-        """Cleanup any long-lived clients and/or resources"""
-        await self._get_provider().close()
-
 
 def _print_materialization_log(
     start_date, end_date, num_feature_views: int, online_store: str
diff --git a/sdk/python/feast/infra/offline_stores/offline_store.py b/sdk/python/feast/infra/offline_stores/offline_store.py
@@ -35,6 +35,7 @@
 from feast.feature_logging import LoggingConfig, LoggingSource
 from feast.feature_view import FeatureView
 from feast.infra.registry.base_registry import BaseRegistry
+from feast.infra.supported_async_methods import SupportedAsyncMethods
 from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.repo_config import RepoConfig
 from feast.saved_dataset import SavedDatasetStorage
@@ -223,6 +224,10 @@ class OfflineStore(ABC):
     the SnowflakeOfflineStore can handle SnowflakeSources but not FileSources.
     """
 
+    @property
+    def async_supported(self) -> SupportedAsyncMethods:
+        return SupportedAsyncMethods()
+
     @staticmethod
     def pull_latest_from_table_or_query(
         config: RepoConfig,
diff --git a/sdk/python/feast/infra/online_stores/dynamodb.py b/sdk/python/feast/infra/online_stores/dynamodb.py
@@ -12,19 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import asyncio
-import contextlib
 import itertools
 import logging
 from datetime import datetime
 from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union
 
-from aiobotocore.config import AioConfig
 from pydantic import StrictBool, StrictStr
 
 from feast import Entity, FeatureView, utils
 from feast.infra.infra_object import DYNAMODB_INFRA_OBJECT_CLASS_TYPE, InfraObject
 from feast.infra.online_stores.helpers import compute_entity_id
 from feast.infra.online_stores.online_store import OnlineStore
+from feast.infra.supported_async_methods import SupportedAsyncMethods
 from feast.protos.feast.core.DynamoDBTable_pb2 import (
     DynamoDBTable as DynamoDBTableProto,
 )
@@ -76,9 +75,6 @@ class DynamoDBOnlineStoreConfig(FeastConfigBaseModel):
     session_based_auth: bool = False
     """AWS session based client authentication"""
 
-    max_pool_connections: int = 10
-    """Max number of connections for async Dynamodb operations"""
-
 
 class DynamoDBOnlineStore(OnlineStore):
     """
@@ -91,14 +87,11 @@ class DynamoDBOnlineStore(OnlineStore):
 
     _dynamodb_client = None
     _dynamodb_resource = None
+    _aioboto_session = None
 
-    async def initialize(self, config: RepoConfig):
-        await _get_aiodynamodb_client(
-            config.online_store.region, config.online_store.max_pool_connections
-        )
-
-    async def close(self):
-        await _aiodynamodb_close()
+    @property
+    def async_supported(self) -> SupportedAsyncMethods:
+        return SupportedAsyncMethods(read=True)
 
     def update(
         self,
@@ -333,17 +326,15 @@ def to_tbl_resp(raw_client_response):
             batches.append(batch)
             entity_id_batches.append(entity_id_batch)
 
-        client = await _get_aiodynamodb_client(
-            online_config.region, online_config.max_pool_connections
-        )
-        response_batches = await asyncio.gather(
-            *[
-                client.batch_get_item(
-                    RequestItems=entity_id_batch,
-                )
-                for entity_id_batch in entity_id_batches
-            ]
-        )
+        async with self._get_aiodynamodb_client(online_config.region) as client:
+            response_batches = await asyncio.gather(
+                *[
+                    client.batch_get_item(
+                        RequestItems=entity_id_batch,
+                    )
+                    for entity_id_batch in entity_id_batches
+                ]
+            )
 
         result_batches = []
         for batch, response in zip(batches, response_batches):
@@ -358,6 +349,14 @@ def to_tbl_resp(raw_client_response):
 
         return list(itertools.chain(*result_batches))
 
+    def _get_aioboto_session(self):
+        if self._aioboto_session is None:
+            self._aioboto_session = session.get_session()
+        return self._aioboto_session
+
+    def _get_aiodynamodb_client(self, region: str):
+        return self._get_aioboto_session().create_client("dynamodb", region_name=region)
+
     def _get_dynamodb_client(
         self,
         region: str,
@@ -490,38 +489,6 @@ def _to_client_batch_get_payload(online_config, table_name, batch):
         }
 
 
-_aioboto_session = None
-_aioboto_client = None
-
-
-def _get_aioboto_session():
-    global _aioboto_session
-    if _aioboto_session is None:
-        logger.debug("initializing the aiobotocore session")
-        _aioboto_session = session.get_session()
-    return _aioboto_session
-
-
-async def _get_aiodynamodb_client(region: str, max_pool_connections: int):
-    global _aioboto_client
-    if _aioboto_client is None:
-        logger.debug("initializing the aiobotocore dynamodb client")
-        client_context = _get_aioboto_session().create_client(
-            "dynamodb",
-            region_name=region,
-            config=AioConfig(max_pool_connections=max_pool_connections),
-        )
-        context_stack = contextlib.AsyncExitStack()
-        _aioboto_client = await context_stack.enter_async_context(client_context)
-    return _aioboto_client
-
-
-async def _aiodynamodb_close():
-    global _aioboto_client
-    if _aioboto_client:
-        await _aioboto_client.close()
-
-
 def _initialize_dynamodb_client(
     region: str,
     endpoint_url: Optional[str] = None,
diff --git a/sdk/python/feast/infra/online_stores/online_store.py b/sdk/python/feast/infra/online_stores/online_store.py
@@ -22,6 +22,7 @@
 from feast.feature_view import FeatureView
 from feast.infra.infra_object import InfraObject
 from feast.infra.registry.base_registry import BaseRegistry
+from feast.infra.supported_async_methods import SupportedAsyncMethods
 from feast.online_response import OnlineResponse
 from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
@@ -36,6 +37,10 @@ class OnlineStore(ABC):
     The interface that Feast uses to interact with the storage system that handles online features.
     """
 
+    @property
+    def async_supported(self) -> SupportedAsyncMethods:
+        return SupportedAsyncMethods()
+
     @abstractmethod
     def online_write_batch(
         self,
@@ -390,9 +395,3 @@ def retrieve_online_documents(
         raise NotImplementedError(
             f"Online store {self.__class__.__name__} does not support online retrieval"
         )
-
-    async def initialize(self, config: RepoConfig) -> None:
-        pass
-
-    async def close(self) -> None:
-        pass
diff --git a/sdk/python/feast/infra/passthrough_provider.py b/sdk/python/feast/infra/passthrough_provider.py
@@ -35,6 +35,7 @@
 from feast.infra.online_stores.helpers import get_online_store_from_config
 from feast.infra.provider import Provider
 from feast.infra.registry.base_registry import BaseRegistry
+from feast.infra.supported_async_methods import ProviderAsyncMethods
 from feast.online_response import OnlineResponse
 from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
@@ -79,6 +80,13 @@ def offline_store(self):
             )
         return self._offline_store
 
+    @property
+    def async_supported(self) -> ProviderAsyncMethods:
+        return ProviderAsyncMethods(
+            online=self.online_store.async_supported,
+            offline=self.offline_store.async_supported,
+        )
+
     @property
     def batch_engine(self) -> BatchMaterializationEngine:
         if self._batch_engine:
@@ -473,9 +481,3 @@ def get_table_column_names_and_types_from_data_source(
         return self.offline_store.get_table_column_names_and_types_from_data_source(
             config=config, data_source=data_source
         )
-
-    async def initialize(self, config: RepoConfig) -> None:
-        await self.online_store.initialize(config)
-
-    async def close(self) -> None:
-        await self.online_store.close()
diff --git a/sdk/python/feast/infra/provider.py b/sdk/python/feast/infra/provider.py
@@ -27,6 +27,7 @@
 from feast.infra.infra_object import Infra
 from feast.infra.offline_stores.offline_store import RetrievalJob
 from feast.infra.registry.base_registry import BaseRegistry
+from feast.infra.supported_async_methods import ProviderAsyncMethods
 from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.online_response import OnlineResponse
 from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
@@ -55,6 +56,10 @@ class Provider(ABC):
     def __init__(self, config: RepoConfig):
         pass
 
+    @property
+    def async_supported(self) -> ProviderAsyncMethods:
+        return ProviderAsyncMethods()
+
     @abstractmethod
     def update_infra(
         self,
@@ -429,14 +434,6 @@ def get_table_column_names_and_types_from_data_source(
         """
         pass
 
-    @abstractmethod
-    async def initialize(self, config: RepoConfig) -> None:
-        pass
-
-    @abstractmethod
-    async def close(self) -> None:
-        pass
-
 
 def get_provider(config: RepoConfig) -> Provider:
     if "." not in config.provider:
diff --git a/sdk/python/feast/infra/supported_async_methods.py b/sdk/python/feast/infra/supported_async_methods.py
@@ -0,0 +1,11 @@
+from pydantic import BaseModel, Field
+
+
+class SupportedAsyncMethods(BaseModel):
+    read: bool = Field(default=False)
+    write: bool = Field(default=False)
+
+
+class ProviderAsyncMethods(BaseModel):
+    online: SupportedAsyncMethods = Field(default_factory=SupportedAsyncMethods)
+    offline: SupportedAsyncMethods = Field(default_factory=SupportedAsyncMethods)
diff --git a/sdk/python/tests/foo_provider.py b/sdk/python/tests/foo_provider.py
@@ -184,9 +184,3 @@ async def get_online_features_async(
         full_feature_names: bool = False,
     ) -> OnlineResponse:
         pass
-
-    async def initialize(self, config: RepoConfig) -> None:
-        pass
-
-    async def close(self) -> None:
-        pass
