Better cpubound

Author: lsbardel

.dev/install

@@ -1,3 +1,3 @@
 #!/usr/bin/env bash
 pip install -U pip poetry
-poetry install --all-extras
+poetry install --all-extras --with docs

docs/index.md

@@ -0,0 +1,17 @@
+# Welcome to MkDocs
+
+For full documentation visit [mkdocs.org](https://www.mkdocs.org).
+
+## Commands
+
+* `mkdocs new [dir-name]` - Create a new project.
+* `mkdocs serve` - Start the live-reloading docs server.
+* `mkdocs build` - Build the documentation site.
+* `mkdocs -h` - Print help message and exit.
+
+## Project layout
+
+    mkdocs.yml    # The configuration file.
+    docs/
+        index.md  # The documentation homepage.
+        ...       # Other markdown pages, images and other files.

fluid/scheduler/readme.md docs/scheduler.md

@@ -1,18 +1,18 @@
-# Distrubuted Task Producer/Consumer
+# Distributed Task Producer/Consumer
 
 This module has a lightweight implementation of a distributed task producer (TaskScheduler) and consumer (TaskConsumer).
 The middleware for distributing tasks can be configured via the Broker interface.
 A redis broker is provided for convenience.
 
 ## Tasks
 
-Tasks are standard python async functions decorated with the `task` or `cpu_task` decorators.
+Tasks are standard python async functions decorated with the `task` decorator.
 
 ```python
-from fluid.scheduler import task, TaskContext
+from fluid.scheduler import task, TaskRun
 
 @task
-async def say_hi(ctx: TaskContext):
+async def say_hi(ctx: TaskRun):
     return "Hi!"
 ```
 
@@ -21,10 +21,10 @@ There are two types of tasks implemented
 * **Simple concurrent tasks** - they run concurrently with the task consumer - thy must be IO type tasks (no heavy CPU bound operations)
 
   ```python
-    from fluid.scheduler import task, TaskContext
+    from fluid.scheduler import task, TaskRun
 
     @task
-    async def fecth_data(ctx: TaskContext):
+    async def fecth_data(ctx: TaskRun):
         # fetch data
         data = await http_cli.get("https://...")
         data_id = await datastore_cli.stote(data)
@@ -35,10 +35,10 @@ There are two types of tasks implemented
 * **CPU bound tasks** - they run on a subprocess
 
   ```python
-    from fluid.scheduler import cpu_task, TaskContext
+    from fluid.scheduler import task, TaskRun
 
-    @cpu_task
-    async def heavy_calculation(ctx: TaskContext):
+    @task(cpu_bound=True)
+    async def heavy_calculation(ctx: TaskRun):
         # perform some heavy calculation
         data = await datastore_cli.get(ctx.params["data_id"])
         ...

fluid/scheduler/__init__.py

@@ -1,21 +1,9 @@
 from .broker import Broker, QueuedTask
-from .constants import TaskPriority, TaskState
 from .consumer import TaskConsumer, TaskManager
 from .crontab import Scheduler, crontab
 from .every import every
+from .models import Task, TaskInfo, TaskPriority, TaskRun, TaskState, task
 from .scheduler import TaskScheduler
-from .task import (
-    Task,
-    TaskConstructor,
-    TaskContext,
-    TaskDecoratorError,
-    TaskExecutor,
-    TaskRunError,
-    create_task_app,
-    task,
-)
-from .task_info import TaskInfo
-from .task_run import TaskRun
 
 __all__ = [
     "Scheduler",
@@ -24,18 +12,12 @@
     "TaskScheduler",
     "task",
     "Task",
-    "TaskContext",
-    "TaskRunError",
-    "TaskExecutor",
-    "TaskConstructor",
-    "TaskDecoratorError",
     "TaskManager",
     "TaskConsumer",
     "TaskInfo",
     "TaskRun",
     "QueuedTask",
     "Broker",
-    "create_task_app",
     "crontab",
     "every",
 ]

fluid/scheduler/broker.py

@@ -2,53 +2,36 @@
 
 from abc import ABC, abstractmethod, abstractproperty
 from functools import cached_property
-from typing import Any, Dict, Iterable, List, NamedTuple, Optional
+from typing import TYPE_CHECKING, Any, Iterable
 from uuid import uuid4
 
+from redis.asyncio.lock import Lock
 from yarl import URL
 
-from fluid import json
-from fluid.tools.redis import FluidRedis
-from fluid.tools.timestamp import Timestamp
+from fluid import settings
+from fluid.utils.redis import Redis, FluidRedis
+import json
+from .errors import UnknownTaskError
 
-from . import settings
-from .constants import TaskPriority, TaskState
-from .task import Task
-from .task_info import TaskInfo
-from .task_run import TaskRun
-
-_brokers: dict[str, type[Broker]] = {}
-
-
-def broker_url_from_env() -> URL:
-    return URL(settings.SCHEDULER_BROKER_URL)
+from .models import QueuedTask, Task, TaskInfo, TaskPriority, TaskRun
 
+if TYPE_CHECKING:  # pragma: no cover
+    from .consumer import TaskManager
 
-class TaskError(RuntimeError):
-    pass
 
-
-class UnknownTask(TaskError):
-    pass
+_brokers: dict[str, type[Broker]] = {}
 
 
-class DisabledTask(TaskError):
-    pass
+def broker_url_from_env() -> URL:
+    return URL(settings.BROKER_URL)
 
 
-class TaskRegistry(Dict[str, Task]):
+class TaskRegistry(dict[str, Task]):
     def periodic(self) -> Iterable[Task]:
         for task in self.values():
             yield task
 
 
-class QueuedTask(NamedTuple):
-    run_id: str
-    task: str
-    params: Dict[str, Any]
-    priority: Optional[TaskPriority] = None
-
-
 class Broker(ABC):
     def __init__(self, url: URL) -> None:
         self.url: URL = url
@@ -59,15 +42,17 @@ def task_queue_names(self) -> tuple[str, ...]:
         """Names of the task queues"""
 
     @abstractmethod
-    async def queue_task(self, queued_task: QueuedTask) -> TaskRun:
+    async def queue_task(
+        self, task_manager: TaskManager, queued_task: QueuedTask
+    ) -> TaskRun:
         """Queue a task"""
 
     @abstractmethod
-    async def get_task_run(self) -> Optional[TaskRun]:
+    async def get_task_run(self, task_manager: TaskManager) -> TaskRun | None:
         """Get a Task run from the task queue"""
 
     @abstractmethod
-    async def queue_length(self) -> Dict[str, int]:
+    async def queue_length(self) -> dict[str, int]:
         """Length of task queues"""
 
     @abstractmethod
@@ -78,15 +63,22 @@ async def get_tasks_info(self, *task_names: str) -> list[TaskInfo]:
     async def update_task(self, task: Task, params: dict[str, Any]) -> TaskInfo:
         """Update a task dynamic parameters"""
 
+    @abstractmethod
     async def close(self) -> None:
         """Close the broker on shutdown"""
 
+    @abstractmethod
+    def lock(self, name: str, timeout: float | None = None) -> Lock:
+        """Create a lock"""
+
     def new_uuid(self) -> str:
         return uuid4().hex
 
     async def filter_tasks(
-        self, scheduled: Optional[bool] = None, enabled: Optional[bool] = None
-    ) -> List[Task]:
+        self,
+        scheduled: bool | None = None,
+        enabled: bool | None = None,
+    ) -> list[Task]:
         task_info = await self.get_tasks_info()
         task_map = {info.name: info for info in task_info}
         tasks = []
@@ -105,7 +97,7 @@ def task_from_registry(self, task: str | Task) -> Task:
         else:
             if task_ := self.registry.get(task):
                 return task_
-            raise UnknownTask(task)
+            raise UnknownTaskError(task)
 
     def register_task(self, task: Task) -> None:
         self.registry[task.name] = task
@@ -114,41 +106,15 @@ async def enable_task(self, task_name: str, enable: bool = True) -> TaskInfo:
         """Enable or disable a registered task"""
         task = self.registry.get(task_name)
         if not task:
-            raise UnknownTask(task_name)
+            raise UnknownTaskError(task_name)
         return await self.update_task(task, dict(enabled=enable))
 
-    def task_run_from_data(self, data: Dict[str, Any]) -> TaskRun:
-        """Build a TaskRun object from its metadata"""
-        data = data.copy()
-        name = data.pop("name")
-        data["task"] = self.task_from_registry(name)
-        return TaskRun(**data)
-
-    def task_run_data(
-        self, queued_task: QueuedTask, state: TaskState
-    ) -> Dict[str, Any]:
-        """Create a dictionary of metadata required by a task run
-
-        This dictionary must be serializable by the broker
-        """
-        task = self.task_from_registry(queued_task.task)
-        priority = queued_task.priority or task.priority
-        return dict(
-            id=queued_task.run_id,
-            name=task.name,
-            priority=priority.name,
-            state=state.name,
-            params=queued_task.params,
-            queued=Timestamp.utcnow(),
-        )
-
     @classmethod
     def from_url(cls, url: str = "") -> Broker:
         p = URL(url or broker_url_from_env())
-        Factory = _brokers.get(p.scheme)
-        if not Factory:
-            raise RuntimeError(f"Invalid broker {p}")
-        return Factory(p)
+        if factory := _brokers.get(p.scheme):
+            return factory(p)
+        raise RuntimeError(f"Invalid broker {p}")
 
     @classmethod
     def register_broker(cls, name: str, factory: type[Broker]) -> None:
@@ -160,7 +126,11 @@ class RedisBroker(Broker):
 
     @cached_property
     def redis(self) -> FluidRedis:
-        return FluidRedis(str(self.url.with_query({})), name=self.name)
+        return FluidRedis.create(str(self.url.with_query({})), name=self.name)
+
+    @property
+    def redis_cli(self) -> Redis:
+        return self.redis.redis_cli
 
     @property
     def name(self) -> str:
@@ -185,7 +155,7 @@ def task_queue_name(self, priority: TaskPriority) -> str:
         return f"{self.name}-queue-{priority.name}"
 
     async def get_tasks_info(self, *task_names: str) -> list[TaskInfo]:
-        pipe = self.redis.cli.pipeline()
+        pipe = self.redis_cli.pipeline()
         names = task_names or self.registry
         requested_task_names = []
         for name in names:
@@ -199,7 +169,7 @@ async def get_tasks_info(self, *task_names: str) -> list[TaskInfo]:
         ]
 
     async def update_task(self, task: Task, params: dict[str, Any]) -> TaskInfo:
-        pipe = self.redis.cli.pipeline()
+        pipe = self.redis_cli.pipeline()
         pipe.hset(
             self.task_hash_name(task.name),
             mapping={name: json.dumps(value) for name, value in params.items()},
@@ -208,41 +178,53 @@ async def update_task(self, task: Task, params: dict[str, Any]) -> TaskInfo:
         _, info = await pipe.execute()
         return self._decode_task(task, info)
 
-    async def queue_length(self) -> Dict[str, int]:
+    async def queue_length(self) -> dict[str, int]:
         if self.task_queue_names:
-            pipe = self.redis.cli.pipeline()
+            pipe = self.redis_cli.pipeline()
             for name in self.task_queue_names:
                 pipe.llen(name)
             result = await pipe.execute()
-            return {p.name: r for p, r in zip(TaskPriority, result)}
+            return dict(zip(TaskPriority, result))
         return {}
 
     async def close(self) -> None:
         """Close the broker on shutdown"""
         await self.redis.close()
 
-    async def get_task_run(self) -> Optional[TaskRun]:
+    async def get_task_run(self, task_manager: TaskManager) -> TaskRun | None:
         if self.task_queue_names:
-            data = await self.redis.cli.brpop(self.task_queue_names, timeout=1)
-            if data:
-                data_str = data[1].decode("utf-8")
-                return self.task_run_from_data(json.loads(data_str))
+            if redis_data := await self.redis_cli.brpop(  # type: ignore [misc]
+                self.task_queue_names,  # type: ignore [arg-type]
+                timeout=1,
+            ):
+                data = json.loads(redis_data[1])
+                data.update(
+                    task=self.task_from_registry(data["task"]),
+                    task_manager=task_manager,
+                )
+                return TaskRun(**data)
         return None
 
-    async def queue_task(self, queued_task: QueuedTask) -> TaskRun:
-        task = self.task_from_registry(queued_task.task)
-        priority = queued_task.priority or task.priority
-        data = self.task_run_data(queued_task, TaskState.queued)
-        await self.redis.cli.lpush(self.task_queue_name(priority), json.dumps(data))
-        return self.task_run_from_data(data)
+    async def queue_task(
+        self, task_manager: TaskManager, queued_task: QueuedTask
+    ) -> TaskRun:
+        task_run = self.create_task_run(task_manager, queued_task)
+        await self.redis_cli.lpush(  # type: ignore [misc]
+            self.task_queue_name(task_run.priority),
+            task_run.model_dump_json(),
+        )
+        return task_run
+
+    def lock(self, name: str, timeout: float | None = None) -> Lock:
+        return self.redis_cli.lock(name, timeout=timeout)
 
     def _decode_task(self, task: Task, data: dict[bytes, Any]) -> TaskInfo:
         info = {name.decode(): json.loads(value) for name, value in data.items()}
         return TaskInfo(
             name=task.name,
             description=task.description,
             schedule=str(task.schedule) if task.schedule else None,
-            priority=task.priority.name,
+            priority=task.priority,
             enabled=info.get("enabled", True),
             last_run_duration=info.get("last_run_duration"),
             last_run_end=info.get("last_run_end"),
@@ -251,3 +233,4 @@ def _decode_task(self, task: Task, data: dict[bytes, Any]) -> TaskInfo:
 
 
 Broker.register_broker("redis", RedisBroker)
+Broker.register_broker("rediss", RedisBroker)