nat-engine 1__py3-none-any.whl

mannf/core/agents/worker_pool.py

@@ -0,0 +1,366 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+# Commercial licensing available — see COMMERCIAL_LICENSE.md.
+
+"""Worker pool orchestration for large-scale distributed NAT execution.
+
+The :class:`WorkerPool` manages a pool of concurrent asyncio workers that
+pull jobs from a shared :class:`~mannf.product.scheduling.queue.JobQueue` and
+execute them.  Workers are implemented as asyncio coroutines (lightweight),
+enabling hundreds of parallel scan/test jobs on a single host.
+
+For multi-host / cluster deployments, use :class:`RedisJobQueue` as the
+backend — all nodes then compete for jobs via the same distributed queue.
+
+Usage::
+
+    from mannf.core.agents.worker_pool import WorkerPool
+    from mannf.product.scheduling.queue import InMemoryJobQueue
+
+    queue = InMemoryJobQueue()
+    pool = WorkerPool(queue=queue, worker_count=4)
+
+    await pool.start()
+
+    job_id = await queue.enqueue("scan", {"scan_id": "abc", …})
+    # Workers will pick up the job automatically.
+
+    await pool.stop()
+    info = pool.info()          # dict with worker states, queue stats, etc.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Callable, Coroutine, Dict, List, Optional
+
+from mannf.product.scheduling.queue import Job, JobQueue
+
+logger = logging.getLogger(__name__)
+
+# Type alias for job handler coroutine functions.
+JobHandler = Callable[[Job], Coroutine[Any, Any, None]]
+
+# ---------------------------------------------------------------------------
+# Worker state
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class WorkerInfo:
+    """Snapshot of a single worker's state."""
+
+    worker_id: str
+    status: str = "idle"          # idle | running | error | stopped
+    current_job_id: Optional[str] = None
+    jobs_completed: int = 0
+    jobs_failed: int = 0
+    last_heartbeat: float = field(default_factory=time.time)
+    started_at: float = field(default_factory=time.time)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "worker_id": self.worker_id,
+            "status": self.status,
+            "current_job_id": self.current_job_id,
+            "jobs_completed": self.jobs_completed,
+            "jobs_failed": self.jobs_failed,
+            "last_heartbeat": self.last_heartbeat,
+            "started_at": self.started_at,
+        }
+
+
+# ---------------------------------------------------------------------------
+# WorkerPool
+# ---------------------------------------------------------------------------
+
+
+class WorkerPool:
+    """Asyncio worker pool that consumes jobs from a :class:`JobQueue`.
+
+    Parameters
+    ----------
+    queue:
+        The job queue to pull jobs from.
+    worker_count:
+        Number of concurrent worker coroutines to run.
+    poll_interval:
+        Seconds to sleep when the queue is empty before polling again.
+    max_attempts:
+        Maximum number of retry attempts for a failed job.
+    heartbeat_interval:
+        Seconds between worker heartbeat updates.
+    handlers:
+        Mapping of ``job_type`` → async callable.  The callable receives the
+        :class:`~mannf.product.scheduling.queue.Job` and should raise on
+        failure.  If a job type has no registered handler the job is nack'd
+        as failed immediately.
+    """
+
+    def __init__(
+        self,
+        queue: JobQueue,
+        worker_count: int = 4,
+        poll_interval: float = 1.0,
+        max_attempts: int = 3,
+        heartbeat_interval: float = 10.0,
+        handlers: Optional[Dict[str, JobHandler]] = None,
+    ) -> None:
+        self._queue = queue
+        self._worker_count = worker_count
+        self._poll_interval = poll_interval
+        self._max_attempts = max_attempts
+        self._heartbeat_interval = heartbeat_interval
+        self._handlers: Dict[str, JobHandler] = handlers or {}
+
+        self._workers: Dict[str, WorkerInfo] = {}
+        self._tasks: List[asyncio.Task[None]] = []
+        self._running = False
+        self._pool_id = str(uuid.uuid4())[:8]
+
+    # ------------------------------------------------------------------
+    # Handler registration
+    # ------------------------------------------------------------------
+
+    def register_handler(self, job_type: str, handler: JobHandler) -> None:
+        """Register an async *handler* for jobs of *job_type*."""
+        self._handlers[job_type] = handler
+        logger.debug("WorkerPool[%s]: registered handler for job_type=%r", self._pool_id, job_type)
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    async def start(self) -> None:
+        """Launch all worker coroutines.  Returns immediately."""
+        if self._running:
+            logger.warning("WorkerPool[%s] already running", self._pool_id)
+            return
+
+        self._running = True
+        for i in range(self._worker_count):
+            worker_id = f"worker-{self._pool_id}-{i}"
+            info = WorkerInfo(worker_id=worker_id)
+            self._workers[worker_id] = info
+            task = asyncio.create_task(
+                self._worker_loop(worker_id),
+                name=f"nat-worker-{worker_id}",
+            )
+            self._tasks.append(task)
+
+        logger.info(
+            "WorkerPool[%s] started with %d workers",
+            self._pool_id,
+            self._worker_count,
+        )
+
+    async def stop(self, timeout: float = 30.0) -> None:
+        """Signal workers to stop and wait for in-flight jobs to finish.
+
+        Parameters
+        ----------
+        timeout:
+            Maximum seconds to wait for a clean drain.  Workers still running
+            after the timeout are cancelled.
+        """
+        self._running = False
+        logger.info("WorkerPool[%s] stopping (timeout=%.1fs)…", self._pool_id, timeout)
+
+        try:
+            await asyncio.wait_for(
+                asyncio.gather(*self._tasks, return_exceptions=True),
+                timeout=timeout,
+            )
+        except asyncio.TimeoutError:
+            logger.warning(
+                "WorkerPool[%s] drain timeout exceeded; cancelling remaining workers",
+                self._pool_id,
+            )
+            for task in self._tasks:
+                if not task.done():
+                    task.cancel()
+            await asyncio.gather(*self._tasks, return_exceptions=True)
+
+        self._tasks.clear()
+        for info in self._workers.values():
+            info.status = "stopped"
+        logger.info("WorkerPool[%s] stopped", self._pool_id)
+
+    @property
+    def is_running(self) -> bool:
+        """``True`` while the pool's worker tasks are active."""
+        return self._running and any(not t.done() for t in self._tasks)
+
+    # ------------------------------------------------------------------
+    # Scale up / down
+    # ------------------------------------------------------------------
+
+    async def scale(self, target_count: int) -> None:
+        """Adjust the number of running workers to *target_count*.
+
+        Workers are added or gracefully removed to reach the target.
+        """
+        current = sum(1 for t in self._tasks if not t.done())
+        delta = target_count - current
+
+        if delta > 0:
+            for _ in range(delta):
+                i = len(self._workers)
+                worker_id = f"worker-{self._pool_id}-{i}"
+                info = WorkerInfo(worker_id=worker_id)
+                self._workers[worker_id] = info
+                task = asyncio.create_task(
+                    self._worker_loop(worker_id),
+                    name=f"nat-worker-{worker_id}",
+                )
+                self._tasks.append(task)
+            logger.info(
+                "WorkerPool[%s] scaled up to %d workers", self._pool_id, target_count
+            )
+
+        elif delta < 0:
+            # Cancel the newest idle workers first
+            removed = 0
+            for task in reversed(self._tasks):
+                if removed >= abs(delta):
+                    break
+                if not task.done():
+                    task.cancel()
+                    removed += 1
+            logger.info(
+                "WorkerPool[%s] scaled down to ~%d workers", self._pool_id, target_count
+            )
+
+        self._worker_count = target_count
+
+    # ------------------------------------------------------------------
+    # Info / status
+    # ------------------------------------------------------------------
+
+    def info(self) -> Dict[str, Any]:
+        """Return a serialisable snapshot of pool state."""
+        active = sum(1 for t in self._tasks if not t.done())
+        return {
+            "pool_id": self._pool_id,
+            "worker_count": self._worker_count,
+            "active_tasks": active,
+            "is_running": self.is_running,
+            "workers": [w.to_dict() for w in self._workers.values()],
+        }
+
+    async def full_status(self) -> Dict[str, Any]:
+        """Return pool info combined with live queue stats."""
+        queue_stats = await self._queue.status()
+        data = self.info()
+        data["queue"] = queue_stats.to_dict()
+        return data
+
+    # ------------------------------------------------------------------
+    # Internal worker loop
+    # ------------------------------------------------------------------
+
+    async def _worker_loop(self, worker_id: str) -> None:
+        """Main coroutine for a single worker.
+
+        Continuously dequeues jobs and dispatches them to the registered
+        handler.  Backs off with `_poll_interval` when the queue is empty.
+        """
+        info = self._workers[worker_id]
+        logger.debug("Worker %s started", worker_id)
+
+        while self._running:
+            info.last_heartbeat = time.time()
+
+            try:
+                job = await self._queue.dequeue(worker_id=worker_id)
+            except Exception:  # noqa: BLE001
+                logger.exception("Worker %s: error dequeuing job", worker_id)
+                await asyncio.sleep(self._poll_interval)
+                continue
+
+            if job is None:
+                # Queue is empty — back off
+                info.status = "idle"
+                info.current_job_id = None
+                await asyncio.sleep(self._poll_interval)
+                continue
+
+            info.status = "running"
+            info.current_job_id = job.job_id
+            logger.info(
+                "Worker %s: picked up job %s (type=%s, attempt=%d)",
+                worker_id,
+                job.job_id,
+                job.job_type,
+                job.attempt,
+            )
+
+            try:
+                await self._dispatch(job)
+                await self._queue.ack(job.job_id)
+                info.jobs_completed += 1
+                logger.info("Worker %s: job %s completed", worker_id, job.job_id)
+            except asyncio.CancelledError:
+                # Worker is being stopped — nack without retry
+                await self._queue.nack(
+                    job.job_id,
+                    error="Worker cancelled",
+                    retry=False,
+                )
+                raise
+            except Exception as exc:  # noqa: BLE001
+                info.jobs_failed += 1
+                error_msg = str(exc)
+                should_retry = job.attempt < self._max_attempts
+                logger.warning(
+                    "Worker %s: job %s failed (attempt %d/%d, retry=%s): %s",
+                    worker_id,
+                    job.job_id,
+                    job.attempt,
+                    self._max_attempts,
+                    should_retry,
+                    error_msg,
+                )
+                await self._queue.nack(
+                    job.job_id,
+                    error=error_msg,
+                    retry=should_retry,
+                    max_attempts=self._max_attempts,
+                )
+            finally:
+                info.current_job_id = None
+                info.status = "idle"
+
+        info.status = "stopped"
+        logger.debug("Worker %s exited", worker_id)
+
+    async def _dispatch(self, job: Job) -> None:
+        """Dispatch a job to its registered handler."""
+        handler = self._handlers.get(job.job_type)
+        if handler is None:
+            raise ValueError(f"No handler registered for job_type={job.job_type!r}")
+        await handler(job)
+
+
+# ---------------------------------------------------------------------------
+# Module-level singleton
+# ---------------------------------------------------------------------------
+
+_default_pool: Optional[WorkerPool] = None
+
+
+def get_default_pool() -> Optional[WorkerPool]:
+    """Return the process-wide default :class:`WorkerPool`, or ``None``."""
+    return _default_pool
+
+
+def set_default_pool(pool: Optional[WorkerPool]) -> None:
+    """Set the process-wide default :class:`WorkerPool`."""
+    global _default_pool
+    _default_pool = pool

mannf/core/anomaly/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+# Commercial licensing available — see COMMERCIAL_LICENSE.md.
+
+"""Anomaly detection sub-package for NAT."""
+
+from mannf.core.anomaly.enhanced_detector import AnomalyAlert, AnomalyType, EnhancedAnomalyDetector
+
+__all__ = [
+    "EnhancedAnomalyDetector",
+    "AnomalyAlert",
+    "AnomalyType",
+]