Bugfix/optuna

[anay-rfai]

Changes

Adapted interim-vs-interim Optuna pruning. Vanilla Optuna compares a running trial's intermediate values against completed trials' intermediate values (interim-vs-complete), which is unusable in RapidFire's sharded pipeline because no trial is "complete" until the entire sweep finishes. This PR rewires the pruner to compare each running trial against the median of its simultaneously-running peers' intermediate values (interim-vs-interim). Because roughly half the population sits below the median of the rest by construction at every evaluation, the adapted MedianPruner behaves like hyperband / successive halving: triggered on shard completion, the population halves (and is replaced if budget allows) at each successive shard. This is what produces the 7-of-8 and 5-of-6 prune rates shown in the tutorial screenshots. Pruning remains a heuristic compute-savings mechanism — it does not change Optuna's (already absent) guarantee that study.best_trial is the global optimum, vanilla or adapted.

Pruning comparison anchored on cumulative shards completed. The peer-median comparison is anchored on the cumulative number of shards completed, for both fit and evals modes (cumulative shard count in fit, raw shard id in evals), so trials are always compared at the same progress point regardless of when they started. A per-prune-check trace line ([RFOptuna prune-check] trial=… step=… dir=… current=… peers=[…] median=… -> PRUNE/continue (reason)) is emitted so the runtime ordering / peer-availability gap — the fastest pipeline reaching each shard first and finding no peers (no_peers_at_step) — is visible in the dashboard/logs, alongside current_is_nan, worse_than_median, and better_than_median.

Front-loaded batches for fit-mode shards. In fit mode, batches are now front-loaded for each shard of data so that pruning comparisons are fair between configs that have different effective batch sizes. Without front-loading, a config with a larger effective batch would have processed more samples at the same wall-clock checkpoint and look "better" purely from batching, biasing the peer-median comparison. Front-loading puts every config on the same samples-processed footing at each shard boundary before the prune decision runs.

Prune → replacement flow and the build_all_indexes flag (evals mode). When a trial is pruned, Optuna is asked for a replacement trial. A replacement config needs its RAG index to already exist, which is governed by the new build_all_indexes flag on RFOptuna (evals mode only; ignored in fit mode):

• build_all_indexes=True (default): every RAG index the search space can reach is built up front, so any replacement Optuna suggests during the run already has its retriever available. The cost is that indexes Optuna never visits are built too — the same index count RFGridSearch would build for the equivalent space.
• build_all_indexes=False: only the indexes needed by the n_initial initial configs are built. A replacement suggestion that needs an unbuilt index is rejected and resampled — cheaper up front, but it narrows the space Optuna can actually explore. A rejected candidate is told FAIL (not COMPLETE), which keeps it out of best_trial and out of the sampler's model so TPE is not taught that this region scored anything; only accepted suggestions consume budget.

Index-affecting Range knobs are discretised either way (each drawn down to Range.sample_n distinct values by sample(n)), since the set of indexes must be finite and known before any query runs; the constructor seed selects which part of each range this run explores.

Search/sampler knobs. granularity ("chunk" default, or "epoch") controls when pruning is evaluated in fit mode (ignored in evals). A single seed (default 42) governs the algorithm's own stochastic state — every Range generator, the global RNG used by List.sample() / fallback draws, and the Optuna study sampler — so RFOptuna() is reproducible out of the box; the run-level seed passed to run_evals / run_fit is ignored for the algorithm's draws and only governs surrounding infrastructure (dataset sharding, etc.).

Fixes. Fixes three tracked Optuna issues — RF-OPT-1, RF-OPT-2, and RF-OPT-4 — and an online-aggregation bug in rapidfireai/evals/metrics/online_strategies.py (with corresponding controller updates in rapidfireai/evals/scheduling/controller.py).

Tutorials. Adds a global seed variable across all tutorial notebooks for reproducibility, adds a new Optuna RAG/scifact tutorial notebook (tutorial_notebooks/rag-contexteng/rf-tutorial-optuna-rag-scifact.ipynb), and beefs up the existing Optuna tutorial notebooks with richer logging and walkthrough content.

No breaking changes to public APIs.

Changelog Content

Additions

• New Optuna RAG/scifact tutorial notebook (tutorial_notebooks/rag-contexteng/rf-tutorial-optuna-rag-scifact.ipynb).
• Beefed-up Optuna tutorial notebooks with richer logging and walkthrough content.
• Global seed variable for all tutorial notebooks for reproducibility.
• build_all_indexes flag on RFOptuna (evals mode) to control whether every reachable RAG index is built up front (True, default) or only the n_initial configs' indexes (False, with missing-index suggestions rejected and resampled).

Changes

• Optuna pruning is now interim-vs-interim: each running trial is compared against the median of its simultaneously-running peers (hyperband / successive-halving style), instead of vanilla Optuna's interim-vs-complete comparison.
• Pruning comparison is anchored on cumulative shards completed, for both fit and evals modes.
• Fit mode now front-loads batches for each shard of data so pruning comparisons are fair between configs with different effective batch sizes.

Fixes

• Fixed RF-OPT-1, RF-OPT-2, RF-OPT-4.
• Fixed online aggregation bug in rapidfireai/evals/metrics/online_strategies.py.

Testing

• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have tested this change manually
• I have tested this change in the following environments:
  • Local development
  • Docker environment
  • Other: GCE

Screenshots (if applicable)

Add screenshots to help explain your changes.

Checklist

• My code follows the project's style guidelines
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• Any dependent changes have been merged and published

Performance Impact

If this PR affects performance, describe the impact and any optimizations made.

Related Issues

Fixes #303 #302 #301 #300

---

Note

High Risk
Changes touch eval orchestration, RAG context lifecycle, Optuna trial sampling/pruning, and fit shard batching—misaligned hashes or coverage could break pipelines at runtime despite extensive new tests.

Overview
RFOptuna now prunes by comparing each trial’s intermediate metric to the median of concurrently running peers (not Optuna’s completed-trial median), with per-check [RFOptuna prune-check] logging and dashboard wiring via set_logger. Fit callbacks report one value per chunk at a cumulative chunks-completed step; pruner=None skips pruning via NopPruner. Evals replacements can be rejected when their RAG index was never built (set_context_feasibility, FAIL trials, build_all_indexes default True), with get_context_coverage_leaves enumerating index-affecting combos before get_runs and the evals controller pre-building those contexts.

Hyperparameter plumbing: Range is a seeded pure sampler (sample(n), sample_n, set_seed); RFOptuna owns a range value cache so coverage enumeration and Optuna suggest share the same discretized values for index-affecting paths (retrieval-only search_cfg / reranker_cfg stay continuous). Nested List of configs registers conditional Optuna params (api_config[idx].…); unreachable Range/List in literals raises at sampling time. RFGridSearch errors on Range; RFRandomSearch / RFOptuna use constructor seed (run-level seed ignored for algorithm draws). Legacy create_model_fn removed from AutoMLAlgorithm.

Evals controller centralizes context hashing (combine_context_hash, leaf_context_hash) and fails registration/launch when a pipeline’s context was not built. Fit chunking assigns extra batches to the first chunks (fairer for Optuna batch-size comparisons). Online metrics clamp means and guard sqrt on negative variance.

^{Reviewed by Cursor Bugbot for commit 0a84538. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit d72ee1c. Configure here.}

[chethan-rfai]

LGTM

[arun-rfai]

LGTM