haystack-ml-stack 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

haystack_ml_stack/__init__.py

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

haystack_ml_stack/dynamo.py

@@ -94,9 +94,11 @@ async def set_stream_features(
     for f in stream_features:
         for s in streams:
             key = f"{s['streamUrl']}{cache_sep}{f}"
-            cached = features_cache.get(key)
-            if cached is not None:
-                s[f] = cached["value"]
+            if key in features_cache:
+                # Only set if value is not None
+                cached = features_cache.get(key)
+                if cached["value"] is not None:
+                    s[f] = cached["value"]
             else:
                 cache_miss[key] = s
 
@@ -120,6 +122,7 @@ async def set_stream_features(
             logger.error("DynamoDB batch_get failed: %s", e)
             return
 
+    updated_keys = set()
     for item in items:
         stream_url = item["pk"]["S"].removeprefix(stream_pk_prefix)
         feature_name = item["sk"]["S"]
@@ -132,3 +135,10 @@ async def set_stream_features(
         }
         if cache_key in cache_miss:
             cache_miss[cache_key][feature_name] = parsed.get("value")
+            updated_keys.add(cache_key)
+
+    # Save keys that were not found in DynamoDB with None value
+    if len(updated_keys) < len(cache_miss):
+        missing_keys = set(cache_miss.keys()) - updated_keys
+        for k in missing_keys:
+            features_cache[k] = {"value": None, "cache_ttl_in_seconds": 300}

{haystack_ml_stack-0.1.1.dist-info → haystack_ml_stack-0.1.2.dist-info}/METADATA

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

haystack_ml_stack-0.1.2.dist-info/RECORD

@@ -0,0 +1,16 @@
+haystack_ml_stack/__init__.py,sha256=HuhQrhFWhx1QcGQLZlFHHa6OmSyPxvU-cF1jEBvSzBk,75
+haystack_ml_stack/app.py,sha256=GF2F3Qg0m0xyZ04tI2iW2Xmzr18NyyLI7dsfJ_LkByM,5221
+haystack_ml_stack/cache.py,sha256=X5sWARTvxbkCnl6NdILa4lwDt7iAm5Wl9CcgoEv7E6s,506
+haystack_ml_stack/dynamo.py,sha256=JltS6Wi0vGwsAGa5lz_qMz5Lb4TMotq-TTE2hpINOkE,4994
+haystack_ml_stack/model_store.py,sha256=ON1R1gl_TWhP20l383DwDPcuKEJJpmC1ziqSL4Mt-rw,1179
+haystack_ml_stack/settings.py,sha256=2JH-H7NtKaVbc11jq4Yl3h5TjCve98g--8pqNvQjfj4,574
+haystack_test_package/__init__.py,sha256=w6S607EnLThZyW699W5dGxjG4UarthbE-hm_PDX-KBM,75
+haystack_test_package/app.py,sha256=GF2F3Qg0m0xyZ04tI2iW2Xmzr18NyyLI7dsfJ_LkByM,5221
+haystack_test_package/cache.py,sha256=X5sWARTvxbkCnl6NdILa4lwDt7iAm5Wl9CcgoEv7E6s,506
+haystack_test_package/dynamo.py,sha256=Qx8iLIxA09ub-qkPYDP-3RyCEetGMtFyDp-6MYqLANE,4736
+haystack_test_package/model_store.py,sha256=ON1R1gl_TWhP20l383DwDPcuKEJJpmC1ziqSL4Mt-rw,1179
+haystack_test_package/settings.py,sha256=2JH-H7NtKaVbc11jq4Yl3h5TjCve98g--8pqNvQjfj4,574
+haystack_ml_stack-0.1.2.dist-info/METADATA,sha256=LBVA_iqpNk6zis26FoLc1aM3Q-gy4s8ctdfbKBVS_Cc,3901
+haystack_ml_stack-0.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+haystack_ml_stack-0.1.2.dist-info/top_level.txt,sha256=N89LQwtVIpTKJhq2cnAqMSk6rbmWRGqinL1yaBZmfpk,40
+haystack_ml_stack-0.1.2.dist-info/RECORD,,

haystack_ml_stack-0.1.2.dist-info/top_level.txt

@@ -0,0 +1,2 @@
+haystack_ml_stack
+haystack_test_package

haystack_test_package/__init__.py

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

haystack_test_package/app.py

@@ -0,0 +1,158 @@
+import logging
+import os
+import random
+import sys
+from http import HTTPStatus
+from typing import Any, Dict, List, Optional
+
+import aiobotocore.session
+from fastapi import FastAPI, HTTPException, Request, Response
+from fastapi.encoders import jsonable_encoder
+
+from .cache import make_features_cache
+from .dynamo import set_stream_features
+from .model_store import download_and_load_model
+from .settings import Settings
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="[%(levelname)s] [%(process)d] %(name)s : %(message)s",
+    handlers=[logging.StreamHandler(sys.stdout)],
+    force=True,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def create_app(
+    settings: Optional[Settings] = None,
+    *,
+    preloaded_model: Optional[Dict[str, Any]] = None,
+) -> FastAPI:
+    """
+    Build a FastAPI app with injectable settings and model.
+    If `preloaded_model` is None, the app will load from S3 on startup.
+    """
+    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] = {
+        "model": preloaded_model,
+        "session": aiobotocore.session.get_session(),
+        "model_name": (
+            os.path.basename(cfg.s3_model_path) if cfg.s3_model_path else None
+        ),
+    }
+
+    @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"]
+            )
+            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)
+
+    @app.get("/health", status_code=HTTPStatus.OK)
+    async def health():
+        model_ok = state["model"] is not None
+        if not model_ok:
+            raise HTTPException(
+                status_code=HTTPStatus.SERVICE_UNAVAILABLE,
+                detail="ML Model not loaded",
+            )
+        return {
+            "status": "ok",
+            "model_loaded": True,
+            "cache_size": len(features_cache),
+            "model_name": state.get("model_name"),
+            "stream_features": state.get("stream_features", []),
+        }
+
+    @app.post("/score", status_code=HTTPStatus.OK)
+    async def score_stream(request: Request, response: Response):
+        if state["model"] is None:
+            raise HTTPException(
+                status_code=HTTPStatus.SERVICE_UNAVAILABLE,
+                detail="ML Model not loaded",
+            )
+
+        try:
+            data = await request.json()
+        except Exception:
+            raise HTTPException(
+                status_code=HTTPStatus.BAD_REQUEST, detail="Invalid JSON payload"
+            )
+
+        user = data.get("user", {})
+        streams: List[Dict[str, Any]] = data.get("streams", [])
+        playlist = data.get("playlist", {})
+
+        if not streams:
+            logger.warning("No streams provided for user %s", user.get("userid", ""))
+            return {}
+
+        # Feature fetch (optional based on model)
+        model = state["model"]
+        stream_features = model.get("stream_features", []) or []
+        if stream_features:
+            logger.info("Fetching stream features for user %s", user.get("userid", ""))
+            await set_stream_features(
+                aio_session=state["session"],
+                streams=streams,
+                stream_features=stream_features,
+                features_cache=features_cache,
+                features_table=cfg.features_table,
+                stream_pk_prefix=cfg.stream_pk_prefix,
+                cache_sep=cfg.cache_separator,
+            )
+
+        # Sampling logs
+        if random.random() < cfg.logs_fraction:
+            logger.info("User %s streams: %s", user.get("userid", ""), streams)
+
+        # Synchronous model execution (user code)
+        try:
+            model_input = model["preprocess"](
+                user, streams, playlist, model.get("params")
+            )
+            model_output = model["predict"](model_input, model.get("params"))
+        except Exception as e:
+            logger.error("Model prediction failed: %s", e)
+            raise HTTPException(
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+                detail="Model prediction failed",
+            )
+
+        if model_output:
+            return jsonable_encoder(model_output)
+
+        raise HTTPException(
+            status_code=HTTPStatus.NOT_FOUND, detail="No model output generated"
+        )
+
+    @app.get("/", status_code=HTTPStatus.OK)
+    async def root():
+        return {
+            "message": "ML Scoring Service is running.",
+            "model_name": state.get("model_name"),
+        }
+
+    return app

haystack_test_package/cache.py

@@ -0,0 +1,19 @@
+from typing import Any
+
+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
+    try:
+        ttl = int(value.get("cache_ttl_in_seconds", -1))
+        if ttl > 0:
+            return now + ttl
+    except Exception:
+        pass
+    return now + ONE_YEAR
+
+
+def make_features_cache(maxsize: int) -> TLRUCache:
+    return TLRUCache(maxsize=maxsize, ttu=_ttu)

haystack_test_package/dynamo.py

@@ -0,0 +1,137 @@
+from typing import Any, Dict, List
+import logging
+
+import aiobotocore.session
+
+logger = logging.getLogger(__name__)
+
+
+async def async_batch_get(
+    dynamo_client, table_name: str, keys: List[Dict[str, Any]]
+) -> List[Dict[str, Any]]:
+    """
+    Asynchronous batch_get_item with chunking for requests > 100 keys
+    and handling for unprocessed keys.
+    """
+    all_items: List[Dict[str, Any]] = []
+    # DynamoDB's BatchGetItem has a 100-item limit per request.
+    CHUNK_SIZE = 100
+
+    # Split the keys into chunks of 100
+    for i in range(0, len(keys), CHUNK_SIZE):
+        chunk_keys = keys[i : i + CHUNK_SIZE]
+        to_fetch = {table_name: {"Keys": chunk_keys}}
+
+        # Inner loop to handle unprocessed keys for the current chunk
+        # Max retries of 3
+        retries = 3
+        while to_fetch and retries > 0:
+            retries -= 1
+            try:
+                resp = await dynamo_client.batch_get_item(RequestItems=to_fetch)
+
+                if "Responses" in resp and table_name in resp["Responses"]:
+                    all_items.extend(resp["Responses"][table_name])
+
+                unprocessed = resp.get("UnprocessedKeys", {})
+                # If there are unprocessed keys, set them to be fetched in the next iteration
+                if unprocessed and unprocessed.get(table_name):
+                    logger.warning(
+                        "Retrying %d unprocessed keys.",
+                        len(unprocessed[table_name]["Keys"]),
+                    )
+                    to_fetch = unprocessed
+                else:
+                    # All keys in the chunk were processed, exit the inner loop
+                    to_fetch = {}
+
+            except Exception as e:
+                logger.error("Error during batch_get_item for a chunk: %s", e)
+                # Stop trying to process this chunk on error and move to the next
+                to_fetch = {}
+
+    return all_items
+
+
+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] = {}
+    for k, v in item.items():
+        if "N" in v:
+            out[k] = float(v["N"])
+        elif "S" in v:
+            out[k] = v["S"]
+        elif "SS" in v:
+            out[k] = v["SS"]
+        elif "NS" in v:
+            out[k] = [float(n) for n in v["NS"]]
+        elif "BOOL" in v:
+            out[k] = v["BOOL"]
+        elif "NULL" in v:
+            out[k] = None
+        elif "L" in v:
+            out[k] = [parse_dynamo_item({"value": i})["value"] for i in v["L"]]
+        elif "M" in v:
+            out[k] = parse_dynamo_item(v["M"])
+    return out
+
+
+async def set_stream_features(
+    *,
+    streams: List[Dict[str, Any]],
+    stream_features: List[str],
+    features_cache,
+    features_table: str,
+    stream_pk_prefix: str,
+    cache_sep: str,
+    aio_session: aiobotocore.session.Session | None = None,
+) -> None:
+    """Fetch missing features for streams from DynamoDB and fill them into streams."""
+    if not streams or not stream_features:
+        return
+
+    cache_miss: Dict[str, Dict[str, Any]] = {}
+    for f in stream_features:
+        for s in streams:
+            key = f"{s['streamUrl']}{cache_sep}{f}"
+            cached = features_cache.get(key)
+            if cached is not None:
+                s[f] = cached["value"]
+            else:
+                cache_miss[key] = s
+
+    if not cache_miss:
+        return
+
+    logger.info("Cache miss for %d items", len(cache_miss))
+
+    # Prepare keys
+    keys = []
+    for k in cache_miss.keys():
+        stream_url, sk = k.split(cache_sep, 1)
+        pk = f"{stream_pk_prefix}{stream_url}"
+        keys.append({"pk": {"S": pk}, "sk": {"S": sk}})
+    logger.info("Keys prepared for DynamoDB: %s", keys)
+
+    session = aio_session or aiobotocore.session.get_session()
+    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)
+            return
+        logger.info("DynamoDB returned %d items", len(items))
+
+    for item in items:
+        stream_url = item["pk"]["S"].removeprefix(stream_pk_prefix)
+        feature_name = item["sk"]["S"]
+        cache_key = f"{stream_url}{cache_sep}{feature_name}"
+        parsed = parse_dynamo_item(item)
+        logger.info("DynamoDB item parsed: %s for %s", parsed, cache_key)
+
+        features_cache[cache_key] = {
+            "value": parsed.get("value"),
+            "cache_ttl_in_seconds": int(parsed.get("cache_ttl_in_seconds", -1)),
+        }
+        if cache_key in cache_miss:
+            cache_miss[cache_key][feature_name] = parsed.get("value")

haystack_test_package/model_store.py

@@ -0,0 +1,36 @@
+import logging
+import os
+from typing import Any, Dict
+
+import aiobotocore.session
+import cloudpickle
+
+logger = logging.getLogger(__name__)
+
+
+async def download_and_load_model(
+    s3_url: str, aio_session: aiobotocore.session.Session | None = None
+) -> Dict[str, Any]:
+    """
+    Downloads cloudpickled model dict from S3 and loads it.
+    Expected keys: 'preprocess', 'predict', 'params', optional 'stream_features'.
+    """
+    if not s3_url or not s3_url.startswith("s3://"):
+        raise ValueError("S3_MODEL_PATH must be a valid s3:// URL")
+
+    bucket, key = s3_url.replace("s3://", "").split("/", 1)
+    pid = os.getpid()
+    local_path = f"/tmp/model_{pid}.pkl"
+
+    session = aio_session or aiobotocore.session.get_session()
+    async with session.create_client("s3") as s3:
+        logger.info("Downloading model from %s...", s3_url)
+        resp = await s3.get_object(Bucket=bucket, Key=key)
+        data = await resp["Body"].read()
+        with open(local_path, "wb") as f:
+            f.write(data)
+        logger.info("Model downloaded to %s", local_path)
+
+    with open(local_path, "rb") as f:
+        model: Dict[str, Any] = cloudpickle.load(f)
+    return model

haystack_test_package/settings.py

@@ -0,0 +1,22 @@
+from pydantic_settings import BaseSettings
+from pydantic import Field
+
+class Settings(BaseSettings):
+    # Logging
+    logs_fraction: float = Field(0.01, alias="LOGS_FRACTION")
+
+    # Model (S3)
+    s3_model_path: str | None = Field(default=None, alias="S3_MODEL_PATH")
+
+    # DynamoDB
+    features_table: str = Field("features", alias="FEATURES_TABLE")
+    stream_pk_prefix: str = "STREAM#"
+
+    # Cache
+    cache_maxsize: int = 50_000
+    cache_separator: str = "--"
+
+    class Config:
+        env_file = ".env"
+        env_file_encoding = "utf-8"
+        extra = "ignore"

haystack_ml_stack-0.1.1.dist-info/RECORD

@@ -1,10 +0,0 @@
-haystack_ml_stack/__init__.py,sha256=8agfBhexMCMsnLx9WmLeKaLEILI5064mTlMnICXyTfU,75
-haystack_ml_stack/app.py,sha256=GF2F3Qg0m0xyZ04tI2iW2Xmzr18NyyLI7dsfJ_LkByM,5221
-haystack_ml_stack/cache.py,sha256=X5sWARTvxbkCnl6NdILa4lwDt7iAm5Wl9CcgoEv7E6s,506
-haystack_ml_stack/dynamo.py,sha256=npYCoYwLs2PyxxjHGpZWkk-aEwKoD6rKx_3snrcO3YA,4544
-haystack_ml_stack/model_store.py,sha256=ON1R1gl_TWhP20l383DwDPcuKEJJpmC1ziqSL4Mt-rw,1179
-haystack_ml_stack/settings.py,sha256=2JH-H7NtKaVbc11jq4Yl3h5TjCve98g--8pqNvQjfj4,574
-haystack_ml_stack-0.1.1.dist-info/METADATA,sha256=jgRcnCn2AYXTIzOFItk2yv8ZnRLaGR7c3GH5ZlwhdkM,3901
-haystack_ml_stack-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-haystack_ml_stack-0.1.1.dist-info/top_level.txt,sha256=S3g0eH9BeMKygOIwmfB25jtCiAfj0b3CbwPTalcIwvc,18
-haystack_ml_stack-0.1.1.dist-info/RECORD,,

haystack_ml_stack-0.1.1.dist-info/top_level.txt

@@ -1 +0,0 @@
-haystack_ml_stack