From 3d4c3037afb327f997eb414089a5b8bec99f9729 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Tue, 8 Sep 2026 13:24:37 +0000 Subject: [PATCH 1/3] Scale a stacked video context on several threads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A video-conditioned policy sends a temporal stack — 25 frames per camera for a 1.6 s window at 15 Hz — and `RestrictImageSize` scaled all 50 of them one after another, on the thread that drives the arm. The frames are independent and Pillow drops the GIL for a resize, so they do not have to queue. Measured on one round's real 1280x720 frames, scaling a 25-frame two-camera stack to 512x288: 267 ms serial, 83 ms threaded. The output is byte-identical either way — a test asserts a stack over the bar equals the same frames scaled one at a time, so the model is shown exactly what it was before. A stack shorter than four frames stays serial: below that the pool costs more to raise than it saves. Ticket: Positronic-Robotics/internal#1168 #refs --- positronic/policy/codec.py | 16 +++++++++++++++- positronic/policy/tests/test_layers.py | 12 ++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 78c0f6bb2..73216b070 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -10,6 +10,7 @@ """ import collections.abc as cabc +from concurrent.futures import ThreadPoolExecutor from dataclasses import replace from functools import partial from typing import Any, final, overload @@ -459,6 +460,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 +476,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 +484,15 @@ def _restrict(self, key: str, value: Any) -> Any: return type(value)(self._restrict(key, v) for v in value) return value + def _scaled_frames(self, stack: np.ndarray) -> list[np.ndarray]: + """Every frame of one stack, scaled. Threaded above ``_PARALLEL_FROM``: the frames are + independent and Pillow drops the GIL for a resize.""" + if len(stack) < self._PARALLEL_FROM: + return [_scaled(frame, self._width, self._height) for frame in stack] + scale = partial(_scaled, width=self._width, height=self._height) + with ThreadPoolExecutor(max_workers=min(len(stack), self._MAX_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..6e6bb2832 100644 --- a/positronic/policy/tests/test_layers.py +++ b/positronic/policy/tests/test_layers.py @@ -636,6 +636,18 @@ 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_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) From 26b238bd00bf39369c037e283f75788ddbb90c6d Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Tue, 8 Sep 2026 13:59:59 +0000 Subject: [PATCH 2/3] Bound the resize pool by the CPUs the process may use A pool sized only by frame count raises workers on a rig whose process has one usable CPU, where threading wins nothing. `_workers` reads `os.process_cpu_count()`, so a cgroup quota or an affinity mask bounds it rather than the machine's core count, and one usable CPU takes the serial path. Ticket: Positronic-Robotics/internal#1168 #refs --- positronic/policy/codec.py | 21 ++++++++++++++++----- positronic/policy/tests/test_layers.py | 12 ++++++++++++ 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 73216b070..7fa2b1348 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -10,6 +10,7 @@ """ import collections.abc as cabc +import os from concurrent.futures import ThreadPoolExecutor from dataclasses import replace from functools import partial @@ -484,13 +485,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 + # FOOTGUN: `cpu_count` reports the machine, not a cgroup quota, so a container pinned to one + # core of many still reads as many and takes the threaded path. + return max(1, min(frames, self._MAX_WORKERS, os.cpu_count() or 1)) + def _scaled_frames(self, stack: np.ndarray) -> list[np.ndarray]: - """Every frame of one stack, scaled. Threaded above ``_PARALLEL_FROM``: the frames are - independent and Pillow drops the GIL for a resize.""" - if len(stack) < self._PARALLEL_FROM: - return [_scaled(frame, self._width, self._height) for frame in stack] + """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) - with ThreadPoolExecutor(max_workers=min(len(stack), self._MAX_WORKERS)) as pool: + 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): diff --git a/positronic/policy/tests/test_layers.py b/positronic/policy/tests/test_layers.py index 6e6bb2832..415253b02 100644 --- a/positronic/policy/tests/test_layers.py +++ b/positronic/policy/tests/test_layers.py @@ -1,5 +1,6 @@ """Unit tests for Layer composition, ChunkedSchedule, TemporalStack, and the policy-pipeline algebra.""" +import os from typing import Any import numpy as np @@ -644,6 +645,17 @@ def test_a_threaded_stack_scales_to_the_same_pixels_as_one_thread(self): 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(os, 'cpu_count', lambda: 1) + codec = RestrictImageSize(64, 48) + assert codec._workers(codec._PARALLEL_FROM + 4) == 1 + + def test_the_pool_is_bounded_by_the_core_count(self, monkeypatch): + monkeypatch.setattr(os, 'cpu_count', 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) From f0da1ee47ff2e4930b4f9af55d53fd9dbf44561b Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Tue, 8 Sep 2026 14:08:18 +0000 Subject: [PATCH 3/3] Size the resize pool from the process's CPU affinity `os.cpu_count()` reports the host, so a process pinned to one core of many still read as many and took the threaded path. `_usable_cpus` reads the affinity mask where the platform publishes one. A cgroup CPU quota is still invisible to both, and the docstring says so. Ticket: Positronic-Robotics/internal#1168 #refs --- positronic/policy/codec.py | 14 +++++++++++--- positronic/policy/tests/test_layers.py | 8 ++++---- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 7fa2b1348..75e35d96b 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -438,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) @@ -490,9 +500,7 @@ def _workers(self, frames: int) -> int: usable CPU, and below ``_PARALLEL_FROM`` it costs more to raise than the frames take.""" if frames < self._PARALLEL_FROM: return 1 - # FOOTGUN: `cpu_count` reports the machine, not a cgroup quota, so a container pinned to one - # core of many still reads as many and takes the threaded path. - return max(1, min(frames, self._MAX_WORKERS, os.cpu_count() or 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 diff --git a/positronic/policy/tests/test_layers.py b/positronic/policy/tests/test_layers.py index 415253b02..d8070104d 100644 --- a/positronic/policy/tests/test_layers.py +++ b/positronic/policy/tests/test_layers.py @@ -1,6 +1,5 @@ """Unit tests for Layer composition, ChunkedSchedule, TemporalStack, and the policy-pipeline algebra.""" -import os from typing import Any import numpy as np @@ -11,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 @@ -647,12 +647,12 @@ def test_a_threaded_stack_scales_to_the_same_pixels_as_one_thread(self): 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(os, 'cpu_count', lambda: 1) + 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_core_count(self, monkeypatch): - monkeypatch.setattr(os, 'cpu_count', lambda: 2) + 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