diff --git a/docs/reference/batch-materialization/bytewax.md b/docs/reference/batch-materialization/bytewax.md index 2e28937f50..6a97bd391d 100644 --- 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:::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 index 991eafa641..b222128bbb 100644 --- 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",