sdks/python/apache_beam/runners/worker/logger.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Python worker logging."""

# pytype: skip-file
# mypy: disallow-untyped-defs

# ruff: noqa: UP006
import contextlib
import json
import logging
import threading
import traceback
from typing import Any
from typing import Dict
from typing import Iterator
from typing import List

from apache_beam.runners.worker import statesampler


# Per-thread worker information. This is used only for logging to set
# context information that changes while work items get executed:
# work_item_id, step_name, stage_name.
class _PerThreadWorkerData(threading.local):
  def __init__(self) -> None:
    super().__init__()
    # in the list, as going up and down all the way to zero incurs several
    # reallocations.
    self.stack: List[Dict[str, Any]] = []

  def get_data(self) -> Dict[str, Any]:
    all_data = {}
    for datum in self.stack:
      all_data.update(datum)
    return all_data


per_thread_worker_data = _PerThreadWorkerData()


@contextlib.contextmanager
def PerThreadLoggingContext(**kwargs: Any) -> Iterator[None]:
  """A context manager to add per thread attributes."""
  stack = per_thread_worker_data.stack
  stack.append(kwargs)
  try:
    yield
  finally:
    stack.pop()


class JsonLogFormatter(logging.Formatter):
  """A JSON formatter class as expected by the logging standard module."""
  def __init__(self, job_id: str, worker_id: str) -> None:
    super().__init__()
    self.job_id = job_id
    self.worker_id = worker_id

  def format(self, record: logging.LogRecord) -> str:
    """Returns a JSON string based on a LogRecord instance.

    Args:
      record: A LogRecord instance. See below for details.

    Returns:
      A JSON string representing the record.

    A LogRecord instance has the following attributes and is used for
    formatting the final message.

    Attributes:
      created: A double representing the timestamp for record creation
        (e.g., 1438365207.624597). Note that the number contains also msecs and
        microsecs information. Part of this is also available in the 'msecs'
        attribute.
      msecs: A double representing the msecs part of the record creation
        (e.g., 624.5970726013184).
      msg: Logging message containing formatting instructions or an arbitrary
        object. This is the first argument of a log call.
      args: A tuple containing the positional arguments for the logging call.
      levelname: A string. Possible values are: INFO, WARNING, ERROR, etc.
      exc_info: None or a 3-tuple with exception information as it is
        returned by a call to sys.exc_info().
      name: Logger's name. Most logging is done using the default root logger
        and therefore the name will be 'root'.
      filename: Basename of the file where logging occurred.
      funcName: Name of the function where logging occurred.
      process: The PID of the process running the worker.
      thread:  An id for the thread where the record was logged. This is not a
        real TID (the one provided by OS) but rather the id (address) of a
        Python thread object. Nevertheless having this value can allow to
        filter log statement from only one specific thread.
    """
    output: Dict[str, Any] = {}
    output['timestamp'] = {
        'seconds': int(record.created), 'nanos': int(record.msecs * 1000000)
    }
    # ERROR. INFO, DEBUG log levels translate into the same for severity
    # property. WARNING becomes WARN.
    output['severity'] = (
        record.levelname if record.levelname != 'WARNING' else 'WARN')

    # msg could be an arbitrary object, convert it to a string first.
    record_msg = str(record.msg)

    # Prepare the actual message using the message formatting string and the
    # positional arguments as they have been used in the log call.
    if record.args:
      try:
        output['message'] = record_msg % record.args
      except (TypeError, ValueError):
        output['message'] = '%s with args (%s)' % (record_msg, record.args)
    else:
      output['message'] = record_msg

    # The thread ID is logged as a combination of the process ID and thread ID
    # since workers can run in multiple processes.
    output['thread'] = '%s:%s' % (record.process, record.thread)
    # job ID and worker ID. These do not change during the lifetime of a worker.
    output['job'] = self.job_id
    output['worker'] = self.worker_id
    # Stage, step and work item ID come from thread local storage since they
    # change with every new work item leased for execution. If there is no
    # work item ID then we make sure the step is undefined too.
    data = per_thread_worker_data.get_data()
    if 'work_item_id' in data:
      output['work'] = data['work_item_id']

    tracker = statesampler.get_current_tracker()
    if tracker:
      output['stage'] = tracker.stage_name

      if tracker.current_state() and tracker.current_state().name_context:
        output['step'] = tracker.current_state().name_context.logging_name()

    # All logging happens using the root logger. We will add the basename of the
    # file and the function name where the logging happened to make it easier
    # to identify who generated the record.
    output['logger'] = '%s:%s:%s' % (
        record.name, record.filename, record.funcName)
    # Add exception information if any is available.
    if record.exc_info:
      output['exception'] = ''.join(
          traceback.format_exception(*record.exc_info))

    return json.dumps(output)


def initialize(
    job_id: str,
    worker_id: str,
    log_path: str,
    log_level: int = logging.INFO) -> None:
  """Initialize root logger so that we log JSON to a file and text to stdout."""

  file_handler = logging.FileHandler(log_path)
  file_handler.setFormatter(JsonLogFormatter(job_id, worker_id))
  logging.getLogger().addHandler(file_handler)

  # Default level is set to INFO to avoid logging various DEBUG level log calls
  # sprinkled throughout the code.
  logging.getLogger().setLevel(log_level)