broccoli-workers 0.1.0__py3-none-any.whl

broccoli/workers/hybrid_worker.py

@@ -0,0 +1,269 @@
+# broccoli/workers/hybrid_worker.py
+# See async_worker.py: this defers evaluation of the `X | None` /
+# lowercase-generic hints below so they don't raise on Python 3.9.
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import time
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime
+
+import redis
+
+from broccoli.core.task.task import Task
+from broccoli.workers.base_worker import BaseWorker
+
+logger = logging.getLogger(__name__)
+
+
+class HybridWorker(BaseWorker):
+    """
+    Worker that combines asyncio (for concurrency) with a thread pool
+    (for CPU-bound / blocking task handlers).
+
+    Result storage and user callbacks are handled entirely inside
+    ``_handle_hybrid_result``, which means ``post_process`` must NOT repeat
+    them.  This class overrides ``post_process`` to run only the registered
+    post-process handlers and nothing else, preventing double-fired callbacks
+    and double-stored results.
+    """
+
+    def __init__(
+        self,
+        redis_url: str = "redis://localhost:6379",
+        worker_id: str = None,
+        thread_workers: int = 4,
+        async_tasks: int = 10,
+        result_ttl: int = 86400,  # 24 hours
+    ):
+        super().__init__(redis_url, worker_id)
+        self.thread_pool = ThreadPoolExecutor(max_workers=thread_workers)
+        self.thread_workers = thread_workers
+        self.async_tasks = async_tasks
+        self.active_tasks: set[str] = set()
+        self.loop: asyncio.AbstractEventLoop | None = None
+        self.result_ttl = result_ttl
+        # Semaphore created lazily inside the running loop.
+        self._semaphore: asyncio.Semaphore | None = None
+
+    # ------------------------------------------------------------------
+    # post_process override — prevent double-firing
+    # ------------------------------------------------------------------
+
+    def post_process(self, task: Task, success: bool) -> None:
+        """
+        HybridWorker handles result storage, hash cleanup, and completion /
+        failure callbacks inside ``_handle_hybrid_result``.  This override
+        runs only the registered post-process handlers so that none of those
+        actions fire a second time.
+        """
+        self._run_post_process_handlers(task, success)
+
+    # ------------------------------------------------------------------
+    # Async / hybrid task processing
+    # ------------------------------------------------------------------
+
+    async def process_task_hybrid(self, task: Task) -> None:
+        """
+        Acquire a concurrency slot, run the task handler in the thread pool,
+        then handle the result (queue transitions, result storage, callbacks).
+        """
+        async with self._semaphore:
+            try:
+                # pre_process is checked before entering the thread pool so a
+                # skip (returns False) can be routed to queue.fail() directly,
+                # rather than being indistinguishable from a real processing
+                # failure and consuming retry budget in _handle_hybrid_result.
+                if not self.pre_process(task):
+                    logger.info(f"Task {task.task_id} skipped by pre_process")
+                    self.queue.fail(task)
+                    return
+
+                # Use get_running_loop() — get_event_loop() is deprecated in 3.10+.
+                loop = asyncio.get_running_loop()
+                try:
+                    success = await asyncio.wait_for(
+                        loop.run_in_executor(self.thread_pool, 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
+                self._handle_hybrid_result(task, success)
+            except Exception as e:
+                logger.error(f"Hybrid task {task.task_id} failed: {e}", exc_info=True)
+                task.error = str(e)
+                self._handle_hybrid_result(task, False)
+            finally:
+                self.active_tasks.discard(task.task_id)
+
+    def _handle_hybrid_result(self, task: Task, success: bool) -> None:
+        """
+        Manage queue transitions and persist results.
+
+        Queue state is handled by TaskQueue methods (complete / fail / requeue),
+        not by direct Redis calls.  Result saving and hash cleanup are kept
+        separate so each concern has one place.
+
+        Completion / failure handlers are fired here (not in post_process) to
+        keep the full lifecycle in one readable sequence.
+        """
+        if success:
+            task.status = "completed"
+            task.progress = 100.0
+            # Update Redis before releasing dependents (same ordering as BaseWorker).
+            self._update_task(task)
+            self.queue.complete(task)  # releases dependents, removes from processing
+            logger.info(f"Task {task.task_id} completed")
+        else:
+            task.retries += 1
+            if task.retries >= task.max_retries:
+                task.status = "failed"
+                if not task.error:
+                    task.error = "Max retries exceeded"
+                self._update_task(task)
+                self.queue.fail(task)  # removes from processing set
+            else:
+                task.status = "pending"
+                self._update_task(task)
+                self.queue.requeue(task.task_id)  # back to runnable queue
+                logger.info(
+                    f"Task {task.task_id} requeued "
+                    f"(attempt {task.retries}/{task.max_retries})"
+                )
+                # Requeued — don't store a partial result or fire handlers.
+                # post_process still runs for registered post-process handlers.
+                self.post_process(task, success)
+                return
+
+        # On permanent failure, record in the dead-letter set before result
+        # storage/cleanup so the task stays inspectable even if _save_result
+        # raises.
+        if task.status == "failed":
+            try:
+                self._redis.zadd(
+                    f"{self.task_prefix}:dead_letter", {task.task_id: time.time()}
+                )
+            except Exception as e:
+                logger.error(
+                    f"Failed to record {task.task_id} in dead-letter set: {e}",
+                    exc_info=True,
+                )
+
+        # Persist extended result metadata with TTL.
+        self._save_result(task)
+
+        # Remove task metadata hash (queue sets already updated above).
+        self._cleanup_task(task)
+
+        # Fire user-facing callbacks.
+        if success:
+            self._run_completion_handlers(task, task.result)
+        else:
+            error = Exception(task.error) if task.error else Exception("Task failed")
+            self._run_failure_handlers(task, error)
+
+        # post_process runs only registered post-process handlers (overridden above).
+        self.post_process(task, success)
+
+    # ------------------------------------------------------------------
+    # Result storage & cleanup
+    # ------------------------------------------------------------------
+
+    def _save_result(self, task: Task) -> None:
+        """Persist full result metadata under ``result:<task_id>`` with a TTL."""
+        result_data = {
+            "id": task.task_id,
+            "task_type": task.task_type,
+            "status": task.status,
+            "result": task.result,
+            "error": task.error,
+            "chain": bool(task.payload.get("__chain_id")),
+            "chain_id": task.payload.get("__chain_id"),
+            "worker_id": self.worker_id,
+            "completed_at": datetime.now().isoformat(),
+            "retries": task.retries,
+        }
+        self._redis.setex(
+            f"result:{task.task_id}", self.result_ttl, json.dumps(result_data)
+        )
+        logger.debug(f"Result saved for task {task.task_id} (TTL={self.result_ttl}s)")
+
+    def _cleanup_task(self, task: Task) -> None:
+        """
+        Remove the task metadata hash from Redis.
+
+        Queue sorted-set membership is already managed by
+        ``queue.complete()`` / ``queue.fail()`` — no ``zrem`` calls here.
+        """
+        self._redis.delete(f"{self.task_prefix}:{task.task_id}")
+
+        if task.payload.get("__chain_id"):
+            self._redis.delete(f"chain:{task.task_id}")
+
+        logger.info(f"Task {task.task_id} metadata cleaned up")
+
+    # ------------------------------------------------------------------
+    # Worker loop
+    # ------------------------------------------------------------------
+
+    def start(self):
+        """Create a fresh event loop and run until stopped."""
+        self._register_signal_handlers()
+        self.running = True
+        self.loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self.loop)
+        logger.info(
+            f"HybridWorker {self.worker_id} started "
+            f"(threads={self.thread_workers}, async_slots={self.async_tasks})"
+        )
+        try:
+            self.loop.run_until_complete(self._run_hybrid())
+        finally:
+            self.loop.close()
+
+    async def _run_hybrid(self):
+        self._semaphore = asyncio.Semaphore(self.async_tasks)
+
+        backoff = 1
+        while self.running:
+            if len(self.active_tasks) >= self.async_tasks:
+                await asyncio.sleep(0.05)
+                continue
+
+            try:
+                loop = asyncio.get_running_loop()
+                task = await loop.run_in_executor(None, self.queue.pop_with_timeout, 1)
+                backoff = 1  # reset after any successful Redis round-trip
+            except redis.exceptions.RedisError as e:
+                logger.error(
+                    f"HybridWorker {self.worker_id} Redis error: {e}, "
+                    f"retrying in {backoff}s",
+                    exc_info=True,
+                )
+                await asyncio.sleep(backoff)
+                backoff = min(backoff * 2, 60)
+                continue
+
+            if task is None:
+                await asyncio.sleep(0.05)
+                continue
+
+            self.active_tasks.add(task.task_id)
+            asyncio.create_task(self.process_task_hybrid(task))
+
+        # Drain in-flight tasks before shutting down.
+        while self.active_tasks:
+            await asyncio.sleep(0.05)
+
+        self.thread_pool.shutdown(wait=True)
+        logger.info(f"HybridWorker {self.worker_id} stopped")
+
+    def stop(self):
+        self.running = False
+        logger.info(f"HybridWorker {self.worker_id} stopping...")

broccoli/workers/threaded_worker.py

@@ -0,0 +1,154 @@
+# broccoli/workers/threaded_worker.py
+# See async_worker.py: defers evaluation of the `dict[str, Task]` hint below
+# so it doesn't raise on Python 3.9 (lowercase generics are 3.9+ for
+# subscripting builtins, but the annotation deferral keeps this consistent
+# across the whole worker family regardless of exact minimum version).
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import TimeoutError as FutureTimeoutError
+
+import redis
+
+from broccoli.core.task.task import Task
+from broccoli.workers.base_worker import BaseWorker
+
+logger = logging.getLogger(__name__)
+
+
+class ThreadedWorker(BaseWorker):
+    """
+    Worker that processes multiple tasks concurrently using a thread pool.
+
+    ``_handle_task_result`` is inherited from ``BaseWorker`` — retry logic,
+    dependency release, and queue-set transitions all live in one place.
+    """
+
+    def __init__(
+        self,
+        redis_url: str = "redis://localhost:6379",
+        worker_id: str = None,
+        max_workers: int = 4,
+    ):
+        super().__init__(redis_url, worker_id)
+        self.max_workers = max_workers
+        self.executor = ThreadPoolExecutor(max_workers=max_workers)
+        self.active_tasks: dict[str, Task] = {}
+        self.task_lock = threading.Lock()
+        # Separate executor to run process() under a timeout guard. Kept
+        # distinct from self.executor (which process_task itself runs on) to
+        # avoid the deadlock risk of a task waiting on a future submitted to
+        # the same bounded pool it's already occupying a slot in.
+        self._timeout_executor = ThreadPoolExecutor(
+            max_workers=max_workers, thread_name_prefix="timeout-guard"
+        )
+
+    def process_task(self, task: Task) -> None:
+        """
+        Process one task inside a thread-pool thread.
+
+        Uses the inherited ``_handle_task_result`` for all state transitions,
+        then calls ``post_process`` for result storage and user callbacks.
+        """
+        try:
+            logger.info(
+                f"Thread {threading.current_thread().name} "
+                f"processing task {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
+
+            # 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:
+                future = self._timeout_executor.submit(self.process, task)
+                try:
+                    success = future.result(timeout=self.task_timeout)
+                except FutureTimeoutError:
+                    future.cancel()  # best-effort; thread keeps running if already started
+                    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) + _update_task
+            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 (e.g. a Redis error) — the task's queue state may be
+            # inconsistent, so we log rather than re-running the state
+            # machine on it a second time.
+            logger.error(
+                f"Task {task.task_id} failed outside handler: {e}", exc_info=True
+            )
+        finally:
+            with self.task_lock:
+                self.active_tasks.pop(task.task_id, None)
+
+    def start(self):
+        """Main loop: pop tasks and submit them to the thread pool."""
+        self._register_signal_handlers()
+        self.running = True
+        logger.info(
+            f"ThreadedWorker {self.worker_id} started (max_workers={self.max_workers})"
+        )
+
+        backoff = 1
+        while self.running:
+            try:
+                with self.task_lock:
+                    active_count = len(self.active_tasks)
+
+                if active_count >= self.max_workers:
+                    time.sleep(0.05)
+                    continue
+
+                task = self.queue.pop()
+                backoff = 1  # reset after any successful Redis round-trip
+                if task is None:
+                    time.sleep(0.05)
+                    continue
+
+                with self.task_lock:
+                    self.active_tasks[task.task_id] = task
+
+                self.executor.submit(self.process_task, task)
+                logger.info(
+                    f"Submitted {task.task_id} to thread pool "
+                    f"(active: {active_count + 1}/{self.max_workers})"
+                )
+
+            except redis.exceptions.RedisError as e:
+                logger.error(
+                    f"ThreadedWorker {self.worker_id} Redis error: {e}, "
+                    f"retrying in {backoff}s",
+                    exc_info=True,
+                )
+                time.sleep(backoff)
+                backoff = min(backoff * 2, 60)
+            except Exception as e:
+                logger.error(
+                    f"ThreadedWorker {self.worker_id} loop error: {e}", exc_info=True
+                )
+                time.sleep(1)
+
+        self.executor.shutdown(wait=True)
+        self._timeout_executor.shutdown(wait=False)
+        logger.info(f"ThreadedWorker {self.worker_id} stopped")

broccoli/workers/worker_pool.py

@@ -0,0 +1,99 @@
+# broccoli/workers/worker_pool.py
+import logging
+import signal
+import threading
+from typing import List
+
+from broccoli.workers.base_worker import BaseWorker
+
+logger = logging.getLogger(__name__)
+
+
+class WorkerPool:
+    """
+    Manages a fixed-size pool of workers, each running in its own daemon thread.
+    """
+
+    def __init__(
+        self,
+        worker_type: type = BaseWorker,
+        num_workers: int = 4,
+        redis_url: str = "redis://localhost:6379",
+        **worker_kwargs,
+    ):
+        self.worker_type = worker_type
+        self.num_workers = num_workers
+        self.redis_url = redis_url
+        self.worker_kwargs = worker_kwargs  # e.g. max_workers=, max_concurrent=
+        self.workers: List[BaseWorker] = []  # initialised before _create_workers runs
+        self.threads: List[threading.Thread] = []
+        self.running = False
+        self.shutdown_flag = threading.Event()
+        self._create_workers()  # populate self.workers once
+
+    def _create_workers(self) -> None:
+        """Instantiate ``num_workers`` workers and append them to ``self.workers``."""
+        for i in range(self.num_workers):
+            worker = self.worker_type(
+                redis_url=self.redis_url,
+                worker_id=f"worker-{i + 1}",
+                **self.worker_kwargs,
+            )
+            self.workers.append(worker)
+
+    def start(self):
+        """
+        Spawn one daemon thread per worker, then block until a shutdown signal
+        arrives (``stop()`` or a signal handler calling ``_signal_handler``).
+        """
+        # Reset per-cycle state so repeated start()/stop() cycles (e.g. in
+        # tests) don't accumulate dead threads from a previous run or get
+        # stuck on an already-set flag from a prior shutdown.
+        self.threads = []
+        self.shutdown_flag.clear()
+        self.running = True
+
+        try:
+            signal.signal(signal.SIGINT, self._signal_handler)
+            signal.signal(signal.SIGTERM, self._signal_handler)
+        except (ValueError, OSError) as e:
+            logger.debug(f"Could not register signal handlers: {e}")
+
+        for i, worker in enumerate(self.workers):
+            thread = threading.Thread(
+                target=worker.start,
+                name=f"Worker-{i + 1}",
+                daemon=True,
+            )
+            self.threads.append(thread)
+            thread.start()
+            logger.info(f"Started worker {i + 1}/{self.num_workers}")
+
+        # Block the calling thread until stop() or a signal fires.
+        self.shutdown_flag.wait()
+        self.stop()
+
+    def _signal_handler(self, signum, frame):
+        """Handle SIGINT / SIGTERM by unblocking the main thread."""
+        logger.info("Received stop signal, shutting down worker pool...")
+        self.shutdown_flag.set()
+
+    def stop(self):
+        """Gracefully stop all workers and join their threads."""
+        logger.info("Stopping all workers...")
+
+        # Set first so a direct stop() call (e.g. from a test or a completion
+        # handler) also unblocks start(), which may be parked on this wait().
+        self.shutdown_flag.set()
+
+        for worker in self.workers:
+            try:
+                worker.stop()
+            except Exception as e:
+                logger.error(f"Error stopping worker: {e}")
+
+        for thread in self.threads:
+            thread.join(timeout=3)
+
+        self.running = False
+        logger.info("All workers stopped")