[DataLoader2] Handle empty queue exception

ghstack-source-id: 892eb01d219786e5dd209e1b34dd06552ae8df4f
Pull Request resolved: #785

torchdata/dataloader2/communication/protocol.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from queue import Empty as EmptyException
+
 from torchdata.dataloader2 import communication
 
 
@@ -69,7 +71,7 @@ def get_new_request(self, block=False):
             raise Exception("Trying to get next request, while having one unserved")
         try:
             response = self.request_queue.get(block=block)
-        except Exception:  # TODO(625): Catch only timeout exceptions
+        except EmptyException:
             raise EmptyQueue("queue is empty")
         self._req_received = response
         return response
@@ -216,7 +218,7 @@ def request_next(self):
     def get_response_reset_iterator(self, block=False):
         try:
             response = self.response_queue.get(block=block)
-        except Exception:  # TODO(627): Catch only timeout exceptions
+        except EmptyException:
             raise EmptyQueue("queue is empty")
         self.request_served(response)
 
@@ -228,7 +230,7 @@ def get_response_next(self, block=False, timeout=None):
             raise Exception("Can not expect any response without submitted request")
         try:
             response = self.response_queue.get(block=block, timeout=timeout)
-        except Exception:  # TODO(628): Catch only timeout exceptions
+        except EmptyException:
             raise EmptyQueue("queue is empty")
         self.request_served(response)