fix: seed timer at training start to avoid AttributeError

[zhixiangli]

Problem

When resuming training from a checkpoint mid-epoch, the workload fails with the following AttributeError:

[rank6]: AttributeError: 'StepTimeCallback' object has no attribute 'start_time'
[rank6]: ^^^^^^^^^^^^^^^
[rank6]: step_time = time.perf_counter() - self.start_time - self.ckpt_time
[rank6]: File "/workload/configs/llama_3_1_8b_cpu_sim.py", line 302, in on_train_batch_end

Root Cause

The StepTimeCallback was initializing self.start_time in the on_train_epoch_start hook. However, when PyTorch Lightning resumes training mid-epoch (e.g., from a checkpoint saved at step=1), the on_train_epoch_start hook is skipped for
that epoch because the epoch is already considered in progress. As a result, self.start_time is never initialized, causing a crash on the first call to on_train_batch_end.

Solution

Initialize self.start_time and self.ckpt_time in the on_train_start hook. This hook is guaranteed to run at the beginning of trainer.fit(), regardless of whether it is a fresh start or a resume.

The on_train_epoch_start hook is retained to reset the timer at the start of subsequent epochs.

[gemini-code-assist]

Code Review

This pull request adds the on_train_start hook to initialize start_time and ckpt_time in the PyTorch Lightning simulation script. The reviewer suggested adding an explanatory comment to this method to clarify that it prevents potential AttributeError issues during mid-epoch resumption, ensuring future developers understand its necessity.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 89.68%. Comparing base (7e658e0) to head (5dd0b53).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #941   +/-   ##
=======================================
  Coverage   89.68%   89.68%           
=======================================
  Files          16       16           
  Lines        3579     3579           
=======================================
  Hits         3210     3210           
  Misses        369      369           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.