feat: adopt structlog for structured logging

kubeflow/common/structured_logging.py

@@ -0,0 +1,130 @@
+"""Structured logging utilities for the Kubeflow SDK."""
+
+from __future__ import annotations
+
+import logging
+import os
+import threading
+import warnings
+
+import structlog
+
+_CONFIGURED = False
+_CONFIGURE_LOCK = threading.Lock()
+
+
+def configure_logging(level: str = "INFO", *, json_output: bool | None = None) -> None:
+    """Configure structlog processors and stdlib logging for the kubeflow namespace.
+
+    This is opt-in. Libraries should not call this at import time. Application
+    entrypoints (scripts, CLI tools) should call it once at startup.
+
+    Safe to call multiple times; only the first call takes effect. Subsequent
+    calls emit a warning and are ignored.
+
+    Args:
+        level: Logging level name (DEBUG, INFO, WARNING, ERROR).
+        json_output: If True, emit JSON lines; if False, human-readable console
+            output. When None, JSON is used when the ``CI`` env var is set
+            to a non-empty value other than "0" or "false".
+
+    Raises:
+        ValueError: If level is not a valid Python logging level name.
+    """
+    global _CONFIGURED
+    with _CONFIGURE_LOCK:
+        if _CONFIGURED:
+            warnings.warn(
+                "configure_logging() has already been called; ignoring subsequent call.",
+                stacklevel=2,
+            )
+            return
+        numeric_level = getattr(logging, level.upper(), None)
+        if not isinstance(numeric_level, int):
+            raise ValueError(f"Invalid log level: {level!r}")
+
+        if json_output is None:
+            ci_val = os.environ.get("CI", "").lower()
+            json_output = ci_val not in ("", "0", "false")
+
+        shared_processors: list[structlog.types.Processor] = [
+            structlog.contextvars.merge_contextvars,
+            structlog.stdlib.add_log_level,
+            structlog.stdlib.add_logger_name,
+            structlog.processors.TimeStamper(fmt="iso"),
+            structlog.processors.StackInfoRenderer(),
+            structlog.processors.format_exc_info,
+            structlog.processors.UnicodeDecoder(),
+        ]
+
+        renderer: structlog.types.Processor = (
+            structlog.processors.JSONRenderer() if json_output else structlog.dev.ConsoleRenderer()
+        )
+
+        # Configure structlog globally. This is intentional here because
+        # configure_logging() is an opt-in call by the application, not
+        # triggered at library import time.
+        structlog.configure(
+            processors=[
+                *shared_processors,
+                structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
+            ],
+            logger_factory=structlog.stdlib.LoggerFactory(),
+            wrapper_class=structlog.stdlib.BoundLogger,
+            cache_logger_on_first_use=True,
+        )
+        formatter = structlog.stdlib.ProcessorFormatter(
+            processors=[
+                structlog.stdlib.ProcessorFormatter.remove_processors_meta,
+                renderer,
+            ],
+            foreign_pre_chain=shared_processors,
+        )
+
+        handler = logging.StreamHandler()
+        handler.setFormatter(formatter)
+
+        ns_logger = logging.getLogger("kubeflow")
+        if ns_logger.handlers:
+            warnings.warn(
+                "The 'kubeflow' logger already has handlers; "
+                "configure_logging() will append a structured handler.",
+                stacklevel=2,
+            )
+        ns_logger.addHandler(handler)
+        ns_logger.setLevel(numeric_level)
+        ns_logger.propagate = False
+
+        _CONFIGURED = True
+
+
+def get_logger(name: str | None = None) -> structlog.stdlib.BoundLogger:
+    """Return a structlog logger wrapping a stdlib logger.
+
+    Without a prior :func:`configure_logging` call, the returned logger routes
+    through stdlib's logging infrastructure with no additional processors. This
+    means it inherits whatever handler configuration the application has set up
+    (safe library default).
+
+    After :func:`configure_logging` is called, loggers use the configured
+    structlog processors and renderer.
+
+    Args:
+        name: Logger name, typically ``__name__``.
+
+    Returns:
+        A structlog bound logger compatible with stdlib logging levels.
+    """
+    if _CONFIGURED:
+        # After configure_logging(), structlog.get_logger() returns loggers
+        # that use the configured processors and renderer.
+        return structlog.get_logger(name)
+    # Before configure_logging(), wrap a stdlib logger directly. We cannot use
+    # structlog.get_logger() here because cache_logger_on_first_use=True (set
+    # by configure_logging) would freeze these early loggers with the
+    # pre-configuration factory, making them ignore a later configure_logging()
+    # call. wrap_logger() bypasses that cache entirely.
+    return structlog.wrap_logger(
+        logging.getLogger(name),
+        wrapper_class=structlog.stdlib.BoundLogger,
+    )

kubeflow/common/structured_logging_test.py

@@ -0,0 +1,15 @@
+
+from kubeflow.common import structured_logging
+
+
+def test_configure_logging_and_get_logger(monkeypatch):
+    """Test that configure_logging succeeds and get_logger returns a structlog logger."""
+    monkeypatch.setattr(structured_logging, "_CONFIGURED", False)
+
+    structured_logging.configure_logging()
+
+    logger = structured_logging.get_logger(__name__)
+
+    assert logger is not None
+
+    logger.info("test message")

kubeflow/optimizer/api/optimizer_client.py

@@ -13,9 +13,9 @@
 # limitations under the License.
 
 from collections.abc import Callable, Iterator
-import logging
 from typing import Any
 
+from kubeflow.common.structured_logging import get_logger
 from kubeflow.common.types import KubernetesBackendConfig
 import kubeflow.common.utils as common_utils
 from kubeflow.optimizer.backends.kubernetes.backend import KubernetesBackend
@@ -29,7 +29,7 @@
 )
 from kubeflow.trainer.types.types import Event, TrainJobTemplate
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class OptimizerClient:

kubeflow/optimizer/backends/kubernetes/backend.py

@@ -14,7 +14,6 @@
 
 from collections.abc import Callable, Iterator
 import copy
-import logging
 import multiprocessing
 import random
 import string
@@ -26,6 +25,7 @@
 from kubernetes import client, config
 
 import kubeflow.common.constants as common_constants
+from kubeflow.common.structured_logging import get_logger
 from kubeflow.common.types import KubernetesBackendConfig
 import kubeflow.common.utils as common_utils
 from kubeflow.optimizer.backends.base import RuntimeBackend
@@ -44,7 +44,7 @@
 import kubeflow.trainer.constants.constants as trainer_constants
 from kubeflow.trainer.types.types import Event, TrainJobTemplate
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class KubernetesBackend(RuntimeBackend):

kubeflow/spark/backends/kubernetes/backend.py

@@ -36,6 +36,7 @@
 from pyspark.sql import SparkSession
 
 from kubeflow.common import constants as common_constants
+from kubeflow.common.structured_logging import get_logger
 from kubeflow.common.types import KubernetesBackendConfig
 from kubeflow.spark.backends.base import RuntimeBackend
 from kubeflow.spark.backends.kubernetes import constants
@@ -62,7 +63,7 @@
     SparkJobStatus,
 )
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 _spark_debug_logging_enabled = False
 

kubeflow/trainer/api/trainer_client.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 from collections.abc import Callable, Iterator
-import logging
 
+from kubeflow.common.structured_logging import get_logger
 from kubeflow.common.types import KubernetesBackendConfig
 import kubeflow.common.utils as common_utils
 from kubeflow.trainer.backends.container.backend import ContainerBackend
@@ -27,7 +27,7 @@
 from kubeflow.trainer.constants import constants
 from kubeflow.trainer.types import types
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class TrainerClient:
@@ -235,7 +235,7 @@ def get_job_events(self, name: str) -> list[types.Event]:
     def wait_for_job_status(
         self,
         name: str,
-        status: set[str] = {constants.TRAINJOB_COMPLETE},
+        status: set[str] | None = None,
         timeout: int = 600,
         polling_interval: int = 2,
         callbacks: list[Callable[[types.TrainJob], None]] | None = None,

kubeflow/trainer/backends/base.py

@@ -75,7 +75,7 @@ def get_job_events(self, name: str) -> list[types.Event]:
     def wait_for_job_status(
         self,
         name: str,
-        status: set[str] = {constants.TRAINJOB_COMPLETE},
+        status: set[str] | None = None,
         timeout: int = 600,
         polling_interval: int = 2,
         callbacks: list[Callable[[types.TrainJob], None]] | None = None,

kubeflow/trainer/backends/container/backend.py

@@ -40,13 +40,13 @@
 from collections.abc import Callable, Iterator
 import concurrent.futures
 from datetime import datetime
-import logging
 import os
 import random
 import shutil
 import string
 import uuid
 
+from kubeflow.common.structured_logging import get_logger
 from kubeflow.trainer.backends.base import RuntimeBackend
 from kubeflow.trainer.backends.container import utils as container_utils
 from kubeflow.trainer.backends.container.adapters.base import (
@@ -62,7 +62,7 @@
 from kubeflow.trainer.constants import constants
 from kubeflow.trainer.types import types
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class ContainerBackend(RuntimeBackend):
@@ -842,13 +842,16 @@ def _build_failure_message(self, name: str) -> str:
     def wait_for_job_status(
         self,
         name: str,
-        status: set[str] = {constants.TRAINJOB_COMPLETE},
+        status: set[str] | None = None,
         timeout: int = 600,
         polling_interval: int = 2,
         callbacks: list[Callable[[types.TrainJob], None]] | None = None,
     ) -> types.TrainJob:
         import time
 
+        if status is None:
+            status = {constants.TRAINJOB_COMPLETE}
+
         end = time.time() + timeout
         while time.time() < end:
             tj = self.get_job(name)

kubeflow/trainer/backends/container/runtime_loader.py

@@ -28,18 +28,18 @@
 
 from datetime import datetime, timedelta
 import json
-import logging
 from pathlib import Path
 from typing import Any
 import urllib.error
 import urllib.request
 
 import yaml
 
+from kubeflow.common.structured_logging import get_logger
 from kubeflow.trainer.constants import constants
 from kubeflow.trainer.types import types as base_types
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 TRAINING_RUNTIMES_DIR = Path(__file__).parents[2] / "config" / "training_runtimes"
 CACHE_DIR = Path.home() / ".kubeflow" / "trainer" / "cache"

kubeflow/trainer/backends/container/utils.py

@@ -17,16 +17,16 @@
 """
 
 from dataclasses import dataclass
-import logging
 import os
 from pathlib import Path
 import shlex
 
 from kubeflow.common.constants import UNKNOWN
+from kubeflow.common.structured_logging import get_logger
 from kubeflow.trainer.constants import constants
 from kubeflow.trainer.types import types
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 def create_workdir(job_name: str) -> str:

kubeflow/trainer/backends/kubernetes/backend.py

@@ -14,7 +14,6 @@
 
 from collections.abc import Callable, Iterator
 import copy
-import logging
 import multiprocessing
 import os
 import random
@@ -28,14 +27,15 @@
 from kubernetes import client, config, watch
 
 import kubeflow.common.constants as common_constants
+from kubeflow.common.structured_logging import get_logger
 from kubeflow.common.types import KubernetesBackendConfig
 import kubeflow.common.utils as common_utils
 from kubeflow.trainer.backends.base import RuntimeBackend
 import kubeflow.trainer.backends.kubernetes.utils as utils
 from kubeflow.trainer.constants import constants
 from kubeflow.trainer.types import types
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class KubernetesBackend(RuntimeBackend):

kubeflow/trainer/backends/localprocess/backend.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 from collections.abc import Callable, Iterator
 from datetime import datetime
-import logging
 import random
 import string
 import tempfile
 import time
 import uuid
 
+from kubeflow.common.structured_logging import get_logger
 from kubeflow.trainer.backends.base import RuntimeBackend
 from kubeflow.trainer.backends.localprocess import utils as local_utils
 from kubeflow.trainer.backends.localprocess.constants import local_runtimes
@@ -32,7 +32,7 @@
 from kubeflow.trainer.constants import constants
 from kubeflow.trainer.types import types
 
-logger = logging.getLogger(__name__)
+logger = get_logger(__name__)
 
 
 class LocalProcessBackend(RuntimeBackend):

kubeflow/trainer/backends/localprocess/job.py

@@ -12,15 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from datetime import datetime
-import logging
 import os
 import subprocess
 import threading
 
+from kubeflow.common.structured_logging import get_logger
 from kubeflow.trainer.constants import constants
 
-logger = logging.getLogger(__name__)
-
+logger = get_logger(__name__)
 
 class LocalJob(threading.Thread):
     def __init__(