ml_goodput_measurement_recorder_training_step_logging.py

1import time
2from ml_goodput_measurement import GoodputRecorder
3
4# 1. Initialize the GoodputRecorder
5# In a real scenario, provide the job name and worker information.
6# For local testing, these can often be inferred or mocked.
7recorder = GoodputRecorder(
8    job_name="example-training-job",
9    logger_name="goodput_logger"
10)
11
12def train_model():
13    total_steps = 10
14    print(f"Starting training for {total_steps} steps...")
15
16    for step in range(total_steps):
17        # Record the start of the step
18        recorder.record_step_start(step)
19        
20        # Simulate training work (e.g., forward/backward pass)
21        time.sleep(0.5)
22        
23        # 2. Record the end of the step
24        # This logs the time taken and contributes to the goodput calculation
25        recorder.record_step_end(step)
26        
27        print(f"Step {step} completed.")
28
29    print("Training complete.")
30
31if __name__ == "__main__":
32    train_model()