Merge pull request #10 from starwit/feature/AB#1745-multiple-streams

Feature/ab#1745 support multiple streams

Author: flonix8

README.md

@@ -1,6 +1,6 @@
 # SAE Redis Writer
 
-This component is part of the Starwit Awareness Engine (SAE). See umbrella repo here: https://github.com/starwit/vision-pipeline-k8s
+This component is part of the Starwit Awareness Engine (SAE). See umbrella repo here: https://github.com/starwit/starwit-awareness-engine
 
 The intention of this stage is to write received messages to a different redis/[valkey](https://valkey.io/) instance. This means data created by SAE can be transfered to a backend.
 It transparently forwards all messages it receives, regardless of their type, to the configured Redis instance.

deployment/valkey/valkey.conf

@@ -0,0 +1 @@
+protected-mode no

doc/DEV_README.md

@@ -2,7 +2,7 @@
 
 ## Docker
 
-You can use [docker_build.sh](docker_build.sh) and [docker_publish.sh](docker_publish.sh) to build and publish Docker images. 
+You can use [docker_build.sh](docker_build.sh) to build an image for local testing. 
 
 Once build you can run Docker image locally like so:
 ```bash

docker_build.sh

@@ -1,3 +1,3 @@
 #!/bin/bash
 
-docker build -t starwitorg/sae-redis-writer:$(poetry version --short) .
+docker build -t starwitorg/sae-redis-writer:local .

docker_push.sh

@@ -10,7 +10,7 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = "^3.10"
 pydantic = "^2.0.0"
-visionapi = { git = "https://github.com/starwit/vision-api.git", subdirectory = "python/visionapi", tag = "3.2.0" }
+visionapi = { git = "https://github.com/starwit/vision-api.git", subdirectory = "python/visionapi", tag = "3.4.0" }
 visionlib = { git = "https://github.com/starwit/vision-lib.git", subdirectory = "python", tag = "0.12.0" }
 redis = "^5.0.0"
 pydantic-settings = "^2.0.3"

poetry.lock

@@ -9,23 +9,25 @@
 class RedisConfig(BaseModel):
     host: str = 'localhost'
     port: Annotated[int, Field(ge=1, le=65536)] = 6379
-    input_stream_prefix: str = 'objecttracker'
 
 class TargetRedisConfig(BaseModel):
     host: str
     port: Annotated[int, Field(ge=1, le=65536)]
-    output_stream_prefix: str = 'output'
     buffer_length: Annotated[int, Field(ge=1)] = 100
     target_stream_maxlen: Annotated[int, Field(ge=1)] = 100
     tls: bool = False
+    
+class MappingConfig(BaseModel):
+    source: str = None
+    target: str = None
 
 class RedisWriterConfig(BaseSettings):
     log_level: LogLevel = LogLevel.WARNING
     redis: RedisConfig = RedisConfig()
     target_redis: TargetRedisConfig
-    stream_ids: List[str] = ['stream1']
     remove_frame_data: bool = True
     prometheus_port: Annotated[int, Field(gt=1024, le=65536)] = 8000
+    mapping_config: List[MappingConfig]
 
     model_config = SettingsConfigDict(env_nested_delimiter='__')
 

pyproject.toml

@@ -35,7 +35,6 @@ class BufferEntry(NamedTuple):
 class Sender:
     def __init__(self, config: RedisWriterConfig) -> None:
         self._config = config.target_redis
-        self._stream_ids = config.stream_ids
         logger.setLevel(config.log_level.value)
 
         self._buffer: Deque[BufferEntry] = deque(maxlen=self._config.buffer_length)

rediswriter/config.py

@@ -5,6 +5,7 @@
 from typing import Dict
 
 from prometheus_client import Counter, start_http_server
+from visionapi.common_pb2 import TypeMessage, MessageType
 from visionapi.sae_pb2 import SaeMessage
 from visionlib.pipeline.consumer import RedisConsumer
 from visionlib.pipeline.formats import is_sae_message
@@ -17,10 +18,6 @@
 
 FRAME_COUNTER = Counter('redis_writer_frame_counter', 'How many frames have been consumed from the Redis input stream')
 
-class MessageType(Enum):
-    SAE = 1
-    OTHER = 2
-
 def run_stage():
 
     stop_event = threading.Event()
@@ -46,9 +43,11 @@ def sig_handler(signum, _):
     logger.info(f'Starting redis writer stage. Config: {CONFIG.model_dump_json(indent=2)}')
 
     redis_writer = RedisWriter(CONFIG)
-
-    consumer = RedisConsumer(CONFIG.redis.host, CONFIG.redis.port, 
-                            stream_keys=[f'{CONFIG.redis.input_stream_prefix}:{id}' for id in CONFIG.stream_ids])
+    
+    stream_mapping = map_config(CONFIG.mapping_config)
+    
+    consumer = RedisConsumer(CONFIG.redis.host, CONFIG.redis.port, stream_mapping.keys())
+    logger.debug(f"Listening to stream keys {stream_mapping.keys()}")
 
     sender = Sender(CONFIG)
 
@@ -62,20 +61,23 @@ def sig_handler(signum, _):
             if stream_key is None:
                 continue
 
-            stream_id = stream_key.split(':')[1]
+            stream_id = stream_key.split(':')[1]            
 
             FRAME_COUNTER.inc()
+            logger.debug(f'Received message on stream {stream_id}')
 
             # Detect stream type by analyzing first message
-            if not stream_id in message_type_by_stream:
-                msg = SaeMessage()
+            if not stream_id in message_type_by_stream:                
+                msg = TypeMessage()
                 msg.ParseFromString(proto_data)
-                if is_sae_message(msg):
-                    msg_type = MessageType.SAE
+                if msg.type != MessageType.UNSPECIFIED:
+                    message_type_by_stream[stream_id] = msg.type
                 else:
-                    msg_type = MessageType.OTHER
-                message_type_by_stream[stream_id] = msg_type
-                logger.info(f'Detected message type {msg_type.name} on stream {stream_id}')
+                    # if message type can't be determined, messages are NOT forwarded
+                    continue
+                
+                type = MessageType.Name(msg.type)
+                logger.info(f'Detected message type {type} on stream {stream_id}')
 
             # Only process SaeMessage messages, otherwise pass verbatim
             if message_type_by_stream[stream_id] == MessageType.SAE:
@@ -86,4 +88,17 @@ def sig_handler(signum, _):
             if output_proto_data is None:
                 continue
 
-            send(f'{CONFIG.target_redis.output_stream_prefix}:{stream_id}', output_proto_data)
+            target_stream = stream_mapping.get(stream_key)
+            send(target_stream, output_proto_data)
+
+            
+def map_config(mapping_config):
+    stream_mapping = {}
+    for mapping in mapping_config:
+        if mapping.source == None:
+            continue
+        if mapping.target == None:
+            stream_mapping[mapping.source] = mapping.source
+        else:            
+            stream_mapping[mapping.source] = mapping.target
+    return stream_mapping