Add logging for dataset input and checkpoint saving by linxiulei · Pull Request #1190 · apple/axlearn

linxiulei · 2025-05-18T20:15:22Z

This PR adds:

logging utility class to log when an operation's execution time exceeds a specified threshold.
use mentioned utility class to around dataset input fetch (per iteration) and checkpoint saving to allow us to observe execution in poor performance

findmyway · 2025-05-19T01:42:24Z

axlearn/common/trainer.py


-        # Checkpointer policy will decide if we should save.
-        self.save_checkpoint(evaler_summaries=evaler_summaries)
+        with LoggingContext("save_checkpoint", timeout=180, threshold=120):


Can we make this configurable?

Sure, but I am thinking which way is better to make it configurable, env var or flags?

The typical convention we use is that if we want to make the child configurable, then the parent config should contain a config for the child.

apghml · 2025-05-19T18:03:50Z

axlearn/common/logging_utils.py

+import time
+
+
+class LoggingContext:


Could you explain the difference of this from the other monitor functionality that exists like GPUMonitor, etc? If they serve overlapping purposes, should they be unified?

Can you please clarify what other monitor functionality we already have? If they overlap, we can unify them

TfSummaryMonitor, GPUMonitor, GoodputMonitor. I might be forgetting some additional ones.

Thanks, I checked those functionality and believe there don't overlap because this PR is mostly for logging specific operations' execution time (the utility function allows you to monitor any operations you want) and those functionality you mentioned is to monitor with specific purposes (goodputs, GPU hardware status etc)

apghml · 2025-05-19T18:04:16Z

axlearn/common/trainer.py

-        self.save_checkpoint(evaler_summaries=evaler_summaries)
+        with LoggingContext("save_checkpoint", timeout=180, threshold=120):
+            # Checkpointer policy will decide if we should save.
+            self.save_checkpoint(evaler_summaries=evaler_summaries)


Since checkpoint syncing is asynchronous, do we expect to see timeouts inside this block?

Yes, IIUC checkpoint saving is asynchronous if the previous saving is complete or it may stuck in https://github.com/apple/axlearn/blob/main/axlearn/common/array_serialization.py#L554

Ethanlm · 2025-05-19T18:13:37Z

axlearn/common/trainer.py

                stop_trace_step = None

-                input_iterator = self.input.batches(self._input_iter)
+                input_iterator = LoggingIterator(


Since it is already recorded in the Goodput, I wonder if we need a separate logging here.

If I am reading it correctly, this LoggingIterator creates one timer thread every time next is called, which seems to be adding a lot of thread creation/cancellation overhead to the main process

Assuming you mean this code

self._maybe_record_event(measurement.Event.START_DATA_LOADING) try: input_batch = next(input_iterator) self._maybe_record_event(measurement.Event.END_DATA_LOADING)

If the next() somehow hangs and then the whole process aborts, would Goodput measurement help catch that?

My understanding is this iterator invokes next() every a few seconds depending on the training scale, so creating and cleaning a thread for every a few seconds should have only very minimal overhead (empirically one thread takes 10 to 100 usec to create).

However, if you still have concerns on overhead, we might create a reusable thread for timer purpose.

ruomingp · 2025-05-21T13:52:58Z

axlearn/common/logging_utils.py

+  """A context manager for monitoring the execution time of operations and
+  logging warnings based on defined timeouts and thresholds.
+
+  This context manager can be used to:
+  - Log a warning if an operation exceeds a specified timeout while still
+  running.
+  - Log a warning if an operation's total execution time exceeds a specified
+  threshold.


It's not clear that we need logging-upon-timeout, given the threading concerns raised by @Ethanlm. I wonder if a simpler semantic would be sufficient:

LoggingContext always wait until the op finishes and log a warning if the execution time exceeds the threshold;

We rely on the watchdog

axlearn/axlearn/common/trainer.py

Line 190 in 2c4ec6e

watchdog_timeout_seconds: Optional[float] = None

to tell us where the threads are stuck

WDYT?

I agree with starting something simpler so I changed this PR to only implement logging upon op completion.

However, I feel logging-upon-timeout is important because program hangs are difficult to debug by just looking through tracebacks of so many workers printed by watchdog timeout. A clear logging message may help identify workers that cause the hangs.

Introduce logging utilities designed to proactively identify and flag operations exceeding performance expectations.

Ethanlm

Thanks

Ethanlm · 2025-05-30T23:48:07Z

axlearn/common/trainer.py

                    self._maybe_record_event(measurement.Event.START_DATA_LOADING)
                    try:
-                        input_batch = next(input_iterator)
+                        with LoggingContext("input_iterator", threshold=60):


Adding "step" information in the log can be useful when we try to debug

Ethanlm · 2025-06-10T17:23:45Z

axlearn/common/trainer.py


-        # Checkpointer policy will decide if we should save.
-        self.save_checkpoint(evaler_summaries=evaler_summaries)
+        with LoggingContext("save_checkpoint", threshold=120):


Can we make it configurable?
We should be able to patch this ad-hoc in jobs we want to debug and monitor closely. But the hard-coded 120s here makes it less flexible to apply in general jobs, hence not a good candidate for PR merge

github-actions · 2025-10-21T02:10:42Z

This pull request has been automatically marked as stale because it has been inactive for 60 days. It will be closed in 7 days if no further activity occurs. If you would like to continue working on this, please remove the stale label or leave a comment.

github-actions · 2025-12-31T02:17:23Z

This pull request has been automatically marked as stale because it has been inactive for 60 days. It will be closed in 7 days if no further activity occurs. If you would like to continue working on this, please remove the stale label or leave a comment.

github-actions · 2026-01-07T02:17:57Z

This pull request was closed because it has been inactive for more than 7 days since being marked as stale. Please feel free to reopen it if you would like to continue.

linxiulei requested review from a team, markblee and ruomingp as code owners May 18, 2025 20:15

findmyway reviewed May 19, 2025

View reviewed changes

apghml suggested changes May 19, 2025

View reviewed changes

Ethanlm reviewed May 19, 2025

View reviewed changes

ruomingp reviewed May 21, 2025

View reviewed changes

linxiulei added 2 commits May 22, 2025 15:58

Introduce LoggingContext

4a3b4a8

Introduce logging utilities designed to proactively identify and flag operations exceeding performance expectations.

Add logging for fetching input dataset and saving checkpoint

c9a34cb

linxiulei force-pushed the logging branch from 73a85d3 to c9a34cb Compare May 22, 2025 16:00

Ethanlm reviewed Jun 10, 2025

View reviewed changes

github-actions bot added the stale label Oct 21, 2025

findmyway removed the stale label Oct 30, 2025

github-actions bot added the stale label Dec 31, 2025

github-actions bot closed this Jan 7, 2026

		import time


		class LoggingContext:

Conversation

linxiulei commented May 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Ethanlm left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Oct 21, 2025

Uh oh!

github-actions bot commented Dec 31, 2025

Uh oh!

github-actions bot commented Jan 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

linxiulei commented May 18, 2025 •

edited

Loading