Allow to add tags on registration (#102)

* Allow to add tags on registration

* Fix log config

Author: lsbardel

docs/release-notes.md

@@ -1,5 +1,19 @@
 # Release Notes
 
+## v2.4.2
+
+Tasks can now be tagged at registration time.
+
+- The task registration methods accept an optional `tags` argument. When
+  provided, the extra tags are merged into each task's own tags as it is
+  registered — applied across `register_task`, `register_from_module`, and
+  `register_from_dict`.
+  ([#102](https://github.com/quantmind/aio-fluid/pull/102))
+- Bumped `python-json-logger` to `>= 4.1.0` and switched JSON logging to the
+  new `pythonjsonlogger.json` formatter path (the old `pythonjsonlogger.jsonlogger`
+  module is deprecated).
+  ([#102](https://github.com/quantmind/aio-fluid/pull/102))
+
 ## v2.4.1
 
 Task history can now be filtered by task tags.

docs/tutorials/task_app.md

@@ -28,6 +28,15 @@ if __name__ == "__main__":
     asyncio.run(consumer.run())
 ```
 
+Pass `tags` to [register_from_module][fluid.scheduler.TaskManager.register_from_module]
+(or [register_from_dict][fluid.scheduler.TaskManager.register_from_dict] /
+[register_task][fluid.scheduler.TaskManager.register_task]) to add extra tags to
+every registered task on top of the tags already declared on each task:
+
+```python
+consumer.register_from_module(task_module_a, tags=["module-a"])
+```
+
 ## FastAPI Integration
 
 A [TaskManager][fluid.scheduler.TaskManager] can be integrated with FastAPI so that

docs/tutorials/task_k8s.md

@@ -1,6 +1,6 @@
 # K8s Jobs
 
-When the [TaskConsumer][fluid.scheduler.TaskConsumer] runs inside a Kubernetes cluster, [CPU bound tasks](/tutorials/task_queue/#cpu-bound-tasks) can be dispatched as [Kubernetes Jobs](https://kubernetes.io/docs/concepts/workloads/controllers/job/) instead of local subprocesses.
+When the [TaskConsumer][fluid.scheduler.TaskConsumer] runs inside a Kubernetes cluster, [CPU bound tasks](task_queue.md#cpu-bound-tasks) can be dispatched as [Kubernetes Jobs](https://kubernetes.io/docs/concepts/workloads/controllers/job/) instead of local subprocesses.
 This offloads heavy computation to dedicated pods and keeps the consumer event loop free.
 
 ## How it works

fluid/scheduler/consumer.py

@@ -7,7 +7,7 @@
 from functools import partial
 from time import monotonic
 from types import ModuleType
-from typing import Any, Awaitable, Callable, Self
+from typing import Any, Awaitable, Callable, Self, Sequence
 
 from starlette.datastructures import State
 from typing_extensions import Annotated, Doc
@@ -136,8 +136,17 @@ def type(self) -> str:
         """The type of the task manager"""
         return snake_case(self.__class__.__name__)
 
-    def register_task(self, task: Annotated[Task, Doc("Task to register")]) -> None:
+    def register_task(
+        self,
+        task: Annotated[Task, Doc("Task to register")],
+        tags: Annotated[
+            Sequence[str] | None,
+            Doc("Extra tags to add to the task before registering it"),
+        ] = None,
+    ) -> None:
         """Register a task with the task manager"""
+        if tags:
+            task = task._replace(tags=task.tags | frozenset(tags))
         self.broker.register_task(task)
 
     async def execute(
@@ -308,13 +317,17 @@ def register_from_module(
                 "- can contain any object, only instances of Task are registered"
             ),
         ],
+        tags: Annotated[
+            Sequence[str] | None,
+            Doc("Extra tags to add to every registered task"),
+        ] = None,
     ) -> None:
         """Register tasks from a python module"""
         for name in dir(module):
             if name.startswith("_"):
                 continue
             if isinstance(obj := getattr(module, name), Task):
-                self.register_task(obj)
+                self.register_task(obj, tags=tags)
 
     def register_from_dict(
         self,
@@ -325,13 +338,17 @@ def register_from_dict(
                 "- can contain any object, only instances of Task are registered"
             ),
         ],
+        tags: Annotated[
+            Sequence[str] | None,
+            Doc("Extra tags to add to every registered task"),
+        ] = None,
     ) -> None:
         """Register tasks from a python dictionary"""
         for name, obj in data.items():
             if name.startswith("_"):
                 continue
             if isinstance(obj, Task):
-                self.register_task(obj)
+                self.register_task(obj, tags=tags)
 
     def register_async_handler(
         self,

fluid/utils/log.py

@@ -59,11 +59,11 @@ def _log_config(
         )
         log_formatters.update(
             json={
-                "()": "pythonjsonlogger.jsonlogger.JsonFormatter",
+                "()": "pythonjsonlogger.json.JsonFormatter",
                 "format": log_format,
             },
             nicejson={
-                "()": "pythonjsonlogger.jsonlogger.JsonFormatter",
+                "()": "pythonjsonlogger.json.JsonFormatter",
                 "format": log_format,
                 "json_indent": 2,
             },

fluid/utils/worker.py

@@ -477,16 +477,9 @@ class Workers(Worker):
     starts each child worker and monitors their health — if any child stops
     unexpectedly the whole group is gracefully stopped.
 
-    On shutdown all child workers are stopped concurrently via
-    [_wait_for_workers][fluid.utils.worker.Workers._wait_for_workers].
+    On shutdown all child workers are stopped concurrently.
     Workers that do not exit within `stopping_grace_period` seconds are
     force-cancelled.
-
-    !!! note "Shutdown ordering"
-        By default all child workers are shut down concurrently. If a subclass
-        needs a specific shutdown order (e.g. an event dispatcher that must
-        outlive the workers that produce events), override
-        [_wait_for_workers][fluid.utils.worker.Workers._wait_for_workers].
     """
 
     def __init__(

mkdocs.yml

@@ -2,6 +2,7 @@ site_name: Aio Fluid
 site_url: https://fluid.quantmind.com
 repo_name: quantmind/aio-fluid
 repo_url: https://github.com/quantmind/aio-fluid
+strict: true
 theme:
   name: material
   palette:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "aio-fluid"
-version = "2.4.1"
+version = "2.4.2"
 description = "Tools for backend python services"
 authors = [
     { name = "Luca Sbardella", email = "luca@quantmind.com" },
@@ -71,7 +71,7 @@ k8s = [
     "python-slugify >= 8.0.4",
 ]
 log = [
-    "python-json-logger >= 3.2.1",
+    "python-json-logger >= 4.1.0",
 ]
 
 [dependency-groups]

tests/scheduler/test_config.py

@@ -50,6 +50,14 @@ def test_register_from_module() -> None:
     assert task_manager.registry["dummy"]
 
 
+def test_register_from_module_with_tags() -> None:
+    task_manager = TaskManager()
+    task_manager.register_from_module(example_tasks, tags=["extra", "module"])
+    assert task_manager.registry
+    dummy = task_manager.registry["dummy"]
+    assert {"extra", "module"} <= dummy.tags
+
+
 def test_cpu_bount_params() -> None:
     cpu_bound = example_tasks.cpu_bound
     assert cpu_bound.params_model is example_tasks.Sleep

tests/utils/test_log.py

@@ -0,0 +1,61 @@
+import json
+import logging
+import warnings
+
+import pytest
+
+from fluid.utils import log
+
+pytestmark = pytest.mark.asyncio(loop_scope="module")
+
+
+async def test_get_level_num_from_int() -> None:
+    assert log.get_level_num(logging.DEBUG) == logging.DEBUG
+
+
+async def test_get_level_num_from_string() -> None:
+    assert log.get_level_num("info") == logging.INFO
+    assert log.get_level_num("WARNING") == logging.WARNING
+
+
+async def test_config_plain() -> None:
+    cfg = log.config(log_handler="plain", level="INFO")
+    assert cfg["handlers"]["plain"]["formatter"] == "plain"
+    assert cfg["loggers"]["fluid"]["level"] == logging.INFO
+
+
+async def test_config_other_level_floored_by_level() -> None:
+    cfg = log.config(level="ERROR", other_level=logging.WARNING)
+    # other loggers cannot be more verbose than the app level
+    assert cfg["root"]["level"] == logging.ERROR
+
+
+async def test_config_app_names() -> None:
+    cfg = log.config(app_names=["myapp"], level="INFO")
+    # the configured app and the default app name are both present
+    assert "myapp" in cfg["loggers"]
+    assert "fluid" in cfg["loggers"]
+
+
+async def test_config_extra_formatters() -> None:
+    cfg = log.config(formatters={"custom": {"format": "%(message)s"}})
+    assert "custom" in cfg["formatters"]
+
+
+@pytest.mark.parametrize("handler", ["json", "nicejson"])
+async def test_config_json_handlers(
+    handler: str, capsys: pytest.CaptureFixture
+) -> None:
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", DeprecationWarning)
+        log.config(log_handler=handler, level="INFO")
+    logging.getLogger("fluid").info("hello %s", handler)
+    captured = capsys.readouterr()
+    record = json.loads(captured.err)
+    assert record["message"] == f"hello {handler}"
+    assert record["levelname"] == "INFO"
+
+
+async def test_config_json_formatter_uses_new_path() -> None:
+    cfg = log.config(log_handler="json")
+    assert cfg["formatters"]["json"]["()"] == "pythonjsonlogger.json.JsonFormatter"