add thread name for dataloader (PaddlePaddle#46990)

Amberain · Oct 13, 2022 · 770501b · 770501b
1 parent f856fc8
commit 770501b
Show file tree

Hide file tree

Showing 6 changed files with 11 additions and 1 deletion.
diff --git a/paddle/fluid/platform/os_info.cc b/paddle/fluid/platform/os_info.cc
@@ -27,6 +27,7 @@ limitations under the License. */
 #else
 #include <unistd.h>
 #endif
+#include "glog/logging.h"
 #include "paddle/fluid/framework/new_executor/workqueue/thread_data_registry.h"
 #include "paddle/fluid/platform/macros.h"  // import DISABLE_COPY_AND_ASSIGN
 
@@ -115,6 +116,7 @@ bool SetCurrentThreadName(const std::string& name) {
     return false;
   }
   instance.SetCurrentThreadData(name);
+  VLOG(4) << __func__ << " " << name;
   return true;
 }
 

diff --git a/paddle/fluid/platform/os_info.h b/paddle/fluid/platform/os_info.h
@@ -57,7 +57,7 @@ ThreadId GetCurrentThreadId();
 // create/destory when using it.
 std::unordered_map<uint64_t, ThreadId> GetAllThreadIds();
 
-static constexpr const char* kDefaultThreadName = "unset";
+static constexpr const char* kDefaultThreadName = "unnamed";
 // Returns kDefaultThreadName if SetCurrentThreadName is never called.
 std::string GetCurrentThreadName();
 

diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -854,6 +854,8 @@ PYBIND11_MODULE(libpaddle, m) {
 
   m.def("_set_paddle_lib_path", &paddle::platform::dynload::SetPaddleLibPath);
 
+  m.def("set_current_thread_name", &paddle::platform::SetCurrentThreadName);
+
   m.def("_promote_types_if_complex_exists",
         &paddle::framework::PromoteTypesIfComplexExists);
 

diff --git a/python/paddle/fluid/dataloader/dataloader_iter.py b/python/paddle/fluid/dataloader/dataloader_iter.py
@@ -205,6 +205,7 @@ def _thread_loop(self, legacy_expected_place):
         # If we do not set cudaDeviceId in new thread, the default cudaDeviceId will be 0,
         # Which may cost hundreds of MB of GPU memory on CUDAPlace(0) if calling some cuda
         # APIs in this thread.
+        core.set_current_thread_name("Dataloader_" + str(id(self)))
         _set_expected_place(legacy_expected_place)
 
         while not self._thread_done_event.is_set():
@@ -530,6 +531,7 @@ def _thread_loop(self, legacy_expected_place):
         # If we do not set cudaDeviceId in new thread, the default cudaDeviceId will be 0,
         # Which may cost hundreds of MB of GPU memory on CUDAPlace(0) if calling some cuda
         # APIs in this thread.
+        core.set_current_thread_name("Dataloader_" + str(id(self)))
         _set_expected_place(legacy_expected_place)
 
         while not self._thread_done_event.is_set():

diff --git a/python/paddle/fluid/layers/io.py b/python/paddle/fluid/layers/io.py
@@ -477,6 +477,7 @@ def start_provide_thread(func):
         def __provider_thread__(legacy_expected_place):
             try:
                 # See _DataLoaderIterSingleProcess._thread_loop() for why set expected place here.
+
                 _set_expected_place(legacy_expected_place)
 
                 for tensors in func():

diff --git a/python/paddle/fluid/reader.py b/python/paddle/fluid/reader.py
@@ -1126,6 +1126,7 @@ def _exit_thread_unexpectedly(self):
 
     def _reader_thread_loop_for_multiprocess(self, legacy_expected_place):
         # See _DataLoaderIterSingleProcess._thread_loop() for why set expected place here.
+        core.set_current_thread_name("Dataloader_" + str(id(self)))
         _set_expected_place(legacy_expected_place)
 
         while not self._thread_done_event.is_set():
@@ -1169,6 +1170,7 @@ def _reader_thread_loop_for_multiprocess(self, legacy_expected_place):
     def _reader_thread_loop_for_singleprocess(self, legacy_expected_place):
         try:
             # See _DataLoaderIterSingleProcess._thread_loop() for why set expected place here.
+            core.set_current_thread_name("Dataloader_" + str(id(self)))
             _set_expected_place(legacy_expected_place)
 
             for sample in self._batch_reader():
@@ -1419,6 +1421,7 @@ def _start(self):
         def __thread_main__(legacy_expected_place):
             try:
                 # See _DataLoaderIterSingleProcess._thread_loop() for why set expected place here.
+                core.set_current_thread_name("Dataloader_" + str(id(self)))
                 _set_expected_place(legacy_expected_place)
 
                 while not self._queue.wait_for_inited(1):