treewalker-exp: one package for data preparation, with manifests and new workloads

.gitignore

@@ -1,15 +1,16 @@
 /target
-/benchmarks/target
+/experiments/benchmarks/target
 /.venv
 __pycache__/
 .DS_Store
 
 # Generated by the scripts
-paper/experiments/artifacts/
-paper/experiments/figures/
-paper/experiments/data/expedia.parquet
-paper/experiments/data/*_summary.md
-paper/experiments/data/*.log
+experiments/artifacts/
+experiments/figures/
+experiments/data/expedia.parquet
+experiments/data/raw/
+experiments/data/*_summary.md
+experiments/data/*.log
 
 # Terraform state and caches
 infra/.terraform/

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace]
-members = ["benchmarks"]
+members = ["experiments/benchmarks"]
 resolver = "3"
 
 [workspace.package]
@@ -46,7 +46,7 @@ rust-version.workspace = true
 description = "Grouped tree-ensemble inference via recursive partial evaluation"
 license.workspace = true
 repository.workspace = true
-include = ["src/**", "tests/**", "docs/treelite-loading.md", "docs/treelite-loading-validation.md", "Cargo.toml", "README.md", "CHANGELOG.md", "LICENSE"]
+include = ["/src/**", "/tests/**", "/docs/treelite-loading.md", "/docs/treelite-loading-validation.md", "/Cargo.toml", "/README.md", "/CHANGELOG.md", "/LICENSE"]
 
 [features]
 research = []  # ablations, work counters, prediction stages, full walk, introspection

README.md

@@ -23,9 +23,11 @@ found [under the neurips2026 tag](https://github.com/Shopify/treewalker/releases
 | Path | Contents |
 |---|---|
 | `src/` | inference library and treelite model parser |
-| `benchmarks/` | separate unpublished crate: benchmark harness, `sweep_bench`, artifact correctness tests |
-| `paper/experiments/scripts/` | dataset preparation, baseline compilation, figures, tables |
-| `paper/experiments/data/` | released result CSVs and machine descriptions |
+| `experiments/benchmarks/` | separate unpublished crate: benchmark harness, `sweep_bench`, artifact correctness tests |
+| `experiments/treewalker_exp/` | `treewalker-exp`: dataset preparation, workloads, training, compiled baselines |
+| `experiments/grids.toml` | the benchmark suites, resolved by `treewalker-exp` into execution manifests |
+| `experiments/scripts/` | figures, tables and the chunked-G run, on the released CSVs |
+| `experiments/data/` | released result CSVs and machine descriptions |
 | `infra/` | Terraform and VM startup script for the two GCE benchmark machines |
 
 ## Building and packaging the library
@@ -54,7 +56,7 @@ manifest and the shared `Cargo.lock`. Build it explicitly; `--target-dir target`
 preserves the binary path used by the experiment scripts:
 
 ```bash
-cargo build --manifest-path benchmarks/Cargo.toml --target-dir target \
+cargo build --manifest-path experiments/benchmarks/Cargo.toml --target-dir target \
   --release --features external-bench
 ```
 
@@ -99,21 +101,23 @@ output and precision policy, errors, limits, and the caller's grouping contract.
 
 ## Datasets
 
-The scripts download the public datasets; none are redistributed here.
+None of the datasets are redistributed here. `uv run treewalker-exp fetch`
+downloads the public ones into `experiments/data/raw/` and checks each against
+a pinned SHA-256; `prepare` fetches them on demand.
 
 | Dataset | Source | Fetched by |
 |---|---|---|
-| SUPPORT | DeepSurv repository (`support_train_test.h5`) | `prepare.py` |
-| FLCHAIN | R `survival` package via Rdatasets | `prepare.py` |
-| Expedia (ICDM 2013) | Kaggle competition `expedia-personalized-sort` | `fetch_expedia.py`, from the `data.zip` you download |
-| UCI Default of Credit Card Clients | OpenML 42477, CC BY 4.0, DOI 10.24432/C55S3H | `prepare_scenario.py` |
+| SUPPORT | DeepSurv repository (`support_train_test.h5`) | `treewalker-exp fetch` |
+| FLCHAIN | R `survival` package via Rdatasets | `treewalker-exp fetch` |
+| Expedia (ICDM 2013) | Kaggle competition `expedia-personalized-sort` | `treewalker-exp fetch-expedia`, from the `data.zip` you download |
+| UCI Default of Credit Card Clients | OpenML 42477, CC BY 4.0, DOI 10.24432/C55S3H | `treewalker-exp fetch` |
 
 The Expedia data may not be redistributed. Accept the competition rules on
 Kaggle, download `data.zip` (for example
 `kaggle competitions download -c expedia-personalized-sort -f data.zip`), then
-run `fetch_expedia.py --train-csv data.zip`; extracting the zip needs Info-ZIP
+run `treewalker-exp fetch-expedia --train-csv data.zip`; extracting the zip needs Info-ZIP
 `unzip`, because it uses Deflate64. It writes
-`paper/experiments/data/expedia.parquet` and checks its content fingerprint
+`experiments/data/expedia.parquet` and checks its content fingerprint
 against the file the paper used.
 
 ## Released results
@@ -138,14 +142,14 @@ benchmark machines. From the repository root:
 
 ```bash
 uv sync
-uv run python3 paper/experiments/scripts/plot.py                    # Figures 3, 4, 6-9
-uv run python3 paper/experiments/scripts/gen_heatmap_tex.py         # Figure 2
-uv run python3 paper/experiments/scripts/decomposition_validation.py  # Table 1, Figure 5
-uv run python3 paper/experiments/scripts/summarize_scenario.py --arch arm   # Tables 3, 10 (also --arch intel)
-uv run python3 paper/experiments/scripts/paper_numbers.py           # every in-text number
+uv run python3 experiments/scripts/plot.py                    # Figures 3, 4, 6-9
+uv run python3 experiments/scripts/gen_heatmap_tex.py         # Figure 2
+uv run python3 experiments/scripts/decomposition_validation.py  # Table 1, Figure 5
+uv run python3 experiments/scripts/summarize_scenario.py --arch arm   # Tables 3, 10 (also --arch intel)
+uv run python3 experiments/scripts/paper_numbers.py           # every in-text number
 ```
 
-Outputs go to `paper/experiments/figures/`. `paper_numbers.py` prints each
+Outputs go to `experiments/figures/`. `paper_numbers.py` prints each
 number next to the value printed in the paper.
 
 | Paper item | Source |
@@ -168,51 +172,65 @@ LightGBM 4.6.0 and XGBoost 3.2.0 were built from source with `-march=native`,
 TreeWalker with `-C target-cpu=native` (`.cargo/config.toml`). The factorial
 grid used Rust 1.94.1; the scenario and chunked runs used Rust 1.97.1, the
 version now pinned in `rust-toolchain.toml`. `infra/scripts/startup.sh` is the
-full machine recipe.
+full machine recipe; it now builds the versions `uv.lock` installs, LightGBM
+4.7.0 and XGBoost 3.4.1.
 
 ### On GCE with Terraform
 
 ```bash
 cd infra
 terraform init
-terraform apply -var project=YOUR_PROJECT              # factorial grid (bench_suite = "paper")
-terraform apply -var project=YOUR_PROJECT -var bench_suite=rebuttal   # scenario + chunked runs
+terraform apply -var project=YOUR_PROJECT -var git_ref=REF              # factorial grid (bench_suite = "paper")
+terraform apply -var project=YOUR_PROJECT -var git_ref=REF -var bench_suite=rebuttal   # scenario + chunked runs
 ```
 
-The factorial grid needs `paper/experiments/data/expedia.parquet` from
-`fetch_expedia.py`; Terraform uploads it for the trainer VM, which checks its
+The factorial grid needs `experiments/data/expedia.parquet` from
+`treewalker-exp fetch-expedia`; Terraform uploads it for the trainer VM, which checks its
 fingerprint before training. Terraform uploads `git archive` of `var.git_ref`
-(default `neurips2026`) to a new bucket. The Intel VM trains the models and shares them, so both machines
+to a new bucket. The ref has no default: it must contain `experiments/` and
+`treewalker-exp`, which the `neurips2026` tag predates. The Intel VM trains the models and shares them, so both machines
 evaluate identical models. Results land in the bucket under `results/`. The
 scenario and chunked runs take about 25 minutes per VM.
 
 ### By hand on a prepared Linux machine
 
 ```bash
-uv sync
-uv run python3 paper/experiments/scripts/fetch_expedia.py --train-csv PATH/data.zip
-uv run python3 paper/experiments/scripts/prepare.py --grid all --prepare-groups --skip-compiled
-uv run python3 paper/experiments/scripts/prepare.py --grid all --compile-only   # tl2cgen, lleaves, sweep_bench
-cargo test --manifest-path benchmarks/Cargo.toml --target-dir target \
+uv sync --group baselines
+uv run --group baselines treewalker-exp fetch-expedia --train-csv PATH/data.zip
+uv run --group baselines treewalker-exp prepare --suite factorial
+uv run --group baselines treewalker-exp compile-baselines --suite factorial   # tl2cgen, lleaves
+uv run --group baselines treewalker-exp build-bench
+cargo test --manifest-path experiments/benchmarks/Cargo.toml --target-dir target \
   --release --features research
 
-taskset -c 0 ./target/release/sweep_bench paper/experiments/artifacts --grid all \
-  --output-dir paper/experiments/data --warmup 3 --iters 21 --min-iters 11 \
+taskset -c 0 ./target/release/sweep_bench experiments/artifacts --grid all \
+  --output-dir experiments/data --warmup 3 --iters 21 --min-iters 11 \
   --max-time-secs 30 --lgb-lib /usr/local/lib/lib_lightgbm.so --xgb-lib PATH/libxgboost.so
 
-uv run python3 paper/experiments/scripts/prepare_scenario.py
-uv run python3 paper/experiments/scripts/prepare_chunked.py --stage prepare
-taskset -c 0 ./target/release/sweep_bench paper/experiments/artifacts --grid scen \
-  --output-dir paper/experiments/data --warmup 3 --iters 21 --min-iters 11 --max-time-secs 10
-taskset -c 0 uv run python3 paper/experiments/scripts/prepare_chunked.py --stage timing
+uv run --group baselines treewalker-exp prepare --suite scenario-v1 --validate
+uv run --group baselines python3 experiments/scripts/prepare_chunked.py --stage prepare
+taskset -c 0 ./target/release/sweep_bench experiments/artifacts --grid scen \
+  --output-dir experiments/data --warmup 3 --iters 21 --min-iters 11 --max-time-secs 10
+taskset -c 0 uv run --group baselines python3 experiments/scripts/prepare_chunked.py --stage timing
 ```
 
-The compiled baselines are installed separately, as in `startup.sh`:
-`uv pip install 'tl2cgen>=1.0.0'` and
-`uv pip install 'lleaves @ git+https://github.com/siboehm/lleaves.git@v1.4.1'`,
-with LLVM 20 for lleaves. `sweep_bench` loads the native LightGBM and XGBoost
+tl2cgen comes from the locked `baselines` dependency group. It has no aarch64
+Linux wheel, so there it builds from source with
+`CXX=$PWD/infra/scripts/cxx-cstdint`. lleaves runs in
+`experiments/compile.py`'s own script lock, with `llc` and `clang` from LLVM 20
+on `PATH` or given by `--llc` and `--clang`. Every `uv run` passes
+`--group baselines`, because `uv run` removes packages outside the groups it
+syncs. `sweep_bench` loads the native LightGBM and XGBoost
 libraries given by `--lgb-lib` and `--xgb-lib`.
 
+`uv.lock`, `experiments/compile.py.lock` and
+`tests/fixtures/import/generate.py.lock` resolve against PyPI through the
+`pypi` index that the project and both scripts declare, which takes priority
+over any index in a user's uv configuration. Relock with `uv lock`,
+`uv lock --script experiments/compile.py` and
+`uv lock --script tests/fixtures/import/generate.py`; `uv lock --check` (with
+`--script` for a script) verifies a lock without changing it.
+
 ### What to expect from a rerun
 
 A clean-room rerun (September 2026, fresh VMs of the same machine types in
@@ -238,7 +256,7 @@ order, layout and every ablation except `disable_exact_sums` leave it unchanged.
 The margin and the link are ordinary rounded `f64` operations.
 
 LightGBM (f64) predictions match Treelite GTIL within 1e-13, and XGBoost (f32)
-predictions match native XGBoost within 1e-5 (`benchmarks/tests/correctness.rs`).
+predictions match native XGBoost within 1e-5 (`experiments/benchmarks/tests/correctness.rs`).
 GTIL and the full walk add leaves in `f64` in tree order, so they differ from the
 exact sum by their rounding, which grows with the partial sums. Measured on
 2026-10-04, the full walk's outputs differ by at most 2.2e-15 on the FLCHAIN and

docs/treelite-loading-validation.md

@@ -45,7 +45,7 @@ already present in the local Cargo cache.
 | `cargo test --release --offline` | 14 unit, 11 import tests; 1 compiling doctest |
 | `cargo fmt --check` | Passed |
 | `cargo clippy --offline --all-targets -- -D warnings` | Passed |
-| `cargo build --offline --release --manifest-path benchmarks/Cargo.toml` | Passed |
+| `cargo build --offline --release --manifest-path experiments/benchmarks/Cargo.toml` | Passed |
 | Benchmark artifact `correctness` suite, `--features test-helpers` | 18 tests passed |
 | Benchmark artifact `parse_test` suite | 6 tests passed |
 | `cargo package --list --allow-dirty --offline` | Library tests, fixtures, and loading documentation included |

docs/treelite-loading.md

@@ -11,7 +11,7 @@ thresholds. Treelite can emit invalid JSON for nonfinite values; TreeWalker
 rejects those dumps rather than repairing them.
 
 ```python
-import treelite  # fixture/reference version: 4.7.0
+import treelite  # fixture/reference version: 4.7.2
 model = treelite.frontend.load_lightgbm_model("model.txt")
 model.serialize("model.bin")
 # Optional, for debugging: model.dump_as_json()
@@ -68,7 +68,7 @@ output dimensions instead of producing a scalar prediction.
 | Property | Support |
 |---|---|
 | Binary | Little-endian Treelite v4 checkpoints, producer versions 4.0–4.7, no extensions |
-| JSON | Strict Treelite scalar dump schema, as exercised with 4.7.0; no syntax repairs or unknown fields |
+| JSON | Strict Treelite scalar dump schema, as exercised with 4.7.2; no syntax repairs or unknown fields |
 | Tasks | Scalar binary classification, regression, and ranking |
 | Outputs | One target, one class/output slot, scalar leaves, nonempty forest |
 | Tree assignments | Every tree contributes to the sole output; target and class IDs must both be 0 |
@@ -208,8 +208,9 @@ truncated binary fields, size overflow declarations, and unsupported exports.
 
 Optional fixture generation is documented in
 [`tests/fixtures/import/generate.py`](../tests/fixtures/import/generate.py).
-It pins Treelite 4.7.0, NumPy 2.2.6, sklearn 1.6.1, and SciPy 1.15.3; the
-manifest records input dtype and per-fixture absolute/relative tolerances.
+It pins Treelite 4.7.2, NumPy 2.5.3, sklearn 1.9.1, and SciPy 1.18.1, and
+`generate.py.lock` pins their dependencies; the manifest records input dtype
+and per-fixture absolute/relative tolerances.
 Float64 fixtures use 1e-14; float32 fixtures allow 1e-7 for the different
 accumulation precision. Existing artifact gates remain f64 1e-14 and native
 XGBoost f32 1e-5.
@@ -220,5 +221,5 @@ not update the archived paper's benchmark results.
 
 Format references: [Treelite v4 serialization](https://treelite.readthedocs.io/en/latest/serialization/v4.html),
 [postprocessors](https://treelite.readthedocs.io/en/latest/knobs/postprocessor.html),
-and the [pinned 4.7.0 reference evaluator](https://github.com/dmlc/treelite/blob/4.7.0/src/gtil/predict.cc).
+and the [pinned 4.7.2 reference evaluator](https://github.com/dmlc/treelite/blob/4.7.2/src/gtil/predict.cc).
 The `latest` documentation may describe a newer producer version.

benchmarks/Cargo.toml → experiments/benchmarks/Cargo.toml

@@ -13,7 +13,7 @@ quickscorer-bench = ["dep:quickscorer"]  # QuickScorer baseline in the legacy CL
 external-bench = ["dep:libloading", "dep:quickscorer"]  # C FFI baselines + QuickScorer
 
 [dependencies]
-treewalker-gbdt = { path = "..", features = ["research"] }
+treewalker-gbdt = { path = "../..", features = ["research"] }
 serde.workspace = true
 simd-json.workspace = true
 quickscorer = { version = "=0.1.0", git = "https://github.com/leomanne/quickscorer-rs.git", rev = "8314d03b5ef2d15385f8d0cd86d8f281d11b2d49", optional = true }
@@ -40,4 +40,4 @@ name = "correctness"
 path = "tests/correctness.rs"
 required-features = ["research"]
 
-# Artifact tests: cargo test --manifest-path benchmarks/Cargo.toml --release --features research
+# Artifact tests: cargo test --manifest-path experiments/benchmarks/Cargo.toml --release --features research

benchmarks/src/bin/sweep_bench.rs → experiments/benchmarks/src/bin/sweep_bench.rs

@@ -853,7 +853,7 @@ fn main() {
 
     // -----------------------------------------------------------------------
     // Validate mode: run TreeWalker on a cell and compare to an f64 reference.
-    // Triggered by --validate FILE. Used by prepare_scenario.py as the
+    // Triggered by --validate FILE. Used by `treewalker-exp prepare --validate` as the
     // per-cell correctness gate (TOL_F64 = 1e-14 for LightGBM).
     // -----------------------------------------------------------------------
     if args.iter().any(|a| a == "--validate") {

benchmarks/src/grid.rs → experiments/benchmarks/src/grid.rs

@@ -1,12 +1,13 @@
 //! Grid definitions and artifact directory discovery.
 //!
-//! Replaces the Python grid constants and dataset scanning logic from
-//! `sweep.py` and `utils.py`.
+//! The released grids' constants and the artifact directory scan. The grids
+//! now live in `experiments/grids.toml`; the runner redesign reads them from
+//! the execution manifest instead.
 
 use std::path::{Path, PathBuf};
 
 // ---------------------------------------------------------------------------
-// Grid combo constants (from Python utils.py)
+// Grid combo constants (the released grids)
 // ---------------------------------------------------------------------------
 
 /// Grid 1: full factorial `T × L × H`.

benchmarks/src/lib.rs → experiments/benchmarks/src/lib.rs

@@ -713,7 +713,8 @@ fn collect_grid3_stats(config: &RunConfig, cells: &[GridCell], arch: &str) {
 // Grid 4: Group distribution experiments (Expedia only)
 // ---------------------------------------------------------------------------
 
-/// Known distribution subdirectory names (from prepare.py --prepare-groups).
+/// Distribution subdirectory names of the released Grid 4. treewalker-exp no
+/// longer writes them; only `empirical`, the parent directory, remains.
 const GROUP_DISTRIBUTIONS: &[&str] = &["empirical", "fixed8", "fixed16", "fixed32", "geom8"];
 
 fn run_grid4_distributions(config: &RunConfig) {
@@ -940,7 +941,7 @@ fn count_work(cd: &CellData, ablation: Ablation, results: &mut [f64]) -> WorkCou
 /// over all `(k, G)` cells, with work counters and an inline
 /// correctness check against the GTIL f64 reference.
 ///
-/// Layout (written by `prepare_scenario.py`):
+/// Layout (written by `treewalker-exp prepare --suite scenario-v1`):
 ///   `<artifacts>/scenario_credit/lightgbm/`  — shared trained model
 ///   `<artifacts>/scenario_credit/cells/k{K}_G{G}/` — per-cell data + config
 ///        + `reference.bin` (GTIL f64 reference for the correctness gate)
@@ -1300,8 +1301,7 @@ mod tests {
 
     #[test]
     fn test_discover_cells() {
-        let artifacts =
-            PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../paper/experiments/artifacts");
+        let artifacts = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../artifacts");
         if !artifacts.exists() {
             return;
         }
@@ -1320,8 +1320,7 @@ mod tests {
 
     #[test]
     fn test_run_single_cell() {
-        let artifacts =
-            PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../paper/experiments/artifacts");
+        let artifacts = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../artifacts");
         if !artifacts.exists() {
             return;
         }

benchmarks/tests/config_test.rs → experiments/benchmarks/tests/config_test.rs

@@ -4,7 +4,7 @@ fn test_dir() -> std::path::PathBuf {
     std::path::PathBuf::from(std::env::var("TEST_ARTIFACTS").unwrap_or_else(|_| {
         concat!(
             env!("CARGO_MANIFEST_DIR"),
-            "/../paper/experiments/artifacts/expedia/nt500_md8"
+            "/../artifacts/expedia/nt500_md8"
         )
         .into()
     }))