surestop 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
surestop-0.2.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Adam Noonan
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,173 @@
1
+ Metadata-Version: 2.4
2
+ Name: surestop
3
+ Version: 0.2.0
4
+ Summary: Stop bad runs early, with a conformal bound on how often you stop one that would have ended well
5
+ Author: Adam Noonan
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/ACNoonan/surestop
8
+ Project-URL: Issues, https://github.com/ACNoonan/surestop/issues
9
+ Keywords: hyperparameter-optimization,early-stopping,pruning,conformal,optuna
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Science/Research
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
14
+ Requires-Python: >=3.10
15
+ Description-Content-Type: text/markdown
16
+ License-File: LICENSE
17
+ Requires-Dist: numpy>=1.24
18
+ Requires-Dist: scipy>=1.10
19
+ Provides-Extra: optuna
20
+ Requires-Dist: optuna>=4; extra == "optuna"
21
+ Provides-Extra: test
22
+ Requires-Dist: pytest; extra == "test"
23
+ Requires-Dist: optuna>=4; extra == "test"
24
+ Dynamic: license-file
25
+
26
+ # SureStop
27
+
28
+ **Stop bad training runs early, with a stated bound on how often you stop one that would have
29
+ ended well.** You pick α. The rule kills runs so that the expected fraction of runs that were
30
+ good *and* got killed stays at or below α.
31
+
32
+ Every pruner in common use (median, percentile, successive halving, Hyperband) is a heuristic,
33
+ and none says how often it kills a run that would have finished well. This one does, under a
34
+ condition you can check: the runs you calibrate on and the runs you judge must come from the same
35
+ random process.
36
+
37
+ It was measured on 400 real LoRA fine-tuning runs under a pre-registered protocol. The curves are
38
+ in this repository, and a script checks that the package reproduces the recorded numbers exactly.
39
+
40
+ ## Install
41
+
42
+ ```
43
+ pip install surestop # numpy + scipy
44
+ pip install "surestop[optuna]" # adds the Optuna pruner
45
+ ```
46
+
47
+ ## Use
48
+
49
+ Plain numpy, on any training loop:
50
+
51
+ ```python
52
+ from surestop import KillRule
53
+
54
+ # completed_curves: array of shape (runs, evaluations), lower is better
55
+ rule = KillRule(alpha=0.05, good_threshold=2.16, min_peek=17).fit(completed_curves)
56
+
57
+ rule.should_kill(curve_so_far) # True -> stop this run now
58
+ rule.kill_step(full_curve) # where it would have fired, or None
59
+ rule.evaluate(held_out_curves) # false kills, joint rate, savings
60
+ ```
61
+
62
+ With Optuna, the first `n_calibration` trials run unpruned to calibrate, then the rule freezes:
63
+
64
+ ```python
65
+ import optuna
66
+ from surestop import ConformalPruner
67
+
68
+ study = optuna.create_study(sampler=optuna.samplers.RandomSampler(),
69
+ pruner=ConformalPruner(n_calibration=60, alpha=0.05))
70
+ ```
71
+
72
+ Pass `maximize=True` (or a `direction="maximize"` study) when higher is better.
73
+
74
+ ## The rule
75
+
76
+ A run is *good* if its final value ends at or below a threshold τ. At every evaluation, a
77
+ predictor forecasts the final value from the curve so far. The default predictor is the latest
78
+ observed value.
79
+
80
+ To calibrate, take n runs that have already finished. For each good one, record its worst moment:
81
+ the largest gap between forecast and τ at any evaluation. Runs that are not good contribute
82
+ nothing. Set λ to the ⌈(1−α)(n+1)⌉-th smallest of these scores. A new run is killed at the first
83
+ evaluation where its forecast exceeds τ + λ.
84
+
85
+ This is conformal risk control on the loss "good and killed". Taking the worst moment across all
86
+ evaluations means one threshold covers every look at a run, so there is no correction for
87
+ multiple looks. The guarantee needs the calibration runs and the new runs to be exchangeable.
88
+ Random sampling gives that. Adaptive samplers such as TPE do not.
89
+
90
+ ## What was measured
91
+
92
+ Experiment CC-03c, pre-registered before any run finished:
93
+
94
+ - **Runs:** 400 LoRA fine-tunes of Qwen2.5-0.5B on dolly-15k, 750 steps each, validation loss
95
+ measured 50 times per run.
96
+ - **Configs:** i.i.d. draws from a frozen hyperparameter prior (`docs/hyperparameter-prior.md`).
97
+ - **τ:** pinned from an earlier pool before the first run landed.
98
+ - **Splits:** two, one random and one past-to-future by night, each about half calibration and
99
+ half test.
100
+
101
+ The pre-registered rule used a fitted power-law predictor. On both splits it kept the bound and
102
+ saved about half the evaluation compute:
103
+
104
+ | split | false kills | joint rate (95% CI) | eval compute saved |
105
+ |---|---|---|---|
106
+ | random | 13 of 200 | 6.5% (3.5–10.9%) | 52.2% |
107
+ | past-to-future | 11 of 203 | 5.4% (2.7–9.5%) | 52.8% |
108
+
109
+ A rule that kills every run immediately was scored by the same gate and failed on both splits.
110
+
111
+ **The flaw that run found:** every kill fired at the 3rd of 50 evaluations. Early rankings only
112
+ match final rankings from about evaluation 18. The rule stayed within budget because the bound
113
+ held. The fitted power law ranked runs worse than the latest observed value at every early
114
+ evaluation.
115
+
116
+ **What this package ships instead,** and what the reproduction script checks:
117
+
118
+ | predictor | hold until | random split | past-to-future split | 2000 resplits |
119
+ |---|---|---|---|---|
120
+ | last value | none | 18 of 200 (9.0%), 78% saved | 7 of 203 (3.4%), 75% saved | joint 5.0%, breach 8.5% |
121
+ | last value | eval 18 | 18 of 200 (9.0%), 57% saved | 7 of 203 (3.4%), 56% saved | joint 5.0%, breach 8.6% |
122
+
123
+ The random split's 18 of 200 is a breach. It sits at the 95th percentile of that rule's own
124
+ resplit distribution, whose mean is 5.0%, so it reads as a split tail rather than a leak. A
125
+ correct rule at this sample size lands on a breach in about 7% of splits. The held version was
126
+ chosen *after* seeing these 400 runs, so its numbers are a hypothesis. It halves
127
+ how often the single best run is killed at the same budget. Its confirmation needs a fresh pool.
128
+
129
+ Full account: `docs/writeup.md`. Run the check yourself:
130
+
131
+ ```
132
+ python reproduce/reproduce_cc03c.py
133
+ ```
134
+
135
+ ## What the bound does not cover
136
+
137
+ - **It needs exchangeable runs.** With a random or quasi-random sampler the bound holds. With
138
+ TPE, Gaussian-process or CMA-ES samplers it does not. The Optuna pruner warns once and carries
139
+ on, but do not rely on the bound there.
140
+ - **It bounds the joint rate rather than the share of good runs killed.** When good runs are rare, a large
141
+ share of them can still die: about a third in the measurement above, mostly borderline ones.
142
+ Lower α if that matters.
143
+ - **An estimated τ gives a measured bound rather than a proven one.** `good_quantile` estimates τ from the
144
+ calibration runs. Simulations put the realized rate at 0.88 to 0.97 times α. Pass a real
145
+ `good_threshold` when you have one.
146
+ - **α = 0.05 needs at least 19 calibration runs**, and the certified threshold gets loose with
147
+ few good runs. Sweeps of a few dozen trials are too small for this.
148
+ - **One setting only so far.** One model, one dataset, one horizon.
149
+
150
+ ## Layout
151
+
152
+ | path | what |
153
+ |---|---|
154
+ | `src/surestop/` | `KillRule` (numpy/scipy) and `ConformalPruner` (Optuna) |
155
+ | `tests/` | the bound holds on synthetic exchangeable runs, and the same assertion fails on a rule that leaks 2× its budget |
156
+ | `reproduce/` | the 400 curves, the pre-registered splits and targets, the script, and its committed output |
157
+ | `docs/` | the write-up, the prior-art sweep, and the frozen hyperparameter prior the curves were drawn from |
158
+ | `optunahub/` | the package as an OptunaHub registry entry |
159
+
160
+ ## Related work
161
+
162
+ Calibrating a stop threshold on a predictor with a conformal guarantee is an established pattern. Xie et al.
163
+ (arXiv:2602.13935) use the same device, a threshold on the maximum of a running score calibrated
164
+ on one class only, to stop LLM reasoning traces. Related work does the same for mixed-integer
165
+ solvers and LLM agent episodes. As far as a search found, nobody has applied it to pruning
166
+ training runs, and no major tuning library ships a pruner with a stated error guarantee.
167
+ `docs/prior-art.md` has the sweep.
168
+
169
+ ## Citing
170
+
171
+ Concept DOI: [10.5281/zenodo.22726440](https://doi.org/10.5281/zenodo.22726440), which always resolves
172
+ to the newest version. `CITATION.cff` has the full entry. The curves are a measurement, released
173
+ under the same MIT licence as the code.
@@ -0,0 +1,148 @@
1
+ # SureStop
2
+
3
+ **Stop bad training runs early, with a stated bound on how often you stop one that would have
4
+ ended well.** You pick α. The rule kills runs so that the expected fraction of runs that were
5
+ good *and* got killed stays at or below α.
6
+
7
+ Every pruner in common use (median, percentile, successive halving, Hyperband) is a heuristic,
8
+ and none says how often it kills a run that would have finished well. This one does, under a
9
+ condition you can check: the runs you calibrate on and the runs you judge must come from the same
10
+ random process.
11
+
12
+ It was measured on 400 real LoRA fine-tuning runs under a pre-registered protocol. The curves are
13
+ in this repository, and a script checks that the package reproduces the recorded numbers exactly.
14
+
15
+ ## Install
16
+
17
+ ```
18
+ pip install surestop # numpy + scipy
19
+ pip install "surestop[optuna]" # adds the Optuna pruner
20
+ ```
21
+
22
+ ## Use
23
+
24
+ Plain numpy, on any training loop:
25
+
26
+ ```python
27
+ from surestop import KillRule
28
+
29
+ # completed_curves: array of shape (runs, evaluations), lower is better
30
+ rule = KillRule(alpha=0.05, good_threshold=2.16, min_peek=17).fit(completed_curves)
31
+
32
+ rule.should_kill(curve_so_far) # True -> stop this run now
33
+ rule.kill_step(full_curve) # where it would have fired, or None
34
+ rule.evaluate(held_out_curves) # false kills, joint rate, savings
35
+ ```
36
+
37
+ With Optuna, the first `n_calibration` trials run unpruned to calibrate, then the rule freezes:
38
+
39
+ ```python
40
+ import optuna
41
+ from surestop import ConformalPruner
42
+
43
+ study = optuna.create_study(sampler=optuna.samplers.RandomSampler(),
44
+ pruner=ConformalPruner(n_calibration=60, alpha=0.05))
45
+ ```
46
+
47
+ Pass `maximize=True` (or a `direction="maximize"` study) when higher is better.
48
+
49
+ ## The rule
50
+
51
+ A run is *good* if its final value ends at or below a threshold τ. At every evaluation, a
52
+ predictor forecasts the final value from the curve so far. The default predictor is the latest
53
+ observed value.
54
+
55
+ To calibrate, take n runs that have already finished. For each good one, record its worst moment:
56
+ the largest gap between forecast and τ at any evaluation. Runs that are not good contribute
57
+ nothing. Set λ to the ⌈(1−α)(n+1)⌉-th smallest of these scores. A new run is killed at the first
58
+ evaluation where its forecast exceeds τ + λ.
59
+
60
+ This is conformal risk control on the loss "good and killed". Taking the worst moment across all
61
+ evaluations means one threshold covers every look at a run, so there is no correction for
62
+ multiple looks. The guarantee needs the calibration runs and the new runs to be exchangeable.
63
+ Random sampling gives that. Adaptive samplers such as TPE do not.
64
+
65
+ ## What was measured
66
+
67
+ Experiment CC-03c, pre-registered before any run finished:
68
+
69
+ - **Runs:** 400 LoRA fine-tunes of Qwen2.5-0.5B on dolly-15k, 750 steps each, validation loss
70
+ measured 50 times per run.
71
+ - **Configs:** i.i.d. draws from a frozen hyperparameter prior (`docs/hyperparameter-prior.md`).
72
+ - **τ:** pinned from an earlier pool before the first run landed.
73
+ - **Splits:** two, one random and one past-to-future by night, each about half calibration and
74
+ half test.
75
+
76
+ The pre-registered rule used a fitted power-law predictor. On both splits it kept the bound and
77
+ saved about half the evaluation compute:
78
+
79
+ | split | false kills | joint rate (95% CI) | eval compute saved |
80
+ |---|---|---|---|
81
+ | random | 13 of 200 | 6.5% (3.5–10.9%) | 52.2% |
82
+ | past-to-future | 11 of 203 | 5.4% (2.7–9.5%) | 52.8% |
83
+
84
+ A rule that kills every run immediately was scored by the same gate and failed on both splits.
85
+
86
+ **The flaw that run found:** every kill fired at the 3rd of 50 evaluations. Early rankings only
87
+ match final rankings from about evaluation 18. The rule stayed within budget because the bound
88
+ held. The fitted power law ranked runs worse than the latest observed value at every early
89
+ evaluation.
90
+
91
+ **What this package ships instead,** and what the reproduction script checks:
92
+
93
+ | predictor | hold until | random split | past-to-future split | 2000 resplits |
94
+ |---|---|---|---|---|
95
+ | last value | none | 18 of 200 (9.0%), 78% saved | 7 of 203 (3.4%), 75% saved | joint 5.0%, breach 8.5% |
96
+ | last value | eval 18 | 18 of 200 (9.0%), 57% saved | 7 of 203 (3.4%), 56% saved | joint 5.0%, breach 8.6% |
97
+
98
+ The random split's 18 of 200 is a breach. It sits at the 95th percentile of that rule's own
99
+ resplit distribution, whose mean is 5.0%, so it reads as a split tail rather than a leak. A
100
+ correct rule at this sample size lands on a breach in about 7% of splits. The held version was
101
+ chosen *after* seeing these 400 runs, so its numbers are a hypothesis. It halves
102
+ how often the single best run is killed at the same budget. Its confirmation needs a fresh pool.
103
+
104
+ Full account: `docs/writeup.md`. Run the check yourself:
105
+
106
+ ```
107
+ python reproduce/reproduce_cc03c.py
108
+ ```
109
+
110
+ ## What the bound does not cover
111
+
112
+ - **It needs exchangeable runs.** With a random or quasi-random sampler the bound holds. With
113
+ TPE, Gaussian-process or CMA-ES samplers it does not. The Optuna pruner warns once and carries
114
+ on, but do not rely on the bound there.
115
+ - **It bounds the joint rate rather than the share of good runs killed.** When good runs are rare, a large
116
+ share of them can still die: about a third in the measurement above, mostly borderline ones.
117
+ Lower α if that matters.
118
+ - **An estimated τ gives a measured bound rather than a proven one.** `good_quantile` estimates τ from the
119
+ calibration runs. Simulations put the realized rate at 0.88 to 0.97 times α. Pass a real
120
+ `good_threshold` when you have one.
121
+ - **α = 0.05 needs at least 19 calibration runs**, and the certified threshold gets loose with
122
+ few good runs. Sweeps of a few dozen trials are too small for this.
123
+ - **One setting only so far.** One model, one dataset, one horizon.
124
+
125
+ ## Layout
126
+
127
+ | path | what |
128
+ |---|---|
129
+ | `src/surestop/` | `KillRule` (numpy/scipy) and `ConformalPruner` (Optuna) |
130
+ | `tests/` | the bound holds on synthetic exchangeable runs, and the same assertion fails on a rule that leaks 2× its budget |
131
+ | `reproduce/` | the 400 curves, the pre-registered splits and targets, the script, and its committed output |
132
+ | `docs/` | the write-up, the prior-art sweep, and the frozen hyperparameter prior the curves were drawn from |
133
+ | `optunahub/` | the package as an OptunaHub registry entry |
134
+
135
+ ## Related work
136
+
137
+ Calibrating a stop threshold on a predictor with a conformal guarantee is an established pattern. Xie et al.
138
+ (arXiv:2602.13935) use the same device, a threshold on the maximum of a running score calibrated
139
+ on one class only, to stop LLM reasoning traces. Related work does the same for mixed-integer
140
+ solvers and LLM agent episodes. As far as a search found, nobody has applied it to pruning
141
+ training runs, and no major tuning library ships a pruner with a stated error guarantee.
142
+ `docs/prior-art.md` has the sweep.
143
+
144
+ ## Citing
145
+
146
+ Concept DOI: [10.5281/zenodo.22726440](https://doi.org/10.5281/zenodo.22726440), which always resolves
147
+ to the newest version. `CITATION.cff` has the full entry. The curves are a measurement, released
148
+ under the same MIT licence as the code.
@@ -0,0 +1,36 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "surestop"
7
+ version = "0.2.0"
8
+ description = "Stop bad runs early, with a conformal bound on how often you stop one that would have ended well"
9
+ readme = "README.md"
10
+ license = "MIT"
11
+ license-files = ["LICENSE"]
12
+ authors = [{ name = "Adam Noonan" }]
13
+ requires-python = ">=3.10"
14
+ dependencies = ["numpy>=1.24", "scipy>=1.10"]
15
+ keywords = ["hyperparameter-optimization", "early-stopping", "pruning", "conformal", "optuna"]
16
+ classifiers = [
17
+ "Development Status :: 3 - Alpha",
18
+ "Intended Audience :: Science/Research",
19
+ "Programming Language :: Python :: 3",
20
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
21
+ ]
22
+
23
+ [project.optional-dependencies]
24
+ optuna = ["optuna>=4"]
25
+ test = ["pytest", "optuna>=4"]
26
+
27
+ [project.urls]
28
+ Homepage = "https://github.com/ACNoonan/surestop"
29
+ Issues = "https://github.com/ACNoonan/surestop/issues"
30
+
31
+ [tool.setuptools.packages.find]
32
+ where = ["src"]
33
+ include = ["surestop*"]
34
+
35
+ [tool.pytest.ini_options]
36
+ testpaths = ["tests"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,17 @@
1
+ """surestop — stop bad runs early, with a bound on how often you stop one that would have ended well."""
2
+
3
+ from .core import KillRule, NotEnoughCalibration, last_value
4
+
5
+ __version__ = "0.2.0"
6
+ __all__ = ["KillRule", "NotEnoughCalibration", "last_value", "ConformalPruner"]
7
+
8
+
9
+ from typing import Any
10
+
11
+
12
+ def __getattr__(name: str) -> Any:
13
+ # Optuna is optional: import the pruner only when someone asks for it.
14
+ if name == "ConformalPruner":
15
+ from .optuna_pruner import ConformalPruner
16
+ return ConformalPruner
17
+ raise AttributeError(name)
@@ -0,0 +1,201 @@
1
+ """A kill rule for training runs that bounds how often it kills a run that would have ended well.
2
+
3
+ The construction is the one CC-03c pre-registered and measured on 400 real LoRA runs
4
+ (`experiments/2026-08-17-CC-03c-validity-gate-that-can-fire/`):
5
+
6
+ 1. Calibrate on n **completed** runs, each a learning curve on a shared step grid.
7
+ 2. A run is *good* if its final value is at or below a threshold τ (lower is better).
8
+ 3. A predictor forecasts the final value from the curve so far, at every step.
9
+ 4. Each good calibration run scores max over eligible steps of (prediction − τ).
10
+ Runs that are not good score −∞.
11
+ 5. λ̂ is the ⌈(1−α)(n+1)⌉-th smallest score. This is conformal risk control on the loss
12
+ "this run is good AND the rule kills it".
13
+ 6. A new run is killed at the first eligible step where prediction − λ̂ > τ.
14
+
15
+ **Guarantee.** If the calibration runs and the new runs are exchangeable — for example, drawn
16
+ i.i.d. by the same random sampler — then the expected fraction of new runs that are good and
17
+ killed is at most α. One kill decision covers every step at once, with no multiplicity
18
+ correction.
19
+
20
+ **What it does not promise.**
21
+ - **It bounds the joint rate, not the share of good runs killed.** With a small good fraction,
22
+ the conditional rate can be large: ~30% on CC-03c.
23
+ - **It needs exchangeability.** An adaptive sampler such as TPE shifts later runs away from the
24
+ calibration runs, and the guarantee no longer applies.
25
+ - **With τ estimated from the calibration finals** (`good_quantile`), the bound is measured, not
26
+ proven: 0.88–0.97× α in CC-03c's simulations. Pass `good_threshold` when you have a real target.
27
+ """
28
+
29
+ from __future__ import annotations
30
+
31
+ import math
32
+ import warnings
33
+ from dataclasses import dataclass, field
34
+ from typing import Any, Callable
35
+
36
+ import numpy as np
37
+ from numpy.typing import ArrayLike
38
+ from scipy.stats import rankdata
39
+
40
+ Predictor = Callable[[np.ndarray], np.ndarray]
41
+
42
+
43
+ def last_value(curves: ArrayLike) -> np.ndarray:
44
+ """Forecast the final value as the latest observed value.
45
+
46
+ A predictor maps curves of shape (n, T) to forecasts of shape (n, T) and must be causal:
47
+ column t may use only columns 0..t. On CC-03c's pool this beat a fitted pow3 curve at every
48
+ early step.
49
+ """
50
+ return np.asarray(curves, dtype=float)
51
+
52
+
53
+ class NotEnoughCalibration(ValueError):
54
+ """Too few calibration runs for any threshold to certify at this α."""
55
+
56
+
57
+ def _spearman_cols(M: np.ndarray, y: np.ndarray) -> np.ndarray:
58
+ rm = rankdata(M, axis=0)
59
+ rm -= rm.mean(axis=0)
60
+ ry = rankdata(y)
61
+ ry -= ry.mean()
62
+ denom = np.sqrt((rm ** 2).sum(axis=0) * (ry ** 2).sum())
63
+ with np.errstate(invalid="ignore", divide="ignore"):
64
+ return np.where(denom > 0, (rm * ry[:, None]).sum(axis=0) / denom, 0.0)
65
+
66
+
67
+ @dataclass
68
+ class KillRule:
69
+ """Certified early-kill rule. Fit on completed curves, then ask `should_kill` on partial ones.
70
+
71
+ Parameters
72
+ ----------
73
+ alpha : bound on the expected fraction of runs that are good and killed.
74
+ good_threshold : a run is good if its final value is at or below this. Overrides `good_quantile`.
75
+ good_quantile : if no threshold is given, τ is this quantile of the calibration finals.
76
+ min_peek : index of the earliest step at which a kill may fire.
77
+ hold_until_resolved : if set (e.g. 0.90), also hold kills until the first step at or after
78
+ `min_peek` where Spearman(prediction, final) on the calibration runs reaches this value.
79
+ Estimated per calibration set, so it varies between sets; pin `min_peek` instead when a
80
+ prior pool tells you where the ranking settles.
81
+ maximize : set True when higher values are better (accuracy, reward).
82
+ predictor : causal forecaster of the final value. Default: last observed value.
83
+ """
84
+
85
+ alpha: float = 0.05
86
+ good_threshold: float | None = None
87
+ good_quantile: float = 0.20
88
+ min_peek: int = 0
89
+ hold_until_resolved: float | None = None
90
+ maximize: bool = False
91
+ predictor: Predictor = last_value
92
+
93
+ tau_: float = field(init=False, default=math.nan)
94
+ lambda_: float = field(init=False, default=math.nan)
95
+ first_peek_: int = field(init=False, default=-1)
96
+ n_steps_: int = field(init=False, default=0)
97
+ n_calibration_: int = field(init=False, default=0)
98
+ n_good_: int = field(init=False, default=0)
99
+ resolution_: np.ndarray | None = field(init=False, default=None)
100
+
101
+ # ── fitting ────────────────────────────────────────────────────────────────
102
+ def _oriented(self, curves: ArrayLike) -> np.ndarray:
103
+ Y = np.asarray(curves, dtype=float)
104
+ return -Y if self.maximize else Y
105
+
106
+ def fit(self, curves: ArrayLike) -> "KillRule":
107
+ Y = self._oriented(curves)
108
+ if Y.ndim != 2 or Y.shape[1] < 2:
109
+ raise ValueError("curves must be a 2-D array (runs × steps) with at least 2 steps")
110
+ if not np.all(np.isfinite(Y)):
111
+ raise ValueError("calibration curves must be finite; drop incomplete runs first")
112
+ if not 0 < self.alpha < 1:
113
+ raise ValueError("alpha must lie in (0, 1)")
114
+ n, T = Y.shape
115
+ order = math.ceil((1 - self.alpha) * (n + 1))
116
+ if order > n:
117
+ raise NotEnoughCalibration(
118
+ f"alpha={self.alpha} needs at least {math.ceil(1 / self.alpha) - 1} calibration "
119
+ f"runs for any threshold to certify; got {n}")
120
+
121
+ finals = Y[:, -1]
122
+ if self.good_threshold is not None:
123
+ tau = -self.good_threshold if self.maximize else float(self.good_threshold)
124
+ else:
125
+ tau = float(np.sort(finals)[math.ceil(self.good_quantile * n) - 1])
126
+ P = np.asarray(self.predictor(Y), dtype=float)
127
+ if P.shape != Y.shape:
128
+ raise ValueError(f"predictor returned shape {P.shape}, expected {Y.shape}")
129
+
130
+ last_eligible = T - 2 # a kill at the final step saves nothing
131
+ j0 = self.min_peek
132
+ if self.hold_until_resolved is not None:
133
+ rho = _spearman_cols(P[:, : T - 1], finals)
134
+ self.resolution_ = rho
135
+ hit = np.flatnonzero(rho[self.min_peek:] >= self.hold_until_resolved)
136
+ j0 = self.min_peek + int(hit[0]) if hit.size else T - 1
137
+
138
+ good = finals <= tau
139
+ if j0 <= last_eligible:
140
+ scores = np.where(good, P[:, j0: last_eligible + 1].max(axis=1) - tau, -np.inf)
141
+ else:
142
+ scores = np.full(n, -np.inf)
143
+ lam = float(np.sort(scores)[order - 1])
144
+ if lam == -np.inf and j0 <= last_eligible:
145
+ warnings.warn(
146
+ f"only {int(good.sum())} good calibration runs: at alpha={self.alpha} no finite "
147
+ "threshold certifies, so the rule kills every run at its first eligible step. "
148
+ "Raise good_quantile or add calibration runs.", RuntimeWarning, stacklevel=2)
149
+
150
+ self.tau_, self.lambda_, self.first_peek_ = tau, lam, j0
151
+ self.n_steps_, self.n_calibration_, self.n_good_ = T, n, int(good.sum())
152
+ return self
153
+
154
+ # ── deciding ───────────────────────────────────────────────────────────────
155
+ def _check_fitted(self) -> None:
156
+ if self.n_steps_ == 0:
157
+ raise RuntimeError("call fit() first")
158
+
159
+ def should_kill(self, partial_curve: ArrayLike) -> bool:
160
+ """True if a run whose curve so far is `partial_curve` should be killed at its latest step."""
161
+ self._check_fitted()
162
+ y = self._oriented(partial_curve)
163
+ t = y.shape[0] - 1
164
+ if t < self.first_peek_ or t > self.n_steps_ - 2:
165
+ return False
166
+ pred = float(np.asarray(self.predictor(y[None, :]))[0, t])
167
+ return bool(pred - self.lambda_ > self.tau_)
168
+
169
+ def kill_step(self, curve: ArrayLike) -> int | None:
170
+ """Index of the first step at which the rule would kill this curve, or None."""
171
+ self._check_fitted()
172
+ Y = self._oriented(curve)[None, :]
173
+ P = np.asarray(self.predictor(Y), dtype=float)[0]
174
+ lo, hi = self.first_peek_, min(self.n_steps_, Y.shape[1]) - 2
175
+ if lo > hi:
176
+ return None
177
+ fire = np.flatnonzero(P[lo: hi + 1] - self.lambda_ > self.tau_)
178
+ return lo + int(fire[0]) if fire.size else None
179
+
180
+ def evaluate(self, curves: ArrayLike) -> dict[str, Any]:
181
+ """Score the rule on held-out COMPLETED curves: false kills, joint rate, savings."""
182
+ self._check_fitted()
183
+ Y = self._oriented(curves)
184
+ n, T = Y.shape
185
+ P = np.asarray(self.predictor(Y), dtype=float)
186
+ lo, hi = self.first_peek_, T - 2
187
+ if lo <= hi:
188
+ fire = P[:, lo: hi + 1] - self.lambda_ > self.tau_
189
+ killed = fire.any(axis=1)
190
+ step = np.where(killed, lo + np.argmax(fire, axis=1), T - 1)
191
+ else:
192
+ killed, step = np.zeros(n, bool), np.full(n, T - 1)
193
+ good = Y[:, -1] <= self.tau_
194
+ fk = int((killed & good).sum())
195
+ return {
196
+ "n": n, "n_good": int(good.sum()), "n_killed": int(killed.sum()),
197
+ "false_kills": fk, "joint_false_kill_rate": fk / n,
198
+ "conditional_false_kill_rate": fk / int(good.sum()) if good.any() else None,
199
+ "savings": float(1 - (step + 1).sum() / (T * n)),
200
+ "median_kill_step": float(np.median(step[killed])) if killed.any() else None,
201
+ }
@@ -0,0 +1,99 @@
1
+ """Optuna integration: a pruner that runs in shadow mode first, then kills with a certified bound.
2
+
3
+ pruner = ConformalPruner(n_calibration=60, alpha=0.05)
4
+ study = optuna.create_study(sampler=optuna.samplers.RandomSampler(), pruner=pruner)
5
+
6
+ **How it behaves.**
7
+ - **Shadow mode first.** The first `n_calibration` completed trials, by trial number, run
8
+ unpruned. Their reported curves calibrate the rule once.
9
+ - **Then frozen.** The rule never recalibrates on later trials: those are survivors of the
10
+ rule's own kills, and fitting on them would bias it.
11
+ - **The final outcome is the last reported intermediate value,** not `trial.value`.
12
+ - **Every calibration trial must report on the same step grid.**
13
+
14
+ **The guarantee needs the later trials to be exchangeable with the calibration trials.** With
15
+ `RandomSampler`, `QMCSampler` or a shuffled `GridSampler` they are. With `TPESampler`, `GPSampler`
16
+ or `CmaEsSampler` they are not, so the pruner warns once and the bound should not be relied on.
17
+ """
18
+
19
+ from __future__ import annotations
20
+
21
+ import warnings
22
+ from typing import Any
23
+
24
+ import numpy as np
25
+ import optuna
26
+ from optuna.pruners import BasePruner
27
+ from optuna.study import StudyDirection
28
+ from optuna.trial import TrialState
29
+
30
+ from .core import KillRule, Predictor, last_value
31
+
32
+ ArrayLike = Any
33
+
34
+ _EXCHANGEABLE = (optuna.samplers.RandomSampler, optuna.samplers.QMCSampler,
35
+ optuna.samplers.GridSampler, optuna.samplers.BruteForceSampler)
36
+
37
+
38
+ class ConformalPruner(BasePruner):
39
+ def __init__(self, n_calibration: int = 60, alpha: float = 0.05, *,
40
+ good_threshold: float | None = None, good_quantile: float = 0.20,
41
+ min_peek: int = 0, hold_until_resolved: float | None = None,
42
+ predictor: Predictor = last_value) -> None:
43
+ self.n_calibration = n_calibration
44
+ self._kw: dict[str, Any] = dict(alpha=alpha, good_threshold=good_threshold, good_quantile=good_quantile,
45
+ min_peek=min_peek, hold_until_resolved=hold_until_resolved,
46
+ predictor=predictor)
47
+ self.rule_: KillRule | None = None
48
+ self.steps_: list[int] | None = None
49
+ self.calibration_trials_: list[int] = []
50
+ self._warned = False
51
+
52
+ @classmethod
53
+ def from_curves(cls, curves: ArrayLike, steps: Any, *, maximize: bool = False,
54
+ **kw: Any) -> "ConformalPruner":
55
+ """A pruner already calibrated offline, e.g. on a previous sweep's completed curves."""
56
+ p = cls(n_calibration=len(curves), **kw)
57
+ p.rule_ = KillRule(maximize=maximize, **p._kw).fit(curves)
58
+ p.steps_ = list(steps)
59
+ return p
60
+
61
+ def _calibrate(self, study: optuna.study.Study) -> bool:
62
+ done = sorted((t for t in study.get_trials(deepcopy=False, states=(TrialState.COMPLETE,))
63
+ if t.intermediate_values), key=lambda t: t.number)[: self.n_calibration]
64
+ if len(done) < self.n_calibration:
65
+ return False
66
+ grid = sorted(done[0].intermediate_values)
67
+ rows = [[t.intermediate_values[s] for s in grid] for t in done
68
+ if sorted(t.intermediate_values) == grid]
69
+ if len(rows) < len(done):
70
+ warnings.warn(f"{len(done) - len(rows)} calibration trials reported a different step "
71
+ "grid and were skipped", RuntimeWarning, stacklevel=3)
72
+ return False
73
+ maximize = study.direction == StudyDirection.MAXIMIZE
74
+ self.rule_ = KillRule(maximize=maximize, **self._kw).fit(np.asarray(rows))
75
+ self.steps_ = grid
76
+ self.calibration_trials_ = [t.number for t in done]
77
+ return True
78
+
79
+ def prune(self, study: optuna.study.Study, trial: optuna.trial.FrozenTrial) -> bool:
80
+ if not self._warned and not isinstance(study.sampler, _EXCHANGEABLE):
81
+ warnings.warn(f"{type(study.sampler).__name__} adapts to earlier trials, so later trials "
82
+ "are not exchangeable with the calibration trials and the false-kill "
83
+ "bound does not hold. Use RandomSampler or QMCSampler for a certified "
84
+ "bound.", RuntimeWarning, stacklevel=2)
85
+ self._warned = True
86
+ step = trial.last_step
87
+ if step is None:
88
+ return False
89
+ if self.rule_ is None and not self._calibrate(study):
90
+ return False
91
+ assert self.rule_ is not None and self.steps_ is not None
92
+ if trial.number in self.calibration_trials_ or step not in self.steps_:
93
+ return False
94
+ idx = self.steps_.index(step)
95
+ values = trial.intermediate_values
96
+ prefix = [values.get(s) for s in self.steps_[: idx + 1]]
97
+ if any(v is None for v in prefix):
98
+ return False
99
+ return self.rule_.should_kill(np.asarray(prefix, dtype=float))
@@ -0,0 +1,173 @@
1
+ Metadata-Version: 2.4
2
+ Name: surestop
3
+ Version: 0.2.0
4
+ Summary: Stop bad runs early, with a conformal bound on how often you stop one that would have ended well
5
+ Author: Adam Noonan
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/ACNoonan/surestop
8
+ Project-URL: Issues, https://github.com/ACNoonan/surestop/issues
9
+ Keywords: hyperparameter-optimization,early-stopping,pruning,conformal,optuna
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Science/Research
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
14
+ Requires-Python: >=3.10
15
+ Description-Content-Type: text/markdown
16
+ License-File: LICENSE
17
+ Requires-Dist: numpy>=1.24
18
+ Requires-Dist: scipy>=1.10
19
+ Provides-Extra: optuna
20
+ Requires-Dist: optuna>=4; extra == "optuna"
21
+ Provides-Extra: test
22
+ Requires-Dist: pytest; extra == "test"
23
+ Requires-Dist: optuna>=4; extra == "test"
24
+ Dynamic: license-file
25
+
26
+ # SureStop
27
+
28
+ **Stop bad training runs early, with a stated bound on how often you stop one that would have
29
+ ended well.** You pick α. The rule kills runs so that the expected fraction of runs that were
30
+ good *and* got killed stays at or below α.
31
+
32
+ Every pruner in common use (median, percentile, successive halving, Hyperband) is a heuristic,
33
+ and none says how often it kills a run that would have finished well. This one does, under a
34
+ condition you can check: the runs you calibrate on and the runs you judge must come from the same
35
+ random process.
36
+
37
+ It was measured on 400 real LoRA fine-tuning runs under a pre-registered protocol. The curves are
38
+ in this repository, and a script checks that the package reproduces the recorded numbers exactly.
39
+
40
+ ## Install
41
+
42
+ ```
43
+ pip install surestop # numpy + scipy
44
+ pip install "surestop[optuna]" # adds the Optuna pruner
45
+ ```
46
+
47
+ ## Use
48
+
49
+ Plain numpy, on any training loop:
50
+
51
+ ```python
52
+ from surestop import KillRule
53
+
54
+ # completed_curves: array of shape (runs, evaluations), lower is better
55
+ rule = KillRule(alpha=0.05, good_threshold=2.16, min_peek=17).fit(completed_curves)
56
+
57
+ rule.should_kill(curve_so_far) # True -> stop this run now
58
+ rule.kill_step(full_curve) # where it would have fired, or None
59
+ rule.evaluate(held_out_curves) # false kills, joint rate, savings
60
+ ```
61
+
62
+ With Optuna, the first `n_calibration` trials run unpruned to calibrate, then the rule freezes:
63
+
64
+ ```python
65
+ import optuna
66
+ from surestop import ConformalPruner
67
+
68
+ study = optuna.create_study(sampler=optuna.samplers.RandomSampler(),
69
+ pruner=ConformalPruner(n_calibration=60, alpha=0.05))
70
+ ```
71
+
72
+ Pass `maximize=True` (or a `direction="maximize"` study) when higher is better.
73
+
74
+ ## The rule
75
+
76
+ A run is *good* if its final value ends at or below a threshold τ. At every evaluation, a
77
+ predictor forecasts the final value from the curve so far. The default predictor is the latest
78
+ observed value.
79
+
80
+ To calibrate, take n runs that have already finished. For each good one, record its worst moment:
81
+ the largest gap between forecast and τ at any evaluation. Runs that are not good contribute
82
+ nothing. Set λ to the ⌈(1−α)(n+1)⌉-th smallest of these scores. A new run is killed at the first
83
+ evaluation where its forecast exceeds τ + λ.
84
+
85
+ This is conformal risk control on the loss "good and killed". Taking the worst moment across all
86
+ evaluations means one threshold covers every look at a run, so there is no correction for
87
+ multiple looks. The guarantee needs the calibration runs and the new runs to be exchangeable.
88
+ Random sampling gives that. Adaptive samplers such as TPE do not.
89
+
90
+ ## What was measured
91
+
92
+ Experiment CC-03c, pre-registered before any run finished:
93
+
94
+ - **Runs:** 400 LoRA fine-tunes of Qwen2.5-0.5B on dolly-15k, 750 steps each, validation loss
95
+ measured 50 times per run.
96
+ - **Configs:** i.i.d. draws from a frozen hyperparameter prior (`docs/hyperparameter-prior.md`).
97
+ - **τ:** pinned from an earlier pool before the first run landed.
98
+ - **Splits:** two, one random and one past-to-future by night, each about half calibration and
99
+ half test.
100
+
101
+ The pre-registered rule used a fitted power-law predictor. On both splits it kept the bound and
102
+ saved about half the evaluation compute:
103
+
104
+ | split | false kills | joint rate (95% CI) | eval compute saved |
105
+ |---|---|---|---|
106
+ | random | 13 of 200 | 6.5% (3.5–10.9%) | 52.2% |
107
+ | past-to-future | 11 of 203 | 5.4% (2.7–9.5%) | 52.8% |
108
+
109
+ A rule that kills every run immediately was scored by the same gate and failed on both splits.
110
+
111
+ **The flaw that run found:** every kill fired at the 3rd of 50 evaluations. Early rankings only
112
+ match final rankings from about evaluation 18. The rule stayed within budget because the bound
113
+ held. The fitted power law ranked runs worse than the latest observed value at every early
114
+ evaluation.
115
+
116
+ **What this package ships instead,** and what the reproduction script checks:
117
+
118
+ | predictor | hold until | random split | past-to-future split | 2000 resplits |
119
+ |---|---|---|---|---|
120
+ | last value | none | 18 of 200 (9.0%), 78% saved | 7 of 203 (3.4%), 75% saved | joint 5.0%, breach 8.5% |
121
+ | last value | eval 18 | 18 of 200 (9.0%), 57% saved | 7 of 203 (3.4%), 56% saved | joint 5.0%, breach 8.6% |
122
+
123
+ The random split's 18 of 200 is a breach. It sits at the 95th percentile of that rule's own
124
+ resplit distribution, whose mean is 5.0%, so it reads as a split tail rather than a leak. A
125
+ correct rule at this sample size lands on a breach in about 7% of splits. The held version was
126
+ chosen *after* seeing these 400 runs, so its numbers are a hypothesis. It halves
127
+ how often the single best run is killed at the same budget. Its confirmation needs a fresh pool.
128
+
129
+ Full account: `docs/writeup.md`. Run the check yourself:
130
+
131
+ ```
132
+ python reproduce/reproduce_cc03c.py
133
+ ```
134
+
135
+ ## What the bound does not cover
136
+
137
+ - **It needs exchangeable runs.** With a random or quasi-random sampler the bound holds. With
138
+ TPE, Gaussian-process or CMA-ES samplers it does not. The Optuna pruner warns once and carries
139
+ on, but do not rely on the bound there.
140
+ - **It bounds the joint rate rather than the share of good runs killed.** When good runs are rare, a large
141
+ share of them can still die: about a third in the measurement above, mostly borderline ones.
142
+ Lower α if that matters.
143
+ - **An estimated τ gives a measured bound rather than a proven one.** `good_quantile` estimates τ from the
144
+ calibration runs. Simulations put the realized rate at 0.88 to 0.97 times α. Pass a real
145
+ `good_threshold` when you have one.
146
+ - **α = 0.05 needs at least 19 calibration runs**, and the certified threshold gets loose with
147
+ few good runs. Sweeps of a few dozen trials are too small for this.
148
+ - **One setting only so far.** One model, one dataset, one horizon.
149
+
150
+ ## Layout
151
+
152
+ | path | what |
153
+ |---|---|
154
+ | `src/surestop/` | `KillRule` (numpy/scipy) and `ConformalPruner` (Optuna) |
155
+ | `tests/` | the bound holds on synthetic exchangeable runs, and the same assertion fails on a rule that leaks 2× its budget |
156
+ | `reproduce/` | the 400 curves, the pre-registered splits and targets, the script, and its committed output |
157
+ | `docs/` | the write-up, the prior-art sweep, and the frozen hyperparameter prior the curves were drawn from |
158
+ | `optunahub/` | the package as an OptunaHub registry entry |
159
+
160
+ ## Related work
161
+
162
+ Calibrating a stop threshold on a predictor with a conformal guarantee is an established pattern. Xie et al.
163
+ (arXiv:2602.13935) use the same device, a threshold on the maximum of a running score calibrated
164
+ on one class only, to stop LLM reasoning traces. Related work does the same for mixed-integer
165
+ solvers and LLM agent episodes. As far as a search found, nobody has applied it to pruning
166
+ training runs, and no major tuning library ships a pruner with a stated error guarantee.
167
+ `docs/prior-art.md` has the sweep.
168
+
169
+ ## Citing
170
+
171
+ Concept DOI: [10.5281/zenodo.22726440](https://doi.org/10.5281/zenodo.22726440), which always resolves
172
+ to the newest version. `CITATION.cff` has the full entry. The curves are a measurement, released
173
+ under the same MIT licence as the code.
@@ -0,0 +1,13 @@
1
+ LICENSE
2
+ README.md
3
+ pyproject.toml
4
+ src/surestop/__init__.py
5
+ src/surestop/core.py
6
+ src/surestop/optuna_pruner.py
7
+ src/surestop.egg-info/PKG-INFO
8
+ src/surestop.egg-info/SOURCES.txt
9
+ src/surestop.egg-info/dependency_links.txt
10
+ src/surestop.egg-info/requires.txt
11
+ src/surestop.egg-info/top_level.txt
12
+ tests/test_core.py
13
+ tests/test_optuna_pruner.py
@@ -0,0 +1,9 @@
1
+ numpy>=1.24
2
+ scipy>=1.10
3
+
4
+ [optuna]
5
+ optuna>=4
6
+
7
+ [test]
8
+ pytest
9
+ optuna>=4
@@ -0,0 +1 @@
1
+ surestop
@@ -0,0 +1,88 @@
1
+ """The bound, a negative control that proves the bound check can fail, and the rule's edges."""
2
+
3
+ import math
4
+
5
+ import numpy as np
6
+ import pytest
7
+
8
+ from surestop import KillRule, NotEnoughCalibration
9
+
10
+ T = 30
11
+
12
+
13
+ def synth(n, rng):
14
+ """Exchangeable learning curves: a final value, plus a decaying offset and noise that shrink with step."""
15
+ finals = rng.normal(2.17, 0.013, n)
16
+ t = np.arange(T)
17
+ offset = rng.uniform(0.05, 0.3, n)[:, None] * np.exp(-t / 6.0)
18
+ noise = rng.normal(0, 1, (n, T)) * 0.012 * np.exp(-t / 10.0)
19
+ Y = finals[:, None] + offset + noise
20
+ Y[:, -1] = finals
21
+ return Y
22
+
23
+
24
+ def mean_joint_rate(alpha_fit, reps=400, n_cal=200, n_test=200, seed=0, **kw):
25
+ rng = np.random.default_rng(seed)
26
+ rates = []
27
+ for _ in range(reps):
28
+ cal, test = synth(n_cal, rng), synth(n_test, rng)
29
+ rule = KillRule(alpha=alpha_fit, good_threshold=2.159, **kw).fit(cal)
30
+ rates.append(rule.evaluate(test)["joint_false_kill_rate"])
31
+ rates = np.array(rates)
32
+ return rates.mean(), rates.std(ddof=1) / math.sqrt(reps)
33
+
34
+
35
+ def test_joint_false_kill_rate_is_bounded():
36
+ mean, se = mean_joint_rate(0.05)
37
+ assert mean <= 0.05 + 3 * se, (mean, se)
38
+
39
+
40
+ def test_negative_control_the_bound_check_can_fail():
41
+ # A rule that leaks 2x its budget: certify at 0.10, hold the result to 0.05. The assertion
42
+ # the real test makes must fail here. (At 0.25 the synthetic pool has too few good runs for a
43
+ # finite threshold, so the rule kills everything — too easy a control to prove anything.)
44
+ mean, se = mean_joint_rate(0.10, reps=400)
45
+ assert not (mean <= 0.05 + 3 * se), (mean, se)
46
+
47
+
48
+ def test_the_rule_actually_saves_compute():
49
+ rng = np.random.default_rng(1)
50
+ rule = KillRule(alpha=0.05, good_threshold=2.159).fit(synth(300, rng))
51
+ assert rule.evaluate(synth(300, rng))["savings"] > 0.2
52
+
53
+
54
+ def test_min_peek_blocks_earlier_kills():
55
+ rng = np.random.default_rng(2)
56
+ rule = KillRule(alpha=0.05, good_threshold=2.159, min_peek=12).fit(synth(300, rng))
57
+ steps = [rule.kill_step(c) for c in synth(300, rng)]
58
+ assert all(s is None or s >= 12 for s in steps)
59
+ assert any(s is not None for s in steps)
60
+
61
+
62
+ def test_should_kill_agrees_with_kill_step():
63
+ rng = np.random.default_rng(3)
64
+ rule = KillRule(alpha=0.05, good_threshold=2.159, min_peek=4).fit(synth(300, rng))
65
+ for c in synth(50, rng):
66
+ first = next((t for t in range(T) if rule.should_kill(c[: t + 1])), None)
67
+ assert first == rule.kill_step(c)
68
+
69
+
70
+ def test_hold_until_resolved_moves_the_first_peek_later():
71
+ rng = np.random.default_rng(4)
72
+ cal = synth(300, rng)
73
+ held = KillRule(alpha=0.05, good_threshold=2.159, hold_until_resolved=0.9).fit(cal)
74
+ assert held.first_peek_ > 0
75
+ assert held.resolution_[held.first_peek_] >= 0.9
76
+
77
+
78
+ def test_too_few_calibration_runs_raises():
79
+ with pytest.raises(NotEnoughCalibration):
80
+ KillRule(alpha=0.05).fit(synth(18, np.random.default_rng(5)))
81
+
82
+
83
+ def test_maximize_mirrors_minimize():
84
+ rng = np.random.default_rng(6)
85
+ cal, test = synth(200, rng), synth(100, rng)
86
+ lo = KillRule(alpha=0.05, good_threshold=2.159).fit(cal)
87
+ hi = KillRule(alpha=0.05, good_threshold=-2.159, maximize=True).fit(-cal)
88
+ assert [lo.kill_step(c) for c in test] == [hi.kill_step(-c) for c in test]
@@ -0,0 +1,48 @@
1
+ """The Optuna pruner: shadow mode first, pruning after, a warning for adaptive samplers."""
2
+
3
+ import numpy as np
4
+ import pytest
5
+
6
+ optuna = pytest.importorskip("optuna")
7
+ from surestop import ConformalPruner # noqa: E402
8
+
9
+ optuna.logging.set_verbosity(optuna.logging.WARNING)
10
+ T = 20
11
+
12
+
13
+ def objective(trial):
14
+ x = trial.suggest_float("x", 0.0, 1.0)
15
+ rng = np.random.default_rng(trial.number)
16
+ final = 2.15 + 0.05 * x + rng.normal(0, 0.005)
17
+ for step in range(T):
18
+ value = final + 0.2 * np.exp(-step / 4) + rng.normal(0, 0.004) * np.exp(-step / 8)
19
+ if step == T - 1:
20
+ value = final
21
+ trial.report(float(value), step)
22
+ if trial.should_prune():
23
+ raise optuna.TrialPruned()
24
+ return final
25
+
26
+
27
+ def test_calibrates_in_shadow_mode_then_prunes():
28
+ pruner = ConformalPruner(n_calibration=40, alpha=0.05)
29
+ study = optuna.create_study(sampler=optuna.samplers.RandomSampler(seed=0), pruner=pruner)
30
+ study.optimize(objective, n_trials=160)
31
+ states = [t.state for t in study.trials]
32
+ assert all(s == optuna.trial.TrialState.COMPLETE for s in states[:40])
33
+ assert pruner.rule_ is not None and np.isfinite(pruner.rule_.lambda_)
34
+ assert sum(s == optuna.trial.TrialState.PRUNED for s in states) > 0
35
+
36
+
37
+ def test_warns_for_an_adaptive_sampler():
38
+ pruner = ConformalPruner(n_calibration=20)
39
+ study = optuna.create_study(sampler=optuna.samplers.TPESampler(seed=0), pruner=pruner)
40
+ with pytest.warns(RuntimeWarning, match="not exchangeable"):
41
+ study.optimize(objective, n_trials=2)
42
+
43
+
44
+ def test_from_curves_is_calibrated_before_the_first_trial():
45
+ rng = np.random.default_rng(0)
46
+ curves = 2.17 + rng.normal(0, 0.01, (60, 1)) + 0.2 * np.exp(-np.arange(T) / 4)
47
+ pruner = ConformalPruner.from_curves(curves, steps=range(T))
48
+ assert pruner.rule_.n_calibration_ == 60 and pruner.steps_ == list(range(T))