fix(pt/pd): fix incompatibility between AutoBatchSize and eval hooks

deepmd/pd/infer/deep_eval.py

@@ -64,6 +64,9 @@
     to_numpy_array,
     to_paddle_tensor,
 )
+from deepmd.utils.batch_size import (
+    RetrySignal,
+)
 from deepmd.utils.econf_embd import (
     sort_element_type,
 )
@@ -830,19 +833,30 @@ def eval_descriptor(
         model = (
             self.dp.model["Default"] if isinstance(self.dp, ModelWrapper) else self.dp
         )
-        model.set_eval_descriptor_hook(True)
-        self.eval(
-            coords,
-            cells,
-            atom_types,
-            atomic=False,
-            fparam=fparam,
-            aparam=aparam,
-            **kwargs,
-        )
-        descriptor = model.eval_descriptor()
-        model.set_eval_descriptor_hook(False)
-        return to_numpy_array(descriptor)
+        while True:
+            if self.auto_batch_size is not None:
+                self.auto_batch_size.set_oom_retry_mode(True)
+            model.set_eval_descriptor_hook(True)
+            retry = False
+            try:
+                self.eval(
+                    coords,
+                    cells,
+                    atom_types,
+                    atomic=False,
+                    fparam=fparam,
+                    aparam=aparam,
+                    **kwargs,
+                )
+                descriptor = model.eval_descriptor()
+            except RetrySignal:
+                retry = True
+            finally:
+                model.set_eval_descriptor_hook(False)
+                if self.auto_batch_size is not None:
+                    self.auto_batch_size.set_oom_retry_mode(False)
+            if not retry:
+                return to_numpy_array(descriptor)
 
     def eval_fitting_last_layer(
         self,
@@ -885,16 +899,27 @@ def eval_fitting_last_layer(
             Fitting output before last layer.
         """
         model = self.dp.model["Default"]
-        model.set_eval_fitting_last_layer_hook(True)
-        self.eval(
-            coords,
-            cells,
-            atom_types,
-            atomic=False,
-            fparam=fparam,
-            aparam=aparam,
-            **kwargs,
-        )
-        fitting_net = model.eval_fitting_last_layer()
-        model.set_eval_fitting_last_layer_hook(False)
-        return to_numpy_array(fitting_net)
+        while True:
+            if self.auto_batch_size is not None:
+                self.auto_batch_size.set_oom_retry_mode(True)
+            model.set_eval_fitting_last_layer_hook(True)
+            retry = False
+            try:
+                self.eval(
+                    coords,
+                    cells,
+                    atom_types,
+                    atomic=False,
+                    fparam=fparam,
+                    aparam=aparam,
+                    **kwargs,
+                )
+                fitting_net = model.eval_fitting_last_layer()
+            except RetrySignal:
+                retry = True
+            finally:
+                model.set_eval_fitting_last_layer_hook(False)
+                if self.auto_batch_size is not None:
+                    self.auto_batch_size.set_oom_retry_mode(False)
+            if not retry:
+                return to_numpy_array(fitting_net)

deepmd/pt/infer/deep_eval.py

@@ -67,6 +67,9 @@
     to_numpy_array,
     to_torch_tensor,
 )
+from deepmd.utils.batch_size import (
+    RetrySignal,
+)
 from deepmd.utils.econf_embd import (
     sort_element_type,
 )
@@ -847,19 +850,30 @@ def eval_descriptor(
             Descriptors.
         """
         model = self.dp.model["Default"]
-        model.set_eval_descriptor_hook(True)
-        self.eval(
-            coords,
-            cells,
-            atom_types,
-            atomic=False,
-            fparam=fparam,
-            aparam=aparam,
-            **kwargs,
-        )
-        descriptor = model.eval_descriptor()
-        model.set_eval_descriptor_hook(False)
-        return to_numpy_array(descriptor)
+        while True:
+            if self.auto_batch_size is not None:
+                self.auto_batch_size.set_oom_retry_mode(True)
+            model.set_eval_descriptor_hook(True)
+            retry = False
+            try:
+                self.eval(
+                    coords,
+                    cells,
+                    atom_types,
+                    atomic=False,
+                    fparam=fparam,
+                    aparam=aparam,
+                    **kwargs,
+                )
+                descriptor = model.eval_descriptor()
+            except RetrySignal:
+                retry = True
+            finally:
+                model.set_eval_descriptor_hook(False)
+                if self.auto_batch_size is not None:
+                    self.auto_batch_size.set_oom_retry_mode(False)
+            if not retry:
+                return to_numpy_array(descriptor)
 
     def eval_fitting_last_layer(
         self,
@@ -902,16 +916,27 @@ def eval_fitting_last_layer(
             Fitting output before last layer.
         """
         model = self.dp.model["Default"]
-        model.set_eval_fitting_last_layer_hook(True)
-        self.eval(
-            coords,
-            cells,
-            atom_types,
-            atomic=False,
-            fparam=fparam,
-            aparam=aparam,
-            **kwargs,
-        )
-        fitting_net = model.eval_fitting_last_layer()
-        model.set_eval_fitting_last_layer_hook(False)
-        return to_numpy_array(fitting_net)
+        while True:
+            if self.auto_batch_size is not None:
+                self.auto_batch_size.set_oom_retry_mode(True)
+            model.set_eval_fitting_last_layer_hook(True)
+            retry = False
+            try:
+                self.eval(
+                    coords,
+                    cells,
+                    atom_types,
+                    atomic=False,
+                    fparam=fparam,
+                    aparam=aparam,
+                    **kwargs,
+                )
+                fitting_net = model.eval_fitting_last_layer()
+            except RetrySignal:
+                retry = True
+            finally:
+                model.set_eval_fitting_last_layer_hook(False)
+                if self.auto_batch_size is not None:
+                    self.auto_batch_size.set_oom_retry_mode(False)
+            if not retry:
+                return to_numpy_array(fitting_net)

deepmd/utils/batch_size.py

@@ -22,6 +22,10 @@
 log = logging.getLogger(__name__)
 
 
+class RetrySignal(Exception):
+    """Signal to retry execution after OOM error."""
+
+
 class AutoBatchSize(ABC):
     """This class allows DeePMD-kit to automatically decide the maximum
     batch size that will not cause an OOM error.
@@ -85,6 +89,7 @@ def __init__(
                     )
 
         self.factor = factor
+        self.oom_retry_mode = False
 
     def execute(
         self, callable: Callable, start_index: int, natoms: int
@@ -135,6 +140,8 @@ def execute(
                 ) from e
             # adjust the next batch size
             self._adjust_batch_size(1.0 / self.factor)
+            if self.oom_retry_mode:
+                raise RetrySignal from e
             return 0, None
         else:
             n_tot = n_batch * natoms
@@ -292,3 +299,20 @@ def is_oom_error(self, e: Exception) -> bool:
         bool
             True if the exception is an OOM error
         """
+
+    def set_oom_retry_mode(self, enable: bool) -> None:
+        """Set OOM retry mode.
+
+        In OOM retry mode, an OOM during execution may reduce the current
+        batch size and raise :class:`RetrySignal` to indicate that execution
+        should be retried.
+
+        Callers that want all data to be re-executed must catch
+        :class:`RetrySignal` and restart the full evaluation themselves.
+
+        Parameters
+        ----------
+        enable : bool
+            True to enable OOM retry mode
+        """
+        self.oom_retry_mode = enable

source/tests/common/test_oom_retry.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+from deepmd.utils.batch_size import (
+    AutoBatchSize,
+    RetrySignal,
+)
+from deepmd.utils.errors import (
+    OutOfMemoryError,
+)
+
+
+class CustomizedAutoBatchSizeGPU(AutoBatchSize):
+    def is_gpu_available(self) -> bool:
+        return True
+
+    def is_oom_error(self, e):
+        return isinstance(e, OutOfMemoryError)
+
+
+class TestOOMRetry(unittest.TestCase):
+    def test_execute_oom_retry_mode_raises_retry_signal(self) -> None:
+        auto_batch_size = CustomizedAutoBatchSizeGPU(256, 2.0)
+
+        oom = OutOfMemoryError("oom")
+
+        def executor(batch_size: int, start_index: int) -> tuple[int, None]:
+            raise oom
+
+        auto_batch_size.set_oom_retry_mode(True)
+        with self.assertRaises(RetrySignal) as context:
+            auto_batch_size.execute(executor, 0, 1)
+        self.assertIs(context.exception.__cause__, oom)
+        self.assertEqual(auto_batch_size.current_batch_size, 128)
+
+    def test_execute_oom_retry_mode_false_returns_zero(self) -> None:
+        auto_batch_size = CustomizedAutoBatchSizeGPU(256, 2.0)
+
+        def executor(batch_size: int, start_index: int) -> tuple[int, None]:
+            raise OutOfMemoryError("oom")
+
+        auto_batch_size.set_oom_retry_mode(False)
+        n_batch, result = auto_batch_size.execute(executor, 0, 1)
+        self.assertEqual(n_batch, 0)
+        self.assertIsNone(result)
+        self.assertEqual(auto_batch_size.current_batch_size, 128)
+
+
+if __name__ == "__main__":
+    unittest.main()