From 4771bb0d4118c35943ddba5874aaaa5c26355de6 Mon Sep 17 00:00:00 2001
From: Jiaming Yuan <jm.yuan@outlook.com>
Date: Mon, 21 Oct 2019 17:03:38 +0800
Subject: [PATCH] Catch exception in transform function omp context. (#4960)

---
 src/common/transform.h                   |  6 +++++-
 tests/cpp/common/test_transform_range.cc | 17 +++++++++++++++--
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/src/common/transform.h b/src/common/transform.h
index 09888338a..b62ac47cd 100644
--- a/src/common/transform.h
+++ b/src/common/transform.h
@@ -5,6 +5,8 @@
 #define XGBOOST_COMMON_TRANSFORM_H_
 
 #include <dmlc/omp.h>
+#include <dmlc/common.h>
+
 #include <xgboost/data.h>
 #include <utility>
 #include <vector>
@@ -148,10 +150,12 @@ class Transform {
     template <typename... HDV>
     void LaunchCPU(Functor func, HDV*... vectors) const {
       omp_ulong end = static_cast<omp_ulong>(*(range_.end()));
+      dmlc::OMPException omp_exc;
 #pragma omp parallel for schedule(static)
       for (omp_ulong idx = 0; idx < end; ++idx) {
-        func(idx, UnpackHDV(vectors)...);
+        omp_exc.Run(func, idx, UnpackHDV(vectors)...);
       }
+      omp_exc.Rethrow();
     }
 
    private:
diff --git a/tests/cpp/common/test_transform_range.cc b/tests/cpp/common/test_transform_range.cc
index dbf6de3b3..81bc73962 100644
--- a/tests/cpp/common/test_transform_range.cc
+++ b/tests/cpp/common/test_transform_range.cc
@@ -50,13 +50,26 @@ TEST(Transform, DeclareUnifiedTest(Basic)) {
   out_vec.Fill(0);
 
   Transform<>::Init(TestTransformRange<bst_float>{},
-	                Range{0, static_cast<Range::DifferenceType>(size)},
-	                TRANSFORM_GPU)
+                    Range{0, static_cast<Range::DifferenceType>(size)},
+                    TRANSFORM_GPU)
       .Eval(&out_vec, &in_vec);
   std::vector<bst_float> res = out_vec.HostVector();
 
   ASSERT_TRUE(std::equal(h_sol.begin(), h_sol.end(), res.begin()));
 }
 
+#if !defined(__CUDACC__)
+TEST(Transform, Exception) {
+  size_t const kSize {16};
+  std::vector<bst_float> h_in(kSize);
+  const HostDeviceVector<bst_float> in_vec{h_in, -1};
+  EXPECT_ANY_THROW({
+    Transform<>::Init([](size_t idx, common::Span<float const> _in) { _in[idx + 1]; },
+                      Range(0, static_cast<Range::DifferenceType>(kSize)), -1)
+        .Eval(&in_vec);
+  });
+}
+#endif
+
 } // namespace common
 } // namespace xgboost