From 9d975ea98d2d3e21360da22b93a11c8358beb324 Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Wed, 14 Jul 2021 13:52:31 +0800
Subject: [PATCH] Improve exception handling.

---
 python-package/xgboost/core.py | 53 +++++++++++++++++++---------------
 1 file changed, 30 insertions(+), 23 deletions(-)

diff --git a/python-package/xgboost/core.py b/python-package/xgboost/core.py
index 1cf229971a6e..2db850ac03a7 100644
--- a/python-package/xgboost/core.py
+++ b/python-package/xgboost/core.py
@@ -350,23 +350,41 @@ def proxy(self) -> "_ProxyDMatrix":
         """Handle of DMatrix proxy."""
         return self._handle
 
-    def _rethrow(self) -> None:
+    def _handle_exception(self, fn: Callable, dft_ret):
+        try:
+            return fn()
+        except Exception as e:
+            # Defer the exception in order to return 0 and stop the iteration.
+            # Exception inside a ctype callback function has no effect except
+            # for printing to stderr (doesn't stop the execution).
+            tb = sys.exc_info()[2]
+            # On dask, the worker is restarted and somehow the information is
+            # lost.
+            self._exception = e.with_traceback(tb)
+        return dft_ret
+
+    def _reraise(self) -> None:
+        self._temporary_data = None
         if self._exception is not None:
             #  pylint 2.7.0 believes `self._exception` can be None even with `assert
             #  isinstace`
-            raise self._exception  # pylint: disable=raising-bad-type
+            exc = self._exception
+            self._exception = None
+            raise exc  # pylint: disable=raising-bad-type
 
     def __del__(self) -> None:
         assert self._temporary_data is None, self._temporary_data
+        assert self._exception is None
 
-    def _reset_wrapper(self, this):  # pylint: disable=unused-argument
+    def _reset_wrapper(self, this) -> None:  # pylint: disable=unused-argument
         """A wrapper for user defined `reset` function."""
-        if self._temporary_data is not None:
-            # free the data
-            self._temporary_data = None
-        self.reset()
+        if self._exception is not None:
+            return
+        # free the data
+        self._temporary_data = None
+        self._handle_exception(self.reset, None)
 
-    def _next_wrapper(self, this):  # pylint: disable=unused-argument
+    def _next_wrapper(self, this) -> int:  # pylint: disable=unused-argument
         """A wrapper for user defined `next` function.
 
         `this` is not used in Python.  ctypes can handle `self` of a Python
@@ -396,19 +414,8 @@ def data_handle(data, *, feature_names=None, feature_types=None, **kwargs):
                 feature_types=feature_types,
                 **kwargs,
             )
-
-        try:
-            # Defer the exception in order to return 0 and stop the iteration.
-            # Exception inside a ctype callback function has no effect except
-            # for printing to stderr (doesn't stop the execution).
-            ret = self.next(data_handle)  # pylint: disable=not-callable
-        except Exception as e:  # pylint: disable=broad-except
-            tb = sys.exc_info()[2]
-            # On dask, the worker is restarted and somehow the information is
-            # lost.
-            self._exception = e.with_traceback(tb)
-            return 0
-        return ret
+        # pylint: disable=not-callable
+        return self._handle_exception(lambda: self.next(data_handle), 0)
 
     def reset(self) -> None:
         """Reset the data iterator.  Prototype for user defined function."""
@@ -638,7 +645,7 @@ def _init_from_iter(self, iterator: DataIter, enable_categorical: bool):
             ctypes.byref(handle),
         )
         # pylint: disable=protected-access
-        it._rethrow()
+        it._reraise()
         # delay check_call to throw intermediate exception first
         _check_call(ret)
         self.handle = handle
@@ -1212,7 +1219,7 @@ def _init(self, data, enable_categorical, **meta):
             ctypes.byref(handle),
         )
         # pylint: disable=protected-access
-        it._rethrow()
+        it._reraise()
         # delay check_call to throw intermediate exception first
         _check_call(ret)
         self.handle = handle