feat(logging): stamp run/action context via LogRecordFactory (#1038)

## Summary

Replaces the per-logger `logging.Filter` approach for context stamping
with a global `LogRecordFactory`. The factory runs on every `LogRecord`
regardless of which logger created it, so any user-defined logger
automatically picks up `run_name`, `action_name`, and
`is_flyte_internal` attributes — no per-handler filter wiring required.

This also fixes a latent bug in main: the old `ContextFilter` and
`FlyteInternalFilter` mutated `record.msg` in place, so a record passing
through multiple handlers could get its `[run][action]` and `[flyte]`
prefixes prepended multiple times. Stamping attributes on the record
(and rendering them at format time) sidesteps that entirely.

## What changed

- **LogRecordFactory** (`src/flyte/_logging.py`) — stamps the active
flyte action context on every record, lazily importing `flyte._context`
to avoid import-time cycles.
- **`ContextFormatter`** replaces `ContextFilter` +
`FlyteInternalFilter`. The `[flyte]` internal-prefix marker is now a
formatter flag (`internal_prefix=True`), so the same record can be
formatted by multiple handlers without prefixes compounding.
- **`ContextFormatter(inner=...)`** — optional kwarg to delegate base
formatting to an existing `logging.Formatter`. Used to wrap pre-existing
root-handler formatters without losing their layout.
- **`initialize_logger(reset_root_logger=False)`** now wraps each
existing root-handler formatter with `ContextFormatter(inner=existing)`,
preserving main's "third-party log lines through root get
`[run][action]`" behavior. Idempotent on repeated init calls. Same
temporal limitation as main: handlers added after `flyte.init()` won't
be wrapped, but the factory still stamps the attrs so callers can format
them themselves.
- **`get_rich_handler(internal_prefix=True)`** — opt-out parameter; the
user-logger call site passes `False` so user-emitted lines don't carry
the `[flyte]` marker even under rich logging.

## Behavior parity vs main

| | main | this PR |
|---|---|---|
| `[run][action]` on flyte/flyte.user logs | ✅ | ✅ |
| `[run][action]` on third-party logs through root (when handler exists
at init) | ✅ (via filter) | ✅ (via formatter wrap) |
| `[run][action]` on third-party logs through root (handler added after
init) | ❌ | ❌ — but attrs still on record |
| `[flyte]` marker on flyte internal logs | ✅ | ✅ |
| Compounding prefixes when record passes through multiple handlers | ❌
(latent bug) | ✅ fixed |
| `record.run_name` / `action_name` available everywhere | partial | ✅
universal (factory) |

## Why this matters for users

`flyte.logger` and any child of `flyte.user` continue to work as before.
New: anyone using a stdlib logger (`logging.getLogger("myapp")`) outside
the flyte namespace also gets `run_name` / `action_name` /
`is_flyte_internal` stamped on every record by the factory — they just
need a formatter that references them (e.g.
`"[%(run_name)s][%(action_name)s] %(message)s"`).

## Example

`examples/basics/reuse_concurrent_logging.py` exercises three logger
configs with a reusable container running concurrent tasks, so log
output from multiple actions interleaves on a single stderr — each line
correctly prefixed by its own `[run][action]`:

1. `flyte.logger` — canonical user-facing logger.
2. `logging.getLogger("flyte.user.myapp")` — child of `flyte.user`,
inherits the formatter via propagation.
3. `logging.getLogger("myapp")` — fully independent logger with its own
handler that references the factory-stamped attrs in its format string.

## Test plan

- [x] `pytest tests/user_api/test_logging.py
tests/flyte/test_logging.py` passes (23 tests)
- [x] `test_user_logger_no_flyte_prefix` updated to check the formatter
flag instead of the removed `FlyteInternalFilter`
- [x] `test_user_logger_no_flyte_prefix_after_rich_init` — regression
test for rich-handler internal prefix bleeding into user logger
- [x] `test_initialize_logger_wraps_existing_root_handlers` — covers
wrap path + idempotency on repeated init
- [ ] Manual: run `examples/basics/reuse_concurrent_logging.py` against
a cluster and confirm all three worker variants emit
`[run][action]`-prefixed log lines on the shared reused container

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Signed-off-by: Haytham Abuelfutuh <haytham@afutuh.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Samhita Alla <aallasamhita@gmail.com>

Author: 2 people

examples/basics/reuse_concurrent_logging.py

@@ -0,0 +1,110 @@
+"""
+Demo: reusable container + concurrent tasks, so log lines from different
+actions are interleaved in the same worker process. Each line should be
+prefixed with its own [run][action] thanks to the LogRecord factory in
+flyte._logging.
+
+Workers exercise three logger configs to confirm context stamping works
+regardless of how the user obtains a logger:
+
+  1. flyte.logger                           — the canonical user-facing logger
+  2. logging.getLogger("flyte.user.myapp")  — child of flyte.user, inherits handler
+  3. logging.getLogger("myapp")             — fully independent logger w/ its own handler
+"""
+
+import asyncio
+import logging
+import sys
+
+import flyte
+
+# One worker process, many concurrent invocations on it. With replicas=1 and
+# concurrency=8, all worker tasks below land in the same Python process, so
+# their log output is genuinely interleaved on a single stderr.
+env = flyte.TaskEnvironment(
+    name="reuse_concurrent_logging",
+    resources=flyte.Resources(cpu="1", memory="500Mi"),
+    reusable=flyte.ReusePolicy(
+        replicas=1,
+        concurrency=8,
+        idle_ttl=60,
+        scaledown_ttl=60,
+    ),
+    image=flyte.Image.from_debian_base().with_pip_packages("unionai-reuse"),
+)
+
+# Variant 2: child of flyte.user. No handler/level setup needed — propagation
+# carries records up to flyte.user's handler, so [run][action] prefixing and
+# the user log level both apply automatically.
+inherited_logger = logging.getLogger("flyte.user.myapp")
+
+# Variant 3: independent logger with its own StreamHandler. The flyte
+# LogRecordFactory still stamps run_name/action_name on every record, so the
+# formatter below can reference them via %(run_name)s / %(action_name)s even
+# though this logger lives outside the flyte.* namespace.
+independent_logger = logging.getLogger("myapp")
+if not independent_logger.handlers:
+    _h = logging.StreamHandler(sys.stderr)
+    _h.setFormatter(logging.Formatter("[%(run_name)s][%(action_name)s] myapp: %(message)s"))
+    independent_logger.addHandler(_h)
+    independent_logger.setLevel(logging.INFO)
+    independent_logger.propagate = False
+
+
+@env.task
+async def worker_flyte_logger(label: str, ticks: int = 3) -> str:
+    """Variant 1: the canonical flyte.logger."""
+    flyte.logger.info("starting label=%s", label)
+    for i in range(ticks):
+        await asyncio.sleep(1)
+        flyte.logger.info("label=%s tick=%d/%d", label, i + 1, ticks)
+    flyte.logger.info("done label=%s", label)
+    return label
+
+
+@env.task
+async def worker_inherited_logger(label: str, ticks: int = 3) -> str:
+    """Variant 2: child of flyte.user — formatting is inherited."""
+    inherited_logger.info("starting label=%s", label)
+    for i in range(ticks):
+        await asyncio.sleep(1)
+        inherited_logger.info("label=%s tick=%d/%d", label, i + 1, ticks)
+    inherited_logger.info("done label=%s", label)
+    return label
+
+
+@env.task
+async def worker_independent_logger(label: str, ticks: int = 3) -> str:
+    """Variant 3: a fully independent stdlib logger — record factory still stamps context."""
+    independent_logger.info("starting label=%s", label)
+    for i in range(ticks):
+        await asyncio.sleep(1)
+        independent_logger.info("label=%s tick=%d/%d", label, i + 1, ticks)
+    independent_logger.info("done label=%s", label)
+    return label
+
+
+# The fan-out parent itself doesn't need to share the reuse pool — clone_with
+# turns reuse off for it and depends on the worker env for image + registration.
+@env.clone_with(name="reuse_concurrent_main", reusable=None, depends_on=[env]).task
+async def main(n: int = 2) -> list[str]:
+    """
+    Dispatches a mix of all three worker variants so logs from each logger
+    config interleave on the reused container's stderr.
+    """
+    flyte.logger.info("dispatching %d of each worker variant", n)
+    coros: list = []
+    for i in range(n):
+        coros.append(worker_flyte_logger(label=f"flyte-{i}"))
+        coros.append(worker_inherited_logger(label=f"inherited-{i}"))
+        coros.append(worker_independent_logger(label=f"independent-{i}"))
+    results = await asyncio.gather(*coros)
+    flyte.logger.info("all workers finished: %s", results)
+    return results
+
+
+if __name__ == "__main__":
+    flyte.init_from_config()
+    run = flyte.run(main, n=2)
+    print(run.url)
+    run.wait()

src/flyte/_logging.py

@@ -3,13 +3,44 @@
 import logging
 import os
 from datetime import datetime
-from typing import Literal, Optional
+from typing import Any, Literal, Optional
 
 import flyte
 
 from ._tools import ipython_check
 
 LogFormat = Literal["console", "json"]
+
+_orig_record_factory = logging.getLogRecordFactory()
+
+
+def _flyte_record_factory(*args: Any, **kwargs: Any) -> logging.LogRecord:
+    record = _orig_record_factory(*args, **kwargs)
+
+    # Stamp the active flyte action context, if any. Imported lazily because
+    # this factory runs on every record, including during flyte's own import.
+    try:
+        from flyte._context import ctx as _flyte_ctx
+
+        c = _flyte_ctx()
+    except Exception:
+        c = None
+    if c is not None:
+        record.run_name = c.action.run_name
+        record.action_name = c.action.name
+    else:
+        record.run_name = None
+        record.action_name = None
+
+    record.is_flyte_internal = record.name == "flyte" or (
+        record.name.startswith("flyte.") and not record.name.startswith("flyte.user")
+    )
+    return record
+
+
+logging.setLogRecordFactory(_flyte_record_factory)
+
+
 _LOG_LEVEL_MAP = {
     "critical": logging.CRITICAL,  # 50
     "error": logging.ERROR,  # 40
@@ -91,7 +122,7 @@ def _get_console():
     return Console(width=width)
 
 
-def get_rich_handler(log_level: int) -> Optional[logging.Handler]:
+def get_rich_handler(log_level: int, internal_prefix: bool = True) -> Optional[logging.Handler]:
     """
     Upgrades the global loggers to use Rich logging.
     """
@@ -117,7 +148,7 @@ def get_rich_handler(log_level: int) -> Optional[logging.Handler]:
         markup=True,
     )
 
-    formatter = logging.Formatter(fmt="%(filename)s:%(lineno)d - %(message)s")
+    formatter = ContextFormatter(fmt="%(filename)s:%(lineno)d - %(message)s", internal_prefix=internal_prefix)
     handler.setFormatter(formatter)
     return handler
 
@@ -190,9 +221,16 @@ def initialize_logger(
     if reset_root_logger:
         _setup_root_logger(use_json=use_json, use_rich=use_rich, log_level=log_level)
     else:
+        # Wrap each existing root-handler formatter so third-party log lines
+        # routed through root render with [run][action]. Captures handlers
+        # registered before this call; handlers added later won't be wrapped
+        # (factory still stamps the attrs, so callers can format them).
         root_logger = logging.getLogger()
         for h in root_logger.handlers:
-            h.addFilter(ContextFilter())
+            existing = h.formatter
+            if isinstance(existing, ContextFormatter):
+                continue
+            h.setFormatter(ContextFormatter(inner=existing))
 
     # Set up Flyte logger handler
     flyte_handler: logging.Handler | None = None
@@ -206,12 +244,7 @@ def initialize_logger(
     if flyte_handler is None:
         flyte_handler = logging.StreamHandler()
         flyte_handler.setLevel(log_level)
-        formatter = logging.Formatter(fmt="%(message)s")
-        flyte_handler.setFormatter(formatter)
-
-    # Add both filters to Flyte handler
-    flyte_handler.addFilter(FlyteInternalFilter())
-    flyte_handler.addFilter(ContextFilter())
+        flyte_handler.setFormatter(ContextFormatter(fmt="%(message)s", internal_prefix=True))
 
     flyte_logger.addHandler(flyte_handler)
     flyte_logger.setLevel(log_level)
@@ -231,17 +264,16 @@ def initialize_logger(
         user_handler.setLevel(user_log_level)
         user_handler.setFormatter(JSONFormatter())
     elif use_rich:
-        rich_handler = get_rich_handler(user_log_level)
+        rich_handler = get_rich_handler(user_log_level, internal_prefix=False)
         user_handler = rich_handler if rich_handler is not None else logging.StreamHandler()
         user_handler.setLevel(user_log_level)
         if not rich_handler:
-            user_handler.setFormatter(logging.Formatter(fmt="%(message)s"))
+            user_handler.setFormatter(ContextFormatter(fmt="%(message)s"))
     else:
         user_handler = logging.StreamHandler()
         user_handler.setLevel(user_log_level)
-        user_handler.setFormatter(logging.Formatter(fmt="%(message)s"))
+        user_handler.setFormatter(ContextFormatter(fmt="%(message)s"))
 
-    user_handler.addFilter(ContextFilter())
     user_flyte_logger.addHandler(user_handler)
     user_flyte_logger.setLevel(user_log_level)
     user_flyte_logger.propagate = False
@@ -274,42 +306,38 @@ def wrapper(*args, **kwargs):
     return decorator(fn)
 
 
-class ContextFilter(logging.Filter):
+class ContextFormatter(logging.Formatter):
     """
-    A logging filter that adds the current action's run name and name to all log records.
-    Applied globally to capture context for both user and Flyte internal logging.
+    Console formatter that prefixes records with action context and an optional
+    [flyte] marker, both pulled from attributes stamped by _flyte_record_factory.
+    Does not mutate record state, so the same record can be formatted by
+    multiple handlers without compounding prefixes.
     """
 
-    def filter(self, record: logging.LogRecord) -> bool:
-        from flyte._context import ctx
-
-        c = ctx()
-        if c:
-            action = c.action
-            # Add as attributes for structured logging (JSON)
-            record.run_name = action.run_name
-            record.action_name = action.name
-            # Also modify message for console/Rich output
-            record.msg = f"[{action.run_name}][{action.name}] {record.msg}"
-        else:
-            record.run_name = None
-            record.action_name = None
-        return True
+    def __init__(
+        self,
+        fmt: str = "%(message)s",
+        *,
+        internal_prefix: bool = False,
+        inner: logging.Formatter | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(fmt=fmt, **kwargs)
+        self._internal_prefix = internal_prefix
+        self._inner = inner
 
-
-class FlyteInternalFilter(logging.Filter):
-    """
-    A logging filter that adds [flyte] prefix to internal Flyte logging only.
-    """
-
-    def filter(self, record: logging.LogRecord) -> bool:
-        is_internal = record.name.startswith("flyte")
-        # Add as attribute for structured logging (JSON)
-        record.is_flyte_internal = is_internal
-        # Also modify message for console/Rich output
-        if is_internal:
-            record.msg = f"[flyte] {record.msg}"
-        return True
+    def format(self, record: logging.LogRecord) -> str:
+        base = self._inner.format(record) if self._inner is not None else super().format(record)
+        parts: list[str] = []
+        run_name = getattr(record, "run_name", None)
+        action_name = getattr(record, "action_name", None)
+        if run_name and action_name:
+            parts.append(f"[{run_name}][{action_name}]")
+        if self._internal_prefix and getattr(record, "is_flyte_internal", False):
+            parts.append("[flyte]")
+        if not parts:
+            return base
+        return f"{' '.join(parts)} {base}"
 
 
 def _setup_root_logger(use_json: bool, use_rich: bool, log_level: int):
@@ -330,28 +358,25 @@ def _setup_root_logger(use_json: bool, use_rich: bool, log_level: int):
     # get_rich_handler can return None in some environments
     if not root_handler:
         root_handler = logging.StreamHandler()
+        root_handler.setFormatter(ContextFormatter(fmt="%(message)s"))
 
-    # Add context filter to ALL logging
-    root_handler.addFilter(ContextFilter())
     root_handler.setLevel(log_level)
-
     root.addHandler(root_handler)
     root.setLevel(log_level)
 
 
 def _create_user_logger() -> logging.Logger:
     """
     Create the user-facing logger. Defaults to INFO so user logs are visible by default.
-    Does not use FlyteInternalFilter — no [flyte] prefix on user messages.
+    No [flyte] prefix on user messages.
     """
     user_flyte_logger = logging.getLogger("flyte.user")
     user_log_level = get_env_user_log_level()
     user_flyte_logger.setLevel(user_log_level)
 
     handler = logging.StreamHandler()
     handler.setLevel(user_log_level)
-    handler.addFilter(ContextFilter())
-    handler.setFormatter(logging.Formatter(fmt="%(message)s"))
+    handler.setFormatter(ContextFormatter(fmt="%(message)s"))
 
     user_flyte_logger.propagate = False
     user_flyte_logger.addHandler(handler)
@@ -366,14 +391,9 @@ def _create_flyte_logger() -> logging.Logger:
     flyte_logger = logging.getLogger("flyte")
     flyte_logger.setLevel(get_env_log_level())
 
-    # Add a handler specifically for flyte logging with the prefix filter
     handler = logging.StreamHandler()
     handler.setLevel(get_env_log_level())
-    handler.addFilter(FlyteInternalFilter())
-    handler.addFilter(ContextFilter())
-
-    formatter = logging.Formatter(fmt="%(message)s")
-    handler.setFormatter(formatter)
+    handler.setFormatter(ContextFormatter(fmt="%(message)s", internal_prefix=True))
 
     # Prevent propagation to root to avoid double logging
     flyte_logger.propagate = False