From 135591835cde4e3b83806abde588b1710f0e8252 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Wed, 26 Aug 2026 21:03:11 +0800 Subject: [PATCH] perf(observations): reduce batch pipeline copies (#1319) --- src/unilab/managers/_noise/noise_cfg.py | 15 ++- src/unilab/managers/observation_manager.py | 72 +++++++++++-- .../test_observation_buffers_noise.py | 102 ++++++++++++++++++ 3 files changed, 179 insertions(+), 10 deletions(-) diff --git a/src/unilab/managers/_noise/noise_cfg.py b/src/unilab/managers/_noise/noise_cfg.py index d5cf2d8ec..13383b018 100644 --- a/src/unilab/managers/_noise/noise_cfg.py +++ b/src/unilab/managers/_noise/noise_cfg.py @@ -68,14 +68,21 @@ def apply(self, data: np.ndarray, *, rng: np.random.Generator | None = None) -> n_min = self._as_array(self.n_min, data.dtype) n_max = self._as_array(self.n_max, data.dtype) - # Generate uniform noise in [0, 1) and scale to [n_min, n_max). + # Generate uniform noise in [0, 1) and transform the generated array + # in place. The pre-existing expression allocated one array for each + # multiply, add, and final data operation; keeping the same float32 + # cast and ufunc order preserves bit-level results while returning the + # transformed noise buffer itself. noise = rng.random(data.shape).astype(data.dtype, copy=False) - noise = noise * (n_max - n_min) + n_min + np.multiply(noise, n_max - n_min, out=noise) + np.add(noise, n_min, out=noise) if self.operation == "add": - return data + noise + np.add(data, noise, out=noise) + return noise elif self.operation == "scale": - return data * noise + np.multiply(data, noise, out=noise) + return noise elif self.operation == "abs": return noise else: diff --git a/src/unilab/managers/observation_manager.py b/src/unilab/managers/observation_manager.py index 7bf1b4dae..8cfa6af60 100644 --- a/src/unilab/managers/observation_manager.py +++ b/src/unilab/managers/observation_manager.py @@ -335,7 +335,7 @@ def _row_env_ids(mask: np.ndarray) -> list[int]: ) # Sanitize (applies to both "warn" and "sanitize" policies). - return np.nan_to_num(tensor, nan=0.0, posinf=0.0, neginf=0.0) + return np.nan_to_num(tensor, copy=False, nan=0.0, posinf=0.0, neginf=0.0) def compute( self, @@ -381,6 +381,17 @@ def compute_group( group_term_names = self._group_obs_term_names[group_name] group_obs: dict[str, np.ndarray] = {} obs_terms = zip(group_term_names, self._group_obs_term_cfgs[group_name], strict=False) + # In the strict default policy a finite result is by far the common + # case. For concatenated groups, scan the assembled output once and + # only inspect individual slices when an error is actually found; this + # retains the per-term diagnostic while removing repeated full scans + # from the hot path. + defer_error_nan_check = ( + self._group_obs_concatenate[group_name] + and not self._group_obs_temporal[group_name] + and group_cfg.nan_check_per_term + and group_cfg.nan_policy == "error" + ) # Reset path (issue #1259 R2): when no term in this group uses delay or # history buffers, everything downstream of the term call is row # independent, so only the reset rows are processed. Term calls and @@ -409,10 +420,28 @@ def compute_group( # NoiseModel.__call__ likewise returns a new array. obs = self._group_obs_class_instances[group_name][term_name](obs) fresh = True - if not row_scoped and not fresh: - # Terms may return backend/command-owned buffers; copy before the - # in-place clip/scale below. Skipped when noise already produced - # a fresh array (issue #1296). + sanitizes_per_term = group_cfg.nan_check_per_term and group_cfg.nan_policy in ( + "warn", + "sanitize", + ) + exposes_term_output = ( + not self._group_obs_concatenate[group_name] + and term_cfg.delay_max_lag == 0 + and term_cfg.history_length == 0 + ) + if ( + not row_scoped + and not fresh + and ( + term_cfg.clip is not None + or term_cfg.scale is not None + or sanitizes_per_term + or exposes_term_output + ) + ): + # Concatenation and temporal buffers already copy their inputs. + # Only take a defensive copy when this pipeline may mutate the + # term or expose it directly to callers. obs = obs.copy() if row_scoped: # Fresh row copy; safe for the in-place clip/scale below. @@ -425,7 +454,11 @@ def compute_group( np.multiply(obs, scale, out=obs) # Check for NaN/Inf before delay/history buffers (per-term checking). - if group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled": + if ( + group_cfg.nan_check_per_term + and group_cfg.nan_policy != "disabled" + and not defer_error_nan_check + ): obs = self._check_and_handle_nans( obs, context=f"{group_name}/{term_name}", @@ -474,6 +507,33 @@ def compute_group( result = np.concatenate( list(group_obs.values()), axis=self._group_obs_concatenate_dim[group_name] ) + if defer_error_nan_check: + finite = np.isfinite(result) + if not finite.all(): + axis = self._group_obs_concatenate_dim[group_name] + axis = axis if axis >= 0 else result.ndim + axis + offset = 0 + for term_name, term_dims in zip( + group_term_names, + self._group_obs_term_dim[group_name], + strict=True, + ): + width = int(term_dims[axis - 1]) + selectors = [slice(None)] * result.ndim + selectors[axis] = slice(offset, offset + width) + selector = tuple(selectors) + if not finite[selector].all(): + # Reuse the established diagnostic path so the + # error still names the first offending term and + # reports reset-row IDs when applicable. + self._check_and_handle_nans( + result[selector], + context=f"{group_name}/{term_name}", + policy=group_cfg.nan_policy, + env_ids=env_ids if row_scoped else None, + ) + break + offset += width # Final check for concatenated result (non-per-term checking). if not group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled": result = self._check_and_handle_nans( diff --git a/tests/managers/test_observation_buffers_noise.py b/tests/managers/test_observation_buffers_noise.py index 0688dbfe7..818fea0a0 100644 --- a/tests/managers/test_observation_buffers_noise.py +++ b/tests/managers/test_observation_buffers_noise.py @@ -95,6 +95,28 @@ def test_noise_configs_use_supplied_generator() -> None: np.testing.assert_array_equal(ConstantNoiseCfg(bias=2.0, operation="abs").apply(data), 2.0) +@pytest.mark.parametrize("operation", ["add", "scale", "abs"]) +def test_uniform_noise_inplace_matches_reference_expression(operation: str) -> None: + data = np.arange(24, dtype=np.float32).reshape(8, 3) + n_min = np.asarray([-0.2, -0.1, -0.05], dtype=np.float32) + n_max = np.asarray([0.3, 0.4, 0.5], dtype=np.float32) + cfg = UniformNoiseCfg(n_min=tuple(n_min), n_max=tuple(n_max), operation=operation) + + reference_rng = np.random.default_rng(1702) + unit = reference_rng.random(data.shape).astype(data.dtype, copy=False) + noise = unit * (n_max - n_min) + n_min + if operation == "add": + expected = data + noise + elif operation == "scale": + expected = data * noise + else: + expected = noise + + actual = cfg.apply(data, rng=np.random.default_rng(1702)) + np.testing.assert_array_equal(actual, expected) + np.testing.assert_array_equal(data, np.arange(24, dtype=np.float32).reshape(8, 3)) + + def test_additive_bias_noise_supports_scalar_terms() -> None: from unilab.managers._noise import NoiseModelWithAdditiveBias @@ -143,6 +165,86 @@ def test_observation_groups_pipeline_order_and_history(fake_env: FakeEnv) -> Non assert manager.get_active_iterable_terms(0)[0][0] == "policy-state" +def test_concatenated_result_owns_each_result_and_protects_term_buffers( + fake_env: FakeEnv, +) -> None: + source = fake_env.obs + manager = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={ + "source": ObservationTermCfg(func=lambda env: env.obs), + "constant": ObservationTermCfg( + func=lambda env: np.ones((env.num_envs, 1), dtype=np.float32) + ), + } + ) + }, + fake_env, + ) + + first = manager.compute(update_history=True)["policy"] + assert isinstance(first, np.ndarray) + first_address = first.ctypes.data + expected_first = np.concatenate((source.copy(), np.ones((fake_env.num_envs, 1))), axis=1) + np.testing.assert_array_equal(first, expected_first) + + fake_env.obs += 100.0 + second = manager.compute(update_history=True)["policy"] + assert isinstance(second, np.ndarray) + assert second.ctypes.data != first_address + np.testing.assert_array_equal(first, expected_first) + np.testing.assert_array_equal(second[:, :2], fake_env.obs) + assert not np.shares_memory(second, fake_env.obs) + + +def test_concatenated_nan_sanitize_does_not_mutate_term_owned_input( + fake_env: FakeEnv, +) -> None: + source = fake_env.obs.copy() + source[0, 0] = np.nan + manager = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={"state": ObservationTermCfg(func=lambda env: source)}, + nan_policy="sanitize", + nan_check_per_term=False, + ) + }, + fake_env, + ) + + result = manager.compute(update_history=True)["policy"] + assert isinstance(result, np.ndarray) + assert np.isfinite(result).all() + assert np.isnan(source[0, 0]) + + +def test_concatenated_nan_error_still_identifies_offending_term(fake_env: FakeEnv) -> None: + def invalid(env: FakeEnv) -> np.ndarray: + result = np.ones((env.num_envs, 1), dtype=np.float32) + result[2, 0] = np.nan + return result + + manager = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={ + "finite": ObservationTermCfg(func=lambda env: env.obs), + "invalid": ObservationTermCfg(func=invalid), + } + ) + }, + fake_env, + ) + + with pytest.raises( + ValueError, + match=r"NaN detected.*'policy/invalid'.*environments: \[2\]", + ): + manager.compute(update_history=True) + + def test_observation_noise_model_delay_and_seed_reproducibility() -> None: cfg = { "policy": ObservationGroupCfg(