chore: Upgrade to Numaflow 0.11 and Fix: Trainer Pod getting the slav…

…e redis client connection (#173)

numaprom/clients/sentinel.py

@@ -21,6 +21,7 @@ def get_redis_client(
     mastername: str,
     recreate: bool = False,
     master_node: bool = True,
+    reset: bool = False,
 ) -> redis_client_t:
     """Return a master redis client for sentinel connections, with retry.
 
@@ -40,7 +41,12 @@ def get_redis_client(
     """
     global SENTINEL_CLIENT
 
+    if reset:
+        LOGGER.info("Reset Sentinel Client to None")
+        SENTINEL_CLIENT = None
+
     if not recreate and SENTINEL_CLIENT:
+        LOGGER.info("Reusing Existing Sentinel Client")
         return SENTINEL_CLIENT
 
     retry = Retry(
@@ -70,14 +76,17 @@ def get_redis_client(
         **conn_kwargs
     )
     if master_node:
+        LOGGER.info("Creating Master Sentinel Redis Client")
         SENTINEL_CLIENT = sentinel.master_for(mastername)
     else:
+        LOGGER.info("Creating Slave Sentinel Redis Client")
         SENTINEL_CLIENT = sentinel.slave_for(mastername)
     LOGGER.info(
         "Sentinel redis params: {args}, master_node: {is_master}",
         args=conn_kwargs,
         is_master=master_node,
     )
+
     return SENTINEL_CLIENT
 
 

numaprom/factory.py

@@ -1,7 +1,7 @@
 from collections.abc import Callable
 
-from pynumaflow.function import Messages
-from pynumaflow.sink import Responses
+from pynumaflow.mapper import Messages
+from pynumaflow.sinker import Responses
 
 from numaprom.udf import preprocess, postprocess, window, inference, threshold
 from numaprom.udsink import train, train_rollout

numaprom/tools.py

@@ -10,7 +10,7 @@
 import pytz
 from numalogic.config import PostprocessFactory
 from numalogic.models.threshold import SigmoidThreshold
-from pynumaflow.function import Messages, Message
+from pynumaflow.mapper import Messages, Message
 from numaprom import LOGGER, MetricConf
 from numaprom.clients.prometheus import Prometheus
 from numaprom.entities import TrainerPayload, StreamPayload

numaprom/udf/inference.py

@@ -8,7 +8,7 @@
 from numalogic.tools.data import StreamingDataset
 from numalogic.tools.exceptions import RedisRegistryError
 from orjson import orjson
-from pynumaflow.function import Datum
+from pynumaflow.mapper import Datum
 from torch.utils.data import DataLoader
 
 from numaprom import LOGGER

numaprom/udf/postprocess.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 from orjson import orjson
-from pynumaflow.function import Datum
+from pynumaflow.mapper import Datum
 from redis.exceptions import RedisError, RedisClusterException
 from redis.sentinel import MasterNotFoundError
 

numaprom/udf/preprocess.py

@@ -5,7 +5,8 @@
 import orjson
 from numalogic.registry import RedisRegistry, LocalLRUCache
 from numalogic.tools.exceptions import RedisRegistryError
-from pynumaflow.function import Datum
+from pynumaflow.mapper import Datum
+
 
 from numaprom import LOGGER
 from numaprom.clients.sentinel import get_redis_client

numaprom/udf/threshold.py

@@ -6,7 +6,7 @@
 from numalogic.registry import RedisRegistry, LocalLRUCache
 from numalogic.tools.exceptions import RedisRegistryError
 from orjson import orjson
-from pynumaflow.function import Datum
+from pynumaflow.mapper import Datum
 
 from numaprom import LOGGER
 from numaprom._constants import TRAIN_VTX_KEY, POSTPROC_VTX_KEY

numaprom/udf/window.py

@@ -6,7 +6,7 @@
 import numpy as np
 import numpy.typing as npt
 from orjson import orjson
-from pynumaflow.function import Datum
+from pynumaflow.mapper import Datum
 from redis.exceptions import RedisError, RedisClusterException
 
 from numaprom import LOGGER

numaprom/udsink/train.py

@@ -10,7 +10,7 @@
 from numalogic.tools.exceptions import RedisRegistryError
 from numalogic.tools.types import redis_client_t
 from orjson import orjson
-from pynumaflow.sink import Datum, Responses, Response
+from pynumaflow.sinker import Datum, Responses, Response
 from sklearn.pipeline import make_pipeline
 from torch.utils.data import DataLoader
 

numaprom/udsink/train_rollout.py

@@ -12,7 +12,7 @@
 from numalogic.tools.exceptions import RedisRegistryError
 from numalogic.tools.types import redis_client_t
 from orjson import orjson
-from pynumaflow.sink import Datum, Responses, Response
+from pynumaflow.sinker import Datum, Responses, Response
 from sklearn.pipeline import make_pipeline
 from torch.utils.data import DataLoader
 
@@ -23,6 +23,8 @@
 from numaprom.watcher import ConfigManager
 
 REQUEST_EXPIRY = int(os.getenv("REQUEST_EXPIRY", "300"))
+# REDIS_CLIENT = get_redis_client_from_conf(master_node=True, recreate=True)
+REDIS_CLIENT_MASTER = get_redis_client_from_conf(master_node=True, reset=True)
 
 
 # TODO: extract all good hashes, including when there are 2 hashes at a time
@@ -107,8 +109,9 @@ def get_model_config(metric_config):
 
 
 def train_rollout(datums: Iterator[Datum]) -> Responses:
+    global REDIS_CLIENT_MASTER
+    redis_client = REDIS_CLIENT_MASTER
     responses = Responses()
-    redis_client = get_redis_client_from_conf()
 
     for _datum in datums:
         payload = TrainerPayload(**orjson.loads(_datum.value))