streamlit-octostar-utils 0.4.2.dev24__tar.gz → 0.5.0__tar.gz

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

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

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

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

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

@@ -9,12 +9,10 @@ import subprocess
 from fastapi import Query
 import time
 import os
-import pickle
 import atexit
 import redis
 import uuid
 import json
-import hashlib
 import shutil
 import threading
 from pottery import Redlock
@@ -57,18 +55,40 @@ 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
 
 
+class TaskResult:
+    """Wrapper for task results that include binary parts alongside JSON data.
+    Tasks returning binary data (e.g. images) should return a TaskResult
+    so that serialized_io writes them as multipart parts instead of attempting
+    JSON serialization on bytes."""
+
+    def __init__(self, data, part=None):
+        self.data = data
+        self.part = part
+
+
 class CelerySerialized:
+    """Serializes task data to a boundary-delimited multipart file.
+
+    Format: metadata JSON part followed by optional binary/streamed parts,
+    separated by boundary markers (the task_id). Replaces pickle entirely.
+    """
+
+    BOUNDARY_PREFIX = b"--"
+    BOUNDARY_SUFFIX = b"\r\n"
+    BOUNDARY_END = b"--\r\n"
+    CONTENT_TYPE_JSON = b"Content-Type: application/json\r\n"
+    CONTENT_TYPE_BYTES = b"Content-Type: application/octet-stream\r\n"
+    HEADER_END = b"\r\n"
+
     def __init__(self, folder, redis_client, data=None):
         self.folder = folder
         self.data = data
@@ -77,18 +97,70 @@ class CelerySerialized:
     def set_task_id(self, task_id):
         self.task_id = task_id
 
-    def dump(self):
+    def _boundary(self):
+        return self.task_id.encode()
+
+    def _write_boundary(self, f):
+        f.write(self.BOUNDARY_PREFIX + self._boundary() + self.BOUNDARY_SUFFIX)
+
+    def _write_end_boundary(self, f):
+        f.write(self.BOUNDARY_PREFIX + self._boundary() + self.BOUNDARY_END)
+
+    def dump(self, parts=None, part_is_list=False):
+        """Write metadata + optional parts in multipart format.
+        parts: optional list of bytes objects to write as additional parts.
+        part_is_list: whether the original part was a list (preserves type on read).
+        """
         assert self.task_id
+        if isinstance(self.data, dict):
+            metadata = self.data
+        else:
+            metadata = {"data": self.data}
+        metadata["part_count"] = len(parts) if parts else 0
+        metadata["part_is_list"] = part_is_list
         with RedisFileLock(self.redis_client, os.path.join(self.folder, self.task_id)):
-            with open(os.path.join(self.folder, self.task_id), "wb") as target_file:
-                pickle.dump(self.data, file=target_file, protocol=pickle.HIGHEST_PROTOCOL)
+            with open(os.path.join(self.folder, self.task_id), "wb") as f:
+                self._write_boundary(f)
+                f.write(self.CONTENT_TYPE_JSON)
+                f.write(self.HEADER_END)
+                f.write(json.dumps(metadata).encode())
+                f.write(b"\r\n")
+                if parts:
+                    for part in parts:
+                        self._write_boundary(f)
+                        f.write(self.CONTENT_TYPE_BYTES)
+                        f.write(self.HEADER_END)
+                        f.write(part)
+                        f.write(b"\r\n")
+                self._write_end_boundary(f)
 
     def load(self):
+        """Read multipart file. Returns (metadata_dict, list_of_bytes_parts)."""
         assert self.task_id
+        boundary = self.BOUNDARY_PREFIX + self._boundary()
+        end_boundary = self.BOUNDARY_PREFIX + self._boundary() + b"--"
         with RedisFileLock(self.redis_client, os.path.join(self.folder, self.task_id)):
-            with open(os.path.join(self.folder, self.task_id), "rb") as source_file:
-                data = pickle.load(source_file)
-        return data
+            with open(os.path.join(self.folder, self.task_id), "rb") as f:
+                raw = f.read()
+        sections = raw.split(boundary)
+        metadata = None
+        parts = []
+        for section in sections:
+            section = section.strip(b"\r\n")
+            if not section or section == b"--":
+                continue
+            header_end = section.find(b"\r\n\r\n")
+            if header_end == -1:
+                continue
+            header = section[:header_end]
+            body = section[header_end + 4:]
+            if body.endswith(b"\r\n"):
+                body = body[:-2]
+            if b"application/json" in header:
+                metadata = json.loads(body)
+            else:
+                parts.append(body)
+        return metadata or {}, parts
 
 
 class CeleryExecutor(object):
@@ -129,7 +201,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()
@@ -147,7 +222,7 @@ class CeleryExecutor(object):
         self.app = Celery(self.filename)
         self.app.conf.broker_url = f"redis://{self.redis_host}:{self.redis_port}/0"
         self.app.conf.result_backend = f"redis://{self.redis_host}:{self.redis_port}/0"
-        self.app.conf.track_started = True
+        self.app.conf.track_started = False
         self.app.conf.task_serializer = "json"
         self.app.conf.result_serializer = "json"
         self.app.conf.accept_content = ["application/json"]
@@ -179,11 +254,11 @@ 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._is_owner = False
         self.set_cleanup_task()
         self.register_state_signals()
 
@@ -213,28 +288,44 @@ 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)
+        if extended_ttl > int(self.app.conf.result_expires):
+            self.redis_client.eval(
+                "for _, k in ipairs(KEYS) do redis.call('expire', k, ARGV[1]) end",
+                2,
+                f"{CeleryExecutor.CELERY_BROKER_PREFIX}{task_id}",
+                f"task:queue:{task_id}",
+                extended_ttl,
+            )
 
     def register_worker_initialization(self):
         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...")
+        if not os.path.isdir(in_dir) or not os.path.isdir(out_dir):
+            logger.warning(f"Data directories missing (in={os.path.isdir(in_dir)}, out={os.path.isdir(out_dir)}), skipping cleanup")
+            return
         redis_client = redis.StrictRedis(host=redis_host, port=redis_port)
         for file_name in os.listdir(in_dir):
             file_path = os.path.join(in_dir, file_name)
@@ -304,6 +395,8 @@ class CeleryExecutor(object):
         }
 
     def start(self):
+        self._is_owner = True
+        atexit.register(self.close)
         logger.info("Initializing data folders...")
         shutil.rmtree(self.root_folder, ignore_errors=True)
         for folder in [self.in_folder, self.out_folder]:
@@ -330,6 +423,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}"
@@ -404,7 +500,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}")
@@ -439,44 +538,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_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")
 
-                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")
-
-                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:
@@ -485,6 +596,9 @@ class CeleryExecutor(object):
                 logger.error(f"Error checking stall for queue '{queue_name}': {e}")
 
     def close(self):
+        if not self._is_owner:
+            return
+        self._is_owner = False
         self.stop_event.set()
         if self.worker_health_check_thread and self.worker_health_check_thread.is_alive():
             self.worker_health_check_thread.join(timeout=2)
@@ -520,9 +634,16 @@ class CeleryExecutor(object):
             task_id = task.request.id
             serialized_data = CelerySerialized(folder=self.in_folder, redis_client=self.redis_client)
             serialized_data.set_task_id(task_id)
-            data = serialized_data.load()
+            metadata, parts = serialized_data.load()
             del serialized_data
-            args, kwargs = data.get("args", []), data.get("kwargs", {})
+            args, kwargs = metadata.get("args", []), metadata.get("kwargs", {})
+
+            part_count = metadata.get("part_count", 0)
+            if part_count > 0:
+                if metadata.get("part_is_list", part_count > 1):
+                    args = [parts] + args
+                else:
+                    args = [parts[0]] + args
 
             if self.app.conf.task_always_eager:
                 queue = task.request.delivery_info.get("routing_key", self.app.conf.task_default_routing_key)
@@ -533,9 +654,25 @@ class CeleryExecutor(object):
             queue = task.request.delivery_info.get("routing_key", self.app.conf.task_default_routing_key)
             task.request.resources = (self.resource_registry or {}).get(queue, {})
             out_data = task_fn(task, *args, **kwargs)
-            serialized_data = CelerySerialized(folder=self.out_folder, data=out_data, redis_client=self.redis_client)
+            if isinstance(out_data, TaskResult):
+                out_part_is_list = isinstance(out_data.part, list)
+                if out_data.part is None:
+                    out_parts = None
+                elif out_part_is_list:
+                    out_parts = out_data.part
+                else:
+                    out_parts = [out_data.part]
+                serialized_data = CelerySerialized(
+                    folder=self.out_folder, data=out_data.data, redis_client=self.redis_client
+                )
+            else:
+                out_parts = None
+                out_part_is_list = False
+                serialized_data = CelerySerialized(
+                    folder=self.out_folder, data=out_data, redis_client=self.redis_client
+                )
             serialized_data.set_task_id(task_id)
-            serialized_data.dump()
+            serialized_data.dump(parts=out_parts, part_is_list=out_part_is_list)
             del serialized_data
             if os.path.isfile(os.path.join(self.in_folder, task_id)):
                 with RedisFileLock(self.redis_client, os.path.join(self.in_folder, task_id)):
@@ -576,36 +713,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, part=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(
@@ -620,28 +744,46 @@ 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
         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
-            )
-            await asyncio.get_running_loop().run_in_executor(
-                self.io_thread_pool,
-                _write_task_data,
-                self.in_folder,
-                args,
-                kwargs,
-                task_id,
-            )
+            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
+                )
+            else:
+                await asyncio.get_running_loop().run_in_executor(
+                    self.io_thread_pool,
+                    _write_task_data,
+                    self.in_folder,
+                    args,
+                    kwargs,
+                    task_id,
+                )
+
             await asyncio.get_running_loop().run_in_executor(
                 self.set_thread_pool, _send_task, task_fn, task_id, options
             )
@@ -652,17 +794,68 @@ 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 _write_task_data_with_part(self, task_id, args, kwargs, part):
+        """Write task data with a streamed part to the multipart file.
+        The part becomes the first arg on the worker side.
+        """
+        boundary = CelerySerialized.BOUNDARY_PREFIX + task_id.encode()
+        boundary_line = boundary + CelerySerialized.BOUNDARY_SUFFIX
+        end_boundary_line = boundary + CelerySerialized.BOUNDARY_END
+
+        is_list = isinstance(part, list)
+        items = part if is_list else [part]
+        part_count = len(items)
+
+        metadata = {"args": args, "kwargs": kwargs, "part_count": part_count, "part_is_list": is_list}
+        metadata_bytes = json.dumps(metadata).encode()
+
+        file_path = os.path.join(self.in_folder, task_id)
+        with open(file_path, "wb") as f:
+            f.write(boundary_line)
+            f.write(CelerySerialized.CONTENT_TYPE_JSON)
+            f.write(CelerySerialized.HEADER_END)
+            f.write(metadata_bytes)
+            f.write(b"\r\n")
+
+            for item in items:
+                f.write(boundary_line)
+                f.write(CelerySerialized.CONTENT_TYPE_BYTES)
+                f.write(CelerySerialized.HEADER_END)
+                if hasattr(item, "__aiter__"):
+                    async for chunk in item:
+                        f.write(chunk if isinstance(chunk, bytes) else chunk.encode())
+                elif hasattr(item, "read"):
+                    while True:
+                        chunk = await item.read(65536)
+                        if not chunk:
+                            break
+                        f.write(chunk if isinstance(chunk, bytes) else chunk.encode())
+                elif isinstance(item, bytes):
+                    f.write(item)
+                else:
+                    raise TypeError(f"Unsupported part item type: {type(item)}")
+                f.write(b"\r\n")
+
+            f.write(end_boundary_line)
+
     async def terminate_task(self, task_id):
         def _terminate_task(celery_app, task_id):
             celery_app.control.revoke(task_id, terminate=True)
 
         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))
@@ -698,6 +891,13 @@ class CeleryExecutor(object):
             self.get_thread_pool, _poll_task_state, self.app, task_id
         )
 
+    async def get_task_info(self, task_id):
+        def _get_info(celery_app, task_id):
+            return celery_app.AsyncResult(task_id).info
+        return await asyncio.get_running_loop().run_in_executor(
+            self.get_thread_pool, _get_info, self.app, task_id,
+        )
+
     async def get_task_result(self, task_id, remove=False):
         def _try_get_task_data(celery_app, task_id):
             async_result = celery_app.AsyncResult(task_id)
@@ -710,11 +910,19 @@ class CeleryExecutor(object):
         def _read_task_data(out_folder, task_id):
             serialized_data = CelerySerialized(folder=out_folder, redis_client=self.redis_client)
             serialized_data.set_task_id(task_id)
-            result = serialized_data.load()
-            return result
+            metadata, parts = serialized_data.load()
+            data = metadata.get("data", metadata)
+            part_count = metadata.get("part_count", 0)
+            if part_count > 0:
+                if metadata.get("part_is_list", part_count > 1):
+                    return TaskResult(data=data, part=parts)
+                else:
+                    return TaskResult(data=data, part=parts[0])
+            return data
 
         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))
@@ -737,8 +945,10 @@ class CeleryExecutor(object):
             )
         return result
 
-    async def send_and_wait_task(self, task_fn, args=[], kwargs={}, timeout=60, **options):
-        task_id = await self.send_task(task_fn, args, kwargs, **options)
+    async def send_and_wait_task(self, task_fn, args=None, kwargs=None, part=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, part=part, **options)
         ready = False
         state = None
         start_time = time.time()
@@ -773,6 +983,11 @@ class FastAPICeleryTaskRoute(Route):
             path="/task/{task_id}",
             methods=["DELETE"],
             summary="Cancel a queued or running task.",
+            description=(
+                "Terminate a task by its ID. If the task is still queued "
+                "(AWAITING) it is revoked; if it is running (STARTED) the "
+                "worker process is interrupted."
+            ),
             status_code=200,
             responses=DefaultErrorRoute.error_responses,
         )
@@ -785,7 +1000,14 @@ class FastAPICeleryTaskRoute(Route):
             self,
             path="/task/{task_id}",
             methods=["GET"],
-            summary="Get task status (and result if available).",
+            summary="Get task status, progress, and result.",
+            description=(
+                "Poll a running or completed task. Returns task_state "
+                "(AWAITING, STARTED, SUCCESS, FAILURE, etc.), and when "
+                "STARTED may include a 'progress' dict with task-specific "
+                "metrics (e.g. nodes_done / nodes_total). When SUCCESS, "
+                "the 'data' field contains the task result."
+            ),
             status_code=200,
             responses=DefaultErrorRoute.error_responses,
         )
@@ -815,10 +1037,12 @@ class FastAPICeleryTaskRoute(Route):
                 data = {"task_state": "UNKNOWN", "task_id": task_id}
             elif state in ["AWAITING", "STARTED"]:
                 data = {"task_state": state, "task_id": task_id}
+                if state == "STARTED":
+                    info = await self.celery_executor.get_task_info(task_id)
+                    if info and isinstance(info, dict):
+                        data["progress"] = info
             elif state == "SUCCESS":
                 data = {"task_state": state, "task_id": task_id, "data": result}
-            elif state == "STARTED":
-                data = {"task_status": state, "task_id": task_id}
             else:
                 raise ValueError(f"Unknown task state {state}!")
             return CommonModels.DataResponseModel(data=data)