streamlit-octostar-utils 0.5.0.dev19__tar.gz → 0.5.1.dev1__tar.gz

{streamlit_octostar_utils-0.5.0.dev19 → streamlit_octostar_utils-0.5.1.dev1}/PKG-INFO

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

{streamlit_octostar_utils-0.5.0.dev19 → streamlit_octostar_utils-0.5.1.dev1}/pyproject.toml

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

{streamlit_octostar_utils-0.5.0.dev19 → streamlit_octostar_utils-0.5.1.dev1}/streamlit_octostar_utils/api_crafter/celery.py

@@ -48,6 +48,65 @@ class RedisFileLock:
         return self.lock.__exit__(exc_type, exc_val, exc_tb)
 
 
+def _install_queue_limit_patch(queue_limits):
+    """Monkey-patch kombu's Redis Channel._put for atomic queue limits.
+
+    The original _put does a bare LPUSH.  The patched version wraps
+    limited queues in a Lua script (LLEN check + LPUSH) so the
+    "is there room?" test and the "enqueue" are a single atomic
+    Redis operation.  Queues without a configured limit are
+    forwarded to the original _put unchanged.
+    """
+    from kombu.transport.redis import Channel
+    from kombu.utils.json import dumps as kombu_dumps
+
+    _ATOMIC_ENQUEUE_LUA = """
+local total = 0
+for i = 1, #KEYS do
+    total = total + redis.call('LLEN', KEYS[i])
+end
+if total >= tonumber(ARGV[1]) then
+    return 0
+end
+redis.call('LPUSH', KEYS[1], ARGV[2])
+return 1
+"""
+
+    if not hasattr(Channel, "_queue_limits"):
+        Channel._queue_limits = {}
+    Channel._queue_limits.update(queue_limits)
+
+    if getattr(Channel, "_queue_limit_patched", False):
+        return
+
+    _original_put = Channel._put
+
+    def _put_with_limit(self, queue, message, **kwargs):
+        max_tasks = Channel._queue_limits.get(queue)
+        if max_tasks is not None:
+            pri = self._get_message_priority(message, reverse=False)
+            target_key = self._q_for_pri(queue, pri)
+            all_keys = [self._q_for_pri(queue, p) for p in self.priority_steps]
+            if target_key in all_keys:
+                all_keys.remove(target_key)
+            all_keys.insert(0, target_key)
+            with self.conn_or_acquire() as client:
+                accepted = client.eval(
+                    _ATOMIC_ENQUEUE_LUA, len(all_keys), *all_keys,
+                    max_tasks, kombu_dumps(message),
+                )
+                if not accepted:
+                    from streamlit_octostar_utils.api_crafter.celery import CeleryExecutor
+                    raise CeleryExecutor.QueueFullException(
+                        f"Queue '{queue}' has reached its limit of {max_tasks} tasks!"
+                    )
+        else:
+            _original_put(self, queue, message, **kwargs)
+
+    Channel._put = _put_with_limit
+    Channel._queue_limit_patched = True
+
+
 class CeleryQueueConfig:
     def __init__(
         self,
@@ -201,10 +260,11 @@ class CeleryExecutor(object):
         self.get_thread_pool = None
         self.set_thread_pool = None
         self.io_thread_pool = None
-        self.queue_semaphores = {
-            k: threading.Semaphore(v.max_tasks_in_queue) if v.max_tasks_in_queue else None
-            for k, v in self.queue_config.items()
-        }
+
+        # Atomic queue-depth limits
+        queue_limits = {k: v.max_tasks_in_queue for k, v in self.queue_config.items() if v.max_tasks_in_queue}
+        if queue_limits:
+            _install_queue_limit_patch(queue_limits)
 
         # Folder setup
         self.base_folder = Path(base_folder).resolve()
@@ -720,17 +780,6 @@ class CeleryExecutor(object):
             self.preload_on_worker_init()
             self.app.conf.dev_preload = True
 
-        def _check_queue_llen(queue_name):
-            if self._queue_stalled.get(queue_name, False):
-                raise CeleryExecutor.QueueStalledException(
-                    f"Queue '{queue_name}' is stalled. Service temporarily unavailable."
-                )
-            if self.redis_client.llen(queue_name) >= self.queue_config[queue_name].max_tasks_in_queue:
-                raise CeleryExecutor.QueueFullException(
-                    f"Queue '{queue_name}' has reached its limit of "
-                    f"{self.queue_config[queue_name].max_tasks_in_queue} tasks!"
-                )
-
         def _write_task_data(in_folder, task_args, task_kwargs, task_id):
             serialized_data = CelerySerialized(
                 folder=in_folder,
@@ -754,22 +803,12 @@ class CeleryExecutor(object):
         queue_name = getattr(task_fn, "queue", queue_name)
         queue_name = options.get("queue", queue_name)
 
-        sem = self.queue_semaphores.get(queue_name)
-        acquired = False
-        if sem is not None:
-            if not sem.acquire(blocking=False):
-                raise CeleryExecutor.QueueFullException(
-                    f"Queue '{queue_name}' has reached its limit of "
-                    f"{self.queue_config[queue_name].max_tasks_in_queue} tasks!"
-                )
-            acquired = True
+        if self._queue_stalled.get(queue_name, False):
+            raise CeleryExecutor.QueueStalledException(
+                f"Queue '{queue_name}' is stalled. Service temporarily unavailable."
+            )
 
         try:
-            if acquired:
-                await asyncio.get_running_loop().run_in_executor(
-                    self.set_thread_pool, _check_queue_llen, queue_name
-                )
-
             if part is not None:
                 await self._write_task_data_with_part(
                     task_id, args, kwargs, part
@@ -800,9 +839,6 @@ class CeleryExecutor(object):
             except Exception:
                 pass
             raise
-        finally:
-            if acquired:
-                sem.release()
         return task_id
 
     async def _write_task_data_with_part(self, task_id, args, kwargs, part):