autogluon.timeseries 1.3.2b20250712__py3-none-any.whl → 1.4.1b20251116__py3-none-any.whl

autogluon/timeseries/{trainer.py → trainer/trainer.py}

@@ -5,14 +5,14 @@ import time
 import traceback
 from collections import defaultdict
 from pathlib import Path
-from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union
+from typing import Any, Literal, Optional, Union
 
 import networkx as nx
 import numpy as np
 import pandas as pd
 from tqdm import tqdm
 
-from autogluon.common.utils.utils import hash_pandas_df, seed_everything
+from autogluon.common.utils.utils import seed_everything
 from autogluon.core.trainer.abstract_trainer import AbstractTrainer
 from autogluon.core.utils.exceptions import TimeLimitExceeded
 from autogluon.core.utils.loaders import load_pkl
@@ -20,23 +20,25 @@ from autogluon.core.utils.savers import save_pkl
 from autogluon.timeseries import TimeSeriesDataFrame
 from autogluon.timeseries.metrics import TimeSeriesScorer, check_get_evaluation_metric
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel, TimeSeriesModelBase
-from autogluon.timeseries.models.ensemble import AbstractTimeSeriesEnsembleModel, GreedyEnsemble
+from autogluon.timeseries.models.ensemble import AbstractTimeSeriesEnsembleModel
 from autogluon.timeseries.models.multi_window import MultiWindowBacktestingModel
-from autogluon.timeseries.models.presets import contains_searchspace, get_preset_models
 from autogluon.timeseries.splitter import AbstractWindowSplitter, ExpandingWindowSplitter
+from autogluon.timeseries.trainer.ensemble_composer import EnsembleComposer, validate_ensemble_hyperparameters
 from autogluon.timeseries.utils.features import (
     ConstantReplacementFeatureImportanceTransform,
     CovariateMetadata,
     PermutationFeatureImportanceTransform,
 )
-from autogluon.timeseries.utils.warning_filters import disable_tqdm, warning_filter
+from autogluon.timeseries.utils.warning_filters import disable_tqdm
+
+from .model_set_builder import TrainableModelSetBuilder, contains_searchspace
+from .prediction_cache import PredictionCache, get_prediction_cache
+from .utils import log_scores_and_times
 
 logger = logging.getLogger("autogluon.timeseries.trainer")
 
 
 class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
-    _cached_predictions_filename = "cached_predictions.pkl"
-
     max_rel_importance_score: float = 1e5
     eps_abs_importance_score: float = 1e-5
     max_ensemble_time_limit: float = 600.0
@@ -50,11 +52,11 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         skip_model_selection: bool = False,
         enable_ensemble: bool = True,
         verbosity: int = 2,
-        val_splitter: Optional[AbstractWindowSplitter] = None,
+        num_val_windows: Optional[int] = None,
+        val_step_size: Optional[int] = None,
         refit_every_n_windows: Optional[int] = 1,
         # TODO: Set cache_predictions=False by default once all models in default presets have a reasonable inference speed
         cache_predictions: bool = True,
-        ensemble_model_type: Optional[Type] = None,
         **kwargs,
     ):
         super().__init__(
@@ -71,32 +73,27 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         self.skip_model_selection = skip_model_selection
         # Ensemble cannot be fit if val_scores are not computed
         self.enable_ensemble = enable_ensemble and not skip_model_selection
-        if ensemble_model_type is None:
-            ensemble_model_type = GreedyEnsemble
-        else:
+        if kwargs.get("ensemble_model_type") is not None:
             logger.warning(
-                "Using a custom `ensemble_model_type` is experimental functionality that may break in future versions."
+                "Using a custom `ensemble_model_type` is no longer supported. Use the `ensemble_hyperparameters` "
+                "argument to `fit` instead."
             )
-        self.ensemble_model_type: Type[AbstractTimeSeriesEnsembleModel] = ensemble_model_type
 
         self.verbosity = verbosity
 
-        #: Dict of normal model -> FULL model. FULL models are produced by
+        #: dict of normal model -> FULL model. FULL models are produced by
         #: self.refit_single_full() and self.refit_full().
         self.model_refit_map = {}
 
         self.eval_metric = check_get_evaluation_metric(eval_metric, prediction_length=prediction_length)
-        if val_splitter is None:
-            val_splitter = ExpandingWindowSplitter(prediction_length=self.prediction_length)
-        assert isinstance(val_splitter, AbstractWindowSplitter), "val_splitter must be of type AbstractWindowSplitter"
-        self.val_splitter = val_splitter
+
+        self.num_val_windows = num_val_windows
+        self.val_step_size = val_step_size
         self.refit_every_n_windows = refit_every_n_windows
-        self.cache_predictions = cache_predictions
         self.hpo_results = {}
 
-        if self._cached_predictions_path.exists():
-            logger.debug(f"Removing existing cached predictions file {self._cached_predictions_path}")
-            self._cached_predictions_path.unlink()
+        self.prediction_cache: PredictionCache = get_prediction_cache(cache_predictions, self.path)
+        self.prediction_cache.clear()
 
     @property
     def path_pkl(self) -> str:
@@ -121,7 +118,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         else:
             return None
 
-    def load_data(self) -> Tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:
+    def load_data(self) -> tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:
         train_data = self.load_train_data()
         val_data = self.load_val_data()
         return train_data, val_data
@@ -136,7 +133,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
         self.models = models
 
-    def _get_model_oof_predictions(self, model_name: str) -> List[TimeSeriesDataFrame]:
+    def _get_model_oof_predictions(self, model_name: str) -> list[TimeSeriesDataFrame]:
         model_path = os.path.join(self.path, self.get_model_attribute(model=model_name, attribute="path"))
         model_type = self.get_model_attribute(model=model_name, attribute="type")
         return model_type.load_oof_predictions(path=model_path)
@@ -144,16 +141,16 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
     def _add_model(
         self,
         model: TimeSeriesModelBase,
-        base_models: Optional[List[str]] = None,
+        base_models: Optional[list[str]] = None,
     ):
         """Add a model to the model graph of the trainer. If the model is an ensemble, also add
         information about dependencies to the model graph (list of models specified via ``base_models``).
 
         Parameters
         ----------
-        model : TimeSeriesModelBase
+        model
             The model to be added to the model graph.
-        base_models : List[str], optional, default None
+        base_models
             If the model is an ensemble, the list of base model names that are included in the ensemble.
             Expected only when ``model`` is a ``AbstractTimeSeriesEnsembleModel``.
 
@@ -176,7 +173,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
             for base_model in base_models:
                 self.model_graph.add_edge(base_model, model.name)
 
-    def _get_model_levels(self) -> Dict[str, int]:
+    def _get_model_levels(self) -> dict[str, int]:
         """Get a dictionary mapping each model to their level in the model graph"""
 
         # get nodes without a parent
@@ -197,7 +194,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
         return levels
 
-    def get_models_attribute_dict(self, attribute: str, models: Optional[List[str]] = None) -> Dict[str, Any]:
+    def get_models_attribute_dict(self, attribute: str, models: Optional[list[str]] = None) -> dict[str, Any]:
         """Get an attribute from the `model_graph` for each of the model names
         specified. If `models` is none, the attribute will be returned for all models"""
         results = {}
@@ -230,13 +227,13 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
             key=lambda mns: (mns[1], -mns[2]),  # (score, -level)
         )[0]
 
-    def get_model_names(self, level: Optional[int] = None) -> List[str]:
+    def get_model_names(self, level: Optional[int] = None) -> list[str]:
         """Get model names that are registered in the model graph"""
         if level is not None:
             return list(node for node, l in self._get_model_levels().items() if l == level)  # noqa: E741
         return list(self.model_graph.nodes)
 
-    def get_info(self, include_model_info: bool = False) -> Dict[str, Any]:
+    def get_info(self, include_model_info: bool = False) -> dict[str, Any]:
         num_models_trained = len(self.get_model_names())
         if self.model_best is not None:
             best_model = self.model_best
@@ -261,25 +258,6 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
         return info
 
-    def _train_single(
-        self,
-        train_data: TimeSeriesDataFrame,
-        model: AbstractTimeSeriesModel,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[float] = None,
-    ) -> AbstractTimeSeriesModel:
-        """Train the single model and return the model object that was fitted. This method
-        does not save the resulting model."""
-        model.fit(
-            train_data=train_data,
-            val_data=val_data,
-            time_limit=time_limit,
-            verbosity=self.verbosity,
-            val_splitter=self.val_splitter,
-            refit_every_n_windows=self.refit_every_n_windows,
-        )
-        return model
-
     def tune_model_hyperparameters(
         self,
         model: AbstractTimeSeriesModel,
@@ -302,7 +280,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
                 hyperparameter_tune_kwargs=hyperparameter_tune_kwargs,
                 time_limit=time_limit,
                 default_num_trials=default_num_trials,
-                val_splitter=self.val_splitter,
+                val_splitter=self._get_val_splitter(),
                 refit_every_n_windows=self.refit_every_n_windows,
             )
         total_tuning_time = time.time() - tuning_start_time
@@ -339,12 +317,13 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         model: AbstractTimeSeriesModel,
         val_data: Optional[TimeSeriesDataFrame] = None,
         time_limit: Optional[float] = None,
-    ) -> List[str]:
+    ) -> list[str]:
         """Fit and save the given model on given training and validation data and save the trained model.
 
         Returns
         -------
-        model_names_trained: the list of model names that were successfully trained
+        model_names_trained
+            the list of model names that were successfully trained
         """
         fit_start_time = time.time()
         model_names_trained = []
@@ -354,7 +333,15 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
                     logger.info(f"\tSkipping {model.name} due to lack of time remaining.")
                     return model_names_trained
 
-            model = self._train_single(train_data, model, val_data=val_data, time_limit=time_limit)
+            model.fit(
+                train_data=train_data,
+                val_data=val_data,
+                time_limit=time_limit,
+                verbosity=self.verbosity,
+                val_splitter=self._get_val_splitter(),
+                refit_every_n_windows=self.refit_every_n_windows,
+            )
+
             fit_end_time = time.time()
             model.fit_time = model.fit_time or (fit_end_time - fit_start_time)
 
@@ -365,7 +352,12 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
                     val_data, store_val_score=True, store_predict_time=True, time_limit=time_limit
                 )
 
-            self._log_scores_and_times(model.val_score, model.fit_time, model.predict_time)
+            log_scores_and_times(
+                val_score=model.val_score,
+                fit_time=model.fit_time,
+                predict_time=model.predict_time,
+                eval_metric_name=self.eval_metric.name_with_sign,
+            )
 
             self.save_model(model=model)
         except TimeLimitExceeded:
@@ -381,31 +373,51 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
         return model_names_trained
 
-    def _log_scores_and_times(
-        self,
-        val_score: Optional[float] = None,
-        fit_time: Optional[float] = None,
-        predict_time: Optional[float] = None,
-    ):
-        if val_score is not None:
-            logger.info(f"\t{val_score:<7.4f}".ljust(15) + f"= Validation score ({self.eval_metric.name_with_sign})")
-        if fit_time is not None:
-            logger.info(f"\t{fit_time:<7.2f} s".ljust(15) + "= Training runtime")
-        if predict_time is not None:
-            logger.info(f"\t{predict_time:<7.2f} s".ljust(15) + "= Validation (prediction) runtime")
-
-    def _train_multi(
+    def fit(
         self,
         train_data: TimeSeriesDataFrame,
-        hyperparameters: Union[str, Dict],
+        hyperparameters: Union[str, dict[Any, dict]],
         val_data: Optional[TimeSeriesDataFrame] = None,
+        ensemble_hyperparameters: Optional[dict] = None,
         hyperparameter_tune_kwargs: Optional[Union[str, dict]] = None,
-        excluded_model_types: Optional[List[str]] = None,
+        excluded_model_types: Optional[list[str]] = None,
         time_limit: Optional[float] = None,
         random_seed: Optional[int] = None,
-    ) -> List[str]:
+    ):
+        """Fit a set of timeseries models specified by the `hyperparameters`
+        dictionary that maps model names to their specified hyperparameters.
+
+        Parameters
+        ----------
+        train_data
+            Training data for fitting time series timeseries models.
+        hyperparameters
+            A dictionary mapping selected model names, model classes or model factory to hyperparameter
+            settings. Model names should be present in `trainer.presets.DEFAULT_MODEL_NAMES`. Optionally,
+            the user may provide one of "default", "light" and "very_light" to specify presets.
+        val_data
+            Optional validation data set to report validation scores on.
+        ensemble_hyperparameters
+            A dictionary mapping ensemble names to their specified hyperparameters. Ensemble names
+            should be defined in the models.ensemble namespace. defaults to `{"GreedyEnsemble": {}}`
+            which only fits a greedy weighted ensemble with default hyperparameters. Providing an
+            empty dictionary disables ensemble training.
+        hyperparameter_tune_kwargs
+            Args for hyperparameter tuning
+        excluded_model_types
+            Names of models that should not be trained, even if listed in `hyperparameters`.
+        time_limit
+            Time limit for training
+        random_seed
+            Random seed that will be set to each model during training
+        """
         logger.info(f"\nStarting training. Start time is {time.strftime('%Y-%m-%d %H:%M:%S')}")
 
+        # Handle ensemble hyperparameters
+        if ensemble_hyperparameters is None:
+            ensemble_hyperparameters = {"GreedyEnsemble": {}}
+        ensemble_hyperparameters = validate_ensemble_hyperparameters(ensemble_hyperparameters)
+
         time_start = time.time()
         hyperparameters = copy.deepcopy(hyperparameters)
 
@@ -415,11 +427,11 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
                 self.save_val_data(val_data)
             self.is_data_saved = True
 
-        models = self.construct_model_templates(
+        models = self.get_trainable_base_models(
             hyperparameters=hyperparameters,
             hyperparameter_tune=hyperparameter_tune_kwargs is not None,  # TODO: remove hyperparameter_tune
             freq=train_data.freq,
-            multi_window=self.val_splitter.num_val_windows > 0,
+            multi_window=self._get_val_splitter().num_val_windows > 0,
             excluded_model_types=excluded_model_types,
         )
 
@@ -439,8 +451,6 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         num_base_models = len(models)
         model_names_trained = []
         for i, model in enumerate(models):
-            assert isinstance(model, AbstractTimeSeriesModel)
-
             if time_limit is None:
                 time_left = None
                 time_left_for_model = None
@@ -490,42 +500,13 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
                     train_data, model=model, val_data=val_data, time_limit=time_left_for_model
                 )
 
-        if self.enable_ensemble:
-            models_available_for_ensemble = self.get_model_names(level=0)
-
-            time_left_for_ensemble = None
-            if time_limit is not None:
-                time_left_for_ensemble = time_limit - (time.time() - time_start)
-
-            if time_left_for_ensemble is not None and time_left_for_ensemble <= 0:
-                logger.info(
-                    "Not fitting ensemble due to lack of time remaining. "
-                    f"Time left: {time_left_for_ensemble:.1f} seconds"
-                )
-            elif len(models_available_for_ensemble) <= 1:
-                logger.info(
-                    "Not fitting ensemble as "
-                    + (
-                        "no models were successfully trained."
-                        if not models_available_for_ensemble
-                        else "only 1 model was trained."
-                    )
-                )
-            else:
-                try:
-                    model_names_trained.append(
-                        self.fit_ensemble(
-                            data_per_window=self._get_ensemble_oof_data(train_data=train_data, val_data=val_data),
-                            model_names=models_available_for_ensemble,
-                            time_limit=time_left_for_ensemble,
-                        )
-                    )
-                except Exception as err:  # noqa
-                    logger.error(
-                        "\tWarning: Exception caused ensemble to fail during training... Skipping this model."
-                    )
-                    logger.error(f"\t{err}")
-                    logger.debug(traceback.format_exc())
+        ensemble_names = self._fit_ensembles(
+            train_data=train_data,
+            val_data=val_data,
+            time_limit=None if time_limit is None else time_limit - (time.time() - time_start),
+            ensemble_hyperparameters=ensemble_hyperparameters,
+        )
+        model_names_trained.extend(ensemble_names)
 
         logger.info(f"Training complete. Models trained: {model_names_trained}")
         logger.info(f"Total runtime: {time.time() - time_start:.2f} s")
@@ -539,82 +520,72 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
         return model_names_trained
 
-    def _get_ensemble_oof_data(
-        self, train_data: TimeSeriesDataFrame, val_data: Optional[TimeSeriesDataFrame]
-    ) -> List[TimeSeriesDataFrame]:
-        if val_data is None:
-            return [val_fold for _, val_fold in self.val_splitter.split(train_data)]
-        else:
-            return [val_data]
-
-    def _get_ensemble_model_name(self) -> str:
-        """Ensure we don't have name collisions in the ensemble model name"""
-        ensemble_name = "WeightedEnsemble"
-        increment = 1
-        while ensemble_name in self._get_banned_model_names():
-            increment += 1
-            ensemble_name = f"WeightedEnsemble_{increment}"
-        return ensemble_name
-
-    def fit_ensemble(
+    def _fit_ensembles(
         self,
-        data_per_window: List[TimeSeriesDataFrame],
-        model_names: List[str],
-        time_limit: Optional[float] = None,
-    ) -> str:
-        logger.info("Fitting simple weighted ensemble.")
+        *,
+        train_data: TimeSeriesDataFrame,
+        val_data: Optional[TimeSeriesDataFrame],
+        time_limit: Optional[float],
+        ensemble_hyperparameters: dict,
+    ) -> list[str]:
+        if not self.enable_ensemble or not ensemble_hyperparameters:
+            logger.warning("Ensemble training is disabled. Skipping ensemble training.")
+            return []
+
+        ensemble_composer = self._get_ensemble_composer(ensemble_hyperparameters).fit(
+            train_data,
+            val_data,
+            time_limit,
+        )
 
-        predictions_per_window: Dict[str, List[TimeSeriesDataFrame]] = {}
-        base_model_scores = self.get_models_attribute_dict(attribute="val_score", models=self.get_model_names(0))
+        ensembles_trained = []
+        for _, model, base_models in ensemble_composer.iter_ensembles():
+            self._add_model(model=model, base_models=base_models)
+            self.save_model(model=model)
+            ensembles_trained.append(model.name)
 
-        for model_name in model_names:
-            predictions_per_window[model_name] = self._get_model_oof_predictions(model_name=model_name)
+        return ensembles_trained if ensembles_trained else []
 
-        time_start = time.time()
-        ensemble = self.ensemble_model_type(
-            name=self._get_ensemble_model_name(),
+    def _get_val_splitter(self) -> AbstractWindowSplitter:
+        if self.num_val_windows is None:
+            val_splitter = ExpandingWindowSplitter(prediction_length=self.prediction_length)
+        else:
+            val_splitter = ExpandingWindowSplitter(
+                prediction_length=self.prediction_length,
+                num_val_windows=self.num_val_windows,
+                val_step_size=self.val_step_size,
+            )
+        return val_splitter
+
+    def _get_ensemble_composer(self, ensemble_hyperparameters: dict) -> "EnsembleComposer":
+        """Create an ensemble composer instance for delegation."""
+        return EnsembleComposer(
+            path=self.path,
+            prediction_length=self.prediction_length,
             eval_metric=self.eval_metric,
             target=self.target,
-            prediction_length=self.prediction_length,
-            path=self.path,
-            freq=data_per_window[0].freq,
             quantile_levels=self.quantile_levels,
-            covariate_metadata=self.covariate_metadata,
+            model_graph=self.model_graph,
+            ensemble_hyperparameters=ensemble_hyperparameters,
+            window_splitter=self._get_val_splitter(),
         )
-        with warning_filter():
-            ensemble.fit(
-                predictions_per_window=predictions_per_window,
-                data_per_window=data_per_window,
-                model_scores=base_model_scores,
-                time_limit=time_limit,
-            )
-        ensemble.fit_time = time.time() - time_start
-
-        predict_time = 0
-        for m in ensemble.model_names:
-            predict_time += self.get_model_attribute(model=m, attribute="predict_time")
-        ensemble.predict_time = predict_time
-
-        score_per_fold = []
-        for window_idx, data in enumerate(data_per_window):
-            predictions = ensemble.predict({n: predictions_per_window[n][window_idx] for n in ensemble.model_names})
-            score_per_fold.append(self._score_with_predictions(data, predictions))
-        ensemble.val_score = float(np.mean(score_per_fold, dtype=np.float64))
-
-        self._log_scores_and_times(
-            val_score=ensemble.val_score,
-            fit_time=ensemble.fit_time,
-            predict_time=ensemble.predict_time,
-        )
-        self._add_model(model=ensemble, base_models=ensemble.model_names)
-        self.save_model(model=ensemble)
-        return ensemble.name
+
+    def _get_validation_windows(
+        self, train_data: TimeSeriesDataFrame, val_data: Optional[TimeSeriesDataFrame]
+    ) -> list[TimeSeriesDataFrame]:
+        """If validation data is provided, return this as a single validation window. If not,
+        use the validation splitter to create a list of validation splits.
+        """
+        if val_data is None:
+            return [val_fold for _, val_fold in self._get_val_splitter().split(train_data)]
+        else:
+            return [val_data]
 
     def leaderboard(
         self,
         data: Optional[TimeSeriesDataFrame] = None,
         extra_info: bool = False,
-        extra_metrics: Optional[List[Union[str, TimeSeriesScorer]]] = None,
+        extra_metrics: Optional[list[Union[str, TimeSeriesScorer]]] = None,
         use_cache: bool = True,
     ) -> pd.DataFrame:
         logger.debug("Generating leaderboard for all models trained")
@@ -704,8 +675,8 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         return df[explicit_column_order]
 
     def persist(
-        self, model_names: Union[Literal["all", "best"], List[str]] = "all", with_ancestors: bool = False
-    ) -> List[str]:
+        self, model_names: Union[Literal["all", "best"], list[str]] = "all", with_ancestors: bool = False
+    ) -> list[str]:
         if model_names == "all":
             model_names = self.get_model_names()
         elif model_names == "best":
@@ -729,7 +700,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
         return model_names
 
-    def unpersist(self, model_names: Union[Literal["all"], List[str]] = "all") -> List[str]:
+    def unpersist(self, model_names: Union[Literal["all"], list[str]] = "all") -> list[str]:
         if model_names == "all":
             model_names = list(self.models.keys())
         if not isinstance(model_names, list):
@@ -826,9 +797,9 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         self,
         data: TimeSeriesDataFrame,
         model: Optional[Union[str, TimeSeriesModelBase]] = None,
-        metrics: Optional[Union[str, TimeSeriesScorer, List[Union[str, TimeSeriesScorer]]]] = None,
+        metrics: Optional[Union[str, TimeSeriesScorer, list[Union[str, TimeSeriesScorer]]]] = None,
         use_cache: bool = True,
-    ) -> Dict[str, float]:
+    ) -> dict[str, float]:
         past_data, known_covariates = data.get_model_inputs_for_scoring(
             prediction_length=self.prediction_length, known_covariates_names=self.covariate_metadata.known_covariates
         )
@@ -846,7 +817,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
     def get_feature_importance(
         self,
         data: TimeSeriesDataFrame,
-        features: List[str],
+        features: list[str],
         model: Optional[Union[str, TimeSeriesModelBase]] = None,
         metric: Optional[Union[str, TimeSeriesScorer]] = None,
         time_limit: Optional[float] = None,
@@ -996,7 +967,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         self,
         model: Union[str, TimeSeriesModelBase],
         data: TimeSeriesDataFrame,
-        model_pred_dict: Dict[str, Optional[TimeSeriesDataFrame]],
+        model_pred_dict: dict[str, Optional[TimeSeriesDataFrame]],
         known_covariates: Optional[TimeSeriesDataFrame] = None,
     ) -> TimeSeriesDataFrame:
         """Generate predictions using the given model.
@@ -1012,8 +983,8 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         self,
         model: Union[str, TimeSeriesModelBase],
         data: TimeSeriesDataFrame,
-        model_pred_dict: Dict[str, Optional[TimeSeriesDataFrame]],
-    ) -> Union[TimeSeriesDataFrame, Dict[str, Optional[TimeSeriesDataFrame]]]:
+        model_pred_dict: dict[str, Optional[TimeSeriesDataFrame]],
+    ) -> Union[TimeSeriesDataFrame, dict[str, Optional[TimeSeriesDataFrame]]]:
         """Get the first argument that should be passed to model.predict.
 
         This method assumes that model_pred_dict contains the predictions of all base models, if model is an ensemble.
@@ -1029,13 +1000,13 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
     def get_model_pred_dict(
         self,
-        model_names: List[str],
+        model_names: list[str],
         data: TimeSeriesDataFrame,
         known_covariates: Optional[TimeSeriesDataFrame] = None,
         raise_exception_if_failed: bool = True,
         use_cache: bool = True,
         random_seed: Optional[int] = None,
-    ) -> Tuple[Dict[str, Optional[TimeSeriesDataFrame]], Dict[str, float]]:
+    ) -> tuple[dict[str, Optional[TimeSeriesDataFrame]], dict[str, float]]:
         """Return a dictionary with predictions of all models for the given dataset.
 
         Parameters
@@ -1055,12 +1026,13 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         use_cache
             If False, will ignore the cache even if it's available.
         """
-        if self.cache_predictions and use_cache:
-            dataset_hash = self._compute_dataset_hash(data=data, known_covariates=known_covariates)
-            model_pred_dict, pred_time_dict_marginal = self._get_cached_pred_dicts(dataset_hash)
+        if use_cache:
+            model_pred_dict, pred_time_dict_marginal = self.prediction_cache.get(
+                data=data, known_covariates=known_covariates
+            )
         else:
             model_pred_dict = {}
-            pred_time_dict_marginal: Dict[str, Any] = {}
+            pred_time_dict_marginal: dict[str, Any] = {}
 
         model_set = set()
         for model_name in model_names:
@@ -1093,9 +1065,11 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
         if len(failed_models) > 0 and raise_exception_if_failed:
             raise RuntimeError(f"Following models failed to predict: {failed_models}")
-        if self.cache_predictions and use_cache:
-            self._save_cached_pred_dicts(
-                dataset_hash,  # type: ignore
+
+        if use_cache:
+            self.prediction_cache.put(
+                data=data,
+                known_covariates=known_covariates,
                 model_pred_dict=model_pred_dict,
                 pred_time_dict=pred_time_dict_marginal,
             )
@@ -1106,7 +1080,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
         return final_model_pred_dict, final_pred_time_dict_total
 
-    def _get_total_pred_time_from_marginal(self, pred_time_dict_marginal: Dict[str, float]) -> Dict[str, float]:
+    def _get_total_pred_time_from_marginal(self, pred_time_dict_marginal: dict[str, float]) -> dict[str, float]:
         pred_time_dict_total = defaultdict(float)
         for model_name in pred_time_dict_marginal.keys():
             for base_model in self.get_minimum_model_set(model_name):
@@ -1114,62 +1088,6 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
                     pred_time_dict_total[model_name] += pred_time_dict_marginal[base_model]
         return dict(pred_time_dict_total)
 
-    @property
-    def _cached_predictions_path(self) -> Path:
-        return Path(self.path) / self._cached_predictions_filename
-
-    @staticmethod
-    def _compute_dataset_hash(
-        data: TimeSeriesDataFrame, known_covariates: Optional[TimeSeriesDataFrame] = None
-    ) -> str:
-        """Compute a unique string that identifies the time series dataset."""
-        combined_hash = hash_pandas_df(data) + hash_pandas_df(known_covariates) + hash_pandas_df(data.static_features)
-        return combined_hash
-
-    def _load_cached_predictions(self) -> dict[str, dict[str, dict[str, Any]]]:
-        """Load cached predictions from disk. If loading fails, an empty dictionary is returned."""
-        if self._cached_predictions_path.exists():
-            try:
-                cached_predictions = load_pkl.load(str(self._cached_predictions_path))
-            except Exception:
-                cached_predictions = {}
-        else:
-            cached_predictions = {}
-        return cached_predictions
-
-    def _get_cached_pred_dicts(
-        self, dataset_hash: str
-    ) -> Tuple[Dict[str, Optional[TimeSeriesDataFrame]], Dict[str, float]]:
-        """Load cached predictions for given dataset_hash from disk, if possible.
-
-        If loading fails for any reason, empty dicts are returned.
-        """
-        cached_predictions = self._load_cached_predictions()
-        if dataset_hash in cached_predictions:
-            try:
-                model_pred_dict = cached_predictions[dataset_hash]["model_pred_dict"]
-                pred_time_dict = cached_predictions[dataset_hash]["pred_time_dict"]
-                assert model_pred_dict.keys() == pred_time_dict.keys()
-                return model_pred_dict, pred_time_dict
-            except Exception:
-                logger.warning("Cached predictions are corrupted. Predictions will be made from scratch.")
-        return {}, {}
-
-    def _save_cached_pred_dicts(
-        self,
-        dataset_hash: str,
-        model_pred_dict: Dict[str, Optional[TimeSeriesDataFrame]],
-        pred_time_dict: Dict[str, float],
-    ) -> None:
-        cached_predictions = self._load_cached_predictions()
-        # Do not save results for models that failed
-        cached_predictions[dataset_hash] = {
-            "model_pred_dict": {k: v for k, v in model_pred_dict.items() if v is not None},
-            "pred_time_dict": {k: v for k, v in pred_time_dict.items() if v is not None},
-        }
-        save_pkl.save(str(self._cached_predictions_path), object=cached_predictions)
-        logger.debug(f"Cached predictions saved to {self._cached_predictions_path}")
-
     def _merge_refit_full_data(
         self, train_data: TimeSeriesDataFrame, val_data: Optional[TimeSeriesDataFrame]
     ) -> TimeSeriesDataFrame:
@@ -1183,8 +1101,8 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         self,
         train_data: Optional[TimeSeriesDataFrame] = None,
         val_data: Optional[TimeSeriesDataFrame] = None,
-        models: Optional[List[str]] = None,
-    ) -> List[str]:
+        models: Optional[list[str]] = None,
+    ) -> list[str]:
         train_data = train_data or self.load_train_data()
         val_data = val_data or self.load_val_data()
         refit_full_data = self._merge_refit_full_data(train_data, val_data)
@@ -1228,7 +1146,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         self.save()
         return models_trained_full
 
-    def refit_full(self, model: str = "all") -> Dict[str, str]:
+    def refit_full(self, model: str = "all") -> dict[str, str]:
         time_start = time.time()
         existing_models = self.get_model_names()
         if model == "all":
@@ -1260,70 +1178,27 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         logger.info(f"Total runtime: {time.time() - time_start:.2f} s")
         return copy.deepcopy(self.model_refit_map)
 
-    def construct_model_templates(
+    def get_trainable_base_models(
         self,
-        hyperparameters: Union[str, Dict[str, Any]],
+        hyperparameters: Union[str, dict[str, Any]],
         *,
         multi_window: bool = False,
         freq: Optional[str] = None,
-        excluded_model_types: Optional[List[str]] = None,
+        excluded_model_types: Optional[list[str]] = None,
         hyperparameter_tune: bool = False,
-    ) -> List[TimeSeriesModelBase]:
-        return get_preset_models(
+    ) -> list[AbstractTimeSeriesModel]:
+        return TrainableModelSetBuilder(
+            freq=freq,
+            prediction_length=self.prediction_length,
             path=self.path,
             eval_metric=self.eval_metric,
-            prediction_length=self.prediction_length,
-            freq=freq,
-            hyperparameters=hyperparameters,
-            hyperparameter_tune=hyperparameter_tune,
             quantile_levels=self.quantile_levels,
-            all_assigned_names=self._get_banned_model_names(),
             target=self.target,
             covariate_metadata=self.covariate_metadata,
-            excluded_model_types=excluded_model_types,
-            # if skip_model_selection = True, we skip backtesting
             multi_window=multi_window and not self.skip_model_selection,
-        )
-
-    def fit(
-        self,
-        train_data: TimeSeriesDataFrame,
-        hyperparameters: Union[str, Dict[Any, Dict]],
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        hyperparameter_tune_kwargs: Optional[Union[str, Dict]] = None,
-        excluded_model_types: Optional[List[str]] = None,
-        time_limit: Optional[float] = None,
-        random_seed: Optional[int] = None,
-    ):
-        """
-        Fit a set of timeseries models specified by the `hyperparameters`
-        dictionary that maps model names to their specified hyperparameters.
-
-        Parameters
-        ----------
-        train_data: TimeSeriesDataFrame
-            Training data for fitting time series timeseries models.
-        hyperparameters: str or Dict
-            A dictionary mapping selected model names, model classes or model factory to hyperparameter
-            settings. Model names should be present in `trainer.presets.DEFAULT_MODEL_NAMES`. Optionally,
-            the user may provide one of "default", "light" and "very_light" to specify presets.
-        val_data: TimeSeriesDataFrame
-            Optional validation data set to report validation scores on.
-        hyperparameter_tune_kwargs
-            Args for hyperparameter tuning
-        excluded_model_types
-            Names of models that should not be trained, even if listed in `hyperparameters`.
-        time_limit
-            Time limit for training
-        random_seed
-            Random seed that will be set to each model during training
-        """
-        self._train_multi(
-            train_data,
-            val_data=val_data,
+        ).get_model_set(
             hyperparameters=hyperparameters,
-            hyperparameter_tune_kwargs=hyperparameter_tune_kwargs,
+            hyperparameter_tune=hyperparameter_tune,
             excluded_model_types=excluded_model_types,
-            time_limit=time_limit,
-            random_seed=random_seed,
+            banned_model_names=self._get_banned_model_names(),
         )