From 64ab8c771fe5e2545eecfbaa0a9429754385afde Mon Sep 17 00:00:00 2001 From: Ali Arda Eker Date: Fri, 17 Jul 2026 11:08:25 -0700 Subject: [PATCH] Log app_finish_time and app_train_loop_finish_time on early-exit path The exit-duration/exit-interval/signal exit path sys.exit()s before the normal-path one-logger logging, leaving both timestamps unset on otherwise COMPLETED jobs. This mis-attributes end-of-run time (shutdown -> uncategorized). Log both timestamps on the should_exit path so shutdown time is recorded. Signed-off-by: Ali Arda Eker --- megatron/training/training.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/megatron/training/training.py b/megatron/training/training.py index f67397fe889..0a74de1c6dc 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -4001,6 +4001,13 @@ def trace_handler(p): if should_exit: break + # Early-exit paths (exit-duration / exit-interval / signal handler) sys.exit() + # below before the normal-path logging, so record the train-loop finish time here. + if should_exit: + one_logger and one_logger.log_metrics( + {'app_train_loop_finish_time': one_logger_utils.get_timestamp_in_ms()} + ) + # Destroy CUDA Graphs. if args.cuda_graph_impl == "transformer_engine" and cuda_graph_helper.graphs_created(): cuda_graph_helper.delete_cuda_graphs() @@ -4048,6 +4055,9 @@ def trace_handler(p): for buf in model_module.buffers + model_module.expert_parallel_buffers: if getattr(buf, 'nccl_mem_pool', None) is not None: nccl_allocator.deregister_mem_pool(buf.nccl_mem_pool, buf.data_parallel_group) + one_logger and one_logger.log_metrics( + {'app_finish_time': one_logger_utils.get_timestamp_in_ms()} + ) wandb_writer = get_wandb_writer() if wandb_writer: wandb_writer.finish()