Training engine: implement eval_step API

[SurbhiJainUSC]

Description

• Implements eval_step() API in MaxTextTrainingEngine mirroring tunix's peft_trainer_v2.
• MetricsRecorder takes a mode argument and the engine owns a second mode="eval" instance, keeping eval numbers out of the train buffer and out of the checkpoint; MetricsBuffer.mode then flows through write_metrics into both the console line and the TensorBoard tag, which becomes {mode}/{name}.
• eval_context() accumulates across every micro-batch into one buffer (WeightedMetric appends sum and denominator, so compute() is the weighted mean over the whole eval set) and writes it exactly once on exit.

Tests

CI tests

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request implements evaluation functionality in the MaxText training engine, introducing an evaluation context, an evaluation step, and separate metrics recording and logging for training and evaluation modes. The review feedback highlights two key improvement opportunities: safely retrieving metrics in _log_metrics to prevent potential KeyError exceptions while also adding console logging for the evaluation mode, and ensuring that _eval_metrics_recorder is properly cleaned up alongside other recorders in the engine's close method.

[codecov]

Codecov Report

❌ Patch coverage is 79.16667% with 25 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/training_engine/maxtext_engine.py	82.22%	10 Missing and 6 partials ⚠️
src/maxtext/training_engine/metrics.py	67.85%	6 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!