Add feature server configuration for AWS lambda (feast-dev#1865)

felixwang9817 · web-flow · commit b3d58f606174 · 2021-09-15T16:57:09.000-07:00
* Add feature server configuration for AWS lambda

Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;

* Rename app.py to config.py and add defaults to config

Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;

* Change asserts to if/raise checks

Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;

* Always generate feature server config

Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;

* Make feature server config optional

Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;

* Small fix

Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;
diff --git a/sdk/python/feast/errors.py b/sdk/python/feast/errors.py
@@ -93,6 +93,25 @@ def __init__(self, provider_name):
         super().__init__(f"Provider '{provider_name}' is not implemented")
 
 
+class FeastProviderNotSetError(Exception):
+    def __init__(self):
+        super().__init__("Provider is not set, but is required")
+
+
+class FeastFeatureServerTypeSetError(Exception):
+    def __init__(self, feature_server_type: str):
+        super().__init__(
+            f"Feature server type was set to {feature_server_type}, but the type should be determined by the provider"
+        )
+
+
+class FeastFeatureServerTypeInvalidError(Exception):
+    def __init__(self, feature_server_type: str):
+        super().__init__(
+            f"Feature server type was set to {feature_server_type}, but this type is invalid"
+        )
+
+
 class FeastModuleImportError(Exception):
     def __init__(self, module_name: str, module_type: str):
         super().__init__(f"Could not import {module_type} module '{module_name}'")
diff --git a/sdk/python/feast/infra/feature_servers/aws_lambda/config.py b/sdk/python/feast/infra/feature_servers/aws_lambda/config.py
@@ -0,0 +1,23 @@
+from pydantic import StrictBool, StrictStr
+from pydantic.typing import Literal
+
+from feast.repo_config import FeastConfigBaseModel
+
+
+class AwsLambdaFeatureServerConfig(FeastConfigBaseModel):
+    """Feature server config for AWS Lambda."""
+
+    type: Literal["aws_lambda"] = "aws_lambda"
+    """Feature server type selector."""
+
+    enabled: StrictBool = False
+    """Whether the feature server should be launched."""
+
+    public: StrictBool = True
+    """Whether the endpoint should be publicly accessible."""
+
+    auth: Literal["none", "api-key"] = "none"
+    """Authentication method for the endpoint."""
+
+    execution_role_name: StrictStr
+    """The execution role for the AWS Lambda function."""
diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py
@@ -13,6 +13,11 @@
 from pydantic.error_wrappers import ErrorWrapper
 from pydantic.typing import Dict, Optional, Union
 
+from feast.errors import (
+    FeastFeatureServerTypeInvalidError,
+    FeastFeatureServerTypeSetError,
+    FeastProviderNotSetError,
+)
 from feast.importer import get_class_from_type
 from feast.usage import log_exceptions
 
@@ -32,6 +37,10 @@
     "redshift": "feast.infra.offline_stores.redshift.RedshiftOfflineStore",
 }
 
+FEATURE_SERVER_CONFIG_CLASS_FOR_TYPE = {
+    "aws_lambda": "feast.infra.feature_servers.aws_lambda.config.AwsLambdaFeatureServerConfig",
+}
+
 
 class FeastBaseModel(BaseModel):
     """ Feast Pydantic Configuration Class """
@@ -86,6 +95,9 @@ class RepoConfig(FeastBaseModel):
     offline_store: Any
     """ OfflineStoreConfig: Offline store configuration (optional depending on provider) """
 
+    feature_server: Optional[Any]
+    """ FeatureServerConfig: Feature server configuration (optional depending on provider) """
+
     repo_path: Optional[Path] = None
 
     def __init__(self, **data: Any):
@@ -105,6 +117,11 @@ def __init__(self, **data: Any):
         elif isinstance(self.offline_store, str):
             self.offline_store = get_offline_config_from_type(self.offline_store)()
 
+        if isinstance(self.feature_server, Dict):
+            self.feature_server = get_feature_server_config_from_type(
+                self.feature_server["type"]
+            )(**self.feature_server)
+
     def get_registry_config(self):
         if isinstance(self.registry, str):
             return RegistryConfig(path=self.registry)
@@ -190,6 +207,43 @@ def _validate_offline_store_config(cls, values):
 
         return values
 
+    @root_validator(pre=True)
+    def _validate_feature_server_config(cls, values):
+        # Having no feature server is the default.
+        if "feature_server" not in values:
+            return values
+
+        # Skip if we aren't creating the configuration from a dict
+        if not isinstance(values["feature_server"], Dict):
+            return values
+
+        # Make sure that the provider configuration is set. We need it to set the defaults
+        if "provider" not in values:
+            raise FeastProviderNotSetError()
+
+        # Make sure that the type is not set, since we will set it based on the provider.
+        if "type" in values["feature_server"]:
+            raise FeastFeatureServerTypeSetError(values["feature_server"]["type"])
+
+        # Set the default type. We only support AWS Lambda for now.
+        if values["provider"] == "aws":
+            values["feature_server"]["type"] = "aws_lambda"
+
+        feature_server_type = values["feature_server"]["type"]
+
+        # Validate the dict to ensure one of the union types match
+        try:
+            feature_server_config_class = get_feature_server_config_from_type(
+                feature_server_type
+            )
+            feature_server_config_class(**values["feature_server"])
+        except ValidationError as e:
+            raise ValidationError(
+                [ErrorWrapper(e, loc="feature_server")], model=RepoConfig,
+            )
+
+        return values
+
     @validator("project")
     def _validate_project_name(cls, v):
         from feast.repo_operations import is_valid_name
@@ -244,6 +298,16 @@ def get_offline_config_from_type(offline_store_type: str):
     return get_class_from_type(module_name, config_class_name, config_class_name)
 
 
+def get_feature_server_config_from_type(feature_server_type: str):
+    # We do not support custom feature servers right now.
+    if feature_server_type not in FEATURE_SERVER_CONFIG_CLASS_FOR_TYPE:
+        raise FeastFeatureServerTypeInvalidError(feature_server_type)
+
+    feature_server_type = FEATURE_SERVER_CONFIG_CLASS_FOR_TYPE[feature_server_type]
+    module_name, config_class_name = feature_server_type.rsplit(".", 1)
+    return get_class_from_type(module_name, config_class_name, config_class_name)
+
+
 def load_repo_config(repo_path: Path) -> RepoConfig:
     config_path = repo_path / "feature_store.yaml"
 
