From fc7ff248ff9c5bdced34c4391aa024f5501b1398 Mon Sep 17 00:00:00 2001 From: Ivy Zhou Date: Fri, 21 Aug 2026 18:30:59 -0700 Subject: [PATCH] Bridge checkpoint metrics into structured logging Summary: `torch_checkpointing` emits `EventLogger` metadata on ordinary Python logging records. TorchTitan kept its structured handlers on a separate non-propagating logger, so those backend records reached the console but not the structured log. Install one private handler on the root logger during `init_structured_logger()`. It forwards records that carry `log_type_name` to TorchTitan's existing structured handlers and ignores plain text. This matches the llama4x logging model without a package-specific public forwarding API. Async writers start in a fresh process, so preserve the resolved parent source, output directory, and rank in the backend config. The child runs any existing initializer first, restores TorchTitan structured logging, and admits backend INFO events only when the backend logger is unset and inherits a stricter level. Explicit DEBUG, WARNING, and ERROR settings remain unchanged. Preserve the backend's `context` and `measured_from_start_time_ms` fields in JSONL output, and stamp `checkpoint_logging_context` with the save step before the backend snapshots it for asynchronous work. Test Plan: `PYTHONPATH=/tmp/torchtitan-additional-packages.iGrJ91/checkpointing /home/ivyzhou/local/torchtitan-venv/bin/python -m pytest -q tests/unit_tests/observability/ tests/unit_tests/test_torch_checkpointing.py tests/unit_tests/test_checkpoint.py` 162 passed. `PYTHONPATH=/tmp/torchtitan-additional-packages.iGrJ91/checkpointing:$PYTHONPATH PATH=/home/ivyzhou/local/torchtitan-venv/bin:$PATH /home/ivyzhou/local/torchtitan-venv/bin/pre-commit run --files tests/unit_tests/observability/test_structured_logging.py tests/unit_tests/test_torch_checkpointing.py torchtitan/components/checkpointer/torch_checkpointing.py torchtitan/observability/structured_logger/__init__.py torchtitan/observability/structured_logger/structured_logging.py` All hooks passed. --- .../observability/test_structured_logging.py | 134 +++++++++++++++++- tests/unit_tests/test_torch_checkpointing.py | 104 ++++++++++++++ .../checkpointer/torch_checkpointing.py | 59 ++++++++ .../structured_logger/jsonl_handler.py | 10 ++ .../structured_logger/structured_logging.py | 38 ++++- 5 files changed, 340 insertions(+), 5 deletions(-) diff --git a/tests/unit_tests/observability/test_structured_logging.py b/tests/unit_tests/observability/test_structured_logging.py index a49f286738..b384728f5b 100644 --- a/tests/unit_tests/observability/test_structured_logging.py +++ b/tests/unit_tests/observability/test_structured_logging.py @@ -33,6 +33,7 @@ set_step, ) from torchtitan.observability.structured_logger.structured_logging import ( + _get_structured_logger_init_args, _structured_logger, event_extra, ExtraFields, @@ -73,16 +74,43 @@ def structured_logger_fixture(): import torchtitan.observability.structured_logger.structured_logging as sl_mod tl = _structured_logger - orig = (tl.handlers[:], tl.level, tl.propagate) + root_logger = logging.getLogger() + orig = ( + tl.handlers[:], + tl.level, + tl.propagate, + sl_mod._disabled, + sl_mod._structured_logger_init_args, + root_logger.handlers[:], + ) # Reset the module-level init sentinel so init_structured_logger re-runs # for each test (otherwise the second call short-circuits as "already # initialized"). sl_mod._is_initialized = False + sl_mod._disabled = False + sl_mod._structured_logger_init_args = None yield tl - tl.handlers, tl.level, tl.propagate = orig + ( + tl.handlers, + tl.level, + tl.propagate, + sl_mod._disabled, + init_args, + root_logger.handlers, + ) = orig + sl_mod._structured_logger_init_args = init_args sl_mod._is_initialized = False +@pytest.fixture +def external_logger(): + source_logger = logging.getLogger("external_library") + orig = (source_logger.handlers[:], source_logger.level, source_logger.propagate) + source_logger.handlers = [] + yield source_logger + source_logger.handlers, source_logger.level, source_logger.propagate = orig + + # --------------------------------------------------------------------------- # Step context tests (hybrid ContextVar) # --------------------------------------------------------------------------- @@ -506,6 +534,45 @@ def test_has_time_us(self): assert "time_us" in parsed assert isinstance(parsed["time_us"], int) + def test_checkpoint_context_fields(self): + fmt = TraceJsonlFormatter(rank=0, source="test") + record = logging.LogRecord( + name="test", + level=logging.INFO, + pathname="test.py", + lineno=1, + msg="test", + args=None, + exc_info=None, + ) + for key, value in event_extra("log_metric").items(): + setattr(record, key, value) + record.context = [] + record.measured_from_start_time_ms = 123456 + + parsed = json.loads(fmt.format(record)) + + assert parsed["context"] == [] + assert parsed["measured_from_start_time_ms"] == 123456 + + def test_omits_missing_context(self): + fmt = TraceJsonlFormatter(rank=0, source="test") + record = logging.LogRecord( + name="test", + level=logging.INFO, + pathname="test.py", + lineno=1, + msg="test", + args=None, + exc_info=None, + ) + for key, value in event_extra("log_metric").items(): + setattr(record, key, value) + + parsed = json.loads(fmt.format(record)) + + assert "context" not in parsed + # --------------------------------------------------------------------------- # TraceEventsOnlyFilter @@ -607,6 +674,15 @@ def test_second_call_is_noop(self, tmp_path, structured_logger_fixture): assert len(structured_logger_fixture.handlers) == handler_count + def test_records_resolved_init_args(self, tmp_path, structured_logger_fixture): + init_structured_logger(rank=17, source="rl_trainer", output_dir=str(tmp_path)) + + assert _get_structured_logger_init_args() == ( + "rl_trainer", + str(tmp_path), + 17, + ) + class TestFactoryMechanism: def test_default_creates_jsonl(self, tmp_path, structured_logger_fixture): @@ -646,6 +722,60 @@ def fake_factory(*, structured_logger, rank, source, output_dir, **kw): ) +# --------------------------------------------------------------------------- +# External structured logging +# --------------------------------------------------------------------------- + + +class TestExternalStructuredLogging: + def test_init_forwards_only_structured_records_from_other_loggers( + self, tmp_path, structured_logger_fixture, external_logger + ): + init_structured_logger(rank=0, source="trainer", output_dir=str(tmp_path)) + external_logger.setLevel(logging.INFO) + + logging.getLogger("external_library.worker").info( + "external metric", + extra={ + "log_type": "event", + "log_type_name": "log_metric", + "event_name": "train.step.e2e.latency_ms", + "step": 7, + "value": 12.5, + "context": ["source:test"], + }, + ) + logging.getLogger("external_library.worker").info("plain text") + + for handler in structured_logger_fixture.handlers: + handler.flush() + trace_dir = os.path.join(str(tmp_path), "structured_logs") + jsonl_files = [f for f in os.listdir(trace_dir) if f.endswith(".jsonl")] + with open(os.path.join(trace_dir, jsonl_files[0])) as f: + lines = [json.loads(line) for line in f if line.strip()] + + assert len(lines) == 1 + assert lines[0]["logger_name"] == "external_library.worker" + assert lines[0]["log_type_name"] == "log_metric" + assert lines[0]["event_name"] == "train.step.e2e.latency_ms" + assert lines[0]["step"] == 7 + assert lines[0]["value"] == 12.5 + assert lines[0]["context"] == ["source:test"] + assert structured_logger_fixture.propagate is False + + def test_second_init_restores_a_removed_root_forwarder( + self, tmp_path, structured_logger_fixture + ): + root_logger = logging.getLogger() + init_structured_logger(rank=0, source="trainer", output_dir=str(tmp_path)) + forwarder = root_logger.handlers[-1] + root_logger.handlers = [] + + init_structured_logger(rank=0, source="trainer", output_dir=str(tmp_path)) + + assert root_logger.handlers == [forwarder] + + # --------------------------------------------------------------------------- # No-op flag # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/test_torch_checkpointing.py b/tests/unit_tests/test_torch_checkpointing.py index 9b19eda95a..76db114c99 100644 --- a/tests/unit_tests/test_torch_checkpointing.py +++ b/tests/unit_tests/test_torch_checkpointing.py @@ -6,6 +6,7 @@ import dataclasses import json +import logging import queue import unittest from concurrent.futures import Future @@ -27,6 +28,7 @@ SyncCheckpointSaverConfig, ) from torch_checkpointing.default_resharder import DefaultResharder +from torch_checkpointing.logging_utils import checkpoint_logging_context from torchtitan.components.checkpointer import ( BaseCheckpointManager, CheckpointManager, @@ -548,3 +550,105 @@ def test_last_step_uses_synchronous_manager_and_model_only_payload(self) -> None sync_config = build.call_args.args[0] self.assertIsNone(sync_config.pre_finalize_callback) manager.close() + + def test_save_stamps_the_step_on_backend_events(self) -> None: + # The backend reads this context when it builds its own events and + # exports it to the async save subprocess. Without it every forwarded + # backend metric carries step=None, which makes them hard to line up + # against the training step they belong to. + config = TorchCheckpointingManager.Config( + enable=True, + interval=1, + keep_latest_k=0, + initial_load_model_only=False, + ) + manager, backend_manager = self._build_manager(config) + self.addCleanup(checkpoint_logging_context.import_context, {}) + + self.assertTrue(manager.save(curr_step=7)) + + self.assertEqual(7, checkpoint_logging_context.get("step")) + backend_manager.save_result.set_result(None) + manager.close() + + def test_subprocess_logging_initializes_and_delegates(self) -> None: + calls = [] + init_fn = mock.Mock(side_effect=lambda *_args: calls.append("existing")) + with mock.patch.object( + manager_module.sl, + "init_structured_logger", + side_effect=lambda **_kwargs: calls.append("structured"), + ) as init_structured_logger: + manager_module._init_subprocess_logging( + ("rl_trainer", "/tmp/output", 17), + init_fn, + ("argument",), + ) + + init_structured_logger.assert_called_once_with( + source="rl_trainer", + output_dir="/tmp/output", + rank=17, + ) + init_fn.assert_called_once_with("argument") + self.assertEqual(["existing", "structured"], calls) + + def _init_subprocess_logging(self) -> None: + with mock.patch.object(manager_module.sl, "init_structured_logger"): + manager_module._init_subprocess_logging( + ("training", "/tmp/output", 0), None, () + ) + + def test_subprocess_logging_only_overrides_suppressed_inherited_level(self) -> None: + root_logger = logging.getLogger() + backend_logger = logging.getLogger(manager_module._BACKEND_LOGGER_NAME) + self.addCleanup(root_logger.setLevel, root_logger.level) + self.addCleanup(backend_logger.setLevel, backend_logger.level) + for root_level, backend_level, expected_level in ( + (logging.WARNING, logging.NOTSET, logging.INFO), + (logging.WARNING, logging.DEBUG, logging.DEBUG), + (logging.WARNING, logging.WARNING, logging.WARNING), + (logging.DEBUG, logging.NOTSET, logging.NOTSET), + ): + with self.subTest(root_level=root_level, backend_level=backend_level): + root_logger.setLevel(root_level) + backend_logger.setLevel(backend_level) + + self._init_subprocess_logging() + + self.assertEqual(expected_level, backend_logger.level) + + def test_async_manager_composes_subprocess_logging_initializer(self) -> None: + original_init_fn = mock.Mock() + config = TorchCheckpointingManager.Config( + enable=True, + keep_latest_k=0, + initial_load_model_only=False, + ) + backend_config = _default_backend_config() + backend_config.subprocess_init_fn = original_init_fn + backend_config.subprocess_init_args = ("argument",) + + with mock.patch.object( + manager_module, + "_get_structured_logger_init_args", + return_value=("rl_trainer", "/tmp/output", 17), + ): + manager, _ = self._build_manager( + config, + backend_config=backend_config, + ) + + self.assertIs( + manager._manager_config.subprocess_init_fn, + manager_module._init_subprocess_logging, + ) + self.assertEqual( + ( + ("rl_trainer", "/tmp/output", 17), + original_init_fn, + ("argument",), + ), + manager._manager_config.subprocess_init_args, + ) + manager.close() diff --git a/torchtitan/components/checkpointer/torch_checkpointing.py b/torchtitan/components/checkpointer/torch_checkpointing.py index ec6b27962d..22f4609fc7 100644 --- a/torchtitan/components/checkpointer/torch_checkpointing.py +++ b/torchtitan/components/checkpointer/torch_checkpointing.py @@ -7,10 +7,12 @@ from __future__ import annotations import copy +import logging import os import queue import re import threading +from collections.abc import Callable from concurrent.futures import Future from dataclasses import dataclass, replace from pathlib import Path @@ -33,6 +35,7 @@ from torch_checkpointing.distributed_metadata import ( METADATA_FILE_NAME as TORCH_CHECKPOINTING_METADATA_FILE_NAME, ) +from torch_checkpointing.logging_utils import checkpoint_logging_context from torch_checkpointing.schema import ItemSpec from torch_checkpointing.staging import CheckpointStagerConfig from torch_checkpointing.storage.base_storage import Storage, StorageConfig @@ -41,6 +44,9 @@ from torchtitan.components.optimizer import LRSchedulersContainer, OptimizersContainer from torchtitan.config import TORCH_DTYPE_MAP from torchtitan.observability import structured_logger as sl +from torchtitan.observability.structured_logger.structured_logging import ( + _get_structured_logger_init_args, +) from torchtitan.protocols.state_dict_adapter import BaseStateDictAdapter from torchtitan.tools import filesystem from torchtitan.tools.logging import logger @@ -60,6 +66,9 @@ _DEFAULT_BARRIER_INIT_TIMEOUT_SEC = 60 _DEFAULT_BARRIER_TIMEOUT_SEC = 600 +# Logger the backend emits its checkpoint events and metrics on. +_BACKEND_LOGGER_NAME = "torch_checkpointing" + def _step_dir_pattern(temp_dir_prefix: str) -> re.Pattern[str]: """Match published and in-flight checkpoint directory names.""" @@ -101,6 +110,51 @@ def remove(self, path: str) -> None: self._storage.rmdir(Path(path)) +def _init_subprocess_logging( + structured_logger_init_args: tuple[str, str, int], + init_fn: Callable[..., None] | None, + init_args: tuple[Any, ...], +) -> None: + """Re-establish structured logging inside the async save subprocess. + + The subprocess does not inherit the parent's logging handlers, so its + checkpoint records would otherwise be lost. + """ + if init_fn is not None: + init_fn(*init_args) + + source, output_dir, rank = structured_logger_init_args + sl.init_structured_logger(source=source, output_dir=output_dir, rank=rank) + + backend_logger = logging.getLogger(_BACKEND_LOGGER_NAME) + if backend_logger.level == logging.NOTSET and not backend_logger.isEnabledFor( + logging.INFO + ): + backend_logger.setLevel(logging.INFO) + + +def _with_structured_logging( + config: BackendCheckpointManager.Config, +) -> BackendCheckpointManager.Config: + """Forward the backend's own log records into TorchTitan's structured log. + + No-op when structured logging is not active, or when saves are synchronous. + Any existing ``subprocess_init_fn`` is chained rather than replaced. + """ + init_args = _get_structured_logger_init_args() + if init_args is None or not isinstance(config.save, AsyncCheckpointSaverConfig): + return config + return replace( + config, + subprocess_init_fn=_init_subprocess_logging, + subprocess_init_args=( + init_args, + config.subprocess_init_fn, + config.subprocess_init_args, + ), + ) + + def _item_specs() -> dict[str, ItemSpec]: resharder = DefaultResharder() return { @@ -254,6 +308,7 @@ def __init__( # so saves and loads use it too, not just our own path probes. if storage_config is not None: manager_config = replace(manager_config, storage_config=storage_config) + manager_config = _with_structured_logging(manager_config) self._manager_config = manager_config self._step_dir_pattern = _step_dir_pattern( manager_config.save.writer_config.temp_dir_prefix @@ -308,6 +363,10 @@ def _save(self, curr_step: int, last_step: bool = False) -> bool: return False sl.add_step_tag("checkpoint_save") + # The backend stamps its own events from this context and carries it + # into the async save subprocess, so without it every forwarded backend + # metric reports step=None. + checkpoint_logging_context.update(step=curr_step) self.maybe_wait_for_saving() # Purge before issuing this step's save, while the folder holds only # settled state: the previous save has been awaited and the next has not diff --git a/torchtitan/observability/structured_logger/jsonl_handler.py b/torchtitan/observability/structured_logger/jsonl_handler.py index e55a50694a..929a6b8abb 100644 --- a/torchtitan/observability/structured_logger/jsonl_handler.py +++ b/torchtitan/observability/structured_logger/jsonl_handler.py @@ -121,11 +121,21 @@ def _log_dict(self, record: logging.LogRecord) -> dict[str, Any]: if isinstance(value, (float, int)): log_dict["value"] = float(value) + context = getattr(record, str(ExtraFields.CONTEXT), None) + if context is not None: + log_dict["context"] = context + # task_name pairs start/end records task_name = getattr(record, str(ExtraFields.TASK_NAME), None) if task_name is not None: log_dict["task_name"] = task_name + measured_from_start_time_ms = getattr( + record, "measured_from_start_time_ms", None + ) + if measured_from_start_time_ms is not None: + log_dict["measured_from_start_time_ms"] = measured_from_start_time_ms + # Caller field for source traceability (file:line:function) log_dict[ "caller" diff --git a/torchtitan/observability/structured_logger/structured_logging.py b/torchtitan/observability/structured_logger/structured_logging.py index dd69c954a3..b89abf5652 100644 --- a/torchtitan/observability/structured_logger/structured_logging.py +++ b/torchtitan/observability/structured_logger/structured_logging.py @@ -40,6 +40,7 @@ # Used to check if handler has been already initialized. If so, re-initializing # is a no-op _is_initialized: bool = False +_structured_logger_init_args: tuple[str, str, int] | None = None # Set by ``init_structured_logger(enable=False)`` to make all trace calls no-ops. _disabled: bool = False @@ -78,11 +79,39 @@ class ExtraFields(enum.StrEnum): LOG_TYPE_NAME = "log_type_name" EVENT_NAME = "event_name" STEP = "step" + CONTEXT = "context" VALUE = "value" RELATIVE_STEP = "relative_step" TASK_NAME = "task_name" +class _StructuredRecordForwarder(logging.Handler): + """Forward structured records from the root logger to trace handlers.""" + + def emit(self, record: logging.LogRecord) -> None: + if record.name == _structured_logger.name: + return + if getattr(record, str(ExtraFields.LOG_TYPE_NAME), None) is None: + return + _structured_logger.handle(record) + + +_root_forwarder = _StructuredRecordForwarder() + + +def _ensure_root_forwarder() -> None: + root_logger = logging.getLogger() + if _root_forwarder not in root_logger.handlers: + root_logger.addHandler(_root_forwarder) + + +def _get_structured_logger_init_args() -> tuple[str, str, int] | None: + if _disabled or not _is_initialized or not _structured_logger.handlers: + return None + assert _structured_logger_init_args is not None + return _structured_logger_init_args + + def event_extra( event_type: str, event_name: str | None = None, @@ -158,8 +187,7 @@ def init_structured_logger( JSONL handler is registered; when set, ONLY the listed factories run. ``rank`` defaults to ``$RANK`` (set by torchrun), so this can run - before ``torch.distributed`` init. Idempotent: second and later calls - are a no-op. + before ``torch.distributed`` init. Repeated calls do not duplicate handlers. When ``enable=False``, all subsequent ``log_trace_*`` calls become no-ops (no handlers are attached). @@ -171,10 +199,11 @@ def init_structured_logger( init_structured_logger(source="trainer", output_dir="./outputs") log_trace_instant("structured_logger_started") """ - global _is_initialized, _disabled + global _is_initialized, _disabled, _structured_logger_init_args if not enable: _disabled = True + _structured_logger_init_args = None console_logger.info( "Structured logging disabled via DebugConfig.enable_structured_logging=False" ) @@ -182,6 +211,7 @@ def init_structured_logger( # Avoids re-initializing if _is_initialized: + _ensure_root_forwarder() return if rank is None: @@ -209,7 +239,9 @@ def init_structured_logger( ): _structured_logger.setLevel(logging.INFO) + _structured_logger_init_args = (source, output_dir, rank) _is_initialized = True + _ensure_root_forwarder() def log_trace_scalar(scalars: dict[str, float | int], *, stacklevel: int = 2) -> None: