haystack-ml-stack 0.3.4__tar.gz → 0.4.1__tar.gz

{haystack_ml_stack-0.3.4 → haystack_ml_stack-0.4.1}/PKG-INFO

@@ -1,12 +1,13 @@
 Metadata-Version: 2.4
 Name: haystack-ml-stack
-Version: 0.3.4
+Version: 0.4.1
 Summary: Functions related to Haystack ML
 Author-email: Oscar Vega <oscar@haystack.tv>
 License: MIT
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 Requires-Dist: protobuf==6.33.2
+Requires-Dist: orjson==3.11.7
 Provides-Extra: server
 Requires-Dist: pydantic==2.5.0; extra == "server"
 Requires-Dist: cachetools==5.5.2; extra == "server"
@@ -15,6 +16,7 @@ Requires-Dist: aioboto3==12.0.0; extra == "server"
 Requires-Dist: fastapi==0.104.1; extra == "server"
 Requires-Dist: pydantic-settings==2.2; extra == "server"
 Requires-Dist: newrelic==11.1.0; extra == "server"
+Requires-Dist: confluent-kafka==2.13.0; extra == "server"
 
 # Haystack ML Stack
 

{haystack_ml_stack-0.3.4 → haystack_ml_stack-0.4.1}/pyproject.toml

@@ -5,13 +5,13 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "haystack-ml-stack"
-version = "0.3.4"
+version = "0.4.1"
 description = "Functions related to Haystack ML"
 readme = "README.md"
 authors = [{ name = "Oscar Vega", email = "oscar@haystack.tv" }]
 requires-python = ">=3.11"
 dependencies = [
-    "protobuf==6.33.2",
+    "protobuf==6.33.2", "orjson==3.11.7"
 ]
 license = { text = "MIT" }
 
@@ -24,4 +24,5 @@ server = [
     "fastapi==0.104.1",
     "pydantic-settings==2.2",
     "newrelic==11.1.0",
+    "confluent-kafka==2.13.0"
 ]

{haystack_ml_stack-0.3.4 → haystack_ml_stack-0.4.1}/src/haystack_ml_stack/__init__.py

@@ -4,11 +4,11 @@ try:
     from .app import create_app
 
     __all__ = ["create_app"]
-except ImportError:
+except ImportError as e:
     pass
 
+
 from ._serializers import SerializerRegistry, FeatureRegistryId
+from ._version import __version__
 
 __all__ = [*__all__, "SerializerRegistry", "FeatureRegistryId"]
-
-__version__ = "0.3.4"

haystack_ml_stack-0.4.1/src/haystack_ml_stack/_kafka.py

@@ -0,0 +1,88 @@
+from confluent_kafka.aio import AIOProducer
+import orjson
+from google.protobuf.message import Message
+import base64
+import os
+import logging
+from .settings import Settings
+from ._version import __version__
+import hashlib
+
+logger = logging.getLogger(__name__)
+SECURITY_PROTOCOL = "SASL_SSL"
+SASL_MECHANISM = "SCRAM-SHA-512"
+
+
+async def send_to_kafka(
+    producer: AIOProducer,
+    topic: str,
+    user: dict,
+    streams: list[dict],
+    playlist: dict,
+    state: dict,
+    model_output: dict,
+    monitoring_meta: dict,
+) -> None:
+    if topic is None or producer is None:
+        return
+    message = {
+        "userid": user.get("userid"),
+        "client_os": playlist.get("clientOs"),
+        "model_input": {"user": user, "streams": streams, "playlist": playlist},
+        "model_output": model_output,
+        "model_name": state["model_name"].replace(".pkl", "")
+        if state["model_name"]
+        else None,
+        "model_type": "streams",
+        "meta": {
+            "monitoring": monitoring_meta,
+            "haystack_ml_stack_version": __version__,
+            "playlist_category": playlist.get("category"),
+            "user_features": state.get("user_features", []),
+            "stream_features": state.get("stream_features", []),
+        },
+    }
+    delivery_future = await producer.produce(
+        topic, orjson.dumps(message, default=default_serialization)
+    )
+    await delivery_future
+    return
+
+
+def default_serialization(obj):
+    if isinstance(obj, Message):
+        return {
+            "version": obj.version,
+            "proto": base64.b64encode(obj.SerializeToString()).decode("ascii"),
+        }
+    raise orjson.JSONEncodeError("Unknown data type to serialize!")
+
+
+def initialize_kafka_producer(app_config: Settings) -> AIOProducer:
+    secret_keys = orjson.loads(os.getenv("SECRET_KEYS") or "{}")
+    if not secret_keys:
+        raise ValueError("No Kafka credentials found.")
+    with open("/tmp/ca.pem", "w") as f:
+        f.write(base64.b64decode(secret_keys["KAFKA_BROKER_CA_CERTIFICATE"]).decode())
+    kafka_config = {
+        "bootstrap.servers": app_config.kafka_bootstrap_servers,
+        "security.protocol": SECURITY_PROTOCOL,
+        "sasl.username": secret_keys["KAFKA_WRITER_USER"],
+        "sasl.password": secret_keys["KAFKA_WRITER_PASSWORD"],
+        "sasl.mechanism": SASL_MECHANISM,
+        "ssl.ca.location": "/tmp/ca.pem",
+        "compression.type": "lz4",
+    }
+    logger.info(
+        "Initializing kafka producer pushing to topic %s", app_config.kafka_topic
+    )
+    producer = AIOProducer(kafka_config)
+    logger.info("Producer initialized!")
+    return producer
+
+
+def should_log_user(userid: str, kafka_fraction: float) -> bool:
+    if not userid:
+        return False
+    hash_value = int(hashlib.sha256(userid.encode()).hexdigest(), 16) / (2**256)
+    return hash_value < kafka_fraction

haystack_ml_stack-0.4.1/src/haystack_ml_stack/_version.py

@@ -0,0 +1 @@
+__version__ = "0.4.1"

{haystack_ml_stack-0.3.4 → haystack_ml_stack-0.4.1}/src/haystack_ml_stack/app.py

@@ -8,20 +8,21 @@ import time
 from contextlib import asynccontextmanager, AsyncExitStack
 import traceback
 import json
+import asyncio
 
 import aiobotocore.session
 from aiobotocore.config import AioConfig
-from fastapi import FastAPI, HTTPException, Request, Response
+from fastapi import FastAPI, HTTPException, Request, Response, BackgroundTasks
 from fastapi.encoders import jsonable_encoder
 import newrelic.agent
 
-
 from .cache import make_features_cache
 from .dynamo import set_all_features, FeatureRetrievalMeta
 from .model_store import download_and_load_model
 from .settings import Settings
 from . import exceptions
 from ._serializers import SerializerRegistry
+from ._kafka import send_to_kafka, initialize_kafka_producer, should_log_user
 from google.protobuf import text_format
 
 logging.basicConfig(
@@ -123,6 +124,10 @@ def create_app(
         # 1. Load ML Model
         if state["model"] is None:
             await load_model(state, cfg)
+        kafka_producer = None
+        if cfg.kafka_bootstrap_servers is not None:
+            kafka_producer = initialize_kafka_producer(app_config=cfg)
+        state["kafka_producer"] = kafka_producer
         async with AsyncExitStack() as stack:
             # 2. Initialize DynamoDB Client (Persistent Pool)
             session = state["session"]
@@ -139,6 +144,17 @@ def create_app(
             # 3. Shutdown Logic
             # The AsyncExitStack automatically closes the DynamoDB client pool here
             logger.info("Shutting down: Connection pools closed.")
+        logger.info("Shutting down: Flushing Kafka queue.")
+        if kafka_producer is not None:
+            try:
+                await kafka_producer.flush()
+            except Exception:
+                logger.error(
+                    "Unknown exception while flushing kafka queue, shutting down producer.\n%s",
+                    traceback.format_exc(),
+                )
+            finally:
+                await kafka_producer.close()
 
     app = FastAPI(
         title="ML Stream Scorer",
@@ -166,7 +182,9 @@ def create_app(
         }
 
     @app.post("/score", status_code=HTTPStatus.OK)
-    async def score_stream(request: Request, response: Response):
+    async def score_stream(
+        request: Request, response: Response, background_tasks: BackgroundTasks
+    ):
         if state["model"] is None:
             raise HTTPException(
                 status_code=HTTPStatus.SERVICE_UNAVAILABLE,
@@ -187,7 +205,8 @@ def create_app(
             ) from e
         except Exception as e:
             logger.error(
-                "Unexpected exception when parsing request.\n %s", traceback.format_exc()
+                "Unexpected exception when parsing request.\n %s",
+                traceback.format_exc(),
             )
             raise HTTPException(
                 status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail="Unknown exception"
@@ -274,22 +293,24 @@ def create_app(
                 status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
                 detail="Model prediction failed",
             ) from e
-
+        monitoring_meta = {
+            "cache_misses": retrieval_meta.cache_misses,
+            "user_cache_misses": retrieval_meta.user_cache_misses,
+            "stream_cache_misses": retrieval_meta.stream_cache_misses,
+            "user_cache_size": len(user_features_cache),
+            "stream_cache_size": len(stream_features_cache),
+            "retrieval_success": int(retrieval_meta.success),
+            "cache_delay_minutes": retrieval_meta.cache_delay_minutes,
+            "dynamo_ms": retrieval_meta.dynamo_ms,
+            "dynamo_parse_ms": retrieval_meta.parsing_ms,
+            "retrieval_ms": retrieval_meta.retrieval_ms,
+            "preprocess_ms": (predict_start - preprocess_start) * 1e-6,
+            "predict_ms": (predict_end - predict_start) * 1e-6,
+            "total_streams": len(model_output),
+        }
         newrelic.agent.record_custom_event(
             "Inference",
-            {
-                "cache_misses": retrieval_meta.cache_misses,
-                "user_cache_misses": retrieval_meta.user_cache_misses,
-                "stream_cache_misses": retrieval_meta.stream_cache_misses,
-                "retrieval_success": int(retrieval_meta.success),
-                "cache_delay_minutes": retrieval_meta.cache_delay_minutes,
-                "dynamo_ms": retrieval_meta.dynamo_ms,
-                "dynamo_parse_ms": retrieval_meta.parsing_ms,
-                "retrieval_ms": retrieval_meta.retrieval_ms,
-                "preprocess_ms": (predict_start - preprocess_start) * 1e-6,
-                "predict_ms": (predict_end - predict_start) * 1e-6,
-                "total_streams": len(model_output),
-            },
+            monitoring_meta,
         )
         if model_output:
             if random_number < cfg.logs_fraction:
@@ -298,6 +319,18 @@ def create_app(
                     userid,
                     model_output,
                 )
+            if should_log_user(userid=userid, kafka_fraction=cfg.kafka_fraction):
+                background_tasks.add_task(
+                    send_to_kafka,
+                    producer=state["kafka_producer"],
+                    topic=cfg.kafka_topic,
+                    user=user,
+                    streams=streams,
+                    playlist=playlist,
+                    state=state,
+                    model_output=model_output,
+                    monitoring_meta=monitoring_meta,
+                )
             return jsonable_encoder(model_output)
 
         raise HTTPException(

{haystack_ml_stack-0.3.4 → haystack_ml_stack-0.4.1}/src/haystack_ml_stack/settings.py

@@ -1,9 +1,15 @@
 from pydantic_settings import BaseSettings
 from pydantic import Field
 
+
 class Settings(BaseSettings):
     # Logging
     logs_fraction: float = Field(0.01, alias="LOGS_FRACTION")
+    kafka_bootstrap_servers: str | None = Field(
+        default=None, alias="KAFKA_BOOTSTRAP_SERVERS"
+    )
+    kafka_fraction: float = Field(0.01, alias="KAFKA_FRACTION")
+    kafka_topic: str = Field(default=None, alias="KAFKA_TOPIC")
 
     # Model (S3)
     s3_model_path: str | None = Field(default=None, alias="S3_MODEL_PATH")
@@ -14,10 +20,10 @@ class Settings(BaseSettings):
 
     # Cache
     stream_cache_maxsize: int = 50_000
-    user_cache_maxsize: int = 500_000
+    user_cache_maxsize: int = 80_000
     cache_separator: str = "--"
 
     class Config:
         env_file = ".env"
         env_file_encoding = "utf-8"
-        extra = "ignore"
+        extra = "ignore"

{haystack_ml_stack-0.3.4 → haystack_ml_stack-0.4.1}/src/haystack_ml_stack.egg-info/PKG-INFO

@@ -1,12 +1,13 @@
 Metadata-Version: 2.4
 Name: haystack-ml-stack
-Version: 0.3.4
+Version: 0.4.1
 Summary: Functions related to Haystack ML
 Author-email: Oscar Vega <oscar@haystack.tv>
 License: MIT
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 Requires-Dist: protobuf==6.33.2
+Requires-Dist: orjson==3.11.7
 Provides-Extra: server
 Requires-Dist: pydantic==2.5.0; extra == "server"
 Requires-Dist: cachetools==5.5.2; extra == "server"
@@ -15,6 +16,7 @@ Requires-Dist: aioboto3==12.0.0; extra == "server"
 Requires-Dist: fastapi==0.104.1; extra == "server"
 Requires-Dist: pydantic-settings==2.2; extra == "server"
 Requires-Dist: newrelic==11.1.0; extra == "server"
+Requires-Dist: confluent-kafka==2.13.0; extra == "server"
 
 # Haystack ML Stack
 

{haystack_ml_stack-0.3.4 → haystack_ml_stack-0.4.1}/src/haystack_ml_stack.egg-info/SOURCES.txt

@@ -1,7 +1,9 @@
 README.md
 pyproject.toml
 src/haystack_ml_stack/__init__.py
+src/haystack_ml_stack/_kafka.py
 src/haystack_ml_stack/_serializers.py
+src/haystack_ml_stack/_version.py
 src/haystack_ml_stack/app.py
 src/haystack_ml_stack/cache.py
 src/haystack_ml_stack/dynamo.py

{haystack_ml_stack-0.3.4 → haystack_ml_stack-0.4.1}/src/haystack_ml_stack.egg-info/requires.txt

@@ -1,4 +1,5 @@
 protobuf==6.33.2
+orjson==3.11.7
 
 [server]
 pydantic==2.5.0
@@ -8,3 +9,4 @@ aioboto3==12.0.0
 fastapi==0.104.1
 pydantic-settings==2.2
 newrelic==11.1.0
+confluent-kafka==2.13.0