From e00c7e15c1b751a1c473d078dac6b6fbad6af6b0 Mon Sep 17 00:00:00 2001 From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:02:20 +0800 Subject: [PATCH 01/15] feat(pt_expt): add rank-aware TensorBoard training observer Introduce a TrainingObserver lifecycle on the shared trainer loop and a chief-owned TensorBoardObserver for pt_expt so tensorboard / tensorboard_freq / tensorboard_log_dir honor rank awareness, independent logging frequency, and clean writer teardown without paying import or host-sync costs when disabled. Refs #5821 #5755 --- deepmd/dpmodel/train/__init__.py | 12 + deepmd/dpmodel/train/observer.py | 207 +++++++++++++ deepmd/dpmodel/train/trainer.py | 72 ++++- deepmd/pt_expt/train/tensorboard.py | 287 ++++++++++++++++++ deepmd/pt_expt/train/training.py | 11 +- deepmd/utils/argcheck.py | 6 +- .../common/dpmodel/test_train_observer.py | 172 +++++++++++ .../train/test_tensorboard_observer.py | 249 +++++++++++++++ 8 files changed, 1008 insertions(+), 8 deletions(-) create mode 100644 deepmd/dpmodel/train/observer.py create mode 100644 deepmd/pt_expt/train/tensorboard.py create mode 100644 source/tests/common/dpmodel/test_train_observer.py create mode 100644 source/tests/pt_expt/train/test_tensorboard_observer.py diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py index 1463dad928..ce42bd8444 100644 --- a/deepmd/dpmodel/train/__init__.py +++ b/deepmd/dpmodel/train/__init__.py @@ -19,6 +19,13 @@ from .metrics import ( TrainingMetricAccumulator, ) +from .observer import ( + CheckpointObservation, + DisplayObservation, + StepObservation, + TrainingObserver, + TrainingObserverList, +) from .schedule import ( StepSchedule, resolve_step_schedule, @@ -47,16 +54,21 @@ "DEFAULT_TASK_KEY", "AbstractTrainEntrypoint", "AbstractTrainer", + "CheckpointObservation", "CheckpointStore", "DisplayInterval", + "DisplayObservation", "LearningCurveWriter", "RankContext", "ShardingPolicy", + "StepObservation", "StepSchedule", "TrainEntrypointOptions", "TrainStepResult", "TrainerConfig", "TrainingMetricAccumulator", + "TrainingObserver", + "TrainingObserverList", "TrainingTask", "TrainingTaskCollection", "TrainingTaskConfig", diff --git a/deepmd/dpmodel/train/observer.py b/deepmd/dpmodel/train/observer.py new file mode 100644 index 0000000000..6b7da5d88b --- /dev/null +++ b/deepmd/dpmodel/train/observer.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Composable observation hooks for the common training loop. + +Observers sit outside the optimizer step and outside any compiled model +graph. They receive rank context and already-produced step or display +payloads; each observer decides whether a given boundary is worth acting +on. The training loop never imports logging backends on their behalf. +""" + +from __future__ import ( + annotations, +) + +from abc import ( + ABC, +) +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +if TYPE_CHECKING: + from collections.abc import ( + Iterator, + Sequence, + ) + + from .timing import ( + DisplayInterval, + ) + from .trainer import ( + RankContext, + TrainingTaskCollection, + TrainStepResult, + ) + +__all__ = [ + "CheckpointObservation", + "DisplayObservation", + "StepObservation", + "TrainingObserver", + "TrainingObserverList", +] + + +@dataclass(frozen=True) +class StepObservation: + """Payload delivered after one optimizer step. + + Attributes + ---------- + step: + Zero-based index of the step that just finished. + display_step: + One-based index used for logging and checkpoint cadence. + task_key: + Task selected for the step. + learning_rate: + Learning rate associated with ``step``. + step_result: + Backend payload returned by :meth:`AbstractTrainer.train_step`. + rank_context: + Rank metadata for the calling process. + """ + + step: int + display_step: int + task_key: str + learning_rate: float + step_result: TrainStepResult + rank_context: RankContext + + +@dataclass(frozen=True) +class DisplayObservation: + """Payload delivered at a learning-curve display boundary. + + Attributes + ---------- + step: + Zero-based index of the step that just finished. + display_step: + One-based index used for logging. + task_key: + Task selected for the step. + learning_rate: + Learning rate associated with ``step``. + train_results: + Training metrics prepared for the learning curve. + valid_results: + Validation metrics prepared for the learning curve, if any. + timing: + Wall-clock summary of the interval ending at ``display_step``. + rank_context: + Rank metadata for the calling process. + """ + + step: int + display_step: int + task_key: str + learning_rate: float + train_results: Any + valid_results: Any | None + timing: DisplayInterval + rank_context: RankContext + + +@dataclass(frozen=True) +class CheckpointObservation: + """Payload delivered when a checkpoint boundary is reached.""" + + display_step: int + rank_context: RankContext + + +class TrainingObserver(ABC): + """Optional side-effect hooks around the common training loop. + + Default implementations are no-ops so an observer can override only the + boundaries it cares about. Implementations must keep work outside + compiled model execution and must not force device-to-host traffic when + they choose not to emit. + """ + + def wants_step(self, display_step: int) -> bool: + """Return whether :meth:`on_step_end` should run for this step.""" + return False + + def on_train_begin( + self, + tasks: TrainingTaskCollection, + *, + rank_context: RankContext, + ) -> None: + """Called once before the first optimizer step.""" + return None + + def on_step_end(self, observation: StepObservation) -> None: + """Called after an optimizer step when :meth:`wants_step` is true.""" + return None + + def on_display(self, observation: DisplayObservation) -> None: + """Called after learning-curve metrics are collected on the chief.""" + return None + + def on_checkpoint(self, observation: CheckpointObservation) -> None: + """Called after a checkpoint boundary has been handled.""" + return None + + def on_train_end( + self, + tasks: TrainingTaskCollection, + *, + rank_context: RankContext, + ) -> None: + """Called after the run finishes or aborts, once resources close.""" + return None + + +class TrainingObserverList(TrainingObserver): + """Fan-out wrapper that preserves observer order.""" + + def __init__(self, observers: Sequence[TrainingObserver] | None = None) -> None: + self._observers = tuple(observer for observer in (observers or ()) if observer) + + def __bool__(self) -> bool: + return bool(self._observers) + + def __iter__(self) -> Iterator[TrainingObserver]: + return iter(self._observers) + + def wants_step(self, display_step: int) -> bool: + return any(observer.wants_step(display_step) for observer in self._observers) + + def on_train_begin( + self, + tasks: TrainingTaskCollection, + *, + rank_context: RankContext, + ) -> None: + for observer in self._observers: + observer.on_train_begin(tasks, rank_context=rank_context) + + def on_step_end(self, observation: StepObservation) -> None: + for observer in self._observers: + if observer.wants_step(observation.display_step): + observer.on_step_end(observation) + + def on_display(self, observation: DisplayObservation) -> None: + for observer in self._observers: + observer.on_display(observation) + + def on_checkpoint(self, observation: CheckpointObservation) -> None: + for observer in self._observers: + observer.on_checkpoint(observation) + + def on_train_end( + self, + tasks: TrainingTaskCollection, + *, + rank_context: RankContext, + ) -> None: + for observer in self._observers: + observer.on_train_end(tasks, rank_context=rank_context) diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py index 531bc7d07c..a1333c0c78 100644 --- a/deepmd/dpmodel/train/trainer.py +++ b/deepmd/dpmodel/train/trainer.py @@ -49,6 +49,13 @@ format_training_message_per_task, ) +from .observer import ( + CheckpointObservation, + DisplayObservation, + StepObservation, + TrainingObserver, + TrainingObserverList, +) from .timing import ( DisplayInterval, TrainingTimer, @@ -538,10 +545,19 @@ def __init__( *, rank_context: RankContext | None = None, metric_accumulator: TrainingMetricAccumulator | None = None, + observers: Sequence[TrainingObserver] | TrainingObserver | None = None, ) -> None: self.trainer_config = trainer_config self.rank_context = rank_context or RankContext() self.metric_accumulator = metric_accumulator + if observers is None: + self.observers = TrainingObserverList() + elif isinstance(observers, TrainingObserverList): + self.observers = observers + elif isinstance(observers, TrainingObserver): + self.observers = TrainingObserverList((observers,)) + else: + self.observers = TrainingObserverList(observers) self.lcurve_writer = LearningCurveWriter() def run(self, tasks: TrainingTaskCollection) -> None: @@ -554,6 +570,7 @@ def run(self, tasks: TrainingTaskCollection) -> None: if self.metric_accumulator is not None: self.metric_accumulator.reset() self.on_train_begin(tasks) + self.observers.on_train_begin(tasks, rank_context=self.rank_context) fout = self._open_learning_curve() timer = TrainingTimer( start_step=start_step, @@ -570,6 +587,20 @@ def run(self, tasks: TrainingTaskCollection) -> None: ) self.metric_accumulator.add(task.key, step_result.train_results) display_step = step + 1 + current_lr: float | None = None + + if self.observers.wants_step(display_step): + current_lr = self.learning_rate(step) + self.observers.on_step_end( + StepObservation( + step=step, + display_step=display_step, + task_key=task.key, + learning_rate=current_lr, + step_result=step_result, + rank_context=self.rank_context, + ) + ) if self._should_display(display_step): if self.rank_context.is_chief: @@ -579,7 +610,9 @@ def run(self, tasks: TrainingTaskCollection) -> None: step=step, step_result=step_result, ) - current_lr = self.learning_rate(step) + if current_lr is None: + current_lr = self.learning_rate(step) + interval = timer.record(display_step) self.lcurve_writer.log_results( step=display_step, learning_rate=current_lr, @@ -595,7 +628,7 @@ def run(self, tasks: TrainingTaskCollection) -> None: else () ), ) - self._log_interval(timer.record(display_step)) + self._log_interval(interval) if fout is not None: if fout.tell() == 0: self.lcurve_writer.write_header( @@ -610,13 +643,29 @@ def run(self, tasks: TrainingTaskCollection) -> None: train_results=train_results, valid_results=valid_results, ) + self.observers.on_display( + DisplayObservation( + step=step, + display_step=display_step, + task_key=task.key, + learning_rate=current_lr, + train_results=train_results, + valid_results=valid_results, + timing=interval, + rank_context=self.rank_context, + ) + ) if self.metric_accumulator is not None: self.metric_accumulator.reset() + if current_lr is None: + learning_rate = self.learning_rate(step) + else: + learning_rate = current_lr self.run_full_validation( step=step, display_step=display_step, - learning_rate=self.learning_rate(step), + learning_rate=learning_rate, ) if ( @@ -625,14 +674,29 @@ def run(self, tasks: TrainingTaskCollection) -> None: and display_step % self.trainer_config.save_freq == 0 ): self.save_checkpoint(display_step) + self.observers.on_checkpoint( + CheckpointObservation( + display_step=display_step, + rank_context=self.rank_context, + ) + ) if self._should_save_final_checkpoint(): self.save_checkpoint(num_steps) + self.observers.on_checkpoint( + CheckpointObservation( + display_step=num_steps, + rank_context=self.rank_context, + ) + ) self._log_average_step_time(timer) finally: if fout is not None: fout.close() - self.on_train_end(tasks) + try: + self.on_train_end(tasks) + finally: + self.observers.on_train_end(tasks, rank_context=self.rank_context) def select_task(self, tasks: TrainingTaskCollection) -> TrainingTask: """Select the task for the next optimizer step.""" diff --git a/deepmd/pt_expt/train/tensorboard.py b/deepmd/pt_expt/train/tensorboard.py new file mode 100644 index 0000000000..e080ddbec8 --- /dev/null +++ b/deepmd/pt_expt/train/tensorboard.py @@ -0,0 +1,287 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Rank-aware TensorBoard observation for pt_expt training. + +The observer owns the SummaryWriter lifecycle. Only the chief rank creates a +writer; non-chief ranks and disabled runs perform no import, filesystem, or +metric host sync work. Legacy PT may construct the same observer. +""" + +from __future__ import ( + annotations, +) + +import logging +from collections.abc import ( + Mapping, +) +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.train.observer import ( + CheckpointObservation, + DisplayObservation, + StepObservation, + TrainingObserver, +) +from deepmd.dpmodel.train.trainer import ( + DEFAULT_TASK_KEY, + RankContext, + TrainingTaskCollection, +) + +log = logging.getLogger(__name__) + +__all__ = [ + "TensorBoardObserver", + "create_tensorboard_observer", +] + + +def create_tensorboard_observer( + training_params: Mapping[str, Any], + *, + rank_context: RankContext, + multi_task: bool = False, +) -> TensorBoardObserver | None: + """Return a TensorBoard observer when enabled; otherwise ``None``. + + A disabled configuration returns ``None`` so the trainer installs no + observer and never imports TensorBoard. + """ + if not bool(training_params.get("tensorboard", False)): + return None + return TensorBoardObserver( + log_dir=str(training_params.get("tensorboard_log_dir", "log")), + freq=int(training_params.get("tensorboard_freq", 1)), + rank_context=rank_context, + multi_task=multi_task, + ) + + +class TensorBoardObserver(TrainingObserver): + """Emit training scalars through a chief-owned SummaryWriter.""" + + def __init__( + self, + *, + log_dir: str, + freq: int, + rank_context: RankContext, + multi_task: bool = False, + writer_factory: Any | None = None, + ) -> None: + self._log_dir = str(log_dir) + self._freq = int(freq) + self._rank_context = rank_context + self._multi_task = bool(multi_task) + self._writer_factory = writer_factory + self._writer: Any | None = None + self._closed = False + + @property + def enabled(self) -> bool: + """Whether this rank would create a writer.""" + return self._freq > 0 and self._rank_context.is_chief + + @property + def log_dir(self) -> str: + return self._log_dir + + @property + def writer(self) -> Any | None: + return self._writer + + def wants_step(self, display_step: int) -> bool: + return self.enabled and self._is_due(display_step) + + def on_train_begin( + self, + tasks: TrainingTaskCollection, + *, + rank_context: RankContext, + ) -> None: + if not self.enabled: + return + writer_factory = self._writer_factory + if writer_factory is None: + from torch.utils.tensorboard import ( + SummaryWriter, + ) + + writer_factory = SummaryWriter + + Path(self._log_dir).mkdir(parents=True, exist_ok=True) + # A restart keeps existing event files; TensorBoard merges by step, so + # continuing with restored global steps does not overwrite history. + self._writer = writer_factory(log_dir=self._log_dir) + self._closed = False + log.info("TensorBoard events will be written under %s", self._log_dir) + + def on_step_end(self, observation: StepObservation) -> None: + writer = self._writer + if writer is None or not self._is_due(observation.display_step): + return + writer.add_scalar( + "learning_rate", + float(observation.learning_rate), + observation.display_step, + ) + self._write_train_from_step(observation) + + def on_display(self, observation: DisplayObservation) -> None: + writer = self._writer + if writer is None or not self._is_due(observation.display_step): + return + # Validation and timing exist only at display boundaries. Emit them + # when that boundary also lands on tensorboard_freq so every tag + # family shares the configured cadence. + writer.add_scalar( + "learning_rate", + float(observation.learning_rate), + observation.display_step, + ) + self._write_metric_tree( + "train", + observation.train_results, + observation.display_step, + ) + self._write_metric_tree( + "valid", + observation.valid_results, + observation.display_step, + ) + timing = observation.timing + writer.add_scalar( + "timing/interval_wall_time", + float(timing.wall_time), + observation.display_step, + ) + writer.add_scalar( + "timing/interval_steps", + float(timing.steps), + observation.display_step, + ) + if timing.eta is not None: + writer.add_scalar( + "timing/eta", + float(timing.eta), + observation.display_step, + ) + + def on_checkpoint(self, observation: CheckpointObservation) -> None: + writer = self._writer + if writer is not None: + writer.flush() + + def on_train_end( + self, + tasks: TrainingTaskCollection, + *, + rank_context: RankContext, + ) -> None: + self.close() + + def close(self) -> None: + """Flush and close the writer if this rank owns one.""" + writer = self._writer + if writer is None or self._closed: + self._writer = None + self._closed = True + return + try: + writer.flush() + writer.close() + finally: + self._writer = None + self._closed = True + + def _is_due(self, display_step: int) -> bool: + if self._freq <= 0: + return False + return display_step == 1 or display_step % self._freq == 0 + + def _write_train_from_step(self, observation: StepObservation) -> None: + step_result = observation.step_result + metrics = step_result.train_results + if metrics is None and isinstance(step_result.payload, Mapping): + more_loss = step_result.payload.get("more_loss") + if isinstance(more_loss, Mapping): + metrics = { + key: value + for key, value in more_loss.items() + if "l2_" not in key + } + if not metrics: + return + if self._multi_task: + self._write_metric_mapping( + f"train/{observation.task_key}", + metrics, + observation.display_step, + ) + else: + self._write_metric_mapping("train", metrics, observation.display_step) + + def _write_metric_tree( + self, + prefix: str, + results: Any, + display_step: int, + ) -> None: + if results is None: + return + if self._is_task_results(results): + for task_key, task_metrics in results.items(): + if not task_metrics: + continue + self._write_metric_mapping( + f"{prefix}/{task_key}", + task_metrics, + display_step, + ) + return + if isinstance(results, Mapping): + tag_prefix = ( + f"{prefix}/{DEFAULT_TASK_KEY}" if self._multi_task else prefix + ) + self._write_metric_mapping(tag_prefix, results, display_step) + + def _write_metric_mapping( + self, + prefix: str, + metrics: Mapping[str, Any], + display_step: int, + ) -> None: + writer = self._writer + if writer is None: + return + for key, value in metrics.items(): + scalar = _as_float(value) + if scalar is None: + continue + writer.add_scalar(f"{prefix}/{key}", scalar, display_step) + + @staticmethod + def _is_task_results(results: Any) -> bool: + if not isinstance(results, Mapping) or not results: + return False + first = next(iter(results.values())) + return first is None or isinstance(first, Mapping) + + +def _as_float(value: Any) -> float | None: + """Convert a metric to float only when an event is being written.""" + if value is None: + return None + if torch.is_tensor(value): + return float(value.detach().float().cpu().item()) + try: + return float(value) + except (TypeError, ValueError): + return None diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 9130f5753b..a3fce537dc 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -120,6 +120,9 @@ NonFiniteGradGuard, clip_grad_norm_, ) +from deepmd.pt_expt.train.tensorboard import ( + create_tensorboard_observer, +) from deepmd.pt_expt.train.utils import ( MatmulPrecisionPolicy, count_parameters, @@ -2384,6 +2387,7 @@ def update_finetune_bias( self._compile_model(compile_opts) self.training_tasks = self._make_training_tasks() + rank_context = RankContext(rank=self.rank, world_size=self.world_size) super().__init__( TrainerConfig.from_training_params( training_params, @@ -2391,7 +2395,7 @@ def update_finetune_bias( start_step=self.start_step, restart_training=self.restart_training, ), - rank_context=RankContext(rank=self.rank, world_size=self.world_size), + rank_context=rank_context, metric_accumulator=( TrainingMetricAccumulator( { @@ -2402,6 +2406,11 @@ def update_finetune_bias( if training_params.get("disp_avg", False) else None ), + observers=create_tensorboard_observer( + training_params, + rank_context=rank_context, + multi_task=self.multi_task, + ), ) self.full_validator, self.ema_full_validator = self._create_full_validators( validating_params=validating_params, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 982c5d2a26..ff531df09f 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -6253,21 +6253,21 @@ def training_args( bool, optional=True, default=False, - doc=supported_backends("tf", "pt", "pd", "tf2") + doc_tensorboard, + doc=supported_backends("tf", "pt", "pd", "tf2", "pt_expt") + doc_tensorboard, ), Argument( "tensorboard_log_dir", str, optional=True, default="log", - doc=supported_backends("tf", "pt", "pd", "tf2") + doc_tensorboard_log_dir, + doc=supported_backends("tf", "pt", "pd", "tf2", "pt_expt") + doc_tensorboard_log_dir, ), Argument( "tensorboard_freq", int, optional=True, default=1, - doc=supported_backends("tf", "pt", "pd", "tf2") + doc_tensorboard_freq, + doc=supported_backends("tf", "pt", "pd", "tf2", "pt_expt") + doc_tensorboard_freq, ), Argument( "gradient_max_norm", diff --git a/source/tests/common/dpmodel/test_train_observer.py b/source/tests/common/dpmodel/test_train_observer.py new file mode 100644 index 0000000000..47a9c03fe3 --- /dev/null +++ b/source/tests/common/dpmodel/test_train_observer.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from pathlib import ( + Path, +) + +from deepmd.dpmodel.train import ( + AbstractTrainer, + CheckpointObservation, + DisplayObservation, + RankContext, + StepObservation, + TrainerConfig, + TrainingObserver, + TrainingTask, + TrainingTaskCollection, + TrainStepResult, +) + + +class RecordingObserver(TrainingObserver): + def __init__(self, *, freq: int = 1) -> None: + self.freq = freq + self.events: list[str] = [] + + def wants_step(self, display_step: int) -> bool: + return display_step == 1 or ( + self.freq > 0 and display_step % self.freq == 0 + ) + + def on_train_begin(self, tasks, *, rank_context: RankContext) -> None: + self.events.append(f"begin:{rank_context.rank}") + + def on_step_end(self, observation: StepObservation) -> None: + self.events.append(f"step:{observation.display_step}:{observation.task_key}") + + def on_display(self, observation: DisplayObservation) -> None: + self.events.append( + f"display:{observation.display_step}:{observation.timing.steps}" + ) + + def on_checkpoint(self, observation: CheckpointObservation) -> None: + self.events.append(f"ckpt:{observation.display_step}") + + def on_train_end(self, tasks, *, rank_context: RankContext) -> None: + self.events.append(f"end:{rank_context.rank}") + + +class DummyData: + def __init__(self, values: list[float]) -> None: + self.values = values + self.index = 0 + + def get_batch(self) -> dict[str, float]: + value = self.values[self.index % len(self.values)] + self.index += 1 + return {"value": value} + + +class DummyTrainer(AbstractTrainer): + def train_step(self, task: TrainingTask, step: int) -> TrainStepResult: + batch = task.training_data.get_batch() + return TrainStepResult( + task_key=task.key, + step=step, + payload=batch, + train_results={"rmse": batch["value"]}, + ) + + def evaluate_training( + self, + task: TrainingTask, + step: int, + step_result: TrainStepResult | None, + ) -> dict[str, float]: + if step_result is None or step_result.task_key != task.key: + return {"rmse": 0.0} + return {"rmse": float(step_result.payload["value"])} + + def evaluate_validation( + self, + task: TrainingTask, + step: int, + step_result: TrainStepResult | None, + ) -> dict[str, float] | None: + if task.validation_data is None: + return None + return {"rmse": float(task.validation_data.get_batch()["value"])} + + def learning_rate(self, step: int) -> float: + return 0.1 / (step + 1) + + def save_checkpoint(self, step: int) -> None: + return None + + +def test_observer_hooks_follow_independent_frequencies(tmp_path: Path) -> None: + observer = RecordingObserver(freq=2) + trainer = DummyTrainer( + TrainerConfig( + num_steps=4, + disp_file=str(tmp_path / "lcurve.out"), + disp_freq=4, + save_freq=4, + timing_in_training=True, + ), + observers=observer, + ) + trainer.run( + TrainingTaskCollection.single(DummyData([1.0, 2.0, 3.0, 4.0]), DummyData([9.0])) + ) + + assert observer.events[0] == "begin:0" + assert "step:1:Default" in observer.events + assert "step:2:Default" in observer.events + assert "step:3:Default" not in observer.events + assert "step:4:Default" in observer.events + # Display still follows disp_freq (plus the opening step); TB step hooks + # follow tensorboard-style frequency independently. + assert "display:1:1" in observer.events + assert "display:4:3" in observer.events + assert "ckpt:4" in observer.events + assert observer.events[-1] == "end:0" + + +def test_observer_still_closes_after_training_failure(tmp_path: Path) -> None: + observer = RecordingObserver(freq=1) + + class FailingTrainer(DummyTrainer): + def train_step(self, task: TrainingTask, step: int) -> TrainStepResult: + raise RuntimeError("boom") + + trainer = FailingTrainer( + TrainerConfig( + num_steps=1, + disp_file=str(tmp_path / "lcurve.out"), + disp_freq=1, + save_freq=1, + ), + observers=observer, + ) + try: + trainer.run(TrainingTaskCollection.single(DummyData([1.0]))) + except RuntimeError: + pass + assert observer.events[0] == "begin:0" + assert observer.events[-1] == "end:0" + + +def test_step_observation_is_skipped_when_no_observer_wants_it( + tmp_path: Path, +) -> None: + calls: list[int] = [] + + class QuietObserver(TrainingObserver): + def wants_step(self, display_step: int) -> bool: + return False + + def on_step_end(self, observation: StepObservation) -> None: + calls.append(observation.display_step) + + trainer = DummyTrainer( + TrainerConfig( + num_steps=3, + disp_file=str(tmp_path / "lcurve.out"), + disp_freq=3, + save_freq=3, + display_in_training=False, + ), + observers=QuietObserver(), + ) + trainer.run(TrainingTaskCollection.single(DummyData([1.0, 2.0, 3.0]))) + assert calls == [] diff --git a/source/tests/pt_expt/train/test_tensorboard_observer.py b/source/tests/pt_expt/train/test_tensorboard_observer.py new file mode 100644 index 0000000000..f538be90ac --- /dev/null +++ b/source/tests/pt_expt/train/test_tensorboard_observer.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from __future__ import ( + annotations, +) + +import datetime +from typing import ( + TYPE_CHECKING, +) +from unittest.mock import ( + MagicMock, +) + +from deepmd.dpmodel.train.observer import ( + CheckpointObservation, + DisplayObservation, + StepObservation, +) +from deepmd.dpmodel.train.timing import ( + DisplayInterval, +) +from deepmd.dpmodel.train.trainer import ( + RankContext, + TrainingTaskCollection, + TrainStepResult, +) +from deepmd.pt_expt.train.tensorboard import ( + TensorBoardObserver, + create_tensorboard_observer, +) + +if TYPE_CHECKING: + from pathlib import ( + Path, + ) + + +def _interval(display_step: int = 2) -> DisplayInterval: + return DisplayInterval( + display_step=display_step, + wall_time=1.5, + steps=2, + eta=10, + timestamp=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + ) + + +def test_disabled_factory_returns_none_without_filesystem(tmp_path: Path) -> None: + observer = create_tensorboard_observer( + {"tensorboard": False, "tensorboard_log_dir": str(tmp_path / "log")}, + rank_context=RankContext(), + ) + assert observer is None + assert not (tmp_path / "log").exists() + + +def test_non_chief_never_creates_writer(tmp_path: Path) -> None: + writer_factory = MagicMock() + observer = TensorBoardObserver( + log_dir=str(tmp_path / "log"), + freq=1, + rank_context=RankContext(rank=1, world_size=2), + writer_factory=writer_factory, + ) + assert not observer.enabled + observer.on_train_begin( + TrainingTaskCollection.single(object()), + rank_context=RankContext(rank=1, world_size=2), + ) + writer_factory.assert_not_called() + assert list(tmp_path.iterdir()) == [] + + +def test_chief_writes_at_configured_frequency_and_closes(tmp_path: Path) -> None: + log_dir = tmp_path / "tb" + writer = MagicMock() + writer_factory = MagicMock(return_value=writer) + observer = TensorBoardObserver( + log_dir=str(log_dir), + freq=2, + rank_context=RankContext(), + multi_task=False, + writer_factory=writer_factory, + ) + observer.on_train_begin( + TrainingTaskCollection.single(object()), + rank_context=RankContext(), + ) + assert log_dir.is_dir() + writer_factory.assert_called_once_with(log_dir=str(log_dir)) + + for display_step, value in ((1, 1.0), (2, 2.0), (3, 3.0), (4, 4.0)): + if not observer.wants_step(display_step): + continue + observer.on_step_end( + StepObservation( + step=display_step - 1, + display_step=display_step, + task_key="Default", + learning_rate=0.01 * display_step, + step_result=TrainStepResult( + task_key="Default", + step=display_step - 1, + train_results={"rmse": value}, + ), + rank_context=RankContext(), + ) + ) + + observer.on_display( + DisplayObservation( + step=1, + display_step=2, + task_key="Default", + learning_rate=0.02, + train_results={"rmse": 2.0}, + valid_results={"rmse": 9.0}, + timing=_interval(2), + rank_context=RankContext(), + ) + ) + observer.on_checkpoint( + CheckpointObservation(display_step=2, rank_context=RankContext()) + ) + observer.on_train_end( + TrainingTaskCollection.single(object()), + rank_context=RankContext(), + ) + + scalar_tags = [call.args[0] for call in writer.add_scalar.call_args_list] + steps = [call.args[2] for call in writer.add_scalar.call_args_list] + assert "learning_rate" in scalar_tags + assert "train/rmse" in scalar_tags + assert "valid/rmse" in scalar_tags + assert "timing/interval_wall_time" in scalar_tags + assert 1 in steps and 2 in steps and 4 in steps + assert 3 not in steps + writer.flush.assert_called() + writer.close.assert_called_once() + assert observer.writer is None + + +def test_restart_keeps_existing_event_directory(tmp_path: Path) -> None: + log_dir = tmp_path / "tb" + log_dir.mkdir() + prior = log_dir / "events.out.tfevents.prior" + prior.write_text("keep-me", encoding="utf-8") + + writer = MagicMock() + writer_factory = MagicMock(return_value=writer) + observer = TensorBoardObserver( + log_dir=str(log_dir), + freq=1, + rank_context=RankContext(), + writer_factory=writer_factory, + ) + observer.on_train_begin( + TrainingTaskCollection.single(object()), + rank_context=RankContext(), + ) + writer_factory.assert_called_once_with(log_dir=str(log_dir)) + assert prior.read_text(encoding="utf-8") == "keep-me" + + +def test_multi_task_tags_include_task_key(tmp_path: Path) -> None: + writer = MagicMock() + observer = TensorBoardObserver( + log_dir=str(tmp_path / "tb"), + freq=1, + rank_context=RankContext(), + multi_task=True, + writer_factory=MagicMock(return_value=writer), + ) + observer._writer = writer + observer.on_step_end( + StepObservation( + step=0, + display_step=1, + task_key="water", + learning_rate=0.1, + step_result=TrainStepResult( + task_key="water", + step=0, + train_results={"rmse": 1.25}, + ), + rank_context=RankContext(), + ) + ) + tags = [call.args[0] for call in writer.add_scalar.call_args_list] + assert "train/water/rmse" in tags + + +def test_multi_process_only_chief_emits(tmp_path: Path) -> None: + """Simulate two ranks: only the chief installs an active writer.""" + chiefs = [] + writers = [] + for rank in (0, 1): + writer = MagicMock() + writers.append(writer) + observer = create_tensorboard_observer( + { + "tensorboard": True, + "tensorboard_log_dir": str(tmp_path / f"rank{rank}"), + "tensorboard_freq": 1, + }, + rank_context=RankContext(rank=rank, world_size=2), + ) + assert observer is not None + observer._writer_factory = MagicMock(return_value=writer) + observer.on_train_begin( + TrainingTaskCollection.single(object()), + rank_context=RankContext(rank=rank, world_size=2), + ) + if observer.wants_step(1): + observer.on_step_end( + StepObservation( + step=0, + display_step=1, + task_key="Default", + learning_rate=0.1, + step_result=TrainStepResult( + task_key="Default", + step=0, + train_results={"rmse": 1.0}, + ), + rank_context=RankContext(rank=rank, world_size=2), + ) + ) + observer.close() + chiefs.append(observer.enabled) + + assert chiefs == [True, False] + assert writers[0].add_scalar.called + assert not writers[1].add_scalar.called + assert not writers[1].close.called + + +def test_exceptional_close_is_idempotent(tmp_path: Path) -> None: + writer = MagicMock() + observer = TensorBoardObserver( + log_dir=str(tmp_path / "tb"), + freq=1, + rank_context=RankContext(), + writer_factory=MagicMock(return_value=writer), + ) + observer._writer = writer + observer.close() + observer.close() + writer.close.assert_called_once() From f72059105c7396e5f35e9aa05330f49eaa97081c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:03:34 +0000 Subject: [PATCH 02/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt_expt/train/tensorboard.py | 8 ++------ deepmd/utils/argcheck.py | 9 ++++++--- source/tests/common/dpmodel/test_train_observer.py | 4 +--- 3 files changed, 9 insertions(+), 12 deletions(-) diff --git a/deepmd/pt_expt/train/tensorboard.py b/deepmd/pt_expt/train/tensorboard.py index e080ddbec8..703f5cf684 100644 --- a/deepmd/pt_expt/train/tensorboard.py +++ b/deepmd/pt_expt/train/tensorboard.py @@ -213,9 +213,7 @@ def _write_train_from_step(self, observation: StepObservation) -> None: more_loss = step_result.payload.get("more_loss") if isinstance(more_loss, Mapping): metrics = { - key: value - for key, value in more_loss.items() - if "l2_" not in key + key: value for key, value in more_loss.items() if "l2_" not in key } if not metrics: return @@ -247,9 +245,7 @@ def _write_metric_tree( ) return if isinstance(results, Mapping): - tag_prefix = ( - f"{prefix}/{DEFAULT_TASK_KEY}" if self._multi_task else prefix - ) + tag_prefix = f"{prefix}/{DEFAULT_TASK_KEY}" if self._multi_task else prefix self._write_metric_mapping(tag_prefix, results, display_step) def _write_metric_mapping( diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index ff531df09f..48cf0da1c5 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -6253,21 +6253,24 @@ def training_args( bool, optional=True, default=False, - doc=supported_backends("tf", "pt", "pd", "tf2", "pt_expt") + doc_tensorboard, + doc=supported_backends("tf", "pt", "pd", "tf2", "pt_expt") + + doc_tensorboard, ), Argument( "tensorboard_log_dir", str, optional=True, default="log", - doc=supported_backends("tf", "pt", "pd", "tf2", "pt_expt") + doc_tensorboard_log_dir, + doc=supported_backends("tf", "pt", "pd", "tf2", "pt_expt") + + doc_tensorboard_log_dir, ), Argument( "tensorboard_freq", int, optional=True, default=1, - doc=supported_backends("tf", "pt", "pd", "tf2", "pt_expt") + doc_tensorboard_freq, + doc=supported_backends("tf", "pt", "pd", "tf2", "pt_expt") + + doc_tensorboard_freq, ), Argument( "gradient_max_norm", diff --git a/source/tests/common/dpmodel/test_train_observer.py b/source/tests/common/dpmodel/test_train_observer.py index 47a9c03fe3..16465e03b4 100644 --- a/source/tests/common/dpmodel/test_train_observer.py +++ b/source/tests/common/dpmodel/test_train_observer.py @@ -23,9 +23,7 @@ def __init__(self, *, freq: int = 1) -> None: self.events: list[str] = [] def wants_step(self, display_step: int) -> bool: - return display_step == 1 or ( - self.freq > 0 and display_step % self.freq == 0 - ) + return display_step == 1 or (self.freq > 0 and display_step % self.freq == 0) def on_train_begin(self, tasks, *, rank_context: RankContext) -> None: self.events.append(f"begin:{rank_context.rank}") From 0d5ddfd9c5afadc390a91ec1ca72513af7ff0a7b Mon Sep 17 00:00:00 2001 From: OutisLi-Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:06:07 +0800 Subject: [PATCH 03/15] feat(pt_expt): add TorchProfilerObserver for profiler lifecycle Implement training.enable_profiler / profiling / profiling_file through a dedicated TrainingObserver that owns start, per-step profiler.step(), Chrome/TensorBoard export, and release on completion or failure. Stack on the TrainingObserver interface from #6066. Disabled runs install no observer. Distributed Chrome paths and TensorBoard worker names use deterministic rank suffixes. Refs #5820 #5755 --- deepmd/pt_expt/train/profiler.py | 222 ++++++++++++++ deepmd/pt_expt/train/training.py | 24 +- deepmd/utils/argcheck.py | 6 +- .../pt_expt/train/test_profiler_observer.py | 284 ++++++++++++++++++ 4 files changed, 528 insertions(+), 8 deletions(-) create mode 100644 deepmd/pt_expt/train/profiler.py create mode 100644 source/tests/pt_expt/train/test_profiler_observer.py diff --git a/deepmd/pt_expt/train/profiler.py b/deepmd/pt_expt/train/profiler.py new file mode 100644 index 0000000000..02be0ec1d5 --- /dev/null +++ b/deepmd/pt_expt/train/profiler.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Torch profiler lifecycle for pt_expt training. + +Owns create/enter, per-step ``profiler.step()``, Chrome/TensorBoard export, and +release on completion or failure. Control stays outside the compiled model +graph. Disabled configurations install no observer. +""" + +from __future__ import ( + annotations, +) + +import logging +from collections.abc import ( + Callable, + Mapping, +) +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +from deepmd.dpmodel.train.observer import ( + StepObservation, + TrainingObserver, +) +from deepmd.dpmodel.train.trainer import ( + RankContext, + TrainingTaskCollection, +) + +log = logging.getLogger(__name__) + +__all__ = [ + "TorchProfilerObserver", + "create_profiler_observer", + "resolve_chrome_trace_path", +] + +# Match legacy PT so enable_profiler / profiling keep their established meaning. +_DEFAULT_SCHEDULE = {"wait": 1, "warmup": 15, "active": 3, "repeat": 1} + + +def resolve_chrome_trace_path( + path: str, + *, + rank: int, + world_size: int, +) -> str: + """Return a collision-free Chrome trace path for the calling rank. + + Single-process training keeps ``path`` unchanged so ``profiling_file`` is + honored literally. Distributed runs insert a deterministic ``.rankN`` + suffix before the extension. + """ + if world_size <= 1: + return path + target = Path(path) + return str(target.with_name(f"{target.stem}.rank{rank}{target.suffix}")) + + +def create_profiler_observer( + training_params: Mapping[str, Any], + *, + rank_context: RankContext, +) -> TorchProfilerObserver | None: + """Return a profiler observer when enabled; otherwise ``None``. + + A disabled configuration returns ``None`` so the trainer installs no + observer and never constructs a profiler session. + """ + enable_profiler = bool(training_params.get("enable_profiler", False)) + profiling = bool(training_params.get("profiling", False)) + if not enable_profiler and not profiling: + return None + return TorchProfilerObserver( + enable_profiler=enable_profiler, + profiling=profiling, + profiling_file=str(training_params.get("profiling_file", "timeline.json")), + tensorboard_log_dir=str(training_params.get("tensorboard_log_dir", "log")), + rank_context=rank_context, + ) + + +class TorchProfilerObserver(TrainingObserver): + """Drive one Torch profiler session across the training observer hooks.""" + + def __init__( + self, + *, + enable_profiler: bool, + profiling: bool, + profiling_file: str, + tensorboard_log_dir: str, + rank_context: RankContext, + profiler_factory: Callable[..., Any] | None = None, + schedule_factory: Callable[..., Any] | None = None, + tensorboard_handler_factory: Callable[..., Any] | None = None, + ) -> None: + self._enable_profiler = bool(enable_profiler) + self._profiling = bool(profiling) + self._profiling_file = str(profiling_file) + self._tensorboard_log_dir = str(tensorboard_log_dir) + self._rank_context = rank_context + self._profiler_factory = profiler_factory + self._schedule_factory = schedule_factory + self._tensorboard_handler_factory = tensorboard_handler_factory + self._profiler: Any | None = None + self._closed = False + self._chrome_trace_path = resolve_chrome_trace_path( + self._profiling_file, + rank=rank_context.rank, + world_size=rank_context.world_size, + ) + + @property + def enabled(self) -> bool: + """Whether this observer owns an active profiling session.""" + return self._enable_profiler or self._profiling + + @property + def chrome_trace_path(self) -> str: + """Resolved Chrome JSON path for this rank.""" + return self._chrome_trace_path + + @property + def profiler(self) -> Any | None: + return self._profiler + + def wants_step(self, display_step: int) -> bool: + # Every completed optimizer step must call profiler.step() once. + return self.enabled + + def on_train_begin( + self, + tasks: TrainingTaskCollection, + *, + rank_context: RankContext, + ) -> None: + if not self.enabled: + return + profiler_factory = self._profiler_factory + schedule_factory = self._schedule_factory + handler_factory = self._tensorboard_handler_factory + if profiler_factory is None or schedule_factory is None: + import torch.profiler as torch_profiler + + if profiler_factory is None: + profiler_factory = torch_profiler.profile + if schedule_factory is None: + schedule_factory = torch_profiler.schedule + if handler_factory is None: + handler_factory = torch_profiler.tensorboard_trace_handler + + on_trace_ready = None + if self._enable_profiler: + assert handler_factory is not None + Path(self._tensorboard_log_dir).mkdir(parents=True, exist_ok=True) + worker_name = None + if self._rank_context.world_size > 1: + worker_name = f"rank{self._rank_context.rank}" + on_trace_ready = handler_factory( + self._tensorboard_log_dir, + worker_name=worker_name, + ) + + self._profiler = profiler_factory( + schedule=schedule_factory(**_DEFAULT_SCHEDULE), + on_trace_ready=on_trace_ready, + record_shapes=True, + with_stack=True, + ) + self._profiler.start() + self._closed = False + log.info( + "Torch profiler started (enable_profiler=%s, profiling=%s)", + self._enable_profiler, + self._profiling, + ) + + def on_step_end(self, observation: StepObservation) -> None: + profiler = self._profiler + if profiler is None: + return + profiler.step() + + def on_train_end( + self, + tasks: TrainingTaskCollection, + *, + rank_context: RankContext, + ) -> None: + self.close() + + def close(self) -> None: + """Stop the profiler, export Chrome traces when requested, and release.""" + profiler = self._profiler + if profiler is None or self._closed: + self._profiler = None + self._closed = True + return + try: + profiler.stop() + if self._enable_profiler: + log.info( + "Profiler TensorBoard traces saved under %s", + self._tensorboard_log_dir, + ) + if self._profiling: + parent = Path(self._chrome_trace_path).parent + if str(parent) not in ("", "."): + parent.mkdir(parents=True, exist_ok=True) + profiler.export_chrome_trace(self._chrome_trace_path) + log.info( + "Profiler Chrome trace saved to: %s", + self._chrome_trace_path, + ) + finally: + self._profiler = None + self._closed = True diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index a3fce537dc..e1ede0cc3f 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -120,6 +120,9 @@ NonFiniteGradGuard, clip_grad_norm_, ) +from deepmd.pt_expt.train.profiler import ( + create_profiler_observer, +) from deepmd.pt_expt.train.tensorboard import ( create_tensorboard_observer, ) @@ -2406,11 +2409,22 @@ def update_finetune_bias( if training_params.get("disp_avg", False) else None ), - observers=create_tensorboard_observer( - training_params, - rank_context=rank_context, - multi_task=self.multi_task, - ), + observers=[ + observer + for observer in ( + create_tensorboard_observer( + training_params, + rank_context=rank_context, + multi_task=self.multi_task, + ), + create_profiler_observer( + training_params, + rank_context=rank_context, + ), + ) + if observer is not None + ] + or None, ) self.full_validator, self.ema_full_validator = self._create_full_validators( validating_params=validating_params, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 48cf0da1c5..6bdbbaa002 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -6232,21 +6232,21 @@ def training_args( bool, optional=True, default=False, - doc=supported_backends("tf", "pt", "pd") + doc_profiling, + doc=supported_backends("tf", "pt", "pd", "pt_expt") + doc_profiling, ), Argument( "profiling_file", str, optional=True, default="timeline.json", - doc=supported_backends("tf", "pt") + doc_profiling_file, + doc=supported_backends("tf", "pt", "pt_expt") + doc_profiling_file, ), Argument( "enable_profiler", bool, optional=True, default=False, - doc=supported_backends("tf", "pt", "pd") + doc_enable_profiler, + doc=supported_backends("tf", "pt", "pd", "pt_expt") + doc_enable_profiler, ), Argument( "tensorboard", diff --git a/source/tests/pt_expt/train/test_profiler_observer.py b/source/tests/pt_expt/train/test_profiler_observer.py new file mode 100644 index 0000000000..8cb37c067a --- /dev/null +++ b/source/tests/pt_expt/train/test_profiler_observer.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, +) +from unittest.mock import ( + MagicMock, +) + +from deepmd.dpmodel.train.observer import ( + StepObservation, +) +from deepmd.dpmodel.train.trainer import ( + RankContext, + TrainingTaskCollection, + TrainStepResult, +) +from deepmd.pt_expt.train.profiler import ( + TorchProfilerObserver, + create_profiler_observer, + resolve_chrome_trace_path, +) + +if TYPE_CHECKING: + from pathlib import ( + Path, + ) + + +def _step(display_step: int, *, rank_context: RankContext | None = None) -> StepObservation: + ctx = rank_context or RankContext() + return StepObservation( + step=display_step - 1, + display_step=display_step, + task_key="Default", + learning_rate=0.1, + step_result=TrainStepResult( + task_key="Default", + step=display_step - 1, + train_results={"rmse": 1.0}, + ), + rank_context=ctx, + ) + + +def test_resolve_chrome_trace_path_single_process() -> None: + assert ( + resolve_chrome_trace_path("timeline.json", rank=0, world_size=1) + == "timeline.json" + ) + assert ( + resolve_chrome_trace_path("out/trace.json", rank=3, world_size=1) + == "out/trace.json" + ) + + +def test_resolve_chrome_trace_path_distributed_rank_suffix() -> None: + assert ( + resolve_chrome_trace_path("timeline.json", rank=0, world_size=4) + == "timeline.rank0.json" + ) + assert ( + resolve_chrome_trace_path("dir/chrome.json", rank=2, world_size=4) + == "dir/chrome.rank2.json" + ) + + +def test_disabled_factory_returns_none_without_profiler(tmp_path: Path) -> None: + observer = create_profiler_observer( + { + "enable_profiler": False, + "profiling": False, + "profiling_file": str(tmp_path / "timeline.json"), + "tensorboard_log_dir": str(tmp_path / "log"), + }, + rank_context=RankContext(), + ) + assert observer is None + assert list(tmp_path.iterdir()) == [] + + +def test_profiling_only_exports_chrome_and_steps_once(tmp_path: Path) -> None: + chrome = tmp_path / "custom_timeline.json" + profiler = MagicMock() + schedule = MagicMock(return_value="sched") + handler = MagicMock() + factory = MagicMock(return_value=profiler) + + observer = TorchProfilerObserver( + enable_profiler=False, + profiling=True, + profiling_file=str(chrome), + tensorboard_log_dir=str(tmp_path / "log"), + rank_context=RankContext(), + profiler_factory=factory, + schedule_factory=schedule, + tensorboard_handler_factory=handler, + ) + assert observer.wants_step(1) + assert observer.wants_step(17) + + observer.on_train_begin( + TrainingTaskCollection.single(object()), + rank_context=RankContext(), + ) + schedule.assert_called_once_with(wait=1, warmup=15, active=3, repeat=1) + factory.assert_called_once() + kwargs = factory.call_args.kwargs + assert kwargs["on_trace_ready"] is None + assert kwargs["schedule"] == "sched" + profiler.start.assert_called_once() + handler.assert_not_called() + assert not (tmp_path / "log").exists() + + for display_step in (1, 2, 3): + observer.on_step_end(_step(display_step)) + assert profiler.step.call_count == 3 + + observer.on_train_end( + TrainingTaskCollection.single(object()), + rank_context=RankContext(), + ) + profiler.stop.assert_called_once() + profiler.export_chrome_trace.assert_called_once_with(str(chrome)) + assert observer.profiler is None + + +def test_enable_profiler_only_uses_tensorboard_handler(tmp_path: Path) -> None: + log_dir = tmp_path / "tb" + profiler = MagicMock() + schedule = MagicMock(return_value="sched") + ready = object() + handler = MagicMock(return_value=ready) + factory = MagicMock(return_value=profiler) + + observer = TorchProfilerObserver( + enable_profiler=True, + profiling=False, + profiling_file=str(tmp_path / "timeline.json"), + tensorboard_log_dir=str(log_dir), + rank_context=RankContext(), + profiler_factory=factory, + schedule_factory=schedule, + tensorboard_handler_factory=handler, + ) + observer.on_train_begin( + TrainingTaskCollection.single(object()), + rank_context=RankContext(), + ) + assert log_dir.is_dir() + handler.assert_called_once_with(str(log_dir), worker_name=None) + assert factory.call_args.kwargs["on_trace_ready"] is ready + + observer.on_step_end(_step(1)) + observer.close() + profiler.stop.assert_called_once() + profiler.export_chrome_trace.assert_not_called() + + +def test_combined_modes_one_session_both_outputs(tmp_path: Path) -> None: + log_dir = tmp_path / "tb" + chrome = tmp_path / "both.json" + profiler = MagicMock() + schedule = MagicMock(return_value="sched") + ready = object() + handler = MagicMock(return_value=ready) + factory = MagicMock(return_value=profiler) + + observer = create_profiler_observer( + { + "enable_profiler": True, + "profiling": True, + "profiling_file": str(chrome), + "tensorboard_log_dir": str(log_dir), + }, + rank_context=RankContext(), + ) + assert observer is not None + observer._profiler_factory = factory + observer._schedule_factory = schedule + observer._tensorboard_handler_factory = handler + + observer.on_train_begin( + TrainingTaskCollection.single(object()), + rank_context=RankContext(), + ) + assert factory.call_count == 1 + assert factory.call_args.kwargs["on_trace_ready"] is ready + observer.on_step_end(_step(1)) + observer.on_step_end(_step(2)) + observer.close() + + assert profiler.step.call_count == 2 + profiler.export_chrome_trace.assert_called_once_with(str(chrome)) + profiler.stop.assert_called_once() + + +def test_distributed_rank_suffix_and_worker_name(tmp_path: Path) -> None: + log_dir = tmp_path / "tb" + chrome = tmp_path / "timeline.json" + profiler = MagicMock() + schedule = MagicMock(return_value="sched") + handler = MagicMock(return_value=object()) + factory = MagicMock(return_value=profiler) + ctx = RankContext(rank=1, world_size=4) + + observer = TorchProfilerObserver( + enable_profiler=True, + profiling=True, + profiling_file=str(chrome), + tensorboard_log_dir=str(log_dir), + rank_context=ctx, + profiler_factory=factory, + schedule_factory=schedule, + tensorboard_handler_factory=handler, + ) + assert observer.chrome_trace_path == str(tmp_path / "timeline.rank1.json") + + observer.on_train_begin( + TrainingTaskCollection.single(object()), + rank_context=ctx, + ) + handler.assert_called_once_with(str(log_dir), worker_name="rank1") + observer.on_step_end(_step(1, rank_context=ctx)) + observer.close() + profiler.export_chrome_trace.assert_called_once_with( + str(tmp_path / "timeline.rank1.json") + ) + + +def test_exceptional_exit_still_releases(tmp_path: Path) -> None: + profiler = MagicMock() + observer = TorchProfilerObserver( + enable_profiler=False, + profiling=True, + profiling_file=str(tmp_path / "timeline.json"), + tensorboard_log_dir=str(tmp_path / "log"), + rank_context=RankContext(), + profiler_factory=MagicMock(return_value=profiler), + schedule_factory=MagicMock(return_value="sched"), + tensorboard_handler_factory=MagicMock(), + ) + observer.on_train_begin( + TrainingTaskCollection.single(object()), + rank_context=RankContext(), + ) + # Simulate exceptional exit via the same finally path AbstractTrainer uses. + observer.on_train_end( + TrainingTaskCollection.single(object()), + rank_context=RankContext(), + ) + profiler.stop.assert_called_once() + assert observer.profiler is None + # Second close is a no-op. + observer.close() + assert profiler.stop.call_count == 1 + + +def test_close_releases_even_when_export_fails(tmp_path: Path) -> None: + profiler = MagicMock() + profiler.export_chrome_trace.side_effect = RuntimeError("disk full") + observer = TorchProfilerObserver( + enable_profiler=False, + profiling=True, + profiling_file=str(tmp_path / "timeline.json"), + tensorboard_log_dir=str(tmp_path / "log"), + rank_context=RankContext(), + profiler_factory=MagicMock(return_value=profiler), + schedule_factory=MagicMock(return_value="sched"), + tensorboard_handler_factory=MagicMock(), + ) + observer.on_train_begin( + TrainingTaskCollection.single(object()), + rank_context=RankContext(), + ) + try: + observer.close() + except RuntimeError: + pass + assert observer.profiler is None + profiler.stop.assert_called_once() From 67cda60369473436251c7ca87c0a947cbd91b35f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:07:53 +0000 Subject: [PATCH 04/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/pt_expt/train/test_profiler_observer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/source/tests/pt_expt/train/test_profiler_observer.py b/source/tests/pt_expt/train/test_profiler_observer.py index 8cb37c067a..0574f3030a 100644 --- a/source/tests/pt_expt/train/test_profiler_observer.py +++ b/source/tests/pt_expt/train/test_profiler_observer.py @@ -30,7 +30,9 @@ ) -def _step(display_step: int, *, rank_context: RankContext | None = None) -> StepObservation: +def _step( + display_step: int, *, rank_context: RankContext | None = None +) -> StepObservation: ctx = rank_context or RankContext() return StepObservation( step=display_step - 1, From 01918f621140548b4f89f085b042431ee1a6dffd Mon Sep 17 00:00:00 2001 From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:08:43 +0800 Subject: [PATCH 05/15] fix(pt_expt): keep step-owned train tags on coinciding display MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit At disp ∩ tensorboard_freq, on_display used to rewrite learning_rate and train/* from DisplayObservation, which may carry interval averages when disp_avg=True. Leave those tags to on_step_end; on_display only emits valid/* and timing/*. --- deepmd/pt_expt/train/tensorboard.py | 15 +---- .../train/test_tensorboard_observer.py | 66 +++++++++++++++++++ 2 files changed, 69 insertions(+), 12 deletions(-) diff --git a/deepmd/pt_expt/train/tensorboard.py b/deepmd/pt_expt/train/tensorboard.py index 703f5cf684..f928a32e6e 100644 --- a/deepmd/pt_expt/train/tensorboard.py +++ b/deepmd/pt_expt/train/tensorboard.py @@ -139,18 +139,9 @@ def on_display(self, observation: DisplayObservation) -> None: if writer is None or not self._is_due(observation.display_step): return # Validation and timing exist only at display boundaries. Emit them - # when that boundary also lands on tensorboard_freq so every tag - # family shares the configured cadence. - writer.add_scalar( - "learning_rate", - float(observation.learning_rate), - observation.display_step, - ) - self._write_metric_tree( - "train", - observation.train_results, - observation.display_step, - ) + # when that boundary also lands on tensorboard_freq. + # learning_rate and train/* stay owned by on_step_end so a coinciding + # disp_freq step cannot overwrite per-step points with interval averages. self._write_metric_tree( "valid", observation.valid_results, diff --git a/source/tests/pt_expt/train/test_tensorboard_observer.py b/source/tests/pt_expt/train/test_tensorboard_observer.py index f538be90ac..c882cf4bff 100644 --- a/source/tests/pt_expt/train/test_tensorboard_observer.py +++ b/source/tests/pt_expt/train/test_tensorboard_observer.py @@ -247,3 +247,69 @@ def test_exceptional_close_is_idempotent(tmp_path: Path) -> None: observer.close() observer.close() writer.close.assert_called_once() + +def test_coinciding_display_keeps_step_train_metrics(tmp_path: Path) -> None: + """At disp ∩ tensorboard_freq, train/* is written once from the step hook. + + DisplayObservations may carry interval averages (disp_avg=True). Those must + not overwrite the per-step train/* / learning_rate points already emitted by + on_step_end under the same tags and step. + """ + writer = MagicMock() + observer = TensorBoardObserver( + log_dir=str(tmp_path / "tb"), + freq=2, + rank_context=RankContext(), + multi_task=False, + writer_factory=MagicMock(return_value=writer), + ) + observer._writer = writer + + step_rmse = 1.5 + step_lr = 0.02 + observer.on_step_end( + StepObservation( + step=1, + display_step=2, + task_key="Default", + learning_rate=step_lr, + step_result=TrainStepResult( + task_key="Default", + step=1, + train_results={"rmse": step_rmse}, + ), + rank_context=RankContext(), + ) + ) + observer.on_display( + DisplayObservation( + step=1, + display_step=2, + task_key="Default", + learning_rate=0.99, + train_results={"rmse": 9.9}, + valid_results={"rmse": 3.0}, + timing=_interval(2), + rank_context=RankContext(), + ) + ) + + train_calls = [ + call + for call in writer.add_scalar.call_args_list + if call.args[0] == "train/rmse" + ] + lr_calls = [ + call + for call in writer.add_scalar.call_args_list + if call.args[0] == "learning_rate" + ] + assert len(train_calls) == 1 + assert train_calls[0].args[1] == step_rmse + assert train_calls[0].args[2] == 2 + assert len(lr_calls) == 1 + assert lr_calls[0].args[1] == step_lr + tags = [call.args[0] for call in writer.add_scalar.call_args_list] + assert "valid/rmse" in tags + assert "timing/interval_wall_time" in tags + From 8393ffe00fa6e83780ab7ea99e379a258ed3176a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:10:48 +0000 Subject: [PATCH 06/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/pt_expt/train/test_tensorboard_observer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/pt_expt/train/test_tensorboard_observer.py b/source/tests/pt_expt/train/test_tensorboard_observer.py index c882cf4bff..fc1d924714 100644 --- a/source/tests/pt_expt/train/test_tensorboard_observer.py +++ b/source/tests/pt_expt/train/test_tensorboard_observer.py @@ -248,6 +248,7 @@ def test_exceptional_close_is_idempotent(tmp_path: Path) -> None: observer.close() writer.close.assert_called_once() + def test_coinciding_display_keeps_step_train_metrics(tmp_path: Path) -> None: """At disp ∩ tensorboard_freq, train/* is written once from the step hook. @@ -312,4 +313,3 @@ def test_coinciding_display_keeps_step_train_metrics(tmp_path: Path) -> None: tags = [call.args[0] for call in writer.add_scalar.call_args_list] assert "valid/rmse" in tags assert "timing/interval_wall_time" in tags - From 18a708395dea35a8392bb5302d583daa3ec81a97 Mon Sep 17 00:00:00 2001 From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:12:42 +0800 Subject: [PATCH 07/15] fix(pt_expt): export combined profiler traces once for both sinks When enable_profiler and profiling are both set, tensorboard_trace_handler already calls export_chrome_trace via on_trace_ready. close() must not export again or Kineto raises "Trace is already saved". Use a custom on_trace_ready that exports once, then copies into the TensorBoard log dir and the rank-resolved profiling_file. close() only Chrome-exports when that ready-handler save never ran (short runs). Profiling-only keeps end-of-run export; enable-profiler-only stays TB-only. --- deepmd/pt_expt/train/profiler.py | 119 +++++++++++++--- .../pt_expt/train/test_profiler_observer.py | 127 +++++++++++++++--- 2 files changed, 214 insertions(+), 32 deletions(-) diff --git a/deepmd/pt_expt/train/profiler.py b/deepmd/pt_expt/train/profiler.py index 02be0ec1d5..16e606a4d1 100644 --- a/deepmd/pt_expt/train/profiler.py +++ b/deepmd/pt_expt/train/profiler.py @@ -4,6 +4,14 @@ Owns create/enter, per-step ``profiler.step()``, Chrome/TensorBoard export, and release on completion or failure. Control stays outside the compiled model graph. Disabled configurations install no observer. + +When both ``enable_profiler`` and ``profiling`` are set, one session serves both +sinks: a custom ``on_trace_ready`` exports the Kineto trace **once**, then places +the file into the TensorBoard log dir and the rank-resolved ``profiling_file``. +``close()`` only performs a Chrome export if that ready-handler save never ran +(short runs that never reach ``RECORD_AND_SAVE``). Profiling-only keeps +end-of-run Chrome export with ``on_trace_ready=None``; enable-profiler-only uses +the TensorBoard handler and never Chrome-exports at close. """ from __future__ import ( @@ -11,6 +19,10 @@ ) import logging +import os +import shutil +import socket +import time from collections.abc import ( Callable, Mapping, @@ -109,6 +121,9 @@ def __init__( self._tensorboard_handler_factory = tensorboard_handler_factory self._profiler: Any | None = None self._closed = False + # True after a ready-handler (or short-run close fallback) has already + # called export_chrome_trace once for this session. + self._trace_saved = False self._chrome_trace_path = resolve_chrome_trace_path( self._profiling_file, rank=rank_context.rank, @@ -129,6 +144,11 @@ def chrome_trace_path(self) -> str: def profiler(self) -> Any | None: return self._profiler + @property + def trace_saved(self) -> bool: + """Whether a Chrome/Kineto export has already completed this session.""" + return self._trace_saved + def wants_step(self, display_step: int) -> bool: # Every completed optimizer step must call profiler.step() once. return self.enabled @@ -154,17 +174,7 @@ def on_train_begin( if handler_factory is None: handler_factory = torch_profiler.tensorboard_trace_handler - on_trace_ready = None - if self._enable_profiler: - assert handler_factory is not None - Path(self._tensorboard_log_dir).mkdir(parents=True, exist_ok=True) - worker_name = None - if self._rank_context.world_size > 1: - worker_name = f"rank{self._rank_context.rank}" - on_trace_ready = handler_factory( - self._tensorboard_log_dir, - worker_name=worker_name, - ) + on_trace_ready = self._build_on_trace_ready(handler_factory) self._profiler = profiler_factory( schedule=schedule_factory(**_DEFAULT_SCHEDULE), @@ -174,6 +184,7 @@ def on_train_begin( ) self._profiler.start() self._closed = False + self._trace_saved = False log.info( "Torch profiler started (enable_profiler=%s, profiling=%s)", self._enable_profiler, @@ -195,28 +206,102 @@ def on_train_end( self.close() def close(self) -> None: - """Stop the profiler, export Chrome traces when requested, and release.""" + """Stop the profiler, export Chrome traces when still needed, and release.""" profiler = self._profiler if profiler is None or self._closed: self._profiler = None self._closed = True return try: + # stop() may invoke on_trace_ready for a pending RECORD_AND_SAVE cycle. profiler.stop() - if self._enable_profiler: + if self._enable_profiler and not self._profiling: log.info( "Profiler TensorBoard traces saved under %s", self._tensorboard_log_dir, ) - if self._profiling: - parent = Path(self._chrome_trace_path).parent - if str(parent) not in ("", "."): - parent.mkdir(parents=True, exist_ok=True) + elif self._profiling and not self._enable_profiler: + # Profiling-only: always end-of-run Chrome export. + self._ensure_parent_dir(self._chrome_trace_path) profiler.export_chrome_trace(self._chrome_trace_path) + self._trace_saved = True log.info( "Profiler Chrome trace saved to: %s", self._chrome_trace_path, ) + elif self._enable_profiler and self._profiling: + if not self._trace_saved: + # Short run: schedule never fired ready-handler; one save, both sinks. + self._export_combined_trace(profiler) + else: + log.info( + "Profiler traces placed under %s and at %s", + self._tensorboard_log_dir, + self._chrome_trace_path, + ) finally: self._profiler = None self._closed = True + + def _build_on_trace_ready( + self, + handler_factory: Callable[..., Any] | None, + ) -> Callable[[Any], None] | None: + """Choose the ready callback for the configured sinks. + + - enable_profiler only → stock TensorBoard handler + - profiling only → ``None`` (export at close) + - combined → export once, then place into TB dir and profiling_file + """ + if self._enable_profiler and self._profiling: + Path(self._tensorboard_log_dir).mkdir(parents=True, exist_ok=True) + return self._combined_on_trace_ready + if self._enable_profiler: + assert handler_factory is not None + Path(self._tensorboard_log_dir).mkdir(parents=True, exist_ok=True) + return handler_factory( + self._tensorboard_log_dir, + worker_name=self._tensorboard_worker_name(), + ) + return None + + def _combined_on_trace_ready(self, prof: Any) -> None: + """Save the Kineto trace once and fan out to both configured sinks.""" + self._export_combined_trace(prof) + + def _export_combined_trace(self, prof: Any) -> None: + """Call ``export_chrome_trace`` once; copy into TB log dir and profiling_file.""" + chrome_path = self._chrome_trace_path + self._ensure_parent_dir(chrome_path) + Path(self._tensorboard_log_dir).mkdir(parents=True, exist_ok=True) + + # Kineto allows only one export per saved cycle. Write the user-facing + # profiling_file first, then copy into the TensorBoard log directory + # with the same naming convention as tensorboard_trace_handler. + prof.export_chrome_trace(chrome_path) + tb_path = self._tensorboard_trace_path() + shutil.copy2(chrome_path, tb_path) + self._trace_saved = True + log.info( + "Profiler Chrome trace saved to: %s (TensorBoard copy: %s)", + chrome_path, + tb_path, + ) + + def _tensorboard_worker_name(self) -> str | None: + if self._rank_context.world_size > 1: + return f"rank{self._rank_context.rank}" + return None + + def _tensorboard_trace_path(self) -> str: + worker = self._tensorboard_worker_name() + if not worker: + worker = f"{socket.gethostname()}_{os.getpid()}" + file_name = f"{worker}.{time.time_ns()}.pt.trace.json" + return str(Path(self._tensorboard_log_dir) / file_name) + + @staticmethod + def _ensure_parent_dir(path: str) -> None: + parent = Path(path).parent + if str(parent) not in ("", "."): + parent.mkdir(parents=True, exist_ok=True) diff --git a/source/tests/pt_expt/train/test_profiler_observer.py b/source/tests/pt_expt/train/test_profiler_observer.py index 0574f3030a..d28eb77aee 100644 --- a/source/tests/pt_expt/train/test_profiler_observer.py +++ b/source/tests/pt_expt/train/test_profiler_observer.py @@ -3,8 +3,8 @@ annotations, ) -from typing import ( - TYPE_CHECKING, +from pathlib import ( + Path, ) from unittest.mock import ( MagicMock, @@ -24,15 +24,8 @@ resolve_chrome_trace_path, ) -if TYPE_CHECKING: - from pathlib import ( - Path, - ) - -def _step( - display_step: int, *, rank_context: RankContext | None = None -) -> StepObservation: +def _step(display_step: int, *, rank_context: RankContext | None = None) -> StepObservation: ctx = rank_context or RankContext() return StepObservation( step=display_step - 1, @@ -162,15 +155,20 @@ def test_enable_profiler_only_uses_tensorboard_handler(tmp_path: Path) -> None: profiler.export_chrome_trace.assert_not_called() -def test_combined_modes_one_session_both_outputs(tmp_path: Path) -> None: +def test_combined_short_run_exports_once_to_both_sinks_at_close(tmp_path: Path) -> None: + """Schedule never fires ready-handler; close() does one export for both sinks.""" log_dir = tmp_path / "tb" chrome = tmp_path / "both.json" profiler = MagicMock() schedule = MagicMock(return_value="sched") - ready = object() - handler = MagicMock(return_value=ready) + handler = MagicMock() factory = MagicMock(return_value=profiler) + def _export(path: str) -> None: + Path(path).write_text("{}", encoding="utf-8") + + profiler.export_chrome_trace.side_effect = _export + observer = create_profiler_observer( { "enable_profiler": True, @@ -190,7 +188,10 @@ def test_combined_modes_one_session_both_outputs(tmp_path: Path) -> None: rank_context=RankContext(), ) assert factory.call_count == 1 - assert factory.call_args.kwargs["on_trace_ready"] is ready + on_ready = factory.call_args.kwargs["on_trace_ready"] + assert callable(on_ready) + # Combined mode must not install the stock TB handler (would double-save). + handler.assert_not_called() observer.on_step_end(_step(1)) observer.on_step_end(_step(2)) observer.close() @@ -198,6 +199,65 @@ def test_combined_modes_one_session_both_outputs(tmp_path: Path) -> None: assert profiler.step.call_count == 2 profiler.export_chrome_trace.assert_called_once_with(str(chrome)) profiler.stop.assert_called_once() + assert chrome.is_file() + tb_traces = list(log_dir.glob("*.pt.trace.json")) + assert len(tb_traces) == 1 + assert tb_traces[0].read_text(encoding="utf-8") == "{}" + assert observer.trace_saved + + +def test_combined_ready_handler_single_save_no_double_export_at_close( + tmp_path: Path, +) -> None: + """Ready-handler already saved: close() must not call export_chrome_trace again. + + Mimics Kineto's destructive second export (RuntimeError: Trace is already saved). + """ + log_dir = tmp_path / "tb" + chrome = tmp_path / "both.json" + export_count = 0 + + def _export(path: str) -> None: + nonlocal export_count + export_count += 1 + if export_count > 1: + raise RuntimeError("Trace is already saved") + Path(path).write_text("{}", encoding="utf-8") + + profiler = MagicMock() + profiler.export_chrome_trace.side_effect = _export + schedule = MagicMock(return_value="sched") + factory = MagicMock(return_value=profiler) + + observer = TorchProfilerObserver( + enable_profiler=True, + profiling=True, + profiling_file=str(chrome), + tensorboard_log_dir=str(log_dir), + rank_context=RankContext(), + profiler_factory=factory, + schedule_factory=schedule, + tensorboard_handler_factory=MagicMock(), + ) + observer.on_train_begin( + TrainingTaskCollection.single(object()), + rank_context=RankContext(), + ) + on_ready = factory.call_args.kwargs["on_trace_ready"] + assert callable(on_ready) + + # Simulate schedule completing an active window (RECORD_AND_SAVE). + on_ready(profiler) + assert export_count == 1 + assert chrome.is_file() + assert list(log_dir.glob("*.pt.trace.json")) + assert observer.trace_saved + + # close() must not attempt a second Kineto export. + observer.close() + assert export_count == 1 + profiler.stop.assert_called_once() + assert observer.profiler is None def test_distributed_rank_suffix_and_worker_name(tmp_path: Path) -> None: @@ -209,6 +269,11 @@ def test_distributed_rank_suffix_and_worker_name(tmp_path: Path) -> None: factory = MagicMock(return_value=profiler) ctx = RankContext(rank=1, world_size=4) + def _export(path: str) -> None: + Path(path).write_text("{}", encoding="utf-8") + + profiler.export_chrome_trace.side_effect = _export + observer = TorchProfilerObserver( enable_profiler=True, profiling=True, @@ -225,12 +290,44 @@ def test_distributed_rank_suffix_and_worker_name(tmp_path: Path) -> None: TrainingTaskCollection.single(object()), rank_context=ctx, ) - handler.assert_called_once_with(str(log_dir), worker_name="rank1") + # Combined uses custom ready handler; stock TB factory unused. + handler.assert_not_called() + on_ready = factory.call_args.kwargs["on_trace_ready"] + on_ready(profiler) observer.on_step_end(_step(1, rank_context=ctx)) observer.close() + profiler.export_chrome_trace.assert_called_once_with( str(tmp_path / "timeline.rank1.json") ) + tb_traces = list(log_dir.glob("rank1.*.pt.trace.json")) + assert len(tb_traces) == 1 + + +def test_enable_profiler_distributed_passes_worker_name(tmp_path: Path) -> None: + log_dir = tmp_path / "tb" + profiler = MagicMock() + handler = MagicMock(return_value=object()) + factory = MagicMock(return_value=profiler) + ctx = RankContext(rank=2, world_size=4) + + observer = TorchProfilerObserver( + enable_profiler=True, + profiling=False, + profiling_file=str(tmp_path / "timeline.json"), + tensorboard_log_dir=str(log_dir), + rank_context=ctx, + profiler_factory=factory, + schedule_factory=MagicMock(return_value="sched"), + tensorboard_handler_factory=handler, + ) + observer.on_train_begin( + TrainingTaskCollection.single(object()), + rank_context=ctx, + ) + handler.assert_called_once_with(str(log_dir), worker_name="rank2") + observer.close() + profiler.export_chrome_trace.assert_not_called() def test_exceptional_exit_still_releases(tmp_path: Path) -> None: From d7a006eb2ae2beca6725373b0127b44fdde165dd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:13:41 +0000 Subject: [PATCH 08/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/pt_expt/train/test_profiler_observer.py | 4 +++- source/tests/pt_expt/train/test_tensorboard_observer.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/source/tests/pt_expt/train/test_profiler_observer.py b/source/tests/pt_expt/train/test_profiler_observer.py index d28eb77aee..58157d4406 100644 --- a/source/tests/pt_expt/train/test_profiler_observer.py +++ b/source/tests/pt_expt/train/test_profiler_observer.py @@ -25,7 +25,9 @@ ) -def _step(display_step: int, *, rank_context: RankContext | None = None) -> StepObservation: +def _step( + display_step: int, *, rank_context: RankContext | None = None +) -> StepObservation: ctx = rank_context or RankContext() return StepObservation( step=display_step - 1, diff --git a/source/tests/pt_expt/train/test_tensorboard_observer.py b/source/tests/pt_expt/train/test_tensorboard_observer.py index c882cf4bff..fc1d924714 100644 --- a/source/tests/pt_expt/train/test_tensorboard_observer.py +++ b/source/tests/pt_expt/train/test_tensorboard_observer.py @@ -248,6 +248,7 @@ def test_exceptional_close_is_idempotent(tmp_path: Path) -> None: observer.close() writer.close.assert_called_once() + def test_coinciding_display_keeps_step_train_metrics(tmp_path: Path) -> None: """At disp ∩ tensorboard_freq, train/* is written once from the step hook. @@ -312,4 +313,3 @@ def test_coinciding_display_keeps_step_train_metrics(tmp_path: Path) -> None: tags = [call.args[0] for call in writer.add_scalar.call_args_list] assert "valid/rmse" in tags assert "timing/interval_wall_time" in tags - From fec294cbdc48d27977e00106c86d592db9cd0145 Mon Sep 17 00:00:00 2001 From: OutisLi-Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:41:18 +0800 Subject: [PATCH 09/15] test(train): annotate expected RuntimeError in observer cleanup test Satisfy empty-except lint on the intentional best-effort catch. --- source/tests/common/dpmodel/test_train_observer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/source/tests/common/dpmodel/test_train_observer.py b/source/tests/common/dpmodel/test_train_observer.py index 16465e03b4..0496d5efdf 100644 --- a/source/tests/common/dpmodel/test_train_observer.py +++ b/source/tests/common/dpmodel/test_train_observer.py @@ -139,6 +139,7 @@ def train_step(self, task: TrainingTask, step: int) -> TrainStepResult: try: trainer.run(TrainingTaskCollection.single(DummyData([1.0]))) except RuntimeError: + # Expected: DummyTrainer.train_step raises; assert cleanup below. pass assert observer.events[0] == "begin:0" assert observer.events[-1] == "end:0" From d8c2013967e0c6006efadd2722c418dfce932c77 Mon Sep 17 00:00:00 2001 From: OutisLi-Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:41:22 +0800 Subject: [PATCH 10/15] fix(train): harden observer cleanup and profiler factory fallbacks Continue on_train_end for remaining observers after one fails so an active profiler still stops. Import torch.profiler when only the TensorBoard handler factory is missing. Document pt_expt profiling behavior in argcheck. --- deepmd/dpmodel/train/observer.py | 11 ++++++++++- deepmd/pt_expt/train/profiler.py | 6 +++++- deepmd/utils/argcheck.py | 12 +++++++++--- source/tests/common/dpmodel/test_train_observer.py | 1 + source/tests/pt_expt/train/test_profiler_observer.py | 1 + 5 files changed, 26 insertions(+), 5 deletions(-) diff --git a/deepmd/dpmodel/train/observer.py b/deepmd/dpmodel/train/observer.py index 6b7da5d88b..d65632d2fd 100644 --- a/deepmd/dpmodel/train/observer.py +++ b/deepmd/dpmodel/train/observer.py @@ -203,5 +203,14 @@ def on_train_end( *, rank_context: RankContext, ) -> None: + # Cleanup must run for every observer even if an earlier one fails + # (e.g. TensorBoard close must not skip stopping an active profiler). + first_error: Exception | None = None for observer in self._observers: - observer.on_train_end(tasks, rank_context=rank_context) + try: + observer.on_train_end(tasks, rank_context=rank_context) + except Exception as error: + if first_error is None: + first_error = error + if first_error is not None: + raise first_error diff --git a/deepmd/pt_expt/train/profiler.py b/deepmd/pt_expt/train/profiler.py index 16e606a4d1..1fb02e773f 100644 --- a/deepmd/pt_expt/train/profiler.py +++ b/deepmd/pt_expt/train/profiler.py @@ -164,7 +164,11 @@ def on_train_begin( profiler_factory = self._profiler_factory schedule_factory = self._schedule_factory handler_factory = self._tensorboard_handler_factory - if profiler_factory is None or schedule_factory is None: + if ( + profiler_factory is None + or schedule_factory is None + or handler_factory is None + ): import torch.profiler as torch_profiler if profiler_factory is None: diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 6bdbbaa002..fb69275892 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -6028,12 +6028,18 @@ def training_args( ) doc_profiling = ( "Enable performance profiling. TensorFlow and PyTorch can export a Chrome " - "JSON trace; PaddlePaddle starts its Nsight Systems profiling flow." + "JSON trace; PaddlePaddle starts its Nsight Systems profiling flow. " + "PyTorch Exportable (pt_expt) writes Chrome traces and, under distributed " + "training, appends a `.rankN` suffix to the profiling file stem." + ) + doc_profiling_file = ( + "Output file for the TensorFlow or PyTorch Chrome JSON trace. " + "For pt_expt distributed runs the stem receives a `.rankN` suffix." ) - doc_profiling_file = "Output file for the TensorFlow or PyTorch Chrome JSON trace." doc_enable_profiler = ( "Enable the backend profiler. TensorFlow and PyTorch write profiler data " - "under `tensorboard_log_dir`; PaddlePaddle starts Nsight Systems profiling." + "under `tensorboard_log_dir`; PaddlePaddle starts Nsight Systems profiling. " + "pt_expt uses TorchProfilerObserver with the same TensorBoard log directory." ) doc_tensorboard = "Enable tensorboard" doc_tensorboard_log_dir = "The log directory of tensorboard outputs" diff --git a/source/tests/common/dpmodel/test_train_observer.py b/source/tests/common/dpmodel/test_train_observer.py index 16465e03b4..0496d5efdf 100644 --- a/source/tests/common/dpmodel/test_train_observer.py +++ b/source/tests/common/dpmodel/test_train_observer.py @@ -139,6 +139,7 @@ def train_step(self, task: TrainingTask, step: int) -> TrainStepResult: try: trainer.run(TrainingTaskCollection.single(DummyData([1.0]))) except RuntimeError: + # Expected: DummyTrainer.train_step raises; assert cleanup below. pass assert observer.events[0] == "begin:0" assert observer.events[-1] == "end:0" diff --git a/source/tests/pt_expt/train/test_profiler_observer.py b/source/tests/pt_expt/train/test_profiler_observer.py index 58157d4406..f5785303ea 100644 --- a/source/tests/pt_expt/train/test_profiler_observer.py +++ b/source/tests/pt_expt/train/test_profiler_observer.py @@ -380,6 +380,7 @@ def test_close_releases_even_when_export_fails(tmp_path: Path) -> None: try: observer.close() except RuntimeError: + # Expected: export path raises; stop/cleanup assertions follow. pass assert observer.profiler is None profiler.stop.assert_called_once() From b80efbb65ea9ffcf3e86e5ed5e5006910c0826c6 Mon Sep 17 00:00:00 2001 From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:55:01 +0800 Subject: [PATCH 11/15] style(profiler): move type-only imports into TYPE_CHECKING Satisfy ruff TC001/TC003 on pre-commit.ci for PR #6067. --- deepmd/pt_expt/train/profiler.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/deepmd/pt_expt/train/profiler.py b/deepmd/pt_expt/train/profiler.py index 1fb02e773f..1cf8f20d79 100644 --- a/deepmd/pt_expt/train/profiler.py +++ b/deepmd/pt_expt/train/profiler.py @@ -23,14 +23,11 @@ import shutil import socket import time -from collections.abc import ( - Callable, - Mapping, -) from pathlib import ( Path, ) from typing import ( + TYPE_CHECKING, Any, ) @@ -38,10 +35,17 @@ StepObservation, TrainingObserver, ) -from deepmd.dpmodel.train.trainer import ( - RankContext, - TrainingTaskCollection, -) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + Mapping, + ) + + from deepmd.dpmodel.train.trainer import ( + RankContext, + TrainingTaskCollection, + ) log = logging.getLogger(__name__) From 08d896c483c454eabc8c792678b828316c93c987 Mon Sep 17 00:00:00 2001 From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:57:52 +0800 Subject: [PATCH 12/15] =?UTF-8?q?fix(train):=20break=20observer=E2=86=94tr?= =?UTF-8?q?ainer=20CodeQL=20cyclic=20import?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move RankContext, TrainingTask, TrainingTaskCollection, and TrainStepResult into a leaf types module so observer annotations no longer TYPE_CHECKING-import trainer while trainer runtime-imports observer. Public re-exports from trainer and the package stay identical. --- deepmd/dpmodel/train/__init__.py | 10 +- deepmd/dpmodel/train/observer.py | 2 +- deepmd/dpmodel/train/trainer.py | 181 ++------------------------ deepmd/dpmodel/train/types.py | 209 +++++++++++++++++++++++++++++++ 4 files changed, 223 insertions(+), 179 deletions(-) create mode 100644 deepmd/dpmodel/train/types.py diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py index ce42bd8444..adb0b43f9e 100644 --- a/deepmd/dpmodel/train/__init__.py +++ b/deepmd/dpmodel/train/__init__.py @@ -38,16 +38,18 @@ TrainingTimer, ) from .trainer import ( - DEFAULT_TASK_KEY, AbstractTrainer, LearningCurveWriter, - RankContext, TrainerConfig, + change_model_out_bias, + change_model_out_bias_by_task, +) +from .types import ( + DEFAULT_TASK_KEY, + RankContext, TrainingTask, TrainingTaskCollection, TrainStepResult, - change_model_out_bias, - change_model_out_bias_by_task, ) __all__ = [ diff --git a/deepmd/dpmodel/train/observer.py b/deepmd/dpmodel/train/observer.py index 6b7da5d88b..f650b9301b 100644 --- a/deepmd/dpmodel/train/observer.py +++ b/deepmd/dpmodel/train/observer.py @@ -31,7 +31,7 @@ from .timing import ( DisplayInterval, ) - from .trainer import ( + from .types import ( RankContext, TrainingTaskCollection, TrainStepResult, diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py index a1333c0c78..742ca260cc 100644 --- a/deepmd/dpmodel/train/trainer.py +++ b/deepmd/dpmodel/train/trainer.py @@ -18,7 +18,6 @@ ) from collections.abc import ( Callable, - Iterator, Mapping, MutableMapping, Sequence, @@ -28,7 +27,6 @@ ) from dataclasses import ( dataclass, - field, ) from pathlib import ( Path, @@ -60,14 +58,19 @@ DisplayInterval, TrainingTimer, ) +from .types import ( + DEFAULT_TASK_KEY, + RankContext, + TrainingTask, + TrainingTaskCollection, + TrainStepResult, +) if TYPE_CHECKING: from .metrics import ( TrainingMetricAccumulator, ) -DEFAULT_TASK_KEY = "Default" - log = logging.getLogger(__name__) LossResults = dict[str, float] @@ -122,23 +125,6 @@ def change_model_out_bias_by_task( return models -@dataclass(frozen=True) -class RankContext: - """Rank metadata used by a trainer. - - A single-process run is represented as rank 0 in a world of size 1. This - makes single-rank training a special case of multi-rank training. - """ - - rank: int = 0 - world_size: int = 1 - - @property - def is_chief(self) -> bool: - """Whether this rank is responsible for user-visible side effects.""" - return self.rank == 0 - - @dataclass(frozen=True) class TrainerConfig: """Common trainer configuration shared by backend implementations.""" @@ -182,159 +168,6 @@ def from_training_params( ) -@dataclass -class TrainingTask: - """One training task. - - Single-task training is represented by a collection containing one task. - """ - - key: str - training_data: Any - validation_data: Any | None = None - valid_numb_batch: int = 1 - data_requirements: list[Any] = field(default_factory=list) - - def __post_init__(self) -> None: - self.valid_numb_batch = max(int(self.valid_numb_batch), 1) - - def add_data_requirements(self) -> None: - """Attach data requirements to train and validation data if possible.""" - if not self.data_requirements: - return - for data in (self.training_data, self.validation_data): - if data is not None and hasattr(data, "add_data_requirements"): - data.add_data_requirements(self.data_requirements) - - -class TrainingTaskCollection: - """Ordered collection of training tasks with optional sampling weights.""" - - def __init__( - self, - tasks: Mapping[str, TrainingTask] | Sequence[TrainingTask], - probabilities: Mapping[str, float] | Sequence[float] | None = None, - ) -> None: - if isinstance(tasks, Mapping): - task_dict = dict(tasks) - else: - task_list = list(tasks) - task_dict = {task.key: task for task in task_list} - if len(task_dict) != len(task_list): - raise ValueError("Training task keys must be unique.") - if not task_dict: - raise ValueError("At least one training task is required.") - for key, task in task_dict.items(): - if key != task.key: - raise ValueError( - f"Task mapping key {key!r} does not match task key {task.key!r}." - ) - self._tasks = task_dict - self._keys = list(task_dict) - self._probabilities = self._normalize_probabilities(probabilities) - - @classmethod - def single( - cls, - training_data: Any, - validation_data: Any | None = None, - *, - key: str = DEFAULT_TASK_KEY, - valid_numb_batch: int = 1, - data_requirements: list[Any] | None = None, - ) -> TrainingTaskCollection: - """Build a task collection for single-task training.""" - task = TrainingTask( - key=key, - training_data=training_data, - validation_data=validation_data, - valid_numb_batch=valid_numb_batch, - data_requirements=list(data_requirements or []), - ) - return cls([task]) - - @property - def keys(self) -> list[str]: - """Task keys in iteration order.""" - return list(self._keys) - - @property - def probabilities(self) -> np.ndarray: - """Normalized task sampling probabilities.""" - return self._probabilities.copy() - - @property - def is_multitask(self) -> bool: - """Whether more than one task is present.""" - return len(self._tasks) > 1 - - def __len__(self) -> int: - return len(self._tasks) - - def __iter__(self) -> Iterator[TrainingTask]: - for key in self._keys: - yield self._tasks[key] - - def __getitem__(self, key: str) -> TrainingTask: - return self._tasks[key] - - def select( - self, - choice: Callable[..., Any] | None = None, - ) -> TrainingTask: - """Select a task according to the configured probabilities.""" - if len(self._keys) == 1: - return self._tasks[self._keys[0]] - chooser = choice or np.random.choice - index = int( - chooser(np.arange(len(self._keys), dtype=np.int_), p=self._probabilities) - ) - return self._tasks[self._keys[index]] - - def _normalize_probabilities( - self, - probabilities: Mapping[str, float] | Sequence[float] | None, - ) -> np.ndarray: - if probabilities is None: - prob = np.ones(len(self._keys), dtype=np.float64) - elif isinstance(probabilities, Mapping): - missing = [key for key in self._keys if key not in probabilities] - if missing: - raise ValueError(f"Missing task probabilities for {missing}.") - unknown = [key for key in probabilities if key not in self._tasks] - if unknown: - raise ValueError(f"Unknown task probabilities for {unknown}.") - prob = np.asarray( - [probabilities[key] for key in self._keys], dtype=np.float64 - ) - else: - prob = np.asarray(probabilities, dtype=np.float64) - if prob.ndim != 1 or prob.shape[0] != len(self._keys): - raise ValueError("Task probabilities must match the number of tasks.") - if not np.all(np.isfinite(prob)): - raise ValueError("Task probabilities must be finite.") - if np.any(prob < 0.0): - raise ValueError("Task probabilities must be non-negative.") - prob_sum = float(np.sum(prob)) - if prob_sum <= 0.0: - raise ValueError("Task probabilities must sum to a positive value.") - return prob / prob_sum - - -@dataclass -class TrainStepResult: - """Backend payload returned from one optimizer step. - - Averaged display consumes ``train_results`` as detached scalar metrics. - These may remain backend arrays until they are converted for display. - """ - - task_key: str - step: int - payload: Any = None - train_results: Mapping[str, Any] | None = None - - class LearningCurveWriter: """Formatter for learning-curve files and per-task training logs.""" diff --git a/deepmd/dpmodel/train/types.py b/deepmd/dpmodel/train/types.py new file mode 100644 index 0000000000..deac0a6bd7 --- /dev/null +++ b/deepmd/dpmodel/train/types.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Shared training-loop types with no observer or trainer dependency. + +These types are the vocabulary that observers annotate against and that the +common trainer owns at runtime. Keeping them in a leaf module breaks the +observer ↔ trainer import cycle that CodeQL flags: observers import here for +annotations; the trainer imports observers for dispatch and re-exports these +names for backward-compatible ``from .trainer import ...`` paths. +""" + +from __future__ import ( + annotations, +) + +from collections.abc import ( + Callable, + Iterator, + Mapping, + Sequence, +) +from dataclasses import ( + dataclass, + field, +) +from typing import ( + Any, +) + +import numpy as np + +DEFAULT_TASK_KEY = "Default" + +__all__ = [ + "DEFAULT_TASK_KEY", + "RankContext", + "TrainingTask", + "TrainingTaskCollection", + "TrainStepResult", +] + + +@dataclass(frozen=True) +class RankContext: + """Rank metadata used by a trainer. + + A single-process run is represented as rank 0 in a world of size 1. This + makes single-rank training a special case of multi-rank training. + """ + + rank: int = 0 + world_size: int = 1 + + @property + def is_chief(self) -> bool: + """Whether this rank is responsible for user-visible side effects.""" + return self.rank == 0 + + +@dataclass +class TrainingTask: + """One training task. + + Single-task training is represented by a collection containing one task. + """ + + key: str + training_data: Any + validation_data: Any | None = None + valid_numb_batch: int = 1 + data_requirements: list[Any] = field(default_factory=list) + + def __post_init__(self) -> None: + self.valid_numb_batch = max(int(self.valid_numb_batch), 1) + + def add_data_requirements(self) -> None: + """Attach data requirements to train and validation data if possible.""" + if not self.data_requirements: + return + for data in (self.training_data, self.validation_data): + if data is not None and hasattr(data, "add_data_requirements"): + data.add_data_requirements(self.data_requirements) + + +class TrainingTaskCollection: + """Ordered collection of training tasks with optional sampling weights.""" + + def __init__( + self, + tasks: Mapping[str, TrainingTask] | Sequence[TrainingTask], + probabilities: Mapping[str, float] | Sequence[float] | None = None, + ) -> None: + if isinstance(tasks, Mapping): + task_dict = dict(tasks) + else: + task_list = list(tasks) + task_dict = {task.key: task for task in task_list} + if len(task_dict) != len(task_list): + raise ValueError("Training task keys must be unique.") + if not task_dict: + raise ValueError("At least one training task is required.") + for key, task in task_dict.items(): + if key != task.key: + raise ValueError( + f"Task mapping key {key!r} does not match task key {task.key!r}." + ) + self._tasks = task_dict + self._keys = list(task_dict) + self._probabilities = self._normalize_probabilities(probabilities) + + @classmethod + def single( + cls, + training_data: Any, + validation_data: Any | None = None, + *, + key: str = DEFAULT_TASK_KEY, + valid_numb_batch: int = 1, + data_requirements: list[Any] | None = None, + ) -> TrainingTaskCollection: + """Build a task collection for single-task training.""" + task = TrainingTask( + key=key, + training_data=training_data, + validation_data=validation_data, + valid_numb_batch=valid_numb_batch, + data_requirements=list(data_requirements or []), + ) + return cls([task]) + + @property + def keys(self) -> list[str]: + """Task keys in iteration order.""" + return list(self._keys) + + @property + def probabilities(self) -> np.ndarray: + """Normalized task sampling probabilities.""" + return self._probabilities.copy() + + @property + def is_multitask(self) -> bool: + """Whether more than one task is present.""" + return len(self._tasks) > 1 + + def __len__(self) -> int: + return len(self._tasks) + + def __iter__(self) -> Iterator[TrainingTask]: + for key in self._keys: + yield self._tasks[key] + + def __getitem__(self, key: str) -> TrainingTask: + return self._tasks[key] + + def select( + self, + choice: Callable[..., Any] | None = None, + ) -> TrainingTask: + """Select a task according to the configured probabilities.""" + if len(self._keys) == 1: + return self._tasks[self._keys[0]] + chooser = choice or np.random.choice + index = int( + chooser(np.arange(len(self._keys), dtype=np.int_), p=self._probabilities) + ) + return self._tasks[self._keys[index]] + + def _normalize_probabilities( + self, + probabilities: Mapping[str, float] | Sequence[float] | None, + ) -> np.ndarray: + if probabilities is None: + prob = np.ones(len(self._keys), dtype=np.float64) + elif isinstance(probabilities, Mapping): + missing = [key for key in self._keys if key not in probabilities] + if missing: + raise ValueError(f"Missing task probabilities for {missing}.") + unknown = [key for key in probabilities if key not in self._tasks] + if unknown: + raise ValueError(f"Unknown task probabilities for {unknown}.") + prob = np.asarray( + [probabilities[key] for key in self._keys], dtype=np.float64 + ) + else: + prob = np.asarray(probabilities, dtype=np.float64) + if prob.ndim != 1 or prob.shape[0] != len(self._keys): + raise ValueError("Task probabilities must match the number of tasks.") + if not np.all(np.isfinite(prob)): + raise ValueError("Task probabilities must be finite.") + if np.any(prob < 0.0): + raise ValueError("Task probabilities must be non-negative.") + prob_sum = float(np.sum(prob)) + if prob_sum <= 0.0: + raise ValueError("Task probabilities must sum to a positive value.") + return prob / prob_sum + + +@dataclass +class TrainStepResult: + """Backend payload returned from one optimizer step. + + Averaged display consumes ``train_results`` as detached scalar metrics. + These may remain backend arrays until they are converted for display. + """ + + task_key: str + step: int + payload: Any = None + train_results: Mapping[str, Any] | None = None From 625b4bb325ef47bc8b9a3293e6e2881c7752eaa1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 10 Oct 2026 02:58:58 +0000 Subject: [PATCH 13/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/train/trainer.py | 2 -- deepmd/dpmodel/train/types.py | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py index 742ca260cc..eabf8e228d 100644 --- a/deepmd/dpmodel/train/trainer.py +++ b/deepmd/dpmodel/train/trainer.py @@ -37,7 +37,6 @@ TextIO, ) -import numpy as np from deepmd.dpmodel.common import ( to_numpy_array, @@ -59,7 +58,6 @@ TrainingTimer, ) from .types import ( - DEFAULT_TASK_KEY, RankContext, TrainingTask, TrainingTaskCollection, diff --git a/deepmd/dpmodel/train/types.py b/deepmd/dpmodel/train/types.py index deac0a6bd7..46ba7b91e4 100644 --- a/deepmd/dpmodel/train/types.py +++ b/deepmd/dpmodel/train/types.py @@ -33,9 +33,9 @@ __all__ = [ "DEFAULT_TASK_KEY", "RankContext", + "TrainStepResult", "TrainingTask", "TrainingTaskCollection", - "TrainStepResult", ] From 04a054384cfe0d731c4971e8e46565e51b2df933 Mon Sep 17 00:00:00 2001 From: OutisLi_Bot <198997215+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:59:56 +0800 Subject: [PATCH 14/15] fix(train): restore DEFAULT_TASK_KEY re-export after pre-commit pre-commit dropped the unused import of DEFAULT_TASK_KEY from trainer, breaking deepmd.dpmodel.train.data and historical trainer imports. Re-export it explicitly and point data.py at the leaf types module. --- deepmd/dpmodel/train/data.py | 2 +- deepmd/dpmodel/train/trainer.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/deepmd/dpmodel/train/data.py b/deepmd/dpmodel/train/data.py index 3e91dcfc53..88898c14c2 100644 --- a/deepmd/dpmodel/train/data.py +++ b/deepmd/dpmodel/train/data.py @@ -19,7 +19,7 @@ StatFileSpec, ) -from .trainer import ( +from .types import ( DEFAULT_TASK_KEY, ) diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py index eabf8e228d..8a785a6fa7 100644 --- a/deepmd/dpmodel/train/trainer.py +++ b/deepmd/dpmodel/train/trainer.py @@ -58,6 +58,7 @@ TrainingTimer, ) from .types import ( + DEFAULT_TASK_KEY as DEFAULT_TASK_KEY, RankContext, TrainingTask, TrainingTaskCollection, From c4c408b29b7a0309dc227a280f5b3ad313f952e6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 10 Oct 2026 03:04:09 +0000 Subject: [PATCH 15/15] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/train/trainer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py index 8a785a6fa7..380fe8aac7 100644 --- a/deepmd/dpmodel/train/trainer.py +++ b/deepmd/dpmodel/train/trainer.py @@ -37,7 +37,6 @@ TextIO, ) - from deepmd.dpmodel.common import ( to_numpy_array, ) @@ -57,8 +56,8 @@ DisplayInterval, TrainingTimer, ) +from .types import DEFAULT_TASK_KEY as DEFAULT_TASK_KEY from .types import ( - DEFAULT_TASK_KEY as DEFAULT_TASK_KEY, RankContext, TrainingTask, TrainingTaskCollection,