From 1a1c588ee992cc5641b23dc373762533147539cd Mon Sep 17 00:00:00 2001 From: zilch Date: Fri, 2 Oct 2026 17:55:35 +0800 Subject: [PATCH] perf: eliminate redundant numba compilation and split pytest CI (#94) --- .github/workflows/README.md | 68 +++++++++++++++++-- .github/workflows/tests.yml | 28 +++++++- .../src/motrix_env_core/numba/fingerprint.py | 40 ++++++++++- .../src/motrix_env_core/numba/kernel.py | 14 +++- .../numba/manager/compiler/compiler.py | 37 ++++++---- .../tests/test_function_fingerprint.py | 44 ++++++++++++ motrix_env_core/tests/test_kernel_clone.py | 48 +++++++++++++ motrix_envs/tests/test_wbt_numba.py | 21 ++++++ pyproject.toml | 7 ++ test/test_all_envs.py | 4 ++ 10 files changed, 288 insertions(+), 23 deletions(-) create mode 100644 motrix_env_core/tests/test_kernel_clone.py diff --git a/.github/workflows/README.md b/.github/workflows/README.md index bf221b43..cdbb0de1 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -16,10 +16,70 @@ access to repository contents and uses no repository secrets. ### File: `tests.yml` -Runs the workspace test suite on pushes to `main` or `stable`, pull requests, and manual -dispatch. It installs the uv workspace with Git LFS assets and executes -`pytest`. The workflow intentionally runs on Linux, which is the supported -public CI environment for the simulator and training dependencies. +Runs on pull requests, pushes to `main` or `stable`, manual dispatch, and daily at 02:00 UTC +(10:00 China Standard Time). Scheduled runs use the default branch. Each job installs the uv +workspace with Git LFS assets on Linux, the supported public CI environment. + +PR and push runs have two independent, parallel jobs: + +- **fast**: `pytest -m "not slow and not integration and not numba"`. +- **numba**: `pytest -m "not integration and (slow or numba)"`. This includes the two expensive + WBT kernel regressions on every PR, so failures are not deferred until the next day. + The entire `test_wbt_numba.py` module stays together: excluding only the first two expensive + tests shifts first-use compilation into later tests rather than making the fast suite cheap. + +These selections partition all non-integration tests, including future `slow` tests. They run on +all PRs rather than relying on path filters that might miss a transitive dependency change. +A failure in either job does not cancel the other job. Both should be required checks for merging; +replace any branch-protection requirement for the old `pytest` job with these two checks. + +Nightly and manual dispatch run the **full** suite, including +`test/test_all_envs.py::test_all_demos`. The full-environment subprocess smoke test is excluded +from PR/push jobs. Each job reports the 20 slowest phases and uploads JUnit results, including on +failure. Numba caches have per-suite write keys to avoid parallel jobs competing for one cache. + +Local commands (activate `.venv` first): + +```bash +# Quick feedback; excludes expensive kernel regressions and full-environment smoke. +python -m pytest -m "not slow and not integration and not numba" +# Dedicated compilation/runtime regressions. +python -m pytest -m "not integration and (slow or numba)" +# Only the full-environment integration tests. +python -m pytest -m integration +# Full coverage remains the default; no tests are silently excluded by addopts. +python -m pytest +``` + +Markers are registered in the root `pyproject.toml`: `slow` describes cost, `numba` identifies +specialized kernel regressions, and `integration` identifies tests reserved for full runs. +Not every test importing Numba needs the `numba` marker: inexpensive core contract tests remain +in the fast suite. + +For new Numba environments, reuse small synthetic-manager tests for shared runtime contracts, +keep task-specific smoke tests minimal, and retain a small set of representative real-environment +regressions. The existing nightly smoke test automatically covers all registered environments. +Do not duplicate a full expensive rollout for every preset. If the dedicated job grows too large, +measure `--durations` and shard by environment family; do not silently drop kernel coverage from PRs. + +#### Diagnosing Numba startup costs + +Compare a fresh `NUMBA_CACHE_DIR` with a second process using the same directory, then compare +another environment in the same process. `NUMBA_DEBUG_CACHE=1` reports actual specialization +cache loads/saves. Manager INFO logs time each evaluate/observe/reset specialization; DEBUG +logs include signatures. Dispatcher preparation is not the actual compilation phase. + +Precompilation uses the real input types, and warmup copies preserve strides and readonly flags, +so warmup and reset/step share signatures. Function fingerprints use explicit code fields rather +than marshal serialization, which can change after Numba inspects an unchanged code object. +Keep these invariants intact when extending the compiler: a warmup must not introduce new +signatures, and constructing the same configuration after execution must reuse the same plan. + +A local G1 WBT play-mode benchmark (`num_envs=2`, init and one step; interpreter imports excluded) +measured about 29 seconds with an empty cache, 0.27 seconds in a second process with disk cache, +and 0.07 seconds for another environment using the in-process cache after these fixes. +The original cold benchmark took about 57 seconds and compiled each kernel twice. +These are diagnostic measurements, not CI timing thresholds; hardware and cache state vary. ### File: `codeql.yml` diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 25fac95e..9705065e 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -5,18 +5,26 @@ on: branches: [main, stable] pull_request: workflow_dispatch: + schedule: + # Daily at 02:00 UTC (10:00 China Standard Time), on the default branch. + - cron: "0 2 * * *" permissions: contents: read concurrency: - group: tests-${{ github.ref }} + group: tests-${{ github.event_name }}-${{ github.ref }} cancel-in-progress: true jobs: pytest: + name: pytest (${{ matrix.suite }}) runs-on: ubuntu-latest timeout-minutes: 45 + strategy: + fail-fast: false + matrix: + include: ${{ fromJSON((github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') && '[{"suite":"full","selection":""}]' || '[{"suite":"fast","selection":"not slow and not integration and not numba"},{"suite":"numba","selection":"not integration and (slow or numba)"}]') }} steps: - name: Resolve numba cache dir run: | @@ -31,8 +39,9 @@ jobs: uses: actions/cache@v4 with: path: ~/.cache/motrixlab - key: numba-kernel-${{ runner.os }}-${{ github.sha }} + key: numba-kernel-${{ runner.os }}-${{ matrix.suite }}-${{ github.sha }} restore-keys: | + numba-kernel-${{ runner.os }}-${{ matrix.suite }}- numba-kernel-${{ runner.os }}- - name: Checkout repository uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4 @@ -52,4 +61,17 @@ jobs: run: sh install.sh --all - name: Run tests - run: .venv/bin/python -m pytest -q --disable-warnings --maxfail=1 + env: + PYTEST_SELECTION: ${{ matrix.selection }} + run: | + .venv/bin/python -m pytest -q --strict-markers --disable-warnings \ + --durations=20 --junitxml=test-results/${{ matrix.suite }}.xml \ + -m "$PYTEST_SELECTION" + + - name: Upload test results + if: always() + uses: actions/upload-artifact@v4 + with: + name: pytest-${{ matrix.suite }} + path: test-results/ + if-no-files-found: warn diff --git a/motrix_env_core/src/motrix_env_core/numba/fingerprint.py b/motrix_env_core/src/motrix_env_core/numba/fingerprint.py index 4db78dfa..49090d42 100644 --- a/motrix_env_core/src/motrix_env_core/numba/fingerprint.py +++ b/motrix_env_core/src/motrix_env_core/numba/fingerprint.py @@ -5,13 +5,49 @@ import hashlib import inspect -import marshal from collections.abc import Callable +from types import CodeType from typing import Any import numpy as np +def _stable_constant(value: Any) -> tuple[Any, ...]: + if isinstance(value, CodeType): + return ("code", _code_fingerprint_parts(value)) + if isinstance(value, tuple): + return ("tuple", tuple(_stable_constant(item) for item in value)) + if isinstance(value, frozenset): + items = sorted((_stable_constant(item) for item in value), key=repr) + return ("frozenset", tuple(items)) + if value is None or isinstance(value, (bool, int, float, complex, str, bytes)): + return (type(value).__name__, value) + return (type(value).__module__, type(value).__qualname__, repr(value)) + + +def _code_fingerprint_parts(code: CodeType) -> tuple[Any, ...]: + # marshal includes string interning/reference flags that can change when + # Numba inspects the same code object. Hash immutable code fields instead. + return ( + code.co_argcount, + code.co_posonlyargcount, + code.co_kwonlyargcount, + code.co_nlocals, + code.co_stacksize, + code.co_flags, + code.co_code, + tuple(_stable_constant(value) for value in code.co_consts), + code.co_names, + code.co_varnames, + code.co_freevars, + code.co_cellvars, + code.co_filename, + code.co_name, + code.co_firstlineno, + code.co_lnotab, + ) + + def function_fingerprint(function: Callable[..., Any]) -> str: """Hash function code, defaults, closures, and referenced global helpers. @@ -33,7 +69,7 @@ def function_fingerprint(function: Callable[..., Any]) -> str: seen.add(id(current)) code = current.__code__ hasher.update(f"{current.__module__}.{current.__qualname__}\0".encode()) - hasher.update(marshal.dumps(code)) + hasher.update(repr(_code_fingerprint_parts(code)).encode()) references = [(name, current.__globals__[name]) for name in code.co_names if name in current.__globals__] references.extend((f"default[{i}]", value) for i, value in enumerate(current.__defaults__ or ())) references.extend(sorted((current.__kwdefaults__ or {}).items())) diff --git a/motrix_env_core/src/motrix_env_core/numba/kernel.py b/motrix_env_core/src/motrix_env_core/numba/kernel.py index e672c524..4ba54acb 100644 --- a/motrix_env_core/src/motrix_env_core/numba/kernel.py +++ b/motrix_env_core/src/motrix_env_core/numba/kernel.py @@ -66,7 +66,19 @@ def validate_kernel_context( def clone_kernel_value(value: T) -> T: """Recursively clone arrays while preserving task-specific tuple types.""" if isinstance(value, np.ndarray): - return cast(T, value.copy()) + # Warmup must use the same Numba array type as real execution: copy() + # normalizes strided views to C layout and drops the readonly flag. + lower = ( + sum(min(0, (size - 1) * stride) for size, stride in zip(value.shape, value.strides)) if value.size else 0 + ) + upper = ( + sum(max(0, (size - 1) * stride) for size, stride in zip(value.shape, value.strides)) if value.size else 0 + ) + storage = np.empty(max(value.itemsize, upper - lower + value.itemsize), dtype=np.uint8) + copied = np.ndarray(value.shape, dtype=value.dtype, buffer=storage, offset=-lower, strides=value.strides) + copied[...] = value + copied.flags.writeable = value.flags.writeable + return cast(T, copied) if isinstance(value, tuple): items = [clone_kernel_value(item) for item in value] if hasattr(type(value), "_fields"): diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py index e4b8de34..264f1072 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py @@ -19,7 +19,6 @@ from numba.extending import register_jitable from motrix_env_core.numba.fingerprint import function_fingerprint -from motrix_env_core.numba.kernel import clone_kernel_value from motrix_env_core.numba.kernel_data import ( KernelDataLayout, KernelDataLowering, @@ -891,9 +890,9 @@ def precompile_reset_kernel(self, inputs: tuple[Any, ...]) -> None: task = self._env._task_program if task is None: return - reset_args = (clone_kernel_value(inputs), np.arange(2, dtype=np.int64), self._env._sim_reset_runtime.buffers) + reset_args = (inputs, np.arange(2, dtype=np.int64), self._env._sim_reset_runtime.buffers) try: - task.reset_kernel.compile(tuple(numba.typeof(arg) for arg in reset_args)) + self._compile_specialization("reset", task.reset_kernel, reset_args) except ( AttributeError, EOFError, @@ -915,17 +914,29 @@ def _compile_specializations_once(self, inputs: tuple[Any, ...]) -> None: # Terms are not compiled standalone: their dispatch bodies are inlined # into the fused kernels below, which type-checks the whole plan in one # compilation instead of paying a separate dispatcher compile per term. - warmup_args = tuple( - clone_kernel_value(value) - for value in (inputs, task.reward_weights, self._env._kernel_buffers, self._env._kernel_outputs) - ) - task.evaluate_kernel.compile(tuple(numba.typeof(arg) for arg in warmup_args)) - observe_args = (clone_kernel_value(inputs), clone_kernel_value(self._env._kernel_outputs)) - task.observe_kernel.compile(tuple(numba.typeof(arg) for arg in observe_args)) + # compile() only consumes types, not values. Copying input arrays here + # changes readonly/strided views into writable contiguous arrays and + # compiles a specialization that real reset/step calls cannot reuse. + warmup_args = (inputs, task.reward_weights, self._env._kernel_buffers, self._env._kernel_outputs) + self._compile_specialization("evaluate", task.evaluate_kernel, warmup_args) + observe_args = (inputs, self._env._kernel_outputs) + self._compile_specialization("observe", task.observe_kernel, observe_args) # Precompile the reset kernel here instead of paying its compilation on # the first reset; env_ids always come from np.flatnonzero (int64). - reset_args = (clone_kernel_value(inputs), np.arange(2, dtype=np.int64), self._env._sim_reset_runtime.buffers) - task.reset_kernel.compile(tuple(numba.typeof(arg) for arg in reset_args)) + reset_args = (inputs, np.arange(2, dtype=np.int64), self._env._sim_reset_runtime.buffers) + self._compile_specialization("reset", task.reset_kernel, reset_args) + + def _compile_specialization(self, kind: str, kernel: Any, args: tuple[Any, ...]) -> None: + signature = tuple(numba.typeof(arg) for arg in args) + started = perf_counter() + kernel.compile(signature) + logger.info( + "Manager startup %s: %s specialization finished in %.3fs", + self._env_name(), + kind, + perf_counter() - started, + ) + logger.debug("Manager %s specialization signature: %s", kind, signature) def _load_kernel(self, kind: str, source: str, plan_key: str) -> tuple[Any, str]: """Load one fused kernel from the in-process or disk cache, else compile it.""" @@ -965,7 +976,7 @@ def _load_kernel(self, kind: str, source: str, plan_key: str) -> tuple[Any, str] kernel = self._compile_kernel(kind, source, filename) _KERNEL_CACHE[plan_key] = kernel logger.info( - "Manager startup %s: %s kernel compile finished in %.3fs (cache=%s)", + "Manager startup %s: %s kernel dispatcher prepared in %.3fs (cache=%s)", self._env_name(), kind, perf_counter() - compile_started, diff --git a/motrix_env_core/tests/test_function_fingerprint.py b/motrix_env_core/tests/test_function_fingerprint.py index 7bbda8e0..d0e9f4d9 100644 --- a/motrix_env_core/tests/test_function_fingerprint.py +++ b/motrix_env_core/tests/test_function_fingerprint.py @@ -1,6 +1,10 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 +import os +import subprocess +import sys + import numba from motrix_env_core.numba.fingerprint import function_fingerprint @@ -67,3 +71,43 @@ def helper(value): return value + 1 assert function_fingerprint(numba.njit(helper)) == function_fingerprint(helper) + + +def test_fingerprint_is_stable_across_hash_seeds(tmp_path): + script = tmp_path / "fingerprint.py" + script.write_text( + "from motrix_env_core.numba.fingerprint import function_fingerprint\n" + "def entry(value):\n" + " return value in {'alpha', 'beta', 'gamma'}\n" + "print(function_fingerprint(entry))\n", + encoding="utf-8", + ) + fingerprints = [] + for seed in ("1", "2"): + result = subprocess.run( + [sys.executable, str(script)], + check=True, + capture_output=True, + text=True, + env={**os.environ, "PYTHONHASHSEED": seed, "PYTHONPATH": "motrix_env_core/src"}, + ) + fingerprints.append(result.stdout.strip()) + assert fingerprints[0] == fingerprints[1] + + +def test_fingerprint_is_stable_after_inline_helper_compilation(): + namespace = {"__name__": "fingerprint_test", "numba": numba} + exec( + "@numba.njit(inline='always')\n" + "def helper(value):\n" + " if value > 0:\n" + " return min(value + 1, 10)\n" + " return max(value - 1, -10)\n" + "def entry(value):\n" + " return helper(value)\n", + namespace, + ) + entry = namespace["entry"] + first = function_fingerprint(entry) + assert numba.njit(entry)(2) == 3 + assert function_fingerprint(entry) == first diff --git a/motrix_env_core/tests/test_kernel_clone.py b/motrix_env_core/tests/test_kernel_clone.py new file mode 100644 index 00000000..8fa78381 --- /dev/null +++ b/motrix_env_core/tests/test_kernel_clone.py @@ -0,0 +1,48 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +from collections import namedtuple + +import numba +import numpy as np +import pytest + +from motrix_env_core.numba.kernel import clone_kernel_value + + +@pytest.mark.parametrize("layout", ["c", "f", "strided", "reversed", "broadcast", "empty", "empty_strided"]) +@pytest.mark.parametrize("readonly", [False, True]) +def test_kernel_clone_preserves_numba_array_type_and_isolates_storage(layout, readonly): + source = np.arange(48, dtype=np.float32).reshape(8, 6) + if layout == "f": + source = np.asfortranarray(source) + elif layout == "strided": + source = source[::2, ::2] + elif layout == "reversed": + source = source[::-1, ::-2] + elif layout == "broadcast": + source = np.broadcast_to(source[:1], (8, 6)) + elif layout == "empty": + source = source[:0] + elif layout == "empty_strided": + source = np.ndarray((0, 2), dtype=np.float32, buffer=bytearray(4), strides=(0, 2**62)) + source.flags.writeable = not readonly + copied = clone_kernel_value(source) + + assert numba.typeof(copied) == numba.typeof(source) + assert copied.strides == source.strides + assert not np.shares_memory(copied, source) + np.testing.assert_array_equal(copied, source) + if not readonly: + copied[...] = -1 + assert np.all(source >= 0) + + +def test_kernel_clone_preserves_nested_namedtuple_and_scalar_values(): + outputs = namedtuple("Outputs", ["value", "count"]) + source = (outputs(np.arange(4, dtype=np.float32), 3), np.float32(0.5)) + copied = clone_kernel_value(source) + assert type(copied[0]) is outputs + assert copied[0].count == 3 + assert copied[1] == np.float32(0.5) + assert not np.shares_memory(source[0].value, copied[0].value) diff --git a/motrix_envs/tests/test_wbt_numba.py b/motrix_envs/tests/test_wbt_numba.py index 1a15de9e..6d818dfd 100644 --- a/motrix_envs/tests/test_wbt_numba.py +++ b/motrix_envs/tests/test_wbt_numba.py @@ -55,6 +55,10 @@ ) from motrix_envs.motion import MotrixMotion, WbtMotionClip # noqa: E402 +# Keep WBT kernels in one suite: deselecting only the first expensive tests +# shifts first-use compilation to later tests instead of removing its cost. +pytestmark = pytest.mark.numba + def _motion_command(env: ManagerEnv) -> WbtMotionCommand: command = env.command_terms["motion"] @@ -215,6 +219,7 @@ def test_numba_wbt_read_plan_reuses_preallocated_arrays() -> None: np.testing.assert_array_equal(action.action_queue[[1, 3], action.action_ptr[0]], 0.0) +@pytest.mark.slow def test_numba_wbt_step_preserves_previous_actor_and_critic_observations() -> None: env = _make_numba_env(_deterministic_manager_cfg(), num_envs=2) initial = env.init_state() @@ -224,7 +229,22 @@ def test_numba_wbt_step_preserves_previous_actor_and_critic_observations() -> No initial_policy_snapshot = initial_policy.copy() initial_value_snapshot = initial_value.copy() + task = env._task_program + assert task is not None + kernels = (task.evaluate_kernel, task.observe_kernel, task.reset_kernel) + signatures = tuple(tuple(kernel.signatures) for kernel in kernels) + assert all(len(items) == 1 for items in signatures) + env.warmup() first = env.step(np.zeros((env.num_envs, *env.action_space.shape), dtype=np.float32)) + assert tuple(tuple(kernel.signatures) for kernel in kernels) == signatures + second_env = _make_numba_env(_deterministic_manager_cfg(), num_envs=2) + second_env.init_state() + assert second_env.manager_layout.plan_keys == env.manager_layout.plan_keys + second_task = second_env._task_program + assert second_task is not None + assert second_task.evaluate_kernel is task.evaluate_kernel + assert second_task.observe_kernel is task.observe_kernel + assert second_task.reset_kernel is task.reset_kernel np.testing.assert_array_equal(initial_policy, initial_policy_snapshot) np.testing.assert_array_equal(initial_value, initial_value_snapshot) @@ -233,6 +253,7 @@ def test_numba_wbt_step_preserves_previous_actor_and_critic_observations() -> No assert not np.shares_memory(initial_value, first.obs.value) +@pytest.mark.slow def test_numba_wbt_observation_noise_bounds_and_determinism() -> None: """Observation noise is uniform(-amp, amp), zero for zero amplitudes, and reproducible for a fresh environment with the same seed and episode step.""" diff --git a/pyproject.toml b/pyproject.toml index dcf598f3..d2b8965b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -136,6 +136,13 @@ triton-rocm = [ { index = "pytorch-rocm72", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, ] +[tool.pytest.ini_options] +markers = [ + "slow: tests with substantial compilation or runtime cost", + "numba: dedicated Numba kernel regression tests", + "integration: full-environment or subprocess integration tests", +] + [tool.mypy] files = [ "motrix_env_core/tests/typecheck/configclass.py", diff --git a/test/test_all_envs.py b/test/test_all_envs.py index 84fc94f4..f20cb579 100644 --- a/test/test_all_envs.py +++ b/test/test_all_envs.py @@ -17,6 +17,8 @@ import subprocess import sys +import pytest + import motrix_envs # noqa: F401 registers built-in environments from motrix_env_core import registry @@ -61,6 +63,8 @@ def _smoke_in_subprocess(case: tuple[str, int]) -> tuple[str, str | None]: return (label, None) +@pytest.mark.slow +@pytest.mark.integration def test_all_demos(): all_envs = sorted(registry.list_registered_envs()) cases = [(env_name, num_env) for num_env in (1, 2) for env_name in all_envs]