Add deps (#54)

* Add deps

* Add deps

* Add deps

Author: lsbardel

LICENSE

@@ -1,6 +1,6 @@
 BSD 3-Clause License
 
-Copyright (c) 2024, Quantmind
+Copyright (c) 2025, Quantmind
 All rights reserved.
 
 Redistribution and use in source and binary forms, with or without

examples/tasks/__init__.py

@@ -1,19 +1,30 @@
 import asyncio
 import os
 import time
+from dataclasses import dataclass, field
 from datetime import timedelta
-from typing import cast
+from typing import Self, cast
 
 from fastapi import FastAPI
 from pydantic import BaseModel, Field
 
 from fluid.scheduler import TaskRun, TaskScheduler, every, task
 from fluid.scheduler.broker import RedisTaskBroker
 from fluid.scheduler.endpoints import setup_fastapi
+from fluid.utils.http_client import HttpxClient
+
+
+@dataclass
+class Deps:
+    http_client: HttpxClient = field(default_factory=HttpxClient)
+
+    @classmethod
+    def get(cls, context: TaskRun) -> Self:
+        return context.deps
 
 
 def task_scheduler() -> TaskScheduler:
-    task_manager = TaskScheduler()
+    task_manager = TaskScheduler(deps=Deps())
     task_manager.register_from_dict(globals())
     return task_manager
 
@@ -63,3 +74,16 @@ async def cpu_bound(context: TaskRun) -> None:
     broker = cast(RedisTaskBroker, context.task_manager.broker)
     redis = broker.redis_cli
     await redis.setex(context.id, os.getpid(), 10)
+
+
+class Scrape(BaseModel):
+    url: str = Field(default="https://httpbin.org/get", description="URL to scrape")
+
+
+@task
+async def scrape(context: TaskRun[Scrape]) -> None:
+    """Scrape a website"""
+    deps = Deps.get(context)
+    response = await deps.http_client.get(context.params.url, callback=True)
+    text = await response.text()
+    context.logger.info(text)

fluid/scheduler/consumer.py

@@ -8,6 +8,7 @@
 from typing import Any, Awaitable, Callable, Self
 
 from inflection import underscore
+from starlette.datastructures import State
 from typing_extensions import Annotated, Doc
 
 from fluid.utils import log
@@ -47,10 +48,32 @@ class TaskManager:
     """The task manager is the main class for managing tasks"""
 
     def __init__(
-        self, *, config: TaskManagerConfig | None = None, **kwargs: Any
+        self,
+        *,
+        deps: Any = None,
+        config: TaskManagerConfig | None = None,
+        **kwargs: Any,
     ) -> None:
-        self.state: dict[str, Any] = {}
-        self.config: TaskManagerConfig = config or TaskManagerConfig(**kwargs)
+        self.deps: Annotated[
+            Any,
+            Doc(
+                """
+                Dependencies for the task manager.
+
+                Production applications requires global dependencies to be
+                available to all tasks. This can be achieved by setting
+                the `deps` attribute of the task manager to an object
+                with the required dependencies.
+
+                Each task can cast the dependencies to the required type.
+                """
+            ),
+        ] = (
+            deps if deps is not None else State()
+        )
+        self.config: Annotated[
+            TaskManagerConfig, Doc("""Task manager configuration""")
+        ] = config or TaskManagerConfig(**kwargs)
         self.dispatcher: Annotated[
             TaskDispatcher,
             Doc(

fluid/scheduler/models.py

@@ -191,7 +191,7 @@ class TaskRun(BaseModel, Generic[TP], arbitrary_types_allowed=True):
     async def execute(self) -> None:
         try:
             self.set_state(TaskState.running)
-            await self.task.executor(self)
+            await self.task.executor(self)  # type: ignore [arg-type]
         except Exception:
             self.set_state(TaskState.failure)
             raise
@@ -247,6 +247,10 @@ def is_done(self) -> bool:
     def is_failure(self) -> bool:
         return self.state.is_failure
 
+    @property
+    def deps(self) -> Any:
+        return self.task_manager.deps
+
     def set_state(
         self,
         state: TaskState,
@@ -293,7 +297,7 @@ def lock(self, timeout: float | None) -> Lock:
         return self.task_manager.broker.lock(self.name, timeout=timeout)
 
     def _dispatch(self) -> None:
-        self.task_manager.dispatcher.dispatch(self.model_copy())
+        self.task_manager.dispatcher.dispatch(self.model_copy())  # type: ignore [arg-type]
 
 
 @dataclass

tests/scheduler/test_endpoints.py

@@ -9,7 +9,7 @@
 
 async def test_get_tasks(cli: TaskClient) -> None:
     data = await cli.get(f"{cli.url}/tasks")
-    assert len(data) == 4
+    assert len(data) == 5
     tasks = {task["name"]: TaskInfo(**task) for task in data}
     dummy = tasks["dummy"]
     assert dummy.name == "dummy"

tests/scheduler/test_scheduler.py

@@ -1,6 +1,6 @@
 import asyncio
 import os
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 
 import pytest
 from redis.asyncio import Redis
@@ -17,17 +17,6 @@
 pytestmark = pytest.mark.asyncio(loop_scope="module")
 
 
-@dataclass
-class WaitFor:
-    name: str
-    times: int = 2
-    runs: list[TaskRun] = field(default_factory=list)
-
-    def __call__(self, task_run: TaskRun) -> None:
-        if task_run.name == self.name:
-            self.runs.append(task_run)
-
-
 async def test_scheduler_manager(task_scheduler: TaskScheduler) -> None:
     assert task_scheduler
     assert task_scheduler.broker.registry
@@ -71,16 +60,6 @@ async def test_dummy_rate_limit(task_scheduler: TaskScheduler) -> None:
     assert s1.state is TaskState.rate_limited or s2.state is TaskState.rate_limited
 
 
-@pytest.mark.flaky
-async def test_scheduled(task_scheduler: TaskScheduler) -> None:
-    handler = WaitFor(name="scheduled")
-    task_scheduler.dispatcher.register_handler("end.scheduled", handler)
-    try:
-        await wait_for(lambda: len(handler.runs) >= 2, timeout=3)
-    finally:
-        task_scheduler.dispatcher.unregister_handler("end.handler")
-
-
 @pytest.mark.flaky
 async def test_cpubound_execution(
     task_scheduler: TaskScheduler, redis: Redis  # type: ignore