autogluon.timeseries 1.2.1b20250224__py3-none-any.whl → 1.4.1b20251215__py3-none-any.whl

autogluon/timeseries/models/gluonts/dataset.py

@@ -0,0 +1,109 @@
+from typing import Any, Iterator, Type
+
+import numpy as np
+import pandas as pd
+from gluonts.dataset.common import Dataset as GluonTSDataset
+from gluonts.dataset.field_names import FieldName
+
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
+from autogluon.timeseries.utils.datetime import norm_freq_str
+
+
+class SimpleGluonTSDataset(GluonTSDataset):
+    """Wrapper for TimeSeriesDataFrame that is compatible with the GluonTS Dataset API."""
+
+    def __init__(
+        self,
+        target_df: TimeSeriesDataFrame,
+        freq: str,
+        target_column: str = "target",
+        feat_static_cat: np.ndarray | None = None,
+        feat_static_real: np.ndarray | None = None,
+        feat_dynamic_cat: np.ndarray | None = None,
+        feat_dynamic_real: np.ndarray | None = None,
+        past_feat_dynamic_cat: np.ndarray | None = None,
+        past_feat_dynamic_real: np.ndarray | None = None,
+        includes_future: bool = False,
+        prediction_length: int | None = None,
+    ):
+        assert target_df is not None
+        # Convert TimeSeriesDataFrame to pd.Series for faster processing
+        self.target_array = target_df[target_column].to_numpy(np.float32)
+        self.feat_static_cat = self._astype(feat_static_cat, dtype=np.int64)
+        self.feat_static_real = self._astype(feat_static_real, dtype=np.float32)
+        self.feat_dynamic_cat = self._astype(feat_dynamic_cat, dtype=np.int64)
+        self.feat_dynamic_real = self._astype(feat_dynamic_real, dtype=np.float32)
+        self.past_feat_dynamic_cat = self._astype(past_feat_dynamic_cat, dtype=np.int64)
+        self.past_feat_dynamic_real = self._astype(past_feat_dynamic_real, dtype=np.float32)
+        self.freq = self._get_freq_for_period(freq)
+
+        # Necessary to compute indptr for known_covariates at prediction time
+        self.includes_future = includes_future
+        self.prediction_length = prediction_length
+
+        # Replace inefficient groupby ITEMID with indptr that stores start:end of each time series
+        self.item_ids = target_df.item_ids
+        self.indptr = target_df.get_indptr()
+        self.start_timestamps = target_df.index[self.indptr[:-1]].to_frame(index=False)[TimeSeriesDataFrame.TIMESTAMP]
+        assert len(self.item_ids) == len(self.start_timestamps)
+
+    @staticmethod
+    def _astype(array: np.ndarray | None, dtype: Type[np.generic]) -> np.ndarray | None:
+        if array is None:
+            return None
+        else:
+            return array.astype(dtype)
+
+    @staticmethod
+    def _get_freq_for_period(freq: str) -> str:
+        """Convert freq to format compatible with pd.Period.
+
+        For example, ME freq must be converted to M when creating a pd.Period.
+        """
+        offset = pd.tseries.frequencies.to_offset(freq)
+        assert offset is not None
+        freq_name = norm_freq_str(offset)
+        if freq_name == "SME":
+            # Replace unsupported frequency "SME" with "2W"
+            return "2W"
+        elif freq_name == "bh":
+            # Replace unsupported frequency "bh" with dummy value "Y"
+            return "Y"
+        else:
+            freq_name_for_period = {"YE": "Y", "QE": "Q", "ME": "M"}.get(freq_name, freq_name)
+            return f"{offset.n}{freq_name_for_period}"
+
+    def __len__(self):
+        return len(self.indptr) - 1  # noqa
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        for j in range(len(self.indptr) - 1):
+            start_idx = self.indptr[j]
+            end_idx = self.indptr[j + 1]
+            # GluonTS expects item_id to be a string
+            ts = {
+                FieldName.ITEM_ID: str(self.item_ids[j]),
+                FieldName.START: pd.Period(self.start_timestamps.iloc[j], freq=self.freq),
+                FieldName.TARGET: self.target_array[start_idx:end_idx],
+            }
+            if self.feat_static_cat is not None:
+                ts[FieldName.FEAT_STATIC_CAT] = self.feat_static_cat[j]
+            if self.feat_static_real is not None:
+                ts[FieldName.FEAT_STATIC_REAL] = self.feat_static_real[j]
+            if self.past_feat_dynamic_cat is not None:
+                ts[FieldName.PAST_FEAT_DYNAMIC_CAT] = self.past_feat_dynamic_cat[start_idx:end_idx].T
+            if self.past_feat_dynamic_real is not None:
+                ts[FieldName.PAST_FEAT_DYNAMIC_REAL] = self.past_feat_dynamic_real[start_idx:end_idx].T
+
+            # Dynamic features that may extend into the future
+            if self.includes_future:
+                assert self.prediction_length is not None, (
+                    "Prediction length must be provided if includes_future is True"
+                )
+                start_idx = start_idx + j * self.prediction_length
+                end_idx = end_idx + (j + 1) * self.prediction_length
+            if self.feat_dynamic_cat is not None:
+                ts[FieldName.FEAT_DYNAMIC_CAT] = self.feat_dynamic_cat[start_idx:end_idx].T
+            if self.feat_dynamic_real is not None:
+                ts[FieldName.FEAT_DYNAMIC_REAL] = self.feat_dynamic_real[start_idx:end_idx].T
+            yield ts

autogluon/timeseries/models/gluonts/{torch/models.py → models.py}

@@ -3,17 +3,18 @@ Module including wrappers for PyTorch implementations of models in GluonTS
 """
 
 import logging
-from typing import Any, Dict, Type
+from typing import Any, Type
 
 from gluonts.model.estimator import Estimator as GluonTSEstimator
 
-from autogluon.timeseries.models.gluonts.abstract_gluonts import AbstractGluonTSModel
 from autogluon.timeseries.utils.datetime import (
     get_lags_for_frequency,
     get_seasonality,
     get_time_features_for_frequency,
 )
 
+from .abstract import AbstractGluonTSModel
+
 # NOTE: We avoid imports for torch and lightning.pytorch at the top level and hide them inside class methods.
 # This is done to skip these imports during multiprocessing (which may cause bugs)
 
@@ -59,7 +60,7 @@ class DeepARModel(AbstractGluonTSModel):
         Distribution output object that defines how the model output is converted to a forecast, and how the loss is computed.
     scaling: bool, default = True
         If True, mean absolute scaling will be applied to each *context window* during training & prediction.
-        Note that this is different from the `target_scaler` that is applied to the *entire time series*.
+        Note that this is different from the ``target_scaler`` that is applied to the *entire time series*.
     max_epochs : int, default = 100
         Number of epochs the model will be trained for
     batch_size : int, default = 64
@@ -80,6 +81,8 @@ class DeepARModel(AbstractGluonTSModel):
 
     # TODO: Replace "scaling: bool" with "window_scaler": {"mean_abs", None} for consistency?
 
+    ag_priority = 40
+
     _supports_known_covariates = True
     _supports_static_features = True
 
@@ -88,13 +91,13 @@ class DeepARModel(AbstractGluonTSModel):
 
         return DeepAREstimator
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_static_cat"] = self.num_feat_static_cat
         init_kwargs["num_feat_static_real"] = self.num_feat_static_real
         init_kwargs["cardinality"] = self.feat_static_cat_cardinality
         init_kwargs["num_feat_dynamic_real"] = self.num_feat_dynamic_real
-        init_kwargs.setdefault("lags_seq", get_lags_for_frequency(self.freq))
+        init_kwargs.setdefault("lags_seq", get_lags_for_frequency(self.freq))  # type: ignore
         init_kwargs.setdefault("time_features", get_time_features_for_frequency(self.freq))
         return init_kwargs
 
@@ -110,7 +113,7 @@ class SimpleFeedForwardModel(AbstractGluonTSModel):
     ----------------
     context_length : int, default = max(10, 2 * prediction_length)
         Number of time units that condition the predictions
-    hidden_dimensions: List[int], default = [20, 20]
+    hidden_dimensions: list[int], default = [20, 20]
         Size of hidden layers in the feedforward network
     distr_output : gluonts.torch.distributions.Output, default = StudentTOutput()
         Distribution output object that defines how the model output is converted to a forecast, and how the loss is computed.
@@ -118,7 +121,7 @@ class SimpleFeedForwardModel(AbstractGluonTSModel):
         Whether to use batch normalization
     mean_scaling : bool, default = True
         If True, mean absolute scaling will be applied to each *context window* during training & prediction.
-        Note that this is different from the `target_scaler` that is applied to the *entire time series*.
+        Note that this is different from the ``target_scaler`` that is applied to the *entire time series*.
     max_epochs : int, default = 100
         Number of epochs the model will be trained for
     batch_size : int, default = 64
@@ -137,6 +140,8 @@ class SimpleFeedForwardModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
+    ag_priority = 10
+
     def _get_estimator_class(self) -> Type[GluonTSEstimator]:
         from gluonts.torch.model.simple_feedforward import SimpleFeedForwardEstimator
 
@@ -198,6 +203,9 @@ class TemporalFusionTransformerModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
+    ag_priority = 45
+    ag_model_aliases = ["TFT"]
+
     _supports_known_covariates = True
     _supports_past_covariates = True
     _supports_cat_covariates = True
@@ -208,12 +216,12 @@ class TemporalFusionTransformerModel(AbstractGluonTSModel):
 
         return TemporalFusionTransformerEstimator
 
-    def _get_default_params(self):
-        return super()._get_default_params() | {
+    def _get_default_hyperparameters(self):
+        return super()._get_default_hyperparameters() | {
             "context_length": min(512, max(64, 2 * self.prediction_length)),
         }
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         if self.num_feat_dynamic_real > 0:
             init_kwargs["dynamic_dims"] = [self.num_feat_dynamic_real]
@@ -260,7 +268,7 @@ class DLinearModel(AbstractGluonTSModel):
         Scaling applied to each *context window* during training & prediction.
         One of ``"mean"`` (mean absolute scaling), ``"std"`` (standardization), ``None`` (no scaling).
 
-        Note that this is different from the `target_scaler` that is applied to the *entire time series*.
+        Note that this is different from the ``target_scaler`` that is applied to the *entire time series*.
     max_epochs : int, default = 100
         Number of epochs the model will be trained for
     batch_size : int, default = 64
@@ -281,8 +289,10 @@ class DLinearModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
-    def _get_default_params(self):
-        return super()._get_default_params() | {
+    ag_priority = 10
+
+    def _get_default_hyperparameters(self):
+        return super()._get_default_hyperparameters() | {
             "context_length": 96,
         }
 
@@ -324,7 +334,7 @@ class PatchTSTModel(AbstractGluonTSModel):
         Scaling applied to each *context window* during training & prediction.
         One of ``"mean"`` (mean absolute scaling), ``"std"`` (standardization), ``None`` (no scaling).
 
-        Note that this is different from the `target_scaler` that is applied to the *entire time series*.
+        Note that this is different from the ``target_scaler`` that is applied to the *entire time series*.
     max_epochs : int, default = 100
         Number of epochs the model will be trained for
     batch_size : int, default = 64
@@ -339,6 +349,8 @@ class PatchTSTModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
+    ag_priority = 30
+
     _supports_known_covariates = True
 
     def _get_estimator_class(self) -> Type[GluonTSEstimator]:
@@ -346,10 +358,10 @@ class PatchTSTModel(AbstractGluonTSModel):
 
         return PatchTSTEstimator
 
-    def _get_default_params(self):
-        return super()._get_default_params() | {"context_length": 96, "patch_len": 16}
+    def _get_default_hyperparameters(self):
+        return super()._get_default_hyperparameters() | {"context_length": 96, "patch_len": 16}
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_dynamic_real"] = self.num_feat_dynamic_real
         return init_kwargs
@@ -415,6 +427,8 @@ class WaveNetModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
+    ag_priority = 25
+
     _supports_known_covariates = True
     _supports_static_features = True
     default_num_samples: int = 100
@@ -424,7 +438,7 @@ class WaveNetModel(AbstractGluonTSModel):
 
         return WaveNetEstimator
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_static_cat"] = self.num_feat_static_cat
         init_kwargs["num_feat_static_real"] = self.num_feat_static_real
@@ -488,7 +502,7 @@ class TiDEModel(AbstractGluonTSModel):
         Scaling applied to each *context window* during training & prediction.
         One of ``"mean"`` (mean absolute scaling), ``"std"`` (standardization), ``None`` (no scaling).
 
-        Note that this is different from the `target_scaler` that is applied to the *entire time series*.
+        Note that this is different from the ``target_scaler`` that is applied to the *entire time series*.
     max_epochs : int, default = 100
         Number of epochs the model will be trained for
     batch_size : int, default = 256
@@ -507,6 +521,8 @@ class TiDEModel(AbstractGluonTSModel):
         If True, ``lightning_logs`` directory will NOT be removed after the model finished training.
     """
 
+    ag_priority = 30
+
     _supports_known_covariates = True
     _supports_static_features = True
 
@@ -515,8 +531,8 @@ class TiDEModel(AbstractGluonTSModel):
 
         return TiDEEstimator
 
-    def _get_default_params(self):
-        return super()._get_default_params() | {
+    def _get_default_hyperparameters(self):
+        return super()._get_default_hyperparameters() | {
             "context_length": min(512, max(64, 2 * self.prediction_length)),
             "encoder_hidden_dim": 64,
             "decoder_hidden_dim": 64,
@@ -531,7 +547,7 @@ class TiDEModel(AbstractGluonTSModel):
             "batch_size": 256,
         }
 
-    def _get_estimator_init_args(self) -> Dict[str, Any]:
+    def _get_estimator_init_args(self) -> dict[str, Any]:
         init_kwargs = super()._get_estimator_init_args()
         init_kwargs["num_feat_static_cat"] = self.num_feat_static_cat
         init_kwargs["num_feat_static_real"] = self.num_feat_static_real

autogluon/timeseries/models/local/__init__.py

@@ -1,5 +1,3 @@
-import joblib.externals.loky
-
 from .naive import AverageModel, NaiveModel, SeasonalAverageModel, SeasonalNaiveModel
 from .npts import NPTSModel
 from .statsforecast import (
@@ -15,8 +13,3 @@ from .statsforecast import (
     ThetaModel,
     ZeroModel,
 )
-
-# By default, joblib w/ loky backend kills processes that take >300MB of RAM assuming that this is caused by a memory
-# leak. This leads to problems for some memory-hungry models like AutoARIMA/Theta.
-# This monkey patch removes this undesired behavior
-joblib.externals.loky.process_executor._MAX_MEMORY_LEAK_SIZE = int(3e10)

autogluon/timeseries/models/local/abstract_local_model.py

@@ -1,26 +1,24 @@
 import logging
 import time
-from multiprocessing import TimeoutError, cpu_count
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from multiprocessing import TimeoutError
+from typing import Any, Callable
 
 import numpy as np
 import pandas as pd
-from joblib import Parallel, delayed
+from joblib import Parallel, cpu_count, delayed
 from scipy.stats import norm
 
 from autogluon.core.utils.exceptions import TimeLimitExceeded
-from autogluon.timeseries.dataset.ts_dataframe import ITEMID, TimeSeriesDataFrame
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
+from autogluon.timeseries.metrics import TimeSeriesScorer
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
+from autogluon.timeseries.utils.constants import AG_DEFAULT_N_JOBS
 from autogluon.timeseries.utils.datetime import get_seasonality
 from autogluon.timeseries.utils.warning_filters import warning_filter
 
 logger = logging.getLogger(__name__)
 
 
-# We use the same default n_jobs across AG-TS to ensure that Joblib reuses the process pool
-AG_DEFAULT_N_JOBS = max(int(cpu_count() * 0.5), 1)
-
-
 class AbstractLocalModel(AbstractTimeSeriesModel):
     """Abstract class for local forecasting models that are trained separately for each time series.
 
@@ -28,50 +26,31 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
 
     Attributes
     ----------
-    allowed_local_model_args : List[str]
+    allowed_local_model_args
         Argument that can be passed to the underlying local model.
-    default_n_jobs : Union[int, float]
-        Default number of CPU cores used to train models. If float, this fraction of CPU cores will be used.
-    default_max_ts_length : Optional[int]
+    default_max_ts_length
         If not None, only the last ``max_ts_length`` time steps of each time series will be used to train the model.
         This significantly speeds up fitting and usually leads to no change in accuracy.
-    init_time_in_seconds : int
+    init_time_in_seconds
         Time that it takes to initialize the model in seconds (e.g., because of JIT compilation by Numba).
         If time_limit is below this number, model won't be trained.
     """
 
-    allowed_local_model_args: List[str] = []
-    default_n_jobs: Union[int, float] = AG_DEFAULT_N_JOBS
-    default_max_ts_length: Optional[int] = 2500
+    allowed_local_model_args: list[str] = []
+    default_max_ts_length: int | None = 2500
     default_max_time_limit_ratio = 1.0
     init_time_in_seconds: int = 0
 
     def __init__(
         self,
-        freq: Optional[str] = None,
+        freq: str | None = None,
         prediction_length: int = 1,
-        path: Optional[str] = None,
-        name: Optional[str] = None,
-        eval_metric: str = None,
-        hyperparameters: Dict[str, Any] = None,
+        path: str | None = None,
+        name: str | None = None,
+        eval_metric: str | TimeSeriesScorer | None = None,
+        hyperparameters: dict[str, Any] | None = None,
         **kwargs,  # noqa
     ):
-        if hyperparameters is None:
-            hyperparameters = {}
-        else:
-            hyperparameters = hyperparameters.copy()
-        # TODO: Replace with 'num_cpus' argument passed to fit (after predictor API is changed)
-        n_jobs = hyperparameters.pop("n_jobs", self.default_n_jobs)
-        if isinstance(n_jobs, float) and 0 < n_jobs <= 1:
-            self.n_jobs = max(int(cpu_count() * n_jobs), 1)
-        elif isinstance(n_jobs, int):
-            self.n_jobs = n_jobs
-        else:
-            raise ValueError(f"n_jobs must be a float between 0 and 1 or an integer (received n_jobs = {n_jobs})")
-        # Default values, potentially overridden inside _fit()
-        self.use_fallback_model = hyperparameters.pop("use_fallback_model", True)
-        self.max_ts_length = hyperparameters.pop("max_ts_length", self.default_max_ts_length)
-
         super().__init__(
             path=path,
             freq=freq,
@@ -82,12 +61,12 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
             **kwargs,
         )
 
-        self._local_model_args: Dict[str, Any] = None
-        self._seasonal_period: Optional[int] = None
-        self._dummy_forecast: Optional[pd.DataFrame] = None
+        self._local_model_args: dict[str, Any]
+        self._seasonal_period: int
+        self._dummy_forecast: pd.DataFrame
 
     @property
-    def allowed_hyperparameters(self) -> List[str]:
+    def allowed_hyperparameters(self) -> list[str]:
         return (
             super().allowed_hyperparameters
             + ["use_fallback_model", "max_ts_length", "n_jobs"]
@@ -97,40 +76,42 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
     def preprocess(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         is_train: bool = False,
         **kwargs,
-    ) -> Tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:
+    ) -> tuple[TimeSeriesDataFrame, TimeSeriesDataFrame | None]:
         if not self._get_tags()["allow_nan"]:
             data = data.fill_missing_values()
         return data, known_covariates
 
-    def _fit(self, train_data: TimeSeriesDataFrame, time_limit: Optional[int] = None, **kwargs):
+    def _get_default_hyperparameters(self) -> dict:
+        return {
+            "n_jobs": AG_DEFAULT_N_JOBS,
+            "use_fallback_model": True,
+            "max_ts_length": self.default_max_ts_length,
+        }
+
+    @staticmethod
+    def _compute_n_jobs(n_jobs: int | float) -> int:
+        if isinstance(n_jobs, float) and 0 < n_jobs <= 1:
+            return max(int(cpu_count() * n_jobs), 1)
+        elif isinstance(n_jobs, int):
+            return n_jobs
+        else:
+            raise ValueError(f"n_jobs must be a float between 0 and 1 or an integer (received n_jobs = {n_jobs})")
+
+    def _fit(self, train_data: TimeSeriesDataFrame, time_limit: int | None = None, **kwargs):
         self._check_fit_params()
 
         if time_limit is not None and time_limit < self.init_time_in_seconds:
             raise TimeLimitExceeded
 
-        # Initialize parameters passed to each local model
-        raw_local_model_args = self._get_model_params().copy()
-
-        unused_local_model_args = []
         local_model_args = {}
-        # TODO: Move filtering logic to AbstractTimeSeriesModel
-        for key, value in raw_local_model_args.items():
+        for key, value in self.get_hyperparameters().items():
             if key in self.allowed_local_model_args:
                 local_model_args[key] = value
-            elif key in self.allowed_hyperparameters:
-                # Quietly ignore params in self.allowed_hyperparameters - they are used by AbstractTimeSeriesModel
-                pass
-            else:
-                unused_local_model_args.append(key)
 
-        if len(unused_local_model_args):
-            logger.warning(
-                f"{self.name} ignores following hyperparameters: {unused_local_model_args}. "
-                f"See the docstring of {self.name} for the list of supported hyperparameters."
-            )
+        self._log_unused_hyperparameters(extra_allowed_hyperparameters=self.allowed_local_model_args)
 
         if "seasonal_period" not in local_model_args or local_model_args["seasonal_period"] is None:
             local_model_args["seasonal_period"] = get_seasonality(self.freq)
@@ -141,35 +122,46 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
         self._dummy_forecast = self._get_dummy_forecast(train_data)
         return self
 
-    def _get_dummy_forecast(self, train_data: TimeSeriesDataFrame) -> pd.DataFrame:
+    def _get_dummy_forecast(self, train_data: TimeSeriesDataFrame, max_num_rows: int = 20_000) -> pd.DataFrame:
         agg_functions = ["mean"] + [get_quantile_function(q) for q in self.quantile_levels]
-        stats_marginal = train_data[self.target].agg(agg_functions)
+        target_series = train_data[self.target]
+        if len(target_series) > max_num_rows:
+            target_series = target_series.sample(max_num_rows, replace=True)
+        stats_marginal = target_series.agg(agg_functions)
         stats_repeated = np.tile(stats_marginal.values, [self.prediction_length, 1])
         return pd.DataFrame(stats_repeated, columns=stats_marginal.index)
 
-    def _update_local_model_args(self, local_model_args: Dict[str, Any]) -> Dict[str, Any]:
+    def _update_local_model_args(self, local_model_args: dict[str, Any]) -> dict[str, Any]:
         return local_model_args
 
     def _predict(self, data: TimeSeriesDataFrame, **kwargs) -> TimeSeriesDataFrame:
-        if self.max_ts_length is not None:
-            logger.debug(f"Shortening all time series to at most {self.max_ts_length}")
-            data = data.groupby(level=ITEMID, sort=False).tail(self.max_ts_length)
+        model_params = self.get_hyperparameters()
+        max_ts_length = model_params["max_ts_length"]
+        if max_ts_length is not None:
+            logger.debug(f"Shortening all time series to at most {max_ts_length}")
+            data = data.slice_by_timestep(-max_ts_length, None)
 
-        df = pd.DataFrame(data).reset_index(level=ITEMID)
-        all_series = (ts for _, ts in df.groupby(by=ITEMID, as_index=False, sort=False)[self.target])
+        indptr = data.get_indptr()
+        target_series = data[self.target].droplevel(level=TimeSeriesDataFrame.ITEMID)
+        all_series = (target_series[indptr[i] : indptr[i + 1]] for i in range(len(indptr) - 1))
 
         # timeout ensures that no individual job takes longer than time_limit
         # TODO: a job started late may still exceed time_limit - how to prevent that?
         time_limit = kwargs.get("time_limit")
-        timeout = None if self.n_jobs == 1 else time_limit
+        # TODO: Take into account num_cpus once the TimeSeriesPredictor API is updated
+        n_jobs = self._compute_n_jobs(model_params["n_jobs"])
+        timeout = None if n_jobs == 1 else time_limit
         # end_time ensures that no new jobs are started after time_limit is exceeded
         end_time = None if time_limit is None else time.time() + time_limit
-        executor = Parallel(self.n_jobs, timeout=timeout)
+        executor = Parallel(n_jobs=n_jobs, timeout=timeout)
 
         try:
             with warning_filter():
                 predictions_with_flags = executor(
-                    delayed(self._predict_wrapper)(ts, end_time=end_time) for ts in all_series
+                    delayed(self._predict_wrapper)(
+                        ts, use_fallback_model=model_params["use_fallback_model"], end_time=end_time
+                    )
+                    for ts in all_series
                 )
         except TimeoutError:
             raise TimeLimitExceeded
@@ -185,7 +177,12 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
         predictions_df.index = self.get_forecast_horizon_index(data)
         return TimeSeriesDataFrame(predictions_df)
 
-    def _predict_wrapper(self, time_series: pd.Series, end_time: Optional[float] = None) -> Tuple[pd.DataFrame, bool]:
+    def _predict_wrapper(
+        self,
+        time_series: pd.Series,
+        use_fallback_model: bool,
+        end_time: float | None = None,
+    ) -> tuple[pd.DataFrame, bool]:
         if end_time is not None and time.time() >= end_time:
             raise TimeLimitExceeded
 
@@ -201,7 +198,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
                 if not np.isfinite(result.values).all():
                     raise RuntimeError("Forecast contains NaN or Inf values.")
             except Exception:
-                if self.use_fallback_model:
+                if use_fallback_model:
                     result = seasonal_naive_forecast(
                         target=time_series.values.ravel(),
                         prediction_length=self.prediction_length,
@@ -222,7 +219,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
 
 
 def seasonal_naive_forecast(
-    target: np.ndarray, prediction_length: int, quantile_levels: List[float], seasonal_period: int
+    target: np.ndarray, prediction_length: int, quantile_levels: list[float], seasonal_period: int
 ) -> pd.DataFrame:
     """Generate seasonal naive forecast, predicting the last observed value from the same period."""