diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index a394f01..126a0e5 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -130,7 +130,7 @@ jobs: run: uv sync --frozen --no-dev - name: Run GIFT-Eval CI subset on Modal - run: uv run modal run -m src.runners.run_modal::run_ci + run: make run-ci env: MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} @@ -145,11 +145,8 @@ jobs: aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }} aws-region: us-east-1 - - name: Sync CI results from S3 - run: make sync-ci-results - - name: Verify replication - run: uv run pytest tests/test_replication.py -n 0 -x + run: make verify-ci env: HF_TOKEN: ${{ secrets.HF_TOKEN }} diff --git a/experiments/gift-eval/Makefile b/experiments/gift-eval/Makefile index a907ef3..4beb6d1 100644 --- a/experiments/gift-eval/Makefile +++ b/experiments/gift-eval/Makefile @@ -1,4 +1,54 @@ -.PHONY: download-gift-eval-data upload-data-to-s3 sync-ci-results sync-results verify-all replication-table +# Gift-Eval experiment — common tasks (run from experiments/gift-eval/). +# +# Modal runs need MODAL_TOKEN_* , HF_TOKEN , and AWS creds in the environment +# (same as GitHub Actions secrets). Long grids: DETACHED=1 make run-replication-full ... + +.DEFAULT_GOAL := help + +.PHONY: help +.PHONY: download-gift-eval-data upload-data-to-s3 +.PHONY: run-ci sync-ci-results verify-ci +.PHONY: run-replication-pilot run-replication-full +.PHONY: sync-replication verify-replication-pilot verify-replication-full +.PHONY: sync-results verify-all replication-table + +RUN_ID ?= +DETACHED ?= +FORCE ?= +MAX_CONTAINERS ?= 20 + +# Pass -d to modal run for detached (background) execution. +MODAL_RUN = uv run modal run $(if $(filter 1 true yes,$(DETACHED)),-d,) + +## help — list targets (default) +help: + @echo "Gift-Eval Makefile" + @echo "" + @echo "Data" + @echo " download-gift-eval-data HF dataset → ./data/gift-eval" + @echo " upload-data-to-s3 sync local dataset to S3" + @echo "" + @echo "CI (configs/ci_subset.yaml — strict per-job HF verify in pytest)" + @echo " run-ci Modal GPU eval + overwrite results/ci on S3" + @echo " sync-ci-results S3 → ./results/ci" + @echo " verify-ci sync + pytest tests/test_replication.py" + @echo "" + @echo "Replication (configs/replication/ — pilot strict, full aggregate verify)" + @echo " run-replication-pilot RUN_ID=... — 8× m4_weekly/short on Modal" + @echo " run-replication-full RUN_ID=... — 16×97 grid (skips existing unless FORCE=1)" + @echo " sync-replication RUN_ID=... — S3 → ./results/replication/\$$RUN_ID" + @echo " verify-replication-pilot RUN_ID=... — per-job MASE/CRPS vs HF" + @echo " verify-replication-full RUN_ID=... — leaderboard aggregate + require complete" + @echo "" + @echo "Full benchmark (configs/models.yaml — production results/ on S3)" + @echo " sync-results S3 → ./results" + @echo " verify-all strict HF verify for all reference models" + @echo " replication-table analysis CSV (verify-only, no assert)" + @echo "" + @echo "Options: DETACHED=1 for Modal -d ; FORCE=1 to re-run all replication jobs" + @echo " MAX_CONTAINERS=N (default 20) caps parallel GPUs on replication full" + +# --- Data ------------------------------------------------------------------- download-gift-eval-data: @hf download Salesforce/GiftEval --repo-type=dataset --local-dir=./data/gift-eval @@ -6,10 +56,44 @@ download-gift-eval-data: upload-data-to-s3: download-gift-eval-data @aws s3 sync ./data/gift-eval s3://foundationforecast-gift-eval/data/gift-eval +# --- CI --------------------------------------------------------------------- +# GitHub Actions: make run-ci && make verify-ci + +run-ci: + @$(MODAL_RUN) -m src.runners.run_modal::run_ci + sync-ci-results: @mkdir -p ./results/ci @aws s3 sync s3://foundationforecast-gift-eval/results/ci ./results/ci +verify-ci: sync-ci-results + @uv run pytest tests/test_replication.py -n 0 -x + +# --- Replication ------------------------------------------------------------ + +run-replication-pilot: + @test -n "$(RUN_ID)" || (echo "Set RUN_ID=..." && exit 1) + @$(MODAL_RUN) -m src.runners.run_modal::run_replication_pilot --run-id $(RUN_ID) + +run-replication-full: + @test -n "$(RUN_ID)" || (echo "Set RUN_ID=..." && exit 1) + @$(MODAL_RUN) -m src.runners.run_modal::run_replication_full --run-id $(RUN_ID) \ + --max-containers $(MAX_CONTAINERS) \ + $(if $(filter 1 true yes,$(FORCE)),--force,) + +sync-replication: + @test -n "$(RUN_ID)" || (echo "Set RUN_ID=..." && exit 1) + @mkdir -p ./results/replication/$(RUN_ID) + @aws s3 sync s3://foundationforecast-gift-eval/results/replication/$(RUN_ID) ./results/replication/$(RUN_ID) + +verify-replication-pilot: sync-replication + @uv run python -m src.runners.run_verify --replication --run-id $(RUN_ID) --replication-pilot + +verify-replication-full: sync-replication + @uv run python -m src.runners.run_verify --replication --run-id $(RUN_ID) --all --require-complete + +# --- Full benchmark --------------------------------------------------------- + sync-results: @mkdir -p ./results @aws s3 sync s3://foundationforecast-gift-eval/results ./results diff --git a/experiments/gift-eval/README.md b/experiments/gift-eval/README.md index b93594b..e651e49 100644 --- a/experiments/gift-eval/README.md +++ b/experiments/gift-eval/README.md @@ -1,183 +1,109 @@ -# FoundationForecast GIFT-Eval Benchmark +# FoundationForecast replicates Salesforce's GIFT-Eval for $37 -End-to-end GIFT-Eval benchmark for [FoundationForecast](https://github.com/AzulGarza/foundationforecast) model wrappers. Evaluation uses [`timecopilot-gift-eval`](https://github.com/TimeCopilot/timecopilot-gift-eval); replication checks compare outputs to official Hugging Face reference CSVs. +Replication of **Salesforce's full GIFT-Eval** forecasting benchmark **under a single API (FoundationForecast)**. The full replication of **16 models** cost **$37** and took **≈34 hours** of GPU time. The best performance among the implemented models is achieved by **TimesFM 3 (Google) for under $8**. -## Layout +This directory also [regression-tests](https://github.com/TimeCopilot/foundationforecast/actions/workflows/ci.yaml) FoundationForecast on every merge to `main`. -``` -src/ -├── eval/ run_gift_eval(), model registry, job config -├── verify/ HF reference loading and replication checks -└── runners/ CLI and Modal entrypoints -tests/ pytest replication checks (imports src.verify) -configs/ models.yaml (full matrix) and ci_subset.yaml (CI) -``` +## Why? -## Setup +Companies and practitioners rely on accurate forecasts to make better decisions. Some of those decisions have to be made **quickly**; as compute capacity grows, that speed requirement is even sharper in nascent, high-frequency settings such as energy markets, cloud capacity, and real-time operations. Other decisions have to be made **under a budget**, so accuracy should be judged together with computational cost. -```bash -cd experiments/gift-eval -uv sync -``` - -Installs the in-repo editable `foundationforecast` package from the monorepo root (`../..`), not PyPI — so local runs and CI always use the current wrapper code. After changing the library (for example model weight caching), refresh the lock metadata with `uv lock` in this directory so `uv sync --frozen` in CI matches the editable source version. +**Time Series Foundation Models (TSFMs)** have changed how forecasting is done. What is still missing is a clear view of the **accuracy–compute trade-off**, and a **unified way** to evaluate the models and deploy them in production. -Requires Python 3.11+. +[GIFT-Eval](https://huggingface.co/spaces/Salesforce/GIFT-Eval) ([Aksu et al., 2024](https://arxiv.org/abs/2410.10393)) is one of the benchmarks practitioners and researchers use to publish SOTA results, by uploading their own implementations. The process does not, by itself, give a unified path to replicability: each model ships with its own API, inference rules, and design. The leaderboard also reports **accuracy only**. It does not measure computational cost on shared infrastructure. -## Dataset - -```bash -make download-gift-eval-data -# optional: make upload-data-to-s3 -``` +We built this reproducible experiment to show that a **unified API** makes it practical to test both **accuracy and cost** on large time series data. **FoundationForecast** is that API. The results below are from running it on the full GIFT-Eval grid. -## Run a single job locally +## Results -```bash -uv run python -m src.runners.run_model \ - --model-key amazon--chronos-bolt-small \ - --dataset-name m4_weekly \ - --term short \ - --storage-path ./data/gift-eval \ - --output-root ./results -``` +We replicated **16 foundation models** on Salesforce's full GIFT-Eval: 23 datasets, ~144k time series, 177M data points, evaluated as **326,490** forecasting tasks (one task = one series test window; multivariate series are forecast per variate, giving **371,330** univariate forecasts) grouped into **97** dataset / frequency / horizon configurations. That is **5.9 billion** probabilistic forecast values (forecasts × horizon × 16 models × 9 quantiles). -## CI subset +**MASE** (point) and **CRPS** (probabilistic) follow the leaderboard convention: each configuration is normalized by Seasonal Naive and aggregated with a **geometric mean** across the 97 configurations. Cost is **not** part of the original evaluation; we estimate it from per-job GPU wall time on a single **NVIDIA A10G** at **$1.10/GPU-hour**. -[`configs/ci_subset.yaml`](configs/ci_subset.yaml) defines **13 jobs**: Chronos on -`m4_weekly/short` and `m4_hourly/short`, plus one representative `model_key` each for -TimesFM 2.5/3.0, TiRex 1.1 and TiRex-2-Zeroshot, Moirai, Toto, FlowState, PatchTST-FM r1, -Granite PatchTST-FM r2, T0, and Tafsut (all on `m4_weekly/short`). -Each job runs on Modal GPU and is **HF-verified** in pytest (metrics must match the -official GIFT-Eval reference CSV). +Every FoundationForecast implementation reproduces the model's aggregated MASE and CRPS **within 2%** of the official leaderboard submission (Chronos-2 small has no published reference yet). -### Local GPU +### Replication -```bash -uv run python -m src.runners.run_ci --local --verify \ - --storage-path ./data/gift-eval \ - --output-root ./results/ci -``` +The table below compares **official** GIFT-Eval aggregates with **FF** (the same aggregate from our FoundationForecast run). Organizations are grouped alphabetically. Shaded MASE and CRPS cells mark the **1st / 2nd / 3rd best** models in each column (darker is better); the ranking is the same on both sides and for both metrics: **TimesFM 3**, **PatchTST-FM r2**, **T0 beta**. -### Modal (CI / GitHub Actions) +GIFT-Eval replication table typeset in LaTeX: official vs FoundationForecast MASE and CRPS, with the top-3 models shaded in both metrics -Always re-runs and overwrites results (no skip-if-exists). Full grid skips jobs that -already have outputs. +### Pareto frontier -```bash -uv run modal run -m src.runners.run_modal::run_ci -make sync-ci-results # download results for local verify / pytest -uv run pytest tests/test_replication.py -n 0 -x -``` +Accuracy alone does not decide which model to deploy. The figure below plots **cost vs accuracy** for the same 16 models: cheaper is to the right, better (lower error) is up. The purple polyline is the **Pareto frontier**. TimesFM 3 is the most accurate model we ran, at **$7.41**; Toto 2 4M is the cheapest plotted point, at **$0.61**. The full grid is **$37 / 34 GPU-hours**. -## Full benchmark grid (Modal) +Cost vs accuracy Pareto frontier on 5.9B probabilistic forecast values for 16 foundation models on GIFT-Eval -One GPU job per `(model_key, dataset, term)`: +## Reproducibility ```bash -uv run modal run -m src.runners.run_modal::main -``` - -## Verify against HF references - -Compare local/S3 results to official GIFT-Eval CSVs. Uses consolidated -`results/{model_key}/all_results.csv` if present, otherwise aggregates -per-job CSVs under `results/{model_key}/`. - -Strict replication asserts **MASE** and **CRPS** only (the GIFT-Eval ranking -metrics), with default tolerances `atol=0.01`, `rtol=0.025`. Other columns in -`all_results.csv` are still written but not compared. - -Every verify run also writes a replication analysis table (CSV) with: +cd experiments/gift-eval +uv sync +make download-gift-eval-data -| Column | Description | -|--------|-------------| -| `dataset` | GIFT-Eval dataset config (e.g. `m4_weekly/W/short`) | -| `model` | Model alias in results CSV | -| `model_key` | Experiment registry key | -| `time_seconds` | Eval wall time (from per-job `timing.json`) | -| `mase` | Our `eval_metrics/MASE[0.5]` | -| `crps` | Our `eval_metrics/mean_weighted_sum_quantile_loss` | -| `reported_gift_eval_mase` | Official HF reference MASE | -| `reported_gift_eval_crps` | Official HF reference CRPS | -| `mase_diff` | `mase - reported_gift_eval_mase` | -| `crps_diff` | `crps - reported_gift_eval_crps` | +# 16 models × 97 configs +make run-replication-full RUN_ID=2026-10-06-nb-v2 +make verify-replication-full RUN_ID=2026-10-06-nb-v2 +``` -```bash -# CI subset (per-job layout under results/ci/) -uv run python -m src.runners.run_verify --ci +Params follow the [official GIFT-Eval notebooks](https://github.com/SalesforceAIResearch/gift-eval/tree/main/notebooks). Results land in `s3://foundationforecast-gift-eval/results/replication//`. The full grid skips jobs that already have outputs unless `FORCE=1`. Long runs: `DETACHED=1`. Parallel GPUs: `MAX_CONTAINERS=20` (default). -# One model -uv run python -m src.runners.run_verify --model-key amazon--chronos-bolt-small +On every merge to `main`, [CI](https://github.com/TimeCopilot/foundationforecast/actions/workflows/ci.yaml) re-runs a 13-job subset (`make run-ci` / `make verify-ci`) with **strict per-job** MASE/CRPS against the Hugging Face references. -# All models with a reference_slug in configs/models.yaml -make sync-results # or: aws s3 sync s3://foundationforecast-gift-eval/results ./results -uv run python -m src.runners.run_verify --all +### Infrastructure -# Table only (no strict assert) — good for exploratory analysis -uv run python -m src.runners.run_verify --all --verify-only \ - --table-output ./results/replication_table.csv -make replication-table +| Need | What | +| --- | --- | +| Python | 3.11+ (`uv sync` from this directory; editable install of repo-root `foundationforecast`) | +| GPU jobs | [Modal](https://modal.com) (`MODAL_TOKEN_ID`, `MODAL_TOKEN_SECRET`) | +| Object store | S3 bucket `foundationforecast-gift-eval` (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`) | +| Weights | `HF_TOKEN` (gated models such as `t0-alpha`); accept Hub licenses | +| Optional | `DETACHED=1`, `FORCE=1`, `MAX_CONTAINERS=N` | -# Require every HF dataset to be present (not just compare overlap) -uv run python -m src.runners.run_verify --all --require-complete -``` +Modal secrets: `aws-secret`, `hf-secret` (`modal secret create hf-secret HF_TOKEN=hf_...`). -Or in one step: +## Conclusion -```bash -make verify-all -``` +GIFT-Eval already tells us **which models are accurate**. Running them through one API, on one GPU SKU, adds the missing axis: **what that accuracy costs**. The replication shows the leaderboard numbers are recoverable (aggregates within 2%). The Pareto plot shows they are not interchangeable: the most accurate model in this run is also one of the more expensive, and several cheaper models sit close enough on CRPS to matter when the budget is the constraint. -**Note:** `time_seconds` is recorded when a job runs via `run_gift_eval` (writes -`timing.json` next to each `all_results.csv`). To backfill timing for jobs that -ran before timing was added: +If you are choosing a TSFM for production, measure **both**. FoundationForecast is the API we used to make that measurement reproducible. -```bash -# Full grid: rerun only jobs with results but no timing.json on S3 -uv run modal run -m src.runners.run_modal::run_missing_timing +## Acknowledgements -# Full grid: force rerun everything (also refreshes metrics) -uv run modal run -m src.runners.run_modal::main --force +Thanks to the teams who trained and released the models in this grid: Amazon (Chronos / Chronos-2 / Chronos Bolt), Datadog (Toto 2), Google (TimesFM 3), Huawei (Tafsut), IBM (PatchTST-FM), NX-AI (TiRex), Salesforce (Moirai), and The Forecasting Company (T0). Thanks also to Salesforce AI Research for GIFT-Eval — the datasets, notebooks, and public leaderboard this experiment replicates. -# CI subset locally -uv run python -m src.runners.run_ci --local --missing-timing-only +## References -# CI on Modal always reruns with force=True (timing included every CI run) -uv run modal run -m src.runners.run_modal::run_ci -``` +- Aksu, T., Woo, G., Liu, J., Liu, X., Liu, C., Savarese, S., Xiong, C., & Sahoo, D. (2024). *GIFT-Eval: A Benchmark For General Time Series Forecasting Model Evaluation*. [arXiv:2410.10393](https://arxiv.org/abs/2410.10393). +- Ansari, A. F., et al. (2024). *Chronos: Learning the Language of Time Series*. [arXiv:2403.07815](https://arxiv.org/abs/2403.07815). +- Das, A., Kong, W., Sen, R., & Zhou, Y. (2023). *A decoder-only foundation model for time-series forecasting* (TimesFM). [arXiv:2310.10688](https://arxiv.org/abs/2310.10688). +- Woo, G., Liu, C., Kumar, A., Xiong, C., Savarese, S., & Sahoo, D. (2024). *Unified Training of Universal Time Series Forecasting Transformers* (Moirai). [arXiv:2402.02592](https://arxiv.org/abs/2402.02592). +- Cohen, B., et al. (2025). *This Time is Different: An Observability Perspective on Time Series Foundation Models* (Toto). [arXiv:2505.14766](https://arxiv.org/abs/2505.14766). +- Auer, A., et al. (2025). *TiRex: Zero-Shot Forecasting Across Long and Short Horizons with Enhanced In-Context Learning*. [arXiv:2505.23719](https://arxiv.org/abs/2505.23719). +- Auer, A., et al. (2026). *TiRex-2*. [arXiv:2607.01204](https://arxiv.org/abs/2607.01204). +- IBM. *PatchTST-FM*. [arXiv:2602.06909](https://arxiv.org/abs/2602.06909). +- The Forecasting Company. [T0 alpha](https://huggingface.co/theforecastingcompany/t0-alpha) / [T0 beta](https://huggingface.co/theforecastingcompany/t0-beta). +- Tafsut-FM. [Tafsut univariate base](https://huggingface.co/Tafsut-FM/tafsut-univariate-base). +- Garza, A., & Rosillo, R. (2026). *FoundationForecast: The API for time series foundation models*. . -Then sync and rebuild the table: +## How to cite -```bash -make sync-results -make replication-table +```bibtex +@software{foundationforecast-gift-eval, + title = {FoundationForecast replicates Salesforce's GIFT-Eval for \$37}, + author = {Garza, Azul and Rosillo, Ren{\'e}e}, + year = {2026}, + url = {https://github.com/TimeCopilot/foundationforecast/tree/main/experiments/gift-eval}, + note = {Replication of 16 time series foundation models on GIFT-Eval under a single API}, +} ``` -## Consolidate S3 results - -```bash -uv run python -m src.runners.download_results --model-key amazon--chronos-bolt-small +```bibtex +@software{foundationforecast, + title = {FoundationForecast: The API for time series foundation models}, + author = {Garza, Azul and Rosillo, Ren{\'e}e}, + year = {2026}, + url = {https://github.com/TimeCopilot/foundationforecast}, + license = {Apache-2.0}, +} ``` - -## Infrastructure - -- **S3 bucket:** `foundationforecast-gift-eval` -- **Modal secrets:** - - `aws-secret` — `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY` - - `hf-secret` — `HF_TOKEN` (required for gated models like `t0-alpha`; create with - `modal secret create hf-secret HF_TOKEN=hf_...`) -- **Modal tokens:** `MODAL_TOKEN_ID`, `MODAL_TOKEN_SECRET` -- **Hugging Face:** accept model licenses on the Hub, then set `HF_TOKEN` in `hf-secret` - -## Adding a model - -1. Add an entry to `configs/models.yaml` with slugified `repo_id` as `model_key` - (`org--model`), `class`, `kwargs.repo_id`, and `reference_slug`. -2. Set `kwargs.repo_id` from the official GIFT-Eval - `results/{reference_slug}/config.json` → `model_link` (many models use - gifteval-specific HF repos, not the default public checkpoint). -3. Set `reference_slug` to the official GIFT-Eval folder name; `alias` defaults from - that in `build_model()` — set `kwargs.alias` explicitly when the CSV `model` column - differs from the folder slug (e.g. `chronos_base` → `Chronos_base`). -4. Set `reference_slug: null` if no public reference exists (verify skips that model). diff --git a/experiments/gift-eval/configs/models.yaml b/experiments/gift-eval/configs/models.yaml index d1dcac1..faa0722 100644 --- a/experiments/gift-eval/configs/models.yaml +++ b/experiments/gift-eval/configs/models.yaml @@ -63,10 +63,13 @@ models: amazon--chronos-2: class: foundationforecast.models.chronos.Chronos - reference_slug: chronos-2 + reference_slug: Chronos-2 + max_length: null + predictor_batch_size: 16 kwargs: repo_id: amazon/chronos-2 - batch_size: 64 + alias: Chronos-2 + batch_size: 16 autogluon--chronos-2-synth: class: foundationforecast.models.chronos.Chronos @@ -76,12 +79,22 @@ models: alias: Chronos-2-Synth batch_size: 64 + autogluon--chronos-2-small: + class: foundationforecast.models.chronos.Chronos + reference_slug: Chronos-2-Small + max_length: null + predictor_batch_size: 16 + kwargs: + repo_id: autogluon/chronos-2-small + alias: Chronos-2-Small + batch_size: 16 + # --- TimesFM --- google--timesfm-1.0-200m: class: foundationforecast.models.timesfm.TimesFM reference_slug: timesfm kwargs: - repo_id: google/timesfm-1.0-200m + repo_id: google/timesfm-1.0-200m-pytorch alias: TimesFM batch_size: 64 diff --git a/experiments/gift-eval/configs/replication/families.yaml b/experiments/gift-eval/configs/replication/families.yaml new file mode 100644 index 0000000..01e5747 --- /dev/null +++ b/experiments/gift-eval/configs/replication/families.yaml @@ -0,0 +1,18 @@ +# One representative model per family for notebook-aligned replication runs. +model_keys: + - amazon--chronos-2 + - amazon--chronos-bolt-base + - amazon--chronos-bolt-small + - autogluon--chronos-2-small + - Datadog--Toto-2.0-313m + - Datadog--Toto-2.0-4m + - google--timesfm-3.0-pytorch + - ibm-granite--granite-timeseries-patchtst-fm-r1 + - ibm-granite--granite-timeseries-patchtst-fm-r2 + - NX-AI--TiRex-1.1-gifteval + - NX-AI--TiRex-2-gifteval-zs + - Salesforce--moirai-1.1-R-large--Moirai_large + - Salesforce--moirai-2.0-R-small + - tafsut-univariate-base + - theforecastingcompany--t0-alpha + - theforecastingcompany--t0-beta diff --git a/experiments/gift-eval/configs/replication/models.yaml b/experiments/gift-eval/configs/replication/models.yaml new file mode 100644 index 0000000..d9e8908 --- /dev/null +++ b/experiments/gift-eval/configs/replication/models.yaml @@ -0,0 +1,151 @@ +# Notebook-aligned params for replication runs only (see configs/replication/notebooks.yaml). +models: + amazon--chronos-2: + class: foundationforecast.models.chronos.Chronos + reference_slug: Chronos-2 + max_length: null + predictor_batch_size: 16 + kwargs: + repo_id: amazon/chronos-2 + alias: Chronos-2 + batch_size: 16 + + amazon--chronos-bolt-base: + class: foundationforecast.models.chronos.Chronos + reference_slug: chronos_bolt_base + max_length: null + predictor_batch_size: 1024 + kwargs: + repo_id: amazon/chronos-bolt-base + alias: chronos_bolt_base + batch_size: 32 + + amazon--chronos-bolt-small: + class: foundationforecast.models.chronos.Chronos + reference_slug: chronos_bolt_small + max_length: null + predictor_batch_size: 1024 + kwargs: + repo_id: amazon/chronos-bolt-small + alias: chronos_bolt_small + batch_size: 64 + + autogluon--chronos-2-small: + class: foundationforecast.models.chronos.Chronos + # HF GIFT-Eval results/{slug} not published yet; enable aggregate verify when live. + reference_slug: Chronos-2-Small + max_length: null + predictor_batch_size: 16 + kwargs: + repo_id: autogluon/chronos-2-small + alias: Chronos-2-Small + batch_size: 16 + + Datadog--Toto-2.0-313m: + class: foundationforecast.models.toto.Toto + reference_slug: Toto-2.0-313m + kwargs: + repo_id: Datadog/Toto-2.0-313m + context_length: 4096 + batch_size: 512 + + Datadog--Toto-2.0-4m: + class: foundationforecast.models.toto.Toto + reference_slug: Toto-2.0-4m + kwargs: + repo_id: Datadog/Toto-2.0-4m + context_length: 4096 + batch_size: 512 + + google--timesfm-3.0-pytorch: + class: foundationforecast.models.timesfm.TimesFM + reference_slug: TimesFM-3 + kwargs: + repo_id: google/timesfm-3.0-pytorch + alias: TimesFM-3 + context_length: 15360 + batch_size: 64 + input_patch_length: 32 + output_patch_length: 64 + use_variate_attention: true + use_sdpa: true + + ibm-granite--granite-timeseries-patchtst-fm-r1: + class: foundationforecast.models.patchtst_fm.PatchTSTFM + reference_slug: Granite-PatchTST-FM-r1 + predictor_batch_size: 1024 + kwargs: + repo_id: ibm-granite/granite-timeseries-patchtst-fm-r1 + alias: Granite-PatchTST-FM-r1 + batch_size: 32 + + ibm-granite--granite-timeseries-patchtst-fm-r2: + class: foundationforecast.models.patchtst_fm.PatchTSTFM + reference_slug: Granite-PatchTST-FM-r2 + max_length: 8192 + kwargs: + repo_id: ibm-granite/granite-timeseries-patchtst-fm-r2 + context_length: 8192 + batch_size: 32 + + NX-AI--TiRex-1.1-gifteval: + class: foundationforecast.models.tirex.TiRex + reference_slug: TiRex + predictor_batch_size: 512 + kwargs: + repo_id: NX-AI/TiRex-1.1-gifteval + alias: TiRex + batch_size: 512 + + NX-AI--TiRex-2-gifteval-zs: + class: foundationforecast.models.tirex.TiRex + reference_slug: TiRex-2-Zeroshot + predictor_batch_size: 512 + kwargs: + repo_id: NX-AI/TiRex-2-gifteval-zs + batch_size: 512 + + Salesforce--moirai-1.1-R-large--Moirai_large: + class: foundationforecast.models.moirai.Moirai + reference_slug: Moirai_large + predictor_batch_size: 512 + kwargs: + repo_id: Salesforce/moirai-1.1-R-large + alias: Moirai_large + context_length: 4000 + batch_size: 16 + + Salesforce--moirai-2.0-R-small: + class: foundationforecast.models.moirai.Moirai + reference_slug: Moirai2 + predictor_batch_size: 512 + kwargs: + repo_id: Salesforce/moirai-2.0-R-small + context_length: 4000 + batch_size: 32 + + tafsut-univariate-base: + class: foundationforecast.models.tafsut.Tafsut + reference_slug: tafsut + kwargs: + repo_id: Tafsut-FM/tafsut-univariate-base + alias: tafsut + batch_size: 128 + + theforecastingcompany--t0-alpha: + class: foundationforecast.models.t0.T0 + reference_slug: t0-alpha + kwargs: + repo_id: theforecastingcompany/t0-alpha + alias: t0-alpha + context_length: 8192 + batch_size: 64 + + theforecastingcompany--t0-beta: + class: foundationforecast.models.t0.T0 + reference_slug: t0-beta + kwargs: + repo_id: theforecastingcompany/t0-beta + alias: t0-beta + context_length: 8192 + batch_size: 64 diff --git a/experiments/gift-eval/configs/replication/notebooks.yaml b/experiments/gift-eval/configs/replication/notebooks.yaml new file mode 100644 index 0000000..33460ea --- /dev/null +++ b/experiments/gift-eval/configs/replication/notebooks.yaml @@ -0,0 +1,51 @@ +# Official GIFT-Eval replication notebooks (SalesforceAIResearch/gift-eval). +# HF results: https://huggingface.co/spaces/Salesforce/GIFT-Eval/tree/main/results/{reference_slug} +sources: + amazon--chronos-2: + notebook: notebooks/chronos-2.ipynb + reference_slug: Chronos-2 + amazon--chronos-bolt-base: + notebook: notebooks/chronos.ipynb + reference_slug: chronos_bolt_base + amazon--chronos-bolt-small: + notebook: notebooks/chronos.ipynb + reference_slug: chronos_bolt_small + autogluon--chronos-2-small: + notebook: notebooks/chronos-2.ipynb + reference_slug: Chronos-2-Small + Datadog--Toto-2.0-313m: + notebook: notebooks/toto_2_0.ipynb + reference_slug: Toto-2.0-313m + Datadog--Toto-2.0-4m: + notebook: notebooks/toto_2_0.ipynb + reference_slug: Toto-2.0-4m + google--timesfm-3.0-pytorch: + notebook: notebooks/timesfm3.ipynb + reference_slug: TimesFM-3 + ibm-granite--granite-timeseries-patchtst-fm-r1: + notebook: notebooks/patchtst_fm.ipynb + reference_slug: Granite-PatchTST-FM-r1 + ibm-granite--granite-timeseries-patchtst-fm-r2: + notebook: notebooks/patchtst_fm_r2.ipynb + reference_slug: Granite-PatchTST-FM-r2 + NX-AI--TiRex-1.1-gifteval: + notebook: notebooks/tirex.ipynb + reference_slug: TiRex + NX-AI--TiRex-2-gifteval-zs: + notebook: notebooks/tirex-2.ipynb + reference_slug: TiRex-2-Zeroshot + Salesforce--moirai-1.1-R-large--Moirai_large: + notebook: notebooks/moirai.ipynb + reference_slug: Moirai_large + Salesforce--moirai-2.0-R-small: + notebook: notebooks/moirai2.ipynb + reference_slug: Moirai2 + tafsut-univariate-base: + notebook: notebooks/tafsut.ipynb + reference_slug: tafsut + theforecastingcompany--t0-alpha: + notebook: notebooks/t0-alpha.ipynb + reference_slug: t0-alpha + theforecastingcompany--t0-beta: + notebook: notebooks/t0-beta.ipynb + reference_slug: t0-beta diff --git a/experiments/gift-eval/src/eval/evaluate.py b/experiments/gift-eval/src/eval/evaluate.py index 9176f86..f632e23 100644 --- a/experiments/gift-eval/src/eval/evaluate.py +++ b/experiments/gift-eval/src/eval/evaluate.py @@ -8,7 +8,7 @@ from timecopilot_gift_eval import GIFTEval, GluonTSPredictor from .jobs import Job, job_output_dir, result_csv, timing_json -from .models import build_model, predictor_batch_size, predictor_max_length +from .models import Registry, build_model, predictor_batch_size, predictor_max_length logger = logging.getLogger(__name__) @@ -23,6 +23,7 @@ def run_gift_eval( storage_path: Path | str, output_root: Path | str = Path("results"), overwrite_results: bool = False, + registry: Registry = "default", ) -> Path: output_path = job_output_dir(job, Path(output_root)) output_path.mkdir(parents=True, exist_ok=True) @@ -34,16 +35,18 @@ def run_gift_eval( job.term, ) - forecaster = build_model(job.model_key) + forecaster = build_model(job.model_key, registry=registry) predictor = GluonTSPredictor( forecaster=forecaster, max_length=predictor_max_length( job.model_key, forecaster, + registry=registry, default=DEFAULT_MAX_LENGTH, ), batch_size=predictor_batch_size( job.model_key, + registry=registry, default=DEFAULT_PREDICTOR_BATCH_SIZE, ), ) diff --git a/experiments/gift-eval/src/eval/jobs.py b/experiments/gift-eval/src/eval/jobs.py index 13a5c7a..e23edbd 100644 --- a/experiments/gift-eval/src/eval/jobs.py +++ b/experiments/gift-eval/src/eval/jobs.py @@ -8,7 +8,13 @@ from timecopilot_gift_eval.utils import DATASETS_WITH_TERMS CONFIGS_DIR = Path(__file__).resolve().parents[2] / "configs" +REPLICATION_CONFIGS_DIR = CONFIGS_DIR / "replication" DEFAULT_RESULTS_ROOT = Path("results") +REPLICATION_PILOT_DATASET = "m4_weekly" +REPLICATION_PILOT_TERM = "short" + +# Families in the replication grid before GIFT-Eval HF results/{slug} exists. +REPLICATION_PENDING_HF_REFERENCE = frozenset({"autogluon--chronos-2-small"}) @dataclass(frozen=True) @@ -28,6 +34,48 @@ def load_models_config() -> dict: return _load_yaml(CONFIGS_DIR / "models.yaml")["models"] +@lru_cache +def load_replication_models_config() -> dict: + return _load_yaml(REPLICATION_CONFIGS_DIR / "models.yaml")["models"] + + +def load_models_config_for_registry(registry: str) -> dict: + if registry == "default": + return load_models_config() + if registry == "replication": + return load_replication_models_config() + raise ValueError(f"Unknown registry {registry!r}") + + +@lru_cache +def load_replication_families() -> tuple[str, ...]: + raw = _load_yaml(REPLICATION_CONFIGS_DIR / "families.yaml")["model_keys"] + return tuple(raw) + + +def load_replication_pilot_jobs() -> list[Job]: + return [ + Job( + model_key=model_key, + dataset_name=REPLICATION_PILOT_DATASET, + term=REPLICATION_PILOT_TERM, + ) + for model_key in load_replication_families() + ] + + +def load_replication_matrix() -> list[Job]: + return [ + Job(model_key=model_key, dataset_name=dataset_name, term=term) + for model_key in load_replication_families() + for dataset_name, term in DATASETS_WITH_TERMS + ] + + +def replication_output_root(run_id: str) -> Path: + return DEFAULT_RESULTS_ROOT / "replication" / run_id + + def load_ci_subset() -> list[Job]: raw = _load_yaml(CONFIGS_DIR / "ci_subset.yaml")["jobs"] return [Job(**job) for job in raw] @@ -58,6 +106,77 @@ def ci_output_root() -> Path: return DEFAULT_RESULTS_ROOT / "ci" +def job_key(job: Job) -> tuple[str, str, str]: + return (job.model_key, job.dataset_name, job.term) + + +def parse_job_dir_suffix(relative_dir: str) -> tuple[str, str, str] | None: + """Parse ``model_key/dataset/term`` (dataset may contain ``/``).""" + parts = relative_dir.strip("/").split("/") + if len(parts) < 3: + return None + term = parts[-1] + model_key = parts[0] + dataset_name = "/".join(parts[1:-1]) + return (model_key, dataset_name, term) + + +def index_s3_job_artifacts( + *, + bucket: str, + prefix: str, +) -> tuple[set[tuple[str, str, str]], set[tuple[str, str, str]]]: + """List ``all_results.csv`` / ``timing.json`` under prefix in one S3 walk.""" + import fsspec + + fs = fsspec.filesystem("s3") + base = f"{bucket}/{prefix.strip('/')}/" + has_results: set[tuple[str, str, str]] = set() + has_timing: set[tuple[str, str, str]] = set() + try: + paths = fs.find(base) + except FileNotFoundError: + return has_results, has_timing + for path in paths: + if not path.startswith(base): + continue + rel = path[len(base) :] + if rel.endswith("/all_results.csv"): + key = parse_job_dir_suffix(rel[: -len("/all_results.csv")]) + if key is not None: + has_results.add(key) + elif rel.endswith("/timing.json"): + key = parse_job_dir_suffix(rel[: -len("/timing.json")]) + if key is not None: + has_timing.add(key) + return has_results, has_timing + + +def filter_jobs_by_s3_mode( + jobs: list[Job], + *, + bucket: str, + prefix: str, + mode: str, +) -> list[Job]: + if mode == "all": + return list(jobs) + if mode not in ("missing", "missing_timing"): + raise ValueError(f"Unknown job selection mode: {mode!r}") + has_results, has_timing = index_s3_job_artifacts(bucket=bucket, prefix=prefix) + selected: list[Job] = [] + for job in jobs: + key = job_key(job) + has_r = key in has_results + has_t = key in has_timing + if mode == "missing": + if not has_r: + selected.append(job) + elif has_r and not has_t: + selected.append(job) + return selected + + def jobs_missing_timing(jobs: list[Job], output_root: Path) -> list[Job]: missing: list[Job] = [] for job in jobs: diff --git a/experiments/gift-eval/src/eval/models.py b/experiments/gift-eval/src/eval/models.py index 8fad6a5..3a53dfa 100644 --- a/experiments/gift-eval/src/eval/models.py +++ b/experiments/gift-eval/src/eval/models.py @@ -1,11 +1,13 @@ from __future__ import annotations import importlib -from typing import Any +from typing import Any, Literal from timecopilot_gift_eval.protocol import ForecasterProtocol -from .jobs import load_models_config +from .jobs import load_models_config_for_registry + +Registry = Literal["default", "replication"] def _import_class(class_path: str) -> type: @@ -14,13 +16,27 @@ def _import_class(class_path: str) -> type: return getattr(module, class_name) -def build_model(model_key: str) -> ForecasterProtocol: - models = load_models_config() +def _model_spec(model_key: str, registry: Registry) -> dict: + models = load_models_config_for_registry(registry) if model_key not in models: available = ", ".join(sorted(models)) - raise KeyError(f"Unknown model_key {model_key!r}. Available: {available}") + raise KeyError( + f"Unknown model_key {model_key!r} in registry {registry!r}. " + f"Available: {available}" + ) + return models[model_key] + + +def model_spec(model_key: str, *, registry: Registry = "default") -> dict: + return _model_spec(model_key, registry) - spec = models[model_key] + +def build_model( + model_key: str, + *, + registry: Registry = "default", +) -> ForecasterProtocol: + spec = _model_spec(model_key, registry) model_cls = _import_class(spec["class"]) kwargs: dict[str, Any] = dict(spec.get("kwargs", {})) reference = spec.get("reference_slug") @@ -29,8 +45,13 @@ def build_model(model_key: str) -> ForecasterProtocol: return model_cls(**kwargs) -def predictor_batch_size(model_key: str, *, default: int = 1024) -> int: - spec = load_models_config()[model_key] +def predictor_batch_size( + model_key: str, + *, + registry: Registry = "default", + default: int = 1024, +) -> int: + spec = _model_spec(model_key, registry) if "predictor_batch_size" in spec: return int(spec["predictor_batch_size"]) return default @@ -40,17 +61,24 @@ def predictor_max_length( model_key: str, forecaster: ForecasterProtocol, *, + registry: Registry = "default", default: int = 4096, ) -> int | None: - spec = load_models_config()[model_key] + spec = _model_spec(model_key, registry) if "max_length" in spec: max_length = spec["max_length"] return None if max_length is None else int(max_length) return int(getattr(forecaster, "context_length", default)) -def reference_slug(model_key: str) -> str | None: - models = load_models_config() - if model_key not in models: - raise KeyError(f"Unknown model_key {model_key!r}") - return models[model_key].get("reference_slug") +def reference_slug(model_key: str, *, registry: Registry = "default") -> str | None: + return _model_spec(model_key, registry).get("reference_slug") + + +def model_keys_with_reference(*, registry: Registry = "default") -> list[str]: + models = load_models_config_for_registry(registry) + return [ + model_key + for model_key, spec in models.items() + if spec.get("reference_slug") is not None + ] diff --git a/experiments/gift-eval/src/runners/run_modal.py b/experiments/gift-eval/src/runners/run_modal.py index 454db0e..fc7bc40 100644 --- a/experiments/gift-eval/src/runners/run_modal.py +++ b/experiments/gift-eval/src/runners/run_modal.py @@ -3,22 +3,32 @@ import modal +_MODAL_GIFT_EVAL = "/root/gift-eval" _MODAL_MONOREPO = "/root/monorepo" -_MODAL_GIFT_EVAL = f"{_MODAL_MONOREPO}/experiments/gift-eval" -def _resolve_paths() -> tuple[Path, Path]: +def _resolve_gift_eval_root() -> Path: here = Path(__file__).resolve() try: gift_eval_root = here.parents[2] if (gift_eval_root / "pyproject.toml").exists(): - return gift_eval_root, gift_eval_root.parent.parent + return gift_eval_root except IndexError: pass - return Path(_MODAL_GIFT_EVAL), Path(_MODAL_MONOREPO) + return Path(_MODAL_GIFT_EVAL) -_GIFT_EVAL_ROOT, _REPO_ROOT = _resolve_paths() +def _resolve_monorepo_root(gift_eval_root: Path) -> Path: + candidate = gift_eval_root.parent.parent + if (candidate / "pyproject.toml").is_file() and ( + candidate / "foundationforecast" + ).is_dir(): + return candidate + return Path(_MODAL_MONOREPO) + + +_GIFT_EVAL_ROOT = _resolve_gift_eval_root() +_REPO_ROOT = _resolve_monorepo_root(_GIFT_EVAL_ROOT) app = modal.App(name="foundationforecast-gift-eval") image = ( @@ -28,6 +38,10 @@ def _resolve_paths() -> tuple[Path, Path]: ) .apt_install("git") .pip_install("uv") + .run_commands( + "uv pip install --system --compile-bytecode " + "'timecopilot-gift-eval>=0.3.1' modal pyyaml s3fs typer", + ) .add_local_file( _REPO_ROOT / "pyproject.toml", remote_path=f"{_MODAL_MONOREPO}/pyproject.toml", @@ -38,25 +52,38 @@ def _resolve_paths() -> tuple[Path, Path]: remote_path=f"{_MODAL_MONOREPO}/README.md", copy=True, ) + .add_local_dir( + _REPO_ROOT / "foundationforecast", + remote_path=f"{_MODAL_MONOREPO}/foundationforecast", + copy=True, + ) + .run_commands( + "uv pip install --system --compile-bytecode -e /root/monorepo", + ) + .add_local_file( + _GIFT_EVAL_ROOT / "pyproject.toml", + remote_path=f"{_MODAL_GIFT_EVAL}/pyproject.toml", + copy=True, + ) .add_local_file( - _REPO_ROOT / "uv.lock", - remote_path=f"{_MODAL_MONOREPO}/uv.lock", + _GIFT_EVAL_ROOT / "README.md", + remote_path=f"{_MODAL_GIFT_EVAL}/README.md", copy=True, ) .add_local_dir( - _REPO_ROOT / "foundationforecast", - remote_path=f"{_MODAL_MONOREPO}/foundationforecast", + _GIFT_EVAL_ROOT / "src", + remote_path=f"{_MODAL_GIFT_EVAL}/src", copy=True, ) .add_local_dir( - _GIFT_EVAL_ROOT, - remote_path=_MODAL_GIFT_EVAL, + _GIFT_EVAL_ROOT / "configs", + remote_path=f"{_MODAL_GIFT_EVAL}/configs", copy=True, ) .workdir(_MODAL_GIFT_EVAL) .env({"PYTHONPATH": _MODAL_GIFT_EVAL}) .run_commands( - "uv pip install --system --compile-bytecode -e .", + "uv pip install --system --no-deps --compile-bytecode -e .", ) ) secret = modal.Secret.from_name( @@ -79,6 +106,10 @@ def _resolve_paths() -> tuple[Path, Path]: S3_CI_RESULTS_PREFIX = "results/ci" +def replication_s3_prefix(run_id: str) -> str: + return f"results/replication/{run_id}" + + @app.function( image=image, volumes=volume, @@ -94,6 +125,7 @@ def run_gift_eval_modal( storage_path: str = "/s3-bucket/data/gift-eval", output_root: str = "/s3-bucket/results", force: bool = False, + registry: str = "default", ) -> None: import logging from pathlib import Path @@ -114,6 +146,7 @@ def run_gift_eval_modal( storage_path=storage_path, output_root=Path(output_root), overwrite_results=force, + registry=registry, # type: ignore[arg-type] ) @@ -127,25 +160,41 @@ def _dispatch_jobs( storage_path: str, output_root: str, force: bool, + registry: str = "default", + max_containers: int | None = None, ) -> None: logging.basicConfig(level=logging.INFO) if not jobs: logging.info("No jobs to run") return args = [ - (*job_tuple, storage_path, output_root, force) + (*job_tuple, storage_path, output_root, force, registry) for job_tuple in _job_tuples(jobs) ] + run_fn = run_gift_eval_modal + if max_containers is not None: + run_fn = run_gift_eval_modal.with_options(max_containers=max_containers) + logging.info("Modal max_containers=%s", max_containers) results = list( - run_gift_eval_modal.starmap( + run_fn.starmap( args, return_exceptions=True, - wrap_returned_exceptions=False, ) ) errors = [result for result in results if isinstance(result, Exception)] if errors: - raise RuntimeError(f"Modal jobs failed: {errors}") + for exc in errors[:10]: + logging.error("Job failed: %s", exc) + if len(errors) > 10: + logging.error("... and %s more failures", len(errors) - 10) + logging.info( + "Modal batch finished: ok=%s failed=%s total=%s", + len(results) - len(errors), + len(errors), + len(results), + ) + if errors and len(errors) == len(results): + raise RuntimeError(f"All Modal jobs failed ({len(errors)} jobs)") def run_ci_modal( @@ -157,31 +206,6 @@ def run_ci_modal( _dispatch_jobs(jobs, storage_path=storage_path, output_root=output_root, force=True) -def _s3_job_paths( - job, - *, - bucket: str, - prefix: str, -) -> tuple[str, str]: - base = f"s3://{bucket}/{prefix}/{job.model_key}/{job.dataset_name}/{job.term}" - return f"{base}/all_results.csv", f"{base}/timing.json" - - -def _job_matches_mode( - *, - mode: str, - has_results: bool, - has_timing: bool, -) -> bool: - if mode == "missing": - return not has_results - if mode == "missing_timing": - return has_results and not has_timing - if mode == "all": - return True - raise ValueError(f"Unknown job selection mode: {mode!r}") - - def _jobs_from_s3( jobs: list, *, @@ -189,21 +213,14 @@ def _jobs_from_s3( prefix: str, mode: str, ) -> list: - import fsspec - - fs = fsspec.filesystem("s3") - selected = [] - for job in jobs: - results_path, timing_path = _s3_job_paths(job, bucket=bucket, prefix=prefix) - has_results = fs.exists(results_path) - has_timing = fs.exists(timing_path) - if _job_matches_mode( - mode=mode, - has_results=has_results, - has_timing=has_timing, - ): - selected.append(job) - return selected + from src.eval.jobs import filter_jobs_by_s3_mode + + return filter_jobs_by_s3_mode( + jobs, + bucket=bucket, + prefix=prefix, + mode=mode, + ) @app.local_entrypoint() @@ -255,3 +272,65 @@ def run_missing_timing() -> None: output_root=f"/s3-bucket/{S3_RESULTS_PREFIX}", force=True, ) + + +@app.local_entrypoint() +def run_replication_pilot(run_id: str, force: bool = True) -> None: + from src.eval.jobs import load_replication_pilot_jobs + + prefix = replication_s3_prefix(run_id) + jobs = load_replication_pilot_jobs() + logging.info( + "Replication pilot: %s jobs → s3://%s/%s", + len(jobs), + S3_BUCKET, + prefix, + ) + _dispatch_jobs( + jobs, + storage_path="/s3-bucket/data/gift-eval", + output_root=f"/s3-bucket/{prefix}", + force=force, + registry="replication", + ) + + +@app.local_entrypoint() +def run_replication_full( + run_id: str, + force: bool = False, + model_key: str = "", + max_containers: int = 20, +) -> None: + from src.eval.jobs import load_replication_matrix + + prefix = replication_s3_prefix(run_id) + jobs = load_replication_matrix() + if model_key: + jobs = [job for job in jobs if job.model_key == model_key] + if force: + selected = jobs + else: + selected = _jobs_from_s3( + jobs, + bucket=S3_BUCKET, + prefix=prefix, + mode="missing", + ) + logging.basicConfig(level=logging.INFO) + logging.info( + "Replication full grid: %s jobs (force=%s, max_containers=%s) → s3://%s/%s", + len(selected), + force, + max_containers, + S3_BUCKET, + prefix, + ) + _dispatch_jobs( + selected, + storage_path="/s3-bucket/data/gift-eval", + output_root=f"/s3-bucket/{prefix}", + force=force, + registry="replication", + max_containers=max_containers, + ) diff --git a/experiments/gift-eval/src/runners/run_model.py b/experiments/gift-eval/src/runners/run_model.py index eb4cd2d..10c8b94 100644 --- a/experiments/gift-eval/src/runners/run_model.py +++ b/experiments/gift-eval/src/runners/run_model.py @@ -26,9 +26,21 @@ def main( Path, typer.Option(help="Path to downloaded GIFT-Eval dataset"), ] = Path("data/gift-eval"), + replication: Annotated[ + bool, + typer.Option( + help="Use configs/replication/models.yaml (notebook-aligned params)" + ), + ] = False, ) -> None: job = Job(model_key=model_key, dataset_name=dataset_name, term=term) - run_gift_eval(job, storage_path=storage_path, output_root=output_root) + registry = "replication" if replication else "default" + run_gift_eval( + job, + storage_path=storage_path, + output_root=output_root, + registry=registry, # type: ignore[arg-type] + ) if __name__ == "__main__": diff --git a/experiments/gift-eval/src/runners/run_verify.py b/experiments/gift-eval/src/runners/run_verify.py index 513cf89..a4e55d4 100644 --- a/experiments/gift-eval/src/runners/run_verify.py +++ b/experiments/gift-eval/src/runners/run_verify.py @@ -6,14 +6,18 @@ import typer -from src.eval.jobs import ci_output_root, load_ci_subset -from src.verify.replication_table import write_replication_table -from src.verify.verify import ( - ReplicationSkip, - model_keys_with_reference, - verify_all, - verify_model, +from src.eval.jobs import ( + ci_output_root, + load_ci_subset, + load_replication_families, + load_replication_matrix, + load_replication_pilot_jobs, + replication_output_root, ) +from src.eval.models import model_keys_with_reference +from src.verify.replication_aggregate import verify_replication_aggregates +from src.verify.replication_table import write_replication_table +from src.verify.verify import ReplicationSkip, verify_all, verify_model logging.basicConfig(level=logging.INFO) app = typer.Typer() @@ -26,22 +30,43 @@ def _resolve_model_keys( model_key: str | None, all_models: bool, ci: bool, -) -> tuple[list[str], Path]: + replication: bool, + run_id: str | None, + replication_pilot: bool, +) -> tuple[list[str], Path, list | None]: + if replication: + if not run_id: + raise typer.BadParameter("--replication requires --run-id RUN_ID") + root = replication_output_root(run_id) + if replication_pilot: + jobs = load_replication_pilot_jobs() + elif all_models or not model_key: + jobs = load_replication_matrix() + else: + jobs = None + if model_key: + keys = [model_key] + elif replication_pilot: + keys = sorted({job.model_key for job in jobs or []}) + else: + keys = list(load_replication_families()) + return keys, root, jobs + if ci: jobs = load_ci_subset() - return sorted({job.model_key for job in jobs}), ci_output_root() + return sorted({job.model_key for job in jobs}), ci_output_root(), jobs if model_key and all_models: raise typer.BadParameter("Use either --model-key or --all, not both") if model_key: - return [model_key], Path("results") + return [model_key], Path("results"), None if all_models: - return model_keys_with_reference(), Path("results") + return model_keys_with_reference(), Path("results"), None raise typer.BadParameter( - "Specify --model-key KEY, --all, or --ci. " + "Specify --model-key KEY, --all, --ci, or --replication --run-id. " "Example: uv run python -m src.runners.run_verify --all" ) @@ -51,6 +76,7 @@ def _run_verify_models( output_root: Path, *, require_complete: bool, + registry: str = "default", ) -> None: passed: list[str] = [] skipped: list[tuple[str, str]] = [] @@ -62,6 +88,7 @@ def _run_verify_models( key, output_root, require_complete=require_complete, + registry=registry, # type: ignore[arg-type] ) passed.append(key) except ReplicationSkip as exc: @@ -97,6 +124,20 @@ def main( bool, typer.Option(help="Verify CI subset jobs (per-job layout under results/ci)"), ] = False, + replication: Annotated[ + bool, + typer.Option(help="Verify replication run under results/replication/RUN_ID"), + ] = False, + run_id: Annotated[ + str | None, + typer.Option(help="Replication run id (required with --replication)"), + ] = None, + replication_pilot: Annotated[ + bool, + typer.Option( + help="With --replication, verify pilot jobs only (m4_weekly/short)" + ), + ] = False, output_root: Annotated[ Path | None, typer.Option(help="Root directory containing benchmark outputs"), @@ -116,23 +157,46 @@ def main( ), ] = False, ) -> None: - model_keys, default_root = _resolve_model_keys( + model_keys, default_root, replication_jobs = _resolve_model_keys( model_key=model_key, all_models=all_models, ci=ci, + replication=replication, + run_id=run_id, + replication_pilot=replication_pilot, ) resolved_output_root = output_root or default_root + registry = "replication" if replication else "default" if ci and not verify_only: verify_all(load_ci_subset(), resolved_output_root) + elif replication and replication_pilot and not verify_only: + verify_all( + replication_jobs or load_replication_pilot_jobs(), + resolved_output_root, + registry=registry, # type: ignore[arg-type] + ) + elif replication and not replication_pilot and not verify_only: + verify_replication_aggregates( + model_keys, + resolved_output_root, + require_complete=require_complete, + registry=registry, # type: ignore[arg-type] + ) elif not verify_only: _run_verify_models( model_keys, resolved_output_root, require_complete=require_complete, + registry=registry, ) - write_replication_table(model_keys, resolved_output_root, table_output) + write_replication_table( + model_keys, + resolved_output_root, + table_output, + registry=registry, # type: ignore[arg-type] + ) if __name__ == "__main__": diff --git a/experiments/gift-eval/src/verify/replication_aggregate.py b/experiments/gift-eval/src/verify/replication_aggregate.py new file mode 100644 index 0000000..bea63db --- /dev/null +++ b/experiments/gift-eval/src/verify/replication_aggregate.py @@ -0,0 +1,206 @@ +"""Leaderboard-style aggregate replication checks (experiment grid only).""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from functools import lru_cache + +import numpy as np +import pandas as pd +from numpy.testing import assert_allclose +from scipy import stats + +from src.eval.jobs import ( + REPLICATION_PENDING_HF_REFERENCE, + Job, + load_replication_matrix, + result_csv, +) +from src.eval.models import Registry, reference_slug +from src.verify.reference import CRPS_COL, MASE_COL, load_reference_results +from src.verify.verify import ReplicationSkip, load_actual_results + +logger = logging.getLogger(__name__) + +SEASONAL_NAIVE_REFERENCE_SLUG = "Seasonal_Naive" + +# Geomean(MASE/SN, CRPS/SN) vs the HF submission on the same dataset set. +# Slightly looser than per-job verify is unnecessary at aggregate level; same +# rtol keeps the experiment aligned with CI semantics while tolerating outlier jobs. +REPLICATION_AGGREGATE_ATOL = 1e-3 +REPLICATION_AGGREGATE_RTOL = 2.5e-2 + + +@dataclass(frozen=True) +class AggregateScores: + norm_mase: float + norm_crps: float + n_configs: int + + +@lru_cache +def seasonal_naive_by_dataset() -> pd.DataFrame: + sn = load_reference_results(SEASONAL_NAIVE_REFERENCE_SLUG) + return sn.rename( + columns={ + MASE_COL: "sn_mase", + CRPS_COL: "sn_crps", + } + )[["dataset", "sn_mase", "sn_crps"]] + + +def geomean_positive(series: pd.Series) -> float: + values = series.astype(float) + values = values[values > 0] + if values.empty: + return float("nan") + return float(stats.gmean(values)) + + +def leaderboard_aggregate(df: pd.DataFrame, sn: pd.DataFrame) -> AggregateScores: + merged = df.merge(sn, on="dataset", how="inner") + if merged.empty: + return AggregateScores(float("nan"), float("nan"), 0) + norm_mase = merged[MASE_COL] / merged["sn_mase"] + norm_crps = merged[CRPS_COL] / merged["sn_crps"] + return AggregateScores( + geomean_positive(norm_mase), + geomean_positive(norm_crps), + len(merged), + ) + + +def missing_matrix_jobs(model_key: str, output_root) -> list[Job]: + from pathlib import Path + + root = Path(output_root) + missing: list[Job] = [] + for job in load_replication_matrix(): + if job.model_key != model_key: + continue + if not result_csv(job, root).exists(): + missing.append(job) + return missing + + +def verify_replication_aggregate( + model_key: str, + output_root, + *, + require_complete: bool = False, + registry: Registry = "replication", + atol: float = REPLICATION_AGGREGATE_ATOL, + rtol: float = REPLICATION_AGGREGATE_RTOL, +) -> AggregateScores: + from pathlib import Path + + root = Path(output_root) + slug = reference_slug(model_key, registry=registry) + if slug is None: + raise ReplicationSkip(f"No reference slug for model_key={model_key!r}") + + if require_complete: + missing = missing_matrix_jobs(model_key, root) + if missing: + sample = ", ".join(f"{j.dataset_name}/{j.term}" for j in missing[:5]) + suffix = "..." if len(missing) > 5 else "" + raise AssertionError( + f"{model_key}: missing {len(missing)} replication matrix jobs " + f"(e.g. {sample}{suffix})" + ) + + actual = load_actual_results(model_key, root) + expected = load_reference_results(slug) + datasets = sorted(set(actual["dataset"]) & set(expected["dataset"])) + if not datasets: + raise AssertionError(f"{model_key}: no overlapping datasets with HF reference") + + sn = seasonal_naive_by_dataset() + actual_sub = actual[actual["dataset"].isin(datasets)] + expected_sub = expected[expected["dataset"].isin(datasets)] + + rep = leaderboard_aggregate(actual_sub, sn) + hf = leaderboard_aggregate(expected_sub, sn) + if rep.n_configs != hf.n_configs or rep.n_configs != len(datasets): + raise AssertionError( + f"{model_key}: aggregate config count mismatch " + f"(rep={rep.n_configs}, hf={hf.n_configs}, overlap={len(datasets)})" + ) + + for label, rep_val, hf_val in ( + ("geomean(MASE/Seasonal_Naive)", rep.norm_mase, hf.norm_mase), + ("geomean(CRPS/Seasonal_Naive)", rep.norm_crps, hf.norm_crps), + ): + try: + assert_allclose( + np.array([rep_val], dtype=float), + np.array([hf_val], dtype=float), + atol=atol, + rtol=rtol, + ) + except AssertionError as exc: + delta_pct = 100 * (rep_val - hf_val) / hf_val if hf_val else float("nan") + raise AssertionError( + f"{model_key}: {label} differs from HF submission aggregate " + f"(rep={rep_val:.6g}, hf={hf_val:.6g}, delta={delta_pct:+.2f}%, " + f"n={rep.n_configs}, atol={atol}, rtol={rtol})" + ) from exc + + logger.info( + "%s: aggregate OK on %s configs — rep MASE/SN=%.4f CRPS/SN=%.4f " + "(HF %.4f / %.4f)", + model_key, + rep.n_configs, + rep.norm_mase, + rep.norm_crps, + hf.norm_mase, + hf.norm_crps, + ) + return rep + + +def verify_replication_aggregates( + model_keys: list[str], + output_root, + *, + require_complete: bool = False, + registry: Registry = "replication", + atol: float = REPLICATION_AGGREGATE_ATOL, + rtol: float = REPLICATION_AGGREGATE_RTOL, +) -> None: + passed: list[str] = [] + skipped: list[tuple[str, str]] = [] + failed: list[tuple[str, str]] = [] + + for key in model_keys: + if key in REPLICATION_PENDING_HF_REFERENCE: + skipped.append((key, "GIFT-Eval HF reference results not published yet")) + logger.warning("Skipped %s: pending HF reference", key) + continue + try: + verify_replication_aggregate( + key, + output_root, + require_complete=require_complete, + registry=registry, + atol=atol, + rtol=rtol, + ) + passed.append(key) + except ReplicationSkip as exc: + skipped.append((key, str(exc))) + logger.warning("Skipped %s: %s", key, exc) + except Exception as exc: + failed.append((key, str(exc))) + logger.error("Failed %s: %s", key, exc) + + logger.info( + "Aggregate verify summary: passed=%s skipped=%s failed=%s", + len(passed), + len(skipped), + len(failed), + ) + if failed: + details = "\n".join(f" {key}: {error}" for key, error in failed) + raise AssertionError(f"Aggregate replication verification failed:\n{details}") diff --git a/experiments/gift-eval/src/verify/replication_table.py b/experiments/gift-eval/src/verify/replication_table.py index 0362bd5..6bdd862 100644 --- a/experiments/gift-eval/src/verify/replication_table.py +++ b/experiments/gift-eval/src/verify/replication_table.py @@ -3,10 +3,12 @@ import json import logging from pathlib import Path +from urllib.error import HTTPError import pandas as pd -from src.eval.models import reference_slug +from src.eval.jobs import REPLICATION_PENDING_HF_REFERENCE +from src.eval.models import Registry, reference_slug from src.verify.reference import CRPS_COL, MASE_COL, load_reference_results logger = logging.getLogger(__name__) @@ -48,12 +50,17 @@ def load_timing_map(model_key: str, output_root: Path) -> dict[str, float]: def build_replication_table( model_keys: list[str], output_root: Path, + *, + registry: Registry = "default", ) -> pd.DataFrame: from src.verify.verify import load_actual_results rows: list[dict] = [] for model_key in model_keys: - slug = reference_slug(model_key) + if model_key in REPLICATION_PENDING_HF_REFERENCE: + logger.warning("Skipping %s: pending HF reference", model_key) + continue + slug = reference_slug(model_key, registry=registry) if slug is None: logger.warning("Skipping %s: no reference_slug", model_key) continue @@ -61,7 +68,7 @@ def build_replication_table( try: actual = load_actual_results(model_key, output_root) expected = load_reference_results(slug) - except FileNotFoundError as exc: + except (FileNotFoundError, HTTPError) as exc: logger.warning("Skipping %s: %s", model_key, exc) continue @@ -105,8 +112,10 @@ def write_replication_table( model_keys: list[str], output_root: Path, table_path: Path, + *, + registry: Registry = "default", ) -> pd.DataFrame: - table = build_replication_table(model_keys, output_root) + table = build_replication_table(model_keys, output_root, registry=registry) table_path.parent.mkdir(parents=True, exist_ok=True) table.to_csv(table_path, index=False) logger.info("Wrote replication table (%s rows) to %s", len(table), table_path) diff --git a/experiments/gift-eval/src/verify/verify.py b/experiments/gift-eval/src/verify/verify.py index 05ff8ba..eb1b578 100644 --- a/experiments/gift-eval/src/verify/verify.py +++ b/experiments/gift-eval/src/verify/verify.py @@ -7,7 +7,7 @@ from timecopilot_gift_eval import GIFTEval from src.eval.jobs import Job, result_csv -from src.eval.models import load_models_config, reference_slug +from src.eval.models import Registry, reference_slug from .reference import ( REPLICATION_ATOL, REPLICATION_METRIC_COLS, @@ -39,8 +39,9 @@ def verify_job( storage_path: Path | str | None = None, atol: float = REPLICATION_ATOL, rtol: float = REPLICATION_RTOL, + registry: Registry = "default", ) -> None: - slug = reference_slug(job.model_key) + slug = reference_slug(job.model_key, registry=registry) if slug is None: raise ReplicationSkip(f"No reference slug for model_key={job.model_key!r}") @@ -74,6 +75,7 @@ def verify_all( storage_path: Path | str | None = None, atol: float = REPLICATION_ATOL, rtol: float = REPLICATION_RTOL, + registry: Registry = "default", ) -> None: for job in jobs: verify_job( @@ -82,6 +84,7 @@ def verify_all( storage_path=storage_path, atol=atol, rtol=rtol, + registry=registry, ) @@ -110,8 +113,9 @@ def verify_model( atol: float = REPLICATION_ATOL, rtol: float = REPLICATION_RTOL, require_complete: bool = False, + registry: Registry = "default", ) -> None: - slug = reference_slug(model_key) + slug = reference_slug(model_key, registry=registry) if slug is None: raise ReplicationSkip(f"No reference slug for model_key={model_key!r}") @@ -145,12 +149,3 @@ def verify_model( len(expected), slug, ) - - -def model_keys_with_reference() -> list[str]: - models = load_models_config() - return [ - model_key - for model_key, spec in models.items() - if spec.get("reference_slug") is not None - ] diff --git a/experiments/gift-eval/tests/test_replication_aggregate.py b/experiments/gift-eval/tests/test_replication_aggregate.py new file mode 100644 index 0000000..c55644e --- /dev/null +++ b/experiments/gift-eval/tests/test_replication_aggregate.py @@ -0,0 +1,135 @@ +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +from scipy import stats +from src.eval.jobs import Job +from src.verify.reference import CRPS_COL, MASE_COL +from src.verify.replication_aggregate import ( + REPLICATION_AGGREGATE_RTOL, + leaderboard_aggregate, + verify_replication_aggregate, +) + + +def test_leaderboard_aggregate_geomean_normalized() -> None: + sn = pd.DataFrame( + { + "dataset": ["cfg/a", "cfg/b"], + "sn_mase": [2.0, 0.5], + "sn_crps": [4.0, 1.0], + } + ) + df = pd.DataFrame( + { + "dataset": ["cfg/a", "cfg/b"], + MASE_COL: [1.0, 0.5], + CRPS_COL: [2.0, 0.5], + } + ) + scores = leaderboard_aggregate(df, sn) + assert scores.n_configs == 2 + assert np.isclose(scores.norm_mase, stats.gmean([0.5, 1.0])) + assert np.isclose(scores.norm_crps, stats.gmean([0.5, 0.5])) + + +def test_verify_replication_aggregate_passes_with_outlier_jobs( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + model_key = "test--model" + datasets = ["cfg/a", "cfg/b"] + sn = pd.DataFrame( + { + "dataset": datasets, + "sn_mase": [1.0, 1.0], + "sn_crps": [1.0, 1.0], + } + ) + # HF submission: geomean MASE/SN = gmean(0.7, 0.7) = 0.7 + hf_rows = pd.DataFrame( + { + "dataset": datasets, + MASE_COL: [0.7, 0.7], + CRPS_COL: [0.5, 0.5], + } + ) + # Per-job values differ (>2.5% rtol) but geomean stays near 0.7 + rep_rows = pd.DataFrame( + { + "dataset": datasets, + MASE_COL: [0.85, 0.58], + CRPS_COL: [0.5, 0.5], + } + ) + + model_root = tmp_path / model_key / "m4_weekly" / "short" + model_root.mkdir(parents=True) + rep_rows.to_csv(model_root / "all_results.csv", index=False) + + monkeypatch.setattr( + "src.verify.replication_aggregate.reference_slug", + lambda _key, registry="replication": "hf-slug", + ) + monkeypatch.setattr( + "src.verify.replication_aggregate.load_reference_results", + lambda slug: hf_rows.copy(), + ) + monkeypatch.setattr( + "src.verify.replication_aggregate.seasonal_naive_by_dataset", + lambda: sn.copy(), + ) + monkeypatch.setattr( + "src.verify.replication_aggregate.load_replication_matrix", + lambda: [ + Job(model_key=model_key, dataset_name="m4_weekly", term="short"), + ], + ) + + verify_replication_aggregate( + model_key, + tmp_path, + require_complete=True, + rtol=REPLICATION_AGGREGATE_RTOL, + ) + + +def test_verify_replication_aggregate_fails_when_aggregate_drifts( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + model_key = "test--model" + datasets = ["cfg/a"] + sn = pd.DataFrame({"dataset": datasets, "sn_mase": [1.0], "sn_crps": [1.0]}) + hf_rows = pd.DataFrame({"dataset": datasets, MASE_COL: [0.5], CRPS_COL: [0.5]}) + rep_rows = pd.DataFrame({"dataset": datasets, MASE_COL: [0.8], CRPS_COL: [0.5]}) + + model_root = tmp_path / model_key / "m4_weekly" / "short" + model_root.mkdir(parents=True) + rep_rows.to_csv(model_root / "all_results.csv", index=False) + + monkeypatch.setattr( + "src.verify.replication_aggregate.reference_slug", + lambda _key, registry="replication": "hf-slug", + ) + monkeypatch.setattr( + "src.verify.replication_aggregate.load_reference_results", + lambda slug: hf_rows.copy(), + ) + monkeypatch.setattr( + "src.verify.replication_aggregate.seasonal_naive_by_dataset", + lambda: sn.copy(), + ) + monkeypatch.setattr( + "src.verify.replication_aggregate.load_replication_matrix", + lambda: [ + Job(model_key=model_key, dataset_name="m4_weekly", term="short"), + ], + ) + + with pytest.raises(AssertionError, match="MASE/Seasonal_Naive"): + verify_replication_aggregate( + model_key, + tmp_path, + require_complete=True, + rtol=REPLICATION_AGGREGATE_RTOL, + ) diff --git a/experiments/gift-eval/tests/test_replication_config.py b/experiments/gift-eval/tests/test_replication_config.py new file mode 100644 index 0000000..441c315 --- /dev/null +++ b/experiments/gift-eval/tests/test_replication_config.py @@ -0,0 +1,84 @@ +import pytest +from src.eval.jobs import ( + REPLICATION_PENDING_HF_REFERENCE, + load_replication_families, + load_replication_pilot_jobs, +) +from src.eval.models import build_model, reference_slug +from src.verify.reference import load_reference_results + + +@pytest.mark.parametrize("model_key", load_replication_families()) +def test_replication_alias_matches_hf_reference_model_column(model_key: str) -> None: + if model_key in REPLICATION_PENDING_HF_REFERENCE: + pytest.skip("GIFT-Eval HF reference results not published yet") + slug = reference_slug(model_key, registry="replication") + assert slug is not None, f"missing reference_slug for {model_key}" + expected_name = load_reference_results(slug)["model"].iloc[0] + forecaster = build_model(model_key, registry="replication") + assert forecaster.alias == expected_name + + +def test_replication_pilot_job_count() -> None: + families = load_replication_families() + pilot = load_replication_pilot_jobs() + assert len(pilot) == len(families) + assert {j.model_key for j in pilot} == set(families) + + +def test_timesfm3_replication_context_length() -> None: + from src.eval.jobs import load_replication_models_config + + spec = load_replication_models_config()["google--timesfm-3.0-pytorch"] + assert spec["kwargs"]["context_length"] == 15360 + + +def test_moirai2_replication_context_and_batch() -> None: + from src.eval.jobs import load_replication_models_config + + spec = load_replication_models_config()["Salesforce--moirai-2.0-R-small"] + assert spec["kwargs"]["context_length"] == 4000 + assert spec["kwargs"]["batch_size"] == 32 + assert spec["predictor_batch_size"] == 512 + + +def test_chronos2_replication_notebook_runner_params() -> None: + from src.eval.jobs import load_replication_models_config + + spec = load_replication_models_config()["amazon--chronos-2"] + assert spec["reference_slug"] == "Chronos-2" + assert spec["predictor_batch_size"] == 16 + assert spec["kwargs"]["batch_size"] == 16 + assert spec["max_length"] is None + + +def test_toto_replication_batch_size() -> None: + from src.eval.jobs import load_replication_models_config + + spec = load_replication_models_config()["Datadog--Toto-2.0-313m"] + assert spec["kwargs"]["batch_size"] == 512 + + +def test_replication_families_count() -> None: + assert len(load_replication_families()) == 16 + + +def test_chronos2_small_replication_params() -> None: + from src.eval.jobs import load_replication_models_config + + spec = load_replication_models_config()["autogluon--chronos-2-small"] + assert spec["kwargs"]["repo_id"] == "autogluon/chronos-2-small" + assert spec["kwargs"]["batch_size"] == 16 + assert spec["max_length"] is None + assert spec["predictor_batch_size"] == 16 + + +def test_moirai11_replication_context_and_batch() -> None: + from src.eval.jobs import load_replication_models_config + + large = load_replication_models_config()[ + "Salesforce--moirai-1.1-R-large--Moirai_large" + ] + assert large["kwargs"]["context_length"] == 4000 + assert large["kwargs"]["batch_size"] == 16 + assert large["predictor_batch_size"] == 512 diff --git a/experiments/gift-eval/tests/test_replication_table.py b/experiments/gift-eval/tests/test_replication_table.py new file mode 100644 index 0000000..7c95deb --- /dev/null +++ b/experiments/gift-eval/tests/test_replication_table.py @@ -0,0 +1,47 @@ +from pathlib import Path +from urllib.error import HTTPError + +import pytest +from src.eval.jobs import REPLICATION_PENDING_HF_REFERENCE +from src.verify import replication_table +from src.verify.replication_table import ( + REPLICATION_TABLE_COLS, + build_replication_table, +) + + +def test_build_replication_table_skips_pending_hf_reference( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pending_key = next(iter(REPLICATION_PENDING_HF_REFERENCE)) + + def _no_network(slug: str): + raise AssertionError(f"reference download attempted for {slug}") + + monkeypatch.setattr(replication_table, "load_reference_results", _no_network) + monkeypatch.setattr( + "src.verify.verify.load_actual_results", lambda key, root: object() + ) + + table = build_replication_table([pending_key], tmp_path, registry="replication") + + assert table.empty + assert list(table.columns) == REPLICATION_TABLE_COLS + + +def test_build_replication_table_skips_missing_remote_reference( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + model_key = "amazon--chronos-2" + + def _not_found(slug: str): + raise HTTPError("https://example/results", 404, "Not Found", {}, None) # type: ignore[arg-type] + + monkeypatch.setattr(replication_table, "load_reference_results", _not_found) + monkeypatch.setattr( + "src.verify.verify.load_actual_results", lambda key, root: object() + ) + + table = build_replication_table([model_key], tmp_path, registry="replication") + + assert table.empty diff --git a/experiments/gift-eval/tests/test_s3_job_selection.py b/experiments/gift-eval/tests/test_s3_job_selection.py new file mode 100644 index 0000000..fcddc37 --- /dev/null +++ b/experiments/gift-eval/tests/test_s3_job_selection.py @@ -0,0 +1,21 @@ +from src.eval.jobs import parse_job_dir_suffix + + +def test_parse_job_dir_suffix_simple_dataset() -> None: + assert parse_job_dir_suffix("amazon--chronos-2/m4_weekly/short") == ( + "amazon--chronos-2", + "m4_weekly", + "short", + ) + + +def test_parse_job_dir_suffix_slash_in_dataset() -> None: + assert parse_job_dir_suffix("amazon--chronos-2/electricity/15T/short") == ( + "amazon--chronos-2", + "electricity/15T", + "short", + ) + + +def test_parse_job_dir_suffix_too_short() -> None: + assert parse_job_dir_suffix("only/two") is None diff --git a/foundationforecast/models/chronos.py b/foundationforecast/models/chronos.py index 650f693..7ffd1e5 100644 --- a/foundationforecast/models/chronos.py +++ b/foundationforecast/models/chronos.py @@ -64,6 +64,47 @@ class ChronosFinetuningConfig: save_path: str | Path | None = None +class _QuantileBatchPredictor: + """Call ``model.predict_quantiles`` per batch, backing off on CUDA OOM. + + When ``adaptive_batch_size`` is set (Chronos-2 pipelines accept a + ``batch_size`` inference argument), a ``torch.cuda.OutOfMemoryError`` halves + the inference batch size and retries. The reduced size is kept for later + batches. Other pipelines, or a batch size that cannot be halved further, + re-raise the error. + """ + + def __init__( + self, + model: Any, + *, + prediction_length: int, + quantile_levels: list[float], + batch_size: int, + adaptive_batch_size: bool, + ) -> None: + self.model = model + self.kwargs: dict[str, Any] = { + "prediction_length": prediction_length, + "quantile_levels": quantile_levels, + } + self.batch_size = batch_size + self.adaptive_batch_size = adaptive_batch_size + + def __call__(self, batch: Any) -> Any: + kwargs = dict(self.kwargs) + if self.adaptive_batch_size: + kwargs["batch_size"] = self.batch_size + while True: + try: + return self.model.predict_quantiles(batch, **kwargs) + except torch.cuda.OutOfMemoryError: + if not self.adaptive_batch_size or self.batch_size < 2: + raise + self.batch_size //= 2 + kwargs["batch_size"] = self.batch_size + + class Chronos(ExogCapableForecaster): """ Chronos models are large pre-trained models for time series forecasting, @@ -340,14 +381,14 @@ def _predict( q_max=q_max, include_median=False, ) - fcsts = [ - model.predict_quantiles( - batch, - prediction_length=h, - quantile_levels=backend_qs, - ) - for batch in tqdm(dataset) - ] # list of tuples + predict_quantiles = _QuantileBatchPredictor( + model, + prediction_length=h, + quantile_levels=backend_qs, + batch_size=self.batch_size, + adaptive_batch_size=isinstance(model, Chronos2Pipeline), + ) + fcsts = [predict_quantiles(batch) for batch in tqdm(dataset)] fcsts_quantiles, fcsts_mean = zip(*fcsts, strict=False) if isinstance(model, Chronos2Pipeline): fcsts_mean = [f_mean for fcst in fcsts_mean for f_mean in fcst] # type: ignore diff --git a/tests/models/test_chronos.py b/tests/models/test_chronos.py index 9eb8183..316aa9a 100644 --- a/tests/models/test_chronos.py +++ b/tests/models/test_chronos.py @@ -2,7 +2,11 @@ import torch from tests.helpers import generate_series -from foundationforecast.models.chronos import Chronos, ChronosFinetuningConfig +from foundationforecast.models.chronos import ( + Chronos, + ChronosFinetuningConfig, + _QuantileBatchPredictor, +) pytestmark = pytest.mark.models @@ -95,3 +99,74 @@ def test_chronos_lora_finetuning_save_and_reuse(tmp_path): assert not fcst_reuse.empty assert "Chronos" in fcst_reuse.columns assert len(fcst_reuse) == n_series * 2 + + +class _FakeOOMPipeline: + """Raises CUDA OOM while ``batch_size`` is above ``fits_at``.""" + + def __init__(self, fits_at: int | None = None, fail_times: int = 0) -> None: + self.fits_at = fits_at + self.fail_times = fail_times + self.calls: list[dict] = [] + + def predict_quantiles(self, batch, **kwargs): + self.calls.append(kwargs) + if self.fits_at is not None and kwargs.get("batch_size", 0) > self.fits_at: + raise torch.cuda.OutOfMemoryError("fake OOM") + if self.fail_times > 0: + self.fail_times -= 1 + raise torch.cuda.OutOfMemoryError("fake OOM") + return ("quantiles", "mean", batch) + + +def _oom_predictor(model, *, batch_size: int, adaptive: bool): + return _QuantileBatchPredictor( + model, + prediction_length=3, + quantile_levels=[0.1, 0.5, 0.9], + batch_size=batch_size, + adaptive_batch_size=adaptive, + ) + + +def test_oom_backoff_halves_batch_size_until_it_fits_and_keeps_it(): + model = _FakeOOMPipeline(fits_at=4) + predict = _oom_predictor(model, batch_size=16, adaptive=True) + + assert predict("b0") == ("quantiles", "mean", "b0") + assert [c["batch_size"] for c in model.calls] == [16, 8, 4] + assert predict.batch_size == 4 + + # Later batches start from the reduced size without retrying. + assert predict("b1") == ("quantiles", "mean", "b1") + assert model.calls[-1]["batch_size"] == 4 + assert len(model.calls) == 4 + + +def test_oom_backoff_passes_prediction_kwargs(): + model = _FakeOOMPipeline() + predict = _oom_predictor(model, batch_size=8, adaptive=True) + predict("b0") + assert model.calls == [ + {"prediction_length": 3, "quantile_levels": [0.1, 0.5, 0.9], "batch_size": 8} + ] + + +def test_oom_backoff_reraises_when_batch_size_cannot_be_halved(): + model = _FakeOOMPipeline(fits_at=0) + predict = _oom_predictor(model, batch_size=2, adaptive=True) + with pytest.raises(torch.cuda.OutOfMemoryError): + predict("b0") + # 2 → 1, then OOM at 1 is fatal. + assert [c["batch_size"] for c in model.calls] == [2, 1] + assert predict.batch_size == 1 + + +def test_oom_backoff_non_adaptive_pipeline_reraises_without_batch_size(): + model = _FakeOOMPipeline(fail_times=1) + predict = _oom_predictor(model, batch_size=16, adaptive=False) + with pytest.raises(torch.cuda.OutOfMemoryError): + predict("b0") + assert len(model.calls) == 1 + assert "batch_size" not in model.calls[0] + assert predict.batch_size == 16