fix ai comment

Author: OutisLi

deepmd/pt/train/validation.py

@@ -6,6 +6,7 @@
 )
 
 import logging
+import traceback
 from dataclasses import (
     dataclass,
 )
@@ -110,6 +111,12 @@ def resolve_full_validation_start_step(
 def parse_validation_metric(metric: str) -> tuple[str, str]:
     """Parse the configured full validation metric."""
     normalized_metric = normalize_full_validation_metric(metric)
+    if normalized_metric not in METRIC_KEY_MAP:
+        supported_metrics = ", ".join(item.upper() for item in METRIC_KEY_MAP)
+        raise ValueError(
+            "validating.validation_metric must be one of "
+            f"{supported_metrics}, got {metric!r}."
+        )
     return normalized_metric, METRIC_KEY_MAP[normalized_metric]
 
 
@@ -255,7 +262,7 @@ def __init__(
         self.enabled = (
             self.full_validation
             and self.start_step is not None
-            and self.start_step < num_steps
+            and self.start_step <= num_steps
         )
         self.step_column_width = max(len("step"), len(str(num_steps)))
         self._write_mode = "a" if restart_training else "w"
@@ -308,21 +315,60 @@ def run(
             dist.barrier()
 
         result: FullValidationResult | None = None
+        caught_exception: Exception | None = None
+        error_message = None
         save_path = [None]
         if self.rank == 0:
-            result = self._evaluate(display_step)
-            save_path[0] = result.saved_best_path
+            try:
+                result = self._evaluate(display_step)
+                save_path[0] = result.saved_best_path
+            except Exception as exc:
+                caught_exception = exc
+                error_message = (
+                    "Full validation failed on rank 0 during evaluation:\n"
+                    f"{traceback.format_exc()}"
+                )
+
+        self._raise_if_distributed_error(error_message, caught_exception)
 
         if self.is_distributed:
             dist.broadcast_object_list(save_path, src=0)
 
         if save_path[0] is not None:
-            save_checkpoint(Path(save_path[0]), lr=lr, step=step_id)
-            if self.rank == 0:
-                self._prune_best_checkpoints(keep_names={Path(save_path[0]).name})
+            try:
+                if not self.is_distributed or self.zero_stage == 0:
+                    if self.rank == 0:
+                        save_checkpoint(Path(save_path[0]), lr=lr, step=step_id)
+                else:
+                    save_checkpoint(Path(save_path[0]), lr=lr, step=step_id)
+                if self.rank == 0:
+                    self._prune_best_checkpoints(keep_names={Path(save_path[0]).name})
+            except Exception as exc:
+                caught_exception = exc
+                error_message = (
+                    "Full validation failed while saving the best checkpoint:\n"
+                    f"{traceback.format_exc()}"
+                )
+            else:
+                error_message = None
+                caught_exception = None
+
+            self._raise_if_distributed_error(error_message, caught_exception)
 
         if self.rank == 0:
-            self._log_result(result)
+            try:
+                self._log_result(result)
+            except Exception as exc:
+                caught_exception = exc
+                error_message = (
+                    "Full validation failed while writing logs:\n"
+                    f"{traceback.format_exc()}"
+                )
+            else:
+                error_message = None
+                caught_exception = None
+
+        self._raise_if_distributed_error(error_message, caught_exception)
 
         if self.is_distributed:
             dist.barrier()
@@ -367,8 +413,12 @@ def evaluate_all_systems(self) -> dict[str, float]:
 
         system_metrics = []
         for dataset in self.validation_data.systems:
-            assert isinstance(dataset, DeepmdDataSetForLoader)
-            system_metrics.append(self._evaluate_system(dataset._data_system))
+            if not isinstance(dataset, DeepmdDataSetForLoader):
+                raise TypeError(
+                    "Full validation expects each dataset in validation_data.systems "
+                    f"to be DeepmdDataSetForLoader, got {type(dataset)!r}."
+                )
+            system_metrics.append(self._evaluate_system(dataset.data_system))
 
         aggregated = weighted_average([metric for metric in system_metrics if metric])
         return {
@@ -555,6 +605,25 @@ def _initialize_best_checkpoints(self, restart_training: bool) -> None:
         else:
             self._prune_best_checkpoints()
 
+    def _raise_if_distributed_error(
+        self,
+        local_error_message: str | None,
+        local_exception: Exception | None = None,
+    ) -> None:
+        """Propagate a local error to all ranks and raise consistently."""
+        error_message = local_error_message
+        if self.is_distributed:
+            gathered_errors = [None] * dist.get_world_size()
+            dist.all_gather_object(gathered_errors, local_error_message)
+            error_message = next(
+                (message for message in gathered_errors if message is not None), None
+            )
+        if error_message is None:
+            return
+        if local_exception is not None:
+            raise RuntimeError(error_message) from local_exception
+        raise RuntimeError(error_message)
+
     def _log_result(self, result: FullValidationResult | None) -> None:
         """Log and persist full validation results on rank 0."""
         assert result is not None

deepmd/pt/utils/dataset.py

@@ -49,6 +49,11 @@ def __init__(
     def __len__(self) -> int:
         return self._data_system.nframes
 
+    @property
+    def data_system(self) -> DeepmdData:
+        """Expose the underlying DeePMD data system."""
+        return self._data_system
+
     def __getitem__(self, index: int) -> dict[str, Any]:
         """Get a frame from the selected system."""
         b_data = self._data_system.get_item_torch(index, max(1, NUM_WORKERS))

deepmd/utils/argcheck.py

@@ -4284,10 +4284,6 @@ def validate_full_validation_config(
     if float(validating.get("full_val_start", 0.0)) == 1.0:
         return
 
-    if multi_task:
-        # Unsupported multi-task mode is rejected during trainer initialization.
-        return
-
     metric = validating["validation_metric"]
     if not is_valid_full_validation_metric(metric):
         valid_metrics = ", ".join(item.upper() for item in FULL_VALIDATION_METRIC_PREFS)
@@ -4296,9 +4292,19 @@ def validate_full_validation_config(
             f"{valid_metrics}, got {metric!r}."
         )
 
+    if multi_task:
+        raise ValueError(
+            "validating.full_validation only supports single-task energy "
+            "training; multi-task training is not supported."
+        )
+
     loss_params = data.get("loss", {})
-    if loss_params.get("type", "ener") != "ener":
-        return
+    loss_type = loss_params.get("type", "ener")
+    if loss_type != "ener":
+        raise ValueError(
+            "validating.full_validation only supports single-task energy "
+            f"training with loss.type='ener'; got loss.type={loss_type!r}."
+        )
 
     if not data.get("training", {}).get("validation_data"):
         raise ValueError(

source/tests/pt/test_training.py

@@ -2,6 +2,7 @@
 import json
 import os
 import shutil
+import tempfile
 import unittest
 from copy import (
     deepcopy,
@@ -761,6 +762,9 @@ def test_fitting_stat_consistency(self) -> None:
 
 class TestFullValidation(unittest.TestCase):
     def setUp(self) -> None:
+        self._cwd = os.getcwd()
+        self._tmpdir = tempfile.TemporaryDirectory()
+        os.chdir(self._tmpdir.name)
         input_json = str(Path(__file__).parent / "water/se_atten.json")
         with open(input_json) as f:
             self.config = json.load(f)
@@ -782,13 +786,8 @@ def setUp(self) -> None:
         }
 
     def tearDown(self) -> None:
-        for f in os.listdir("."):
-            if (f.startswith("model") or f.startswith("best")) and f.endswith(".pt"):
-                os.remove(f)
-            if f in ["lcurve.out", "val.log", "checkpoint"]:
-                os.remove(f)
-            if f.startswith("stat_files"):
-                shutil.rmtree(f)
+        os.chdir(self._cwd)
+        self._tmpdir.cleanup()
 
     @patch("deepmd.pt.train.validation.FullValidator.evaluate_all_systems")
     def test_full_validation_rotates_best_checkpoint(self, mocked_eval) -> None:
@@ -843,12 +842,10 @@ def test_full_validation_rejects_multitask(self) -> None:
             "full_val_file": "val.log",
             "full_val_start": 0.0,
         }
-        config["model"], shared_links = preprocess_shared_params(config["model"])
+        config["model"], _ = preprocess_shared_params(config["model"])
         config = update_deepmd_input(config, warning=False)
-        config = normalize(config, multi_task=True)
-
         with self.assertRaisesRegex(ValueError, "multi-task"):
-            get_trainer(config, shared_links=shared_links)
+            normalize(config, multi_task=True)
 
 
 if __name__ == "__main__":

source/tests/pt/test_validation.py

@@ -7,6 +7,9 @@
 )
 
 import torch
+from dargs.dargs import (
+    ArgumentValueError,
+)
 
 from deepmd.pt.train.validation import (
     FullValidator,
@@ -135,5 +138,5 @@ def test_normalize_rejects_zero_prefactor_metric(self) -> None:
     def test_normalize_rejects_invalid_metric(self) -> None:
         config = _make_single_task_config()
         config["validating"]["validation_metric"] = "X:MAE"
-        with self.assertRaisesRegex(Exception, "validation_metric"):
+        with self.assertRaisesRegex(ArgumentValueError, "validation_metric"):
             normalize(config)