diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 094872c..26b9d26 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -35,7 +35,7 @@ jobs: uses: CodSpeedHQ/action@v5 with: mode: simulation - run: uv run pytest tests/benchmarks/test_core.py --codspeed -m benchmark -o addopts= + run: uv run pytest tests/benchmarks/test_core.py tests/benchmarks/test_orchestration.py --codspeed -m benchmark -o addopts= benchmarks-models: name: CodSpeed models diff --git a/tests/benchmarks/conftest.py b/tests/benchmarks/conftest.py index 885396e..c37e756 100644 --- a/tests/benchmarks/conftest.py +++ b/tests/benchmarks/conftest.py @@ -12,6 +12,16 @@ def panel_df(): ) +@pytest.fixture(scope="session") +def large_panel_df(): + return generate_series( + n_series=100, + freq="D", + min_length=500, + max_length=500, + ) + + @pytest.fixture(scope="session") def chronos_bolt(): from foundationforecast.models.chronos import Chronos diff --git a/tests/benchmarks/test_core.py b/tests/benchmarks/test_core.py index 65e30fc..daf259c 100644 --- a/tests/benchmarks/test_core.py +++ b/tests/benchmarks/test_core.py @@ -12,9 +12,16 @@ def test_maybe_infer_freq(benchmark, panel_df): assert result == "D" -def test_timeseries_dataset_from_df(benchmark, panel_df): +@pytest.mark.parametrize( + "panel_df_fixture", + ["panel_df", "large_panel_df"], + ids=["small", "large"], +) +def test_timeseries_dataset_from_df(benchmark, panel_df_fixture, request): + df = request.getfixturevalue(panel_df_fixture) + def build_dataset(): - return TimeSeriesDataset.from_df(panel_df, batch_size=4) + return TimeSeriesDataset.from_df(df, batch_size=4) dataset = benchmark(build_dataset) assert len(dataset) > 0 diff --git a/tests/benchmarks/test_orchestration.py b/tests/benchmarks/test_orchestration.py new file mode 100644 index 0000000..a35766c --- /dev/null +++ b/tests/benchmarks/test_orchestration.py @@ -0,0 +1,38 @@ +import pytest + +from foundationforecast import FoundationForecast +from foundationforecast.core.forecaster import Forecaster +from foundationforecast.core.utils import TimeSeriesDataset + +pytestmark = pytest.mark.benchmark + + +class DatasetTouchingModel(Forecaster): + """Lightweight model that exercises dataset construction like real forecasters.""" + + batch_size = 32 + + def __init__(self, alias: str = "DatasetTouch"): + self.alias = alias + + def forecast(self, df, h, freq=None, level=None, quantiles=None, **kwargs): + panel = kwargs.get("panel") + ds_kwargs = {"df": df, "batch_size": self.batch_size} + if panel is not None: + ds_kwargs["panel"] = panel + dataset = TimeSeriesDataset.from_df(**ds_kwargs) + fcst = dataset.make_future_dataframe(h=h, freq=freq or "D") + fcst[self.alias] = 1.0 + return fcst + + +@pytest.mark.parametrize("n_models", [3, 5], ids=["3-models", "5-models"]) +def test_foundation_forecast_multi_model(benchmark, large_panel_df, n_models): + models: list[Forecaster] = [ + DatasetTouchingModel(alias=f"DatasetTouch{i}") for i in range(n_models) + ] + forecaster = FoundationForecast(models=models) + result = benchmark(forecaster.forecast, large_panel_df, h=12, freq="D") + assert len(result) == large_panel_df["unique_id"].nunique() * 12 + for model in models: + assert model.alias in result.columns