broccoli-workers 0.1.0__py3-none-any.whl

broccoli/core/task/task.py

@@ -0,0 +1,58 @@
+# video_scheduler/core/task.py
+import json
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Optional
+
+
+@dataclass
+class Task:
+    task_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    task_type: str = ""
+    status: str = "pending"
+    progress: float = 0.0
+    retries: int = 0
+    max_retries: int = 3
+    error: Optional[str] = None
+    secondary_error: Optional[str] = None
+    payload: dict = field(default_factory=dict)
+    result: Any = None
+    depends_on: Optional[str] = None
+    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
+    updated_at: Optional[str] = None
+
+    def to_dict(self) -> dict:
+        return {
+            "task_id": self.task_id,
+            "task_type": self.task_type,
+            "status": self.status,
+            "progress": str(self.progress),
+            "retries": str(self.retries),
+            "max_retries": str(self.max_retries),
+            "error": self.error or "",
+            "secondary_error": self.secondary_error or "",
+            "payload": json.dumps(self.payload),
+            "result": json.dumps(self.result) if self.result is not None else "",
+            "depends_on": self.depends_on or "",
+            "created_at": self.created_at or "",
+            "updated_at": self.updated_at or "",
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "Task":
+        return cls(
+            task_id=data.get("task_id"),
+            task_type=data.get("task_type", ""),
+            status=data.get("status", "pending"),
+            progress=float(data.get("progress", 0.0)),
+            retries=int(data.get("retries", 0)),
+            max_retries=int(data.get("max_retries", 3)),
+            error=data.get("error") or None,
+            secondary_error=data.get("secondary_error") or None,
+            payload=json.loads(data.get("payload", "{}")),
+            result=json.loads(data.get("result")) if data.get("result") else None,
+            depends_on=data.get("depends_on") or None,
+            created_at=data.get("created_at") or datetime.now().isoformat(),
+            updated_at=data.get("updated_at") or None,
+        )

broccoli/core/task/task_queue.py

@@ -0,0 +1,300 @@
+# broccoli/core/task/task_queue.py
+import logging
+import time
+from typing import Optional
+
+from broccoli.core.redis_controller import RedisController
+from broccoli.core.task.task import Task
+
+logger = logging.getLogger(__name__)
+
+# Max tasks per priority tier before FIFO scores would collide.
+# 10^12 tasks per tier is effectively unlimited in practice.
+_FIFO_TIER_SIZE = 1_000_000_000_000
+
+# Redis field used to carry a task's original priority through the waiting state
+# so it can be restored when the dependency is released.
+_PRIORITY_FIELD = "_broccoli_priority"
+
+
+class TaskQueue:
+    """
+    Priority queue backed by three Redis structures:
+
+    ``<base>:queue``       — sorted set of runnable task IDs
+                             score = priority_tier * _FIFO_TIER_SIZE + monotonic_seq
+    ``<base>:processing``  — sorted set of in-flight task IDs
+                             score = Unix timestamp of when the task was popped
+                             (used for crash recovery, not for ordering)
+    ``<base>:sequence``    — monotonic counter for FIFO ordering within a priority tier
+    ``dependency:<id>``    — set of task IDs waiting for task <id> to complete
+
+    Dependency resolution happens at *push time* and *completion time*, never
+    inside ``pop()``.  Workers only ever see runnable tasks.
+    """
+
+    def __init__(
+        self,
+        redis_url: str = "redis://localhost:6379",
+        queue_name: str = "tasks:queue",
+        task_prefix: str = "task",
+    ):
+        self.redis_url = redis_url
+        self._redis = RedisController(redis_url).get_client()
+        self.queue_key = queue_name
+        self.task_prefix = task_prefix
+
+        # Derive sibling keys from queue_name so all keys share a namespace.
+        # e.g. "tasks:queue" → "tasks:processing", "tasks:sequence"
+        base = (
+            queue_name[: -len(":queue")]
+            if queue_name.endswith(":queue")
+            else queue_name
+        )
+        self.base = base
+        self.processing_key = f"{base}:processing"
+        self.sequence_key = f"{base}:sequence"
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def push(self, task: Task, priority: int = 0) -> str:
+        """
+        Persist task metadata then either enqueue it or register it as waiting.
+
+        If the task has no dependency, or its dependency is already completed,
+        it goes straight into the runnable queue.  Otherwise it is stored as
+        ``waiting`` and added to ``dependency:<parent_id>`` so ``complete()``
+        can release it later.
+
+        Race condition handled with a register-then-verify pattern:
+          1. Register in the dependency set first.
+          2. Re-check the parent's status.
+          3. If parent is already completed (raced past us), release immediately.
+        This eliminates the TOCTOU window present in a check-then-register approach.
+        """
+        # Always persist task metadata first so workers can read it after pop().
+        self._redis.hset(f"{self.task_prefix}:{task.task_id}", mapping=task.to_dict())
+
+        if task.depends_on:
+            # Mark as waiting immediately and register in the dependency set.
+            # We do this BEFORE reading the parent's status to close the
+            # check-then-register race window.
+            task.status = "waiting"
+            pipe = self._redis.pipeline()
+            pipe.hset(f"{self.task_prefix}:{task.task_id}", "status", "waiting")
+            # Store the requested priority so complete() can restore it.
+            pipe.hset(
+                f"{self.task_prefix}:{task.task_id}", _PRIORITY_FIELD, str(priority)
+            )
+            pipe.sadd(f"dependency:{task.depends_on}", task.task_id)
+            pipe.execute()
+
+            # Now verify: if the parent already completed before our sadd, release now.
+            dep_status = self._redis.hget(
+                f"{self.task_prefix}:{task.depends_on}", "status"
+            )
+            parent_done = (
+                dep_status
+                and (
+                    dep_status.decode() if isinstance(dep_status, bytes) else dep_status
+                )
+                == "completed"
+            )
+
+            if parent_done:
+                # Parent completed before (or concurrently with) our registration.
+                # Release ourselves immediately and clean up.
+                self._redis.srem(f"dependency:{task.depends_on}", task.task_id)
+                self._enqueue(task.task_id, priority)
+                logger.debug(
+                    f"Task {task.task_id}: dependency {task.depends_on} already "
+                    "completed, enqueued immediately"
+                )
+            else:
+                logger.debug(
+                    f"Task {task.task_id} waiting on dependency {task.depends_on}"
+                )
+        else:
+            self._enqueue(task.task_id, priority)
+
+        return task.task_id
+
+    def pop(self) -> Optional[Task]:
+        """
+        Pop the highest-priority runnable task and move it to the processing set.
+
+        No dependency checking happens here — that was handled at push / complete
+        time.  Everything in the queue is ready to run.
+
+        Returns ``None`` if the queue is empty after the blocking timeout.
+        """
+        result = self._redis.bzpopmin(self.queue_key, timeout=1)
+        if result is None:
+            return None
+
+        _, task_id, _fifo_score = result
+        task_id = task_id.decode() if isinstance(task_id, bytes) else task_id
+
+        task_data = self._redis.hgetall(f"{self.task_prefix}:{task_id}")
+        if not task_data:
+            logger.error(
+                f"Task data missing for {task_id!r}; moving to dead-letter set"
+            )
+            self._redis.sadd(f"{self.task_prefix}:dead_letter", task_id)
+            return None
+
+        # Score in the processing set is a Unix timestamp so recover_stalled()
+        # can identify tasks that have been in-flight longer than expected.
+        now = time.time()
+        self._redis.zadd(self.processing_key, {task_id: now})
+        self._redis.hset(f"{self.task_prefix}:{task_id}", "status", "in_progress")
+
+        task = Task.from_dict(task_data)
+        task.status = "in_progress"
+        return task
+
+    def complete(self, task: Task) -> None:
+        """
+        Called after a task finishes successfully.
+
+        1. Removes the task from the processing set.
+        2. Reads any tasks that were waiting on this one and enqueues them,
+           restoring their original priority.
+        3. Deletes the dependency set.
+        """
+        task_id = task.task_id
+        self._redis.zrem(self.processing_key, task_id)
+
+        dep_key = f"dependency:{task_id}"
+        waiting_ids = self._redis.smembers(dep_key)
+        if waiting_ids:
+            for raw_id in waiting_ids:
+                wid = raw_id.decode() if isinstance(raw_id, bytes) else raw_id
+                # Restore the priority the caller supplied when pushing the waiting task.
+                raw_pri = self._redis.hget(f"{self.task_prefix}:{wid}", _PRIORITY_FIELD)
+                priority = int(raw_pri) if raw_pri else 0
+                self._enqueue(wid, priority)
+                logger.info(
+                    f"Released waiting task {wid} "
+                    f"(dependency {task_id} completed, priority={priority})"
+                )
+            self._redis.delete(dep_key)
+
+    def fail(self, task: Task) -> None:
+        """Remove a permanently failed (or skipped) task from the processing set."""
+        self._redis.zrem(self.processing_key, task.task_id)
+
+    def requeue(self, task_id: str, priority: Optional[int] = None) -> None:
+        """
+        Move a task from the processing set back to the runnable queue (retry).
+
+        A fresh FIFO sequence number is issued so the retry goes to the back of
+        its priority tier rather than jumping the queue.
+
+        If ``priority`` isn't given explicitly, the task's originally-requested
+        priority (stashed under ``_PRIORITY_FIELD`` at push time) is restored,
+        matching the behaviour of ``complete()`` when it releases dependents.
+        Without this, every retry would silently reset to priority 0 and could
+        jump ahead of lower-priority tasks still waiting in the queue.
+        """
+        self._redis.zrem(self.processing_key, task_id)
+        if priority is None:
+            raw_pri = self._redis.hget(f"{self.task_prefix}:{task_id}", _PRIORITY_FIELD)
+            priority = int(raw_pri) if raw_pri else 0
+        self._enqueue(task_id, priority)
+
+    def recover_stalled(self, timeout_seconds: int = 3600) -> int:
+        """
+        Crash-recovery helper: re-enqueue tasks that have been in the processing
+        set for longer than ``timeout_seconds``.
+
+        The processing set uses Unix timestamps as scores (set by ``pop()``),
+        so the threshold is simply ``time.time() - timeout_seconds``.
+
+        Returns the number of tasks recovered.
+        """
+        threshold = time.time() - timeout_seconds
+        stalled = self._redis.zrangebyscore(self.processing_key, "-inf", threshold)
+        for raw_id in stalled:
+            tid = raw_id.decode() if isinstance(raw_id, bytes) else raw_id
+            self.requeue(tid)
+            logger.warning(f"Crash-recovery: re-enqueued stalled task {tid}")
+        return len(stalled)
+
+    # ------------------------------------------------------------------
+    # Convenience helpers
+    # ------------------------------------------------------------------
+
+    def get_task(self, task_id: str) -> Optional[Task]:
+        """Fetch task metadata by ID, or ``None`` if it doesn't exist."""
+        task_data = self._redis.hgetall(f"{self.task_prefix}:{task_id}")
+        if not task_data:
+            return None
+        return Task.from_dict(task_data)
+
+    def get_queue_name(self) -> str:
+        return self.queue_key
+
+    def is_empty(self) -> bool:
+        return self._redis.zcard(self.queue_key) == 0
+
+    def pop_with_timeout(self, timeout: int = 1) -> Optional[Task]:
+        """Alias for ``pop()``; timeout is already handled internally."""
+        return self.pop()
+
+    # ------------------------------------------------------------------
+    # Private helpers
+    # ------------------------------------------------------------------
+
+    def _enqueue(self, task_id: str, priority: int = 0) -> None:
+        """
+        Add *task_id* to the runnable sorted set.
+
+        Score = priority * _FIFO_TIER_SIZE + monotonic_seq
+
+        This gives strict priority ordering across tiers *and* FIFO within a
+        tier, without requiring timestamps or Lua scripts.
+        """
+        seq = self._redis.incr(self.sequence_key)
+        score = priority * _FIFO_TIER_SIZE + seq
+        self._redis.zadd(self.queue_key, {task_id: score})
+        self._redis.hset(
+            f"{self.task_prefix}:{task_id}",
+            mapping={"status": "pending", _PRIORITY_FIELD: str(priority)},
+        )
+        logger.debug(f"Enqueued task {task_id} (priority={priority}, seq={seq})")
+
+    def is_fully_drained(self) -> bool:
+        """True only when both the runnable queue and processing set are empty."""
+        pipe = self._redis.pipeline()
+        pipe.zcard(self.queue_key)
+        pipe.zcard(self.processing_key)
+        runnable, processing = pipe.execute()
+        return runnable == 0 and processing == 0
+
+    def get_waiting_for(self, task_id: str) -> list:
+        """Return the IDs of tasks currently blocked on ``task_id``."""
+        return [
+            tid.decode() if isinstance(tid, bytes) else tid
+            for tid in self._redis.smembers(f"dependency:{task_id}")
+        ]
+
+    def stats(self) -> dict:
+        """
+        Snapshot of queue depth for monitoring/diagnostics.
+
+        ``dead_letter`` reflects tasks whose hash went missing on pop() —
+        see ``pop()`` for how they get there.
+        """
+        pipe = self._redis.pipeline()
+        pipe.zcard(self.queue_key)
+        pipe.zcard(self.processing_key)
+        pipe.scard(f"{self.task_prefix}:dead_letter")
+        runnable, processing, dead_letter = pipe.execute()
+        return {
+            "runnable": runnable,
+            "processing": processing,
+            "dead_letter": dead_letter,
+        }

broccoli/core/task/task_registry.py

@@ -0,0 +1,47 @@
+# video_scheduler/core/task_registry.py
+import logging
+from typing import Any, Callable, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class TaskRegistry:
+    """Registry for task handlers that can be executed by workers."""
+
+    _instance = None
+    _handlers: Dict[str, Callable] = {}
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def register(self, task_type: str):
+        """Decorator to register a task handler."""
+
+        def decorator(func: Callable):
+            self._handlers[task_type] = func
+            logger.info(f"Registered handler for task type: {task_type}")
+            return func
+
+        return decorator
+
+    def register_manually(self, task_type: str, handler: Callable):
+        """Manually register a task handler."""
+        self._handlers[task_type] = handler
+        logger.info(f"Manually registered handler for task type: {task_type}")
+
+    def get_handler(self, task_type: str) -> Optional[Callable]:
+        """Get the handler for a task type."""
+        return self._handlers.get(task_type)
+
+    def get_all_handlers(self) -> Dict[str, Callable]:
+        """Get all registered handlers."""
+        return self._handlers
+
+    def execute(self, task_type: str, payload: Dict[str, Any], **kwargs) -> Any:
+        """Execute a task handler."""
+        handler = self.get_handler(task_type)
+        if not handler:
+            raise ValueError(f"No handler registered for task type: {task_type}")
+        return handler(payload, **kwargs)

broccoli/workers/__init__.py

@@ -0,0 +1,15 @@
+from broccoli.workers.async_worker import AsyncWorker
+from broccoli.workers.base_worker import BaseWorker
+from broccoli.workers.chain_worker import ChainWorker
+from broccoli.workers.hybrid_worker import HybridWorker
+from broccoli.workers.threaded_worker import ThreadedWorker
+from broccoli.workers.worker_pool import WorkerPool
+
+__all__ = [
+    "BaseWorker",
+    "ChainWorker",
+    "ThreadedWorker",
+    "AsyncWorker",
+    "HybridWorker",
+    "WorkerPool",
+]

broccoli/workers/async_worker.py

@@ -0,0 +1,157 @@
+# broccoli/workers/async_worker.py
+# PEP 563 defers annotation evaluation, so the PEP 604 `X | None` and
+# lowercase-generic (`set[str]`) hints below don't raise on Python 3.9,
+# which doesn't support that syntax natively (3.10+ only).
+from __future__ import annotations
+
+import asyncio
+import logging
+
+import redis
+
+from broccoli.core.task.task import Task
+from broccoli.workers.base_worker import BaseWorker
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncWorker(BaseWorker):
+    """
+    Worker that processes tasks concurrently using asyncio.
+
+    CPU-bound / blocking task handlers are offloaded to the default thread-pool
+    executor so they don't stall the event loop.
+
+    State transitions (complete / fail / requeue) are fully delegated to
+    ``BaseWorker._handle_task_result`` — no duplicate logic here.
+    """
+
+    def __init__(
+        self,
+        redis_url: str = "redis://localhost:6379",
+        worker_id: str = None,
+        max_concurrent: int = 10,
+    ):
+        super().__init__(redis_url, worker_id)
+        self.max_concurrent = max_concurrent
+        # Semaphore is created lazily inside the running event loop to avoid
+        # "no running event loop" errors at construction time.
+        self._semaphore: asyncio.Semaphore | None = None
+        self.active_tasks: set[str] = set()
+
+    # ------------------------------------------------------------------
+    # Async task lifecycle
+    # ------------------------------------------------------------------
+
+    async def process_task_async(self, task: Task) -> None:
+        """Acquire a concurrency slot and process one task."""
+        async with self._semaphore:
+            try:
+                logger.info(f"AsyncWorker {self.worker_id} processing {task.task_id}")
+
+                if not self.pre_process(task):
+                    logger.info(f"Task {task.task_id} skipped by pre_process")
+                    self.queue.fail(task)  # remove from processing set
+                    return
+
+                # Use get_running_loop() — get_event_loop() is deprecated in 3.10+
+                # and raises a DeprecationWarning inside a running coroutine.
+                loop = asyncio.get_running_loop()
+                # Only self.process() is allowed to fail the task via exception.
+                # _handle_task_result / post_process must be allowed to propagate
+                # their own errors to the outer except below, rather than being
+                # re-invoked a second time on an already-transitioned task.
+                try:
+                    success = await asyncio.wait_for(
+                        loop.run_in_executor(None, self.process, task),
+                        timeout=self.task_timeout,
+                    )
+                except asyncio.TimeoutError:
+                    logger.error(
+                        f"Task {task.task_id} timed out after {self.task_timeout}s"
+                    )
+                    task.error = f"Task timed out after {self.task_timeout}s"
+                    success = False
+                except Exception as e:
+                    logger.error(
+                        f"Task {task.task_id} handler raised: {e}", exc_info=True
+                    )
+                    task.error = str(e)
+                    success = False
+
+                # Central state machine: complete / requeue / fail
+                self._handle_task_result(task, success)
+
+                # Result storage and user-facing callbacks
+                self.post_process(task, success)
+
+            except Exception as e:
+                # Reached only if _handle_task_result/post_process themselves
+                # raised — the task's queue state may already be inconsistent,
+                # so we log rather than re-running the state machine on it again.
+                logger.error(
+                    f"Task {task.task_id} failed outside handler: {e}", exc_info=True
+                )
+            finally:
+                self.active_tasks.discard(task.task_id)
+
+    # ------------------------------------------------------------------
+    # Event-loop entry points
+    # ------------------------------------------------------------------
+
+    async def start_async(self):
+        """Main async loop: pop tasks and dispatch them as asyncio tasks."""
+        self._register_signal_handlers()
+        # Semaphore must be created inside the running loop.
+        self._semaphore = asyncio.Semaphore(self.max_concurrent)
+        self.running = True
+        logger.info(
+            f"AsyncWorker {self.worker_id} started "
+            f"(max_concurrent={self.max_concurrent})"
+        )
+
+        backoff = 1
+        while self.running:
+            try:
+                if len(self.active_tasks) >= self.max_concurrent:
+                    await asyncio.sleep(0.05)
+                    continue
+
+                loop = asyncio.get_running_loop()
+                task = await loop.run_in_executor(None, self.queue.pop)
+                backoff = 1  # reset after any successful Redis round-trip
+
+                if task is None:
+                    await asyncio.sleep(0.05)
+                    continue
+
+                self.active_tasks.add(task.task_id)
+                asyncio.create_task(self.process_task_async(task))
+                logger.info(
+                    f"Dispatched {task.task_id} "
+                    f"(active: {len(self.active_tasks)}/{self.max_concurrent})"
+                )
+
+            except redis.exceptions.RedisError as e:
+                logger.error(
+                    f"AsyncWorker {self.worker_id} Redis error: {e}, "
+                    f"retrying in {backoff}s",
+                    exc_info=True,
+                )
+                await asyncio.sleep(backoff)
+                backoff = min(backoff * 2, 60)
+            except Exception as e:
+                logger.error(
+                    f"AsyncWorker {self.worker_id} loop error: {e}", exc_info=True
+                )
+                await asyncio.sleep(1)
+
+        # Drain any still-running tasks before exiting.
+        while self.active_tasks:
+            await asyncio.sleep(0.05)
+
+        logger.info(f"AsyncWorker {self.worker_id} stopped")
+
+    def start(self):
+        """Synchronous entry point — runs the async loop until completion."""
+        asyncio.run(self.start_async())