SAASMLOPS-734 bytewax in-cluster config, custom labels, fix worker image deps

james-crabtree-sp · james-crabtree-sp · commit 8985d6e3f4f3 · 2023-09-14T11:10:43.000-05:00
Signed-off-by: James Crabtree &lt;james.crabtree@sailpoint.com&gt;
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
@@ -61,6 +61,9 @@ class BytewaxMaterializationEngineConfig(FeastConfigBaseModel):
     include_security_context_capabilities: bool = True
     """ (optional)  Include security context capabilities in the init and job container spec """
 
+    labels: dict = {}
+    """ (optional) additional labels to append to kubernetes objects """
+
 
 class BytewaxMaterializationEngine(BatchMaterializationEngine):
     def __init__(
@@ -82,7 +85,7 @@ def __init__(
         self.online_store = online_store
 
         # TODO: Configure k8s here
-        k8s_config.load_kube_config()
+        k8s_config.load_config()
 
         self.k8s_client = client.api_client.ApiClient()
         self.v1 = client.CoreV1Api(self.k8s_client)
@@ -196,14 +199,13 @@ def _create_configuration_map(self, job_id, paths, feature_view, namespace):
             {"paths": paths, "feature_view": feature_view.name}
         )
 
+        labels = {"feast-bytewax-materializer": "configmap"}
         configmap_manifest = {
             "kind": "ConfigMap",
             "apiVersion": "v1",
             "metadata": {
                 "name": f"feast-{job_id}",
-                "labels": {
-                    "feast-bytewax-materializer": "configmap",
-                },
+                "labels": {**labels, **self.batch_engine_config.labels},
             },
             "data": {
                 "feature_store.yaml": feature_store_configuration,
@@ -260,15 +262,15 @@ def _create_job_definition(self, job_id, namespace, pods, env):
                 "drop": ["ALL"],
             }
 
+        job_labels = {"feast-bytewax-materializer": "job"}
+        pod_labels = {"feast-bytewax-materializer": "pod"}
         job_definition = {
             "apiVersion": "batch/v1",
             "kind": "Job",
             "metadata": {
                 "name": f"dataflow-{job_id}",
                 "namespace": namespace,
-                "labels": {
-                    "feast-bytewax-materializer": "job",
-                },
+                "labels": {**job_labels, **self.batch_engine_config.labels},
             },
             "spec": {
                 "ttlSecondsAfterFinished": 3600,
@@ -278,9 +280,7 @@ def _create_job_definition(self, job_id, namespace, pods, env):
                 "template": {
                     "metadata": {
                         "annotations": self.batch_engine_config.annotations,
-                        "labels": {
-                            "feast-bytewax-materializer": "pod",
-                        },
+                        "labels": {**pod_labels, **self.batch_engine_config.labels},
                     },
                     "spec": {
                         "restartPolicy": "Never",
diff --git a/setup.py b/setup.py
@@ -98,7 +98,7 @@
     "hiredis>=2.0.0,<3",
 ]
 
-AWS_REQUIRED = ["boto3>=1.17.0,<2", "docker>=5.0.2"]
+AWS_REQUIRED = ["boto3>=1.17.0,<2", "docker>=5.0.2", "s3fs"]
 
 BYTEWAX_REQUIRED = ["bytewax==0.15.1", "docker>=5.0.2", "kubernetes<=20.13.0"]
 

Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,7 @@`
`98`	`98`	`"hiredis>=2.0.0,<3",`
`99`	`99`	`]`
`100`	`100`
`101`		`-AWS_REQUIRED = ["boto3>=1.17.0,<2", "docker>=5.0.2"]`
	`101`	`+AWS_REQUIRED = ["boto3>=1.17.0,<2", "docker>=5.0.2", "s3fs"]`
`102`	`102`
`103`	`103`	`BYTEWAX_REQUIRED = ["bytewax==0.15.1", "docker>=5.0.2", "kubernetes<=20.13.0"]`
`104`	`104`