context.logger for task execute methods

Author: lukasbindreiter

tilebox-workflows/pyproject.toml

@@ -34,9 +34,9 @@ dependencies = [
     "python-dateutil>=2.9.0.post0",
     "obstore>=0.8.2",
     "opentelemetry-proto>=1.30.0",
-    "structlog>=25.5.0",
     # grpcio 1.80.0 contains unwanted log message spam: https://github.com/grpc/grpc/issues/42293
     "grpcio<1.80.0",
+    "opentelemetry-instrumentation-logging>=0.62b1",
 ]
 
 [dependency-groups]

tilebox-workflows/tests/observability/test_tracing.py

@@ -0,0 +1,76 @@
+from collections.abc import Iterator
+
+import pytest
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan, Span, TracerProvider
+from opentelemetry.sdk.trace.export import SpanProcessor
+
+from tilebox.workflows.observability import tracing
+
+
+class RecordingSpanProcessor(SpanProcessor):
+    def __init__(self) -> None:
+        self.span_names: list[str] = []
+
+    def on_start(self, span: Span, parent_context: Context | None = None) -> None:
+        pass
+
+    def on_end(self, span: ReadableSpan) -> None:
+        self.span_names.append(span.name)
+
+    def shutdown(self) -> None:
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:  # noqa: ARG002
+        return True
+
+
+@pytest.fixture(autouse=True)
+def reset_tilebox_tracing() -> Iterator[None]:
+    tracing._set_tilebox_tracer_provider(TracerProvider())
+    tracing._workflow_tracers.clear()
+    yield
+    tracing._set_tilebox_tracer_provider(TracerProvider())
+    tracing._workflow_tracers.clear()
+
+
+@pytest.fixture
+def span_processors(monkeypatch: pytest.MonkeyPatch) -> list[RecordingSpanProcessor]:
+    processors: list[RecordingSpanProcessor] = []
+
+    def create_processor(*args: object, **kwargs: object) -> RecordingSpanProcessor:  # noqa: ARG001
+        processor = RecordingSpanProcessor()
+        processors.append(processor)
+        return processor
+
+    monkeypatch.setattr(tracing, "_otel_span_exporter", create_processor)
+    return processors
+
+
+def test_workflow_tracers_do_not_share_client_span_processors(
+    span_processors: list[RecordingSpanProcessor],
+) -> None:
+    tracers = [tracing.WorkflowTracer(service=None, url="https://api.tilebox.com", token=None) for _ in range(3)]
+
+    for index, tracer in enumerate(tracers):
+        with tracer.start_as_current_span(f"span-{index}"):
+            pass
+
+    assert [processor.span_names for processor in span_processors] == [["span-0"], ["span-1"], ["span-2"]]
+
+
+def test_workflow_tracers_copy_configured_span_processors_once(
+    span_processors: list[RecordingSpanProcessor],
+) -> None:
+    tracing.configure_otel_tracing(endpoint="https://otel.example.com")
+    tracers = [tracing.WorkflowTracer(service=None, url="https://api.tilebox.com", token=None) for _ in range(2)]
+
+    for index, tracer in enumerate(tracers):
+        with tracer.start_as_current_span(f"span-{index}"):
+            pass
+
+    assert [processor.span_names for processor in span_processors] == [
+        ["span-0", "span-1"],
+        ["span-0"],
+        ["span-1"],
+    ]

tilebox-workflows/tests/runner/test_runner.py

@@ -208,9 +208,10 @@ def test_runner_disallow_duplicate_task_identifiers() -> None:
         "dummy-cluster",
         InMemoryCache(),
         NoopWorkflowTracer(),
-        None,
         MagicMock(),
         RunnerContext(NoopWorkflowTracer()),
+        MagicMock(),
+        MagicMock(),
     )
 
     runner.register(FlakyTask)

tilebox-workflows/tilebox/workflows/client.py

@@ -1,5 +1,7 @@
 import logging
 import os
+import warnings
+from uuid import uuid4
 
 from _tilebox.grpc.channel import open_channel, parse_channel_info
 from tilebox.datasets.sync.client import Client as DatasetsClient
@@ -12,6 +14,12 @@
 )
 from tilebox.workflows.jobs.client import JobClient
 from tilebox.workflows.jobs.service import JobService
+from tilebox.workflows.observability.logging import (
+    OTELLoggingHandler,
+    StructuredLogger,
+    _create_tilebox_logger,
+    _create_tilebox_logger_provider,
+)
 from tilebox.workflows.observability.tracing import WorkflowTracer
 from tilebox.workflows.runner.task_runner import TaskRunner, _LeaseRenewer
 from tilebox.workflows.runner.task_service import TaskService
@@ -34,22 +42,27 @@ def __init__(
         token = _token_from_env(url, token)
         self._auth: dict[str, str] = {"token": token, "url": url}
         self._channel = open_channel(url, token)
-        self._tracer = WorkflowTracer(service=name, url=url, token=token)
 
-        self._logger: logging.Logger | None = None
+        # configure logging and tracing
+        self._client_id = uuid4()  # a random uuid to scope loggers to this client instance
+        self._logger_provider = _create_tilebox_logger_provider(service=name, url=url, token=token)
 
-    def configure_tracing(self, tracer: WorkflowTracer) -> None:
-        """
-        Configure the tracer to use for tracing of tasks and jobs within the workflow clients.
+        # task logger is the logger available for users to emit logs from within a Task.execute method, via
+        # context.logger
+        self._task_logger = _create_tilebox_logger(self._client_id, scope="tasks")
+        self._task_logger_handler = OTELLoggingHandler(level=logging.INFO, logger_provider=self._logger_provider)
+        self._task_logger.addHandler(self._task_logger_handler)
 
-        The tracer will be used by all task runners and job clients created by this client.
+        # runner logger is the logger used for logging internal events within a task runner, for example when a
+        # Tilebox API call fails, or when unexpected errors occur. This logger is not exposed to users,
+        # and is only used for logging internal events within the client and task runners.
+        self._runner_logger = _create_tilebox_logger(self._client_id, scope="runner")
+        self._runner_logger_handler = OTELLoggingHandler(level=logging.INFO, logger_provider=self._logger_provider)
+        self._runner_logger.addHandler(self._runner_logger_handler)
 
-        Calling this method multiple times will replace the existing tracing configuration. However, task runners and
-        job clients that have already been created will not be affected by subsequent calls to this method.
-        """
-        self._tracer = tracer
+        self._tracer = WorkflowTracer(service=name, url=url, token=token)
 
-    def configure_logging(self, logger: logging.Logger) -> None:
+    def configure_logging(self, level: int | logging.Logger, runner_level: int | None = None) -> None:
         """
         Configure the logger to use for logging of internal events within workflow clients.
 
@@ -61,7 +74,25 @@ def configure_logging(self, logger: logging.Logger) -> None:
         Args:
             logger: The logger to use for logging.
         """
-        self._logger = logger
+        if not isinstance(level, int):
+            warning_message = (
+                "Configuring a logger instance directly on a client is deprecated and will be removed in a future "
+                "version. If you want to export logs to an external system, configure the tilebox root logger "
+                "instance, which you can get with `tilebox.workflows.observability.logging.get_logger()`."
+            )
+            warnings.warn(
+                warning_message,
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            # to preserve backwards compatibility with the old API where the first argument was a logger
+            self._runner_logger = level
+        else:
+            # always adjust the level of the handler, not the loggers themselves, to make sure that other logger
+            # handlers still receive the logs (for example, if the user configured the tilebox root logger to export
+            # all logs at DEBUG level to a file)
+            self._task_logger_handler.setLevel(level)
+            self._runner_logger_handler.setLevel(runner_level or level)
 
     def jobs(self) -> JobClient:
         """Get a client for the jobs service.
@@ -114,9 +145,10 @@ def runner(
             found_cluster.slug,
             cache,
             self._tracer,
-            self._logger,
             _LeaseRenewer(**self._auth),
             runner_context,
+            task_logger=StructuredLogger(self._task_logger, {}),
+            runner_logger=StructuredLogger(self._runner_logger, {}),
         )
 
         if tasks is not None:

tilebox-workflows/tilebox/workflows/observability/logging.py

@@ -6,18 +6,19 @@
 import re
 import sys
 import traceback
-from datetime import timedelta
+from datetime import datetime, timedelta
 from functools import lru_cache
 from importlib.metadata import PackageNotFoundError, version
-from typing import ClassVar, TextIO
-from uuid import uuid4
+from typing import Any, ClassVar, TextIO
+from uuid import UUID, uuid4
 
 from opentelemetry.exporter.otlp.proto.http._log_exporter import (
     DEFAULT_LOGS_EXPORT_PATH,
     OTLPLogExporter,
     _append_logs_path,
 )
-from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
+from opentelemetry.instrumentation.logging.handler import LoggingHandler
+from opentelemetry.sdk._logs import LoggerProvider
 from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
 from opentelemetry.sdk.resources import (
     HOST_ARCH,
@@ -31,6 +32,7 @@
     Resource,
 )
 from opentelemetry.semconv.attributes import exception_attributes
+from opentelemetry.trace import get_current_span
 from opentelemetry.util.types import _ExtendedAttributes
 
 # prefix for stdlib loggers
@@ -44,6 +46,8 @@
 _OTEL_LOGS_ENDPOINT_ENV_VAR = "OTEL_LOGS_ENDPOINT"
 _OTEL_EXPORT_INTERVAL_ENV_VAR = "OTEL_EXPORT_INTERVAL"
 
+_WORKFLOW_LOG_ATTRIBUTES = "tilebox_structured_log_attributes"
+
 # process-unique identifier to distinguish different instances of the same service running on the same host
 _instance_id = str(uuid4())
 
@@ -87,10 +91,51 @@ def _root_logger() -> logging.Logger:
     return root_logger
 
 
+def _current_span_attributes() -> dict[str, str]:
+    span_context = get_current_span().get_span_context()
+    if not span_context.is_valid:
+        return {}
+
+    return {
+        "trace_id": f"{span_context.trace_id:032x}",
+        "span_id": f"{span_context.span_id:016x}",
+    }
+
+
+def _sanitize_otel_attribute_value(
+    value: Any,
+) -> str | bool | int | float | bytes | list[str | bool | int | float | bytes]:
+    if isinstance(value, str | bool | int | float | bytes):
+        return value
+
+    if isinstance(value, datetime):
+        return value.isoformat()
+
+    if isinstance(value, tuple | list):
+        values = []
+        for item in value:
+            if isinstance(item, str | bool | int | float | bytes):
+                values.append(item)
+            else:
+                values.append(_sanitize_otel_attribute_value(item))
+        return values
+
+    return str(value)
+
+
+def _sanitize_otel_attributes(attributes: dict[str, Any]) -> _ExtendedAttributes:
+    return {str(key): _sanitize_otel_attribute_value(value) for key, value in attributes.items()}
+
+
 class OTELLoggingHandler(LoggingHandler):
-    @staticmethod
-    def _get_attributes(record: logging.LogRecord) -> _ExtendedAttributes:
-        attributes = {}
+    def _get_attributes(self, record: logging.LogRecord) -> _ExtendedAttributes:
+        attributes: dict[str, Any] = {}
+        attributes.update(_current_span_attributes())
+
+        workflow_attributes = getattr(record, _WORKFLOW_LOG_ATTRIBUTES, None)
+        if isinstance(workflow_attributes, dict):
+            attributes.update(workflow_attributes)
+
         # the default implementation returns attributes for the filepath, lineno and function of the log record
         # we don't want that by default, so we override it to return an empty dict
         if record.exc_info:
@@ -104,7 +149,87 @@ def _get_attributes(record: logging.LogRecord) -> _ExtendedAttributes:
                 attributes[exception_attributes.EXCEPTION_STACKTRACE] = "".join(
                     traceback.format_exception(*record.exc_info)
                 )
-        return attributes
+        return _sanitize_otel_attributes(attributes)
+
+
+class StructuredLogger:
+    """A small structured logging wrapper for logs emitted during task execution."""
+
+    def __init__(self, logger: logging.Logger, attributes: dict[str, Any] | None = None) -> None:
+        self._logger = logger
+        self._attributes = attributes or {}
+
+    def bind(self, **attributes: Any) -> "StructuredLogger":
+        """Return a new logger that includes the given attributes in every log record."""
+        return StructuredLogger(self._logger, self._attributes | attributes)
+
+    def log(self, level: int, message: object, /, *args: Any, **attributes: Any) -> None:
+        """Log a message with structured attributes."""
+        self._log(level, message, args, attributes, exc_info=False)
+
+    def debug(self, message: object, /, *args: Any, **attributes: Any) -> None:
+        """Log a debug message with structured attributes."""
+        self._log(logging.DEBUG, message, args, attributes, exc_info=False)
+
+    def info(self, message: object, /, *args: Any, **attributes: Any) -> None:
+        """Log an info message with structured attributes."""
+        self._log(logging.INFO, message, args, attributes, exc_info=False)
+
+    def warning(self, message: object, /, *args: Any, **attributes: Any) -> None:
+        """Log a warning message with structured attributes."""
+        self._log(logging.WARNING, message, args, attributes, exc_info=False)
+
+    def error(self, message: object, /, *args: Any, **attributes: Any) -> None:
+        """Log an error message with structured attributes."""
+        self._log(logging.ERROR, message, args, attributes, exc_info=False)
+
+    def exception(self, message: object, /, *args: Any, **attributes: Any) -> None:
+        """Log an error message with structured attributes and the current exception information."""
+        self._log(logging.ERROR, message, args, attributes, exc_info=True)
+
+    def critical(self, message: object, /, *args: Any, **attributes: Any) -> None:
+        """Log a critical message with structured attributes."""
+        self._log(logging.CRITICAL, message, args, attributes, exc_info=False)
+
+    def _log(
+        self,
+        level: int,
+        message: object,
+        args: tuple[Any, ...],
+        attributes: dict[str, Any],
+        *,
+        exc_info: bool,
+    ) -> None:
+        if not self._logger.isEnabledFor(level):
+            return
+
+        workflow_attributes = self._attributes | attributes | _current_span_attributes()
+        self._logger.log(
+            level,
+            message,
+            *args,
+            exc_info=exc_info,
+            extra={_WORKFLOW_LOG_ATTRIBUTES: workflow_attributes},
+            stacklevel=3,
+        )
+
+
+@lru_cache(maxsize=16)  # reuse logger providers for the same credentials if possible
+def _create_tilebox_logger_provider(service: str | None, url: str, token: str | None) -> LoggerProvider:
+    provider = LoggerProvider(resource=_get_default_resource(service))
+    batch_exporter = _otel_log_exporter(
+        endpoint=url,
+        headers={"Authorization": f"Bearer {token}"} if token is not None else None,
+    )
+    provider.add_log_record_processor(batch_exporter)
+    return provider
+
+
+def _create_tilebox_logger(client_id: UUID, scope: str) -> logging.Logger:
+    logger = logging.getLogger(f"{_LOGGING_NAMESPACE}.clients.{client_id}.{scope}")
+    logger.setLevel(logging.DEBUG)  # always debug, so that other handlers can still filter by that level
+    logger.propagate = True
+    return logger
 
 
 def _otel_log_exporter(