streamlit-octostar-utils 0.4.2.dev24__tar.gz → 0.4.2.dev25__tar.gz

{streamlit_octostar_utils-0.4.2.dev24 → streamlit_octostar_utils-0.4.2.dev25}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamlit-octostar-utils
-Version: 0.4.2.dev24
+Version: 0.4.2.dev25
 Summary: 
 License: MIT
 License-File: LICENSE

{streamlit_octostar_utils-0.4.2.dev24 → streamlit_octostar_utils-0.4.2.dev25}/pyproject.toml

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

{streamlit_octostar_utils-0.4.2.dev24 → streamlit_octostar_utils-0.4.2.dev25}/streamlit_octostar_utils/api_crafter/celery.py

@@ -129,7 +129,10 @@ class CeleryExecutor(object):
         self.get_thread_pool = None
         self.set_thread_pool = None
         self.io_thread_pool = None
-        self.queue_threadlocks = {k: threading.Lock() for k in self.queue_config.keys()}
+        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()
+        }
 
         # Folder setup
         self.base_folder = Path(base_folder).resolve()
@@ -576,36 +579,23 @@ class CeleryExecutor(object):
 
         return decorator
 
-    async def send_task(self, task_fn, args=[], kwargs={}, **options) -> str:
+    async def send_task(self, task_fn, args=None, kwargs=None, **options) -> str:
+        args = args if args is not None else []
+        kwargs = kwargs if kwargs is not None else {}
         if self.app.conf.task_always_eager and "dev_preload" not in self.app.conf:
             self.preload_on_worker_init()
             self.app.conf.dev_preload = True
 
-        def _reserve_queue_slot(queue_name):
+        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."
                 )
-            limit = self.queue_config[queue_name].max_tasks_in_queue
-            if limit:
-                reservation_key = f"queue:reserved:{queue_name}"
-                with self.queue_threadlocks[queue_name]:
-                    queue_count = self.redis_client.llen(queue_name)
-                    reserved_count = int(self.redis_client.get(reservation_key) or 0)
-                    total_count = queue_count + reserved_count
-                    if total_count >= limit:
-                        raise CeleryExecutor.QueueFullException(
-                            f"Queue '{queue_name}' has reached its limit of {limit} tasks!"
-                        )
-                    self.redis_client.incr(reservation_key)
-                return True
-            return False
-
-        def _release_queue_slot(queue_name):
-            limit = self.queue_config[queue_name].max_tasks_in_queue
-            if limit:
-                reservation_key = f"queue:reserved:{queue_name}"
-                self.redis_client.decr(reservation_key)
+            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(
@@ -629,11 +619,22 @@ class CeleryExecutor(object):
         queue_name = self.app.conf.task_default_routing_key
         queue_name = getattr(task_fn, "queue", queue_name)
         queue_name = options.get("queue", queue_name)
-        reserved = False
+
+        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
+
         try:
-            reserved = await asyncio.get_running_loop().run_in_executor(
-                self.set_thread_pool, _reserve_queue_slot, queue_name
-            )
+            if acquired:
+                await asyncio.get_running_loop().run_in_executor(
+                    self.set_thread_pool, _check_queue_llen, queue_name
+                )
             await asyncio.get_running_loop().run_in_executor(
                 self.io_thread_pool,
                 _write_task_data,
@@ -652,9 +653,15 @@ class CeleryExecutor(object):
             logger.info(f"Cancelling task {task_id} due to disconnect!")
             await self.terminate_task(task_id)
             raise
+        except Exception:
+            try:
+                await self.terminate_task(task_id)
+            except Exception:
+                pass
+            raise
         finally:
-            if reserved:
-                await asyncio.get_running_loop().run_in_executor(self.set_thread_pool, _release_queue_slot, queue_name)
+            if acquired:
+                sem.release()
         return task_id
 
     async def terminate_task(self, task_id):
@@ -737,7 +744,9 @@ class CeleryExecutor(object):
             )
         return result
 
-    async def send_and_wait_task(self, task_fn, args=[], kwargs={}, timeout=60, **options):
+    async def send_and_wait_task(self, task_fn, args=None, kwargs=None, timeout=60, **options):
+        args = args if args is not None else []
+        kwargs = kwargs if kwargs is not None else {}
         task_id = await self.send_task(task_fn, args, kwargs, **options)
         ready = False
         state = None