Refactor workers (#62)

* Refactor workers

* Refactor workers

Author: lsbardel

.vscode/launch.json

@@ -13,7 +13,7 @@
             "env": {},
             "args": [
                 "-x",
-                "tests/scheduler/test_scheduler.py::test_disabled_execution"
+                "tests/scheduler/test_scheduler.py"
             ],
             "debugOptions": [
                 "RedirectOutput"
@@ -50,5 +50,16 @@
                 "RedirectOutput"
             ]
         },
+        {
+            "name": "AIO-FLUID simple_worker",
+            "type": "python",
+            "request": "launch",
+            "program": "${workspaceFolder}/docs_src/simple_worker.py",
+            "cwd": "${workspaceFolder}",
+            "justMyCode": false,
+            "debugOptions": [
+                "RedirectOutput"
+            ]
+        },
     ]
 }

docs/index.md

@@ -30,7 +30,7 @@ pip install aio-fluid
 To install all the dependencies:
 
 ```
-pip install aio-fluid[cli, db, http, log]
+pip install aio-fluid[cli, db, http, log, k8s]
 ```
 this includes the following extra dependencies:
 

docs/reference/workers.md

@@ -7,11 +7,9 @@ There are several worker classes which can be imported from `fluid.utils.worker`
 from fluid.utils.worker import StoppingWorker
 ```
 
-::: fluid.utils.worker.Worker
-
-::: fluid.utils.worker.RunningWorker
+::: fluid.utils.worker.WorkerState
 
-::: fluid.utils.worker.StoppingWorker
+::: fluid.utils.worker.Worker
 
 ::: fluid.utils.worker.WorkerFunction
 
@@ -22,5 +20,3 @@ from fluid.utils.worker import StoppingWorker
 ::: fluid.utils.worker.AsyncConsumer
 
 ::: fluid.utils.worker.Workers
-
-::: fluid.utils.worker.DynamicWorkers

docs/tutorials/workers.md

@@ -0,0 +1,25 @@
+# Workers
+
+Workers are the main building block for asynchronous programming with `aio-fluid`. They are responsible for running tasks and managing their lifecycle.
+All workers implemented derive from the base abstract class [Worker][fluid.utils.worker.Worker] where the main method to implement is the [Worker.run][fluid.utils.worker.Worker.run] method.
+
+## Worker Lifecycle
+
+The lifecycle of a worker is managed by the [WorkerState][fluid.utils.worker.WorkerState] class which provides a set of states that a worker can be in. The worker starts in an inital state and than it can be started and stopped.
+
+### Startup
+
+To start a worker one uses the async method [Worker.startup][fluid.utils.worker.Worker.startup] which create the task running the worker. The task will transition the worker from [WorkerState.INIT][fluid.utils.worker.WorkerState.INIT] to the [WorkerState.RUNNING][fluid.utils.worker.WorkerState.RUNNING] state. The worker will then run the [Worker.on_startup][fluid.utils.worker.Worker.on_startup] coroutine method (which by default is a no-op) follow by the main worker coroutine method [Worker.run][fluid.utils.worker.Worker.run] method until it is stopped.
+
+This is a very simple example of a worker that prints a message every second until it is stopped:
+
+```python
+--8<-- "./docs_src/simple_worker.py"
+```
+
+### Shutdown
+
+To shut down a worker there are few possibilities.
+
+* Direct call to the async [Worker.shutdown][fluid.utils.worker.Worker.shutdown] method which will trigger the graceful shutdown and wait for the worker to finish its work.
+* Call the [Worker.gracefully_stop][fluid.utils.worker.Worker.gracefully_stop] method which will trigger the graceful shutdown. Importantly, this method does not wait for the worker to finish its work, ti simply transition from the [WorkerState.RUNNING][fluid.utils.worker.WorkerState.RUNNING] to [WorkerState.STOPPING][fluid.utils.worker.WorkerState.STOPPING] state. To wait for the worker exit one should call the async [Worker.wait_for_shutdown][fluid.utils.worker.Worker.wait_for_shutdown] method (as in the example above)

docs_src/simple_worker.py

@@ -0,0 +1,27 @@
+import asyncio
+
+from fluid.utils.worker import Worker
+
+
+class SimpleWorker(Worker):
+    async def run(self):
+        while self.is_running():
+            self.print_message()
+            await asyncio.sleep(1)
+
+    def print_message(self):
+        print(f"Hello from {self.worker_name} in state {self.worker_state}")
+
+
+
+async def main():
+    worker = SimpleWorker()
+    worker.print_message()
+    await worker.startup()
+    asyncio.get_event_loop().call_later(5, worker.gracefully_stop)
+    await worker.wait_for_shutdown()
+    worker.print_message()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/tasks/__init__.py

@@ -3,7 +3,7 @@
 import time
 from dataclasses import dataclass, field
 from datetime import timedelta
-from typing import Self, cast
+from typing import Any, Self, cast
 
 from fastapi import FastAPI
 from pydantic import BaseModel, Field
@@ -23,8 +23,10 @@ def get(cls, context: TaskRun) -> Self:
         return context.deps
 
 
-def task_scheduler() -> TaskScheduler:
-    task_manager = TaskScheduler(deps=Deps())
+def task_scheduler(*, deps: Deps | None = None, **kwargs: Any) -> TaskScheduler:
+    deps = deps or Deps()
+    task_manager = TaskScheduler(deps=deps, **kwargs)
+    task_manager.add_async_context_manager(deps.http_client)
     task_manager.register_from_dict(globals())
     return task_manager
 
@@ -47,9 +49,10 @@ async def dummy(context: TaskRun[Sleep]) -> None:
 
 
 @task(schedule=every(timedelta(seconds=2)))
-async def scheduled(context: TaskRun) -> None:
-    """A simple scheduled task"""
-    await asyncio.sleep(0.1)
+async def ping(context: TaskRun) -> None:
+    """A simple scheduled task that ping the broker"""
+    redis_cli = cast(RedisTaskBroker, context.task_manager.broker).redis_cli
+    await redis_cli.ping()
 
 
 class AddValues(BaseModel):

fluid/__init__.py

@@ -1,3 +1,3 @@
 """Reusable server side python modules"""
 
-__version__ = "1.3.4"
+__version__ = "1.4.0"

fluid/scheduler/consumer.py

@@ -86,19 +86,28 @@ def __init__(
             ),
         ] = TaskDispatcher()
         self.broker = TaskBroker.from_url(self.config.broker_url)
+        self._async_contexts: list[Any] = []
         self._stack = AsyncExitStack()
 
     async def __aenter__(self) -> Self:
+        for cm in self._async_contexts:
+            await self._stack.enter_async_context(cm)
         return self
 
     async def __aexit__(self, *exc_info: Any) -> None:
         try:
-            await self._stack.aclose()
+            await self._stack.__aexit__(*exc_info)
         finally:
-            await self.on_shutdown()
+            await self.broker.close()
 
-    async def enter_async_context(self, cm: Any) -> Any:
-        return await self._stack.enter_async_context(cm)
+    async def on_startup(self) -> None:
+        await self.__aenter__()
+
+    async def on_shutdown(self) -> None:
+        await self.__aexit__(None, None, None)
+
+    def add_async_context_manager(self, cm: Any) -> None:
+        self._async_contexts.append(cm)
 
     @property
     def registry(self) -> TaskRegistry:
@@ -115,9 +124,6 @@ async def execute(self, task: Task | str, **params: Any) -> TaskRun:
         await task_run.execute()
         return task_run
 
-    async def on_shutdown(self) -> None:
-        await self.broker.close()
-
     def execute_sync(self, task: Task | str, **params: Any) -> TaskRun:
         return asyncio.run(self._execute_and_exit(task, **params))
 
@@ -263,10 +269,9 @@ async def _queue_task(self) -> None:
         try:
             task = self._task_to_queue.pop()
         except IndexError:
-            await asyncio.sleep(0.1)
+            pass
         else:
             await self.queue(task)
-            await asyncio.sleep(0)
 
     async def _consume_tasks(self, worker_name: str) -> None:
         if not self.config.consume_tasks:

fluid/scheduler/endpoints.py

@@ -152,4 +152,7 @@ def setup_fastapi(
     app.state.task_manager = task_manager
     if isinstance(task_manager, Worker):
         app_workers(app).add_workers(task_manager)
+    else:
+        app.add_event_handler("startup", task_manager.on_startup)
+        app.add_event_handler("shutdown", task_manager.on_shutdown)
     return app

fluid/tools_fastapi/__init__.py

@@ -1,17 +1,32 @@
 from __future__ import annotations
 
-from fastapi import FastAPI
+from typing import Annotated
+
+from fastapi import Depends, FastAPI, Request
 
 from fluid.utils.worker import Workers
 
 from .service import FastapiAppWorkers
 
 
+def get_workers_from_request(request: Request) -> Workers | None:
+    """Get workers from request."""
+    return get_workers_from_app(request.app)
+
+
+def get_workers_from_app(app: FastAPI) -> Workers | None:
+    """Get workers from request."""
+    return getattr(app.state, "workers", None)
+
+
 def app_workers(app: FastAPI) -> Workers:
     """Get workers from app state or create new workers."""
-    if workers := getattr(app.state, "workers", None):
+    if workers := get_workers_from_app(app):
         return workers
     else:
         workers = FastapiAppWorkers.setup(app)
         app.state.workers = workers
         return workers
+
+
+WorkersDep = Annotated[Workers | None, Depends(get_workers_from_request)]

fluid/tools_fastapi/service.py

@@ -6,7 +6,7 @@
 
 from fluid import settings
 from fluid.utils import log
-from fluid.utils.worker import Worker, Workers
+from fluid.utils.worker import Workers
 
 logger = log.get_logger(__name__)
 
@@ -21,13 +21,6 @@ def setup(cls, app: FastAPI, **kwargs: Any) -> Self:
         app.add_event_handler("shutdown", workers.shutdown)
         return workers
 
-    def bail_out(self, reason: str, code: int = 1) -> None:
+    def after_shutdown(self, reason: str, code: int = 1) -> None:
         if settings.ENV != "test":
-            logger.warning("shutting down due to %s", reason)
             os.kill(os.getpid(), signal.SIGTERM)
-
-    def get_active_worker(self, *, worker_name: str) -> Worker | None:
-        worker = self._workers.get_worker_by_name(worker_name)
-        if worker and not worker.is_stopping():
-            return worker
-        return None

fluid/utils/errors.py

@@ -22,3 +22,7 @@ def __init__(
 
     def __str__(self) -> str:
         return self.msg
+
+
+class WorkerStartError(FluidError):
+    """Worker start error"""