feat: Add plan name to logs (#1178)

Fixes #883

---------

Co-authored-by: Joseph Ware <[email protected]>

Author: dan-fernandes

docs/reference/openapi.yaml

@@ -224,6 +224,28 @@ components:
       - new_state
       title: StateChangeRequest
       type: object
+    Task:
+      additionalProperties: false
+      description: Task that will run a plan
+      properties:
+        metadata:
+          additionalProperties: true
+          description: Any metadata to apply to all runs within this task
+          title: Metadata
+          type: object
+        name:
+          description: Name of plan to run
+          title: Name
+          type: string
+        params:
+          additionalProperties: true
+          description: Values for parameters to plan, if any
+          title: Params
+          type: object
+      required:
+      - name
+      title: Task
+      type: object
     TaskRequest:
       additionalProperties: false
       description: Request to run a task with related info
@@ -293,7 +315,7 @@ components:
           title: Request Id
           type: string
         task:
-          title: Task
+          $ref: '#/components/schemas/Task'
         task_id:
           title: Task Id
           type: string
@@ -355,7 +377,7 @@ info:
     name: Apache 2.0
     url: https://www.apache.org/licenses/LICENSE-2.0.html
   title: BlueAPI Control
-  version: 1.1.0
+  version: 1.1.1
 openapi: 3.1.0
 paths:
   /config/oidc:

src/blueapi/client/client.py

@@ -25,7 +25,7 @@
     TasksListResponse,
     WorkerTask,
 )
-from blueapi.worker import Task, TrackableTask, WorkerEvent, WorkerState
+from blueapi.worker import TrackableTask, WorkerEvent, WorkerState
 from blueapi.worker.event import ProgressEvent, TaskStatus
 
 from .event_bus import AnyEvent, BlueskyStreamingError, EventBusClient, OnAnyEvent
@@ -159,15 +159,15 @@ def resume(self) -> WorkerState:
         return self._rest.set_state(WorkerState.RUNNING, defer=False)
 
     @start_as_current_span(TRACER, "task_id")
-    def get_task(self, task_id: str) -> TrackableTask[Task]:
+    def get_task(self, task_id: str) -> TrackableTask:
         """
         Get a task stored by the worker
 
         Args:
             task_id: Unique ID for the task
 
         Returns:
-            TrackableTask[Task]: Task details
+            TrackableTask: Task details
         """
         assert task_id, "Task ID not provided!"
         return self._rest.get_task(task_id)

src/blueapi/client/rest.py

@@ -26,7 +26,7 @@
     TasksListResponse,
     WorkerTask,
 )
-from blueapi.worker import Task, TrackableTask, WorkerState
+from blueapi.worker import TrackableTask, WorkerState
 
 T = TypeVar("T")
 
@@ -172,8 +172,8 @@ def set_state(
             data={"new_state": state, "defer": defer},
         )
 
-    def get_task(self, task_id: str) -> TrackableTask[Task]:
-        return self._request_and_deserialize(f"/tasks/{task_id}", TrackableTask[Task])
+    def get_task(self, task_id: str) -> TrackableTask:
+        return self._request_and_deserialize(f"/tasks/{task_id}", TrackableTask)
 
     def get_all_tasks(self) -> TasksListResponse:
         return self._request_and_deserialize("/tasks", TasksListResponse)

src/blueapi/log.py

@@ -1,5 +1,6 @@
 import logging
 import os
+from contextlib import contextmanager
 
 from graypy import GELFTCPHandler
 
@@ -21,6 +22,39 @@ def filter(self, record: logging.LogRecord) -> bool:
         return True
 
 
+class PlanTagFilter(logging.Filter):
+    """Filter to attach name of plan as an attribute to LogRecords.
+
+    Attaches the attribute `plan_name` to all LogRecords that are passed through.
+    """
+
+    def __init__(self, plan_name: str):
+        self.plan_name = plan_name
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        record.plan_name = self.plan_name
+        return True
+
+
+@contextmanager
+def plan_tag_filter_context(plan_name: str, logger: logging.Logger):
+    """Context manager that attaches and removes `PlanTagFilter` to a given logger.
+
+    Creates an instance of PlanTagFilter and attaches it to the given logger for the
+    duration of the context. On exit the filter is removed.
+
+        Args:
+            plan_name: str name of plan being executed
+            logger: logging.Logger to attach filter to
+    """
+    filter = PlanTagFilter(plan_name)
+    try:
+        logger.addFilter(filter)
+        yield
+    finally:
+        logger.removeFilter(filter)
+
+
 def set_up_logging(logging_config: LoggingConfig) -> None:
     """Configure root level logger for blueapi.
 

src/blueapi/service/main.py

@@ -58,7 +58,7 @@
 from .runner import WorkerDispatcher
 
 #: API version to publish in OpenAPI schema
-REST_API_VERSION = "1.1.0"
+REST_API_VERSION = "1.1.1"
 
 LICENSE_INFO: dict[str, str] = {
     "name": "Apache 2.0",

src/blueapi/worker/task_worker.py

@@ -6,7 +6,7 @@
 from functools import partial
 from queue import Full, Queue
 from threading import Event, RLock
-from typing import Any, Generic, TypeVar
+from typing import Any, TypeVar
 
 from bluesky.protocols import Status
 from observability_utils.tracing import (
@@ -31,6 +31,7 @@
     WatchableStatus,
 )
 from blueapi.core.bluesky_event_loop import configure_bluesky_event_loop
+from blueapi.log import plan_tag_filter_context
 from blueapi.utils.base_model import BlueapiBaseModel
 from blueapi.utils.thread_exception import handle_all_exceptions
 
@@ -57,13 +58,13 @@
 T = TypeVar("T")
 
 
-class TrackableTask(BlueapiBaseModel, Generic[T]):
+class TrackableTask(BlueapiBaseModel):
     """
     A representation of a task that the worker recognizes
     """
 
     task_id: str
-    task: T
+    task: Task
     request_id: str | SkipJsonSchema[None] = None
     is_complete: bool = False
     is_pending: bool = True
@@ -229,7 +230,7 @@ def get_tasks_by_status(self, status: TaskStatusEnum) -> list[TrackableTask]:
         return []
 
     @start_as_current_span(TRACER)
-    def get_active_task(self) -> TrackableTask[Task] | None:
+    def get_active_task(self) -> TrackableTask | None:
         """
         Returns the task the worker is currently running
         Returns:
@@ -254,7 +255,8 @@ def begin_task(self, task_id: str) -> None:
         if task is None:
             raise KeyError(f"No pending task with ID {task_id}")
         else:
-            self._submit_trackable_task(task)
+            with plan_tag_filter_context(task.task.name, LOGGER):
+                self._submit_trackable_task(task)
 
     @start_as_current_span(TRACER, "task.name", "task.params")
     def submit_task(self, task: Task) -> str:
@@ -419,19 +421,29 @@ def _cycle(self) -> None:
             LOGGER.info("Awaiting task")
             next_task: TrackableTask | KillSignal = self._task_channel.get()
             if isinstance(next_task, TrackableTask):
-                if self._current_task_otel_context is not None:
-                    with TRACER.start_as_current_span(
-                        "_cycle",
-                        context=self._current_task_otel_context,
-                        kind=SpanKind.SERVER,
-                    ):
-                        LOGGER.info(f"Got new task: {next_task}")
-                        self._current = next_task
-                        self._current_task_otel_context = get_current()
-                        add_span_attributes({"next_task.task_id": next_task.task_id})
-
-                        self._current.is_pending = False
-                        self._current.task.do_task(self._ctx)
+
+                def process_task():
+                    LOGGER.info(f"Got new task: {next_task}")
+                    self._current = next_task
+                    self._current.is_pending = False
+                    self._current.task.do_task(self._ctx)
+
+                with plan_tag_filter_context(next_task.task.name, LOGGER):
+                    if self._current_task_otel_context is not None:
+                        with TRACER.start_as_current_span(
+                            "_cycle",
+                            context=self._current_task_otel_context,
+                            kind=SpanKind.SERVER,
+                        ):
+                            self._current_task_otel_context = get_current()
+                            add_span_attributes(
+                                {"next_task.task_id": next_task.task_id}
+                            )
+
+                            process_task()
+                    else:
+                        process_task()
+
             elif isinstance(next_task, KillSignal):
                 # If we receive a kill signal we begin to shut the worker down.
                 # Note that the kill signal is explicitly not a type of task as we don't

tests/unit_tests/service/test_interface.py

@@ -234,9 +234,9 @@ def test_begin_task_no_task_id(worker_mock: MagicMock):
 
 @patch("blueapi.service.interface.TaskWorker.get_tasks_by_status")
 def test_get_tasks_by_status(get_tasks_by_status_mock: MagicMock):
-    pending_task1 = TrackableTask(task_id="0", task=None)
-    pending_task2 = TrackableTask(task_id="1", task=None)
-    running_task = TrackableTask(task_id="2", task=None)
+    pending_task1 = TrackableTask(task_id="0", task=Task(name="pending_task1"))
+    pending_task2 = TrackableTask(task_id="1", task=Task(name="pending_task2"))
+    running_task = TrackableTask(task_id="2", task=Task(name="running_task"))
 
     def mock_tasks_by_status(status: TaskStatusEnum) -> list[TrackableTask]:
         if status == TaskStatusEnum.PENDING:
@@ -307,9 +307,9 @@ def test_cancel_active_task(cancel_active_task_mock: MagicMock):
 @patch("blueapi.service.interface.TaskWorker.get_tasks")
 def test_get_tasks(get_tasks_mock: MagicMock):
     tasks = [
-        TrackableTask(task_id="0", task=None),
-        TrackableTask(task_id="1", task=None),
-        TrackableTask(task_id="2", task=None),
+        TrackableTask(task_id="0", task=Task(name="0")),
+        TrackableTask(task_id="1", task=Task(name="1")),
+        TrackableTask(task_id="2", task=Task(name="2")),
     ]
     get_tasks_mock.return_value = tasks
 

tests/unit_tests/service/test_rest_api.py

@@ -300,7 +300,7 @@ def test_put_plan_fails_if_not_idle(mock_runner: Mock, client: TestClient) -> No
 
     # Set to non idle
     mock_runner.run.return_value = TrackableTask(
-        task=None, task_id=task_id_current, is_complete=False
+        task=Task(name="none"), task_id=task_id_current, is_complete=False
     )
 
     resp = client.put("/worker/task", json={"task_id": task_id_new})

tests/unit_tests/test_log.py

@@ -1,12 +1,13 @@
 import logging
-from unittest.mock import patch
+from collections.abc import Generator
+from unittest.mock import Mock, patch
 
 import pytest
 from bluesky.log import logger as bluesky_logger
 from dodal.log import LOGGER as dodal_logger
 
 from blueapi.config import GraylogConfig, LoggingConfig
-from blueapi.log import set_up_logging
+from blueapi.log import PlanTagFilter, plan_tag_filter_context, set_up_logging
 
 
 def clear_all_loggers_and_handlers(logger):
@@ -23,7 +24,7 @@ def clear_all_loggers_and_handlers(logger):
 
 
 @pytest.fixture(params=LOGGER_NAMES)
-def logger_with_graylog(request):
+def logger_with_graylog(request) -> Generator[logging.Logger]:
     logger = logging.getLogger(request.param)
     graylog_config = GraylogConfig(enabled=True)
     set_up_logging(LoggingConfig(graylog=graylog_config))
@@ -32,7 +33,7 @@ def logger_with_graylog(request):
 
 
 @pytest.fixture(params=LOGGER_NAMES)
-def logger_without_graylog(request):
+def logger_without_graylog(request) -> Generator[logging.Logger]:
     logger = logging.getLogger(request.param)
     graylog_config = GraylogConfig(enabled=False)
     set_up_logging(LoggingConfig(graylog=graylog_config))
@@ -41,19 +42,19 @@ def logger_without_graylog(request):
 
 
 @pytest.fixture
-def logger(logger_with_graylog):
+def logger(logger_with_graylog) -> logging.Logger:
     return logger_with_graylog
 
 
 @pytest.fixture
-def mock_stream_handler_emit():
+def mock_stream_handler_emit() -> Generator[Mock]:
     with patch("blueapi.log.logging.StreamHandler.emit") as stream_handler_emit:
         # stream_handler_emit.reset_mock()
         yield stream_handler_emit
 
 
 @pytest.fixture
-def mock_graylog_emit():
+def mock_graylog_emit() -> Generator[Mock]:
     with patch("blueapi.log.GELFTCPHandler.emit") as graylog_emit:
         yield graylog_emit
 
@@ -65,7 +66,7 @@ def mock_graylog_emit():
 
 
 @pytest.fixture(params=MOCK_HANDLER_EMIT_STRINGS)
-def mock_handler_emit(request):
+def mock_handler_emit(request) -> Generator[Mock]:
     with patch(request.param) as mock_emit:
         yield mock_emit
 
@@ -101,3 +102,25 @@ def test_library_logger_intergrations(logger, library_logger, mock_handler_emit)
     mock_handler_emit.assert_not_called()
     library_logger.info("FOO")
     mock_handler_emit.assert_called()
+
+
+def test_plan_tag_filter_context_adds_filter(logger):
+    assert logger.filters == []
+    with plan_tag_filter_context("foo", logger):
+        filters = logger.filters
+        assert len(filters) == 1
+        assert isinstance(filters[0], PlanTagFilter)
+
+
+def test_plan_tag_filter_context_removes_filter(logger):
+    with plan_tag_filter_context("foo", logger):
+        filters = logger.filters
+        assert len(filters) == 1
+        assert isinstance(filters[0], PlanTagFilter)
+    assert logger.filters == []
+
+
+def test_filter_tags_plan_name(logger, mock_handler_emit):
+    logger.addFilter(PlanTagFilter("foo"))
+    logger.info("FOO")
+    assert mock_handler_emit.call_args[0][0].plan_name == "foo"