adds to raise exception to main caller thread from prefetch generator

Author: jtsodapop

labelbox/data/annotation_types/collection.py

@@ -172,7 +172,11 @@ class LabelGenerator(PrefetchGenerator):
     than the LabelList but will be much more memory efficient.
     """
 
-    def __init__(self, data: Generator[Label, None, None], *args, **kwargs):
+    def __init__(self,
+                 data: Generator[Label, None, None],
+                 multithread: bool = False,
+                 *args,
+                 **kwargs):
         self._fns = {}
         super().__init__(data, *args, **kwargs)
 

labelbox/data/generator.py

@@ -1,10 +1,8 @@
-import logging
 import threading
 from queue import Queue
 from typing import Any, Iterable
-import threading
 
-logger = logging.getLogger(__name__)
+from labelbox.exceptions import ThreadException
 
 
 class ThreadSafeGen:
@@ -27,13 +25,43 @@ def __next__(self):
             return next(self.iterable)
 
 
+class PrefetchThread(threading.Thread):
+    """Class to override the Thread class. Helps raise 
+    exceptions to the main caller thread
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.exc = None
+
+    def run(self):
+        try:
+            super().run()
+        except BaseException as e:
+            self.exc = e
+
+    def join(self, timeout=None):
+        threading.Thread.join(self)
+        if self.exc:
+            raise self.exc
+
+
 class PrefetchGenerator:
     """
     Applys functions asynchronously to the output of a generator.
     Useful for modifying the generator results based on data from a network
     """
 
-    def __init__(self, data: Iterable[Any], prefetch_limit=20, num_executors=4):
+    #maybe change num exec to just 1, and if 1, make sync
+    #instead of self.get qeue in next, itll return just self._data.next
+    #kwarg on export for multithread, and all other things that use prefetch
+
+    # def __init__(self, data: Iterable[Any], prefetch_limit=20, num_executors=4):
+    def __init__(self,
+                 data: Iterable[Any],
+                 prefetch_limit=20,
+                 num_executors=1,
+                 multithread: bool = False):
         if isinstance(data, (list, tuple)):
             self._data = (r for r in data)
         else:
@@ -44,14 +72,17 @@ def __init__(self, data: Iterable[Any], prefetch_limit=20, num_executors=4):
         self.completed_threads = 0
         # Can only iterate over once it the queue.get hangs forever.
         self.done = False
+        if multithread:
+            num_executors = 4
         self.num_executors = num_executors
         self.threads = [
-            threading.Thread(target=self.fill_queue)
-            for _ in range(num_executors)
+            PrefetchThread(target=self.fill_queue) for _ in range(num_executors)
         ]
         for thread in self.threads:
             thread.daemon = True
             thread.start()
+        for thread in self.threads:
+            thread.join()
 
     def _process(self, value) -> Any:
         raise NotImplementedError("Abstract method needs to be implemented")
@@ -63,11 +94,9 @@ def fill_queue(self):
                 if value is None:
                     raise ValueError("Unexpected None")
                 self.queue.put(value)
-        except Exception as e:
-            logger.warning("Unexpected exception while filling the queue. %r",
-                           e)
-        finally:
-            self.queue.put(None)
+        except:
+            raise ThreadException(
+                "Unexpected exception while filling the queue.")
 
     def __iter__(self):
         return self

labelbox/data/serialization/labelbox_v1/converter.py

@@ -82,7 +82,7 @@ class LBV1VideoIterator(PrefetchGenerator):
     Generator that fetches video annotations in the background to be faster.
     """
 
-    def __init__(self, examples, client):
+    def __init__(self, examples, client, multithread: bool = False):
         self.client = client
         super().__init__(examples)
 

labelbox/exceptions.py

@@ -123,3 +123,8 @@ class MALValidationError(LabelboxError):
 class OperationNotAllowedException(Exception):
     """Raised when user does not have permissions to a resource or has exceeded usage limit"""
     pass
+
+
+class ThreadException(Exception):
+    """Raised when there is an issue with a thread, typically in Prefetch generator"""
+    pass