fix: Bytewax materializer security context (feast-dev#3573)

* fix: Bytewax materializer security context Signed-off-by: adamschmidt <aschmidt1978@gmail.com> * fix: lint Signed-off-by: adamschmidt <aschmidt1978@gmail.com> * chore: docs Signed-off-by: adamschmidt <aschmidt1978@gmail.com> --------- Signed-off-by: adamschmidt <aschmidt1978@gmail.com>
zerafachris · Apr 21, 2023 · 6794338 · 6794338
1 parent 67acc01
commit 6794338
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 10 deletions.
diff --git a/docs/reference/batch-materialization/bytewax.md b/docs/reference/batch-materialization/bytewax.md
@@ -58,6 +58,7 @@ batch_engine:
   image_pull_secrets:
     - my_container_secret
   service_account_name: my-k8s-service-account
+  include_security_context_capabilities: false
   annotations:
     # example annotation you might include if running on AWS EKS
     iam.amazonaws.com/role: arn:aws:iam::<account number>:role/MyBytewaxPlatformRole
@@ -73,8 +74,9 @@ batch_engine:
 **Notes:**
 
 * The `namespace` configuration directive specifies which Kubernetes [namespace](https://kubernetes.io/docs/concepts/overview/working-with-objects/namespaces/) jobs, services and configuration maps will be created in.
-* The `image_pull_secrets` configuration directive specifies the pre-configured secret to use when pulling the image container from your registry
-* The `service_account_name` specifies which Kubernetes service account to run the job under
+* The `image_pull_secrets` configuration directive specifies the pre-configured secret to use when pulling the image container from your registry.
+* The `service_account_name` specifies which Kubernetes service account to run the job under.
+* The `include_security_context_capabilities` flag indicates whether or not `"add": ["NET_BIND_SERVICE"]` and `"drop": ["ALL"]` are included in the job & pod security context capabilities.
 * `annotations` allows you to include additional Kubernetes annotations to the job. This is particularly useful for IAM roles which grant the running pod access to cloud platform resources (for example).
 * The `resources` configuration directive sets the standard Kubernetes [resource requests](https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/) for the job containers to utilise when materializing data.
 

diff --git a/sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_engine.py b/sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_engine.py
@@ -58,6 +58,9 @@ class BytewaxMaterializationEngineConfig(FeastConfigBaseModel):
     annotations: dict = {}
     """ (optional) Annotations to apply to the job container. Useful for linking the service account to IAM roles, operational metadata, etc  """
 
+    include_security_context_capabilities: bool = True
+    """ (optional)  Include security context capabilities in the init and job container spec """
+
 
 class BytewaxMaterializationEngine(BatchMaterializationEngine):
     def __init__(
@@ -198,6 +201,9 @@ def _create_configuration_map(self, job_id, paths, feature_view, namespace):
             "apiVersion": "v1",
             "metadata": {
                 "name": f"feast-{job_id}",
+                "labels": {
+                    "feast-bytewax-materializer": "configmap",
+                },
             },
             "data": {
                 "feature_store.yaml": feature_store_configuration,
@@ -247,12 +253,22 @@ def _create_job_definition(self, job_id, namespace, pods, env):
         # Add any Feast configured environment variables
         job_env.extend(env)
 
+        securityContextCapabilities = None
+        if self.batch_engine_config.include_security_context_capabilities:
+            securityContextCapabilities = {
+                "add": ["NET_BIND_SERVICE"],
+                "drop": ["ALL"],
+            }
+
         job_definition = {
             "apiVersion": "batch/v1",
             "kind": "Job",
             "metadata": {
                 "name": f"dataflow-{job_id}",
                 "namespace": namespace,
+                "labels": {
+                    "feast-bytewax-materializer": "job",
+                },
             },
             "spec": {
                 "ttlSecondsAfterFinished": 3600,
@@ -262,6 +278,9 @@ def _create_job_definition(self, job_id, namespace, pods, env):
                 "template": {
                     "metadata": {
                         "annotations": self.batch_engine_config.annotations,
+                        "labels": {
+                            "feast-bytewax-materializer": "pod",
+                        },
                     },
                     "spec": {
                         "restartPolicy": "Never",
@@ -282,10 +301,7 @@ def _create_job_definition(self, job_id, namespace, pods, env):
                                 "resources": {},
                                 "securityContext": {
                                     "allowPrivilegeEscalation": False,
-                                    "capabilities": {
-                                        "add": ["NET_BIND_SERVICE"],
-                                        "drop": ["ALL"],
-                                    },
+                                    "capabilities": securityContextCapabilities,
                                     "readOnlyRootFilesystem": True,
                                 },
                                 "terminationMessagePath": "/dev/termination-log",
@@ -320,10 +336,7 @@ def _create_job_definition(self, job_id, namespace, pods, env):
                                 "resources": self.batch_engine_config.resources,
                                 "securityContext": {
                                     "allowPrivilegeEscalation": False,
-                                    "capabilities": {
-                                        "add": ["NET_BIND_SERVICE"],
-                                        "drop": ["ALL"],
-                                    },
+                                    "capabilities": securityContextCapabilities,
                                     "readOnlyRootFilesystem": False,
                                 },
                                 "terminationMessagePath": "/dev/termination-log",