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: