streamlit-octostar-utils 0.5.0.dev12__tar.gz → 0.5.0.dev14__tar.gz

{streamlit_octostar_utils-0.5.0.dev12 → streamlit_octostar_utils-0.5.0.dev14}/PKG-INFO

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

{streamlit_octostar_utils-0.5.0.dev12 → streamlit_octostar_utils-0.5.0.dev14}/pyproject.toml

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

{streamlit_octostar_utils-0.5.0.dev12 → streamlit_octostar_utils-0.5.0.dev14}/streamlit_octostar_utils/api_crafter/celery.py

@@ -13,7 +13,6 @@ import atexit
 import redis
 import uuid
 import json
-import hashlib
 import shutil
 import threading
 from pottery import Redlock
@@ -56,14 +55,12 @@ class CeleryQueueConfig:
         max_tasks_in_queue=None,
         max_tasks_per_child=None,
         max_memory_per_child=None,
-        stall_timeout=1200,
         **options,
     ):
         self.n_workers = n_workers
         self.max_tasks_in_queue = max_tasks_in_queue
         self.max_tasks_per_child = max_tasks_per_child
         self.max_memory_per_child = max_memory_per_child  # KiB
-        self.stall_timeout = stall_timeout # seconds; None or 0 to disable
         self.options = options
 
 
@@ -257,9 +254,9 @@ class CeleryExecutor(object):
         self.worker_info = {}
 
         # Queue stall detection
-        self._queue_fingerprints = {}
-        self._queue_fingerprint_changed_at = {}
         self._queue_stalled = {}
+        self._last_stall_check = 0
+        self._stall_check_interval = 60
 
         atexit.register(self.close)
         self.set_cleanup_task()
@@ -291,6 +288,11 @@ class CeleryExecutor(object):
     def set_started_state(self, task_id, task, *args, **kwargs):
         result = AsyncResult(task_id, app=self.app)
         result.backend.store_result(task_id, result=None, state=CeleryExecutor.STARTED)
+        queue = task.request.delivery_info.get(
+            "routing_key", self.app.conf.task_default_routing_key
+        ) if task else None
+        if queue:
+            self.redis_client.set(f"queue:last_started:{queue}", str(time.time()))
         request_timelimit = getattr(getattr(task, "request", None), "timelimit", None) or (None, None)
         time_limit = request_timelimit[0] or getattr(task, "time_limit", None) or 0
         extended_ttl = int(time_limit) + int(self.app.conf.result_expires)
@@ -307,20 +309,17 @@ class CeleryExecutor(object):
         if self.preload_functions:
             celery_signals.worker_process_init.connect(self.preload_on_worker_init)
 
-    def set_last_completed_time(self, sender=None, task_id=None, task=None, **kwargs):
+    def cleanup_task_keys(self, sender=None, task_id=None, **kwargs):
         try:
-            queue = task.request.delivery_info.get(
-                "routing_key", self.app.conf.task_default_routing_key
-            ) if task else None
-            if queue:
-                self.redis_client.set(f"queue:last_completed:{queue}", str(time.time()))
+            if task_id:
+                self.redis_client.delete(f"task:queue:{task_id}")
         except Exception:
             pass
 
     def register_state_signals(self):
         celery_signals.before_task_publish.connect(self.set_awaiting_state)
         celery_signals.task_prerun.connect(self.set_started_state)
-        celery_signals.task_postrun.connect(self.set_last_completed_time)
+        celery_signals.task_postrun.connect(self.cleanup_task_keys)
 
     def cleanup_task_results(in_dir, out_dir, redis_host, redis_port, task_expires, result_expires):
         logger.info("Starting cleanup of expired task results...")
@@ -419,6 +418,9 @@ class CeleryExecutor(object):
             if attempts_done == CeleryExecutor.MAX_STARTUP_CHECKS:
                 raise TimeoutError("Redis not ready after a long wait!")
         self.redis_client.flushall()
+        boot_time = str(time.time())
+        for queue_name in self.queue_config:
+            self.redis_client.set(f"queue:last_started:{queue_name}", boot_time)
         for queue, queue_config in self.queue_config.items():
             for slot in range(queue_config.n_workers):
                 worker_name = f"celery@{self.name}:{queue}:{slot}"
@@ -493,7 +495,10 @@ class CeleryExecutor(object):
         while not self.stop_event.is_set():
             try:
                 self._restart_dead_processes()
-                self._check_queue_stalls()
+                now = time.time()
+                if now - self._last_stall_check >= self._stall_check_interval:
+                    self._last_stall_check = now
+                    self._check_queue_stalls()
                 time.sleep(5)
             except Exception as e:
                 logger.error(f"Error in worker health check: {e}")
@@ -528,44 +533,56 @@ class CeleryExecutor(object):
             logger.info(f"Restarted beat process (PID: {self.beat_process.pid})")
 
     def _check_queue_stalls(self):
-        for queue_name, queue_config in self.queue_config.items():
-            if not queue_config.stall_timeout:
-                continue
+        try:
+            inspector = self.app.control.inspect(timeout=5)
+            active_data = inspector.active() or {}
+        except Exception as e:
+            logger.error(f"Failed to inspect active workers: {e}")
+            return
+
+        now = time.time()
+        grace_period = 3 * self.app.conf.worker_proc_alive_timeout
+
+        active_queues = set()
+        for _worker_name, tasks in active_data.items():
+            for task_info in (tasks or []):
+                queue = (task_info.get("delivery_info") or {}).get(
+                    "routing_key", self.app.conf.task_default_routing_key
+                )
+                active_queues.add(queue)
+
+        for queue_name in self.queue_config:
             try:
-                queue_items = self.redis_client.lrange(queue_name, 0, -1)
-                if len(queue_items) == 0:
+                pending_count = self.redis_client.llen(queue_name)
+                was_stalled = self._queue_stalled.get(queue_name, False)
+
+                if pending_count == 0:
+                    self.redis_client.delete(f"queue:first_enqueued:{queue_name}")
+                    if was_stalled:
+                        logger.info(f"Queue '{queue_name}' has recovered from stall.")
                     self._queue_stalled[queue_name] = False
-                    self._queue_fingerprints.pop(queue_name, None)
-                    self._queue_fingerprint_changed_at.pop(queue_name, None)
                     continue
 
-                fingerprint = hashlib.md5(b"".join(sorted(queue_items))).hexdigest()
-                now_time = time.time()
-                prev_fingerprint = self._queue_fingerprints.get(queue_name)
-
-                if fingerprint != prev_fingerprint:
-                    self._queue_fingerprints[queue_name] = fingerprint
-                    self._queue_fingerprint_changed_at[queue_name] = now_time
+                if queue_name in active_queues:
+                    if was_stalled:
+                        logger.info(f"Queue '{queue_name}' has recovered from stall.")
                     self._queue_stalled[queue_name] = False
                     continue
 
-                fingerprint_age = now_time - self._queue_fingerprint_changed_at.get(queue_name, now_time)
-
-                last_completed_raw = self.redis_client.get(f"queue:last_completed:{queue_name}")
-                last_completed = float(last_completed_raw) if last_completed_raw else 0
-                time_since_completion = (now_time - last_completed) if last_completed else float("inf")
+                last_started_raw = self.redis_client.get(f"queue:last_started:{queue_name}")
+                first_enqueued_raw = self.redis_client.get(f"queue:first_enqueued:{queue_name}")
+                last_started = float(last_started_raw) if last_started_raw else 0
+                first_enqueued = float(first_enqueued_raw) if first_enqueued_raw else 0
+                last_activity = max(last_started, first_enqueued)
+                time_since_activity = (now - last_activity) if last_activity else float("inf")
 
-                was_stalled = self._queue_stalled.get(queue_name, False)
-                is_stalled = (
-                    fingerprint_age >= queue_config.stall_timeout
-                    and time_since_completion >= queue_config.stall_timeout
-                )
+                is_stalled = time_since_activity >= grace_period
                 self._queue_stalled[queue_name] = is_stalled
 
                 if is_stalled and not was_stalled:
                     logger.error(
-                        f"Queue '{queue_name}' is STALLED: {len(queue_items)} task(s) stuck for "
-                        f"{fingerprint_age:.0f}s with no completions in {time_since_completion:.0f}s. "
+                        f"Queue '{queue_name}' is STALLED: {pending_count} task(s) pending, "
+                        f"no active workers, no task started in {time_since_activity:.0f}s. "
                         f"New requests will receive 503."
                     )
                 elif not is_stalled and was_stalled:
@@ -719,10 +736,10 @@ class CeleryExecutor(object):
             task_fn.apply_async(task_id=task_id, **options)
 
         def _store_task_queue_mapping(task_id, queue_name):
-            self.redis_client.set(
-                f"task:queue:{task_id}", queue_name,
-                ex=self.app.conf.result_expires,
-            )
+            pipe = self.redis_client.pipeline()
+            pipe.set(f"task:queue:{task_id}", queue_name, ex=self.app.conf.result_expires)
+            pipe.set(f"queue:first_enqueued:{queue_name}", str(time.time()), nx=True)
+            pipe.execute()
 
         task_id = str(uuid.uuid4())
         queue_name = self.app.conf.task_default_routing_key
@@ -830,6 +847,7 @@ class CeleryExecutor(object):
 
         def _remove_task_data(celery_app, in_folder, out_folder, task_id):
             celery_app.AsyncResult(task_id).forget()
+            self.redis_client.delete(f"task:queue:{task_id}")
             if os.path.isfile(os.path.join(in_folder, task_id)):
                 with RedisFileLock(self.redis_client, os.path.join(self.in_folder, task_id)):
                     os.remove(os.path.join(in_folder, task_id))
@@ -889,6 +907,7 @@ class CeleryExecutor(object):
 
         def _remove_task_data(celery_app, in_folder, out_folder, task_id):
             celery_app.AsyncResult(task_id).forget()
+            self.redis_client.delete(f"task:queue:{task_id}")
             if os.path.isfile(os.path.join(in_folder, task_id)):
                 with RedisFileLock(self.redis_client, os.path.join(self.in_folder, task_id)):
                     os.remove(os.path.join(in_folder, task_id))