task scheduler

Author: luke-beamcloud

examples/tasks.py

@@ -0,0 +1,48 @@
+from fastloop import ExecutorType, FastLoop, RetryPolicy
+
+app = FastLoop(name="task-demo")
+
+
+@app.task(name="add", retry=RetryPolicy(max_attempts=3))
+async def add(a: int, b: int) -> int:
+    return a + b
+
+
+@app.task(name="cpu_work", executor=ExecutorType.PROCESS)
+def cpu_work(n: int) -> int:
+    total = 0
+    for i in range(n):
+        total += i * i
+    return total
+
+
+@app.task(name="flaky", retry=RetryPolicy(max_attempts=5, initial_delay=0.5))
+async def flaky_task() -> str:
+    import random
+
+    if random.random() < 0.5:
+        raise Exception("Random failure")
+    return "success"
+
+
+@app.schedule(name="cleanup", cron="0 * * * *")
+async def hourly_cleanup() -> dict:
+    return {"cleaned": True}
+
+
+@app.schedule(name="heartbeat", interval=60)
+async def heartbeat() -> dict:
+    return {"alive": True}
+
+
+async def demo():
+    task_id = await app.invoke("add", a=2, b=3)
+    result = await app.invoke("add", a=10, b=20, wait=True)
+    schedule_id = await app.schedule_task(
+        "add", cron="*/5 * * * *", args={"a": 1, "b": 2}
+    )
+    print(f"task_id={task_id}, result={result}, schedule_id={schedule_id}")
+
+
+if __name__ == "__main__":
+    app.run(port=8112)

fastloop/__init__.py

@@ -1,23 +1,25 @@
-"""
-FastLoop - A framework for building event-driven loop applications.
-"""
-
 from . import integrations
 from .context import LoopContext
 from .fastloop import FastLoop
 from .loop import Loop
 from .models import LoopEvent, WorkflowBlock
-from .types import BlockPlan, RetryPolicy, ScheduleType
+from .scheduler import Schedule
+from .task import TaskResult
+from .types import BlockPlan, ExecutorType, RetryPolicy, ScheduleType, TaskStatus
 from .workflow import Workflow
 
 __all__ = [
     "BlockPlan",
+    "ExecutorType",
     "FastLoop",
     "Loop",
     "LoopContext",
     "LoopEvent",
     "RetryPolicy",
+    "Schedule",
     "ScheduleType",
+    "TaskResult",
+    "TaskStatus",
     "Workflow",
     "WorkflowBlock",
     "integrations",

fastloop/constants.py

@@ -7,3 +7,4 @@
 DEFAULT_SYSTEM_CONFIG_DIR = "/etc/fastloop.d/"
 LEASE_TTL_S = 30
 LEASE_HEARTBEAT_INTERVAL_S = 10
+MAX_EVENT_HISTORY = 1000

fastloop/exceptions.py

@@ -66,6 +66,10 @@ class WorkflowNotFoundError(Exception):
     pass
 
 
+class TaskNotFoundError(Exception):
+    pass
+
+
 class WorkflowMaxRetriesError(Exception):
     def __init__(
         self,

fastloop/executor.py

@@ -0,0 +1,48 @@
+import asyncio
+from collections.abc import Callable
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
+from functools import partial
+from typing import Any
+
+from .types import ExecutorType
+
+_thread_pool: ThreadPoolExecutor | None = None
+_process_pool: ProcessPoolExecutor | None = None
+
+
+def _get_pool(executor_type: ExecutorType) -> ThreadPoolExecutor | ProcessPoolExecutor:
+    global _thread_pool, _process_pool
+    if executor_type == ExecutorType.THREAD:
+        if _thread_pool is None:
+            _thread_pool = ThreadPoolExecutor(max_workers=4)
+        return _thread_pool
+    if _process_pool is None:
+        _process_pool = ProcessPoolExecutor(max_workers=4)
+    return _process_pool
+
+
+async def run_in_executor(
+    executor_type: ExecutorType,
+    func: Callable[..., Any],
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    if executor_type == ExecutorType.ASYNC:
+        if asyncio.iscoroutinefunction(func):
+            return await func(*args, **kwargs)
+        return func(*args, **kwargs)
+
+    loop = asyncio.get_event_loop()
+    return await loop.run_in_executor(
+        _get_pool(executor_type), partial(func, *args, **kwargs)
+    )
+
+
+def shutdown_executors() -> None:
+    global _thread_pool, _process_pool
+    if _thread_pool:
+        _thread_pool.shutdown(wait=False)
+        _thread_pool = None
+    if _process_pool:
+        _process_pool.shutdown(wait=False)
+        _process_pool = None

fastloop/fastloop.py

@@ -29,15 +29,18 @@
 from .exceptions import (
     LoopAlreadyDefinedError,
     LoopNotFoundError,
+    TaskNotFoundError,
     WorkflowNotFoundError,
 )
 from .integrations import Integration
 from .logging import configure_logging, setup_logger
 from .loop import Loop, LoopManager
 from .models import LoopEvent
 from .monitor import LoopMonitor
+from .scheduler import Schedule, validate_cron
 from .state.state import StateManager, create_state_manager
-from .types import BaseConfig, LoopStatus, RetryPolicy
+from .task import TaskManager, TaskResult
+from .types import BaseConfig, ExecutorType, LoopStatus, RetryPolicy
 from .utils import get_func_import_path, import_func_from_path, infer_application_path
 from .workflow import Workflow, WorkflowManager
 
@@ -83,6 +86,7 @@ async def lifespan(_: FastAPI):
             self._monitor_task.cancel()
             await self.loop_manager.stop_all()
             await self.workflow_manager.stop_all()
+            await self.task_manager.stop_all()
 
         super().__init__(*args, **kwargs, lifespan=lifespan)
 
@@ -102,10 +106,12 @@ async def lifespan(_: FastAPI):
         )
         self.loop_manager: LoopManager = LoopManager(self.config, self.state_manager)
         self.workflow_manager: WorkflowManager = WorkflowManager(self.state_manager)
+        self.task_manager: TaskManager = TaskManager(self.state_manager)
         self._monitor_task: asyncio.Task[None] | None = None
         self._loop_start_func: Callable[[LoopContext], None] | None = None
         self._loop_metadata: dict[str, dict[str, Any]] = {}
         self._workflow_metadata: dict[str, dict[str, Any]] = {}
+        self._task_metadata: dict[str, dict[str, Any]] = {}
 
         configure_logging(
             pretty_print=self.config_manager.get("prettyPrintLogs", False)
@@ -802,3 +808,139 @@ async def restart_workflow(self, workflow_run_id: str) -> bool:
                 extra={"workflow_run_id": workflow_run_id, "error": str(e)},
             )
             return False
+
+    def task(
+        self,
+        name: str,
+        retry: RetryPolicy | None = None,
+        executor: ExecutorType = ExecutorType.ASYNC,
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """Register a task. Creates POST /{name} and GET /{name}/{task_id} endpoints."""
+
+        def _decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+            if name in self._task_metadata:
+                raise LoopAlreadyDefinedError(f"Task {name} already registered")
+
+            self._task_metadata[name] = {
+                "func": func,
+                "retry": retry,
+                "executor": executor,
+            }
+
+            async def _invoke_handler(request: dict[str, Any]):
+                result = await self.task_manager.submit(
+                    func=func,
+                    args=request,
+                    task_name=name,
+                    retry_policy=retry,
+                    executor_type=executor,
+                )
+                return JSONResponse(
+                    content={"task_id": result.task_id, "status": "pending"},
+                    status_code=HTTPStatus.ACCEPTED,
+                )
+
+            async def _status_handler(task_id: str):
+                try:
+                    task = await self.state_manager.get_task(task_id)
+                    return JSONResponse(content=task.to_dict())
+                except TaskNotFoundError as e:
+                    raise HTTPException(
+                        status_code=HTTPStatus.NOT_FOUND, detail=str(e)
+                    ) from e
+
+            self.add_api_route(f"/{name}", _invoke_handler, methods=["POST"])
+            self.add_api_route(f"/{name}/{{task_id}}", _status_handler, methods=["GET"])
+
+            return func
+
+        return _decorator
+
+    async def invoke(
+        self,
+        task_name: str,
+        wait: bool = False,
+        timeout: float | None = None,
+        **kwargs: Any,
+    ) -> str | TaskResult | Any:
+        """Invoke a task. Returns task_id if wait=False, else waits and returns result."""
+        if task_name not in self._task_metadata:
+            raise ValueError(f"Unknown task: {task_name}")
+
+        metadata = self._task_metadata[task_name]
+        handle = await self.task_manager.submit(
+            func=metadata["func"],
+            args=kwargs,
+            task_name=task_name,
+            retry_policy=metadata.get("retry"),
+            executor_type=metadata.get("executor", ExecutorType.ASYNC),
+        )
+
+        if wait:
+            return await handle.result(timeout=timeout)
+
+        return handle.task_id
+
+    def schedule(
+        self,
+        name: str,
+        cron: str | None = None,
+        interval: float | None = None,
+        retry: RetryPolicy | None = None,
+        executor: ExecutorType = ExecutorType.ASYNC,
+        args: dict[str, Any] | None = None,
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """Register a scheduled task. Use cron="*/5 * * * *" or interval=60."""
+        if not cron and not interval:
+            raise ValueError("Must specify either cron or interval")
+
+        if cron and not validate_cron(cron):
+            raise ValueError(f"Invalid cron expression: {cron}")
+
+        def _decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+            self.task(name=name, retry=retry, executor=executor)(func)
+
+            schedule = Schedule(
+                task_name=name,
+                cron=cron,
+                interval_seconds=interval,
+                args=args or {},
+            )
+
+            self._task_metadata[name]["schedule"] = schedule
+
+            return func
+
+        return _decorator
+
+    async def schedule_task(
+        self,
+        task_name: str,
+        cron: str | None = None,
+        interval: float | None = None,
+        args: dict[str, Any] | None = None,
+        schedule_id: str | None = None,
+    ) -> str:
+        """Programmatically schedule a registered task. Returns the schedule_id."""
+        if task_name not in self._task_metadata:
+            raise ValueError(f"Unknown task: {task_name}")
+
+        if not cron and not interval:
+            raise ValueError("Must specify either cron or interval")
+
+        if cron and not validate_cron(cron):
+            raise ValueError(f"Invalid cron expression: {cron}")
+
+        sid = schedule_id or task_name
+        schedule = Schedule(
+            task_name=task_name,
+            cron=cron,
+            interval_seconds=interval,
+            args=args or {},
+        )
+        await self.state_manager.save_schedule(sid, schedule)
+        return sid
+
+    async def unschedule_task(self, schedule_id: str) -> None:
+        """Remove a scheduled task by its schedule_id."""
+        await self.state_manager.delete_schedule(schedule_id)