fix: Refactor/minor cleanup (#334)

* fix: Make max commands per pipe configurable

* fix: Add same changes to redis

* fix: type issue

* fix: typing

* fix: typing

* fix: typing

---------

Co-authored-by: Krishna mohan Pulipati <kpulipati@expediagroup.com>

Author: kpulipati29

sdk/python/feast/infra/online_stores/eg_valkey.py

@@ -14,6 +14,7 @@
 import json
 import logging
 import math
+import os
 import time
 from datetime import datetime, timedelta, timezone
 from enum import Enum
@@ -96,6 +97,9 @@ class EGValkeyOnlineStoreConfig(FeastConfigBaseModel):
     read_batch_size: Optional[int] = 100
     """(Optional) number of keys to read in a single batch for online read requests. Anything < 1 means no batching."""
 
+    max_pipeline_commands: Optional[int] = 500
+    """(Optional) The maximum number of Valkey commands to queue in a pipeline before sending them to Valkey in a single batch."""
+
 
 class EGValkeyOnlineStore(OnlineStore):
     """
@@ -327,10 +331,14 @@ def online_write_batch(
                     )
 
                 sort_key_name = table.sort_keys[0].name
+
                 num_cmds = 0
-                # Picking an arbitrary number to start with and will be tuned after perf testing
-                # TODO : Make this a config as this can be different for different users based on payload size etc..
-                num_cmds_per_pipeline_execute = 500
+                max_pipeline_commands_per_process = (
+                    EGValkeyOnlineStore._get_max_pipeline_commands_per_process(
+                        online_store_config.max_pipeline_commands
+                    )
+                )
+
                 ttl_feature_view = table.ttl
 
                 for entity_key, values, timestamp, _ in data:
@@ -381,7 +389,7 @@ def online_write_batch(
                         pipe.expire(name=hash_key, time=ttl)
                         num_cmds += 1
 
-                    if num_cmds >= num_cmds_per_pipeline_execute:
+                    if num_cmds >= max_pipeline_commands_per_process:
                         # TODO: May be add retries with backoff
                         try:
                             results = pipe.execute()  # flush
@@ -414,14 +422,13 @@ def online_write_batch(
                 # AFTER batch flush: run TTL cleanup
                 if run_cleanup_by_event_time and ttl_feature_view_seconds:
                     cleanup_cmds = 0
-                    cleanup_cmds_per_execute = 500
                     cutoff = (int(time.time()) - ttl_feature_view_seconds) * 1000
                     for zset_key, entity_key_bytes in zsets_to_cleanup:
                         self._run_cleanup_by_event_time(
                             pipe, zset_key, ttl_feature_view_seconds, cutoff
                         )
                         cleanup_cmds += 2
-                        if cleanup_cmds >= cleanup_cmds_per_execute:
+                        if cleanup_cmds >= max_pipeline_commands_per_process:
                             try:
                                 pipe.execute()
                             except ValkeyError:
@@ -539,6 +546,17 @@ def _get_ttl(
         ttl_remaining = timestamp - utils._utc_now() + ttl_offset
         return math.ceil(ttl_remaining.total_seconds())
 
+    @staticmethod
+    def _get_max_pipeline_commands_per_process(
+        max_pipeline_commands: int | None,
+    ) -> int:
+        assert max_pipeline_commands is not None
+        num_processes = int(os.environ.get("NUM_PROCESSES", 1))
+        max_pipeline_commands_per_process = max(
+            1, math.ceil(max_pipeline_commands / num_processes)
+        )
+        return max_pipeline_commands_per_process
+
     def _run_cleanup_by_event_time(
         self, pipe, zset_key: bytes, ttl_seconds: int, cutoff
     ):

sdk/python/feast/infra/online_stores/redis.py

@@ -14,6 +14,7 @@
 import json
 import logging
 import math
+import os
 import time
 from datetime import datetime, timedelta, timezone
 from enum import Enum
@@ -88,6 +89,9 @@ class RedisOnlineStoreConfig(FeastConfigBaseModel):
     read_batch_size: Optional[int] = 100
     """(Optional) number of keys to read in a single batch for online read requests. Anything < 1 means no batching."""
 
+    max_pipeline_commands: Optional[int] = 500
+    """(Optional) The maximum number of Redis commands to queue in a pipeline before sending them to Redis in a single batch.In most cases, limiting the size of the pipeline to 100-1000 operations per shard gives the best results."""
+
 
 class RedisOnlineStore(OnlineStore):
     """
@@ -321,9 +325,12 @@ def online_write_batch(
 
                 sort_key_name = table.sort_keys[0].name
                 num_cmds = 0
-                # Picking an arbitrary number to start with and will be tuned after perf testing
-                # TODO : Make this a config as this can be different for different users based on payload size etc..
-                num_cmds_per_pipeline_execute = 500
+                max_pipeline_commands_per_process = (
+                    RedisOnlineStore._get_max_pipeline_commands_per_process(
+                        online_store_config.max_pipeline_commands
+                    )
+                )
+
                 ttl_feature_view = table.ttl
 
                 for entity_key, values, timestamp, _ in data:
@@ -374,7 +381,7 @@ def online_write_batch(
                         pipe.expire(name=hash_key, time=ttl)
                         num_cmds += 1
 
-                    if num_cmds >= num_cmds_per_pipeline_execute:
+                    if num_cmds >= max_pipeline_commands_per_process:
                         # TODO: May be add retries with backoff
                         try:
                             results = pipe.execute()  # flush
@@ -407,14 +414,13 @@ def online_write_batch(
                 # AFTER batch flush: run TTL cleanup
                 if run_cleanup_by_event_time and ttl_feature_view_seconds:
                     cleanup_cmds = 0
-                    cleanup_cmds_per_execute = 500
                     cutoff = (int(time.time()) - ttl_feature_view_seconds) * 1000
                     for zset_key, entity_key_bytes in zsets_to_cleanup:
                         self._run_cleanup_by_event_time(
                             pipe, zset_key, ttl_feature_view_seconds, cutoff
                         )
                         cleanup_cmds += 2
-                        if cleanup_cmds >= cleanup_cmds_per_execute:
+                        if cleanup_cmds >= max_pipeline_commands_per_process:
                             try:
                                 pipe.execute()
                             except RedisError:
@@ -518,6 +524,17 @@ def sort_key_bytes(sort_key_name: str, sort_val: ValueProto, v: int = 3) -> byte
         sk = EntityKeyProto(join_keys=[sort_key_name], entity_values=[sort_val])
         return serialize_entity_key(sk, entity_key_serialization_version=v)
 
+    @staticmethod
+    def _get_max_pipeline_commands_per_process(
+        max_pipeline_commands: int | None,
+    ) -> int:
+        assert max_pipeline_commands is not None
+        num_processes = int(os.environ.get("NUM_PROCESSES", 1))
+        max_pipeline_commands_per_process = max(
+            1, math.ceil(max_pipeline_commands / num_processes)
+        )
+        return max_pipeline_commands_per_process
+
     @staticmethod
     def _get_ttl(
         ttl_feature_view: timedelta,

sdk/python/feast/infra/passthrough_provider.py

@@ -403,6 +403,8 @@ def ingest_df(
             if table.num_rows < num_processes:
                 num_processes = table.num_rows
 
+            os.environ["NUM_PROCESSES"] = str(num_processes)
+
             # Input table is split into smaller chunks and processed in parallel
             chunks = self.split_table(num_processes, table)
 
@@ -413,6 +415,7 @@ def ingest_df(
             with Pool(processes=num_processes) as pool:
                 pool.starmap(self.process, chunks_to_parallelize)
         else:
+            os.environ["NUM_PROCESSES"] = "1"
             self.process(table, feature_view, join_keys)
 
     @staticmethod