FOEPD-2119: Use SaveContext to write to DB in background thread

fiftyone/core/collections.py

@@ -7,6 +7,7 @@
 """
 
 from collections import defaultdict, deque
+from concurrent.futures import ThreadPoolExecutor
 from copy import copy
 from datetime import datetime
 from operator import itemgetter
@@ -18,6 +19,7 @@
 from packaging.version import Version
 import random
 import string
+import threading
 import timeit
 import warnings
 
@@ -79,6 +81,28 @@ def registrar(func):
 aggregation = _make_registrar()
 
 
+class _DummyFuture:
+    def __init__(self, *, value=None):
+        self.value = value
+
+    def result(self):
+        return self.value
+
+    def done(self):
+        return True
+
+
+class _DummyExecutor:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    def submit(self, fn, *args, **kwargs):
+        return _DummyFuture(value=fn(*args, **kwargs))
+
+
 class SaveContext(object):
     """Context that saves samples from a collection according to a configurable
     batching strategy.
@@ -101,13 +125,16 @@ class SaveContext(object):
             -   ``"latency"``: a target latency, in seconds, between saves
 
             By default, ``fo.config.default_batcher`` is used
+        async_writes (False): whether to perform batch writes asynchronously in
+            a background thread
     """
 
     def __init__(
         self,
         sample_collection,
         batch_size=None,
         batching_strategy=None,
+        async_writes=False,
     ):
         batch_size, batching_strategy = fou.parse_batching_strategy(
             batch_size=batch_size, batching_strategy=batching_strategy
@@ -132,6 +159,19 @@ def __init__(
         self._encoding_ratio = 1.0
         self._last_time = None
 
+        self.samples_lock = threading.Lock()
+        self.frames_lock = threading.Lock()
+        self.batch_ids_lock = threading.Lock()
+        self.reloading_lock = threading.Lock()
+
+        self.executor = (
+            # Using more than one worker will introduce race conditions in the state preserved between DB writes
+            ThreadPoolExecutor(max_workers=1)
+            if async_writes
+            else _DummyExecutor()
+        )
+        self.futures = []
+
     def __enter__(self):
         if self._batching_strategy == "static":
             self._curr_batch_size = 0
@@ -140,11 +180,32 @@ def __enter__(self):
         elif self._batching_strategy == "latency":
             self._last_time = timeit.default_timer()
 
+        self.executor.__enter__()
         return self
 
     def __exit__(self, *args):
         self._save_batch()
 
+        error = None
+        try:
+            # Loop-drain self.futures so any submissions triggered by
+            # self._save_batch() are awaited.
+            while self.futures:
+                futures = self.futures
+                self.futures = []
+                for future in futures:
+                    try:
+                        future.result()
+                    except Exception as e:
+                        if error is None:
+                            error = e
+            self.futures.clear()
+        finally:
+            self.executor.__exit__(*args)
+
+        if error and (not args or args[0] is None):
+            raise error
+
     def save(self, sample):
         """Registers the sample for saving in the next batch.
 
@@ -162,16 +223,20 @@ def save(self, sample):
         updated = sample_ops or frame_ops
 
         if sample_ops:
-            self._sample_ops.extend(sample_ops)
+            with self.samples_lock:
+                self._sample_ops.extend(sample_ops)
 
         if frame_ops:
-            self._frame_ops.extend(frame_ops)
+            with self.frames_lock:
+                self._frame_ops.extend(frame_ops)
 
         if updated and self._is_generated:
-            self._batch_ids.append(sample.id)
+            with self.batch_ids_lock:
+                self._batch_ids.append(sample.id)
 
         if updated and isinstance(sample, fosa.SampleView):
-            self._reload_parents.append(sample)
+            with self.reloading_lock:
+                self._reload_parents.append(sample)
 
         if self._batching_strategy == "static":
             self._curr_batch_size += 1
@@ -200,24 +265,31 @@ def save(self, sample):
                 self._save_batch()
                 self._last_time = timeit.default_timer()
 
-    def _save_batch(self):
+    def _do_save_batch(self):
         encoded_size = -1
         if self._sample_ops:
+            with self.samples_lock:
+                sample_ops = self._sample_ops.copy()
+                self._sample_ops.clear()
             res = foo.bulk_write(
-                self._sample_ops,
+                sample_ops,
                 self._sample_coll,
                 ordered=False,
                 batcher=False,
             )[0]
             encoded_size += res.bulk_api_result.get("nBytes", 0)
-            self._sample_ops.clear()
 
         if self._frame_ops:
+            with self.frames_lock:
+                frame_ops = self._frame_ops.copy()
+                self._frame_ops.clear()
             res = foo.bulk_write(
-                self._frame_ops, self._frame_coll, ordered=False, batcher=False
+                frame_ops,
+                self._frame_coll,
+                ordered=False,
+                batcher=False,
             )[0]
             encoded_size += res.bulk_api_result.get("nBytes", 0)
-            self._frame_ops.clear()
 
         self._encoding_ratio = (
             self._curr_batch_size_bytes / encoded_size
@@ -226,14 +298,32 @@ def _save_batch(self):
         )
 
         if self._batch_ids and self._is_generated:
-            self.sample_collection._sync_source(ids=self._batch_ids)
-            self._batch_ids.clear()
+            with self.batch_ids_lock:
+                batch_ids = self._batch_ids.copy()
+                self._batch_ids.clear()
+            self.sample_collection._sync_source(ids=batch_ids)
 
         if self._reload_parents:
-            for sample in self._reload_parents:
+            with self.reloading_lock:
+                reload_parents = self._reload_parents.copy()
+                self._reload_parents.clear()
+            for sample in reload_parents:
                 sample._reload_parents()
 
-            self._reload_parents.clear()
+    def _save_batch(self):
+        pending = []
+        for future in self.futures:
+            if not future.done():
+                pending.append(future)
+            else:
+                try:
+                    future.result()
+                except Exception:
+                    pending.append(future)  # re-raise in __exit__
+        self.futures = pending
+
+        future = self.executor.submit(self._do_save_batch)
+        self.futures.append(future)
 
 
 class SampleCollection(object):

fiftyone/core/models.py

@@ -471,7 +471,12 @@ 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))
+        ctx = context.enter_context(
+            foc.SaveContext(
+                samples,
+                async_writes=True,
+            )
+        )
 
         for sample_batch, imgs in zip(
             fou.iter_batches(samples, batch_size),
@@ -1207,7 +1212,12 @@ def _compute_image_embeddings_data_loader(
     with contextlib.ExitStack() as context:
         pb = context.enter_context(fou.ProgressBar(samples, progress=progress))
         if embeddings_field is not None:
-            ctx = context.enter_context(foc.SaveContext(samples))
+            ctx = context.enter_context(
+                foc.SaveContext(
+                    samples,
+                    async_writes=True,
+                )
+            )
         else:
             ctx = None