Display rank too.

dmlc · Oct 9, 2020 · e3acf86 · e3acf86
1 parent cbceb99
commit e3acf86
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
@@ -391,8 +391,9 @@ using XGBBaseDeviceAllocator = thrust::device_malloc_allocator<T>;
 
 inline void ThrowOOMError(std::exception const& e, size_t bytes) {
   auto device = CurrentDevice();
+  auto rank = rabit::GetRank();
   std::stringstream ss;
-  ss << "Memory allocation error: " << e.what() << "\n"
+  ss << "Memory allocation error on worker " << rank << ": " << e.what() << "\n"
      << "- Free memory: " << AvailableMemory(device) << "\n"
      << "- Requested memory: " << bytes << std::endl;
   LOG(FATAL) << ss.str();