Limit number of background tasks

Author: exupero

fiftyone/core/models.py

@@ -478,13 +478,6 @@ def _apply_image_model_data_loader(
     with contextlib.ExitStack() as context:
         pb = context.enter_context(fou.ProgressBar(samples, progress=progress))
         ctx = context.enter_context(foc.SaveContext(samples))
-        submit = context.enter_context(
-            fou.async_executor(
-                max_workers=1,
-                skip_failures=skip_failures,
-                warning="Async failure labeling batches",
-            )
-        )
 
         def save_batch(sample_batch, labels_batch):
             with _handle_batch_error(skip_failures, sample_batch):
@@ -499,9 +492,15 @@ def save_batch(sample_batch, labels_batch):
                     )
                     ctx.save(sample)
 
-        for sample_batch, imgs in zip(
-            fou.iter_batches(samples, batch_size),
-            data_loader,
+        for submit, (sample_batch, imgs) in fou.async_iterator(
+            zip(
+                fou.iter_batches(samples, batch_size),
+                data_loader,
+            ),
+            limit=10,
+            max_workers=1,
+            skip_failures=skip_failures,
+            warning="Async failure labeling batches",
         ):
             with _handle_batch_error(skip_failures, sample_batch):
                 if isinstance(imgs, Exception):
@@ -1221,24 +1220,22 @@ def _compute_image_embeddings_data_loader(
         else:
             ctx = None
 
-        submit = context.enter_context(
-            fou.async_executor(
-                max_workers=1,
-                skip_failures=skip_failures,
-                warning="Async failure saving embeddings",
-            )
-        )
-
         def save_batch(sample_batch, embeddings_batch):
             with _handle_batch_error(skip_failures, sample_batch):
                 for sample, embedding in zip(sample_batch, embeddings_batch):
                     sample[embeddings_field] = embedding
                     if ctx:
                         ctx.save(sample)
 
-        for sample_batch, imgs in zip(
-            fou.iter_batches(samples, batch_size),
-            data_loader,
+        for submit, (sample_batch, imgs) in fou.async_iterator(
+            zip(
+                fou.iter_batches(samples, batch_size),
+                data_loader,
+            ),
+            limit=10,
+            max_workers=1,
+            skip_failures=skip_failures,
+            warning="Async failure saving embeddings",
         ):
             embeddings_batch = [None] * len(sample_batch)
 

fiftyone/core/utils.py

@@ -3222,3 +3222,60 @@ def submit(*args, **kwargs):
                 if not skip_failures:
                     raise e
                 logger.warning(warning, exc_info=True)
+
+
+def async_iterator(
+    iterator,
+    *,
+    limit,
+    max_workers,
+    skip_failures=False,
+    warning="Async failure",
+):
+    """
+    Wraps an iterable with a thread pool executor and provides a function to add
+    tasks to a background queue. When the background queue has more than
+    max number of allowed tasks, iteration pauses until a task completes.
+
+    Example::
+
+        for submit, item in async_iterator(iterator, limit=4, max_workers=4):
+            submit(task(item))
+
+    Args:
+        iterator: an iterable to consume
+        limit: the maximum number of background tasks to allow before
+            pausing iteration
+        max_workers: the maximum number of workers to use
+        skip_failures (False): whether to skip exceptions raised by tasks
+        warning ("Async failure"): the warning message to log if a task
+            raises an exception and ``skip_failures == True``
+
+    Raises:
+        Exception: if a task raises an exception and ``skip_failures == False``
+    """
+
+    def wait_for_task(future):
+        try:
+            yield future.result()
+        except Exception as e:
+            if not skip_failures:
+                raise e
+            logger.warning(warning, exc_info=True)
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = []
+
+        def submit(*args, **kwargs):
+            future = executor.submit(*args, **kwargs)
+            futures.append(future)
+            return future
+
+        for item in iterator:
+            yield submit, item
+
+            while len(futures) > limit:
+                wait_for_task(futures.pop(0))
+
+        for future in futures:
+            wait_for_task(future)