streamlit-octostar-utils 0.6.0__tar.gz → 0.6.3.dev1__tar.gz

{streamlit_octostar_utils-0.6.0 → streamlit_octostar_utils-0.6.3.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamlit-octostar-utils
-Version: 0.6.0
+Version: 0.6.3.dev1
 Summary: 
 License: MIT
 License-File: LICENSE

{streamlit_octostar_utils-0.6.0 → streamlit_octostar_utils-0.6.3.dev1}/pyproject.toml

@@ -5,7 +5,7 @@ include = '\.pyi?$'
 
 [tool.poetry]
 name = "streamlit-octostar-utils"
-version = "0.6.0"
+version = "0.6.3-dev.1"
 description = ""
 license = "MIT"
 authors = ["Octostar"]

{streamlit_octostar_utils-0.6.0 → streamlit_octostar_utils-0.6.3.dev1}/streamlit_octostar_utils/api_crafter/celery.py

@@ -590,12 +590,64 @@ class CeleryExecutor(object):
                 }
             logger.info(f"All resources preloaded for queue {queue}")
 
-    def set_awaiting_state(self, sender=None, headers=None, **kwargs):
-        task_id = headers.get("id") if headers else None
+    def _redis_result_ttl_while_awaiting(self, queue_name, hard_time_limit):
+        """TTL (seconds) for result-backend keys while a task may sit AWAITING in the broker.
+
+        Uses ``ceil(T * N / M) + result_expires`` when a positive hard time limit and queue
+        config exist; otherwise falls back to ``task_expires + result_expires`` so broker
+        message lifetime does not outlive the result key.
+        """
+        result_exp = int(self.app.conf.result_expires)
+        task_exp = int(self.app.conf.task_expires)
+        tl = int(hard_time_limit or 0)
+        qc = self.queue_config.get(queue_name)
+        if qc and tl > 0:
+            m_workers = max(int(qc.n_workers), 1)
+            n_cap = qc.max_tasks_in_queue
+            if n_cap is None:
+                n_cap = m_workers
+            n_cap = max(int(n_cap), 1)
+            queue_wait = (tl * n_cap + m_workers - 1) // m_workers
+            return queue_wait + result_exp
+        return max(result_exp, task_exp + result_exp)
+
+    @staticmethod
+    def _hard_time_limit_from_publish(body, headers, sender_task_name, app):
+        """Resolve hard time limit from AMQP publish (v1 body dict or v2 headers list)."""
+        h = headers or {}
+        raw = h.get("timelimit")
+        if not raw and isinstance(body, dict):
+            raw = body.get("timelimit")
+        if raw:
+            hard = raw[0] if isinstance(raw, (list, tuple)) else None
+            if hard is not None:
+                return int(hard)
+        if sender_task_name and sender_task_name in app.tasks:
+            t = getattr(app.tasks[sender_task_name], "time_limit", None)
+            if t is not None:
+                return int(t)
+        return 0
+
+    def set_awaiting_state(
+        self, sender=None, body=None, headers=None, routing_key=None, **kwargs
+    ):
+        headers = headers or {}
+        task_id = headers.get("id")
+        if not task_id and isinstance(body, dict):
+            task_id = body.get("id")
         if not task_id:
             return
         result = AsyncResult(task_id, app=self.app)
         result.backend.store_result(task_id, result=None, state=CeleryExecutor.AWAITING)
+        queue_name = routing_key or self.app.conf.task_default_routing_key
+        hard_tl = self._hard_time_limit_from_publish(body, headers, sender, self.app)
+        await_ttl = self._redis_result_ttl_while_awaiting(queue_name, hard_tl)
+        meta_key = f"{CeleryExecutor.CELERY_BROKER_PREFIX}{task_id}"
+        try:
+            if await_ttl > 0 and self.redis_client.exists(meta_key):
+                self.redis_client.expire(meta_key, await_ttl)
+        except Exception:
+            logger.debug("Could not extend AWAITING result TTL for %s", task_id, exc_info=True)
 
     def set_started_state(self, task_id, task, *args, **kwargs):
         result = AsyncResult(task_id, app=self.app)
@@ -1059,9 +1111,9 @@ class CeleryExecutor(object):
         def _send_task(task_fn, task_id, options):
             task_fn.apply_async(task_id=task_id, **options)
 
-        def _store_task_queue_mapping(task_id, queue_name):
+        def _store_task_queue_mapping(task_id, queue_name, redis_await_ttl):
             pipe = self.redis_client.pipeline()
-            pipe.set(f"task:queue:{task_id}", queue_name, ex=self.app.conf.result_expires)
+            pipe.set(f"task:queue:{task_id}", queue_name, ex=redis_await_ttl)
             pipe.set(f"queue:first_enqueued:{queue_name}", str(time.time()), nx=True)
             pipe.execute()
 
@@ -1075,6 +1127,9 @@ class CeleryExecutor(object):
                 f"Queue '{queue_name}' is stalled. Service temporarily unavailable."
             )
 
+        hard_tl = options.get("time_limit") or getattr(task_fn, "time_limit", None) or 0
+        redis_await_ttl = self._redis_result_ttl_while_awaiting(queue_name, hard_tl)
+
         try:
             if part is not None:
                 await self._write_task_data_with_part(
@@ -1094,7 +1149,7 @@ class CeleryExecutor(object):
                 self.set_thread_pool, _send_task, task_fn, task_id, options
             )
             await asyncio.get_running_loop().run_in_executor(
-                self.set_thread_pool, _store_task_queue_mapping, task_id, queue_name
+                self.set_thread_pool, _store_task_queue_mapping, task_id, queue_name, redis_await_ttl
             )
         except asyncio.CancelledError:
             logger.info(f"Cancelling task {task_id} due to disconnect!")

{streamlit_octostar_utils-0.6.0 → streamlit_octostar_utils-0.6.3.dev1}/streamlit_octostar_utils/api_crafter/nifi.py

@@ -670,6 +670,7 @@ class NifiContextManager(object):
         self.lazy_sync = lazy_sync
         self.client, self.ontology_name = self.get_client(json_data)
         self._ontology = None
+        self._input_loader = input_loader
 
     @property
     def ontology(self):
@@ -863,20 +864,20 @@ class NifiContextManager(object):
         revert it separately by including it in ``entities`` if needed.
 
         """
-        if self._input_loader is None:
+        loader = getattr(self, "_input_loader", None)
+        if loader is None:
             raise RuntimeError(
                 "revert is unavailable: this NifiContextManager was not "
                 "constructed with an input_loader"
             )
         if not entities:
             return
-        metadata, parts = self._input_loader.load()
+        metadata, parts = loader.load()
         if not parts:
             raise RuntimeError(
                 "revert is unavailable: input loader returned no parts"
             )
         body = json.loads(parts[0])
-        # Drop wrapper-side refs to the raw bytes/parts immediately.
         metadata = None
         parts = None
         try:

{streamlit_octostar_utils-0.6.0 → streamlit_octostar_utils-0.6.3.dev1}/streamlit_octostar_utils/api_crafter/parallelism.py

@@ -46,23 +46,34 @@ R = TypeVar('R')
 class ParallelismConfig:
     """
     Configuration for controlling parallelism and throughput.
-    
+
+    `max_parallel` is the single user-facing concurrency knob; its concrete
+    meaning depends on the executor it is paired with:
+
+    * ``ThreadExecutor`` -- maximum number of items processed simultaneously
+      (semaphore + thread-pool size).
+    * ``BatchExecutor`` -- default number of items per ``batch_fn`` call when
+      the caller does not pass an explicit ``batch_size``. The batch dimension
+      is the natural parallelism for a batched-inference model.
+    * ``LoopExecutor`` -- ignored; the executor is intrinsically sequential
+      and only honors ``throttle``.
+
     Attributes:
-        max_parallel: Maximum number of entities being processed simultaneously.
-                     Acts as a concurrency limit (semaphore).
-                     None = unlimited concurrency
-        
+        max_parallel: Maximum number of items processed simultaneously, or
+                     (for ``BatchExecutor``) the default batch dimension.
+                     None = unlimited concurrency / batch_size defaults to 1.
+
         throttle: Maximum rate of entities starting processing per second.
                  Acts as a rate limiter.
                  None = unlimited rate
-    
+
     Examples:
         >>> # At most 4 entities processed at once, no rate limit
         >>> config = ParallelismConfig(max_parallel=4, throttle=None)
-        
+
         >>> # Unlimited concurrency, but max 10 entities/second
         >>> config = ParallelismConfig(max_parallel=None, throttle=10)
-        
+
         >>> # Max 2 concurrent, max 5/second rate
         >>> config = ParallelismConfig(max_parallel=2, throttle=5)
     """
@@ -603,23 +614,29 @@ class LoopExecutor(Executor):
 class BatchExecutor(Executor):
     """
     Process items in batches with parallelism control.
-    
+
     Useful for model inference where batching improves throughput
     (e.g., neural network inference, OCR, NER).
-    
+
+    The user-facing knob is ``ParallelismConfig.max_parallel``, which controls
+    the default batch dimension: each ``batch_fn`` call receives up to
+    ``max_parallel`` items. Callers that need to pin a different batch size
+    (e.g. "one bin per call" when the items are themselves pre-batched, or
+    "all items in a single mega-batch") may pass an explicit ``batch_size``
+    to ``process()`` to override the default.
+
     Example:
-        >>> config = ParallelismConfig(max_parallel=None, throttle=10)
+        >>> # Default: batch_size derived from max_parallel
+        >>> config = ParallelismConfig(max_parallel=8)
         >>> processor = BatchExecutor(config)
-        >>> 
+        >>>
         >>> def process_batch(texts):
-        ...     # Process batch of texts together
         ...     return model.predict(texts)
-        >>> 
-        >>> results = processor.process(
-        ...     all_texts,
-        ...     process_batch,
-        ...     batch_size=10
-        ... )
+        >>>
+        >>> results = processor.process(all_texts, process_batch)
+        >>>
+        >>> # Or override explicitly:
+        >>> results = processor.process(all_texts, process_batch, batch_size=10)
     """
     
     def __init__(self, config: ParallelismConfig):
@@ -633,50 +650,53 @@ class BatchExecutor(Executor):
         self.controller = ParallelismController(config)
     
     def process(
-        self, 
-        items: List[T], 
+        self,
+        items: List[T],
         batch_fn: Callable[[List[T]], List[R]],
-        batch_size: int = 1
+        batch_size: Optional[int] = None
     ) -> List[ProcessResult]:
         """
         Process items in batches.
-        
+
         Args:
             items: List of items to process
-            batch_fn: Function that takes a batch (list) and returns a list of results.
-                     The function should return results in the same order as input.
-                     Length of output must match length of input batch.
-            batch_size: Size of each batch. Should be tuned based on model capacity
-                       and memory constraints.
-        
+            batch_fn: Function that takes a batch (list) and returns a list of
+                     results. The function should return results in the same
+                     order as input. Length of output must match length of
+                     input batch.
+            batch_size: Optional in-code override for the batch dimension. When
+                     ``None`` (the default), the batch dimension is taken from
+                     ``ParallelismConfig.max_parallel`` (or ``1`` when
+                     ``max_parallel`` is also ``None``). Pass an explicit value
+                     only when the caller needs a fixed batch size that does
+                     not track ``max_parallel`` -- e.g. ``batch_size=1`` when
+                     the items are themselves pre-batched units, or
+                     ``batch_size=len(items)`` for "all-in-one-call" semantics.
+
         Returns:
             List of ProcessResult objects in same order as items.
             Call .result() on each to get value or raise stored exception.
-        
+
         Note:
             If a batch fails, all items in that batch will have the same exception.
             The entire batch is processed atomically (acquires permits once per batch,
             not once per item), since batching assumes the batch is processed together.
-        
+
         Example:
-            >>> # Process texts in batches of 10
-            >>> def ner_batch(texts):
-            ...     return ner_model.process_batch(texts)
-            >>> 
-            >>> results = processor.process(
-            ...     texts,
-            ...     ner_batch,
-            ...     batch_size=10
-            ... )
-            >>> for r in results:
-            ...     try:
-            ...         entities = r.result()
-            ...     except Exception as e:
-            ...         print(f"Failed: {e}")
+            >>> # Default batch_size = max_parallel
+            >>> config = ParallelismConfig(max_parallel=10)
+            >>> processor = BatchExecutor(config)
+            >>> results = processor.process(texts, ner_batch)
+            >>>
+            >>> # Explicit override
+            >>> results = processor.process(texts, ner_batch, batch_size=10)
         """
         if not items:
             return []
-        
+
+        if batch_size is None:
+            batch_size = self.config.max_parallel or 1
+
         if batch_size < 1:
             raise ValueError(f"batch_size must be >= 1, got {batch_size}")