perf(core): run visual participants (actor, scene) in parallel within a shared decode stream

src/vidxp/capabilities/visual.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
+import contextvars
+import queue
+import threading
 from dataclasses import dataclass
 from time import perf_counter
-from typing import Any, Protocol, Sequence
+from typing import Any, Callable, Protocol, Sequence
 
 from vidxp.capabilities.contracts import CapabilityIndexResult
 from vidxp.capabilities.registry import CapabilityRegistry
@@ -147,59 +150,177 @@ def _consume_visual_stream(
     progress: ProgressCallback | None,
     timings: dict[str, float],
 ) -> FrameStreamStats:
+    participant_queues = {
+        p.name: queue.Queue(maxsize=4)
+        for p in participants
+    }
+    errors: list[Exception] = []
+    errors_lock = threading.Lock()
     stream_stats = FrameStreamStats()
-    stream = iter(
-        iter_frame_batches(
-            source.path,
-            samplings=tuple(
-                participant.sampling for participant in participants
-            ),
-            batch_size=max(
-                participant.processor.batch_size(config)
-                for participant in participants
-            ),
-            cancellation=cancellation,
-            stats=stream_stats,
-        )
-    )
-    while True:
-        stream_started = perf_counter()
-        try:
-            samples = next(stream)
-        except StopIteration:
-            timings["frame_stream"] += perf_counter() - stream_started
-            break
-        rgb_samples = _rgb_samples(samples)
-        timings["frame_stream"] += perf_counter() - stream_started
-
-        for participant in participants:
-            participant_samples = [
-                sample
-                for sample in rgb_samples
-                if _is_participant_sample(sample, participant)
-            ]
-            if not participant_samples:
-                continue
-            processor_started = perf_counter()
-            participant.processor.process(
-                participant_samples,
-                state=participant.state,
-                info=info,
-                config=config,
-                storage=storage,
-                cancellation=cancellation,
+    done = threading.Event()
+    stop = threading.Event()
+    progress_lock = threading.Lock()
+    batch_ends: list[int] = []
+    finished_batches = dict.fromkeys(participant_queues, 0)
+    reported = 0
+
+    def _record_error(exc: Exception) -> None:
+        with errors_lock:
+            errors.append(exc)
+        stop.set()
+
+    def _finish_batch(name: str) -> None:
+        nonlocal reported
+        # Reported under the lock so updates from different participant
+        # threads arrive in order.
+        with progress_lock:
+            finished_batches[name] += 1
+            completed = min(finished_batches.values())
+            current = batch_ends[completed - 1] if completed else 0
+            if current <= reported:
+                return
+            reported = current
+            report_progress(
+                progress,
+                "visual_indexing",
+                "Indexing the shared sampled-frame stream.",
+                current,
+                expected,
             )
-            timings[participant.name] += (
-                perf_counter() - processor_started
+
+    def _decode_worker() -> None:
+        try:
+            stream = iter(
+                iter_frame_batches(
+                    source.path,
+                    samplings=tuple(
+                        p.sampling for p in participants
+                    ),
+                    batch_size=max(
+                        p.processor.batch_size(config)
+                        for p in participants
+                    ),
+                    cancellation=cancellation,
+                    stats=stream_stats,
+                )
             )
+            while not stop.is_set():
+                stream_started = perf_counter()
+                try:
+                    samples = next(stream)
+                except StopIteration:
+                    timings["frame_stream"] += (
+                        perf_counter() - stream_started
+                    )
+                    break
+                rgb_samples = _rgb_samples(samples)
+                timings["frame_stream"] += (
+                    perf_counter() - stream_started
+                )
+                with progress_lock:
+                    batch_ends.append(
+                        (batch_ends[-1] if batch_ends else 0)
+                        + len(rgb_samples)
+                    )
+                for participant_queue in participant_queues.values():
+                    while True:
+                        try:
+                            participant_queue.put(
+                                rgb_samples, timeout=0.2
+                            )
+                            break
+                        except queue.Full:
+                            if stop.is_set():
+                                return
+                            cancellation.raise_if_cancelled()
+        except Exception as exc:
+            _record_error(exc)
+        finally:
+            done.set()
+            for q in participant_queues.values():
+                try:
+                    q.put_nowait(None)
+                except queue.Full:
+                    pass
+
+    def _participant_worker(
+        participant: _Participant,
+        work_queue: queue.Queue,
+    ) -> None:
+        try:
+            while not stop.is_set():
+                cancellation.raise_if_cancelled()
+                try:
+                    batch = work_queue.get(timeout=0.2)
+                except queue.Empty:
+                    if done.is_set():
+                        break
+                    continue
+                if batch is None or stop.is_set():
+                    break
+                participant_samples = [
+                    sample
+                    for sample in batch
+                    if _is_participant_sample(sample, participant)
+                ]
+                if participant_samples:
+                    started = perf_counter()
+                    participant.processor.process(
+                        participant_samples,
+                        state=participant.state,
+                        info=info,
+                        config=config,
+                        storage=storage,
+                        cancellation=cancellation,
+                    )
+                    timings[participant.name] += (
+                        perf_counter() - started
+                    )
+                _finish_batch(participant.name)
+        except Exception as exc:
+            _record_error(exc)
+
+    def _in_context(
+        target: Callable[..., None],
+        *args: Any,
+    ) -> Callable[[], None]:
+        thread_context = contextvars.copy_context()
+
+        def runner() -> None:
+            thread_context.run(target, *args)
 
-        report_progress(
-            progress,
-            "visual_indexing",
-            "Indexing the shared sampled-frame stream.",
-            stream_stats.frames_materialized,
-            expected,
+        return runner
+
+    decode_thread = threading.Thread(
+        target=_in_context(_decode_worker),
+        name="vidxp-visual-decode",
+    )
+    participant_threads = [
+        threading.Thread(
+            target=_in_context(
+                _participant_worker, p, participant_queues[p.name]
+            ),
+            name=f"vidxp-visual-{p.name}",
         )
+        for p in participants
+    ]
+
+    threads = [decode_thread, *participant_threads]
+    try:
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+    except BaseException:
+        stop.set()
+        for t in threads:
+            if t.is_alive():
+                t.join()
+        raise
+
+    if errors:
+        raise errors[0]
+
     return stream_stats
 
 
@@ -248,7 +369,9 @@ def index_visuals(
         raise ValueError("Visual indexing requires a video path.")
 
     selected = tuple(
-        config.enabled_modalities if modalities is None else modalities
+        dict.fromkeys(
+            config.enabled_modalities if modalities is None else modalities
+        )
     )
     if not selected:
         raise ValueError("At least one visual capability must be selected.")