diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index a92fb72..0e09024 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -15,6 +15,38 @@ permissions: contents: 'read' jobs: + benchmark-smoke: + name: 'Hierarchical comparison correctness (3.12)' + runs-on: 'ubuntu-latest' + steps: + - name: 'Checkout repository' + uses: 'actions/checkout@v6' + with: + persist-credentials: false + - name: 'Setup uv' + uses: 'astral-sh/setup-uv@v7.2.0' + with: + python-version: '3.12' + - name: 'Install optional comparison dependencies explicitly' + run: 'uv sync --dev --extra benchmark' + - name: 'Check tiny trajectory and reverse-mode comparisons' + run: >- + uv run --no-sync python -m scripts.benchmark_hierarchical + --groups 1 --local-dimensions 1 --steps 4 --solvers core external + diffrax-fixed diffrax-adaptive numpy scipy --derivatives trajectory grad-rev + --repeats 1 --timeout 180 --require-success --output /tmp/hierarchical-smoke.jsonl + - name: 'Check supported Diffrax higher-order configuration' + run: >- + uv run --no-sync python -m scripts.benchmark_hierarchical + --groups 1 --local-dimensions 1 --steps 4 --solvers diffrax-fixed + --derivatives hessian logdet-grad --repeats 1 --timeout 180 + --require-success --output /tmp/hierarchical-higher-order.jsonl + - name: 'Retain correctness/protocol evidence' + if: 'always()' + uses: 'actions/upload-artifact@v7' + with: + name: 'hierarchical-comparison-correctness' + path: '/tmp/hierarchical-*.jsonl' backend: name: 'Solver qualification (${{ matrix.backend }}, 3.12)' runs-on: 'ubuntu-latest' diff --git a/CHANGELOG.md b/CHANGELOG.md index 000f8a9..fc50df3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,13 @@ versioning while the project remains pre-1.0. ### Added +- Opt-in serial hierarchical measurements emit correctness-gated JSON with + dimensions, accepted steps, cold tracing/compilation, warm execution, memory + and provenance. Matched-mesh comparisons use public native/external loops + and supported Diffrax AD; adaptive Tsit5/Kvaerno5 and concrete SciPy/NumPy + references are identified separately. Tiny CI comparisons reject failures; + timeout records remain censored without fabricated speedups (#204). + - Required NumPy/JAX/CPU Torch solver jobs reject missing or skipped selected qualification. A structured trait-coupled model verifies every deterministic fixed method against independent matrix-exponential solutions in both diff --git a/benchmarks/hierarchical_measurement.py b/benchmarks/hierarchical_measurement.py new file mode 100644 index 0000000..f7473ee --- /dev/null +++ b/benchmarks/hierarchical_measurement.py @@ -0,0 +1,487 @@ +"""Isolated optional measurements and fair Diffrax configurations for the fixture.""" + +from __future__ import annotations + +import hashlib +import importlib.metadata +import os +import platform +import resource +import shutil +import subprocess # noqa: S404 +import sys +import time +from pathlib import Path +from typing import TYPE_CHECKING, Any, NamedTuple + +import numpy as np + +from benchmarks.hierarchical_ode import ( + HierarchicalProblem, + _flatten_mesh, + initial_state, + loss_function, + nominal_parameters, + numpy_replay_trajectory, + objective, + observations, + reference_objective, + reference_trajectory, + replay_trajectory, + rhs, +) +from op_engine import Array, CoreSolver, ModelCore +from op_engine.core_solver import ( + AdaptiveConfig, + AdaptiveStepSchedule, + DtControllerConfig, + RunConfig, +) +from op_engine.model_core import ModelCoreOptions + +if TYPE_CHECKING: + from argparse import Namespace + from collections.abc import Callable + +OUTPUT_TIMES = np.asarray([0.0, 0.2, 0.5, 1.0]) +REGULARIZATION = 25.0 + + +class _DiffraxSettings(NamedTuple): + solver: str + derivative: str + max_step: float + + +def _version(name: str) -> str | None: + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + return None + + +def _hardware() -> dict[str, object]: + model = platform.processor() + path = Path("/proc/cpuinfo") + if path.exists(): + model = next( + ( + line.split(":", 1)[1].strip() + for line in path.read_text(encoding="utf-8").splitlines() + if line.startswith("model name") + ), + model, + ) + return { + "machine": platform.machine(), + "cpu": model, + "logical_cpus": os.cpu_count(), + "thread_environment": { + name: os.environ.get(name) + for name in [ + "OMP_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "MKL_NUM_THREADS", + "XLA_FLAGS", + ] + }, + } + + +def _provenance() -> dict[str, object]: + root = Path(__file__).resolve().parents[1] + git = shutil.which("git") or "/usr/bin/git" + commit = subprocess.run( # noqa: S603 + [git, "rev-parse", "HEAD"], + cwd=root, + capture_output=True, + text=True, + check=False, + ) + dirty = subprocess.run( # noqa: S603 + [git, "status", "--porcelain"], + cwd=root, + capture_output=True, + text=True, + check=False, + ) + from benchmarks import hierarchical_ode # noqa: PLC0415 + from op_engine import core_solver # noqa: PLC0415 + + files = { + "fixture": Path(hierarchical_ode.__file__), + "measurement": Path(__file__), + "driver": root / "scripts/benchmark_hierarchical.py", + "core_solver": Path(core_solver.__file__), + } + return { + "commit": commit.stdout.strip() if commit.returncode == 0 else None, + "dirty": bool(dirty.stdout) if dirty.returncode == 0 else None, + "source_sha256": { + name: hashlib.sha256(path.read_bytes()).hexdigest() + for name, path in files.items() + }, + } + + +def _discover( + problem: HierarchicalProblem, parameters: np.ndarray, minimum_steps: int +) -> tuple[AdaptiveStepSchedule, RunConfig]: + config = RunConfig( + method="dopri5", + adaptive=True, + adaptive_cfg=AdaptiveConfig( + rtol=2e-8, atol=2e-10, dt_init=min(0.02, 1 / minimum_steps) + ), + dt_controller=DtControllerConfig(dt_max=1 / minimum_steps), + replay_model_id=problem.model_id, + replay_reference_regime="nominal_parameters()", + ) + core = ModelCore( + 2, + problem.groups, + OUTPUT_TIMES, + options=ModelCoreOptions( + other_axes=(problem.local_dimension,), axis_names=("pool", "group", "local") + ), + ) + core.set_initial_state(initial_state(problem, parameters)) + solver = CoreSolver(core) + solver.run(rhs(problem, parameters), config=config) + if solver.last_adaptive_schedule is None: + msg = "Nominal discovery did not produce a mesh" + raise RuntimeError(msg) + return solver.last_adaptive_schedule, config + + +def _diffrax_functions( + problem: HierarchicalProblem, + schedule: AdaptiveStepSchedule, + settings: _DiffraxSettings, + data: Array, +) -> tuple[ + Callable[[Array], Array], + Callable[[Array], Array], + dict[str, object], + Callable[[Array], Any], +]: + import diffrax as d # noqa: PLC0415 + import jax.numpy as jnp # noqa: PLC0415 + + solver_name, derivative, max_step = settings + if derivative in {"trajectory", "grad-rev"}: + adjoint = d.RecursiveCheckpointAdjoint() + elif derivative == "grad-fwd": + adjoint = d.ForwardMode() + else: + adjoint = d.DirectAdjoint() + # Bounded RK-stage scans support forward and higher-order AD. The adjoint + # selection additionally covers saved outputs and the solve iteration loop. + if solver_name == "diffrax-fixed": + solver = d.Dopri5(scan_kind="bounded") + starts, sizes, _ = _flatten_mesh(schedule) + boundaries = np.concatenate((starts, [schedule.output_times[-1]])) + controller = d.StepTo(ts=jnp.asarray(boundaries)) + max_steps = len(sizes) + dt0 = None + mesh_mode = "identical frozen accepted mesh" + else: + solver = ( + d.Kvaerno5(scan_kind="bounded") + if problem.stiff + else d.Tsit5(scan_kind="bounded") + ) + controller = d.PIDController(rtol=2e-8, atol=2e-10, dtmax=max_step) + max_steps = max(256, 4 * sum(map(len, schedule.step_sizes))) + dt0 = min(0.02, max_step) + mesh_mode = "live adaptive PID; separate task from frozen replay" + term = d.ODETerm(lambda t, y, args: rhs(problem, args)(t, y)) + + def solution_at(values: Array) -> Any: # noqa: ANN401 + return d.diffeqsolve( + term, + solver, + t0=0.0, + t1=1.0, + dt0=dt0, + y0=initial_state(problem, values), + args=values, + saveat=d.SaveAt(ts=jnp.asarray(OUTPUT_TIMES)), + stepsize_controller=controller, + adjoint=adjoint, + max_steps=max_steps, + ) + + def trajectory(values: Array) -> Array: + return solution_at(values).ys + + def loss(values: Array) -> Array: + return objective(problem, values, trajectory(values), data) + + settings = { + "solver": type(solver).__name__, + "adjoint": type(adjoint).__name__, + "scan_kind": "bounded", + "max_steps": max_steps, + "mesh_mode": mesh_mode, + "controller": type(controller).__name__, + "rtol": None if solver_name == "diffrax-fixed" else 2e-8, + "atol": None if solver_name == "diffrax-fixed" else 2e-10, + "implicit_root_finder": type(solver.root_finder).__name__ + if hasattr(solver, "root_finder") + else None, + "root_find_max_steps": getattr(solver, "root_find_max_steps", None), + } + return trajectory, loss, settings, solution_at + + +def run_case(args: Namespace) -> dict[str, object]: # noqa: C901, PLR0912, PLR0914, PLR0915 + """Measure one correctness-gated task in a fresh CPU process. + + Returns: + Dimensions, agreement, phase timings, warm samples, memory and provenance. + + Raises: + ValueError: If an unsupported solver/derivative combination is selected. + RuntimeError: If independent accuracy or positive-definiteness checks fail. + """ + total_start = time.perf_counter() + problem = HierarchicalProblem( + args.groups[0], + args.local_dimensions[0], + coupled=args.coupling[0] == "coupled", + stiff=args.regimes[0] == "stiff", + ) + nominal = nominal_parameters(problem) + solver_name, derivative = args.solvers[0], args.derivatives[0] + setup_start = time.perf_counter() + schedule, config = _discover(problem, nominal, args.steps[0]) + discovery_seconds = time.perf_counter() - setup_start + reference_start = time.perf_counter() + independent = reference_trajectory(problem, nominal, OUTPUT_TIMES) + data_np = observations(problem, OUTPUT_TIMES) + reference_seconds = time.perf_counter() - reference_start + base: dict[str, object] = { + "status": "success", + "solver": solver_name, + "derivative": derivative, + "groups": problem.groups, + "local_parameter_dimension": problem.local_dimension, + "parameter_count": problem.parameter_count, + "state_shape": problem.state_shape, + "regime": "stiff" if problem.stiff else "nonstiff", + "coupled": problem.coupled, + "requested_minimum_steps": args.steps[0], + "accepted_steps": sum(map(len, schedule.step_sizes)), + "core_discovery_accepted_steps": sum(map(len, schedule.step_sizes)), + "output_times": OUTPUT_TIMES.tolist(), + "dtype": "float64", + "platform": "cpu", + "rtol": 2e-8, + "atol": 2e-10, + "discovery_dt_max": 1 / args.steps[0], + "tolerances_apply_to": "Discovery only; frozen replay bypasses error control", + "reference_method": "SciPy DOP853", + "reference_rtol": 2e-12, + "reference_atol": 2e-14, + "reference_seconds": reference_seconds, + "core_discovery_seconds": discovery_seconds, + "regularization": REGULARIZATION, + "timeout_seconds": args.timeout, + "timeout_policy": "censored; no measured speedup inferred", + "repeats": args.repeats, + "versions": { + name: _version(name) + for name in ["op-engine", "numpy", "scipy", "jax", "diffrax"] + }, + "python": sys.version.split()[0], + "hardware": _hardware(), + "provenance": _provenance(), + } + if solver_name in {"numpy", "scipy"}: + if derivative != "trajectory": + msg = f"{solver_name} is a trajectory reference, not an AD implementation" + raise ValueError(msg) + function = ( + (lambda: numpy_replay_trajectory(problem, nominal, schedule, config)) + if solver_name == "numpy" + else (lambda: reference_trajectory(problem, nominal, OUTPUT_TIMES)) + ) + start = time.perf_counter() + value = function() + first_seconds = time.perf_counter() - start + samples = [] + for _ in range(args.repeats): + start = time.perf_counter() + function() + samples.append(time.perf_counter() - start) + error = float(np.max(np.abs(value - independent))) + if not np.allclose(value, independent, rtol=5e-7, atol=3e-9): + msg = "Concrete trajectory failed its independent reference" + raise RuntimeError(msg) + base.update({ + "trace_lower_seconds": None, + "compile_seconds": None, + "first_execution_seconds": first_seconds, + "warm_execution_seconds": samples, + "trajectory_max_absolute_error": error, + "xla_buffer_bytes": None, + "method": "dopri5 frozen replay" + if solver_name == "numpy" + else "DOP853 accurate live reference", + }) + else: + import jax # noqa: PLC0415 + import jax.numpy as jnp # noqa: PLC0415 + + jax.config.update("jax_enable_x64", True) # noqa: FBT003 + jax.config.update("jax_enable_compilation_cache", False) # noqa: FBT003 + base["persistent_compilation_cache_enabled"] = ( + jax.config.jax_enable_compilation_cache + ) + backend_start = time.perf_counter() + parameters = jnp.asarray(nominal) + data = jnp.asarray(data_np) + parameters.block_until_ready() + data.block_until_ready() + base["backend_setup_seconds"] = time.perf_counter() - backend_start + if solver_name in {"core", "external"}: + + def trajectory(values: Array) -> Array: + return replay_trajectory( + problem, + values, + schedule, + config, + external=solver_name == "external", + ) + + loss = loss_function( + problem, schedule, config, data, external=solver_name == "external" + ) + settings = { + "solver": "dopri5", + "adjoint": "native JAX AD with checkpointed scan", + "max_steps": sum(map(len, schedule.step_sizes)), + "mesh_mode": "identical frozen accepted mesh", + "fsal": True, + } + else: + trajectory, loss, settings, solution_at = _diffrax_functions( + problem, + schedule, + _DiffraxSettings(solver_name, derivative, 1 / args.steps[0]), + data, + ) + base["numerical_configuration"] = settings + + def logdet(values: Array) -> Array: + matrix = jax.hessian(loss)(values) + REGULARIZATION * jnp.eye( + problem.parameter_count + ) + return 0.5 * jnp.linalg.slogdet(matrix)[1] + + functions = { + "trajectory": trajectory, + "grad-fwd": jax.jacfwd(loss), + "grad-rev": jax.grad(loss), + "hessian": jax.hessian(loss), + "logdet-grad": jax.grad(logdet), + } + trace_start = time.perf_counter() + lowered = jax.jit(functions[derivative]).lower(parameters) + lowered_at = time.perf_counter() + compiled = lowered.compile() + compiled_at = time.perf_counter() + value = jax.block_until_ready(compiled(parameters)) + executed_at = time.perf_counter() + samples = [] + for _ in range(args.repeats): + start = time.perf_counter() + jax.block_until_ready(compiled(parameters)) + samples.append(time.perf_counter() - start) + base["trace_lower_seconds"] = lowered_at - trace_start + # Keep the phase boundaries from the first compilation separate from warm runs. + base["compile_seconds"] = compiled_at - lowered_at + base["first_execution_seconds"] = executed_at - compiled_at + base["warm_execution_seconds"] = samples + actual = np.asarray(jax.jit(trajectory)(parameters)) + if solver_name.startswith("diffrax"): + solution = solution_at(parameters) + stats = {key: int(value) for key, value in solution.stats.items()} + base["solver_step_statistics"] = stats + base["accepted_steps"] = stats["num_accepted_steps"] + error = float(np.max(np.abs(actual - independent))) + if not np.allclose(actual, independent, rtol=5e-7, atol=3e-9): + msg = "JAX trajectory failed the independent reference" + raise RuntimeError(msg) + if not np.all(np.isfinite(np.asarray(value))): + msg = "Derivative output contains nonfinite values" + raise RuntimeError(msg) + base["trajectory_max_absolute_error"] = error + if derivative in {"grad-fwd", "grad-rev"}: + shifts = 1e-4 * np.eye(problem.parameter_count) + expected = np.asarray([ + ( + reference_objective(problem, nominal + shift, OUTPUT_TIMES, data_np) + - reference_objective( + problem, nominal - shift, OUTPUT_TIMES, data_np + ) + ) + / 2e-4 + for shift in shifts + ]) + if not np.allclose(value, expected, rtol=3e-5, atol=3e-6): + msg = "Gradient failed independent finite differences" + raise RuntimeError(msg) + base["derivative_max_absolute_error"] = float( + np.max(np.abs(np.asarray(value) - expected)) + ) + if derivative in {"hessian", "logdet-grad"}: + hessian = np.asarray(jax.jit(jax.hessian(loss))(parameters)) + eigenvalues = np.linalg.eigvalsh( + hessian + REGULARIZATION * np.eye(problem.parameter_count) + ) + if not np.all(eigenvalues > 0): + msg = "Regularized Hessian is not positive definite" + raise RuntimeError(msg) + base["regularized_hessian_min_eigenvalue"] = float(eigenvalues.min()) + shifts = 1e-4 * np.eye(problem.parameter_count) + if derivative == "hessian": + gradient = jax.jit(jax.grad(loss)) + expected = np.column_stack([ + ( + np.asarray(gradient(parameters + shift)) + - np.asarray(gradient(parameters - shift)) + ) + / 2e-4 + for shift in shifts + ]) + else: + evaluate = jax.jit(logdet) + expected = np.asarray([ + (evaluate(parameters + shift) - evaluate(parameters - shift)) / 2e-4 + for shift in shifts + ]) + if not np.allclose(value, expected, rtol=3e-5, atol=3e-6): + msg = "Higher derivative failed centered differences" + raise RuntimeError(msg) + base["derivative_max_absolute_error"] = float( + np.max(np.abs(np.asarray(value) - expected)) + ) + analysis = compiled.memory_analysis() + base["xla_buffer_bytes"] = { + name: getattr(analysis, name, None) + for name in [ + "argument_size_in_bytes", + "output_size_in_bytes", + "temp_size_in_bytes", + "alias_size_in_bytes", + ] + } + rss_divisor = 1024**2 if sys.platform == "darwin" else 1024 + base["peak_process_rss_mib"] = ( + resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / rss_divisor + ) + base["total_worker_seconds"] = time.perf_counter() - total_start + return base diff --git a/docs/guides/hierarchical-measurements.md b/docs/guides/hierarchical-measurements.md new file mode 100644 index 0000000..7ec553c --- /dev/null +++ b/docs/guides/hierarchical-measurements.md @@ -0,0 +1,118 @@ +# Hierarchical scaling and comparison measurements + +The [hierarchical ODE fixture](hierarchical-ode.md) supplies independent cell +references, deterministic observations, shared/local/IC parameters, four +regimes, and the statistical objective. This measurement layer invokes that +same problem in fresh serial CPU workers. Small CI cases qualify the protocol +and supported comparison configurations; larger timing sweeps are opt-in. +There are no speed or memory thresholds in CI. + +## Setup and bounded commands + +```shell +uv sync --dev --extra benchmark +just benchmark-hierarchical --groups 1 --local-dimensions 1 --steps 4 \ + --solvers core external diffrax-fixed numpy scipy \ + --derivatives trajectory grad-fwd grad-rev hessian logdet-grad \ + --repeats 3 --timeout 180 --require-success --output /tmp/hierarchical-small.jsonl +``` + +The optional `benchmark` extra installs Diffrax; it is not a numerical-core +runtime dependency. Python/JAX/Diffrax/SciPy versions are reported, so retain +the resolved environment with results. Concrete NumPy and SciPy references +produce trajectory rows, not derivative rows. A larger sweep can vary: + +```shell +just benchmark-hierarchical --groups 2 4 8 --local-dimensions 1 2 \ + --steps 8 32 128 --regimes nonstiff stiff --coupling independent coupled \ + --solvers core external diffrax-fixed --derivatives grad-rev hessian logdet-grad \ + --repeats 5 --timeout 300 --output /tmp/hierarchical-scaling.jsonl +``` + +`--steps` sets a minimum-step policy by capping discovery/adaptive `dt_max` at +`1/steps`; it does not invent accepted steps. Actual accepted counts are +reported. Output times remain `[0, 0.2, 0.5, 1]`, so increasing step count does +not also increase observation dimension. Increasing local dimension adds rate +and IC parameters per group: total fitted dimension is +`2 + 2*groups*local_dimension`. + +## Compare equivalent tasks + +| Solver label | Task and configuration | +| --- | --- | +| `core` | Built-in Dormand–Prince accepted-mesh replay with JAX scan/checkpointing | +| `external` | Public `fixed_explicit_step` in an external checkpointed scan, including FSAL reuse, on the identical mesh | +| `diffrax-fixed` | Diffrax Dormand–Prince with `StepTo` on that exact frozen mesh and the same output times | +| `diffrax-adaptive` | Separate live-adaptive task: Tsit5 for nonstiff problems, Kvaerno5 for stiff problems, PID tolerances `2e-8/2e-10`, with the same maximum-step policy | +| `numpy` | Eager core replay of the identical nominal mesh; trajectory parity/reference timing | +| `scipy` | Independent cell-loop DOP853 solve at `2e-12/2e-14`; accurate trajectory reference timing | + +Fixed-mesh comparisons hold step times, method, parameters, outputs and objective +constant. Adaptive comparison rows identify their own accepted/rejected steps +and the producing core mesh count separately; their workload is not identical +to frozen replay. Do not combine unlike tasks into a claimed solver speedup. +Reference-RHS loops intentionally differ from the native vectorized RHS, so +SciPy timings identify the chosen trajectory-reference implementation. + +Diffrax configurations follow its +[solver recommendations](https://docs.kidger.site/diffrax/usage/how-to-choose-a-solver/), +[adjoint contracts](https://docs.kidger.site/diffrax/api/adjoints/), and +[Hessian example](https://docs.kidger.site/diffrax/examples/hessian/). +Reverse gradient and trajectory rows use `RecursiveCheckpointAdjoint`, forward +rows use `ForwardMode`, and Hessian/log-determinant-gradient rows use +`DirectAdjoint`. RK stages select `scan_kind="bounded"`, and saved-output loops +use the adjoint appropriate to the transform. DirectAdjoint supports both modes +but costs can depend on `max_steps`; that bound is recorded. Fixed rows set it +to the actual mesh length. Adaptive rows use a stated bound based on discovery, +and failures remain visible rather than implying unsupported AD. + +All current measurements use float64 CPU execution. GPU and other namespaces +are outside this measurement protocol. The regularized Hessian is `H+25*I`; +positive definiteness and nonfinite values are checked before a row qualifies +for interpretation. See the fixture guide for priors and nonzero residuals. + +## JSON fields and timing boundaries + +Each completed row is correctness-gated against the independent trajectory. +Gradient rows compare independent-reference centered differences; Hessian and +log-determinant-gradient rows compare centered differences of their relevant +native derivatives. The fixture's separate required tests also check forward/ +reverse agreement and refinement. Failed agreement produces an error row. + +| Fields | Meaning | +| --- | --- | +| `groups`, `local_parameter_dimension`, `parameter_count`, `state_shape` | Problem dimensions and fitted initial-condition contribution | +| `accepted_steps`, `core_discovery_accepted_steps`, `solver_step_statistics` | Actual measured-solver steps and nominal core discovery counts; Diffrax statistics distinguish rejection | +| `core_discovery_seconds`, `reference_seconds`, `backend_setup_seconds` | Preprocessing/reference/imported-backend setup, outside selected compiled numerical timing | +| `trace_lower_seconds`, `compile_seconds` | First selected transform tracing/lowering and compilation in a fresh process | +| `first_execution_seconds`, `warm_execution_seconds` | Synchronized execution after compilation and full warm samples | +| `total_worker_seconds` | Time inside the measurement function, including setup and validation; process startup/imports occur before this timer | +| `peak_process_rss_mib` | Whole worker peak RSS, including imports, setup, validation and any extra validation compilations; not isolated live solver storage | +| `xla_buffer_bytes` | Selected compiled transform's reported argument/output/temp/alias buffer requirements; may be unavailable | +| `trajectory_max_absolute_error`, `derivative_max_absolute_error`, `regularized_hessian_min_eigenvalue` | Agreement and positive-definiteness evidence | +| `versions`, `python`, `hardware`, `provenance`, `numerical_configuration` | Exact packages, CPU/thread settings, commit/dirty flag/content hashes, adjoint/stage/controller/step bounds | + +A fresh worker prevents previous cases' compilation/memory caches from becoming +an advantage. Workers also disable JAX's persistent compilation cache, even if +the parent configured a cache directory, and report +`persistent_compilation_cache_enabled: false`. See JAX's +[cache setting](https://docs.jax.dev/en/latest/config_options.html#enable-compilation-cache). +First compiled-transform timing occurs before the postmeasurement +accuracy checks; every result is validated before emission. Warm calls block +until completed. RSS includes the full worker lifecycle, while XLA buffer +analysis describes only the selected compiled transform. Concrete reference +rows have no tracing/compilation or XLA-memory claims. +`backend_setup_seconds` covers array/backend initialization after JAX imports; +the process timeout and RSS cover startup/import costs too. + +Timeouts are serialized with `status: timeout`, `censored: true` and no timing +result. They are not measured speedups. Errors include diagnostics and return +codes; `--require-success` makes any failed/censored case fail the command. +Retain failed rows when reporting a sweep rather than selecting only favorable +results. CI uploads its tiny JSON protocol evidence as an artifact; virtual +runner timings are not publication scaling claims. + +For publication, save the JSON, environment resolution, exact source tag/commit +and content hashes with the archive. Report hardware, dimension and mesh policy, +objective/regularization, tolerances, derivative/adjoint mode, repetitions, +timeout policy, and agreement before interpreting runtime or memory trends. diff --git a/justfile b/justfile index a0d6993..b72185e 100644 --- a/justfile +++ b/justfile @@ -102,8 +102,8 @@ build-test-core: uv pip install --python "${CLEANROOM}/venv/bin/python" -r "${CLEANROOM}/dev-requirements.txt" cp pyproject.toml "${CLEANROOM}/pyproject.toml" cp -R tests "${CLEANROOM}/tests" + cp -R scripts "${CLEANROOM}/scripts" cp -R examples "${CLEANROOM}/examples" - cp -R benchmarks "${CLEANROOM}/benchmarks" cd "${CLEANROOM}" "${CLEANROOM}/venv/bin/pytest" --import-mode=importlib tests --quiet --exitfirst @@ -195,3 +195,8 @@ qualify-torch: uv sync --dev uv pip install 'torch>=2.6' --index-url https://download.pytorch.org/whl/cpu uv run --no-sync python scripts/qualify_backend.py torch + + +# Larger serial measurements are opt-in; install the benchmark extra first. +benchmark-hierarchical *args: + uv run --no-sync python -m scripts.benchmark_hierarchical {{args}} diff --git a/mkdocs.yml b/mkdocs.yml index 020fd48..87c9bc5 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -61,6 +61,8 @@ nav: - Replay Provenance: guides/replay-provenance.md - Hierarchical ODE Evidence: guides/hierarchical-ode.md + + - Hierarchical Measurements: guides/hierarchical-measurements.md - Backend Strategies: guides/backends.md - API Reference: - Adaptive Tau: api-reference/adaptive_tau.md diff --git a/pyproject.toml b/pyproject.toml index 998a095..3c4908e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,9 @@ Documentation = "https://accidda.github.io/op_engine/" Issues = "https://github.com/ACCIDDA/op_engine/issues" [project.optional-dependencies] +benchmark = [ + "diffrax>=0.7.2", +] cupy = [ "cupy>=13", ] diff --git a/scripts/benchmark_hierarchical.py b/scripts/benchmark_hierarchical.py new file mode 100644 index 0000000..89f4f3e --- /dev/null +++ b/scripts/benchmark_hierarchical.py @@ -0,0 +1,174 @@ +"""Run correctness-gated hierarchical measurements in isolated serial CPU workers.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import subprocess # noqa: S404 +import sys +from itertools import product +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterator + + +SOLVERS = ("core", "external", "diffrax-fixed", "diffrax-adaptive", "numpy", "scipy") +DERIVATIVES = ("trajectory", "grad-fwd", "grad-rev", "hessian", "logdet-grad") + + +def _positive(value: str) -> int: + result = int(value) + if result < 1: + msg = "dimensions and minimum steps must be positive integers" + raise argparse.ArgumentTypeError(msg) + return result + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--groups", nargs="+", type=_positive, default=[2]) + parser.add_argument("--local-dimensions", nargs="+", type=_positive, default=[1]) + parser.add_argument("--steps", nargs="+", type=_positive, default=[16]) + parser.add_argument( + "--regimes", nargs="+", choices=["nonstiff", "stiff"], default=["nonstiff"] + ) + parser.add_argument( + "--coupling", + nargs="+", + choices=["independent", "coupled"], + default=["independent"], + ) + parser.add_argument( + "--solvers", + nargs="+", + choices=SOLVERS, + default=["core", "external", "diffrax-fixed"], + ) + parser.add_argument( + "--derivatives", nargs="+", choices=DERIVATIVES, default=list(DERIVATIVES) + ) + parser.add_argument("--repeats", type=_positive, default=3) + parser.add_argument("--timeout", type=float, default=180.0) + parser.add_argument("--output", type=Path) + parser.add_argument("--require-success", action="store_true") + parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS) + return parser + + +def _cases(args: argparse.Namespace) -> Iterator[dict[str, object]]: + fields = ( + "groups", + "local_dimensions", + "steps", + "regimes", + "coupling", + "solvers", + "derivatives", + ) + for values in product(*(getattr(args, field) for field in fields)): + case = dict(zip(fields, values, strict=True)) + # Concrete reference solvers measure trajectories; they do not expose AD. + if ( + case["solvers"] in {"numpy", "scipy"} + and case["derivatives"] != "trajectory" + ): + continue + yield case + + +def _isolated(case: dict[str, object], args: argparse.Namespace) -> dict[str, object]: + command = [sys.executable, "-m", "scripts.benchmark_hierarchical", "--worker"] + for field, value in case.items(): + command.extend(("--" + field.replace("_", "-"), str(value))) + command.extend(("--repeats", str(args.repeats), "--timeout", str(args.timeout))) + environment = os.environ.copy() + environment.update({ + "JAX_PLATFORMS": "cpu", + "JAX_ENABLE_COMPILATION_CACHE": "false", + "OMP_NUM_THREADS": "1", + "OPENBLAS_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1", + }) + try: + process = subprocess.run( # noqa: S603 + command, + check=False, + capture_output=True, + text=True, + timeout=args.timeout, + env=environment, + ) + except subprocess.TimeoutExpired: + return { + "status": "timeout", + "case": case, + "timeout_seconds": args.timeout, + "censored": True, + "timing_result": None, + "interpretation": "No completed measurement; not a measured speedup", + } + if process.stdout.strip(): + try: + row = json.loads(process.stdout.strip().splitlines()[-1]) + except json.JSONDecodeError: + row = {"status": "error", "error": "Worker did not return valid JSON"} + else: + row = {"status": "error", "error": "Worker exited without a result"} + if process.returncode: + row.update({ + "status": "error", + "returncode": process.returncode, + "case": case, + "stderr": process.stderr[-2000:], + }) + return row + + +def main(argv: list[str] | None = None) -> int: + """Emit one JSON record per independent measurement, including censored failures. + + Returns: + Nonzero for failed required measurements or invalid timeout configuration. + """ + parser = _parser() + args = parser.parse_args(argv) + if not math.isfinite(args.timeout) or args.timeout <= 0: + parser.error("timeout must be finite and positive") + if args.worker: + try: + from benchmarks.hierarchical_measurement import run_case # noqa: PLC0415 + + result = run_case(args) + except ( + ImportError, + RuntimeError, + ValueError, + TypeError, + OverflowError, + ) as error: + print( + json.dumps({ + "status": "error", + "error": f"{type(error).__name__}: {error}", + }) + ) + return 1 + print(json.dumps(result, allow_nan=False)) + return 0 + rows = [_isolated(case, args) for case in _cases(args)] + text = "\n".join(json.dumps(row, allow_nan=False) for row in rows) + "\n" + if args.output is None: + print(text, end="") + else: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(text, encoding="utf-8") + print(f"Wrote {len(rows)} cases to {args.output}") + return int(args.require_success and any(row["status"] != "success" for row in rows)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/op_engine/test_hierarchical_measurement_protocol.py b/tests/op_engine/test_hierarchical_measurement_protocol.py new file mode 100644 index 0000000..9c0e2aa --- /dev/null +++ b/tests/op_engine/test_hierarchical_measurement_protocol.py @@ -0,0 +1,107 @@ +"""Measurement protocol contracts keep censored runs distinct from results.""" + +from __future__ import annotations + +import json +import subprocess # noqa: S404 +from argparse import Namespace +from typing import TYPE_CHECKING + +import pytest + +from scripts.benchmark_hierarchical import _cases, _isolated, main # noqa: PLC2701 + +if TYPE_CHECKING: + from pathlib import Path + + +def _arguments() -> Namespace: + return Namespace( + groups=[1], + local_dimensions=[1], + steps=[4], + regimes=["nonstiff"], + coupling=["independent"], + solvers=["core", "numpy", "scipy"], + derivatives=["trajectory", "hessian"], + repeats=1, + timeout=0.1, + ) + + +def test_reference_solvers_do_not_claim_derivative_capability() -> None: + """Concrete references emit trajectory cases while the native path emits both.""" + cases = list(_cases(_arguments())) + assert len(cases) == 4 + assert all( + case["derivatives"] == "trajectory" + for case in cases + if case["solvers"] in {"numpy", "scipy"} + ) + + +def test_timeout_is_censored_without_a_timing_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An interrupted worker is never presented as a completed fast measurement.""" + + def timeout(*_args: object, **_kwargs: object) -> None: + command = "worker" + raise subprocess.TimeoutExpired(command, 0.1) + + monkeypatch.setattr(subprocess, "run", timeout) + result = _isolated(next(_cases(_arguments())), _arguments()) + assert result["status"] == "timeout" + assert result["censored"] is True + assert result["timing_result"] is None + + +def test_required_measurement_failure_remains_visible_in_json( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """A failed required case writes its diagnostic and makes the command fail.""" + monkeypatch.setattr( + "scripts.benchmark_hierarchical._isolated", + lambda *_args: {"status": "timeout", "censored": True}, + ) + output = tmp_path / "cases.jsonl" + status = main([ + "--groups", + "1", + "--local-dimensions", + "1", + "--steps", + "4", + "--solvers", + "core", + "--derivatives", + "trajectory", + "--output", + str(output), + "--require-success", + ]) + assert status == 1 + assert json.loads(output.read_text(encoding="utf-8"))["status"] == "timeout" + + +def test_invalid_case_dimension_is_rejected_before_work() -> None: + """Malformed fixture dimensions cannot produce meaningless benchmark rows.""" + with pytest.raises(SystemExit): + main(["--groups", "0"]) + + +def test_workers_override_inherited_persistent_compilation_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Inherited disk caches cannot turn cold compilation into cached loading.""" + monkeypatch.setenv("JAX_ENABLE_COMPILATION_CACHE", "true") + + def worker(*_args: object, **kwargs: object) -> subprocess.CompletedProcess[str]: + environment = kwargs["env"] + assert isinstance(environment, dict) + assert environment["JAX_ENABLE_COMPILATION_CACHE"] == "false" + return subprocess.CompletedProcess("worker", 0, '{"status":"success"}', "") + + monkeypatch.setattr(subprocess, "run", worker) + assert _isolated(next(_cases(_arguments())), _arguments())["status"] == "success"