Scale a stacked video context on several threads

positronic/policy/codec.py

@@ -10,6 +10,8 @@
 """
 
 import collections.abc as cabc
+import os
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import replace
 from functools import partial
 from typing import Any, final, overload
@@ -436,6 +438,16 @@ def to_spec(self):
         return {'name': self.WIRE_NAME}
 
 
+def _usable_cpus() -> int:
+    """CPUs this process may run on — its affinity mask where the platform publishes one, else the
+    host's count. FOOTGUN: neither reads a cgroup CPU quota, so a container limited by `--cpus` and
+    not by a mask still reads the host's cores.
+    """
+    if hasattr(os, 'sched_getaffinity'):
+        return len(os.sched_getaffinity(0))
+    return os.cpu_count() or 1
+
+
 def _scaled(image: np.ndarray, width: int, height: int) -> np.ndarray:
     h, w = image.shape[:2]
     scale = min(1.0, width / w, height / h)
@@ -459,6 +471,10 @@ class RestrictImageSize(Codec):
 
     WIRE_NAME = 'restrict_image_size'
 
+    # Below this a stack scales quicker in one thread than a pool costs to raise.
+    _PARALLEL_FROM = 4
+    _MAX_WORKERS = 8
+
     def __init__(self, width: int = 640, height: int = 640):
         self._width = width
         self._height = height
@@ -471,14 +487,31 @@ def _restrict(self, key: str, value: Any) -> Any:
         if isinstance(value, np.ndarray) and value.ndim in (3, 4) and value.shape[-1] == 3:
             # A TemporalStack emits a (T, H, W, 3) stack, so bound each frame rather than the stack's first axis.
             if value.ndim == 4:
-                return np.stack([_scaled(frame, self._width, self._height) for frame in value])
+                return np.stack(self._scaled_frames(value))
             return _scaled(value, self._width, self._height)
         if isinstance(value, cabc.Mapping):
             return {k: self._restrict(k, v) for k, v in value.items()}
         if isinstance(value, list | tuple):
             return type(value)(self._restrict(key, v) for v in value)
         return value
 
+    def _workers(self, frames: int) -> int:
+        """Threads to scale ``frames`` on. One means the serial path: a pool wins nothing on a single
+        usable CPU, and below ``_PARALLEL_FROM`` it costs more to raise than the frames take."""
+        if frames < self._PARALLEL_FROM:
+            return 1
+        return max(1, min(frames, self._MAX_WORKERS, _usable_cpus()))
+
+    def _scaled_frames(self, stack: np.ndarray) -> list[np.ndarray]:
+        """Every frame of one stack, scaled. Pillow drops the GIL for a resize and the frames are
+        independent, so more than one may run at a time."""
+        scale = partial(_scaled, width=self._width, height=self._height)
+        workers = self._workers(len(stack))
+        if workers == 1:
+            return [scale(frame) for frame in stack]
+        with ThreadPoolExecutor(max_workers=workers) as pool:
+            return list(pool.map(scale, stack))
+
     def decode(self, data):
         return data
 

positronic/policy/tests/test_layers.py

@@ -10,6 +10,7 @@
 from positronic.drivers.roboarm import keys as roboarm_keys
 from positronic.drivers.roboarm.command import Impedance, JointDelta
 from positronic.geom import Rotation, Transform3D
+from positronic.policy import codec as codec_module
 from positronic.policy import spec
 from positronic.policy.action import AbsoluteJointsAction, AbsolutePositionAction, IKJointsAction, JointDeltaAction
 from positronic.policy.base import Layer, Policy, Session
@@ -636,6 +637,29 @@ def test_stacked_frames_are_bounded_per_frame(self):
         stack = np.zeros((3, 480, 640, 3), dtype=np.uint8)
         assert RestrictImageSize(64, 48).encode({'cam': stack})['cam'].shape == (3, 48, 64, 3)
 
+    def test_a_threaded_stack_scales_to_the_same_pixels_as_one_thread(self):
+        """A stack over the parallel bar scales to the same pixels as the frames taken one at a time."""
+        rng = np.random.default_rng(0)
+        stack = rng.integers(0, 256, size=(RestrictImageSize._PARALLEL_FROM + 4, 480, 640, 3), dtype=np.uint8)
+        codec = RestrictImageSize(64, 48)
+        one_at_a_time = np.stack([codec.encode({'cam': frame})['cam'] for frame in stack])
+        np.testing.assert_array_equal(codec.encode({'cam': stack})['cam'], one_at_a_time)
+
+    def test_a_single_usable_cpu_stays_serial(self, monkeypatch):
+        """A pool wins nothing on a single core, and costs threads to raise."""
+        monkeypatch.setattr(codec_module, '_usable_cpus', lambda: 1)
+        codec = RestrictImageSize(64, 48)
+        assert codec._workers(codec._PARALLEL_FROM + 4) == 1
+
+    def test_the_pool_is_bounded_by_the_cpus_the_process_may_run_on(self, monkeypatch):
+        monkeypatch.setattr(codec_module, '_usable_cpus', lambda: 2)
+        codec = RestrictImageSize(64, 48)
+        assert codec._workers(codec._MAX_WORKERS * 4) == 2
+
+    def test_a_stack_under_the_parallel_bar_still_scales(self):
+        stack = np.zeros((RestrictImageSize._PARALLEL_FROM - 1, 480, 640, 3), dtype=np.uint8)
+        assert RestrictImageSize(64, 48).encode({'cam': stack})['cam'].shape[1:] == (48, 64, 3)
+
     def test_nested_images_are_reached(self):
         result = RestrictImageSize(64, 48).encode({'video': {'cam': _image(480, 640)}, 'seq': [_image(480, 640)]})
         assert result['video']['cam'].shape == (48, 64, 3)