agent-learning-kit 0.1.0__py3-none-any.whl

fi/evals/framework/backends/base.py

@@ -0,0 +1,223 @@
+"""
+Base backend abstraction for evaluation execution.
+
+This module defines the abstract interface for evaluation backends,
+enabling pluggable execution strategies (thread pool, Temporal, etc.).
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Dict, Any, Optional, Callable, TypeVar, Generic, List
+from datetime import datetime, timezone
+
+
+T = TypeVar("T")
+
+
+class TaskStatus(Enum):
+    """Status of a backend task."""
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+    TIMEOUT = "timeout"
+
+
+@dataclass
+class BackendConfig:
+    """
+    Base configuration for backends.
+
+    Subclasses add backend-specific options.
+    """
+    max_workers: int = 4
+    timeout_seconds: float = 300.0  # 5 minutes default
+    retry_count: int = 0
+    retry_delay_seconds: float = 1.0
+
+
+@dataclass
+class TaskHandle(Generic[T]):
+    """
+    Handle to a submitted task.
+
+    Allows checking status and retrieving results.
+    """
+    task_id: str
+    backend_name: str
+    submitted_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    # Internal state - set by backend
+    _status: TaskStatus = TaskStatus.PENDING
+    _result: Optional[T] = None
+    _error: Optional[str] = None
+    _completed_at: Optional[datetime] = None
+
+    @property
+    def status(self) -> TaskStatus:
+        """Current task status."""
+        return self._status
+
+    @property
+    def is_done(self) -> bool:
+        """Whether task has completed (success or failure)."""
+        return self._status in (
+            TaskStatus.COMPLETED,
+            TaskStatus.FAILED,
+            TaskStatus.CANCELLED,
+            TaskStatus.TIMEOUT,
+        )
+
+    @property
+    def succeeded(self) -> bool:
+        """Whether task completed successfully."""
+        return self._status == TaskStatus.COMPLETED
+
+    @property
+    def result(self) -> Optional[T]:
+        """Task result if completed successfully."""
+        return self._result
+
+    @property
+    def error(self) -> Optional[str]:
+        """Error message if task failed."""
+        return self._error
+
+
+class Backend(ABC):
+    """
+    Abstract base class for evaluation backends.
+
+    Backends handle the actual execution of evaluation tasks,
+    whether locally (thread pool), distributed (Temporal, Celery),
+    or in a compute cluster (Ray).
+
+    Example implementation:
+        class MyBackend(Backend):
+            name = "my_backend"
+
+            def submit(self, fn, args, kwargs, context):
+                # Submit task to execution system
+                task_id = my_system.submit(fn, *args, **kwargs)
+                return TaskHandle(task_id=task_id, backend_name=self.name)
+
+            def get_result(self, handle, timeout):
+                # Wait for and return result
+                return my_system.wait(handle.task_id, timeout)
+
+    Thread Safety:
+        Backends must be thread-safe. Multiple threads may call
+        submit() and get_result() concurrently.
+    """
+
+    name: str = "base"
+
+    @abstractmethod
+    def submit(
+        self,
+        fn: Callable[..., T],
+        args: tuple = (),
+        kwargs: Optional[Dict[str, Any]] = None,
+        context: Optional[Dict[str, Any]] = None,
+    ) -> TaskHandle[T]:
+        """
+        Submit a task for execution.
+
+        Args:
+            fn: The function to execute
+            args: Positional arguments
+            kwargs: Keyword arguments
+            context: Trace context for propagation (trace_id, span_id, etc.)
+
+        Returns:
+            TaskHandle to track the task
+        """
+        pass
+
+    @abstractmethod
+    def get_result(
+        self,
+        handle: TaskHandle[T],
+        timeout: Optional[float] = None,
+    ) -> T:
+        """
+        Get the result of a submitted task.
+
+        Args:
+            handle: The task handle from submit()
+            timeout: Maximum seconds to wait (None = wait forever)
+
+        Returns:
+            The task result
+
+        Raises:
+            TimeoutError: If timeout exceeded
+            Exception: If task failed
+        """
+        pass
+
+    @abstractmethod
+    def get_status(self, handle: TaskHandle) -> TaskStatus:
+        """
+        Get current status of a task.
+
+        Args:
+            handle: The task handle
+
+        Returns:
+            Current TaskStatus
+        """
+        pass
+
+    @abstractmethod
+    def cancel(self, handle: TaskHandle) -> bool:
+        """
+        Attempt to cancel a task.
+
+        Args:
+            handle: The task handle
+
+        Returns:
+            True if cancelled, False if already running/complete
+        """
+        pass
+
+    def submit_batch(
+        self,
+        tasks: List[tuple],  # List of (fn, args, kwargs, context)
+    ) -> List[TaskHandle]:
+        """
+        Submit multiple tasks.
+
+        Default implementation submits sequentially.
+        Override for backends with batch submission support.
+
+        Args:
+            tasks: List of (function, args, kwargs, context) tuples
+
+        Returns:
+            List of TaskHandles
+        """
+        handles = []
+        for fn, args, kwargs, context in tasks:
+            handle = self.submit(fn, args, kwargs or {}, context)
+            handles.append(handle)
+        return handles
+
+    def shutdown(self, wait: bool = True) -> None:
+        """
+        Shutdown the backend.
+
+        Args:
+            wait: Whether to wait for pending tasks to complete
+        """
+        pass  # Default: no-op
+
+    def __enter__(self) -> "Backend":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.shutdown(wait=True)

fi/evals/framework/backends/celery_backend.py

@@ -0,0 +1,417 @@
+"""
+Celery backend for distributed task execution.
+
+Provides distributed task execution using Celery with Redis/RabbitMQ.
+Requires: pip install 'celery[redis]'
+"""
+
+import logging
+import threading
+import uuid
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, TypeVar
+
+from .base import Backend, BackendConfig, TaskHandle, TaskStatus
+from ._utils import CELERY
+
+T = TypeVar("T")
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CeleryConfig(BackendConfig):
+    """
+    Configuration for Celery backend.
+
+    Attributes:
+        broker_url: Message broker URL (Redis/RabbitMQ)
+        result_backend: Result backend URL
+        task_queue: Queue name for tasks
+        task_priority: Default task priority (0-9, higher = more priority)
+        task_serializer: Serialization format ('json', 'pickle')
+        result_serializer: Result serialization format
+        task_acks_late: Acknowledge tasks after execution
+        task_reject_on_worker_lost: Reject tasks if worker dies
+        task_track_started: Track when tasks start executing
+        worker_prefetch_multiplier: Number of tasks to prefetch
+    """
+
+    broker_url: str = "redis://localhost:6379/0"
+    result_backend: str = "redis://localhost:6379/1"
+    task_queue: str = "eval_tasks"
+    task_priority: int = 0
+    task_serializer: str = "cloudpickle"
+    result_serializer: str = "cloudpickle"
+    task_acks_late: bool = True
+    task_reject_on_worker_lost: bool = True
+    task_track_started: bool = True
+    worker_prefetch_multiplier: int = 1
+
+
+class CeleryBackend(Backend):
+    """
+    Celery backend for distributed task execution.
+
+    Uses Celery for distributed task execution across multiple workers.
+    Integrates with existing Celery infrastructure (Redis, RabbitMQ).
+
+    Example:
+        config = CeleryConfig(
+            broker_url="redis://localhost:6379/0",
+            result_backend="redis://localhost:6379/1",
+            task_queue="eval_tasks",
+        )
+        backend = CeleryBackend(config)
+
+        handle = backend.submit(my_eval_func, args=(input_data,))
+        result = backend.get_result(handle)
+
+    Note:
+        Requires Celery workers to be running to execute tasks.
+        Start workers with: celery -A your_app worker -Q eval_tasks
+    """
+
+    name: str = "celery"
+
+    def __init__(self, config: Optional[CeleryConfig] = None):
+        """
+        Initialize Celery backend.
+
+        Args:
+            config: Celery configuration
+
+        Raises:
+            ImportError: If celery is not installed
+        """
+        CELERY.require()
+
+        self.config = config or CeleryConfig()
+        self._app: Optional[Any] = None
+        self._task: Optional[Any] = None
+        self._handles: Dict[str, TaskHandle] = {}
+        self._async_results: Dict[str, Any] = {}
+        self._lock = threading.Lock()
+
+        self._setup_celery()
+
+    def _setup_celery(self) -> None:
+        """Set up Celery application and task from shared worker module."""
+        from .celery_worker import app, eval_task
+
+        # Override broker/backend if config differs from env defaults
+        app.conf.update(
+            broker_url=self.config.broker_url,
+            result_backend=self.config.result_backend,
+            task_serializer=self.config.task_serializer,
+            result_serializer=self.config.result_serializer,
+            task_acks_late=self.config.task_acks_late,
+            task_reject_on_worker_lost=self.config.task_reject_on_worker_lost,
+            task_track_started=self.config.task_track_started,
+            worker_prefetch_multiplier=self.config.worker_prefetch_multiplier,
+            task_default_queue=self.config.task_queue,
+        )
+
+        self._app = app
+        self._task = eval_task
+
+    def submit(
+        self,
+        fn: Callable[..., T],
+        args: tuple = (),
+        kwargs: Optional[Dict[str, Any]] = None,
+        context: Optional[Dict[str, Any]] = None,
+    ) -> TaskHandle[T]:
+        """
+        Submit a task to Celery.
+
+        Args:
+            fn: The function to execute
+            args: Positional arguments
+            kwargs: Keyword arguments
+            context: Trace context (stored in task metadata)
+
+        Returns:
+            TaskHandle to track the task
+        """
+        kwargs = kwargs or {}
+        task_id = str(uuid.uuid4())
+
+        # Create handle
+        handle = TaskHandle(
+            task_id=task_id,
+            backend_name=self.name,
+            metadata={
+                "queue": self.config.task_queue,
+                "priority": self.config.task_priority,
+                "function": fn.__name__ if hasattr(fn, "__name__") else str(fn),
+                "context": context,
+            },
+        )
+        handle._status = TaskStatus.PENDING
+
+        with self._lock:
+            self._handles[task_id] = handle
+
+        try:
+            # Submit task to Celery
+            async_result = self._task.apply_async(
+                args=(fn, args, kwargs),
+                task_id=task_id,
+                queue=self.config.task_queue,
+                priority=self.config.task_priority,
+            )
+
+            with self._lock:
+                self._async_results[task_id] = async_result
+                handle._status = TaskStatus.PENDING
+
+            logger.debug(f"Submitted task {task_id} to queue {self.config.task_queue}")
+
+        except Exception as e:
+            handle._status = TaskStatus.FAILED
+            handle._error = str(e)
+            logger.error(f"Failed to submit task {task_id}: {e}")
+
+        return handle
+
+    def get_result(
+        self,
+        handle: TaskHandle[T],
+        timeout: Optional[float] = None,
+    ) -> T:
+        """
+        Get the result of a Celery task.
+
+        Args:
+            handle: The task handle from submit()
+            timeout: Maximum seconds to wait
+
+        Returns:
+            The task result
+
+        Raises:
+            TimeoutError: If timeout exceeded
+            Exception: If task failed
+        """
+        timeout = timeout or self.config.timeout_seconds
+
+        with self._lock:
+            async_result = self._async_results.get(handle.task_id)
+
+        if async_result is None:
+            # Try to get result from backend
+            from celery.result import AsyncResult
+            async_result = AsyncResult(handle.task_id, app=self._app)
+
+        try:
+            result = async_result.get(timeout=timeout)
+
+            with self._lock:
+                if handle.task_id in self._handles:
+                    self._handles[handle.task_id]._status = TaskStatus.COMPLETED
+                    self._handles[handle.task_id]._result = result
+
+            return result
+
+        except Exception as e:
+            error_name = type(e).__name__
+            if "TimeoutError" in error_name or "TimeLimitExceeded" in error_name:
+                with self._lock:
+                    if handle.task_id in self._handles:
+                        self._handles[handle.task_id]._status = TaskStatus.TIMEOUT
+                raise TimeoutError(f"Task {handle.task_id} timed out after {timeout}s")
+
+            with self._lock:
+                if handle.task_id in self._handles:
+                    self._handles[handle.task_id]._status = TaskStatus.FAILED
+                    self._handles[handle.task_id]._error = str(e)
+            raise
+
+    def get_status(self, handle: TaskHandle) -> TaskStatus:
+        """
+        Get current status of a Celery task.
+
+        Args:
+            handle: The task handle
+
+        Returns:
+            Current TaskStatus
+        """
+        with self._lock:
+            async_result = self._async_results.get(handle.task_id)
+
+        if async_result is None:
+            from celery.result import AsyncResult
+            async_result = AsyncResult(handle.task_id, app=self._app)
+
+        return self._map_celery_status(async_result.status)
+
+    def _map_celery_status(self, celery_status: str) -> TaskStatus:
+        """Map Celery task status to TaskStatus."""
+        status_map = {
+            "PENDING": TaskStatus.PENDING,
+            "STARTED": TaskStatus.RUNNING,
+            "SUCCESS": TaskStatus.COMPLETED,
+            "FAILURE": TaskStatus.FAILED,
+            "REVOKED": TaskStatus.CANCELLED,
+            "RETRY": TaskStatus.RUNNING,
+        }
+        return status_map.get(celery_status, TaskStatus.PENDING)
+
+    def cancel(self, handle: TaskHandle) -> bool:
+        """
+        Revoke/cancel a Celery task.
+
+        Args:
+            handle: The task handle
+
+        Returns:
+            True if revoked, False otherwise
+        """
+        try:
+            with self._lock:
+                async_result = self._async_results.get(handle.task_id)
+
+            if async_result is None:
+                from celery.result import AsyncResult
+                async_result = AsyncResult(handle.task_id, app=self._app)
+
+            async_result.revoke(terminate=True)
+
+            with self._lock:
+                if handle.task_id in self._handles:
+                    self._handles[handle.task_id]._status = TaskStatus.CANCELLED
+
+            logger.debug(f"Revoked task {handle.task_id}")
+            return True
+
+        except Exception as e:
+            logger.warning(f"Failed to revoke task {handle.task_id}: {e}")
+            return False
+
+    def submit_batch(
+        self,
+        tasks: List[tuple],
+    ) -> List[TaskHandle]:
+        """
+        Submit multiple tasks efficiently using Celery group.
+
+        Args:
+            tasks: List of (function, args, kwargs, context) tuples
+
+        Returns:
+            List of TaskHandles
+        """
+        from celery import group
+
+        handles = []
+        signatures = []
+
+        for fn, args, kwargs, context in tasks:
+            kwargs = kwargs or {}
+            task_id = str(uuid.uuid4())
+
+            handle = TaskHandle(
+                task_id=task_id,
+                backend_name=self.name,
+                metadata={
+                    "queue": self.config.task_queue,
+                    "function": fn.__name__ if hasattr(fn, "__name__") else str(fn),
+                    "context": context,
+                },
+            )
+            handle._status = TaskStatus.PENDING
+            handles.append(handle)
+
+            with self._lock:
+                self._handles[task_id] = handle
+
+            sig = self._task.signature(
+                args=(fn, args, kwargs),
+                task_id=task_id,
+                queue=self.config.task_queue,
+            )
+            signatures.append(sig)
+
+        # Submit as a group
+        try:
+            job = group(signatures)
+            group_result = job.apply_async()
+
+            # Store individual results
+            for handle, async_result in zip(handles, group_result.results):
+                with self._lock:
+                    self._async_results[handle.task_id] = async_result
+
+        except Exception as e:
+            logger.error(f"Failed to submit batch: {e}")
+            for handle in handles:
+                handle._status = TaskStatus.FAILED
+                handle._error = str(e)
+
+        return handles
+
+    def shutdown(self, wait: bool = True) -> None:
+        """
+        Shutdown the Celery backend.
+
+        Revokes only our tracked tasks, not the entire Celery queue.
+
+        Args:
+            wait: Whether to wait for pending tasks
+        """
+        if self._app is not None:
+            with self._lock:
+                task_ids = list(self._async_results.keys())
+
+            for task_id in task_ids:
+                try:
+                    self._app.control.revoke(task_id, terminate=not wait)
+                except Exception:
+                    pass
+
+            with self._lock:
+                self._handles.clear()
+                self._async_results.clear()
+
+            logger.info(f"Celery backend shut down, revoked {len(task_ids)} tasks")
+
+    def get_queue_length(self) -> int:
+        """
+        Get the number of tasks waiting in the queue.
+
+        Returns:
+            Number of pending tasks in the queue
+        """
+        try:
+            with self._app.pool.acquire(block=True) as conn:
+                return conn.default_channel.client.llen(self.config.task_queue)
+        except Exception:
+            return -1
+
+    def get_stats(self) -> dict:
+        """
+        Get Celery backend statistics.
+
+        Returns:
+            Dictionary with queue and worker stats
+        """
+        stats = {
+            "queue": self.config.task_queue,
+            "broker": self.config.broker_url,
+            "pending_tasks": len([
+                h for h in self._handles.values()
+                if h._status == TaskStatus.PENDING
+            ]),
+            "running_tasks": len([
+                h for h in self._handles.values()
+                if h._status == TaskStatus.RUNNING
+            ]),
+        }
+
+        try:
+            stats["queue_length"] = self.get_queue_length()
+        except Exception:
+            stats["queue_length"] = -1
+
+        return stats

fi/evals/framework/backends/celery_worker.py

@@ -0,0 +1,78 @@
+"""
+Shared Celery app and task definition.
+
+This module is importable by both the SDK client (CeleryBackend) and a
+standalone Celery worker process, solving the problem of task registration
+needing to be visible to workers.
+
+Run as a worker::
+
+    celery -A fi.evals.framework.backends.celery_worker worker \\
+        -Q eval_tasks --loglevel=info
+
+Or configure via environment variables::
+
+    CELERY_BROKER_URL=redis://redis:6379/0 \\
+    CELERY_RESULT_BACKEND=redis://redis:6379/1 \\
+    celery -A fi.evals.framework.backends.celery_worker worker \\
+        -Q eval_tasks --loglevel=info
+
+Requires: pip install 'celery[redis]'
+
+Security note:
+    Pickle serialization is used here because Celery tasks transport arbitrary
+    Python callables (evaluation functions). This is the standard approach for
+    Celery task serialization in trusted environments. Never expose the broker
+    to untrusted networks.
+"""
+
+import os
+
+import cloudpickle
+from celery import Celery
+from kombu.serialization import register
+
+# ---------------------------------------------------------------------------
+# Register cloudpickle as a Celery serializer.
+# Unlike stdlib pickle, cloudpickle can serialize lambdas, closures, and
+# functions defined in __main__ — exactly what we need for ad-hoc eval fns.
+# ---------------------------------------------------------------------------
+
+register(
+    "cloudpickle",
+    cloudpickle.dumps,
+    cloudpickle.loads,
+    content_type="application/x-cloudpickle",
+    content_encoding="binary",
+)
+
+# ---------------------------------------------------------------------------
+# App configuration — shared between SDK client and Docker worker
+# ---------------------------------------------------------------------------
+
+BROKER_URL = os.environ.get("CELERY_BROKER_URL", "redis://localhost:6379/0")
+RESULT_BACKEND = os.environ.get("CELERY_RESULT_BACKEND", "redis://localhost:6379/1")
+TASK_QUEUE = os.environ.get("CELERY_TASK_QUEUE", "eval_tasks")
+
+app = Celery(
+    "eval_tasks",
+    broker=BROKER_URL,
+    backend=RESULT_BACKEND,
+)
+
+app.conf.update(
+    task_serializer="cloudpickle",
+    result_serializer="cloudpickle",
+    accept_content=["cloudpickle", "json"],
+    task_acks_late=True,
+    task_reject_on_worker_lost=True,
+    task_track_started=True,
+    worker_prefetch_multiplier=1,
+    task_default_queue=TASK_QUEUE,
+)
+
+
+@app.task(bind=True, name="eval_task")
+def eval_task(self, fn, args, kwargs):
+    """Generic task that executes the provided function."""
+    return fn(*args, **kwargs)