diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 78c0f6bb2..75e35d96b 100644 --- a/positronic/policy/codec.py +++ b/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,7 +487,7 @@ 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()} @@ -479,6 +495,23 @@ def _restrict(self, key: str, value: Any) -> Any: 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 diff --git a/positronic/policy/tests/test_layers.py b/positronic/policy/tests/test_layers.py index 28a5d1007..d8070104d 100644 --- a/positronic/policy/tests/test_layers.py +++ b/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)