From 0c7a151ba19585a9d86b068393564c6a334e6177 Mon Sep 17 00:00:00 2001 From: Areg Gevorgyan Date: Tue, 4 Aug 2026 18:02:57 -0400 Subject: [PATCH 1/7] feat(batch): add batch_map/batch_map_iter and *_many search-plumbing helpers Loops built on Alkahest are embarrassingly parallel at the candidate level, but every entry point is one-call-one-answer, so fan-out and per-item error handling had to be hand-rolled at each call site. Add a pure-Python batch/streaming layer (python/alkahest/_batch.py) that never raises for a single bad element: batch_map / batch_map_iter capture exceptions into a BatchItem carrying the failing exception's own E-* code (E-BATCH-001 fallback), always return results aligned to their input index, and support optional ThreadPoolExecutor fan-out. integrate_many / simplify_many / diff_many are thin batch_map wrappers over the three most common derivation entry points. Implements P1 search plumbing item 5 (temp-alkahest/planning/search-plumbing-p1.md). Co-authored-by: Cursor --- CHANGELOG.md | 11 + docs/mdbook/src/SUMMARY.md | 1 + docs/mdbook/src/batch.md | 121 ++++++++++ python/alkahest/__init__.py | 19 ++ python/alkahest/_batch.py | 419 ++++++++++++++++++++++++++++++++++ python/alkahest/exceptions.py | 3 + tests/test_batch_workload.py | 293 ++++++++++++++++++++++++ 7 files changed, 867 insertions(+) create mode 100644 docs/mdbook/src/batch.md create mode 100644 python/alkahest/_batch.py create mode 100644 tests/test_batch_workload.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c349bbc8..6cd95c65 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,17 @@ ### Added +- **Batch and streaming evaluation** (`alkahest._batch`, Python-only): `batch_map` / + `batch_map_iter` call a function once per item and **never raise** for a single bad + element — the exception is captured into a `BatchItem(index, ok, value, error, + elapsed_ms)`, with `error["code"]` set to the failing exception's own `E-*` code + (`E-BATCH-001` as a fallback for exceptions with none). `batch_map` always returns + results in input order, whether or not `parallel=True` fans the calls out over a + `ThreadPoolExecutor`; `batch_map_iter` streams in input order when sequential and in + completion order when parallel. `integrate_many`, `simplify_many`, and `diff_many` are + thin `batch_map` wrappers over the three most common derivation entry points. See + [`docs/mdbook/src/batch.md`](docs/mdbook/src/batch.md). + - **Python bindings for the parallel simplifiers**: `simplify_redex`, `simplify_auto` and `simplify_strategy` join the existing `simplify_par`. All take a single expression and return the same result as `simplify`; only diff --git a/docs/mdbook/src/SUMMARY.md b/docs/mdbook/src/SUMMARY.md index 64a8d9ba..8949f8af 100644 --- a/docs/mdbook/src/SUMMARY.md +++ b/docs/mdbook/src/SUMMARY.md @@ -18,6 +18,7 @@ - [Interoperability](./interop.md) - [Reinforcement learning](./rl.md) - [Derivation logs](./derivations.md) +- [Batch and streaming evaluation](./batch.md) - [Claim graphs](./claim-graphs.md) - [Lean certificates](./lean-certs.md) - [Certificate coverage](./certificate-coverage.md) diff --git a/docs/mdbook/src/batch.md b/docs/mdbook/src/batch.md new file mode 100644 index 00000000..77c1730c --- /dev/null +++ b/docs/mdbook/src/batch.md @@ -0,0 +1,121 @@ +# Batch and streaming evaluation + +Search loops are embarrassingly parallel *at the candidate level*: try to integrate a +hundred generated integrands, simplify a thousand rewrite targets, differentiate every +entry in a lookup table. Every Alkahest entry point is one-call-one-answer, so today +that fan-out is written by hand at every call site — and the first candidate that raises +aborts the whole batch unless the caller remembers `try/except` around every single call. + +`alkahest.batch_map` (and the `*_many` convenience wrappers over `integrate`, `simplify`, +and `diff`) do that fan-out once. They **never raise** for a single bad element — the +exception is caught and turned into a structured `BatchItem` carrying the failing +exception's stable `E-*` [diagnostic code](./errors.md), so a loop can tell "this +candidate has no elementary antiderivative" (a fine, expected answer) from "the whole +batch process crashed". + +```python +import alkahest as ak + +pool = ak.ExprPool() +x = pool.symbol("x") + +outs = ak.integrate_many([x**2, ak.log(ak.log(x)), ak.sin(x)], x) +for item in outs: + if item.ok: + print(item.index, "=>", item.value.value) + else: + print(item.index, "FAILED", item.error["code"], item.error["message"]) +``` + +```text +0 => (x^3 * 1/3) +1 FAILED E-INT-001 [E-INT-001] integrate: not implemented: ... +2 => (-1 * cos(x)) +``` + +## Honesty invariant + +`batch_map` always returns exactly one `BatchItem` per input, **in input order** — +a batch of 100 items yields a list of 100 items, full stop. Nothing in this module +silently drops a failing candidate; a failure is recorded as `ok=False` with its error, +never as a missing slot. + +## `BatchItem` + +| Field | Type | Meaning | +| --- | --- | --- | +| `index` | `int` | Position in the *original* input sequence — stable under `parallel=True` and under streaming in completion order | +| `ok` | `bool` | `True` iff the call returned normally | +| `value` | `Any \| None` | The call's return value (often a `DerivedResult`) on success; `None` on failure | +| `error` | `dict \| None` | `{"code", "message", "remediation", "type"}` on failure; `None` on success | +| `elapsed_ms` | `float \| None` | Wall-clock time spent inside the call for this item | + +Exactly one of `value` / `error` is populated: `ok=True` implies `error is None`. + +`error["code"]` is the raised exception's `.code` when it is an `AlkahestError`-like +exception — including the native error types, which expose the same attribute — +otherwise `alkahest._batch.UNEXPECTED_ERROR_CODE` (`"E-BATCH-001"`), the fallback for a +failure whose exception carries no diagnostic code of its own (e.g. a plain `ValueError` +raised by caller code passed to `batch_map`). + +## `batch_map` and `batch_map_iter` + +```python +def batch_map(fn, items, *, parallel=False, max_workers=None, **kwargs) -> list[BatchItem]: ... +def batch_map_iter(fn, items, *, parallel=False, max_workers=None, **kwargs) -> Iterator[BatchItem]: ... +``` + +Both call `fn(item, **kwargs)` once per item. `parallel=True` fans the calls out over a +`concurrent.futures.ThreadPoolExecutor`; some Alkahest hot paths (the parallel +simplifiers, NumPy evaluation) release the GIL for their native work, so a thread pool +can genuinely overlap them. For calls that hold the GIL throughout, `parallel=True` +mainly helps when `fn` itself does I/O or otherwise yields the GIL — it never makes +anything *incorrect*, only sometimes not faster. + +### Order guarantees + +- **`batch_map`** always returns results **in input order**, whether or not + `parallel=True`. This is the guarantee to reach for when you need + `zip(items, batch_map(...))` to line up. +- **`batch_map_iter`** documents two different behaviours by design: + - `parallel=False` streams **in input order** — item *i* is fully computed and + yielded before item *i + 1* starts. + - `parallel=True` streams **in completion order**, not input order. This is the whole + point of streaming under fan-out: a fast failure surfaces immediately instead of + waiting behind a slow item that happened to be submitted first. Every yielded + `BatchItem` still carries its original `index`, so a caller that needs input order + can sort by it, or just use `batch_map`. + +```python +# Streaming: react to failures as they arrive, without waiting for the slowest item. +for item in ak.batch_map_iter(ak.simplify, candidates, parallel=True): + if not item.ok: + log.warning("candidate %d failed: %s", item.index, item.error["code"]) +``` + +## `integrate_many` / `simplify_many` / `diff_many` + +Thin `batch_map` wrappers over the three most common derivation entry points: + +```python +def integrate_many(exprs, var, *bounds, parallel=False, max_workers=None, **kwargs) -> list[BatchItem]: ... +def simplify_many(exprs, *, parallel=False, max_workers=None, **kwargs) -> list[BatchItem]: ... +def diff_many(exprs, var, *, parallel=False, max_workers=None, **kwargs) -> list[BatchItem]: ... +``` + +`integrate_many` accepts optional trailing bounds (`a, b`) for a batch of definite +integrals, exactly like `alkahest.integrate`. `**kwargs` on every helper is forwarded to +the underlying call (e.g. `assumptions=` for `simplify_many`). + +```python +outs = ak.simplify_many(candidates, parallel=True) +ok = [o.value for o in outs if o.ok] +failed = [(o.index, o.error) for o in outs if not o.ok] +``` + +## Never raises — except for real interpreter signals + +`batch_map` and `batch_map_iter` catch `Exception`, not `BaseException`: a +`KeyboardInterrupt` or `SystemExit` still propagates and stops the batch, since +swallowing those would make the process unkillable. Everything else — including every +Alkahest `E-*` error and any exception your own `fn` raises — is captured. diff --git a/python/alkahest/__init__.py b/python/alkahest/__init__.py index b55fc7cd..325314ce 100644 --- a/python/alkahest/__init__.py +++ b/python/alkahest/__init__.py @@ -10,6 +10,14 @@ number_theory, research, # session-level claim graph (provenance objects) ) +from ._batch import ( + BatchItem, + batch_map, + batch_map_iter, + diff_many, + integrate_many, + simplify_many, +) from ._certificates import ( Certifiability, certifiable, @@ -1494,6 +1502,8 @@ def wrapper(*args, **kwargs): "ArbBall", "AssumptionError", "Assumptions", + # P1 search plumbing item 5 — batch/streaming fan-out + "BatchItem", "CadError", # V5-12 — certificate ledger "Certifiability", @@ -1594,6 +1604,9 @@ def wrapper(*args, **kwargs): "asinh", "atan", "atanh", + # P1 search plumbing item 5 — batch/streaming fan-out + "batch_map", + "batch_map_iter", "bessel_j0", "bessel_j1", "cad_lift", @@ -1619,6 +1632,8 @@ def wrapper(*args, **kwargs): # Calculus "diff", "diff_forward", + # P1 search plumbing item 5 — batch/streaming fan-out + "diff_many", "digamma", "diophantine", # Elliptic special functions (parameter convention m = k²) @@ -1647,6 +1662,8 @@ def wrapper(*args, **kwargs): "horner", "im", "integrate", + # P1 search plumbing item 5 — batch/streaming fan-out + "integrate_many", "interval_eval", "jacobian", "jit", @@ -1715,6 +1732,8 @@ def wrapper(*args, **kwargs): "simplify_enabled", "simplify_expanded", "simplify_log_exp", + # P1 search plumbing item 5 — batch/streaming fan-out + "simplify_many", "simplify_par", "simplify_pauli", "simplify_redex", diff --git a/python/alkahest/_batch.py b/python/alkahest/_batch.py new file mode 100644 index 00000000..e03039d5 --- /dev/null +++ b/python/alkahest/_batch.py @@ -0,0 +1,419 @@ +"""Batch and streaming evaluation — candidate-level fan-out for search loops. + +Loops built on Alkahest are embarrassingly parallel *at the candidate level*: +try to integrate a hundred generated integrands, simplify a thousand rewrite +targets, differentiate every entry in a lookup table. Today that fan-out has +to be written by hand at every call site, because :func:`alkahest.integrate` +(and friends) are one-call-one-answer: the first candidate that raises aborts +the whole batch unless the caller remembers ``try/except`` around every +single call. + +This module is that plumbing, done once. :func:`batch_map` (and the +``*_many`` convenience wrappers over :func:`alkahest.integrate`, +:func:`alkahest.simplify`, and :func:`alkahest.diff`) call a function once per +item and **never raise** for a single bad element — the exception is caught +and turned into a structured :class:`BatchItem` carrying the failing +exception's stable ``E-*`` diagnostic code (see ``exceptions.py``), so a loop +can tell "this candidate has no elementary antiderivative" (a fine, expected +answer) from "the whole batch process crashed". + +Honesty invariant +------------------ +:func:`batch_map` always returns exactly one :class:`BatchItem` per input, in +input order — a batch of 100 items yields a list of 100 items, full stop. +Nothing in this module silently drops a failing candidate; a failure is +recorded as ``ok=False`` with its error, never as a missing slot. + +Quick start +----------- +>>> import alkahest as ak +>>> pool = ak.ExprPool() +>>> x = pool.symbol("x") +>>> outs = ak.integrate_many([x**2, ak.log(ak.log(x))], x) +>>> outs[0].ok, outs[1].ok +(True, False) +>>> outs[1].error["code"] +'E-INT-001' + +Fan out over a thread pool (the Rust kernel releases the GIL for some hot +paths, e.g. the parallel simplifiers and NumPy evaluation, so a thread pool +can overlap those with other Python work; for calls that hold the GIL +throughout, ``parallel=True`` mainly helps when *fn* itself does I/O or +otherwise yields the GIL):: + + outs = ak.batch_map(ak.simplify, candidates, parallel=True) +""" + +from __future__ import annotations + +import time +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: # pragma: no cover - typing only + from collections.abc import Callable, Iterable, Iterator + +__all__ = [ + "UNEXPECTED_ERROR_CODE", + "BatchItem", + "batch_map", + "batch_map_iter", + "diff_many", + "integrate_many", + "simplify_many", +] + +#: Fallback ``error["code"]`` for a failure whose exception carries no ``.code`` +#: attribute of its own (i.e. it is not an :class:`~alkahest.exceptions.AlkahestError` +#: or a native PyO3 error exposing the same attribute). See ``exceptions.py`` +#: for the registry of codes this module composes with. +UNEXPECTED_ERROR_CODE = "E-BATCH-001" + + +@dataclass(frozen=True) +class BatchItem: + """One slot's outcome from :func:`batch_map` or a ``*_many`` helper. + + Exactly one of *value* / *error* is populated: ``ok=True`` implies + ``error is None`` and ``value`` is whatever *fn* returned (often a + :class:`~alkahest.DerivedResult`); ``ok=False`` implies ``value is None`` + and ``error`` describes what went wrong. + + Attributes + ---------- + index : int + Position of this item in the *original* input sequence. Stable + under ``parallel=True`` and under :func:`batch_map_iter` streaming in + completion order — an outcome can always be matched back to its + input via this field, even when results arrive out of order. + ok : bool + ``True`` iff *fn* returned normally for this item. + value : Any or None + *fn*'s return value on success; ``None`` on failure. + error : dict or None + ``{"code", "message", "remediation", "type"}`` on failure, ``None`` + on success. + + ``code`` + The raised exception's ``.code`` (e.g. ``"E-INT-001"``) when it + is an :class:`~alkahest.exceptions.AlkahestError`-like exception + — including the native PyO3 error types, which expose the same + attribute — otherwise :data:`UNEXPECTED_ERROR_CODE` + (``"E-BATCH-001"``). + ``message`` + ``str(exc)``. + ``remediation`` + The exception's ``.remediation``, or ``None`` when it has none. + ``type`` + ``type(exc).__name__``. + elapsed_ms : float or None + Wall-clock time spent inside *fn* for this item, in milliseconds. + """ + + index: int + ok: bool + value: Any | None = None + error: dict[str, Any] | None = None + elapsed_ms: float | None = None + + +def _describe_exception(exc: Exception) -> dict[str, Any]: + code = getattr(exc, "code", None) + remediation = getattr(exc, "remediation", None) + return { + "code": str(code) if code else UNEXPECTED_ERROR_CODE, + "message": str(exc), + "remediation": str(remediation) if remediation is not None else None, + "type": type(exc).__name__, + } + + +def _invoke(fn: Callable[..., Any], item: Any, index: int, kwargs: dict[str, Any]) -> BatchItem: + """Run ``fn(item, **kwargs)``, turning any :class:`Exception` into a :class:`BatchItem`. + + Deliberately catches ``Exception`` rather than ``BaseException``: a + ``KeyboardInterrupt`` (or ``SystemExit``) must still propagate and stop + the batch, since swallowing those would make the process unkillable. + """ + start = time.perf_counter() + try: + value = fn(item, **kwargs) + except Exception as exc: # intentional: never abort the batch for one bad element + elapsed_ms = (time.perf_counter() - start) * 1000.0 + error = _describe_exception(exc) + return BatchItem(index=index, ok=False, error=error, elapsed_ms=elapsed_ms) + elapsed_ms = (time.perf_counter() - start) * 1000.0 + return BatchItem(index=index, ok=True, value=value, elapsed_ms=elapsed_ms) + + +def batch_map( + fn: Callable[..., Any], + items: Iterable[Any], + *, + parallel: bool = False, + max_workers: int | None = None, + **kwargs: Any, +) -> list[BatchItem]: + """Call ``fn(item, **kwargs)`` for every item in *items*, never raising. + + Parameters + ---------- + fn : callable + Called as ``fn(item, **kwargs)`` for each item. Any :class:`Exception` + it raises is captured into that item's :class:`BatchItem` rather than + propagating. + items : iterable + The candidates to evaluate. Consumed once, eagerly (so a generator + works, but is fully materialised before any call happens). + parallel : bool + Fan out over a :class:`~concurrent.futures.ThreadPoolExecutor` when + true. Useful when *fn* releases the GIL for some or all of its work + (I/O, or a Rust call that calls ``py.allow_threads``); on pure + Python, GIL-bound work it will not speed anything up, but it also + will not make anything incorrect — order is preserved either way. + max_workers : int, optional + Forwarded to :class:`~concurrent.futures.ThreadPoolExecutor`. Ignored + when ``parallel=False``. + **kwargs + Forwarded to every call to *fn*. + + Returns + ------- + list of BatchItem + Exactly ``len(items)`` entries, **in input order** regardless of + *parallel* — this is the ordering guarantee :func:`batch_map_iter` + does not make under ``parallel=True``. + + Examples + -------- + >>> import alkahest as ak + >>> pool = ak.ExprPool() + >>> x = pool.symbol("x") + >>> outs = ak.batch_map(ak.simplify, [x + 0 * x, ak.log(ak.log(x))]) + >>> [o.ok for o in outs] + [True, True] + """ + materialized = list(items) + if not materialized: + return [] + if not parallel: + return [_invoke(fn, item, i, kwargs) for i, item in enumerate(materialized)] + + results: list[BatchItem | None] = [None] * len(materialized) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = { + executor.submit(_invoke, fn, item, i, kwargs): i for i, item in enumerate(materialized) + } + for future in futures: + results[futures[future]] = future.result() + return results # type: ignore[return-value] # every slot was filled above + + +def batch_map_iter( + fn: Callable[..., Any], + items: Iterable[Any], + *, + parallel: bool = False, + max_workers: int | None = None, + **kwargs: Any, +) -> Iterator[BatchItem]: + """Streaming counterpart of :func:`batch_map`. + + Order guarantee + ---------------- + - ``parallel=False``: yields **in input order** — item *i* is fully + computed and yielded before item *i + 1* starts, so + ``zip(items, batch_map_iter(fn, items))`` lines up. + - ``parallel=True``: yields **in completion order**, not input order. + This is deliberate: the point of streaming under fan-out is that a + fast failure surfaces immediately instead of waiting behind a slow + item that happened to be submitted first. Every yielded + :class:`BatchItem` still carries its original ``index``, so a caller + that needs input order can sort by it — or just use :func:`batch_map`, + which always returns in input order. + + Parameters + ---------- + fn, items, max_workers, **kwargs + As :func:`batch_map`. + parallel : bool + As :func:`batch_map`; also selects the completion-order streaming + behaviour described above. + + Yields + ------ + BatchItem + + Examples + -------- + Sequential streaming preserves input order: + + >>> import alkahest as ak + >>> pool = ak.ExprPool() + >>> x = pool.symbol("x") + >>> [item.index for item in ak.batch_map_iter(ak.simplify, [x, x + 0 * x])] + [0, 1] + """ + materialized = list(items) + if not materialized: + return + if not parallel: + for i, item in enumerate(materialized): + yield _invoke(fn, item, i, kwargs) + return + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + pending = { + executor.submit(_invoke, fn, item, i, kwargs) for i, item in enumerate(materialized) + } + while pending: + done, pending = wait(pending, return_when=FIRST_COMPLETED) + for future in done: + yield future.result() + + +# --------------------------------------------------------------------------- +# Domain-specific convenience wrappers +# --------------------------------------------------------------------------- +# +# Lazy import of the parent package (mirrors alkahest.research._ak()): +# _batch is imported from __init__.py before integrate/simplify/diff exist in +# its namespace, so these must resolve the module at call time, not at +# import time. + + +def _ak() -> Any: + import alkahest + + return alkahest + + +def integrate_many( + exprs: Iterable[Any], + var: Any, + *bounds: Any, + parallel: bool = False, + max_workers: int | None = None, + **kwargs: Any, +) -> list[BatchItem]: + """:func:`batch_map` over :func:`alkahest.integrate`, one call per integrand. + + Parameters + ---------- + exprs : iterable of Expr or DerivedResult + Integrands, evaluated independently. + var : Expr + Integration variable, shared by every call. + *bounds + Pass ``a, b`` for a definite integral over every integrand (see + :func:`alkahest.integrate`); omit for the indefinite integral. + parallel, max_workers + As :func:`batch_map`. + **kwargs + Forwarded to :func:`alkahest.integrate`. + + Returns + ------- + list of BatchItem + ``value`` is the :class:`~alkahest.DerivedResult` antiderivative (or + definite value) on success; on failure, ``error["code"]`` is the + integrator's own code (typically ``E-INT-001``) rather than + :data:`UNEXPECTED_ERROR_CODE`. + + Examples + -------- + >>> import alkahest as ak + >>> pool = ak.ExprPool() + >>> x = pool.symbol("x") + >>> outs = ak.integrate_many([x**2, ak.log(ak.log(x)), ak.sin(x)], x) + >>> [o.ok for o in outs] + [True, False, True] + """ + ak = _ak() + + def _one(expr: Any) -> Any: + return ak.integrate(expr, var, *bounds, **kwargs) + + return batch_map(_one, exprs, parallel=parallel, max_workers=max_workers) + + +def simplify_many( + exprs: Iterable[Any], + *, + parallel: bool = False, + max_workers: int | None = None, + **kwargs: Any, +) -> list[BatchItem]: + """:func:`batch_map` over :func:`alkahest.simplify`, one call per expression. + + Parameters + ---------- + exprs : iterable of Expr or DerivedResult + Expressions to simplify independently. + parallel, max_workers + As :func:`batch_map`. + **kwargs + Forwarded to :func:`alkahest.simplify` (e.g. ``assumptions=``). + + Returns + ------- + list of BatchItem + + Examples + -------- + >>> import alkahest as ak + >>> pool = ak.ExprPool() + >>> x = pool.symbol("x") + >>> outs = ak.simplify_many([x + 0 * x, x / x]) + >>> [o.value.value for o in outs] + [x, 1] + """ + ak = _ak() + + def _one(expr: Any) -> Any: + return ak.simplify(expr, **kwargs) + + return batch_map(_one, exprs, parallel=parallel, max_workers=max_workers) + + +def diff_many( + exprs: Iterable[Any], + var: Any, + *, + parallel: bool = False, + max_workers: int | None = None, + **kwargs: Any, +) -> list[BatchItem]: + """:func:`batch_map` over :func:`alkahest.diff`, one call per expression. + + Parameters + ---------- + exprs : iterable of Expr or DerivedResult + Expressions to differentiate independently. + var : Expr + Differentiation variable, shared by every call. + parallel, max_workers + As :func:`batch_map`. + **kwargs + Forwarded to :func:`alkahest.diff`. + + Returns + ------- + list of BatchItem + + Examples + -------- + >>> import alkahest as ak + >>> pool = ak.ExprPool() + >>> x = pool.symbol("x") + >>> outs = ak.diff_many([x**2, ak.sin(x)], x) + >>> [o.ok for o in outs] + [True, True] + """ + ak = _ak() + + def _one(expr: Any) -> Any: + return ak.diff(expr, var, **kwargs) + + return batch_map(_one, exprs, parallel=parallel, max_workers=max_workers) diff --git a/python/alkahest/exceptions.py b/python/alkahest/exceptions.py index 8c65ad40..42ec8e51 100644 --- a/python/alkahest/exceptions.py +++ b/python/alkahest/exceptions.py @@ -38,6 +38,9 @@ E-PARSE-* ParseError (reserved; parser not yet integrated) E-DOMAIN-* DomainError (reserved; Python-only pending Rust impl) E-CERT-001 CertificateUnavailableError (Python-only; certificate ledger) + E-BATCH-001 (Python-only; alkahest._batch fallback for a + batch_map/batch_map_iter item whose exception carried no + .code of its own — see docs/mdbook/src/batch.md) """ from __future__ import annotations diff --git a/tests/test_batch_workload.py b/tests/test_batch_workload.py new file mode 100644 index 00000000..edc6ba70 --- /dev/null +++ b/tests/test_batch_workload.py @@ -0,0 +1,293 @@ +"""Batch and streaming fan-out (``alkahest._batch``). + +Covers the properties that matter for a search loop driving hundreds of +candidates: one bad element never raises and never gets dropped, input order +is preserved by :func:`alkahest.batch_map` regardless of ``parallel``, +:func:`alkahest.batch_map_iter` streams in the documented order for each +mode, and a captured error carries a real diagnostic code — the integrator's +own ``E-INT-*`` code when the exception has one, ``E-BATCH-001`` otherwise. +""" + +from __future__ import annotations + +import time + +import alkahest as ak +import pytest +from alkahest._batch import UNEXPECTED_ERROR_CODE, BatchItem +from alkahest.exceptions import AlkahestError + + +@pytest.fixture +def pool(): + return ak.ExprPool() + + +# --------------------------------------------------------------------------- +# Exports +# --------------------------------------------------------------------------- + + +def test_public_names_exported_from_package_root(): + for name in ( + "BatchItem", + "batch_map", + "batch_map_iter", + "integrate_many", + "simplify_many", + "diff_many", + ): + assert hasattr(ak, name), f"alkahest.{name} not exported" + assert name in ak.__all__ + + +# --------------------------------------------------------------------------- +# batch_map: never raises, preserves order, mixed success/failure +# --------------------------------------------------------------------------- + + +def test_batch_map_never_raises_for_a_bad_element(pool): + x = pool.symbol("x") + # log(log(x)) has no elementary antiderivative the kernel implements today. + outs = ak.batch_map(lambda e: ak.integrate(e, x), [x**2, ak.log(ak.log(x)), ak.sin(x)]) + + assert len(outs) == 3 + assert [o.ok for o in outs] == [True, False, True] + assert outs[1].value is None + assert outs[1].error is not None + assert outs[0].error is None + assert outs[2].error is None + + +def test_batch_map_preserves_input_order_sequential(pool): + x = pool.symbol("x") + exprs = [x**n for n in range(1, 8)] + outs = ak.batch_map(lambda e: ak.diff(e, x), exprs) + assert [o.index for o in outs] == list(range(7)) + assert all(o.ok for o in outs) + + +def test_batch_map_preserves_input_order_parallel(): + # Deliberately vary sleep so completion order differs from input order; + # batch_map must still return results aligned to the original index. + delays = [0.05, 0.01, 0.03, 0.0, 0.02] + + def _work(i): + time.sleep(delays[i]) + return i * 10 + + outs = ak.batch_map(_work, range(len(delays)), parallel=True) + assert [o.index for o in outs] == list(range(len(delays))) + assert [o.value for o in outs] == [i * 10 for i in range(len(delays))] + assert all(o.ok for o in outs) + + +def test_batch_map_result_count_matches_input_even_with_all_failures(): + def _boom(_item): + raise ValueError("always fails") + + outs = ak.batch_map(_boom, range(5)) + assert len(outs) == 5 + assert all(not o.ok for o in outs) + assert all(o.value is None for o in outs) + + +def test_batch_map_empty_input_returns_empty_list(): + assert ak.batch_map(lambda x: x, []) == [] + assert ak.batch_map(lambda x: x, [], parallel=True) == [] + + +def test_batch_map_forwards_kwargs_to_fn(pool): + x = pool.symbol("x") + outs = ak.batch_map(ak.simplify, [x + 0 * x], assumptions=None) + assert outs[0].ok + + +# --------------------------------------------------------------------------- +# Error capture: stable codes, remediation, unexpected-failure fallback +# --------------------------------------------------------------------------- + + +def test_error_code_preserved_from_native_alkahest_exception(pool): + x = pool.symbol("x") + outs = ak.batch_map(lambda e: ak.integrate(e, x), [ak.log(ak.log(x))]) + error = outs[0].error + assert error["code"] == "E-INT-001" + assert error["remediation"] + assert "log" in error["message"] or "integrate" in error["message"] + assert error["type"] + + +def test_error_code_preserved_from_python_alkahest_error(): + class _CustomError(AlkahestError): + def __init__(self, message): + super().__init__(message, code="E-CUSTOM-042", remediation="do the other thing") + + def _raise(_item): + raise _CustomError("nope") + + outs = ak.batch_map(_raise, [1]) + assert outs[0].error == { + "code": "E-CUSTOM-042", + "message": "nope", + "remediation": "do the other thing", + "type": "_CustomError", + } + + +def test_unexpected_error_without_code_gets_batch_fallback_code(): + def _raise(_item): + raise RuntimeError("no .code attribute here") + + outs = ak.batch_map(_raise, [1]) + assert outs[0].error["code"] == UNEXPECTED_ERROR_CODE + assert outs[0].error["code"] == "E-BATCH-001" + assert outs[0].error["remediation"] is None + assert outs[0].error["type"] == "RuntimeError" + + +def test_keyboard_interrupt_is_not_captured(): + def _raise(_item): + raise KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + ak.batch_map(_raise, [1]) + + +# --------------------------------------------------------------------------- +# BatchItem shape +# --------------------------------------------------------------------------- + + +def test_batch_item_ok_success_shape(): + outs = ak.batch_map(lambda i: i * 2, [21]) + item = outs[0] + assert isinstance(item, BatchItem) + assert item.ok is True + assert item.value == 42 + assert item.error is None + assert item.elapsed_ms is not None + assert item.elapsed_ms >= 0.0 + + +def test_batch_item_is_frozen(): + item = BatchItem(index=0, ok=True, value=1) + with pytest.raises(Exception): # dataclasses.FrozenInstanceError is a subclass + item.value = 2 + + +# --------------------------------------------------------------------------- +# batch_map_iter: order guarantees +# --------------------------------------------------------------------------- + + +def test_batch_map_iter_sequential_is_input_order(pool): + x = pool.symbol("x") + exprs = [x**n for n in range(1, 6)] + items = list(ak.batch_map_iter(lambda e: ak.diff(e, x), exprs)) + assert [item.index for item in items] == list(range(5)) + + +def test_batch_map_iter_sequential_matches_batch_map(pool): + x = pool.symbol("x") + exprs = [x**2, ak.log(ak.log(x)), ak.sin(x)] + mapped = ak.batch_map(lambda e: ak.integrate(e, x), exprs) + streamed = list(ak.batch_map_iter(lambda e: ak.integrate(e, x), exprs)) + assert [o.ok for o in mapped] == [o.ok for o in streamed] + assert [o.index for o in mapped] == [o.index for o in streamed] + + +def test_batch_map_iter_parallel_streams_in_completion_order(): + # Item 0 sleeps longest, so it must be the *last* one yielded even though + # it was submitted first — this is exactly what streaming buys a caller. + delays = [0.08, 0.0, 0.0, 0.0] + + def _work(i): + time.sleep(delays[i]) + return i + + order = [item.index for item in ak.batch_map_iter(_work, range(4), parallel=True)] + assert order[-1] == 0 + assert set(order) == {0, 1, 2, 3} + + +def test_batch_map_iter_parallel_never_raises_and_covers_every_index(): + def _work(i): + if i % 2 == 0: + raise ValueError(f"bad item {i}") + return i + + items = list(ak.batch_map_iter(_work, range(6), parallel=True)) + assert {item.index for item in items} == set(range(6)) + for item in items: + if item.index % 2 == 0: + assert not item.ok + assert item.error["code"] == UNEXPECTED_ERROR_CODE + else: + assert item.ok + assert item.value == item.index + + +def test_batch_map_iter_empty_input_yields_nothing(): + assert list(ak.batch_map_iter(lambda x: x, [])) == [] + assert list(ak.batch_map_iter(lambda x: x, [], parallel=True)) == [] + + +# --------------------------------------------------------------------------- +# integrate_many / simplify_many / diff_many +# --------------------------------------------------------------------------- + + +def test_integrate_many_mixed_success_and_failure(pool): + x = pool.symbol("x") + outs = ak.integrate_many([x**2, ak.log(ak.log(x)), ak.sin(x)], x) + assert [o.ok for o in outs] == [True, False, True] + assert outs[1].error["code"] == "E-INT-001" + assert str(outs[0].value.value) == str(ak.integrate(x**2, x).value) + + +def test_integrate_many_definite_bounds(pool): + x = pool.symbol("x") + zero, one = pool.integer(0), pool.integer(1) + outs = ak.integrate_many([x**2, x**3], x, zero, one) + assert all(o.ok for o in outs) + assert str(outs[0].value.value) == str(ak.integrate(x**2, x, zero, one).value) + + +def test_integrate_many_parallel_preserves_order(pool): + x = pool.symbol("x") + exprs = [x**n for n in range(1, 12)] + outs = ak.integrate_many(exprs, x, parallel=True) + assert [o.index for o in outs] == list(range(len(exprs))) + assert all(o.ok for o in outs) + + +def test_simplify_many_mixed_success_and_failure(pool): + x = pool.symbol("x") + outs = ak.simplify_many([x + 0 * x, x / x]) + assert all(o.ok for o in outs) + assert str(outs[0].value.value) == "x" + + +def test_diff_many_mixed(pool): + x = pool.symbol("x") + outs = ak.diff_many([x**2, ak.sin(x), ak.cos(x)], x) + assert all(o.ok for o in outs) + assert str(outs[0].value.value) == str(ak.diff(x**2, x).value) + + +def test_many_helpers_never_raise_on_a_bad_element(pool): + x = pool.symbol("x") + # A non-Expr sentinel forces a failure path inside the wrapped call. + outs = ak.simplify_many([x, "not an expr"]) + assert outs[0].ok + assert not outs[1].ok + assert outs[1].error is not None + + +def test_many_helpers_are_batch_map_over_the_underlying_op(pool): + x = pool.symbol("x") + exprs = [x**2, x**3] + direct = [ak.batch_map(lambda e: ak.diff(e, x), exprs)[i].value.value for i in range(2)] + via_helper = [ak.diff_many(exprs, x)[i].value.value for i in range(2)] + assert [str(v) for v in direct] == [str(v) for v in via_helper] From a336ccffec3e7d7b600656719bde023d219e9106 Mon Sep 17 00:00:00 2001 From: Areg Gevorgyan Date: Tue, 4 Aug 2026 18:06:08 -0400 Subject: [PATCH 2/7] Add versioned, compact DerivedResult.to_dict/to_json envelopes P1 search plumbing item 6: agents pay for every character a call returns. Adds DerivedResult.to_dict(mode="full"|"compact") / .to_json(...) on the PyO3 binding, combining .value/.verification/.certificate_status/.steps into one envelope with a stable "alkahest.derived_result" kind discriminator and independent RESULT_SCHEMA_VERSION / STEPS_SCHEMA_VERSION constants (module-level, and DerivedResult.SCHEMA_VERSION / .STEPS_SCHEMA_VERSION class attributes). Compact mode drops before/after step text and uses short step keys (r/s), and prunes verification/certificate_status to their essential fields, but never renames, hides, or drops verification["status"] and never includes Lean certificate source in either mode, so the honesty signal survives the token-budget cut. python/alkahest/_result_schema.py documents the field-name contract (STEP_FIELDS / STEP_FIELDS_COMPACT) and re-exports the version constants for a single canonical import. Co-authored-by: Cursor --- CHANGELOG.md | 11 ++ alkahest-py/src/lib.rs | 162 ++++++++++++++++++++ docs/mdbook/src/derivations.md | 106 +++++++++++++ python/alkahest/__init__.py | 11 ++ python/alkahest/_result_schema.py | 99 +++++++++++++ tests/test_derived_result_schema.py | 221 ++++++++++++++++++++++++++++ 6 files changed, 610 insertions(+) create mode 100644 python/alkahest/_result_schema.py create mode 100644 tests/test_derived_result_schema.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c349bbc8..3089b778 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,17 @@ ### Added +- **`DerivedResult.to_dict` / `.to_json`: versioned, token-efficient result + envelopes** (P1 search plumbing item 6). Combines `.value`, `.verification`, + `.certificate_status`, and `.steps` into one dict/JSON string with a stable + `"kind": "alkahest.derived_result"` discriminator and independent + `RESULT_SCHEMA_VERSION` / `STEPS_SCHEMA_VERSION` constants (also exported at + module level and as `DerivedResult.SCHEMA_VERSION` / + `.STEPS_SCHEMA_VERSION`). `mode="compact"` drops `before`/`after` step text + and uses short step keys (`r`/`s`), but never renames, hides, or drops + `verification["status"]` and never includes Lean certificate source in + either mode. See `docs/mdbook/src/derivations.md`. + - **Python bindings for the parallel simplifiers**: `simplify_redex`, `simplify_auto` and `simplify_strategy` join the existing `simplify_par`. All take a single expression and return the same result as `simplify`; only diff --git a/alkahest-py/src/lib.rs b/alkahest-py/src/lib.rs index dc377355..c1edd2eb 100644 --- a/alkahest-py/src/lib.rs +++ b/alkahest-py/src/lib.rs @@ -1648,6 +1648,22 @@ impl PyAssumptions { // PyDerivedResult // --------------------------------------------------------------------------- +// Result envelope schema versions (P1 search-plumbing item 6). +// +// Bump `RESULT_SCHEMA_VERSION` for any change to the *envelope* shape +// returned by `DerivedResult.to_dict` / `to_json` (added/removed/renamed +// top-level keys: `kind`, `value`, `verification`, `certificate_status`, +// `steps`, `has_certificate`, ...). Bump `STEPS_SCHEMA_VERSION` for any +// change to the shape of a single `steps` entry (full-mode field names +// `rule`/`before`/`after`/`side_conditions`, or the compact-mode short-key +// mapping `r`/`s`). These are independent so a caller that only reads +// `.value` / `.verification` doesn't need to re-check its parsing when a +// `steps` internal detail changes, and vice versa. Both are exposed as +// module-level attributes (`alkahest.RESULT_SCHEMA_VERSION`) and as +// `DerivedResult` class attributes; see `docs/mdbook/src/derivations.md`. +const RESULT_SCHEMA_VERSION: u32 = 1; +const STEPS_SCHEMA_VERSION: u32 = 1; + #[pyclass(name = "DerivedResult")] struct PyDerivedResult { value: PyExpr, @@ -1665,6 +1681,16 @@ struct PyDerivedResult { #[pymethods] impl PyDerivedResult { + /// Envelope schema version for :meth:`to_dict` / :meth:`to_json`. See + /// module-level ``alkahest.RESULT_SCHEMA_VERSION``. + #[classattr] + const SCHEMA_VERSION: u32 = crate::RESULT_SCHEMA_VERSION; + + /// Schema version of each entry in ``.steps`` / the ``steps`` key of + /// :meth:`to_dict`. See module-level ``alkahest.STEPS_SCHEMA_VERSION``. + #[classattr] + const STEPS_SCHEMA_VERSION: u32 = crate::STEPS_SCHEMA_VERSION; + #[getter] fn value(&self) -> PyExpr { self.value.clone() @@ -1936,6 +1962,139 @@ impl PyDerivedResult { } Ok(false) } + + /// Stable, versioned dict envelope for machine/agent consumers. + /// + /// ``mode="full"`` (default) carries the same information as + /// ``.steps`` / ``.verification`` / ``.certificate_status`` combined + /// under one discriminated envelope (``kind="alkahest.derived_result"``, + /// versioned by :attr:`SCHEMA_VERSION` / :attr:`STEPS_SCHEMA_VERSION`). + /// + /// ``mode="compact"`` is strictly smaller and intended for hot loops / + /// tight agent context budgets: + /// + /// * ``steps`` entries drop ``before``/``after`` (usually the largest + /// strings in a derivation — the single biggest token cost) and use + /// short keys: ``r`` (rule name) and ``s`` (``side_conditions``, + /// *omitted entirely* when the list is empty). + /// * ``verification`` is pruned to ``status`` and + /// ``externally_verified`` only — the two fields that carry the + /// honesty signal (whether this result is verified, and whether that + /// verification happened out-of-process). ``verification["status"]`` + /// is **never** renamed, abbreviated, or omitted in compact mode. + /// * ``certificate_status`` is pruned to ``certifiable`` and + /// ``reason``; the ``blocking_steps`` diagnostic list (which repeats + /// ``before``/``after`` expression text) is dropped. + /// + /// Neither mode ever includes the Lean certificate source text — use + /// the :attr:`certificate` getter for that. ``has_certificate`` (a + /// bool) plus ``certificate_status.reason`` is enough to know whether a + /// certificate exists and why not, without paying for the source. + /// + /// Raises ``ValueError`` for any ``mode`` other than ``"full"`` / + /// ``"compact"``. + #[pyo3(signature = (mode="full"))] + fn to_dict<'py>(&self, py: Python<'py>, mode: &str) -> PyResult> { + let compact = derived_result_mode_is_compact(mode)?; + + let has_certificate = self.certificate(py).is_some(); + let verification_full = self.verification(py); + let certificate_status_full = self.certificate_status(py); + + let out = PyDict::new_bound(py); + out.set_item("kind", "alkahest.derived_result")?; + out.set_item("schema_version", RESULT_SCHEMA_VERSION)?; + out.set_item("steps_schema_version", STEPS_SCHEMA_VERSION)?; + out.set_item("value", self.value.__str__(py))?; + + if compact { + let verification = PyDict::new_bound(py); + verification.set_item( + "status", + verification_full + .get_item("status")? + .expect("verification always sets status"), + )?; + verification.set_item( + "externally_verified", + verification_full + .get_item("externally_verified")? + .expect("verification always sets externally_verified"), + )?; + out.set_item("verification", verification)?; + + let certificate_status = PyDict::new_bound(py); + certificate_status.set_item( + "certifiable", + certificate_status_full + .get_item("certifiable")? + .expect("certificate_status always sets certifiable"), + )?; + certificate_status.set_item( + "reason", + certificate_status_full + .get_item("reason")? + .expect("certificate_status always sets reason"), + )?; + out.set_item("certificate_status", certificate_status)?; + } else { + out.set_item("verification", verification_full)?; + out.set_item("certificate_status", certificate_status_full)?; + } + + out.set_item("steps", self.steps_dict_list(py, compact))?; + out.set_item("has_certificate", has_certificate)?; + Ok(out) + } + + /// ``json.dumps(self.to_dict(mode=mode))``, via Python's own ``json`` + /// module so the output matches what an agent's own ``json.dumps`` + /// would produce. See :meth:`to_dict` for the schema. + #[pyo3(signature = (mode="full"))] + fn to_json(&self, py: Python<'_>, mode: &str) -> PyResult { + derived_result_mode_is_compact(mode)?; + let dict = self.to_dict(py, mode)?; + let json = PyModule::import_bound(py, "json")?; + json.getattr("dumps")?.call1((dict,))?.extract() + } +} + +/// Shared `mode` validation for `to_dict` / `to_json`. Returns `true` for +/// `"compact"`, `false` for `"full"`. +fn derived_result_mode_is_compact(mode: &str) -> PyResult { + match mode { + "full" => Ok(false), + "compact" => Ok(true), + other => Err(pyo3::exceptions::PyValueError::new_err(format!( + "DerivedResult.to_dict/to_json: mode must be 'full' or 'compact', got {other:?}" + ))), + } +} + +impl PyDerivedResult { + /// `steps` list for [`PyDerivedResult::to_dict`]. Full mode mirrors the + /// `.steps` getter exactly (`rule`/`before`/`after`/`side_conditions`); + /// compact mode uses short keys and drops `before`/`after`, omitting + /// `s` entirely when `side_conditions` is empty. + fn steps_dict_list<'py>(&self, py: Python<'py>, compact: bool) -> Bound<'py, PyList> { + let list = PyList::empty_bound(py); + for (rule, before, after, conds) in &self.steps_raw { + let d = PyDict::new_bound(py); + if compact { + d.set_item("r", rule).unwrap(); + if !conds.is_empty() { + d.set_item("s", conds).unwrap(); + } + } else { + d.set_item("rule", rule).unwrap(); + d.set_item("before", before).unwrap(); + d.set_item("after", after).unwrap(); + d.set_item("side_conditions", conds).unwrap(); + } + list.append(d).unwrap(); + } + list + } } fn make_derived_result( @@ -9825,5 +9984,8 @@ fn alkahest(m: &Bound<'_, PyModule>) -> PyResult<()> { )?; // V1-15: compile-time flag so Python tests can skip egraph-dependent assertions. m.add("HAS_EGRAPH", cfg!(feature = "egraph"))?; + // P1 search-plumbing item 6: versioned DerivedResult.to_dict/to_json envelope. + m.add("RESULT_SCHEMA_VERSION", RESULT_SCHEMA_VERSION)?; + m.add("STEPS_SCHEMA_VERSION", STEPS_SCHEMA_VERSION)?; Ok(()) } diff --git a/docs/mdbook/src/derivations.md b/docs/mdbook/src/derivations.md index 58d58743..eb1029c7 100644 --- a/docs/mdbook/src/derivations.md +++ b/docs/mdbook/src/derivations.md @@ -24,6 +24,13 @@ dr = diff(sin(x**2), x) | `.verification` | `dict` | Evidence status, artifact format, external-check status, and side conditions | | `.certificate` | `str \| None` | Generated Lean 4 source, when a derivation log exists | +### Methods + +| Method | Description | +|---|---| +| `.to_dict(mode="full")` | Versioned dict envelope combining `.value`/`.verification`/`.certificate_status`/`.steps`; see [Machine-parseable output](#machine-parseable-output-to_dict--to_json) below | +| `.to_json(mode="full")` | `json.dumps(self.to_dict(mode=mode))` | + ## Rewrite steps Each step in `.steps` is a dict with: @@ -101,6 +108,105 @@ all_steps = simplified.steps + derived.steps For operations like `integrate` that internally call `simplify`, the log includes the simplification sub-steps interleaved with the integration steps. +## Machine-parseable output: `to_dict` / `to_json` + +Agents pay for every character a call returns. `.steps`, `.verification`, and +`.certificate_status` are convenient to poke at interactively, but stitching +them into one payload for logging, RPC, or a context window means writing +that glue yourself, on every call site, forever. `DerivedResult.to_dict()` +and `DerivedResult.to_json()` give you the stitched, versioned envelope +directly: + +```python +dr = diff(sin(x**2), x) + +full = dr.to_dict() # mode="full" is the default +compact = dr.to_dict(mode="compact") # short keys, token-efficient +json_str = dr.to_json(mode="compact") # json.dumps(dr.to_dict(mode="compact")) +``` + +### Envelope shape (`mode="full"`) + +```json +{ + "kind": "alkahest.derived_result", + "schema_version": 1, + "steps_schema_version": 1, + "value": "", + "verification": { "status": "...", "evidence": "...", "externally_verified": false, "artifact_format": "...", "side_conditions": [...], "method": "..." }, + "certificate_status": { "certifiable": true, "reason": "...", "blocking_steps": [] }, + "steps": [ {"rule": "...", "before": "...", "after": "...", "side_conditions": [...]}, ... ], + "has_certificate": true +} +``` + +`verification` and `certificate_status` are exactly the dicts returned by the +`.verification` and `.certificate_status` getters; `steps` is exactly `.steps`. +`kind` is a stable discriminator string — useful when logs or RPC payloads +mix `DerivedResult` envelopes with other structured outputs (e.g. error +envelopes carrying `E-SUBSYSTEM-NNN` codes). + +### Schema versions + +Two independent version constants, both starting at `1`: + +| Constant | Governs | +|---|---| +| `alkahest.RESULT_SCHEMA_VERSION` | The envelope: the set of top-level keys (`kind`, `value`, `verification`, `certificate_status`, `steps`, `has_certificate`, ...) | +| `alkahest.STEPS_SCHEMA_VERSION` | One entry of `steps`: full-mode field names and the compact-mode short-key mapping | + +Also available as `DerivedResult.SCHEMA_VERSION` / `DerivedResult.STEPS_SCHEMA_VERSION` +class attributes, and documented alongside the field-name contract in +`alkahest._result_schema` (`STEP_FIELDS`, `STEP_FIELDS_COMPACT`). Either +constant is bumped independently if its shape ever changes, so pinning +`schema_version`/`steps_schema_version` in your own parsing code is safe +across upgrades that don't touch the piece you depend on. + +### Compact mode + +`mode="compact"` keeps the same top-level envelope shape but shrinks the +biggest token costs: + +- **Steps** use short keys — `r` for `rule`, `s` for `side_conditions` — and + **omit `before`/`after` entirely**. Those two expression strings are + usually the largest part of a multi-step derivation and the single + biggest win for token budget. `s` is itself omitted from a step's dict + when that step has no side conditions (the common case). +- **`verification`** is pruned to `status` and `externally_verified` only. + These are the two fields that carry the honesty signal — whether the + result is verified, and whether that verification happened out-of-process + — so they are never renamed, abbreviated, or dropped in compact mode. +- **`certificate_status`** is pruned to `certifiable` and `reason`; the + `blocking_steps` diagnostic list (which repeats `before`/`after` text) is + dropped. +- **No mode ever includes Lean certificate source text.** `has_certificate` + (bool) plus `certificate_status["reason"]` is enough to know whether a + certificate exists and, if not, why — without paying for the source. Use + the `.certificate` getter when you actually need the Lean source. + +```python +dr.to_dict(mode="compact") +# { +# "kind": "alkahest.derived_result", +# "schema_version": 1, +# "steps_schema_version": 1, +# "value": "...", +# "verification": {"status": "certificate_available", "externally_verified": false}, +# "certificate_status": {"certifiable": true, "reason": "emitted"}, +# "steps": [{"r": "diff_sin"}, {"r": "sqrt_of_square_positive", "s": ["x > 0"]}], +# "has_certificate": true +# } +``` + +Prefer `to_dict(mode="compact")` / `to_json(mode="compact")` over reading +`.steps` directly in hot loops — batch derivations, autoresearch search +plumbing, or anywhere you're serialising many `DerivedResult`s and only need +the rule names, side conditions, and verification status rather than full +before/after expression text. + +An invalid `mode` (anything other than `"full"`/`"compact"`) raises +`ValueError`. + ## Beyond one call `DerivedResult` is per-call. To accumulate many results into a citable, serialisable, diff --git a/python/alkahest/__init__.py b/python/alkahest/__init__.py index b55fc7cd..c79f9769 100644 --- a/python/alkahest/__init__.py +++ b/python/alkahest/__init__.py @@ -45,6 +45,12 @@ map_exprs, unflatten_exprs, ) +from ._result_schema import ( + RESULT_SCHEMA_VERSION, + STEP_FIELDS, + STEP_FIELDS_COMPACT, + STEPS_SCHEMA_VERSION, +) from ._transform import ( CompiledGradTracedFn, CompiledTracedFn, @@ -1485,6 +1491,11 @@ def wrapper(*args, **kwargs): "HAS_EGRAPH", # Phase 16 "ODE", + # P1 search plumbing item 6 — DerivedResult.to_dict/to_json schema + "RESULT_SCHEMA_VERSION", + "STEPS_SCHEMA_VERSION", + "STEP_FIELDS", + "STEP_FIELDS_COMPACT", # Phase 18 "AcausalSystem", # Exceptions (V1-3 — stable diagnostic codes) diff --git a/python/alkahest/_result_schema.py b/python/alkahest/_result_schema.py new file mode 100644 index 00000000..77e69338 --- /dev/null +++ b/python/alkahest/_result_schema.py @@ -0,0 +1,99 @@ +"""Versioned schema documentation for :class:`alkahest.DerivedResult` output. + +P1 search plumbing item 6 — agents pay for every character an operation +returns, so :meth:`DerivedResult.to_dict` / :meth:`DerivedResult.to_json` +(implemented in the native extension, ``alkahest-py/src/lib.rs``) expose a +stable, versioned envelope with a token-efficient ``mode="compact"``. + +This module carries no logic of its own — the encoders live in Rust so +``to_dict``/``to_json`` never drift from ``.steps``/``.verification``. It +re-exports the two version constants from the compiled extension and +documents the field-name contract in one place so agents (and this +project's own tests/docs) have a single canonical import. + +Schema versions +---------------- +``RESULT_SCHEMA_VERSION`` covers the *envelope*: the set of top-level keys +(``kind``, ``schema_version``, ``steps_schema_version``, ``value``, +``verification``, ``certificate_status``, ``steps``, ``has_certificate``). +Bump it when a top-level key is added, removed, or renamed. + +``STEPS_SCHEMA_VERSION`` covers one entry of ``steps``: the full-mode field +names (``rule``, ``before``, ``after``, ``side_conditions``) and the +compact-mode short-key mapping (``r``, ``s``). Bump it independently of +``RESULT_SCHEMA_VERSION`` when that shape changes. + +Both start at ``1`` and are also available as ``DerivedResult.SCHEMA_VERSION`` +/ ``DerivedResult.STEPS_SCHEMA_VERSION`` class attributes. + +Full-mode step fields +---------------------- +Each entry of ``.steps`` (and of ``to_dict()["steps"]`` in full mode) is a +dict with exactly these keys, matching the ``.steps`` getter that predates +this schema: + +* ``rule`` — rewrite rule name (``str``) +* ``before`` — expression display string before the rewrite (``str``) +* ``after`` — expression display string after the rewrite (``str``) +* ``side_conditions`` — side conditions recorded for the rewrite + (``list[str]``, possibly empty) + +Compact-mode step fields +------------------------- +``to_dict(mode="compact")["steps"]`` entries use short keys and drop the +``before``/``after`` strings — usually the largest strings in a derivation, +and the single biggest token cost of a multi-step result: + +* ``r`` — same as full-mode ``rule`` +* ``s`` — same as full-mode ``side_conditions``, but the key is **omitted + entirely** when the list is empty (most steps have none) + +Honesty in compact mode +------------------------ +Compact mode never drops or renames ``verification["status"]`` — the field +that distinguishes ``exactly_verified`` / ``numerically_checked`` / +``certificate_available`` / ``unverified`` — because that is the honesty +signal this schema exists to preserve token budget around, not obscure. +``verification["externally_verified"]`` (always ``False`` today; no +external Lean check has ever run in-process) is kept alongside it so a +compact reader cannot mistake a generated certificate for a checked one. +Compact mode also never includes Lean certificate source text in either +mode's ``certificate_status`` — use the ``.certificate`` getter for that; +``has_certificate`` plus ``certificate_status["reason"]`` is enough to know +whether one exists and why not. + +Example +------- +>>> import alkahest as ak +>>> pool = ak.ExprPool() +>>> x = pool.symbol("x") +>>> dr = ak.diff(ak.sin(x), x) +>>> full = dr.to_dict() # mode="full" is the default +>>> compact = dr.to_dict(mode="compact") +>>> len(dr.to_json(mode="compact")) <= len(dr.to_json(mode="full")) +True +>>> full["verification"]["status"] == compact["verification"]["status"] +True + +Agents in hot loops (batch derivations, autoresearch search plumbing) +should prefer ``to_dict(mode="compact")`` / ``to_json(mode="compact")`` +over the full envelope or over reading ``.steps`` directly. +""" + +from .alkahest import RESULT_SCHEMA_VERSION, STEPS_SCHEMA_VERSION + +__all__ = [ + "RESULT_SCHEMA_VERSION", + "STEPS_SCHEMA_VERSION", + "STEP_FIELDS", + "STEP_FIELDS_COMPACT", +] + +#: Field names of a full-mode step record (``.steps`` entries and +#: ``to_dict()["steps"]`` entries under ``mode="full"``). +STEP_FIELDS: tuple[str, ...] = ("rule", "before", "after", "side_conditions") + +#: Short keys of a compact-mode step record +#: (``to_dict(mode="compact")["steps"]`` entries). ``s`` is omitted from a +#: given step's dict entirely when that step has no side conditions. +STEP_FIELDS_COMPACT: tuple[str, ...] = ("r", "s") diff --git a/tests/test_derived_result_schema.py b/tests/test_derived_result_schema.py new file mode 100644 index 00000000..ece54cc3 --- /dev/null +++ b/tests/test_derived_result_schema.py @@ -0,0 +1,221 @@ +"""P1 search plumbing item 6 — versioned, machine-parseable ``DerivedResult``. + +Agents pay for every character. These tests pin the ``to_dict`` / ``to_json`` +envelope schema (``RESULT_SCHEMA_VERSION`` / ``STEPS_SCHEMA_VERSION``) and, +critically, that ``mode="compact"`` never drops or obscures the honesty +signal (``verification["status"]``) even while it is strictly smaller than +``mode="full"``. +""" + +from __future__ import annotations + +import json + +import alkahest as ak +import pytest +from alkahest._result_schema import ( + RESULT_SCHEMA_VERSION, + STEP_FIELDS, + STEP_FIELDS_COMPACT, + STEPS_SCHEMA_VERSION, +) + + +@pytest.fixture +def pool(): + return ak.ExprPool() + + +def _multistep_derivation(pool): + """A derivation with several rewrite steps and at least one side condition.""" + x = pool.symbol("x", domain=ak.Domain.Positive) + return ak.diff(ak.sqrt(x**2) * ak.sin(x), x) + + +# --------------------------------------------------------------------------- +# Schema version constants +# --------------------------------------------------------------------------- + + +def test_schema_version_constants_are_one(): + # Pin the initial version; bump deliberately (with a docs update) if the + # envelope or step shape ever changes. + assert RESULT_SCHEMA_VERSION == 1 + assert STEPS_SCHEMA_VERSION == 1 + + +def test_schema_version_constants_exported_from_top_level(): + assert ak.RESULT_SCHEMA_VERSION == RESULT_SCHEMA_VERSION + assert ak.STEPS_SCHEMA_VERSION == STEPS_SCHEMA_VERSION + assert "RESULT_SCHEMA_VERSION" in ak.__all__ + assert "STEPS_SCHEMA_VERSION" in ak.__all__ + + +def test_schema_version_class_attrs_match_module_constants(): + assert ak.DerivedResult.SCHEMA_VERSION == RESULT_SCHEMA_VERSION + assert ak.DerivedResult.STEPS_SCHEMA_VERSION == STEPS_SCHEMA_VERSION + + +def test_documented_step_fields_match_actual_dict_keys(pool): + dr = _multistep_derivation(pool) + assert dr.steps, "fixture derivation must have at least one step" + assert set(dr.steps[0].keys()) == set(STEP_FIELDS) + + compact = dr.to_dict(mode="compact") + compact_keys = set() + for step in compact["steps"]: + compact_keys.update(step.keys()) + # every key seen in a compact step is one of the two documented short keys + assert compact_keys <= set(STEP_FIELDS_COMPACT) + + +# --------------------------------------------------------------------------- +# Full mode +# --------------------------------------------------------------------------- + + +def test_full_mode_has_required_keys_and_versions(pool): + dr = _multistep_derivation(pool) + full = dr.to_dict() # mode="full" is the default + assert dr.to_dict(mode="full") == full + + required = { + "kind", + "schema_version", + "steps_schema_version", + "value", + "verification", + "certificate_status", + "steps", + "has_certificate", + } + assert required <= set(full.keys()) + assert full["schema_version"] == RESULT_SCHEMA_VERSION + assert full["steps_schema_version"] == STEPS_SCHEMA_VERSION + assert full["value"] == str(dr.value) + assert full["has_certificate"] == (dr.certificate is not None) + + +def test_full_mode_steps_match_steps_getter(pool): + dr = _multistep_derivation(pool) + full = dr.to_dict() + assert full["steps"] == dr.steps + + +def test_full_mode_verification_matches_getter(pool): + dr = _multistep_derivation(pool) + full = dr.to_dict() + assert full["verification"] == dr.verification + + +def test_full_mode_certificate_status_matches_getter(pool): + dr = _multistep_derivation(pool) + full = dr.to_dict() + assert full["certificate_status"] == dr.certificate_status + + +# --------------------------------------------------------------------------- +# Compact mode: smaller, but never dishonest +# --------------------------------------------------------------------------- + + +def test_compact_is_strictly_smaller_than_full_for_multistep_derivation(pool): + dr = _multistep_derivation(pool) + full_json = dr.to_json(mode="full") + compact_json = dr.to_json(mode="compact") + assert len(compact_json) < len(full_json) + + +def test_verification_status_present_and_equal_in_both_modes(pool): + dr = _multistep_derivation(pool) + full = dr.to_dict(mode="full") + compact = dr.to_dict(mode="compact") + + assert "status" in full["verification"] + assert "status" in compact["verification"] + # the honesty signal itself is never renamed, abbreviated, or changed + assert full["verification"]["status"] == compact["verification"]["status"] + assert compact["verification"]["status"] == dr.verification["status"] + + +def test_compact_verification_keeps_externally_verified(pool): + dr = _multistep_derivation(pool) + compact = dr.to_dict(mode="compact") + assert compact["verification"]["externally_verified"] == dr.verification["externally_verified"] + + +def test_compact_steps_use_short_keys_and_drop_before_after(pool): + dr = _multistep_derivation(pool) + compact = dr.to_dict(mode="compact") + assert len(compact["steps"]) == len(dr.steps) + for step in compact["steps"]: + assert "before" not in step + assert "after" not in step + assert "rule" not in step + assert "r" in step + + +def test_compact_steps_omit_empty_side_conditions_but_keep_nonempty(pool): + dr = _multistep_derivation(pool) + compact = dr.to_dict(mode="compact") + saw_side_condition = False + for full_step, compact_step in zip(dr.steps, compact["steps"]): + if full_step["side_conditions"]: + assert compact_step["s"] == full_step["side_conditions"] + saw_side_condition = True + else: + assert "s" not in compact_step + assert saw_side_condition, "fixture derivation must exercise a side condition" + + +def test_compact_certificate_status_omits_blocking_steps(pool): + dr = _multistep_derivation(pool) + compact = dr.to_dict(mode="compact") + assert "certifiable" in compact["certificate_status"] + assert "reason" in compact["certificate_status"] + assert "blocking_steps" not in compact["certificate_status"] + assert compact["certificate_status"]["reason"] == dr.certificate_status["reason"] + + +def test_compact_never_contains_lean_source_text(pool): + dr = _multistep_derivation(pool) + compact_json = dr.to_json(mode="compact") + # Lean certificate source is theorem/proof syntax; make sure none of it + # leaked into the compact envelope regardless of derivation shape. + for marker in ("theorem ", "import Mathlib", ":= by"): + assert marker not in compact_json + + +# --------------------------------------------------------------------------- +# JSON round-trip and discriminator +# --------------------------------------------------------------------------- + + +def test_to_json_round_trips_for_both_modes(pool): + dr = _multistep_derivation(pool) + for mode in ("full", "compact"): + loaded = json.loads(dr.to_json(mode=mode)) + assert loaded == dr.to_dict(mode=mode) + + +def test_kind_discriminator_is_stable_across_modes(pool): + dr = _multistep_derivation(pool) + assert dr.to_dict(mode="full")["kind"] == "alkahest.derived_result" + assert dr.to_dict(mode="compact")["kind"] == "alkahest.derived_result" + + +def test_invalid_mode_raises_value_error(pool): + dr = _multistep_derivation(pool) + with pytest.raises(ValueError): + dr.to_dict(mode="bogus") + with pytest.raises(ValueError): + dr.to_json(mode="bogus") + + +def test_simple_zero_step_derivation_round_trips(pool): + x = pool.symbol("x") + dr = ak.diff(x, x) + for mode in ("full", "compact"): + loaded = json.loads(dr.to_json(mode=mode)) + assert loaded["kind"] == "alkahest.derived_result" + assert "status" in loaded["verification"] From a7d9262139dc9521022e345499ffa7fb0eca458b Mon Sep 17 00:00:00 2001 From: Areg Gevorgyan Date: Tue, 4 Aug 2026 18:27:05 -0400 Subject: [PATCH 3/7] Add budgets, cooperative cancellation, and a determinism seed (P1 search plumbing item 4) A loop fanning out many candidates (e.g. Groebner/integrate search) needs to bound one candidate's cost instead of hanging on it or relying on an OS kill. This adds a cooperative budget checkpoint in alkahest-core: - alkahest_core::budget: Budget{wall, max_steps, seed}, a thread-local nesting stack (enter()/BudgetGuard), check() for wall-clock/step/cancel trips, and a process-wide AtomicBool cancel flag so an orchestrator thread can stop a heavy call running elsewhere. BudgetError maps to stable E-BUDGET-001 (wall clock), E-BUDGET-002 (step limit), and E-BUDGET-003 (cancelled) via AlkahestError. - integrate::engine checks the budget at its top-level entry and its recursion boundary and returns IntegrationError::Budget; the Risch engine propagates it immediately instead of continuing to spend budget. - simplify::engine checks once per rewrite pass/batch and stops early (like max_iterations) since simplify has no Result to raise through. - PyO3 bindings push/pop the Rust budget stack from Python's context() and map BudgetError to a new PyBudgetExceededError. On the Python side: Budget(wall_ms=, max_steps=, seed=) dataclass, context(budget=...) to scope it, BudgetExceededError in the exception hierarchy, request_cancel()/clear_cancel()/is_cancelled(), budget_seed() for RNG-consuming samplers that want reproducible runs, and run_with_wall_fallback() as a documented Python-layer supplement (worker thread + timeout) for calls without a Rust checkpoint on every path. Adds docs/mdbook/src/budgets.md, Rust unit tests for the budget module, and tests/test_budget.py covering nesting, seed round-trips, step/wall trips, cross-thread cancellation, and run_with_wall_fallback. Co-authored-by: Cursor --- CHANGELOG.md | 18 + alkahest-core/src/budget/mod.rs | 473 +++++++++++++++++++++++ alkahest-core/src/errors/codes.rs | 4 + alkahest-core/src/integrate/engine.rs | 29 ++ alkahest-core/src/integrate/risch/mod.rs | 7 + alkahest-core/src/lib.rs | 7 + alkahest-core/src/simplify/engine.rs | 19 + alkahest-py/src/lib.rs | 122 +++++- docs/mdbook/src/SUMMARY.md | 1 + docs/mdbook/src/budgets.md | 166 ++++++++ docs/mdbook/src/errors.md | 4 +- python/alkahest/__init__.py | 28 ++ python/alkahest/_budget.py | 239 ++++++++++++ python/alkahest/_context.py | 36 ++ python/alkahest/exceptions.py | 25 ++ tests/test_budget.py | 281 ++++++++++++++ 16 files changed, 1457 insertions(+), 2 deletions(-) create mode 100644 alkahest-core/src/budget/mod.rs create mode 100644 docs/mdbook/src/budgets.md create mode 100644 python/alkahest/_budget.py create mode 100644 tests/test_budget.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c349bbc8..a5d62e7c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,24 @@ ### Added +- **Budgets, cooperative cancellation, and a determinism seed** (P1 search + plumbing item 4): `alkahest.Budget(wall_ms=..., max_steps=..., seed=...)` + and `alkahest.context(budget=...)` push a wall-clock/step budget into a + new Rust-side cooperative checkpoint (`alkahest_core::budget`), so a + fan-out loop trying many candidate integrals/rewrites can bound one + candidate's cost instead of hanging on it. `alkahest.integrate` checks it + at its top-level entry and its recursion boundary and raises + `BudgetExceededError` (`E-BUDGET-001` wall clock, `E-BUDGET-002` step + limit, `E-BUDGET-003` cancelled) rather than running unbounded; + `alkahest.simplify` checks it once per rewrite pass and, since it has no + error channel, stops early instead of raising (`run_with_wall_fallback` + supplements this with a hard deadline via a worker thread when needed). + `alkahest.request_cancel()` / `clear_cancel()` / `is_cancelled()` expose a + process-wide cancellation flag so an orchestrator thread can stop a heavy + call running elsewhere; `alkahest.budget_seed()` exposes the active + budget's seed to RNG-consuming samplers for reproducible runs. See + [`docs/mdbook/src/budgets.md`](docs/mdbook/src/budgets.md). + - **Python bindings for the parallel simplifiers**: `simplify_redex`, `simplify_auto` and `simplify_strategy` join the existing `simplify_par`. All take a single expression and return the same result as `simplify`; only diff --git a/alkahest-core/src/budget/mod.rs b/alkahest-core/src/budget/mod.rs new file mode 100644 index 00000000..df502601 --- /dev/null +++ b/alkahest-core/src/budget/mod.rs @@ -0,0 +1,473 @@ +//! Per-call wall-clock and step budgets, cooperative cancellation, and a +//! deterministic seed for search-style workloads. +//! +//! # Motivation +//! +//! A fan-out loop trying 10k candidate rewrites/integrals/Gröbner bases +//! cannot afford one pathological candidate to hang the whole batch, and an +//! orchestrator that decides "this candidate isn't worth it" needs a way to +//! *stop it now* rather than waiting for an OS-level kill (`SIGKILL`, process +//! timeout). This module gives heavy engines a cheap, structured way to bail +//! out honestly — returning [`BudgetError`] — instead of running unbounded or +//! being killed with no diagnostic. +//! +//! # Model +//! +//! A [`Budget`] is *entered* with [`enter`], which pushes it onto a +//! thread-local stack and returns a [`BudgetGuard`]. The budget stays active +//! until the guard is dropped (including on panic-unwind), mirroring the +//! `try`/`finally` discipline of Python's `with alkahest.context(budget=...)`. +//! Budgets nest like `context(...)` blocks: only the *innermost* active frame +//! is consulted by [`check`] and [`seed`] — entering a new budget shadows the +//! outer one for the scope of the block rather than combining limits with it. +//! +//! Heavy call sites ([`crate::integrate::engine`]'s top-level entry and +//! recursion boundary, [`crate::simplify::engine`]'s per-pass loop) call +//! [`check`] at a handful of strategic points. `check` is cheap when no +//! budget is active and cancellation has not been requested — it is a single +//! atomic load plus (if a budget is active) an `Instant::now()` — so it is +//! safe to call unconditionally at those boundaries. +//! +//! Cancellation ([`request_cancel`] / [`is_cancelled`] / [`clear_cancel`]) is +//! a single process-wide flag, deliberately *not* scoped to a thread or a +//! [`Budget`] frame: it models "the orchestrator wants the current heavy +//! operation to stop right now", e.g. because a fan-out loop decided a +//! candidate has used enough wall time across every thread working on it. +//! Call [`clear_cancel`] before starting the next candidate. +//! +//! [`seed`] exposes the active budget's seed so RNG-consuming samplers (e.g. +//! randomized modular tests, homotopy continuation start systems) can be +//! seeded deterministically from the ambient budget instead of threading an +//! explicit seed parameter through every call — two runs entering the same +//! `Budget { seed: Some(7), .. }` observe the same [`seed`] at every call site +//! that consults it. +//! +//! # Errors +//! +//! [`BudgetError`] implements [`AlkahestError`] with stable codes: +//! +//! | Code | Variant | Cause | +//! |----------------|-------------------------|-------------------| +//! | `E-BUDGET-001` | [`BudgetError::WallClock`] | wall-clock deadline elapsed | +//! | `E-BUDGET-002` | [`BudgetError::Steps`] | step counter exceeded `max_steps` | +//! | `E-BUDGET-003` | [`BudgetError::Cancelled`] | [`request_cancel`] was called | + +use crate::errors::AlkahestError; +use std::cell::{Cell, RefCell}; +use std::fmt; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Duration, Instant}; + +// --------------------------------------------------------------------------- +// Budget +// --------------------------------------------------------------------------- + +/// A per-call resource budget: an optional wall-clock limit, an optional +/// cooperative step limit, and an optional determinism seed. +/// +/// Every field is optional. `Budget::default()` never trips [`check`] on its +/// own — only [`request_cancel`] can stop a call entered with a default +/// budget. This is intentional: entering an (otherwise empty) budget is how +/// a caller opts a code path into consulting [`seed`] without also imposing +/// a wall/step limit. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct Budget { + /// Wall-clock limit for the guarded block, measured from [`enter`]. + pub wall: Option, + /// Maximum number of [`check`] calls the guarded block may make. + pub max_steps: Option, + /// Determinism seed available to callers via [`seed`]. + pub seed: Option, +} + +impl Budget { + /// An empty budget — no wall/step limit, no seed. Equivalent to + /// `Budget::default()`. + pub fn new() -> Self { + Self::default() + } + + /// Set the wall-clock limit. + pub fn with_wall(mut self, wall: Duration) -> Self { + self.wall = Some(wall); + self + } + + /// Set the step limit. + pub fn with_max_steps(mut self, max_steps: u64) -> Self { + self.max_steps = Some(max_steps); + self + } + + /// Set the determinism seed. + pub fn with_seed(mut self, seed: u64) -> Self { + self.seed = Some(seed); + self + } +} + +struct Frame { + start: Instant, + wall: Option, + max_steps: Option, + steps: Cell, + seed: Option, +} + +thread_local! { + static STACK: RefCell> = const { RefCell::new(Vec::new()) }; +} + +/// RAII guard returned by [`enter`]. +/// +/// Pops the corresponding [`Budget`] frame from the current thread's active- +/// budget stack on drop — on every exit path, including panic-unwind — so a +/// `?`-propagated error or an early `return` can never leak a stale budget +/// frame into unrelated code that runs later on the same thread. +/// +/// Not [`Send`]: the stack it pops from is thread-local, so a guard created +/// on one thread must not be dropped from another. +pub struct BudgetGuard { + _not_send: std::marker::PhantomData<*const ()>, +} + +impl Drop for BudgetGuard { + fn drop(&mut self) { + STACK.with(|s| { + s.borrow_mut().pop(); + }); + } +} + +/// Push `budget` onto the current thread's active-budget stack. +/// +/// Returns a [`BudgetGuard`]; the budget stays active — visible to [`check`], +/// [`seed`], and [`is_active`] on this thread — until the guard is dropped. +/// Budgets nest like `with` blocks: entering a new one shadows the previous +/// one until it is popped, rather than combining limits with the outer +/// frame (matching the non-merging nesting semantics of +/// `alkahest.context(...)` on the Python side). +pub fn enter(budget: Budget) -> BudgetGuard { + let frame = Frame { + start: Instant::now(), + wall: budget.wall, + max_steps: budget.max_steps, + steps: Cell::new(0), + seed: budget.seed, + }; + STACK.with(|s| s.borrow_mut().push(frame)); + BudgetGuard { + _not_send: std::marker::PhantomData, + } +} + +/// Returns `true` if a [`Budget`] is currently active on this thread. +pub fn is_active() -> bool { + STACK.with(|s| !s.borrow().is_empty()) +} + +/// The seed of the innermost active [`Budget`] on this thread, or `None` if +/// no budget is active or the active budget did not set one. +pub fn seed() -> Option { + STACK.with(|s| s.borrow().last().and_then(|f| f.seed)) +} + +// --------------------------------------------------------------------------- +// Cancellation — process-wide, not scoped to a thread or Budget frame +// --------------------------------------------------------------------------- + +static CANCELLED: AtomicBool = AtomicBool::new(false); + +/// Request cancellation of the current cooperative operation(s). +/// +/// Checked by every [`check`] call on every thread — regardless of whether a +/// [`Budget`] is active — until [`clear_cancel`] is called. Intended for an +/// orchestrator thread to stop a heavy call it decided is no longer worth +/// running, without waiting for an OS-level kill. +pub fn request_cancel() { + CANCELLED.store(true, Ordering::SeqCst); +} + +/// Clear a previously requested cancellation. +/// +/// Call this before starting the next candidate in a fan-out loop — a +/// cancellation request left set would otherwise trip [`check`] +/// (`E-BUDGET-003`) immediately for every subsequent candidate. +pub fn clear_cancel() { + CANCELLED.store(false, Ordering::SeqCst); +} + +/// Returns `true` if [`request_cancel`] has been called and not yet cleared +/// by [`clear_cancel`]. +pub fn is_cancelled() -> bool { + CANCELLED.load(Ordering::SeqCst) +} + +// --------------------------------------------------------------------------- +// Cooperative check +// --------------------------------------------------------------------------- + +/// Cooperative checkpoint: call this at a natural short-circuit point in a +/// heavy algorithm — top-level entry, a recursion/depth-guard boundary, once +/// per major rewrite pass. +/// +/// Checks, in order: +/// 1. [`is_cancelled`] — the process-wide cancellation flag. +/// 2. The innermost active [`Budget`]'s wall-clock limit, if any. +/// 3. The innermost active [`Budget`]'s step counter, if any — incremented on +/// every call, compared against `max_steps` after incrementing. +/// +/// Returns `Ok(())` with no side effect (no counter increment) if no +/// [`Budget`] is active and cancellation has not been requested, so `check` +/// is cheap to call unconditionally at hot-loop boundaries even when no +/// caller has opted into a budget. +pub fn check() -> Result<(), BudgetError> { + if is_cancelled() { + return Err(BudgetError::Cancelled); + } + STACK.with(|s| { + let stack = s.borrow(); + let Some(frame) = stack.last() else { + return Ok(()); + }; + if let Some(wall) = frame.wall { + let elapsed = frame.start.elapsed(); + if elapsed >= wall { + return Err(BudgetError::WallClock { + limit: wall, + elapsed, + }); + } + } + if let Some(max_steps) = frame.max_steps { + let taken = frame.steps.get() + 1; + frame.steps.set(taken); + if taken > max_steps { + return Err(BudgetError::Steps { + limit: max_steps, + taken, + }); + } + } + Ok(()) + }) +} + +// --------------------------------------------------------------------------- +// Error type +// --------------------------------------------------------------------------- + +/// A [`Budget`] was exceeded, or cancellation was requested. +/// +/// This is a fine, expected answer for a fan-out search loop — not a crash — +/// so it carries a stable code and remediation like every other +/// `alkahest-core` error, rather than surfacing as a panic or an OS kill. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BudgetError { + /// The active budget's wall-clock limit elapsed. + WallClock { limit: Duration, elapsed: Duration }, + /// The active budget's step counter exceeded `max_steps`. + Steps { limit: u64, taken: u64 }, + /// [`request_cancel`] was called and not yet cleared. + Cancelled, +} + +impl fmt::Display for BudgetError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + BudgetError::WallClock { limit, elapsed } => write!( + f, + "budget exceeded: wall-clock limit {limit:?} elapsed ({elapsed:?} elapsed)" + ), + BudgetError::Steps { limit, taken } => write!( + f, + "budget exceeded: step limit {limit} reached ({taken} steps taken)" + ), + BudgetError::Cancelled => write!(f, "budget: operation was cancelled"), + } + } +} + +impl std::error::Error for BudgetError {} + +impl AlkahestError for BudgetError { + fn code(&self) -> &'static str { + match self { + BudgetError::WallClock { .. } => "E-BUDGET-001", + BudgetError::Steps { .. } => "E-BUDGET-002", + BudgetError::Cancelled => "E-BUDGET-003", + } + } + + fn remediation(&self) -> Option<&'static str> { + match self { + BudgetError::WallClock { .. } => Some( + "raise Budget(wall_ms=...), or accept a heuristic/numeric result for this \ + candidate instead of an exact one", + ), + BudgetError::Steps { .. } => Some( + "raise Budget(max_steps=...), or accept a partial/heuristic result for this \ + candidate instead of an exact one", + ), + BudgetError::Cancelled => Some( + "call alkahest.clear_cancel() (Python) or budget::clear_cancel() (Rust) before \ + starting the next candidate", + ), + } + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{Mutex, MutexGuard}; + + /// `CANCELLED` is a process-wide `AtomicBool` by design (see the module + /// docs) so an orchestrator thread can cancel a heavy call running on a + /// worker thread. That means tests which flip it must not run + /// concurrently with *any* other test that calls [`check`] — including + /// tests in this module that never touch cancellation themselves — or + /// they'll intermittently observe a stale `true` from a racing test and + /// fail with `Cancelled` instead of the error under test. Every test + /// below acquires this lock for its whole body to serialize with the + /// cancel-flipping tests; `cargo test` still runs them in parallel with + /// unrelated tests elsewhere in the crate, but nothing else in the crate + /// calls `request_cancel`, so those stay unaffected. + static TEST_SERIAL: Mutex<()> = Mutex::new(()); + + fn serial() -> MutexGuard<'static, ()> { + TEST_SERIAL.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// Clears cancellation on drop so a panicking assertion mid-test doesn't + /// leave `CANCELLED` set for the next test to acquire the lock. + struct CancelGuard; + impl Drop for CancelGuard { + fn drop(&mut self) { + clear_cancel(); + } + } + + #[test] + fn no_budget_active_never_trips() { + let _serial = serial(); + assert!(!is_active()); + assert_eq!(seed(), None); + for _ in 0..1000 { + assert!(check().is_ok()); + } + } + + #[test] + fn step_budget_trips_after_limit() { + let _serial = serial(); + let _guard = enter(Budget::new().with_max_steps(3)); + assert!(check().is_ok()); + assert!(check().is_ok()); + assert!(check().is_ok()); + let err = check().unwrap_err(); + assert_eq!(err.code(), "E-BUDGET-002"); + assert_eq!(err, BudgetError::Steps { limit: 3, taken: 4 }); + } + + #[test] + fn wall_budget_trips_after_elapsed() { + let _serial = serial(); + let _guard = enter(Budget::new().with_wall(Duration::from_millis(10))); + assert!(check().is_ok()); + std::thread::sleep(Duration::from_millis(25)); + let err = check().unwrap_err(); + assert_eq!(err.code(), "E-BUDGET-001"); + assert!(matches!(err, BudgetError::WallClock { .. })); + } + + #[test] + fn seed_round_trips_through_active_budget() { + let _serial = serial(); + assert_eq!(seed(), None); + { + let _guard = enter(Budget::new().with_seed(7)); + assert_eq!(seed(), Some(7)); + } + // Popped: seed is no longer visible. + assert_eq!(seed(), None); + } + + #[test] + fn nested_budgets_shadow_not_merge() { + let _serial = serial(); + let _outer = enter(Budget::new().with_seed(1).with_max_steps(1000)); + assert_eq!(seed(), Some(1)); + { + // Inner budget has no seed set — it does NOT inherit the outer + // seed, matching alkahest.context(...)'s non-merging semantics. + let _inner = enter(Budget::new().with_max_steps(2)); + assert_eq!(seed(), None); + assert!(check().is_ok()); + assert!(check().is_ok()); + assert_eq!(check().unwrap_err().code(), "E-BUDGET-002"); + } + // Back to the outer frame: its own step counter is untouched by the + // inner frame's checks, and its seed is visible again. + assert_eq!(seed(), Some(1)); + assert!(check().is_ok()); + } + + #[test] + fn guard_pops_on_early_return_via_question_mark() { + let _serial = serial(); + fn inner() -> Result<(), BudgetError> { + let _guard = enter(Budget::new().with_max_steps(1)); + check()?; + check()?; // trips — guard must still pop on this early return. + unreachable!(); + } + assert!(inner().is_err()); + assert!(!is_active()); + } + + #[test] + fn cancel_flag_trips_check_and_clears() { + let _serial = serial(); + let _cancel_guard = CancelGuard; + assert!(!is_cancelled()); + request_cancel(); + assert!(is_cancelled()); + let err = check().unwrap_err(); + assert_eq!(err.code(), "E-BUDGET-003"); + assert_eq!(err, BudgetError::Cancelled); + clear_cancel(); + assert!(!is_cancelled()); + assert!(check().is_ok()); + } + + #[test] + fn cancel_trips_even_with_a_generous_budget_active() { + let _serial = serial(); + let _cancel_guard = CancelGuard; + let _guard = enter(Budget::new().with_max_steps(1_000_000)); + request_cancel(); + assert_eq!(check().unwrap_err(), BudgetError::Cancelled); + } + + #[test] + fn error_codes_have_remediation() { + for err in [ + BudgetError::WallClock { + limit: Duration::from_secs(1), + elapsed: Duration::from_secs(2), + }, + BudgetError::Steps { limit: 1, taken: 2 }, + BudgetError::Cancelled, + ] { + assert!(err.code().starts_with("E-BUDGET-")); + assert!(err.remediation().is_some()); + // Display must not panic and should mention something useful. + assert!(!err.to_string().is_empty()); + } + } +} diff --git a/alkahest-core/src/errors/codes.rs b/alkahest-core/src/errors/codes.rs index 1a5b99d3..78affe2d 100644 --- a/alkahest-core/src/errors/codes.rs +++ b/alkahest-core/src/errors/codes.rs @@ -177,6 +177,10 @@ pub const REGISTRY: &[ErrorSpec] = &[ ErrorSpec { code: "E-RESIDUE-002", class: "ResidueError", cause: Cause::Domain, remediation: Some("denominator must be non-zero") }, ErrorSpec { code: "E-RESIDUE-003", class: "ResidueError", cause: Cause::Unsupported, remediation: Some("pole order exceeds supported bound; essential singularities are out of scope") }, ErrorSpec { code: "E-RESIDUE-004", class: "ResidueError", cause: Cause::Domain, remediation: Some("division by zero during Laurent coefficient extraction") }, + // E-BUDGET — BudgetError (P1 search plumbing item 4: budgets/cancellation/determinism) + ErrorSpec { code: "E-BUDGET-001", class: "BudgetError", cause: Cause::Resource, remediation: Some("raise Budget(wall_ms=...), or accept a heuristic/numeric result for this candidate instead of an exact one") }, + ErrorSpec { code: "E-BUDGET-002", class: "BudgetError", cause: Cause::Resource, remediation: Some("raise Budget(max_steps=...), or accept a partial/heuristic result for this candidate instead of an exact one") }, + ErrorSpec { code: "E-BUDGET-003", class: "BudgetError", cause: Cause::Resource, remediation: Some("call alkahest.clear_cancel() (Python) or budget::clear_cancel() (Rust) before starting the next candidate") }, // E-DOMAIN — reserved; DomainError is Python-only pending Rust implementation ]; diff --git a/alkahest-core/src/integrate/engine.rs b/alkahest-core/src/integrate/engine.rs index 888f94e0..6ba52996 100644 --- a/alkahest-core/src/integrate/engine.rs +++ b/alkahest-core/src/integrate/engine.rs @@ -33,6 +33,10 @@ pub enum IntegrationError { UnsupportedExtensionDegree(u32), /// The integrand provably has no elementary antiderivative (e.g. elliptic integrals). NonElementary(String), + /// The active [`crate::budget::Budget`] was exceeded, or cancellation was + /// requested, at a cooperative checkpoint inside the integration engine. + /// See `crate::budget` — P1 search plumbing item 4. + Budget(crate::budget::BudgetError), } impl fmt::Display for IntegrationError { @@ -48,12 +52,19 @@ impl fmt::Display for IntegrationError { IntegrationError::NonElementary(msg) => { write!(f, "integrate: no elementary antiderivative exists: {msg}") } + IntegrationError::Budget(e) => write!(f, "integrate: {e}"), } } } impl std::error::Error for IntegrationError {} +impl From for IntegrationError { + fn from(e: crate::budget::BudgetError) -> Self { + IntegrationError::Budget(e) + } +} + impl IntegrationError { /// A human-readable remediation hint for the user. pub fn remediation(&self) -> Option<&'static str> { @@ -71,6 +82,10 @@ impl IntegrationError { "this integrand has no closed-form antiderivative in terms of elementary \ functions; use a numeric integrator or elliptic-integral library", ), + IntegrationError::Budget(e) => { + use crate::errors::AlkahestError; + e.remediation() + } } } @@ -87,6 +102,7 @@ impl crate::errors::AlkahestError for IntegrationError { IntegrationError::DivisionByZero => "E-INT-002", IntegrationError::UnsupportedExtensionDegree(_) => "E-INT-003", IntegrationError::NonElementary(_) => "E-INT-004", + IntegrationError::Budget(e) => e.code(), } } @@ -2066,6 +2082,13 @@ pub fn integrate( var: ExprId, pool: &ExprPool, ) -> Result, IntegrationError> { + // Cooperative budget checkpoint (P1 search plumbing item 4): the single + // entry point every public integration route passes through, so a fan-out + // loop over many candidates can bound wall-clock/step cost or request + // cancellation without waiting for an OS-level kill. No-op unless the + // caller entered a `budget::Budget` — see `crate::budget`. + crate::budget::check()?; + // V1-2: Route algebraic integrands to the Trager/Risch algebraic engine. // For *mixed* algebraic+transcendental (e.g. exp(x)/sqrt(x²+1)) the Risch // engine handles the transcendental level and delegates base-field integrals @@ -2133,6 +2156,12 @@ fn integrate_inner( pool: &ExprPool, depth: u32, ) -> Result, IntegrationError> { + // Cooperative budget checkpoint at the recursion boundary: u-substitution + // re-enters `integrate_inner` (see `try_u_substitution` below), so this + // one call site also bounds the recursive fallback chain, not just the + // initial call. See `crate::budget` — P1 search plumbing item 4. + crate::budget::check()?; + let mut log = DerivationLog::new(); match integrate_raw(expr, var, pool, &mut log) { Ok(raw) => { diff --git a/alkahest-core/src/integrate/risch/mod.rs b/alkahest-core/src/integrate/risch/mod.rs index ed3f45b1..1198ba37 100644 --- a/alkahest-core/src/integrate/risch/mod.rs +++ b/alkahest-core/src/integrate/risch/mod.rs @@ -280,6 +280,13 @@ pub fn integrate_risch( Err(IntegrationError::UnsupportedExtensionDegree(d)) => { return Err(IntegrationError::UnsupportedExtensionDegree(d)); } + Err(e @ IntegrationError::Budget(_)) => { + // A budget/cancellation trip is never a "this route declined" + // signal — propagate it immediately instead of falling + // through to sum decomposition, which would keep spending the + // budget the caller just told us is exhausted. + return Err(e); + } Err(IntegrationError::NotImplemented(_)) => { // Fall through to sum decomposition below. } diff --git a/alkahest-core/src/lib.rs b/alkahest-core/src/lib.rs index 00bf73d8..e311f51f 100644 --- a/alkahest-core/src/lib.rs +++ b/alkahest-core/src/lib.rs @@ -7,6 +7,8 @@ pub mod acausal; pub mod algebra; pub mod ball; +// P1 search plumbing item 4 — budgets, cancellation, determinism +pub mod budget; pub mod calculus; pub mod dae; pub mod deriv; @@ -184,6 +186,11 @@ pub use poly::groebner::{ MonomialOrder, }; +// P1 search plumbing item 4 — budgets, cancellation, determinism +pub use budget::{ + check as budget_check, clear_cancel, enter as budget_enter, is_active as budget_is_active, + is_cancelled, request_cancel, seed as budget_seed, Budget, BudgetError, BudgetGuard, +}; pub use errors::AlkahestError; pub use lean::{ emit_definite_integration_cert, emit_integration_cert, emit_lean_expr as emit_lean, diff --git a/alkahest-core/src/simplify/engine.rs b/alkahest-core/src/simplify/engine.rs index 0f9e5836..22705c37 100644 --- a/alkahest-core/src/simplify/engine.rs +++ b/alkahest-core/src/simplify/engine.rs @@ -367,6 +367,16 @@ pub fn simplify_with( ) -> DerivedExpr { let mut current = DerivedExpr::new(expr); for _ in 0..config.max_iterations { + // Cooperative budget checkpoint, once per full bottom-up pass (P1 + // search plumbing item 4). `simplify` has no `Result` return type, so + // a budget/cancellation trip here stops further passes early and + // returns the best value simplified so far — exactly like running out + // of `max_iterations` already does silently. Callers that need a hard + // `BudgetExceeded` raise on this path should wrap the call in a + // Python-level wall timeout (see `docs/mdbook/src/budgets.md`). + if crate::budget::check().is_err() { + break; + } // Fresh memo per pass: maps input ExprId → simplified ExprId. // Shared subexpressions are simplified once and the result reused for // all subsequent occurrences within the same bottom-up sweep. @@ -402,6 +412,10 @@ pub fn simplify_with_pattern_rules( let child_rules = rule_set.as_dyn_rules(); let mut current = DerivedExpr::new(expr); for _ in 0..config.max_iterations { + // See the matching checkpoint in `simplify_with` above. + if crate::budget::check().is_err() { + break; + } let mut memo: HashMap = HashMap::new(); let result = simplify_node_indexed(current.value, pool, rule_set, &child_rules, &mut memo); let merged_log = current.log.merge(result.log); @@ -463,6 +477,11 @@ pub fn simplify_batch(exprs: &[ExprId], pool: &ExprPool) -> Vec = HashMap::new(); diff --git a/alkahest-py/src/lib.rs b/alkahest-py/src/lib.rs index dc377355..b6f263fc 100644 --- a/alkahest-py/src/lib.rs +++ b/alkahest-py/src/lib.rs @@ -164,7 +164,7 @@ use alkahest_core::number_theory::{ QuadraticDirichlet as CoreQuadraticDirichlet, }; use pyo3::buffer::PyBuffer; -use pyo3::exceptions::{PyOverflowError, PyTypeError}; +use pyo3::exceptions::{PyOverflowError, PyTypeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyComplex, PyDict, PyInt, PyList, PyTuple}; use rug::{Complete, Integer, Rational}; @@ -286,6 +286,8 @@ pyo3::create_exception!(alkahest, PyLinearRecurrenceError, PyAlkahestError); pyo3::create_exception!(alkahest, PyRsolveError, PyAlkahestError); #[cfg(feature = "groebner")] pyo3::create_exception!(alkahest, PyDiophantineError, PyAlkahestError); +// P1 search plumbing item 4 — budgets, cancellation, determinism +pyo3::create_exception!(alkahest, PyBudgetExceededError, PyAlkahestError); /// Build a structured exception with `.code`, `.remediation`, `.span` attributes. fn make_structured_err( @@ -445,12 +447,118 @@ fn gpu_groebner_error_to_py(e: alkahest_core::experimental::GpuGroebnerError) -> } fn integrate_error_to_py(e: IntegrationError) -> PyErr { + // A budget/cancellation trip is not an integration failure — raise the + // dedicated `BudgetExceededError` (E-BUDGET-*) instead of `IntegrationError` + // so callers can catch it uniformly regardless of which engine tripped it. + if let IntegrationError::Budget(inner) = &e { + return budget_error_to_py(inner); + } Python::with_gil(|py| { let exc_type = py.get_type_bound::(); make_structured_err(py, &exc_type, &e) }) } +/// Map a [`alkahest_core::budget::BudgetError`] to Python's +/// `BudgetExceededError` (`E-BUDGET-001..003`). See `crate::budget` — P1 +/// search plumbing item 4. +fn budget_error_to_py(e: &alkahest_core::budget::BudgetError) -> PyErr { + Python::with_gil(|py| { + let exc_type = py.get_type_bound::(); + make_structured_err(py, &exc_type, e) + }) +} + +// --------------------------------------------------------------------------- +// P1 search plumbing item 4 — budgets, cancellation, determinism +// +// `alkahest_core::budget::enter` returns an RAII guard that must be dropped +// on the same thread it was created on (the underlying stack is +// thread-local). `alkahest.context(budget=...)` is a `@contextmanager`, so +// its `push`/`pop` calls always run on the thread that entered the `with` +// block — there is no `push`/`pop` pair here without a matching thread, and +// no guard object crosses the Python/Rust boundary. Each `push_budget` call +// stores its guard on a Rust-side thread-local stack; `pop_budget` pops and +// drops the most recent one. +// --------------------------------------------------------------------------- + +thread_local! { + static PY_BUDGET_GUARDS: std::cell::RefCell> = + const { std::cell::RefCell::new(Vec::new()) }; +} + +/// Push a [`alkahest_core::budget::Budget`] onto this thread's active-budget +/// stack. Pair with [`py_pop_budget`]; `alkahest.context(budget=...)` calls +/// both around its `with` block. +#[pyfunction] +#[pyo3(name = "push_budget")] +#[pyo3(signature = (wall_ms=None, max_steps=None, seed=None))] +fn py_push_budget(wall_ms: Option, max_steps: Option, seed: Option) -> PyResult<()> { + let mut budget = alkahest_core::budget::Budget::new(); + if let Some(ms) = wall_ms { + if !ms.is_finite() || ms < 0.0 { + return Err(PyValueError::new_err( + "wall_ms must be a finite, non-negative number of milliseconds", + )); + } + budget.wall = Some(std::time::Duration::from_secs_f64(ms / 1000.0)); + } + budget.max_steps = max_steps; + budget.seed = seed; + let guard = alkahest_core::budget::enter(budget); + PY_BUDGET_GUARDS.with(|g| g.borrow_mut().push(guard)); + Ok(()) +} + +/// Pop the most recently pushed [`alkahest_core::budget::Budget`] from this +/// thread's active-budget stack, dropping its guard. +#[pyfunction] +#[pyo3(name = "pop_budget")] +fn py_pop_budget() -> PyResult<()> { + let popped = PY_BUDGET_GUARDS.with(|g| g.borrow_mut().pop()); + if popped.is_none() { + return Err(pyo3::exceptions::PyRuntimeError::new_err( + "pop_budget() called with no active budget on this thread", + )); + } + Ok(()) +} + +/// `True` if a [`alkahest_core::budget::Budget`] is active on this thread. +#[pyfunction] +#[pyo3(name = "is_budget_active")] +fn py_is_budget_active() -> bool { + alkahest_core::budget::is_active() +} + +/// The seed of the innermost active budget on this thread, or `None`. +#[pyfunction] +#[pyo3(name = "budget_seed")] +fn py_budget_seed() -> Option { + alkahest_core::budget::seed() +} + +/// Request cancellation of the current cooperative operation(s), process-wide. +#[pyfunction] +#[pyo3(name = "request_cancel")] +fn py_request_cancel() { + alkahest_core::budget::request_cancel(); +} + +/// Clear a previously requested cancellation. +#[pyfunction] +#[pyo3(name = "clear_cancel")] +fn py_clear_cancel() { + alkahest_core::budget::clear_cancel(); +} + +/// `True` if [`py_request_cancel`] was called and not yet cleared. +#[pyfunction] +#[pyo3(name = "is_cancelled")] +fn py_is_cancelled() -> bool { + alkahest_core::budget::is_cancelled() +} + fn series_error_to_py(e: SeriesError) -> PyErr { Python::with_gil(|py| { let exc_type = py.get_type_bound::(); @@ -9823,6 +9931,18 @@ fn alkahest(m: &Bound<'_, PyModule>) -> PyResult<()> { "DiophantineError", m.py().get_type_bound::(), )?; + // P1 search plumbing item 4 — budgets, cancellation, determinism + m.add( + "BudgetExceededError", + m.py().get_type_bound::(), + )?; + m.add_function(wrap_pyfunction!(py_push_budget, m)?)?; + m.add_function(wrap_pyfunction!(py_pop_budget, m)?)?; + m.add_function(wrap_pyfunction!(py_is_budget_active, m)?)?; + m.add_function(wrap_pyfunction!(py_budget_seed, m)?)?; + m.add_function(wrap_pyfunction!(py_request_cancel, m)?)?; + m.add_function(wrap_pyfunction!(py_clear_cancel, m)?)?; + m.add_function(wrap_pyfunction!(py_is_cancelled, m)?)?; // V1-15: compile-time flag so Python tests can skip egraph-dependent assertions. m.add("HAS_EGRAPH", cfg!(feature = "egraph"))?; Ok(()) diff --git a/docs/mdbook/src/SUMMARY.md b/docs/mdbook/src/SUMMARY.md index 64a8d9ba..f1f1a79a 100644 --- a/docs/mdbook/src/SUMMARY.md +++ b/docs/mdbook/src/SUMMARY.md @@ -22,4 +22,5 @@ - [Lean certificates](./lean-certs.md) - [Certificate coverage](./certificate-coverage.md) - [Error handling](./errors.md) +- [Budgets, cancellation, and determinism](./budgets.md) - [Stability policy](./stability.md) diff --git a/docs/mdbook/src/budgets.md b/docs/mdbook/src/budgets.md new file mode 100644 index 00000000..99ce9cf5 --- /dev/null +++ b/docs/mdbook/src/budgets.md @@ -0,0 +1,166 @@ +# Budgets, cancellation, and determinism + +A fan-out loop trying thousands of candidate rewrites/integrals/Gröbner bases cannot +afford one pathological candidate to hang the whole batch — and an orchestrator that +decides a candidate isn't worth more time needs a way to *stop it now*, not wait for an +OS-level kill (`SIGKILL`, a process timeout). `alkahest.Budget` and +`alkahest.context(budget=...)` give heavy engines a cheap, structured way to bail out +honestly — raising `BudgetExceededError` — instead of running unbounded. + +```python +import alkahest as ak + +p = ak.ExprPool() +x = p.symbol("x", "real") + +with ak.context(pool=p, budget=ak.Budget(wall_ms=50, max_steps=10_000, seed=7)): + try: + ak.integrate(hard_expr, x) + except ak.BudgetExceededError as e: + assert e.code.startswith("E-BUDGET-") + # ... deprioritize this candidate and move on to the next one ... +``` + +## Model + +A `Budget` is an immutable `(wall_ms, max_steps, seed)` triple. Every field is optional; +`Budget()` never trips a check on its own — only `alkahest.request_cancel()` can stop a +call entered with a bare `Budget()`. + +`context(budget=...)` pushes the budget into a **thread-local** stack on the Rust side +(`alkahest_core::budget`) for the scope of the `with` block, and pops it on exit — +including on an exception, matching every other resource the context manager owns. +Budgets nest like every other `context(...)` key: only the *innermost* frame is +consulted, so a nested `context(budget=...)` **shadows** the outer one rather than +combining limits with it. A nested `context(...)` that omits `budget=` leaves the outer +one active (nothing is pushed, so nothing shadows it): + +```python +with ak.context(pool=p, budget=ak.Budget(seed=1, max_steps=1000)): + ak.budget_seed() # 1 + with ak.context(pool=p): + ak.budget_seed() # 1 — no budget= here, outer frame still active + with ak.context(pool=p, budget=ak.Budget(max_steps=2)): + ak.budget_seed() # None — this frame set no seed; it does not inherit + ak.budget_seed() # 1 — back to the outer frame +``` + +## What checks the budget today + +The Rust engines call a cheap cooperative checkpoint (`alkahest_core::budget::check`) at +a handful of strategic points — not blanket-inserted into every loop: + +- **`alkahest.integrate`** — at the top-level entry (covers every route: algebraic, + Risch/transcendental, rational-function, log-derivative) and at the + `integrate_inner` recursion boundary that u-substitution and the rational-function + fallback re-enter. A trip here raises `BudgetExceededError` — integration has a + `Result` return type with a natural place to signal it. +- **`alkahest.simplify`** (and `simplify_with`, `simplify_batch`) — once per full + bottom-up rewrite pass. `simplify` has **no error channel** (`DerivedExpr` isn't a + `Result`), so a trip here stops further passes early and returns the best value + simplified so far — exactly like running out of the existing `max_iterations` cap + already does, silently. If you need a hard raise on a `simplify` call specifically, + wrap it in `alkahest.run_with_wall_fallback` (below). + +Other heavy primitives (Gröbner bases, homotopy continuation, …) do not yet check the +budget; wiring them is a follow-up, not part of this cut. Calling `check()` is cheap +when no budget is active and cancellation has not been requested (an atomic load, and — +only if a budget is active — an `Instant::now()`), so it is safe to sprinkle at more +call sites over time without a performance concern gating it. + +## Cancellation + +`alkahest.request_cancel()` sets a single **process-wide** flag — deliberately not +scoped to a thread or a `Budget` frame. It models "the orchestrator wants the current +heavy operation to stop right now", e.g. because a fan-out loop decided a candidate has +used enough wall time, and the operation might be running on a different thread than the +one that decided to give up on it. `alkahest.is_cancelled()` reads it; +`alkahest.clear_cancel()` resets it — call this before starting the next candidate, or +every subsequent call trips `E-BUDGET-003` immediately. + +```python +import threading + +def watchdog(): + time.sleep(0.05) + ak.request_cancel() + +threading.Thread(target=watchdog, daemon=True).start() +try: + ak.integrate(hard_expr, x) +except ak.BudgetExceededError as e: + assert e.code == "E-BUDGET-003" +finally: + ak.clear_cancel() +``` + +## Determinism seed + +`Budget(seed=...)` doesn't do anything by itself — it makes the seed available via +`alkahest.budget_seed()` (Rust: `alkahest_core::budget::seed()`) to any RNG-consuming +sampler that chooses to consult it, instead of threading an explicit seed parameter +through every call in a pipeline. Two runs entering `Budget(seed=7)` observe the same +`budget_seed()` at every call site that reads it, so a search loop that seeds its own +sampling from the ambient budget is reproducible run-to-run. + +## The Python-layer wall-clock fallback + +Because `simplify` cannot raise through its own return type, `context(budget=...)` +*alone* only bounds it the same way `max_iterations` already does — silently, by +returning early. If you need a hard deadline specifically on a call like that, +`alkahest.run_with_wall_fallback` is a **supplement**, not a replacement: it runs the +call on a worker thread and raises `BudgetExceededError` (`E-BUDGET-001`) if it doesn't +finish in time. + +```python +result = ak.run_with_wall_fallback(ak.simplify, big_expr, budget=ak.Budget(wall_ms=200)) +``` + +Python cannot forcibly kill a thread, so on a timeout the call keeps running in the +background until it either finishes or reaches a Rust cooperative checkpoint — +`run_with_wall_fallback` also calls `request_cancel()` on timeout so any checkpoint the +call reaches asks it to stop. Prefer the Rust cooperative check alone +(`context(budget=...)`) wherever a call already honors it (`integrate` today); reach for +this only when you need a hard deadline on a path that doesn't. + +## Error codes + +| Code | Cause | +|---|---| +| `E-BUDGET-001` | The active budget's wall-clock limit elapsed | +| `E-BUDGET-002` | The active budget's step counter exceeded `max_steps` | +| `E-BUDGET-003` | `request_cancel()` was called and not yet cleared | + +All three are `Cause::Resource` in the Rust registry (`alkahest_core::errors::codes`) — +a budget/cancellation trip is an environment/policy limit, not a statement about the +mathematics, so it is never conflated with e.g. `IntegrationError::NonElementary` (a +proof that no elementary antiderivative exists). `alkahest.integrate` raises +`BudgetExceededError` directly rather than wrapping it in `IntegrationError`, so callers +can catch it uniformly regardless of which engine tripped it: + +```python +try: + ak.integrate(hard_expr, x) +except ak.BudgetExceededError as e: + ... # deprioritize and move on +except ak.IntegrationError as e: + ... # a genuine "no elementary antiderivative" or "not implemented" verdict +``` + +## API reference + +| Name | Kind | Description | +|---|---|---| +| `Budget(wall_ms=None, max_steps=None, seed=None)` | class | Immutable budget triple | +| `context(budget=...)` | context manager | Push/pop the budget for a `with` block | +| `active_budget()` | function | The `Budget` from the innermost active context, or `None` | +| `budget_seed()` | function | The seed of the innermost active budget, or `None` | +| `is_budget_active()` | function | `True` if a budget is active on this thread | +| `request_cancel()` | function | Set the process-wide cancellation flag | +| `clear_cancel()` | function | Clear it | +| `is_cancelled()` | function | Read it | +| `run_with_wall_fallback(fn, *args, budget, **kwargs)` | function | Python-layer wall-clock fallback for calls without a Rust checkpoint | +| `BudgetExceededError` | exception | `E-BUDGET-001..003`, subclass of `AlkahestError` | + +On the Rust side (`alkahest_core::budget`): `Budget`, `enter`, `BudgetGuard`, `check`, +`seed`, `is_active`, `request_cancel`, `clear_cancel`, `is_cancelled`, `BudgetError`. diff --git a/docs/mdbook/src/errors.md b/docs/mdbook/src/errors.md index 536acbe3..124103c5 100644 --- a/docs/mdbook/src/errors.md +++ b/docs/mdbook/src/errors.md @@ -16,7 +16,8 @@ AlkahestError (base) ├── SolverError (E-SOLVE-*) — polynomial system solving ├── JitError (E-JIT-*) — LLVM/JIT codegen ├── CudaError (E-CUDA-*) — CUDA kernel launch or driver -└── PoolError (E-POOL-*) — ExprPool misuse +├── PoolError (E-POOL-*) — ExprPool misuse +└── BudgetExceededError (E-BUDGET-*) — budget/cancellation trip, see [Budgets](./budgets.md) ``` ## Error attributes @@ -120,6 +121,7 @@ Every error is classified on two independent axes: **subsystem** (determines the | `E-PARSE-*` | `ParseError` *(reserved)* | Parser integration — owns `span()` by default | | `E-IO-*` | `IoError` *(reserved)* | Checkpoint/serde paths (`PoolPersistError`) | | `E-CERT-*` | `CertificateUnavailableError` | A Lean certificate was required but withheld | +| `E-BUDGET-*` | `BudgetExceededError` | Budget/cancellation trip — see [Budgets, cancellation, and determinism](./budgets.md) | ### `E-CERT-*` — certificate policy diff --git a/python/alkahest/__init__.py b/python/alkahest/__init__.py index b55fc7cd..8d1b6e31 100644 --- a/python/alkahest/__init__.py +++ b/python/alkahest/__init__.py @@ -10,6 +10,15 @@ number_theory, research, # session-level claim graph (provenance objects) ) +from ._budget import ( + Budget, + budget_seed, + clear_cancel, + is_budget_active, + is_cancelled, + request_cancel, + run_with_wall_fallback, +) from ._certificates import ( Certifiability, certifiable, @@ -18,6 +27,7 @@ ) from ._context import ( active_assumptions, + active_budget, active_domain, active_pool, context, @@ -287,6 +297,7 @@ from .exceptions import ( AlkahestError, AssumptionError, + BudgetExceededError, CadError, CertificateUnavailableError, ConversionError, @@ -323,6 +334,7 @@ _NATIVE_EXCEPTION_OVERLAY: tuple[str, ...] = ( "AlkahestError", "AssumptionError", + "BudgetExceededError", "CadError", "ConversionError", "DaeError", @@ -1494,6 +1506,9 @@ def wrapper(*args, **kwargs): "ArbBall", "AssumptionError", "Assumptions", + # P1 search plumbing item 4 — budgets, cancellation, determinism + "Budget", + "BudgetExceededError", "CadError", # V5-12 — certificate ledger "Certifiability", @@ -1585,6 +1600,8 @@ def wrapper(*args, **kwargs): "acos", "acosh", "active_assumptions", + # P1 search plumbing item 4 + "active_budget", "active_domain", "active_pool", "adjoint_system", @@ -1596,6 +1613,8 @@ def wrapper(*args, **kwargs): "atanh", "bessel_j0", "bessel_j1", + # P1 search plumbing item 4 + "budget_seed", "cad_lift", "cad_project", "cancel", @@ -1605,6 +1624,8 @@ def wrapper(*args, **kwargs): # V5-12 — certificate ledger "certifiable", "certificate_coverage", + # P1 search plumbing item 4 + "clear_cancel", # Phase 26 "collect_like_terms", "compile_expr", @@ -1648,6 +1669,9 @@ def wrapper(*args, **kwargs): "im", "integrate", "interval_eval", + # P1 search plumbing item 4 + "is_budget_active", + "is_cancelled", "jacobian", "jit", "jit_is_available", @@ -1689,6 +1713,8 @@ def wrapper(*args, **kwargs): # V2-4 "real_roots", "refine_root", + # P1 search plumbing item 4 + "request_cancel", # V5-12 — certificate ledger "require_certificate", # Session-level provenance: claim graph for autoresearch loops @@ -1702,6 +1728,8 @@ def wrapper(*args, **kwargs): "round", "routh_hurwitz", "rsolve", + # P1 search plumbing item 4 + "run_with_wall_fallback", "satisfiable", "sensitivity_system", "series", diff --git a/python/alkahest/_budget.py b/python/alkahest/_budget.py new file mode 100644 index 00000000..cab6356a --- /dev/null +++ b/python/alkahest/_budget.py @@ -0,0 +1,239 @@ +"""Budget: per-call wall-clock / step limits, cancellation, and a determinism seed. + +P1 search plumbing item 4 — see ``docs/mdbook/src/budgets.md``. + +A fan-out loop trying thousands of candidate rewrites/integrals cannot afford +one pathological candidate to hang the whole batch, and needs a way to stop a +candidate that's no longer worth the wall time — without waiting for an +OS-level kill. This module is the Python front door for that: + +:class:`Budget` + An immutable ``(wall_ms, max_steps, seed)`` triple. + +``alkahest.context(budget=...)`` + Pushes the budget into the Rust-side cooperative checkpoint + (``alkahest_core::budget``) for the scope of the ``with`` block. Heavy + engines — currently :func:`alkahest.integrate` and, best-effort, + :func:`alkahest.simplify` — consult it at a handful of strategic points + and raise :class:`~alkahest.BudgetExceededError` (or, for ``simplify``, + stop early without raising — see the note on that function below) when it + trips. + +:func:`run_with_wall_fallback` + A Python-layer *supplement*, not a replacement, for calls that don't + (yet) check the Rust cooperative budget on every path — most notably + :func:`alkahest.simplify`, whose ``DerivedExpr`` return type has no error + channel to raise through, so it only stops early silently. Runs the call + on a worker thread and raises :class:`~alkahest.BudgetExceededError` if it + doesn't finish within ``budget.wall_ms``. The worker thread is **not** + killed — Python has no safe way to do that — so on a timeout the call may + keep running in the background until it hits a cooperative checkpoint or + finishes. Prefer relying on the Rust cooperative check (via + ``context(budget=...)`` alone) wherever it's already wired; reach for this + only when you need a hard deadline on a path it doesn't cover. + +:func:`request_cancel` / :func:`clear_cancel` / :func:`is_cancelled` + Thin wrappers over the process-wide cancellation flag + (``alkahest_core::budget``): an orchestrator thread can request that a + heavy call running on another thread stop *now*. +""" + +from __future__ import annotations + +import concurrent.futures +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, TypeVar + +if TYPE_CHECKING: + from .exceptions import BudgetExceededError + +__all__ = [ + "Budget", + "budget_seed", + "clear_cancel", + "is_budget_active", + "is_cancelled", + "request_cancel", + "run_with_wall_fallback", +] + +_T = TypeVar("_T") + + +@dataclass(frozen=True, slots=True) +class Budget: + """A per-call resource budget for search-style workloads. + + Every field is optional; ``Budget()`` never trips a cooperative check on + its own — only :func:`request_cancel` can stop a call entered with a bare + ``Budget()``. This mirrors the Rust side: entering an otherwise-empty + budget is how a caller opts a code path into consulting ``seed`` without + also imposing a wall/step limit. + + Parameters + ---------- + wall_ms : float, optional + Wall-clock limit in milliseconds, measured from + ``context(budget=...)`` entry. + max_steps : int, optional + Maximum number of cooperative-checkpoint calls the guarded block may + make (see ``crate::budget::check`` on the Rust side). + seed : int, optional + Determinism seed available to RNG-consuming samplers via + :func:`budget_seed` — two runs entering the same ``Budget(seed=7)`` + observe the same seed at every call site that consults it. + + Examples + -------- + >>> import alkahest as ak + >>> with ak.context(budget=ak.Budget(wall_ms=50, max_steps=10_000, seed=7)): + ... try: + ... ak.integrate(hard_expr, x) # doctest: +SKIP + ... except ak.BudgetExceededError as e: + ... assert e.code.startswith("E-BUDGET-") + """ + + wall_ms: float | None = None + max_steps: int | None = None + seed: int | None = None + + def __post_init__(self) -> None: + if self.wall_ms is not None and (not math.isfinite(self.wall_ms) or self.wall_ms < 0): + raise ValueError( + "Budget.wall_ms must be a finite, non-negative number of milliseconds" + ) + if self.max_steps is not None and self.max_steps < 0: + raise ValueError("Budget.max_steps must be a non-negative integer") + if self.seed is not None and self.seed < 0: + raise ValueError("Budget.seed must be a non-negative integer") + + +def _native(): + from . import alkahest as _alkahest_native + + return _alkahest_native + + +def _budget_exceeded( + message: str, *, remediation: str | None, code: str = "E-BUDGET-001" +) -> BudgetExceededError: + """Build an ``alkahest.BudgetExceededError`` resolved at call time, not import time. + + ``alkahest/__init__.py`` overlays the pure-Python stub in + ``alkahest.exceptions`` with the compiled PyO3 class once the package + finishes initialising (so ``except alkahest.BudgetExceededError`` also + catches errors raised by the Rust cooperative checkpoint). Importing the + stub directly at module load time here — before that overlay runs, since + this module is itself imported from ``alkahest/__init__.py`` — would + construct a *different* class than the one callers catch as + ``alkahest.BudgetExceededError``. + + The compiled PyO3 exception class (unlike the pure-Python stub) has no + ``__init__`` accepting ``code=``/``remediation=`` — Rust sets those as + plain instance attributes after construction (see + ``alkahest-py::make_structured_err``) — so this does the same instead of + calling the constructor with keyword arguments. + """ + from . import BudgetExceededError as _cls + + exc = _cls(message) + exc.code = code + exc.remediation = remediation + exc.span = None + return exc + + +def run_with_wall_fallback( + fn: Callable[..., _T], + /, + *args: Any, + budget: Budget, + **kwargs: Any, +) -> _T: + """Run ``fn(*args, **kwargs)``, enforcing ``budget.wall_ms`` even if ``fn`` + doesn't check the Rust cooperative budget on every path. + + This is a *supplement* to, not a replacement for, entering the budget via + ``context(budget=...)`` — call this from inside such a block (or pass a + budget that also carries ``max_steps``/``seed``) so cooperative call sites + still see it. See the module docstring for why the worker thread is not + forcibly stopped on timeout. + + Parameters + ---------- + fn : callable + The function to call. + *args, **kwargs + Forwarded to ``fn``. + budget : Budget + If ``budget.wall_ms`` is ``None``, this is equivalent to + ``fn(*args, **kwargs)`` — no thread is spawned. + + Raises + ------ + BudgetExceededError + (``E-BUDGET-001``) if ``fn`` does not return within ``budget.wall_ms`` + milliseconds. + + Examples + -------- + >>> import alkahest as ak + >>> b = ak.Budget(wall_ms=5) + >>> ak.run_with_wall_fallback(lambda: ak.simplify(x**2), budget=b) # doctest: +SKIP + """ + if budget.wall_ms is None: + return fn(*args, **kwargs) + + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(fn, *args, **kwargs) + try: + return future.result(timeout=budget.wall_ms / 1000.0) + except concurrent.futures.TimeoutError as exc: + # Best-effort: ask any cooperative Rust checkpoint the call has + # reached (or will reach) to stop, since we can't stop the + # Python thread itself. + request_cancel() + raise _budget_exceeded( + f"[E-BUDGET-001] budget exceeded: wall-clock limit {budget.wall_ms} ms elapsed", + remediation=( + "raise Budget(wall_ms=...), or accept a heuristic/numeric result for this " + "candidate instead of an exact one" + ), + ) from exc + + +def request_cancel() -> None: + """Request cancellation of the current cooperative operation(s), process-wide. + + See ``alkahest_core::budget::request_cancel`` — checked by every + cooperative checkpoint on every thread until :func:`clear_cancel` is + called. + """ + _native().request_cancel() + + +def clear_cancel() -> None: + """Clear a previously requested cancellation. + + Call this before starting the next candidate in a fan-out loop. + """ + _native().clear_cancel() + + +def is_cancelled() -> bool: + """Return ``True`` if :func:`request_cancel` was called and not yet cleared.""" + return bool(_native().is_cancelled()) + + +def is_budget_active() -> bool: + """Return ``True`` if a :class:`Budget` is active on this thread.""" + return bool(_native().is_budget_active()) + + +def budget_seed() -> int | None: + """Return the seed of the innermost active :class:`Budget` on this thread. + + ``None`` if no budget is active or the active budget did not set one. + """ + return _native().budget_seed() diff --git a/python/alkahest/_context.py b/python/alkahest/_context.py index 658822e3..ad4adec7 100644 --- a/python/alkahest/_context.py +++ b/python/alkahest/_context.py @@ -38,6 +38,8 @@ if TYPE_CHECKING: from collections.abc import Generator + from ._budget import Budget + # --------------------------------------------------------------------------- # Thread-local state # --------------------------------------------------------------------------- @@ -71,6 +73,7 @@ def context( precision: int | None = None, assumptions: Any = None, require_certificate: bool | None = None, + budget: Budget | None = None, **extra: Any, ) -> Generator[None, None, None]: """Thread-local context for Alkahest calls. @@ -107,6 +110,15 @@ def context( :func:`alkahest.require_certificate`; use it to stop a research loop from silently accumulating claims it cannot back up. Pass ``False`` in an inner block to opt back out. + budget : Budget, optional + P1 search plumbing item 4. When set, pushes a wall-clock / step + budget and a determinism seed into the Rust-side cooperative + checkpoint (``alkahest_core::budget``) for the scope of this block. + :func:`alkahest.integrate` raises :class:`~alkahest.BudgetExceededError` + (``E-BUDGET-*``) if the budget trips; :func:`alkahest.request_cancel` + trips it from another thread. Like every other context key, a nested + ``context(budget=...)`` shadows this one rather than merging with it — + see :class:`~alkahest.Budget` and ``docs/mdbook/src/budgets.md``. **extra Additional key-value pairs stored in the context and accessible via :func:`get_context_value`. @@ -157,12 +169,24 @@ def context( ctx["assumptions"] = assumptions if require_certificate is not None: ctx["require_certificate"] = require_certificate + if budget is not None: + ctx["budget"] = budget ctx.update(extra) _state.stack.append(ctx) + budget_pushed = False + if budget is not None: + from . import alkahest as _native + + _native.push_budget(wall_ms=budget.wall_ms, max_steps=budget.max_steps, seed=budget.seed) + budget_pushed = True try: yield finally: + if budget_pushed: + from . import alkahest as _native + + _native.pop_budget() _state.stack.pop() @@ -242,6 +266,18 @@ def simplify_enabled() -> bool: return bool(get_context_value("simplify", False)) +def active_budget() -> Any | None: + """Return the :class:`~alkahest.Budget` from the innermost active context, + or ``None`` if no context set one. + + The budget is already active in the Rust-side cooperative checkpoint + while its ``context(budget=...)`` block is open (see :func:`context`); + this accessor is for introspection, e.g. a fan-out loop wanting to + log/adjust its own remaining budget. + """ + return get_context_value("budget") + + def active_assumptions() -> Any | None: """Return the :class:`~alkahest.Assumptions` from the innermost active context, or ``None`` if no context set one. diff --git a/python/alkahest/exceptions.py b/python/alkahest/exceptions.py index 8c65ad40..d3c23268 100644 --- a/python/alkahest/exceptions.py +++ b/python/alkahest/exceptions.py @@ -38,6 +38,7 @@ E-PARSE-* ParseError (reserved; parser not yet integrated) E-DOMAIN-* DomainError (reserved; Python-only pending Rust impl) E-CERT-001 CertificateUnavailableError (Python-only; certificate ledger) + E-BUDGET-001 … E-BUDGET-003 BudgetExceededError (P1 search plumbing item 4) """ from __future__ import annotations @@ -463,6 +464,30 @@ def __init__( super().__init__(message, code="E-CERT-001", remediation=remediation, span=span) +class BudgetExceededError(AlkahestError): + """A :class:`~alkahest.Budget` was exceeded, or cancellation was requested. + + Raised by heavy engines (:func:`alkahest.integrate` today) at a + cooperative checkpoint inside the Rust kernel — see + ``alkahest_core::budget`` and :mod:`alkahest._budget`. A fine, expected + answer for a fan-out search loop, not a crash: ``.code`` distinguishes + the three causes: + + - ``E-BUDGET-001`` — the active budget's wall-clock limit elapsed. + - ``E-BUDGET-002`` — the active budget's step counter exceeded ``max_steps``. + - ``E-BUDGET-003`` — :func:`alkahest.request_cancel` was called and not + yet cleared. + """ + + def __init__( + self, + message: str, + remediation: str | None = None, + span: tuple[int, int] | None = None, + ): + super().__init__(message, code="E-BUDGET-001", remediation=remediation, span=span) + + class ParseError(AlkahestError): """Parse error with source span (reserved; parser not yet integrated).""" diff --git a/tests/test_budget.py b/tests/test_budget.py new file mode 100644 index 00000000..85888f57 --- /dev/null +++ b/tests/test_budget.py @@ -0,0 +1,281 @@ +"""Budgets, cooperative cancellation, and determinism (P1 search plumbing item 4). + +See ``docs/mdbook/src/budgets.md``. These tests exercise the Python-visible +surface: ``Budget``, ``context(budget=...)``, ``BudgetExceededError``, the +process-wide cancellation flag, the determinism seed, and the +``run_with_wall_fallback`` supplement for calls with no Rust checkpoint. +""" + +from __future__ import annotations + +import threading +import time + +import alkahest as ak +import pytest + + +@pytest.fixture +def pool() -> ak.ExprPool: + return ak.ExprPool() + + +@pytest.fixture +def x(pool: ak.ExprPool) -> ak.Expr: + return pool.symbol("x", "real") + + +@pytest.fixture(autouse=True) +def _clear_cancel_before_and_after(): + """Cancellation is a process-wide flag (see the module docs) — never let a + failing assertion in one test leave it set for the next test to inherit.""" + ak.clear_cancel() + yield + ak.clear_cancel() + + +# --------------------------------------------------------------------------- +# Budget dataclass +# --------------------------------------------------------------------------- + + +def test_budget_defaults_are_all_none(): + b = ak.Budget() + assert b.wall_ms is None + assert b.max_steps is None + assert b.seed is None + + +def test_budget_is_immutable(): + b = ak.Budget(wall_ms=1) + with pytest.raises(AttributeError): + b.wall_ms = 2 # type: ignore[misc] + + +@pytest.mark.parametrize( + "kwargs", + [ + {"wall_ms": -1}, + {"wall_ms": float("nan")}, + {"wall_ms": float("inf")}, + {"max_steps": -1}, + {"seed": -1}, + ], +) +def test_budget_rejects_invalid_values(kwargs): + with pytest.raises(ValueError): + ak.Budget(**kwargs) + + +# --------------------------------------------------------------------------- +# context(budget=...) nesting +# --------------------------------------------------------------------------- + + +def test_no_budget_by_default(): + assert not ak.is_budget_active() + assert ak.budget_seed() is None + + +def test_context_budget_activates_and_deactivates(pool): + assert not ak.is_budget_active() + with ak.context(pool=pool, budget=ak.Budget(max_steps=100)): + assert ak.is_budget_active() + assert not ak.is_budget_active() + + +def test_context_budget_round_trips_active_budget(pool): + b = ak.Budget(wall_ms=50, max_steps=100, seed=7) + with ak.context(pool=pool, budget=b): + assert ak.active_budget() == b + assert ak.active_budget() is None + + +def test_seed_round_trips_through_context(pool): + assert ak.budget_seed() is None + with ak.context(pool=pool, budget=ak.Budget(seed=42)): + assert ak.budget_seed() == 42 + assert ak.budget_seed() is None + + +def test_nested_context_without_budget_kw_keeps_outer_active(pool): + """A nested context(...) that omits budget= pushes nothing onto the Rust + stack, so the outer budget (and its seed) stays visible — unlike pool/ + domain, which the inner frame *does* hide (context() doesn't merge).""" + with ak.context(pool=pool, budget=ak.Budget(seed=1, max_steps=1000)), ak.context(pool=pool): + assert ak.budget_seed() == 1 + assert ak.is_budget_active() + + +def test_nested_context_with_budget_kw_shadows_not_merges(pool): + """A nested budget= replaces the outer one for the block — it does not + inherit the outer seed.""" + with ak.context(pool=pool, budget=ak.Budget(seed=1, max_steps=1000)): + with ak.context(pool=pool, budget=ak.Budget(max_steps=2)): + assert ak.budget_seed() is None + # Back to the outer frame on exit from the inner `with`. + assert ak.budget_seed() == 1 + + +def test_context_budget_pops_even_on_exception(pool): + def _raise_inside_budget(): + assert ak.is_budget_active() + raise RuntimeError("boom") + + with pytest.raises(RuntimeError), ak.context(pool=pool, budget=ak.Budget(max_steps=100)): + _raise_inside_budget() + assert not ak.is_budget_active() + + +# --------------------------------------------------------------------------- +# Step budget trips integrate() +# --------------------------------------------------------------------------- + + +def test_step_budget_trips_integrate(pool, x): + """integrate() checks the cooperative budget at its top-level entry, so a + max_steps=0 budget must trip on the very first call.""" + with ak.context(pool=pool, budget=ak.Budget(max_steps=0)): + with pytest.raises(ak.BudgetExceededError) as excinfo: + ak.integrate(x**2, x) + assert excinfo.value.code == "E-BUDGET-002" + + +def test_generous_step_budget_does_not_trip_integrate(pool, x): + with ak.context(pool=pool, budget=ak.Budget(max_steps=1_000_000)): + result = ak.integrate(x**2, x) + assert result.value is not None + + +def test_budget_exceeded_error_is_also_alkahest_error(pool, x): + with ak.context(pool=pool, budget=ak.Budget(max_steps=0)), pytest.raises(ak.AlkahestError): + ak.integrate(x**2, x) + + +def test_budget_exceeded_error_is_also_value_error(pool, x): + """Structured alkahest errors all derive from ValueError, same as every + other exception class in the hierarchy.""" + with ak.context(pool=pool, budget=ak.Budget(max_steps=0)), pytest.raises(ValueError): + ak.integrate(x**2, x) + + +def test_no_budget_active_integrate_still_works(pool, x): + result = ak.integrate(x**2, x) + assert result.value is not None + + +# --------------------------------------------------------------------------- +# Wall-clock budget +# --------------------------------------------------------------------------- + + +def test_wall_budget_trips_after_elapsed(pool, x): + with ak.context(pool=pool, budget=ak.Budget(wall_ms=1)): + time.sleep(0.02) + with pytest.raises(ak.BudgetExceededError) as excinfo: + ak.integrate(x**2, x) + assert excinfo.value.code == "E-BUDGET-001" + + +# --------------------------------------------------------------------------- +# Cancellation +# --------------------------------------------------------------------------- + + +def test_cancel_flag_trips_check_and_clears(pool, x): + assert not ak.is_cancelled() + ak.request_cancel() + assert ak.is_cancelled() + with pytest.raises(ak.BudgetExceededError) as excinfo: + ak.integrate(x**2, x) + assert excinfo.value.code == "E-BUDGET-003" + ak.clear_cancel() + assert not ak.is_cancelled() + result = ak.integrate(x**2, x) + assert result.value is not None + + +def test_cancel_trips_even_without_a_budget_context(pool, x): + """Cancellation is process-wide, not scoped to a Budget frame — it trips + the cooperative checkpoint even with no context(budget=...) active.""" + assert not ak.is_budget_active() + ak.request_cancel() + with pytest.raises(ak.BudgetExceededError): + ak.integrate(x**2, x) + + +def test_cancel_from_another_thread_trips_check_on_this_thread(pool, x): + """The whole point of a process-wide flag: an orchestrator thread can + cancel a heavy call running on a different thread.""" + barrier = threading.Event() + + def watchdog(): + barrier.wait(timeout=2.0) + ak.request_cancel() + + t = threading.Thread(target=watchdog) + t.start() + barrier.set() + t.join(timeout=2.0) + assert ak.is_cancelled() + with pytest.raises(ak.BudgetExceededError): + ak.integrate(x**2, x) + + +# --------------------------------------------------------------------------- +# run_with_wall_fallback — Python-layer supplement for calls with no Rust +# checkpoint on every path (documented use case: simplify). +# --------------------------------------------------------------------------- + + +def test_run_with_wall_fallback_passthrough_when_no_wall_ms(): + assert ak.run_with_wall_fallback(lambda: 1 + 1, budget=ak.Budget()) == 2 + + +def test_run_with_wall_fallback_raises_on_timeout(): + def slow(): + time.sleep(0.2) + return 42 + + with pytest.raises(ak.BudgetExceededError) as excinfo: + ak.run_with_wall_fallback(slow, budget=ak.Budget(wall_ms=10)) + assert excinfo.value.code == "E-BUDGET-001" + + +def test_run_with_wall_fallback_returns_value_when_fast_enough(): + assert ak.run_with_wall_fallback(lambda: 1 + 1, budget=ak.Budget(wall_ms=5_000)) == 2 + + +def test_run_with_wall_fallback_forwards_args_and_kwargs(): + def add(a, b, *, c=0): + return a + b + c + + result = ak.run_with_wall_fallback(add, 1, 2, budget=ak.Budget(wall_ms=5_000), c=3) + assert result == 6 + + +def test_run_with_wall_fallback_propagates_underlying_exception(): + def boom(): + raise KeyError("nope") + + with pytest.raises(KeyError): + ak.run_with_wall_fallback(boom, budget=ak.Budget(wall_ms=5_000)) + + +# --------------------------------------------------------------------------- +# Error codes present and well-formed +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("code", ["E-BUDGET-001", "E-BUDGET-002", "E-BUDGET-003"]) +def test_budget_error_codes_documented(code): + """Lock in the three stable codes from the module docstring / mdbook page.""" + assert code.startswith("E-BUDGET-") + + +def test_budget_exceeded_error_has_remediation(pool, x): + with ak.context(pool=pool, budget=ak.Budget(max_steps=0)): + with pytest.raises(ak.BudgetExceededError) as excinfo: + ak.integrate(x**2, x) + assert excinfo.value.remediation + assert isinstance(excinfo.value.remediation, str) From f0b3f6b96271fd5a5f133afe4850b5a83bc5701f Mon Sep 17 00:00:00 2001 From: Areg Gevorgyan Date: Sat, 8 Aug 2026 15:14:47 -0400 Subject: [PATCH 4/7] fix(integrate): encode budget trips without a new IntegrationError variant Adding IntegrationError::Budget broke cargo-semver-checks on the exhaustive public enum. Carry E-BUDGET-* inside NotImplemented with a marker so Python still raises BudgetExceededError honestly. Co-authored-by: Cursor --- alkahest-core/src/integrate/engine.rs | 79 ++++++++++++++++++++---- alkahest-core/src/integrate/risch/mod.rs | 2 +- alkahest-py/src/lib.rs | 18 +++--- 3 files changed, 75 insertions(+), 24 deletions(-) diff --git a/alkahest-core/src/integrate/engine.rs b/alkahest-core/src/integrate/engine.rs index 6ba52996..0ad9da5d 100644 --- a/alkahest-core/src/integrate/engine.rs +++ b/alkahest-core/src/integrate/engine.rs @@ -26,6 +26,14 @@ use std::fmt; #[derive(Debug, Clone, PartialEq, Eq)] pub enum IntegrationError { /// The expression is outside the supported Risch subset. + /// + /// Also used as a **semver-safe carrier** for budget/cancellation trips + /// (see [`IntegrationError::from`] for [`crate::budget::BudgetError`]): + /// adding a dedicated `Budget` variant would be a major break on this + /// exhaustive enum. Encoded messages start with [`BUDGET_MARKER`]; use + /// [`IntegrationError::is_budget`] / [`IntegrationError::budget_code`] to + /// distinguish them from genuine "not implemented" declines. Python maps + /// these to `BudgetExceededError` (`E-BUDGET-*`). NotImplemented(String), /// Division by zero would occur (e.g. power-rule with n=-1 on a non-x base). DivisionByZero, @@ -33,16 +41,23 @@ pub enum IntegrationError { UnsupportedExtensionDegree(u32), /// The integrand provably has no elementary antiderivative (e.g. elliptic integrals). NonElementary(String), - /// The active [`crate::budget::Budget`] was exceeded, or cancellation was - /// requested, at a cooperative checkpoint inside the integration engine. - /// See `crate::budget` — P1 search plumbing item 4. - Budget(crate::budget::BudgetError), } +/// Prefix for [`IntegrationError::NotImplemented`] messages that encode a +/// [`crate::budget::BudgetError`]. Invisible to casual grepping of user-facing +/// "not implemented" strings; stripped from [`Display`]. +const BUDGET_MARKER: &str = "[[budget]]"; + impl fmt::Display for IntegrationError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { - IntegrationError::NotImplemented(msg) => write!(f, "integrate: not implemented: {msg}"), + IntegrationError::NotImplemented(msg) => { + if let Some(rest) = msg.strip_prefix(BUDGET_MARKER) { + write!(f, "integrate: {rest}") + } else { + write!(f, "integrate: not implemented: {msg}") + } + } IntegrationError::DivisionByZero => write!(f, "integrate: division by zero"), IntegrationError::UnsupportedExtensionDegree(q) => write!( f, @@ -52,7 +67,6 @@ impl fmt::Display for IntegrationError { IntegrationError::NonElementary(msg) => { write!(f, "integrate: no elementary antiderivative exists: {msg}") } - IntegrationError::Budget(e) => write!(f, "integrate: {e}"), } } } @@ -61,13 +75,56 @@ impl std::error::Error for IntegrationError {} impl From for IntegrationError { fn from(e: crate::budget::BudgetError) -> Self { - IntegrationError::Budget(e) + use crate::errors::AlkahestError; + // Encode code + Display body so Python/callers keep E-BUDGET-* without + // a new exhaustive-enum variant (cargo-semver-checks major). + IntegrationError::NotImplemented(format!("{BUDGET_MARKER}[{}] {e}", e.code())) } } impl IntegrationError { + /// `true` when this error encodes a budget/cancellation trip rather than a + /// genuine "outside the Risch subset" decline. + pub fn is_budget(&self) -> bool { + matches!(self, IntegrationError::NotImplemented(msg) if msg.starts_with(BUDGET_MARKER)) + } + + /// The `E-BUDGET-*` code when [`is_budget`](Self::is_budget), else `None`. + pub fn budget_code(&self) -> Option<&'static str> { + let IntegrationError::NotImplemented(msg) = self else { + return None; + }; + let rest = msg.strip_prefix(BUDGET_MARKER)?; + if rest.starts_with("[E-BUDGET-001]") { + Some("E-BUDGET-001") + } else if rest.starts_with("[E-BUDGET-002]") { + Some("E-BUDGET-002") + } else if rest.starts_with("[E-BUDGET-003]") { + Some("E-BUDGET-003") + } else { + None + } + } + /// A human-readable remediation hint for the user. pub fn remediation(&self) -> Option<&'static str> { + if let Some(code) = self.budget_code() { + return match code { + "E-BUDGET-001" => Some( + "raise Budget(wall_ms=...), or accept a heuristic/numeric result for this \ + candidate instead of an exact one", + ), + "E-BUDGET-002" => Some( + "raise Budget(max_steps=...), or accept a partial/heuristic result for this \ + candidate instead of an exact one", + ), + "E-BUDGET-003" => Some( + "call alkahest.clear_cancel() (Python) or budget::clear_cancel() (Rust) before \ + starting the next candidate", + ), + _ => None, + }; + } match self { IntegrationError::NotImplemented(_) => Some( "only power, linearity, sin/cos/exp rules and algebraic (sqrt) rules \ @@ -82,10 +139,6 @@ impl IntegrationError { "this integrand has no closed-form antiderivative in terms of elementary \ functions; use a numeric integrator or elliptic-integral library", ), - IntegrationError::Budget(e) => { - use crate::errors::AlkahestError; - e.remediation() - } } } @@ -97,12 +150,14 @@ impl IntegrationError { impl crate::errors::AlkahestError for IntegrationError { fn code(&self) -> &'static str { + if let Some(code) = self.budget_code() { + return code; + } match self { IntegrationError::NotImplemented(_) => "E-INT-001", IntegrationError::DivisionByZero => "E-INT-002", IntegrationError::UnsupportedExtensionDegree(_) => "E-INT-003", IntegrationError::NonElementary(_) => "E-INT-004", - IntegrationError::Budget(e) => e.code(), } } diff --git a/alkahest-core/src/integrate/risch/mod.rs b/alkahest-core/src/integrate/risch/mod.rs index 1198ba37..68597e6c 100644 --- a/alkahest-core/src/integrate/risch/mod.rs +++ b/alkahest-core/src/integrate/risch/mod.rs @@ -280,7 +280,7 @@ pub fn integrate_risch( Err(IntegrationError::UnsupportedExtensionDegree(d)) => { return Err(IntegrationError::UnsupportedExtensionDegree(d)); } - Err(e @ IntegrationError::Budget(_)) => { + Err(e) if e.is_budget() => { // A budget/cancellation trip is never a "this route declined" // signal — propagate it immediately instead of falling // through to sum decomposition, which would keep spending the diff --git a/alkahest-py/src/lib.rs b/alkahest-py/src/lib.rs index d473d610..920ae9c4 100644 --- a/alkahest-py/src/lib.rs +++ b/alkahest-py/src/lib.rs @@ -450,8 +450,13 @@ fn integrate_error_to_py(e: IntegrationError) -> PyErr { // A budget/cancellation trip is not an integration failure — raise the // dedicated `BudgetExceededError` (E-BUDGET-*) instead of `IntegrationError` // so callers can catch it uniformly regardless of which engine tripped it. - if let IntegrationError::Budget(inner) = &e { - return budget_error_to_py(inner); + // (Budget trips are encoded inside `NotImplemented` for Rust semver; see + // `IntegrationError::is_budget`.) + if e.is_budget() { + return Python::with_gil(|py| { + let exc_type = py.get_type_bound::(); + make_structured_err(py, &exc_type, &e) + }); } Python::with_gil(|py| { let exc_type = py.get_type_bound::(); @@ -459,15 +464,6 @@ fn integrate_error_to_py(e: IntegrationError) -> PyErr { }) } -/// Map a [`alkahest_core::budget::BudgetError`] to Python's -/// `BudgetExceededError` (`E-BUDGET-001..003`). See `crate::budget` — P1 -/// search plumbing item 4. -fn budget_error_to_py(e: &alkahest_core::budget::BudgetError) -> PyErr { - Python::with_gil(|py| { - let exc_type = py.get_type_bound::(); - make_structured_err(py, &exc_type, e) - }) -} // --------------------------------------------------------------------------- // P1 search plumbing item 4 — budgets, cancellation, determinism From f6aadfbaab50d40b3bd79219e7c8c67f9af8d461 Mon Sep 17 00:00:00 2001 From: Areg Gevorgyan Date: Sat, 8 Aug 2026 15:23:53 -0400 Subject: [PATCH 5/7] fix(ci): rustfmt + drop private rustdoc link to budget marker Co-authored-by: Cursor --- alkahest-core/src/integrate/engine.rs | 9 +++++---- alkahest-py/src/lib.rs | 1 - 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/alkahest-core/src/integrate/engine.rs b/alkahest-core/src/integrate/engine.rs index 0ad9da5d..43b96c0e 100644 --- a/alkahest-core/src/integrate/engine.rs +++ b/alkahest-core/src/integrate/engine.rs @@ -30,10 +30,11 @@ pub enum IntegrationError { /// Also used as a **semver-safe carrier** for budget/cancellation trips /// (see [`IntegrationError::from`] for [`crate::budget::BudgetError`]): /// adding a dedicated `Budget` variant would be a major break on this - /// exhaustive enum. Encoded messages start with [`BUDGET_MARKER`]; use - /// [`IntegrationError::is_budget`] / [`IntegrationError::budget_code`] to - /// distinguish them from genuine "not implemented" declines. Python maps - /// these to `BudgetExceededError` (`E-BUDGET-*`). + /// exhaustive enum. Encoded messages start with the internal `[[budget]]` + /// marker; use [`IntegrationError::is_budget`] / + /// [`IntegrationError::budget_code`] to distinguish them from genuine + /// "not implemented" declines. Python maps these to `BudgetExceededError` + /// (`E-BUDGET-*`). NotImplemented(String), /// Division by zero would occur (e.g. power-rule with n=-1 on a non-x base). DivisionByZero, diff --git a/alkahest-py/src/lib.rs b/alkahest-py/src/lib.rs index 920ae9c4..e3273ee6 100644 --- a/alkahest-py/src/lib.rs +++ b/alkahest-py/src/lib.rs @@ -464,7 +464,6 @@ fn integrate_error_to_py(e: IntegrationError) -> PyErr { }) } - // --------------------------------------------------------------------------- // P1 search plumbing item 4 — budgets, cancellation, determinism // From ba7bfc6ff511e85e0f56ed7d878c079f9272f345 Mon Sep 17 00:00:00 2001 From: Areg Gevorgyan Date: Sat, 8 Aug 2026 15:36:48 -0400 Subject: [PATCH 6/7] style: ruff-format _budget.py for CI Co-authored-by: Cursor --- python/alkahest/_budget.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/alkahest/_budget.py b/python/alkahest/_budget.py index cab6356a..9f1f52a5 100644 --- a/python/alkahest/_budget.py +++ b/python/alkahest/_budget.py @@ -100,9 +100,7 @@ class Budget: def __post_init__(self) -> None: if self.wall_ms is not None and (not math.isfinite(self.wall_ms) or self.wall_ms < 0): - raise ValueError( - "Budget.wall_ms must be a finite, non-negative number of milliseconds" - ) + raise ValueError("Budget.wall_ms must be a finite, non-negative number of milliseconds") if self.max_steps is not None and self.max_steps < 0: raise ValueError("Budget.max_steps must be a non-negative integer") if self.seed is not None and self.seed < 0: From 9500e21aa5a6c3c88f9bdeb03d1b1020e78eaa5f Mon Sep 17 00:00:00 2001 From: Areg Gevorgyan Date: Sat, 8 Aug 2026 15:45:56 -0400 Subject: [PATCH 7/7] fix(python): satisfy ty on Budget dataclass and batch_map return type Co-authored-by: Cursor --- python/alkahest/_batch.py | 13 ++++++------- python/alkahest/_budget.py | 2 +- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/python/alkahest/_batch.py b/python/alkahest/_batch.py index e03039d5..6339d44f 100644 --- a/python/alkahest/_batch.py +++ b/python/alkahest/_batch.py @@ -200,14 +200,13 @@ def batch_map( if not parallel: return [_invoke(fn, item, i, kwargs) for i, item in enumerate(materialized)] - results: list[BatchItem | None] = [None] * len(materialized) with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = { - executor.submit(_invoke, fn, item, i, kwargs): i for i, item in enumerate(materialized) - } - for future in futures: - results[futures[future]] = future.result() - return results # type: ignore[return-value] # every slot was filled above + # Submit in input order and collect in the same order so the return + # type stays ``list[BatchItem]`` (no ``None`` placeholders for ty). + futures = [ + executor.submit(_invoke, fn, item, i, kwargs) for i, item in enumerate(materialized) + ] + return [future.result() for future in futures] def batch_map_iter( diff --git a/python/alkahest/_budget.py b/python/alkahest/_budget.py index 9f1f52a5..04dd7f29 100644 --- a/python/alkahest/_budget.py +++ b/python/alkahest/_budget.py @@ -61,7 +61,7 @@ _T = TypeVar("_T") -@dataclass(frozen=True, slots=True) +@dataclass(frozen=True) class Budget: """A per-call resource budget for search-style workloads.