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

haystack_ml_stack/__init__.py

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

haystack_ml_stack/dynamo.py

@@ -9,15 +9,46 @@ 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 unprocessed keys handling."""
+    """
+    Asynchronous batch_get_item with chunking for requests > 100 keys
+    and handling for unprocessed keys.
+    """
     all_items: List[Dict[str, Any]] = []
-    to_fetch = {table_name: {"Keys": keys}}
-
-    while to_fetch:
-        resp = await dynamo_client.batch_get_item(RequestItems=to_fetch)
-        all_items.extend(resp["Responses"].get(table_name, []))
-        unprocessed = resp.get("UnprocessedKeys", {})
-        to_fetch = unprocessed if unprocessed.get(table_name) else {}
+    # 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
 
@@ -100,4 +131,4 @@ async def set_stream_features(
             "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")
+            cache_miss[cache_key][feature_name] = parsed.get("value")

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

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

haystack_ml_stack-0.1.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+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.0.dist-info/RECORD

@@ -1,10 +0,0 @@
-haystack_ml_stack/__init__.py,sha256=_XV9uYskJeiiu8Gj8k2Bxp2IWVgE3iJULbis-28n8D8,75
-haystack_ml_stack/app.py,sha256=GF2F3Qg0m0xyZ04tI2iW2Xmzr18NyyLI7dsfJ_LkByM,5221
-haystack_ml_stack/cache.py,sha256=X5sWARTvxbkCnl6NdILa4lwDt7iAm5Wl9CcgoEv7E6s,506
-haystack_ml_stack/dynamo.py,sha256=1WeQReIZHX3YvjC6abwznqN67-Lax_YEMJ9z_xqw0Z4,3312
-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.0.dist-info/METADATA,sha256=Vs3t-nBxxmlBsZEnTKj0fOn1Ot2jKfD70nrKpQUg9dk,3901
-haystack_ml_stack-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-haystack_ml_stack-0.1.0.dist-info/top_level.txt,sha256=S3g0eH9BeMKygOIwmfB25jtCiAfj0b3CbwPTalcIwvc,18
-haystack_ml_stack-0.1.0.dist-info/RECORD,,