Fix pin_memory_thread not exiting quickly (#23646)

ssnl · facebook-github-bot · commit 0539462ca296 · 2019-08-01T15:24:14.000-07:00
Summary: fixes #23642 Pull Request resolved: #23646 Differential Revision: D16600874 Pulled By: soumith fbshipit-source-id: 50f0828d774a558d6f21e9dd21135906bd5be128
diff --git a/torch/utils/data/_utils/pin_memory.py b/torch/utils/data/_utils/pin_memory.py
@@ -22,7 +22,7 @@ def _pin_memory_loop(in_queue, out_queue, device_id, done_event):
         except queue.Empty:
             continue
         idx, data = r
-        if not isinstance(data, ExceptionWrapper):
+        if not done_event.is_set() and not isinstance(data, ExceptionWrapper):
             try:
                 data = pin_memory(data)
             except Exception:
diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py
@@ -886,9 +886,13 @@ def _shutdown_workers(self):
                 # corrupted data in `worker_result_queue` which `pin_memory_thread`
                 # reads from.
                 if hasattr(self, 'pin_memory_thread'):
-                    self.pin_memory_thread_done_event.set()
                     # Use hasattr in case error happens before we set the attribute.
+                    self.pin_memory_thread_done_event.set()
+                    # Send something to pin_memory_thread in case it is waiting
+                    # so that it can wake up and check `pin_memory_thread_done_event`
+                    self.worker_result_queue.put((None, None))
                     self.pin_memory_thread.join()
+                    self.worker_result_queue.close()
 
                 # Exit workers now.
                 self.workers_done_event.set()
