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

james-crabtree-sp · james-crabtree-sp · commit 2cfe0454508b · 2023-09-01T14:21:56.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
@@ -76,6 +76,7 @@
     "bowler",  # Needed for automatic repo upgrades
     # FastAPI does not correctly pull starlette dependency on httpx see thread(https://github.com/tiangolo/fastapi/issues/5656).
     "httpx>=0.23.3",
+    "pyOpenSSL"
 ]
 
 GCP_REQUIRED = [
@@ -93,7 +94,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
`@@ -76,6 +76,7 @@`
`76`	`76`	`"bowler", # Needed for automatic repo upgrades`
`77`	`77`	`# FastAPI does not correctly pull starlette dependency on httpx see thread(https://github.com/tiangolo/fastapi/issues/5656).`
`78`	`78`	`"httpx>=0.23.3",`
	`79`	`+ "pyOpenSSL"`
`79`	`80`	`]`
`80`	`81`
`81`	`82`	`GCP_REQUIRED = [`
`@@ -93,7 +94,7 @@`
`93`	`94`	`"hiredis>=2.0.0,<3",`
`94`	`95`	`]`
`95`	`96`
`96`		`-AWS_REQUIRED = ["boto3>=1.17.0,<2", "docker>=5.0.2"]`
	`97`	`+AWS_REQUIRED = ["boto3>=1.17.0,<2", "docker>=5.0.2", "s3fs"]`
`97`	`98`
`98`	`99`	`BYTEWAX_REQUIRED = ["bytewax==0.15.1", "docker>=5.0.2", "kubernetes<=20.13.0"]`
`99`	`100`