haystack-ml-stack 0.2.4__tar.gz → 0.2.5__tar.gz

{haystack_ml_stack-0.2.4 → haystack_ml_stack-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haystack-ml-stack
-Version: 0.2.4
+Version: 0.2.5
 Summary: Functions related to Haystack ML
 Author-email: Oscar Vega <oscar@haystack.tv>
 License: MIT

{haystack_ml_stack-0.2.4 → haystack_ml_stack-0.2.5}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "haystack-ml-stack"
-version = "0.2.4"
+version = "0.2.5"
 description = "Functions related to Haystack ML"
 readme = "README.md"
 authors = [{ name = "Oscar Vega", email = "oscar@haystack.tv" }]

{haystack_ml_stack-0.2.4 → haystack_ml_stack-0.2.5}/src/haystack_ml_stack/__init__.py

@@ -1,4 +1,4 @@
 from .app import create_app
 
 __all__ = ["create_app"]
-__version__ = "0.2.4"
+__version__ = "0.2.5"

{haystack_ml_stack-0.2.4 → haystack_ml_stack-0.2.5}/src/haystack_ml_stack/app.py

@@ -5,8 +5,10 @@ import sys
 from http import HTTPStatus
 from typing import Any, Dict, List, Optional
 import time
+from contextlib import asynccontextmanager, AsyncExitStack
 
 import aiobotocore.session
+from aiobotocore.config import AioConfig
 from fastapi import FastAPI, HTTPException, Request, Response
 from fastapi.encoders import jsonable_encoder
 import newrelic.agent
@@ -25,7 +27,7 @@ logging.basicConfig(
 )
 
 logger = logging.getLogger(__name__)
-APP_NAME = os.environ.get("NEW_RELIC_APP_NAME", None)
+MAX_POOL_CONNECTIONS = int(os.environ.get("MAX_POOL_CONNECTIONS", 50))
 
 
 def create_app(
@@ -39,12 +41,6 @@ def create_app(
     """
     cfg = settings or Settings()
 
-    app = FastAPI(
-        title="ML Stream Scorer",
-        description="Scores video streams using a pre-trained ML model and DynamoDB features.",
-        version="1.0.0",
-    )
-
     # Mutable state: cache + model
     features_cache = make_features_cache(cfg.cache_maxsize)
     state: Dict[str, Any] = {
@@ -55,24 +51,59 @@ def create_app(
         ),
     }
 
-    @app.on_event("startup")
-    async def _startup() -> None:
-        if state["model"] is not None:
-            logger.info("Using preloaded model.")
-            return
-
-        if not cfg.s3_model_path:
-            logger.critical("S3_MODEL_PATH not set; service will be unhealthy.")
-            return
-
-        try:
-            state["model"] = await download_and_load_model(
-                cfg.s3_model_path, aio_session=state["session"]
+    @asynccontextmanager
+    async def lifespan(app_server: FastAPI):
+        """
+        Handles startup and shutdown logic.
+        Everything before 'yield' runs on startup.
+        Everything after 'yield' runs on shutdown.
+        """
+        async with AsyncExitStack() as stack:
+            # 1. Initialize DynamoDB Client (Persistent Pool)
+            session = state["session"]
+            state["dynamo_client"] = await stack.enter_async_context(
+                session.create_client(
+                    "dynamodb",
+                    # Ensure the pool is large enough for ML concurrency
+                    config=AioConfig(max_pool_connections=MAX_POOL_CONNECTIONS),
+                )
             )
-            state["stream_features"] = state["model"].get("stream_features", [])
-            logger.info("Model loaded on startup.")
-        except Exception as e:
-            logger.critical("Failed to load model: %s", e)
+            logger.info("DynamoDB persistent client initialized.")
+
+            # 2. Load ML Model
+            if state["model"] is None:
+                if not cfg.s3_model_path:
+                    logger.critical("S3_MODEL_PATH not set; service will be unhealthy.")
+                else:
+                    try:
+                        # Pass the persistent session/client if needed
+                        state["model"] = await download_and_load_model(
+                            cfg.s3_model_path, aio_session=state["session"]
+                        )
+                        state["stream_features"] = state["model"].get(
+                            "stream_features", []
+                        )
+                        state["user_features"] = state["model"].get("user_features", [])
+
+                        newrelic.agent.add_custom_attribute(
+                            "total_stream_features", len(state["stream_features"])
+                        )
+                        logger.info("Model loaded successfully.")
+                    except Exception as e:
+                        logger.critical("Failed to load model: %s", e)
+
+            yield
+
+            # 3. Shutdown Logic
+            # The AsyncExitStack automatically closes the DynamoDB client pool here
+            logger.info("Shutting down: Connection pools closed.")
+
+    app = FastAPI(
+        title="ML Stream Scorer",
+        description="Scores video streams using a pre-trained ML model and DynamoDB features.",
+        version="1.0.0",
+        lifespan=lifespan,
+    )
 
     @app.get("/health", status_code=HTTPStatus.OK)
     async def health():
@@ -130,7 +161,7 @@ def create_app(
         )
         if stream_features:
             retrieval_meta = await set_stream_features(
-                aio_session=state["session"],
+                dynamo_client=state["dynamo_client"],
                 streams=streams,
                 stream_features=stream_features,
                 features_cache=features_cache,
@@ -168,7 +199,6 @@ def create_app(
         newrelic.agent.record_custom_event(
             "Inference",
             {
-                "app_name": APP_NAME,
                 "cache_misses": retrieval_meta.cache_misses,
                 "retrieval_success": int(retrieval_meta.success),
                 "cache_delay_minutes": retrieval_meta.cache_delay_minutes,

{haystack_ml_stack-0.2.4 → haystack_ml_stack-0.2.5}/src/haystack_ml_stack/cache.py

@@ -5,14 +5,14 @@ from cachetools import TLRUCache
 
 def _ttu(_, value: Any, now: float) -> float:
     """Time-To-Use policy: allow per-item TTL via 'cache_ttl_in_seconds' or fallback."""
-    ONE_YEAR = 365 * 24 * 60 * 60
+    ONE_WEEK = 7 * 24 * 60 * 60
     try:
         ttl = int(value.get("cache_ttl_in_seconds", -1))
         if ttl > 0:
             return now + ttl
     except Exception:
         pass
-    return now + ONE_YEAR
+    return now + ONE_WEEK
 
 
 def make_features_cache(maxsize: int) -> TLRUCache:

{haystack_ml_stack-0.2.4 → haystack_ml_stack-0.2.5}/src/haystack_ml_stack/dynamo.py

@@ -2,7 +2,6 @@ from typing import Any, Dict, List, NamedTuple
 import logging
 import time
 import datetime
-import aiobotocore.session
 from boto3.dynamodb.types import TypeDeserializer
 import newrelic.agent
 import asyncio
@@ -38,10 +37,14 @@ async def async_batch_get(
     """
     # DynamoDB's BatchGetItem has a 100-item limit per request.
     CHUNK_SIZE = 100
-    chunks = [keys[i : i + CHUNK_SIZE] for i in range(0, len(keys), CHUNK_SIZE)]
-    tasks = [_fetch_chunk(dynamo_client, table_name, chunk) for chunk in chunks]
-    results = await asyncio.gather(*tasks)
-    all_items = [item for batch in results for item in batch]
+
+    if len(keys) <= CHUNK_SIZE:
+        all_items = await _fetch_chunk(dynamo_client, table_name, keys)
+    else:
+        chunks = [keys[i : i + CHUNK_SIZE] for i in range(0, len(keys), CHUNK_SIZE)]
+        tasks = [_fetch_chunk(dynamo_client, table_name, chunk) for chunk in chunks]
+        results = await asyncio.gather(*tasks)
+        all_items = [item for batch in results for item in batch]
     return all_items
 
 
@@ -76,7 +79,6 @@ async def _fetch_chunk(dynamo_client, table_name: str, chunk_keys):
     return items
 
 
-@newrelic.agent.function_trace()
 def parse_dynamo_item(item: Dict[str, Any]) -> Dict[str, Any]:
     """Parse a DynamoDB attribute map (low-level) to Python types."""
     # out: Dict[str, Any] = {}
@@ -92,7 +94,7 @@ async def set_stream_features(
     features_table: str,
     stream_pk_prefix: str,
     cache_sep: str,
-    aio_session: aiobotocore.session.Session | None = None,
+    dynamo_client,
 ) -> FeatureRetrievalMeta:
     time_start = time.perf_counter_ns()
     """Fetch missing features for streams from DynamoDB and fill them into streams."""
@@ -145,22 +147,20 @@ async def set_stream_features(
         pk = f"{stream_pk_prefix}{stream_url}"
         keys.append({"pk": {"S": pk}, "sk": {"S": sk}})
 
-    session = aio_session or aiobotocore.session.get_session()
     dynamo_start = time.perf_counter_ns()
-    async with session.create_client("dynamodb") as dynamodb:
-        try:
-            items = await async_batch_get(dynamodb, features_table, keys)
-        except Exception as e:
-            logger.error("DynamoDB batch_get failed: %s", e)
-            end_time = time.perf_counter_ns()
-            return FeatureRetrievalMeta(
-                cache_misses=cache_misses,
-                retrieval_ms=(end_time - time_start) * 1e-6,
-                success=False,
-                cache_delay_minutes=cache_delay / 60,
-                dynamo_ms=(end_time - dynamo_start) * 1e-6,
-                parsing_ms=0,
-            )
+    try:
+        items = await async_batch_get(dynamo_client, features_table, keys)
+    except Exception as e:
+        logger.error("DynamoDB batch_get failed: %s", e)
+        end_time = time.perf_counter_ns()
+        return FeatureRetrievalMeta(
+            cache_misses=cache_misses,
+            retrieval_ms=(end_time - time_start) * 1e-6,
+            success=False,
+            cache_delay_minutes=cache_delay / 60,
+            dynamo_ms=(end_time - dynamo_start) * 1e-6,
+            parsing_ms=0,
+        )
     dynamo_end = time.perf_counter_ns()
     updated_keys = set()
     for item in items:

{haystack_ml_stack-0.2.4 → haystack_ml_stack-0.2.5}/src/haystack_ml_stack.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haystack-ml-stack
-Version: 0.2.4
+Version: 0.2.5
 Summary: Functions related to Haystack ML
 Author-email: Oscar Vega <oscar@haystack.tv>
 License: MIT