SAASMLOPS-809 fix bytewax workers so they only process a single file (#6)

* SAASMLOPS-809 fix bytewax workers so they only process a single file

* SAASMLOPS-809 fix newlines

Signed-off-by: James Crabtree <james.crabtree@sailpoint.com>

Author: james-crabtree-sp

sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_dataflow.py

@@ -5,32 +5,37 @@
 import s3fs
 from bytewax.dataflow import Dataflow  # type: ignore
 from bytewax.execution import cluster_main
-from bytewax.inputs import ManualInputConfig, distribute
+from bytewax.inputs import ManualInputConfig
 from bytewax.outputs import ManualOutputConfig
 from tqdm import tqdm
+import logging
 
 from feast import FeatureStore, FeatureView, RepoConfig
 from feast.utils import _convert_arrow_to_proto, _run_pyarrow_field_mapping
 
+logger = logging.getLogger(__name__)
+
 
 class BytewaxMaterializationDataflow:
     def __init__(
         self,
         config: RepoConfig,
         feature_view: FeatureView,
         paths: List[str],
+        worker_index: int
     ):
         self.config = config
         self.feature_store = FeatureStore(config=config)
 
         self.feature_view = feature_view
+        self.worker_index = worker_index
         self.paths = paths
 
         self._run_dataflow()
 
     def process_path(self, path):
         fs = s3fs.S3FileSystem()
-        print(f"Processing path {path}")
+        logger.info(f"Processing path {path}")
         dataset = pq.ParquetDataset(path, filesystem=fs, use_legacy_dataset=False)
         batches = []
         for fragment in dataset.fragments:
@@ -40,11 +45,7 @@ def process_path(self, path):
         return batches
 
     def input_builder(self, worker_index, worker_count, _state):
-        worker_paths = distribute(self.paths, worker_index, worker_count)
-        for path in worker_paths:
-            yield None, path
-
-        return
+        return [(None, self.paths[self.worker_index])]
 
     def output_builder(self, worker_index, worker_count):
         def output_fn(batch):

sdk/python/feast/infra/materialization/contrib/bytewax/bytewax_materialization_engine.py

@@ -202,7 +202,7 @@ def _materialize_one(
                 try:
                     self.batch_v1.delete_namespaced_job(job.job_id(), self.namespace)
                 except ApiException as de:
-                    logger.warning(f"Could not delete job due to API Error: {ae.body}")
+                    logger.warning(f"Could not delete job due to API Error: {de.body}")
                 raise e
             self._print_pod_logs(job.job_id(), feature_view)
         return job
@@ -290,7 +290,7 @@ def _create_job_definition(self, job_id, namespace, pods, env):
             },
             {
                 "name": "BYTEWAX_REPLICAS",
-                "value": "1",
+                "value": f"{pods}",
             },
             {
                 "name": "BYTEWAX_KEEP_CONTAINER_ALIVE",
@@ -324,8 +324,8 @@ def _create_job_definition(self, job_id, namespace, pods, env):
             "spec": {
                 "ttlSecondsAfterFinished": 3600,
                 "backoffLimit": self.batch_engine_config.retry_limit,
-                "completions": 1,
-                "parallelism": 1,
+                "completions": pods,
+                "parallelism": min(pods, self.batch_engine_config.max_parallelism),
                 "completionMode": "Indexed",
                 "template": {
                     "metadata": {
@@ -342,7 +342,7 @@ def _create_job_definition(self, job_id, namespace, pods, env):
                                 "env": [
                                     {
                                         "name": "BYTEWAX_REPLICAS",
-                                        "value": "1",
+                                        "value": f"{pods}",
                                     }
                                 ],
                                 "image": "busybox",

sdk/python/feast/infra/materialization/contrib/bytewax/dataflow.py

@@ -1,3 +1,4 @@
+import os
 import yaml
 
 from feast import FeatureStore, RepoConfig
@@ -19,4 +20,5 @@
                 config,
                 store.get_feature_view(bytewax_config["feature_view"]),
                 bytewax_config["paths"],
+                int(os.environ["JOB_COMPLETION_INDEX"])
             )