feat: add CodSpeed benchmarks for core pipeline and models

.github/workflows/codspeed.yml

@@ -0,0 +1,72 @@
+name: CodSpeed
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+  workflow_dispatch:
+
+concurrency:
+  group: codspeed-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  benchmarks-core:
+    name: CodSpeed core
+    runs-on: ubuntu-latest
+    steps:
+      - name: Clone repo
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Set up uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true
+
+      - name: Install dependencies
+        run: uv sync --frozen
+
+      - name: Run core benchmarks
+        uses: CodSpeedHQ/action@v5
+        with:
+          mode: simulation
+          run: uv run pytest tests/benchmarks/test_core.py --codspeed -m benchmark -o addopts=
+
+  benchmarks-models:
+    name: CodSpeed models
+    runs-on: codspeed-macro
+    steps:
+      - name: Clone repo
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Set up uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true
+
+      - name: Cache HuggingFace Models
+        uses: actions/cache@v4
+        with:
+          path: ~/.cache/huggingface/hub/
+          key: ${{ runner.os }}-huggingface-codspeed-${{ hashFiles('tests/benchmarks/**') }}
+
+      - name: Install dependencies
+        run: uv sync --frozen
+
+      - name: Run model benchmarks
+        uses: CodSpeedHQ/action@v5
+        with:
+          mode: walltime
+          run: uv run pytest tests/benchmarks/test_models.py --codspeed -m benchmark -o addopts=
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}

pyproject.toml

@@ -6,6 +6,7 @@ requires = ["hatchling"]
 dev = [
   "mktestdocs>=0.2.5",
   "pre-commit",
+  "pytest-codspeed>=5.0.3",
   "pytest-cov>=6.0",
   "pytest-mock>=3.15.1",
   "pytest-rerunfailures>=15.1",
@@ -94,8 +95,9 @@ omit = ["tests/*"]
 packages = ["foundationforecast"]
 
 [tool.pytest.ini_options]
-addopts = "-m 'not docs'"
+addopts = "-m 'not docs and not benchmark'"
 markers = [
+  "benchmark: marks CodSpeed performance benchmarks",
   "docs: marks tests related to documentation",
   "models: marks tests that download model weights",
 ]

tests/benchmarks/conftest.py

@@ -0,0 +1,47 @@
+import pytest
+from tests.helpers import generate_series
+
+
+@pytest.fixture(scope="session")
+def panel_df():
+    return generate_series(
+        n_series=10,
+        freq="D",
+        min_length=100,
+        max_length=100,
+    )
+
+
+@pytest.fixture(scope="session")
+def chronos_bolt():
+    from foundationforecast.models.chronos import Chronos
+
+    return Chronos(repo_id="amazon/chronos-bolt-tiny", alias="Chronos-Bolt")
+
+
+@pytest.fixture(scope="session")
+def timesfm():
+    from foundationforecast.models.timesfm import TimesFM
+
+    return TimesFM(
+        repo_id="google/timesfm-1.0-200m-pytorch",
+        context_length=256,
+    )
+
+
+@pytest.fixture(scope="session")
+def toto():
+    from foundationforecast.models.toto import Toto
+
+    return Toto(context_length=256, batch_size=2)
+
+
+@pytest.fixture(scope="session")
+def moirai():
+    from foundationforecast.models.moirai import Moirai
+
+    return Moirai(
+        context_length=256,
+        batch_size=2,
+        repo_id="Salesforce/moirai-1.1-R-small",
+    )

tests/benchmarks/test_core.py

@@ -0,0 +1,38 @@
+import pytest
+from tests.helpers import SeasonalNaiveModel
+
+from foundationforecast.core.forecaster import QuantileConverter, maybe_infer_freq
+from foundationforecast.core.utils import TimeSeriesDataset
+
+pytestmark = pytest.mark.benchmark
+
+
+def test_maybe_infer_freq(benchmark, panel_df):
+    result = benchmark(maybe_infer_freq, panel_df, None)
+    assert result == "D"
+
+
+def test_timeseries_dataset_from_df(benchmark, panel_df):
+    def build_dataset():
+        return TimeSeriesDataset.from_df(panel_df, batch_size=4)
+
+    dataset = benchmark(build_dataset)
+    assert len(dataset) > 0
+    assert len(next(iter(dataset))) <= 4
+
+
+def test_quantile_converter_level_to_quantiles(benchmark):
+    def convert():
+        qc = QuantileConverter(level=[80, 95])
+        return qc.quantiles
+
+    quantiles = benchmark(convert)
+    assert quantiles is not None
+    assert len(quantiles) > 0
+
+
+def test_seasonal_naive_forecast(benchmark, panel_df):
+    model = SeasonalNaiveModel()
+    result = benchmark(model.forecast, panel_df, h=12, freq="D")
+    assert len(result) == panel_df["unique_id"].nunique() * 12
+    assert "SeasonalNaive" in result.columns

tests/benchmarks/test_models.py

@@ -0,0 +1,19 @@
+import pytest
+
+pytestmark = [pytest.mark.benchmark, pytest.mark.models]
+
+
+@pytest.mark.parametrize(
+    "model_fixture,expected_alias",
+    [
+        pytest.param("chronos_bolt", "Chronos-Bolt", id="chronos-bolt"),
+        pytest.param("timesfm", "TimesFM", id="timesfm-1"),
+        pytest.param("toto", "Toto", id="toto"),
+        pytest.param("moirai", "Moirai", id="moirai-1.1"),
+    ],
+)
+def test_model_forecast(benchmark, model_fixture, expected_alias, panel_df, request):
+    model = request.getfixturevalue(model_fixture)
+    result = benchmark(model.forecast, panel_df, h=12, freq="D")
+    assert len(result) == panel_df["unique_id"].nunique() * 12
+    assert expected_alias in result.columns