feat(pt_expt): add rank-aware TensorBoard training observer

[OutisLi-Bot]

Summary

Implements #5821: training.tensorboard / tensorboard_freq / tensorboard_log_dir for pt_expt through a rank-aware TrainingObserver.

• Add a lightweight TrainingObserver lifecycle on the shared AbstractTrainer loop (on_train_begin, on_step_end, on_display, on_checkpoint, on_train_end).
• Add TensorBoardObserver under deepmd/pt_expt/train (reusable by legacy PT). Only the chief rank owns a SummaryWriter.
• Disabled runs install no observer: no TensorBoard import, writer, filesystem access, or host sync.
• tensorboard_freq is independent of disp_freq; step hooks run only when due.
• Writer is flushed/closed on normal completion and exceptional exit; restarts keep existing event files and continue with restored global steps.
• Schema docs now list pt_expt for the TensorBoard options.

Tests

• source/tests/common/dpmodel/test_train_observer.py
• source/tests/pt_expt/train/test_tensorboard_observer.py (single-process + simulated multi-rank chief-only writer, frequency, cleanup, restart)

pytest -q source/tests/common/dpmodel/test_train_observer.py \
  source/tests/pt_expt/train/test_tensorboard_observer.py \
  source/tests/common/dpmodel/test_train_abstract_trainer.py

Refs #5821 #5755

Summary by CodeRabbit

• New Features
  • Added TensorBoard logging for PyTorch Exportable training, with configurable output directory and frequency for training and validation metrics, learning rates, and timing information.
  • Added training observation hooks for step, display, checkpoint, and lifecycle events, allowing integrations to respond to training progress.
• Tests
  • Added coverage for observer callbacks and TensorBoard logging, including frequency settings and rank-specific writer behavior.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: b417d71e-cc1f-4b90-b1af-68f4d0c24ffb

---

📥 Commits

Reviewing files that changed from the base of the PR and between fec294c and e626f00.

---

📒 Files selected for processing (5)

• deepmd/dpmodel/train/__init__.py
• deepmd/dpmodel/train/data.py
• deepmd/dpmodel/train/observer.py
• deepmd/dpmodel/train/trainer.py
• deepmd/dpmodel/train/types.py

---

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

---

📝 Walkthrough

Walkthrough

The training API adds shared training types and observer payloads with lifecycle hooks. The trainer dispatches observations during training. The PyTorch Exportable backend adds a TensorBoard observer for training, validation, and timing metrics.

Changes

Training observers and TensorBoard logging

Layer / File(s)	Summary
Shared training types and observer API deepmd/dpmodel/train/types.py, deepmd/dpmodel/train/observer.py, deepmd/dpmodel/train/__init__.py, deepmd/dpmodel/train/data.py, deepmd/dpmodel/train/trainer.py	Adds shared rank, task, and step-result types, observer payloads and hooks, ordered callback dispatch, and package exports. Training data imports the default task key from the shared types module.
Trainer event integration deepmd/dpmodel/train/trainer.py, source/tests/common/dpmodel/test_train_observer.py	The trainer accepts observers and sends lifecycle, step, display, and checkpoint observations. Tests cover callback frequencies, skipped step events, and train-end callbacks after failure.
TensorBoard observer and pt_expt wiring deepmd/pt_expt/train/tensorboard.py, deepmd/pt_expt/train/training.py, deepmd/utils/argcheck.py, source/tests/pt_expt/train/test_tensorboard_observer.py	Adds TensorBoard scalar logging and writer lifecycle handling, wires the observer into pt_expt training, and adds pt_expt to the supported-backend labels. Tests cover configuration, ranks, frequencies, tags, and writer behavior.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant PtTraining as pt_expt training
  participant Trainer as AbstractTrainer
  participant Observers as TrainingObserverList
  participant TensorBoard as TensorBoardObserver
  participant Writer as SummaryWriter
  PtTraining->>Trainer: Pass observer and RankContext
  Trainer->>Observers: Dispatch training and observation callbacks
  Observers->>TensorBoard: Call observer hooks
  TensorBoard->>Writer: Create writer on enabled chief rank
  TensorBoard->>Writer: Write metrics and timing scalars
  TensorBoard->>Writer: Flush at checkpoint and close at training end

Loading

Suggested reviewers: outisli

---

Merge Risk: ⚪ Minimal · up to e626f

No merge-blocking issue was established; the reviewed change is ready to merge after normal checks.

Pre-merge checks | 4 | 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 28.05% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 82 functions across 10 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely identifies the main change: adding a rank-aware TensorBoard training observer for pt_expt.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

---

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🧪 Generate unit tests (beta)

• Create a new PR

---

🛠️ Fix failing CI checks 💡

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

---

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

❌ Patch coverage is 88.07339% with 39 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.00%. Comparing base (64d3266) to head (e626f00).

Files with missing lines	Patch %	Lines
deepmd/pt_expt/train/tensorboard.py	82.30%	23 Missing ⚠️
deepmd/dpmodel/train/types.py	90.90%	9 Missing ⚠️
deepmd/dpmodel/train/observer.py	92.30%	5 Missing ⚠️
deepmd/dpmodel/train/trainer.py	92.85%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6066      +/-   ##
==========================================
- Coverage   77.85%   76.00%   -1.86%     
==========================================
  Files        1170     1173       +3     
  Lines      140914   141100     +186     
  Branches     5056     5051       -5     
==========================================
- Hits       109707   107241    -2466     
- Misses      29323    32282    +2959     
+ Partials     1884     1577     -307     

☔ 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.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.