We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Add Kubernetes Deployment Options to the Bytewax Materialization Engine
**Is your feature request related to a problem? Please describe.**
The Bytewax materialization engine needs to support configuration options for more advanced Kubernetes deployments (EKS, GCP, etc) to make it usable at scale.

**Describe the solution you'd like**
The main configuration options that are needed for the Bytewax materialization job are:

* setting explicit resource requests and limits (rather than relying on platform defaults which may not be enough)
* supporting service accounts and IAM roles
* specifying an image pull secret to support pulling Docker images from Dockerhub, Artifactory, etc

**Describe alternatives you've considered**
The Kubernetes job that runs is dynamically generated by the bytewax code. Existing configuration options are insufficient.

**Additional context**
I'd really like to test the Bytewax materialization engine on our instance of EKS. In its current implementation, it's operationally too simple and not usable outside of minikube.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_engine.py`
Content:
```
1 import uuid
2 from datetime import datetime
3 from typing import Callable, List, Literal, Sequence, Union
4 
5 import yaml
6 from kubernetes import client
7 from kubernetes import config as k8s_config
8 from kubernetes import utils
9 from kubernetes.utils import FailToCreateError
10 from pydantic import StrictStr
11 from tqdm import tqdm
12 
13 from feast import FeatureView, RepoConfig
14 from feast.batch_feature_view import BatchFeatureView
15 from feast.entity import Entity
16 from feast.infra.materialization.batch_materialization_engine import (
17     BatchMaterializationEngine,
18     MaterializationJob,
19     MaterializationTask,
20 )
21 from feast.infra.offline_stores.offline_store import OfflineStore
22 from feast.infra.online_stores.online_store import OnlineStore
23 from feast.infra.registry.base_registry import BaseRegistry
24 from feast.repo_config import FeastConfigBaseModel
25 from feast.stream_feature_view import StreamFeatureView
26 from feast.utils import _get_column_names, get_default_yaml_file_path
27 
28 from .bytewax_materialization_job import BytewaxMaterializationJob
29 
30 
31 class BytewaxMaterializationEngineConfig(FeastConfigBaseModel):
32     """Batch Materialization Engine config for Bytewax"""
33 
34     type: Literal["bytewax"] = "bytewax"
35     """ Materialization type selector"""
36 
37     namespace: StrictStr = "default"
38     """ (optional) The namespace in Kubernetes to use when creating services, configuration maps and jobs.
39     """
40 
41     image: StrictStr = "bytewax/bytewax-feast:latest"
42     """ (optional) The container image to use when running the materialization job."""
43 
44     env: List[dict] = []
45     """ (optional) A list of environment variables to set in the created Kubernetes pods.
46     These environment variables can be used to reference Kubernetes secrets.
47     """
48 
49 
50 class BytewaxMaterializationEngine(BatchMaterializationEngine):
51     def __init__(
52         self,
53         *,
54         repo_config: RepoConfig,
55         offline_store: OfflineStore,
56         online_store: OnlineStore,
57         **kwargs,
58     ):
59         super().__init__(
60             repo_config=repo_config,
61             offline_store=offline_store,
62             online_store=online_store,
63             **kwargs,
64         )
65         self.repo_config = repo_config
66         self.offline_store = offline_store
67         self.online_store = online_store
68 
69         # TODO: Configure k8s here
70         k8s_config.load_kube_config()
71 
72         self.k8s_client = client.api_client.ApiClient()
73         self.v1 = client.CoreV1Api(self.k8s_client)
74         self.batch_v1 = client.BatchV1Api(self.k8s_client)
75         self.batch_engine_config = repo_config.batch_engine
76         self.namespace = self.batch_engine_config.namespace
77 
78     def update(
79         self,
80         project: str,
81         views_to_delete: Sequence[
82             Union[BatchFeatureView, StreamFeatureView, FeatureView]
83         ],
84         views_to_keep: Sequence[
85             Union[BatchFeatureView, StreamFeatureView, FeatureView]
86         ],
87         entities_to_delete: Sequence[Entity],
88         entities_to_keep: Sequence[Entity],
89     ):
90         """This method ensures that any necessary infrastructure or resources needed by the
91         engine are set up ahead of materialization."""
92         pass
93 
94     def teardown_infra(
95         self,
96         project: str,
97         fvs: Sequence[Union[BatchFeatureView, StreamFeatureView, FeatureView]],
98         entities: Sequence[Entity],
99     ):
100         """This method ensures that any infrastructure or resources set up by ``update()``are torn down."""
101         pass
102 
103     def materialize(
104         self,
105         registry: BaseRegistry,
106         tasks: List[MaterializationTask],
107     ) -> List[MaterializationJob]:
108         return [
109             self._materialize_one(
110                 registry,
111                 task.feature_view,
112                 task.start_time,
113                 task.end_time,
114                 task.project,
115                 task.tqdm_builder,
116             )
117             for task in tasks
118         ]
119 
120     def _materialize_one(
121         self,
122         registry: BaseRegistry,
123         feature_view: Union[BatchFeatureView, StreamFeatureView, FeatureView],
124         start_date: datetime,
125         end_date: datetime,
126         project: str,
127         tqdm_builder: Callable[[int], tqdm],
128     ):
129         entities = []
130         for entity_name in feature_view.entities:
131             entities.append(registry.get_entity(entity_name, project))
132 
133         (
134             join_key_columns,
135             feature_name_columns,
136             timestamp_field,
137             created_timestamp_column,
138         ) = _get_column_names(feature_view, entities)
139 
140         offline_job = self.offline_store.pull_latest_from_table_or_query(
141             config=self.repo_config,
142             data_source=feature_view.batch_source,
143             join_key_columns=join_key_columns,
144             feature_name_columns=feature_name_columns,
145             timestamp_field=timestamp_field,
146             created_timestamp_column=created_timestamp_column,
147             start_date=start_date,
148             end_date=end_date,
149         )
150 
151         paths = offline_job.to_remote_storage()
152         job_id = str(uuid.uuid4())
153         return self._create_kubernetes_job(job_id, paths, feature_view)
154 
155     def _create_kubernetes_job(self, job_id, paths, feature_view):
156         try:
157             # Create a k8s configmap with information needed by bytewax
158             self._create_configuration_map(job_id, paths, feature_view, self.namespace)
159 
160             # Create the k8s job definition
161             self._create_job_definition(
162                 job_id,
163                 self.namespace,
164                 len(paths),  # Create a pod for each parquet file
165                 self.batch_engine_config.env,
166             )
167         except FailToCreateError as failures:
168             return BytewaxMaterializationJob(job_id, self.namespace, error=failures)
169 
170         return BytewaxMaterializationJob(job_id, self.namespace)
171 
172     def _create_configuration_map(self, job_id, paths, feature_view, namespace):
173         """Create a Kubernetes configmap for this job"""
174 
175         repo_path = self.repo_config.repo_path
176         assert repo_path
177         feature_store_path = get_default_yaml_file_path(repo_path)
178         feature_store_configuration = feature_store_path.read_text()
179 
180         materialization_config = yaml.dump(
181             {"paths": paths, "feature_view": feature_view.name}
182         )
183 
184         configmap_manifest = {
185             "kind": "ConfigMap",
186             "apiVersion": "v1",
187             "metadata": {
188                 "name": f"feast-{job_id}",
189             },
190             "data": {
191                 "feature_store.yaml": feature_store_configuration,
192                 "bytewax_materialization_config.yaml": materialization_config,
193             },
194         }
195         self.v1.create_namespaced_config_map(
196             namespace=namespace,
197             body=configmap_manifest,
198         )
199 
200     def _create_job_definition(self, job_id, namespace, pods, env):
201         """Create a kubernetes job definition."""
202         job_env = [
203             {"name": "RUST_BACKTRACE", "value": "full"},
204             {
205                 "name": "BYTEWAX_PYTHON_FILE_PATH",
206                 "value": "/bytewax/dataflow.py",
207             },
208             {"name": "BYTEWAX_WORKDIR", "value": "/bytewax"},
209             {
210                 "name": "BYTEWAX_WORKERS_PER_PROCESS",
211                 "value": "1",
212             },
213             {
214                 "name": "BYTEWAX_POD_NAME",
215                 "valueFrom": {
216                     "fieldRef": {
217                         "apiVersion": "v1",
218                         "fieldPath": "metadata.annotations['batch.kubernetes.io/job-completion-index']",
219                     }
220                 },
221             },
222             {
223                 "name": "BYTEWAX_REPLICAS",
224                 "value": f"{pods}",
225             },
226             {
227                 "name": "BYTEWAX_KEEP_CONTAINER_ALIVE",
228                 "value": "false",
229             },
230             {
231                 "name": "BYTEWAX_STATEFULSET_NAME",
232                 "value": f"dataflow-{job_id}",
233             },
234         ]
235         # Add any Feast configured environment variables
236         job_env.extend(env)
237 
238         job_definition = {
239             "apiVersion": "batch/v1",
240             "kind": "Job",
241             "metadata": {
242                 "name": f"dataflow-{job_id}",
243                 "namespace": namespace,
244             },
245             "spec": {
246                 "ttlSecondsAfterFinished": 3600,
247                 "completions": pods,
248                 "parallelism": pods,
249                 "completionMode": "Indexed",
250                 "template": {
251                     "spec": {
252                         "restartPolicy": "Never",
253                         "subdomain": f"dataflow-{job_id}",
254                         "initContainers": [
255                             {
256                                 "env": [
257                                     {
258                                         "name": "BYTEWAX_REPLICAS",
259                                         "value": f"{pods}",
260                                     }
261                                 ],
262                                 "image": "busybox",
263                                 "imagePullPolicy": "Always",
264                                 "name": "init-hostfile",
265                                 "resources": {},
266                                 "securityContext": {
267                                     "allowPrivilegeEscalation": False,
268                                     "capabilities": {
269                                         "add": ["NET_BIND_SERVICE"],
270                                         "drop": ["ALL"],
271                                     },
272                                     "readOnlyRootFilesystem": True,
273                                 },
274                                 "terminationMessagePath": "/dev/termination-log",
275                                 "terminationMessagePolicy": "File",
276                                 "volumeMounts": [
277                                     {"mountPath": "/etc/bytewax", "name": "hostfile"},
278                                     {
279                                         "mountPath": "/tmp/bytewax/",
280                                         "name": "python-files",
281                                     },
282                                     {
283                                         "mountPath": "/var/feast/",
284                                         "name": f"feast-{job_id}",
285                                     },
286                                 ],
287                             }
288                         ],
289                         "containers": [
290                             {
291                                 "command": ["sh", "-c", "sh ./entrypoint.sh"],
292                                 "env": job_env,
293                                 "image": self.batch_engine_config.image,
294                                 "imagePullPolicy": "Always",
295                                 "name": "process",
296                                 "ports": [
297                                     {
298                                         "containerPort": 9999,
299                                         "name": "process",
300                                         "protocol": "TCP",
301                                     }
302                                 ],
303                                 "resources": {},
304                                 "securityContext": {
305                                     "allowPrivilegeEscalation": False,
306                                     "capabilities": {
307                                         "add": ["NET_BIND_SERVICE"],
308                                         "drop": ["ALL"],
309                                     },
310                                     "readOnlyRootFilesystem": False,
311                                 },
312                                 "terminationMessagePath": "/dev/termination-log",
313                                 "terminationMessagePolicy": "File",
314                                 "volumeMounts": [
315                                     {"mountPath": "/etc/bytewax", "name": "hostfile"},
316                                     {
317                                         "mountPath": "/var/feast/",
318                                         "name": f"feast-{job_id}",
319                                     },
320                                 ],
321                             }
322                         ],
323                         "volumes": [
324                             {"emptyDir": {}, "name": "hostfile"},
325                             {
326                                 "configMap": {
327                                     "defaultMode": 420,
328                                     "name": f"feast-{job_id}",
329                                 },
330                                 "name": "python-files",
331                             },
332                             {
333                                 "configMap": {"name": f"feast-{job_id}"},
334                                 "name": f"feast-{job_id}",
335                             },
336                         ],
337                     }
338                 },
339             },
340         }
341         utils.create_from_dict(self.k8s_client, job_definition)
342
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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
--- a/sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_engine.py
+++ b/sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_engine.py
@@ -46,6 +46,17 @@
     These environment variables can be used to reference Kubernetes secrets.
     """
 
+    image_pull_secrets: List[str] = []
+    """ (optional) The secrets to use when pulling the image to run for the materialization job """
+
+    resources: dict = {}
+    """ (optional) The resource requests and limits for the materialization containers """
+
+    service_account_name: StrictStr = ""
+    """ (optional) The service account name to use when running the job """
+
+    annotations: dict = {}
+    """ (optional) Annotations to apply to the job container. Useful for linking the service account to IAM roles, operational metadata, etc  """
 
 class BytewaxMaterializationEngine(BatchMaterializationEngine):
     def __init__(
@@ -248,9 +259,14 @@
                 "parallelism": pods,
                 "completionMode": "Indexed",
                 "template": {
+                    "metadata": {
+                        "annotations": self.batch_engine_config.annotations,
+                    },
                     "spec": {
                         "restartPolicy": "Never",
                         "subdomain": f"dataflow-{job_id}",
+                        "imagePullSecrets": self.batch_engine_config.image_pull_secrets,
+                        "serviceAccountName": self.batch_engine_config.service_account_name,
                         "initContainers": [
                             {
                                 "env": [
@@ -300,7 +316,7 @@
                                         "protocol": "TCP",
                                     }
                                 ],
-                                "resources": {},
+                                "resources": self.batch_engine_config.resources,
                                 "securityContext": {
                                     "allowPrivilegeEscalation": False,
                                     "capabilities": {

{"golden_diff": "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\n--- a/sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_engine.py\n+++ b/sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_engine.py\n@@ -46,6 +46,17 @@\n     These environment variables can be used to reference Kubernetes secrets.\n     \"\"\"\n \n+    image_pull_secrets: List[str] = []\n+    \"\"\" (optional) The secrets to use when pulling the image to run for the materialization job \"\"\"\n+\n+    resources: dict = {}\n+    \"\"\" (optional) The resource requests and limits for the materialization containers \"\"\"\n+\n+    service_account_name: StrictStr = \"\"\n+    \"\"\" (optional) The service account name to use when running the job \"\"\"\n+\n+    annotations: dict = {}\n+    \"\"\" (optional) Annotations to apply to the job container. Useful for linking the service account to IAM roles, operational metadata, etc  \"\"\"\n \n class BytewaxMaterializationEngine(BatchMaterializationEngine):\n     def __init__(\n@@ -248,9 +259,14 @@\n                 \"parallelism\": pods,\n                 \"completionMode\": \"Indexed\",\n                 \"template\": {\n+                    \"metadata\": {\n+                        \"annotations\": self.batch_engine_config.annotations,\n+                    },\n                     \"spec\": {\n                         \"restartPolicy\": \"Never\",\n                         \"subdomain\": f\"dataflow-{job_id}\",\n+                        \"imagePullSecrets\": self.batch_engine_config.image_pull_secrets,\n+                        \"serviceAccountName\": self.batch_engine_config.service_account_name,\n                         \"initContainers\": [\n                             {\n                                 \"env\": [\n@@ -300,7 +316,7 @@\n                                         \"protocol\": \"TCP\",\n                                     }\n                                 ],\n-                                \"resources\": {},\n+                                \"resources\": self.batch_engine_config.resources,\n                                 \"securityContext\": {\n                                     \"allowPrivilegeEscalation\": False,\n                                     \"capabilities\": {\n", "issue": "Add Kubernetes Deployment Options to the Bytewax Materialization Engine\n**Is your feature request related to a problem? Please describe.**\r\nThe Bytewax materialization engine needs to support configuration options for more advanced Kubernetes deployments (EKS, GCP, etc) to make it usable at scale.\r\n\r\n**Describe the solution you'd like**\r\nThe main configuration options that are needed for the Bytewax materialization job are:\r\n\r\n* setting explicit resource requests and limits (rather than relying on platform defaults which may not be enough)\r\n* supporting service accounts and IAM roles\r\n* specifying an image pull secret to support pulling Docker images from Dockerhub, Artifactory, etc\r\n\r\n**Describe alternatives you've considered**\r\nThe Kubernetes job that runs is dynamically generated by the bytewax code. Existing configuration options are insufficient.\r\n\r\n**Additional context**\r\nI'd really like to test the Bytewax materialization engine on our instance of EKS. In its current implementation, it's operationally too simple and not usable outside of minikube.\r\n\n", "before_files": [{"content": "import uuid\nfrom datetime import datetime\nfrom typing import Callable, List, Literal, Sequence, Union\n\nimport yaml\nfrom kubernetes import client\nfrom kubernetes import config as k8s_config\nfrom kubernetes import utils\nfrom kubernetes.utils import FailToCreateError\nfrom pydantic import StrictStr\nfrom tqdm import tqdm\n\nfrom feast import FeatureView, RepoConfig\nfrom feast.batch_feature_view import BatchFeatureView\nfrom feast.entity import Entity\nfrom feast.infra.materialization.batch_materialization_engine import (\n    BatchMaterializationEngine,\n    MaterializationJob,\n    MaterializationTask,\n)\nfrom feast.infra.offline_stores.offline_store import OfflineStore\nfrom feast.infra.online_stores.online_store import OnlineStore\nfrom feast.infra.registry.base_registry import BaseRegistry\nfrom feast.repo_config import FeastConfigBaseModel\nfrom feast.stream_feature_view import StreamFeatureView\nfrom feast.utils import _get_column_names, get_default_yaml_file_path\n\nfrom .bytewax_materialization_job import BytewaxMaterializationJob\n\n\nclass BytewaxMaterializationEngineConfig(FeastConfigBaseModel):\n    \"\"\"Batch Materialization Engine config for Bytewax\"\"\"\n\n    type: Literal[\"bytewax\"] = \"bytewax\"\n    \"\"\" Materialization type selector\"\"\"\n\n    namespace: StrictStr = \"default\"\n    \"\"\" (optional) The namespace in Kubernetes to use when creating services, configuration maps and jobs.\n    \"\"\"\n\n    image: StrictStr = \"bytewax/bytewax-feast:latest\"\n    \"\"\" (optional) The container image to use when running the materialization job.\"\"\"\n\n    env: List[dict] = []\n    \"\"\" (optional) A list of environment variables to set in the created Kubernetes pods.\n    These environment variables can be used to reference Kubernetes secrets.\n    \"\"\"\n\n\nclass BytewaxMaterializationEngine(BatchMaterializationEngine):\n    def __init__(\n        self,\n        *,\n        repo_config: RepoConfig,\n        offline_store: OfflineStore,\n        online_store: OnlineStore,\n        **kwargs,\n    ):\n        super().__init__(\n            repo_config=repo_config,\n            offline_store=offline_store,\n            online_store=online_store,\n            **kwargs,\n        )\n        self.repo_config = repo_config\n        self.offline_store = offline_store\n        self.online_store = online_store\n\n        # TODO: Configure k8s here\n        k8s_config.load_kube_config()\n\n        self.k8s_client = client.api_client.ApiClient()\n        self.v1 = client.CoreV1Api(self.k8s_client)\n        self.batch_v1 = client.BatchV1Api(self.k8s_client)\n        self.batch_engine_config = repo_config.batch_engine\n        self.namespace = self.batch_engine_config.namespace\n\n    def update(\n        self,\n        project: str,\n        views_to_delete: Sequence[\n            Union[BatchFeatureView, StreamFeatureView, FeatureView]\n        ],\n        views_to_keep: Sequence[\n            Union[BatchFeatureView, StreamFeatureView, FeatureView]\n        ],\n        entities_to_delete: Sequence[Entity],\n        entities_to_keep: Sequence[Entity],\n    ):\n        \"\"\"This method ensures that any necessary infrastructure or resources needed by the\n        engine are set up ahead of materialization.\"\"\"\n        pass\n\n    def teardown_infra(\n        self,\n        project: str,\n        fvs: Sequence[Union[BatchFeatureView, StreamFeatureView, FeatureView]],\n        entities: Sequence[Entity],\n    ):\n        \"\"\"This method ensures that any infrastructure or resources set up by ``update()``are torn down.\"\"\"\n        pass\n\n    def materialize(\n        self,\n        registry: BaseRegistry,\n        tasks: List[MaterializationTask],\n    ) -> List[MaterializationJob]:\n        return [\n            self._materialize_one(\n                registry,\n                task.feature_view,\n                task.start_time,\n                task.end_time,\n                task.project,\n                task.tqdm_builder,\n            )\n            for task in tasks\n        ]\n\n    def _materialize_one(\n        self,\n        registry: BaseRegistry,\n        feature_view: Union[BatchFeatureView, StreamFeatureView, FeatureView],\n        start_date: datetime,\n        end_date: datetime,\n        project: str,\n        tqdm_builder: Callable[[int], tqdm],\n    ):\n        entities = []\n        for entity_name in feature_view.entities:\n            entities.append(registry.get_entity(entity_name, project))\n\n        (\n            join_key_columns,\n            feature_name_columns,\n            timestamp_field,\n            created_timestamp_column,\n        ) = _get_column_names(feature_view, entities)\n\n        offline_job = self.offline_store.pull_latest_from_table_or_query(\n            config=self.repo_config,\n            data_source=feature_view.batch_source,\n            join_key_columns=join_key_columns,\n            feature_name_columns=feature_name_columns,\n            timestamp_field=timestamp_field,\n            created_timestamp_column=created_timestamp_column,\n            start_date=start_date,\n            end_date=end_date,\n        )\n\n        paths = offline_job.to_remote_storage()\n        job_id = str(uuid.uuid4())\n        return self._create_kubernetes_job(job_id, paths, feature_view)\n\n    def _create_kubernetes_job(self, job_id, paths, feature_view):\n        try:\n            # Create a k8s configmap with information needed by bytewax\n            self._create_configuration_map(job_id, paths, feature_view, self.namespace)\n\n            # Create the k8s job definition\n            self._create_job_definition(\n                job_id,\n                self.namespace,\n                len(paths),  # Create a pod for each parquet file\n                self.batch_engine_config.env,\n            )\n        except FailToCreateError as failures:\n            return BytewaxMaterializationJob(job_id, self.namespace, error=failures)\n\n        return BytewaxMaterializationJob(job_id, self.namespace)\n\n    def _create_configuration_map(self, job_id, paths, feature_view, namespace):\n        \"\"\"Create a Kubernetes configmap for this job\"\"\"\n\n        repo_path = self.repo_config.repo_path\n        assert repo_path\n        feature_store_path = get_default_yaml_file_path(repo_path)\n        feature_store_configuration = feature_store_path.read_text()\n\n        materialization_config = yaml.dump(\n            {\"paths\": paths, \"feature_view\": feature_view.name}\n        )\n\n        configmap_manifest = {\n            \"kind\": \"ConfigMap\",\n            \"apiVersion\": \"v1\",\n            \"metadata\": {\n                \"name\": f\"feast-{job_id}\",\n            },\n            \"data\": {\n                \"feature_store.yaml\": feature_store_configuration,\n                \"bytewax_materialization_config.yaml\": materialization_config,\n            },\n        }\n        self.v1.create_namespaced_config_map(\n            namespace=namespace,\n            body=configmap_manifest,\n        )\n\n    def _create_job_definition(self, job_id, namespace, pods, env):\n        \"\"\"Create a kubernetes job definition.\"\"\"\n        job_env = [\n            {\"name\": \"RUST_BACKTRACE\", \"value\": \"full\"},\n            {\n                \"name\": \"BYTEWAX_PYTHON_FILE_PATH\",\n                \"value\": \"/bytewax/dataflow.py\",\n            },\n            {\"name\": \"BYTEWAX_WORKDIR\", \"value\": \"/bytewax\"},\n            {\n                \"name\": \"BYTEWAX_WORKERS_PER_PROCESS\",\n                \"value\": \"1\",\n            },\n            {\n                \"name\": \"BYTEWAX_POD_NAME\",\n                \"valueFrom\": {\n                    \"fieldRef\": {\n                        \"apiVersion\": \"v1\",\n                        \"fieldPath\": \"metadata.annotations['batch.kubernetes.io/job-completion-index']\",\n                    }\n                },\n            },\n            {\n                \"name\": \"BYTEWAX_REPLICAS\",\n                \"value\": f\"{pods}\",\n            },\n            {\n                \"name\": \"BYTEWAX_KEEP_CONTAINER_ALIVE\",\n                \"value\": \"false\",\n            },\n            {\n                \"name\": \"BYTEWAX_STATEFULSET_NAME\",\n                \"value\": f\"dataflow-{job_id}\",\n            },\n        ]\n        # Add any Feast configured environment variables\n        job_env.extend(env)\n\n        job_definition = {\n            \"apiVersion\": \"batch/v1\",\n            \"kind\": \"Job\",\n            \"metadata\": {\n                \"name\": f\"dataflow-{job_id}\",\n                \"namespace\": namespace,\n            },\n            \"spec\": {\n                \"ttlSecondsAfterFinished\": 3600,\n                \"completions\": pods,\n                \"parallelism\": pods,\n                \"completionMode\": \"Indexed\",\n                \"template\": {\n                    \"spec\": {\n                        \"restartPolicy\": \"Never\",\n                        \"subdomain\": f\"dataflow-{job_id}\",\n                        \"initContainers\": [\n                            {\n                                \"env\": [\n                                    {\n                                        \"name\": \"BYTEWAX_REPLICAS\",\n                                        \"value\": f\"{pods}\",\n                                    }\n                                ],\n                                \"image\": \"busybox\",\n                                \"imagePullPolicy\": \"Always\",\n                                \"name\": \"init-hostfile\",\n                                \"resources\": {},\n                                \"securityContext\": {\n                                    \"allowPrivilegeEscalation\": False,\n                                    \"capabilities\": {\n                                        \"add\": [\"NET_BIND_SERVICE\"],\n                                        \"drop\": [\"ALL\"],\n                                    },\n                                    \"readOnlyRootFilesystem\": True,\n                                },\n                                \"terminationMessagePath\": \"/dev/termination-log\",\n                                \"terminationMessagePolicy\": \"File\",\n                                \"volumeMounts\": [\n                                    {\"mountPath\": \"/etc/bytewax\", \"name\": \"hostfile\"},\n                                    {\n                                        \"mountPath\": \"/tmp/bytewax/\",\n                                        \"name\": \"python-files\",\n                                    },\n                                    {\n                                        \"mountPath\": \"/var/feast/\",\n                                        \"name\": f\"feast-{job_id}\",\n                                    },\n                                ],\n                            }\n                        ],\n                        \"containers\": [\n                            {\n                                \"command\": [\"sh\", \"-c\", \"sh ./entrypoint.sh\"],\n                                \"env\": job_env,\n                                \"image\": self.batch_engine_config.image,\n                                \"imagePullPolicy\": \"Always\",\n                                \"name\": \"process\",\n                                \"ports\": [\n                                    {\n                                        \"containerPort\": 9999,\n                                        \"name\": \"process\",\n                                        \"protocol\": \"TCP\",\n                                    }\n                                ],\n                                \"resources\": {},\n                                \"securityContext\": {\n                                    \"allowPrivilegeEscalation\": False,\n                                    \"capabilities\": {\n                                        \"add\": [\"NET_BIND_SERVICE\"],\n                                        \"drop\": [\"ALL\"],\n                                    },\n                                    \"readOnlyRootFilesystem\": False,\n                                },\n                                \"terminationMessagePath\": \"/dev/termination-log\",\n                                \"terminationMessagePolicy\": \"File\",\n                                \"volumeMounts\": [\n                                    {\"mountPath\": \"/etc/bytewax\", \"name\": \"hostfile\"},\n                                    {\n                                        \"mountPath\": \"/var/feast/\",\n                                        \"name\": f\"feast-{job_id}\",\n                                    },\n                                ],\n                            }\n                        ],\n                        \"volumes\": [\n                            {\"emptyDir\": {}, \"name\": \"hostfile\"},\n                            {\n                                \"configMap\": {\n                                    \"defaultMode\": 420,\n                                    \"name\": f\"feast-{job_id}\",\n                                },\n                                \"name\": \"python-files\",\n                            },\n                            {\n                                \"configMap\": {\"name\": f\"feast-{job_id}\"},\n                                \"name\": f\"feast-{job_id}\",\n                            },\n                        ],\n                    }\n                },\n            },\n        }\n        utils.create_from_dict(self.k8s_client, job_definition)\n", "path": "sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_engine.py"}], "after_files": [{"content": "import uuid\nfrom datetime import datetime\nfrom typing import Callable, List, Literal, Sequence, Union\n\nimport yaml\nfrom kubernetes import client\nfrom kubernetes import config as k8s_config\nfrom kubernetes import utils\nfrom kubernetes.utils import FailToCreateError\nfrom pydantic import StrictStr\nfrom tqdm import tqdm\n\nfrom feast import FeatureView, RepoConfig\nfrom feast.batch_feature_view import BatchFeatureView\nfrom feast.entity import Entity\nfrom feast.infra.materialization.batch_materialization_engine import (\n    BatchMaterializationEngine,\n    MaterializationJob,\n    MaterializationTask,\n)\nfrom feast.infra.offline_stores.offline_store import OfflineStore\nfrom feast.infra.online_stores.online_store import OnlineStore\nfrom feast.infra.registry.base_registry import BaseRegistry\nfrom feast.repo_config import FeastConfigBaseModel\nfrom feast.stream_feature_view import StreamFeatureView\nfrom feast.utils import _get_column_names, get_default_yaml_file_path\n\nfrom .bytewax_materialization_job import BytewaxMaterializationJob\n\n\nclass BytewaxMaterializationEngineConfig(FeastConfigBaseModel):\n    \"\"\"Batch Materialization Engine config for Bytewax\"\"\"\n\n    type: Literal[\"bytewax\"] = \"bytewax\"\n    \"\"\" Materialization type selector\"\"\"\n\n    namespace: StrictStr = \"default\"\n    \"\"\" (optional) The namespace in Kubernetes to use when creating services, configuration maps and jobs.\n    \"\"\"\n\n    image: StrictStr = \"bytewax/bytewax-feast:latest\"\n    \"\"\" (optional) The container image to use when running the materialization job.\"\"\"\n\n    env: List[dict] = []\n    \"\"\" (optional) A list of environment variables to set in the created Kubernetes pods.\n    These environment variables can be used to reference Kubernetes secrets.\n    \"\"\"\n\n    image_pull_secrets: List[str] = []\n    \"\"\" (optional) The secrets to use when pulling the image to run for the materialization job \"\"\"\n\n    resources: dict = {}\n    \"\"\" (optional) The resource requests and limits for the materialization containers \"\"\"\n\n    service_account_name: StrictStr = \"\"\n    \"\"\" (optional) The service account name to use when running the job \"\"\"\n\n    annotations: dict = {}\n    \"\"\" (optional) Annotations to apply to the job container. Useful for linking the service account to IAM roles, operational metadata, etc  \"\"\"\n\nclass BytewaxMaterializationEngine(BatchMaterializationEngine):\n    def __init__(\n        self,\n        *,\n        repo_config: RepoConfig,\n        offline_store: OfflineStore,\n        online_store: OnlineStore,\n        **kwargs,\n    ):\n        super().__init__(\n            repo_config=repo_config,\n            offline_store=offline_store,\n            online_store=online_store,\n            **kwargs,\n        )\n        self.repo_config = repo_config\n        self.offline_store = offline_store\n        self.online_store = online_store\n\n        # TODO: Configure k8s here\n        k8s_config.load_kube_config()\n\n        self.k8s_client = client.api_client.ApiClient()\n        self.v1 = client.CoreV1Api(self.k8s_client)\n        self.batch_v1 = client.BatchV1Api(self.k8s_client)\n        self.batch_engine_config = repo_config.batch_engine\n        self.namespace = self.batch_engine_config.namespace\n\n    def update(\n        self,\n        project: str,\n        views_to_delete: Sequence[\n            Union[BatchFeatureView, StreamFeatureView, FeatureView]\n        ],\n        views_to_keep: Sequence[\n            Union[BatchFeatureView, StreamFeatureView, FeatureView]\n        ],\n        entities_to_delete: Sequence[Entity],\n        entities_to_keep: Sequence[Entity],\n    ):\n        \"\"\"This method ensures that any necessary infrastructure or resources needed by the\n        engine are set up ahead of materialization.\"\"\"\n        pass\n\n    def teardown_infra(\n        self,\n        project: str,\n        fvs: Sequence[Union[BatchFeatureView, StreamFeatureView, FeatureView]],\n        entities: Sequence[Entity],\n    ):\n        \"\"\"This method ensures that any infrastructure or resources set up by ``update()``are torn down.\"\"\"\n        pass\n\n    def materialize(\n        self,\n        registry: BaseRegistry,\n        tasks: List[MaterializationTask],\n    ) -> List[MaterializationJob]:\n        return [\n            self._materialize_one(\n                registry,\n                task.feature_view,\n                task.start_time,\n                task.end_time,\n                task.project,\n                task.tqdm_builder,\n            )\n            for task in tasks\n        ]\n\n    def _materialize_one(\n        self,\n        registry: BaseRegistry,\n        feature_view: Union[BatchFeatureView, StreamFeatureView, FeatureView],\n        start_date: datetime,\n        end_date: datetime,\n        project: str,\n        tqdm_builder: Callable[[int], tqdm],\n    ):\n        entities = []\n        for entity_name in feature_view.entities:\n            entities.append(registry.get_entity(entity_name, project))\n\n        (\n            join_key_columns,\n            feature_name_columns,\n            timestamp_field,\n            created_timestamp_column,\n        ) = _get_column_names(feature_view, entities)\n\n        offline_job = self.offline_store.pull_latest_from_table_or_query(\n            config=self.repo_config,\n            data_source=feature_view.batch_source,\n            join_key_columns=join_key_columns,\n            feature_name_columns=feature_name_columns,\n            timestamp_field=timestamp_field,\n            created_timestamp_column=created_timestamp_column,\n            start_date=start_date,\n            end_date=end_date,\n        )\n\n        paths = offline_job.to_remote_storage()\n        job_id = str(uuid.uuid4())\n        return self._create_kubernetes_job(job_id, paths, feature_view)\n\n    def _create_kubernetes_job(self, job_id, paths, feature_view):\n        try:\n            # Create a k8s configmap with information needed by bytewax\n            self._create_configuration_map(job_id, paths, feature_view, self.namespace)\n\n            # Create the k8s job definition\n            self._create_job_definition(\n                job_id,\n                self.namespace,\n                len(paths),  # Create a pod for each parquet file\n                self.batch_engine_config.env,\n            )\n        except FailToCreateError as failures:\n            return BytewaxMaterializationJob(job_id, self.namespace, error=failures)\n\n        return BytewaxMaterializationJob(job_id, self.namespace)\n\n    def _create_configuration_map(self, job_id, paths, feature_view, namespace):\n        \"\"\"Create a Kubernetes configmap for this job\"\"\"\n\n        repo_path = self.repo_config.repo_path\n        assert repo_path\n        feature_store_path = get_default_yaml_file_path(repo_path)\n        feature_store_configuration = feature_store_path.read_text()\n\n        materialization_config = yaml.dump(\n            {\"paths\": paths, \"feature_view\": feature_view.name}\n        )\n\n        configmap_manifest = {\n            \"kind\": \"ConfigMap\",\n            \"apiVersion\": \"v1\",\n            \"metadata\": {\n                \"name\": f\"feast-{job_id}\",\n            },\n            \"data\": {\n                \"feature_store.yaml\": feature_store_configuration,\n                \"bytewax_materialization_config.yaml\": materialization_config,\n            },\n        }\n        self.v1.create_namespaced_config_map(\n            namespace=namespace,\n            body=configmap_manifest,\n        )\n\n    def _create_job_definition(self, job_id, namespace, pods, env):\n        \"\"\"Create a kubernetes job definition.\"\"\"\n        job_env = [\n            {\"name\": \"RUST_BACKTRACE\", \"value\": \"full\"},\n            {\n                \"name\": \"BYTEWAX_PYTHON_FILE_PATH\",\n                \"value\": \"/bytewax/dataflow.py\",\n            },\n            {\"name\": \"BYTEWAX_WORKDIR\", \"value\": \"/bytewax\"},\n            {\n                \"name\": \"BYTEWAX_WORKERS_PER_PROCESS\",\n                \"value\": \"1\",\n            },\n            {\n                \"name\": \"BYTEWAX_POD_NAME\",\n                \"valueFrom\": {\n                    \"fieldRef\": {\n                        \"apiVersion\": \"v1\",\n                        \"fieldPath\": \"metadata.annotations['batch.kubernetes.io/job-completion-index']\",\n                    }\n                },\n            },\n            {\n                \"name\": \"BYTEWAX_REPLICAS\",\n                \"value\": f\"{pods}\",\n            },\n            {\n                \"name\": \"BYTEWAX_KEEP_CONTAINER_ALIVE\",\n                \"value\": \"false\",\n            },\n            {\n                \"name\": \"BYTEWAX_STATEFULSET_NAME\",\n                \"value\": f\"dataflow-{job_id}\",\n            },\n        ]\n        # Add any Feast configured environment variables\n        job_env.extend(env)\n\n        job_definition = {\n            \"apiVersion\": \"batch/v1\",\n            \"kind\": \"Job\",\n            \"metadata\": {\n                \"name\": f\"dataflow-{job_id}\",\n                \"namespace\": namespace,\n            },\n            \"spec\": {\n                \"ttlSecondsAfterFinished\": 3600,\n                \"completions\": pods,\n                \"parallelism\": pods,\n                \"completionMode\": \"Indexed\",\n                \"template\": {\n                    \"metadata\": {\n                        \"annotations\": self.batch_engine_config.annotations,\n                    },\n                    \"spec\": {\n                        \"restartPolicy\": \"Never\",\n                        \"subdomain\": f\"dataflow-{job_id}\",\n                        \"imagePullSecrets\": self.batch_engine_config.image_pull_secrets,\n                        \"serviceAccountName\": self.batch_engine_config.service_account_name,\n                        \"initContainers\": [\n                            {\n                                \"env\": [\n                                    {\n                                        \"name\": \"BYTEWAX_REPLICAS\",\n                                        \"value\": f\"{pods}\",\n                                    }\n                                ],\n                                \"image\": \"busybox\",\n                                \"imagePullPolicy\": \"Always\",\n                                \"name\": \"init-hostfile\",\n                                \"resources\": {},\n                                \"securityContext\": {\n                                    \"allowPrivilegeEscalation\": False,\n                                    \"capabilities\": {\n                                        \"add\": [\"NET_BIND_SERVICE\"],\n                                        \"drop\": [\"ALL\"],\n                                    },\n                                    \"readOnlyRootFilesystem\": True,\n                                },\n                                \"terminationMessagePath\": \"/dev/termination-log\",\n                                \"terminationMessagePolicy\": \"File\",\n                                \"volumeMounts\": [\n                                    {\"mountPath\": \"/etc/bytewax\", \"name\": \"hostfile\"},\n                                    {\n                                        \"mountPath\": \"/tmp/bytewax/\",\n                                        \"name\": \"python-files\",\n                                    },\n                                    {\n                                        \"mountPath\": \"/var/feast/\",\n                                        \"name\": f\"feast-{job_id}\",\n                                    },\n                                ],\n                            }\n                        ],\n                        \"containers\": [\n                            {\n                                \"command\": [\"sh\", \"-c\", \"sh ./entrypoint.sh\"],\n                                \"env\": job_env,\n                                \"image\": self.batch_engine_config.image,\n                                \"imagePullPolicy\": \"Always\",\n                                \"name\": \"process\",\n                                \"ports\": [\n                                    {\n                                        \"containerPort\": 9999,\n                                        \"name\": \"process\",\n                                        \"protocol\": \"TCP\",\n                                    }\n                                ],\n                                \"resources\": self.batch_engine_config.resources,\n                                \"securityContext\": {\n                                    \"allowPrivilegeEscalation\": False,\n                                    \"capabilities\": {\n                                        \"add\": [\"NET_BIND_SERVICE\"],\n                                        \"drop\": [\"ALL\"],\n                                    },\n                                    \"readOnlyRootFilesystem\": False,\n                                },\n                                \"terminationMessagePath\": \"/dev/termination-log\",\n                                \"terminationMessagePolicy\": \"File\",\n                                \"volumeMounts\": [\n                                    {\"mountPath\": \"/etc/bytewax\", \"name\": \"hostfile\"},\n                                    {\n                                        \"mountPath\": \"/var/feast/\",\n                                        \"name\": f\"feast-{job_id}\",\n                                    },\n                                ],\n                            }\n                        ],\n                        \"volumes\": [\n                            {\"emptyDir\": {}, \"name\": \"hostfile\"},\n                            {\n                                \"configMap\": {\n                                    \"defaultMode\": 420,\n                                    \"name\": f\"feast-{job_id}\",\n                                },\n                                \"name\": \"python-files\",\n                            },\n                            {\n                                \"configMap\": {\"name\": f\"feast-{job_id}\"},\n                                \"name\": f\"feast-{job_id}\",\n                            },\n                        ],\n                    }\n                },\n            },\n        }\n        utils.create_from_dict(self.k8s_client, job_definition)\n", "path": "sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_engine.py"}]}