fix: address CodeQL findings in PR #5397

Replace assert with if/raise ValueError for user-facing config
validation (data_stat_protect, finetune branch/head checks).
Wrap train() in try/finally for destroy_process_group cleanup.
Add parents=True, exist_ok=True to stat_file mkdir.
Add strict=True to zip() calls. Fix minor test issues.

Author: Han Wang

deepmd/pt_expt/entrypoints/main.py

@@ -91,7 +91,7 @@ def get_trainer(
                     with h5py.File(stat_file_path, "w"):
                         pass
                 else:
-                    Path(stat_file_path).mkdir()
+                    Path(stat_file_path).mkdir(parents=True, exist_ok=True)
             stat_file_path = DPPath(stat_file_path, "a")
     else:
         # Multi-task: build per-task data systems
@@ -143,7 +143,7 @@ def get_trainer(
                         with h5py.File(_sf, "w"):
                             pass
                     else:
-                        Path(_sf).mkdir(parents=True)
+                        Path(_sf).mkdir(parents=True, exist_ok=True)
                 stat_file_path[model_key] = DPPath(_sf, "a")
             else:
                 stat_file_path[model_key] = None
@@ -290,18 +290,19 @@ def train(
     if os.environ.get("LOCAL_RANK") is not None:
         dist.init_process_group(backend="cuda:nccl,cpu:gloo")
 
-    trainer = get_trainer(
-        config,
-        init_model,
-        restart,
-        finetune_model=finetune,
-        finetune_links=finetune_links,
-        shared_links=shared_links,
-    )
-    trainer.run()
-
-    if dist.is_available() and dist.is_initialized():
-        dist.destroy_process_group()
+    try:
+        trainer = get_trainer(
+            config,
+            init_model,
+            restart,
+            finetune_model=finetune,
+            finetune_links=finetune_links,
+            shared_links=shared_links,
+        )
+        trainer.run()
+    finally:
+        if dist.is_available() and dist.is_initialized():
+            dist.destroy_process_group()
 
 
 def freeze(

deepmd/pt_expt/train/training.py

@@ -611,13 +611,16 @@ def _make_sample(
                     for ii in model_params["model_dict"]
                 ]
             )
-            assert np.allclose(_data_stat_protect, _data_stat_protect[0]), (
-                "Model key 'data_stat_protect' must be the same in each branch when multitask!"
-            )
+            if not np.allclose(_data_stat_protect, _data_stat_protect[0]):
+                raise ValueError(
+                    "Model key 'data_stat_protect' must be the same in each branch when multitask!"
+                )
             self.wrapper.share_params(
                 shared_links,
                 resume=(resuming and not self._finetune_update_stat) or self.rank != 0,
-                model_key_prob_map=dict(zip(self.model_keys, self.model_prob)),
+                model_key_prob_map=dict(
+                    zip(self.model_keys, self.model_prob, strict=True)
+                ),
                 data_stat_protect=_data_stat_protect[0],
             )
 
@@ -825,7 +828,9 @@ def _make_sample(
                 self._unwrapped.share_params(
                     shared_links,
                     resume=True,
-                    model_key_prob_map=dict(zip(self.model_keys, self.model_prob)),
+                    model_key_prob_map=dict(
+                        zip(self.model_keys, self.model_prob, strict=True)
+                    ),
                 )
 
             if optimizer_state_dict is not None:

deepmd/pt_expt/utils/finetune.py

@@ -105,10 +105,11 @@ def get_finetune_rules(
         finetune_links["Default"] = finetune_rule
     else:
         # Multi-task target — mirrors PT's logic
-        assert model_branch == "", (
-            "Multi-task fine-tuning does not support command-line branches chosen!"
-            "Please define the 'finetune_head' in each model params!"
-        )
+        if model_branch != "":
+            raise ValueError(
+                "Multi-task fine-tuning does not support command-line branches chosen! "
+                "Please define the 'finetune_head' in each model params!"
+            )
         if not finetune_from_multi_task:
             pretrained_keys = ["Default"]
         else:
@@ -120,10 +121,11 @@ def get_finetune_rules(
                 and model_config["model_dict"][model_key]["finetune_head"] != "RANDOM"
             ):
                 pretrained_key = model_config["model_dict"][model_key]["finetune_head"]
-                assert pretrained_key in pretrained_keys, (
-                    f"'{pretrained_key}' head chosen to finetune not exist in the pretrained model!"
-                    f"Available heads are: {list(pretrained_keys)}"
-                )
+                if pretrained_key not in pretrained_keys:
+                    raise ValueError(
+                        f"'{pretrained_key}' head chosen to finetune not exist in the pretrained model! "
+                        f"Available heads are: {list(pretrained_keys)}"
+                    )
                 model_branch_from = pretrained_key
             elif (
                 "finetune_head" not in model_config["model_dict"][model_key]

source/tests/pt_expt/test_multitask.py

@@ -936,7 +936,7 @@ def test_multitask_finetune_no_change_model_params(self) -> None:
         # Also set trainable=False to verify it's preserved
         ft_config_true["model"]["shared_dict"]["my_descriptor"]["trainable"] = False
         ft_config_true["model"], _ = preprocess_shared_params(ft_config_true["model"])
-        model_config_true, finetune_links_true = get_finetune_rules(
+        model_config_true, _finetune_links_true = get_finetune_rules(
             ckpt_path, deepcopy(ft_config_true["model"]), change_model_params=True
         )
 
@@ -1164,7 +1164,7 @@ def test_multitask_restart(self) -> None:
     def test_multitask_freeze(self) -> None:
         """Train, then freeze with --head and verify.
 
-        Only runs for se_e2_a descriptor to avoid redundant slow freeze tests.
+        Only runs for dpa3 descriptor to avoid redundant slow freeze tests.
         """
         if self.descriptor.get("type") != "dpa3":
             return
@@ -1948,7 +1948,7 @@ def test_gradient_accumulation(self) -> None:
 
         # Verify descriptor params are aliased (share_params)
         mt_desc_2 = mt_trainer.wrapper.model["model_2"].atomic_model.descriptor
-        for (n1, p1), (n2, p2) in zip(
+        for (n1, p1), (_n2, p2) in zip(
             mt_desc.named_parameters(), mt_desc_2.named_parameters(), strict=True
         ):
             assert p1.data_ptr() == p2.data_ptr(), (