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

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

@@ -5,38 +5,40 @@ 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
 
 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
 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
-from autogluon.timeseries.models.ensemble import AbstractTimeSeriesEnsembleModel, TimeSeriesGreedyEnsemble
+from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel, TimeSeriesModelBase
+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
 
-logger = logging.getLogger("autogluon.timeseries.trainer")
+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[AbstractTimeSeriesModel]):
-    _cached_predictions_filename = "cached_predictions.pkl"
 
+class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
     max_rel_importance_score: float = 1e5
     eps_abs_importance_score: float = 1e-5
     max_ensemble_time_limit: float = 600.0
@@ -45,16 +47,16 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         self,
         path: str,
         prediction_length: int = 1,
-        eval_metric: Union[str, TimeSeriesScorer, None] = None,
-        eval_metric_seasonal_period: Optional[int] = None,
+        eval_metric: str | TimeSeriesScorer | None = None,
         save_data: bool = True,
         skip_model_selection: bool = False,
         enable_ensemble: bool = True,
         verbosity: int = 2,
-        val_splitter: Optional[AbstractWindowSplitter] = None,
-        refit_every_n_windows: Optional[int] = 1,
+        num_val_windows: tuple[int, ...] = (1,),
+        val_step_size: int | None = None,
+        refit_every_n_windows: int | None = 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__(
@@ -66,38 +68,39 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         self.prediction_length = prediction_length
         self.quantile_levels = kwargs.get("quantile_levels", [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])
         self.target = kwargs.get("target", "target")
-        self.metadata = kwargs.get("metadata", CovariateMetadata())
+        self.covariate_metadata = kwargs.get("covariate_metadata", CovariateMetadata())
         self.is_data_saved = False
         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 = TimeSeriesGreedyEnsemble
-        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 = 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: TimeSeriesScorer = check_get_evaluation_metric(eval_metric)
-        self.eval_metric_seasonal_period = eval_metric_seasonal_period
-        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.eval_metric = check_get_evaluation_metric(eval_metric, prediction_length=prediction_length)
+
+        self.num_val_windows = num_val_windows
+
+        # Validate num_val_windows
+        if len(self.num_val_windows) == 0:
+            raise ValueError("num_val_windows cannot be empty")
+        if not all(isinstance(w, int) and w > 0 for w in self.num_val_windows):
+            raise ValueError(f"num_val_windows must contain only positive integers, got {self.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:
@@ -115,14 +118,14 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         path = os.path.join(self.path_data, "train.pkl")
         return load_pkl.load(path=path)
 
-    def load_val_data(self) -> Optional[TimeSeriesDataFrame]:
+    def load_val_data(self) -> TimeSeriesDataFrame | None:
         path = os.path.join(self.path_data, "val.pkl")
         if os.path.exists(path):
             return load_pkl.load(path=path)
         else:
             return None
 
-    def load_data(self) -> Tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:
+    def load_data(self) -> tuple[TimeSeriesDataFrame, TimeSeriesDataFrame | None]:
         train_data = self.load_train_data()
         val_data = self.load_val_data()
         return train_data, val_data
@@ -137,24 +140,24 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
         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)
 
     def _add_model(
         self,
-        model: AbstractTimeSeriesModel,
-        base_models: Optional[List[str]] = None,
+        model: TimeSeriesModelBase,
+        base_models: list[str] | None = 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 : AbstractTimeSeriesModel
+        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``.
 
@@ -177,8 +180,8 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
             for base_model in base_models:
                 self.model_graph.add_edge(base_model, model.name)
 
-    def _get_model_levels(self) -> Dict[str, int]:
-        """Get a dictionary mapping each model to their level in the model graph"""
+    def _get_model_layers(self) -> dict[str, int]:
+        """Get a dictionary mapping each model to their layer in the model graph"""
 
         # get nodes without a parent
         rootset = set(self.model_graph.nodes)
@@ -191,14 +194,14 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
             for dest_node in paths_to:
                 paths_from[dest_node][source_node] = paths_to[dest_node]
 
-        # determine levels
-        levels = {}
+        # determine layers
+        layers = {}
         for n in paths_from:
-            levels[n] = max(paths_from[n].get(src, 0) for src in rootset)
+            layers[n] = max(paths_from[n].get(src, 0) for src in rootset)
 
-        return levels
+        return layers
 
-    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: list[str] | None = 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 = {}
@@ -216,28 +219,28 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         if len(models) == 1:
             return models[0]
         model_performances = self.get_models_attribute_dict(attribute="val_score")
-        model_levels = self._get_model_levels()
-        model_name_score_level_list = [
-            (m, model_performances[m], model_levels.get(m, 0)) for m in models if model_performances[m] is not None
+        model_layers = self._get_model_layers()
+        model_name_score_layer_list = [
+            (m, model_performances[m], model_layers.get(m, 0)) for m in models if model_performances[m] is not None
         ]
 
-        if not model_name_score_level_list:
+        if not model_name_score_layer_list:
             raise ValueError("No fitted models have validation scores computed.")
 
         # rank models in terms of validation score. if two models have the same validation score,
-        # rank them by their level in the model graph (lower level models are preferred).
+        # rank them by their layer in the model graph (lower layer models are preferred).
         return max(
-            model_name_score_level_list,
-            key=lambda mns: (mns[1], -mns[2]),  # (score, -level)
+            model_name_score_layer_list,
+            key=lambda mns: (mns[1], -mns[2]),  # (score, -layer)
         )[0]
 
-    def get_model_names(self, level: Optional[int] = None) -> List[str]:
+    def get_model_names(self, layer: int | None = 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
+        if layer is not None:
+            return list(node for node, l in self._get_model_layers().items() if l == layer)  # 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
@@ -262,32 +265,13 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
         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,
         train_data: TimeSeriesDataFrame,
-        time_limit: Optional[float] = None,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        hyperparameter_tune_kwargs: Union[str, dict] = "auto",
+        time_limit: float | None = None,
+        val_data: TimeSeriesDataFrame | None = None,
+        hyperparameter_tune_kwargs: str | dict = "auto",
     ):
         default_num_trials = None
         if time_limit is None and (
@@ -303,7 +287,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
                 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(use_val_data=val_data is not None),
                 refit_every_n_windows=self.refit_every_n_windows,
             )
         total_tuning_time = time.time() - tuning_start_time
@@ -313,11 +297,21 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         # add each of the trained HPO configurations to the trained models
         for model_hpo_name, model_info in hpo_models.items():
             model_path = os.path.join(self.path, model_info["path"])
+
             # Only load model configurations that didn't fail
-            if Path(model_path).exists():
-                model_hpo = self.load_model(model_hpo_name, path=model_path, model_type=type(model))
-                self._add_model(model_hpo)
-                model_names_trained.append(model_hpo.name)
+            if not Path(model_path).exists():
+                continue
+
+            model_hpo = self.load_model(model_hpo_name, path=model_path, model_type=type(model))
+
+            # override validation score to align evaluations on the final ensemble layer's window
+            if isinstance(model_hpo, MultiWindowBacktestingModel):
+                model_hpo.val_score = float(
+                    np.mean([info["val_score"] for info in model_hpo.info_per_val_window[-self.num_val_windows[-1] :]])
+                )
+
+            self._add_model(model_hpo)
+            model_names_trained.append(model_hpo.name)
 
         logger.info(f"\tTrained {len(model_names_trained)} models while tuning {model.name}.")
 
@@ -338,14 +332,15 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         self,
         train_data: TimeSeriesDataFrame,
         model: AbstractTimeSeriesModel,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[float] = None,
-    ) -> List[str]:
+        val_data: TimeSeriesDataFrame | None = None,
+        time_limit: float | None = None,
+    ) -> 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 = []
@@ -355,26 +350,46 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
                     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=None if isinstance(model, MultiWindowBacktestingModel) else val_data,
+                time_limit=time_limit,
+                verbosity=self.verbosity,
+                val_splitter=self._get_val_splitter(use_val_data=val_data is not None),
+                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)
 
             if time_limit is not None:
-                time_limit = fit_end_time - fit_start_time
-            if val_data is not None and not self.skip_model_selection:
+                time_limit = time_limit - (fit_end_time - fit_start_time)
+            if val_data is not None:
                 model.score_and_cache_oof(
                     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)
+            # by default, MultiWindowBacktestingModel computes validation score on all windows. However,
+            # when doing multi-layer stacking, the trainer only scores on the windows of the last layer.
+            # we override the val_score to align scores.
+            if isinstance(model, MultiWindowBacktestingModel):
+                model.val_score = float(
+                    np.mean([info["val_score"] for info in model.info_per_val_window[-self.num_val_windows[-1] :]])
+                )
+
+            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:
             logger.error(f"\tTime limit exceeded... Skipping {model.name}.")
-        except (Exception, MemoryError) as err:
+        except (Exception, MemoryError):
             logger.error(f"\tWarning: Exception caused {model.name} to fail during training... Skipping this model.")
-            logger.error(f"\t{err}")
-            logger.debug(traceback.format_exc())
+            logger.error(traceback.format_exc())
         else:
             self._add_model(model=model)  # noqa: F821
             model_names_trained.append(model.name)  # noqa: F821
@@ -383,45 +398,75 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
         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],
-        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,
-    ) -> List[str]:
+        hyperparameters: str | dict[Any, dict],
+        val_data: TimeSeriesDataFrame | None = None,
+        ensemble_hyperparameters: dict | list[dict] | None = None,
+        hyperparameter_tune_kwargs: str | dict | None = None,
+        excluded_model_types: list[str] | None = None,
+        time_limit: float | None = None,
+        random_seed: int | None = None,
+    ):
+        """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": {}}]
+        if isinstance(ensemble_hyperparameters, dict):
+            ensemble_hyperparameters = [ensemble_hyperparameters]
+        validate_ensemble_hyperparameters(ensemble_hyperparameters)
+
         time_start = time.time()
         hyperparameters = copy.deepcopy(hyperparameters)
 
+        if val_data is not None:
+            if self.num_val_windows[-1] != 1:
+                raise ValueError(
+                    f"When val_data is provided, the last element of num_val_windows must be 1, "
+                    f"got {self.num_val_windows[-1]}"
+                )
+        multi_window = self._get_val_splitter(use_val_data=val_data is not None).num_val_windows > 0
+
         if self.save_data and not self.is_data_saved:
             self.save_train_data(train_data)
             if val_data is not None:
                 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=multi_window,
             excluded_model_types=excluded_model_types,
         )
 
@@ -433,7 +478,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
                     "When `skip_model_selection=True`, only a single model must be provided via `hyperparameters` "
                     f"but {len(models)} models were given"
                 )
-            if contains_searchspace(models[0].get_user_params()):
+            if contains_searchspace(models[0].get_hyperparameters()):
                 raise ValueError(
                     "When `skip_model_selection=True`, model configuration should contain no search spaces."
                 )
@@ -461,7 +506,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
             if random_seed is not None:
                 seed_everything(random_seed + i)
 
-            if contains_searchspace(model.get_user_params()):
+            if contains_searchspace(model.get_hyperparameters()):
                 fit_log_message = f"Hyperparameter tuning model {model.name}. "
                 if time_left is not None:
                     fit_log_message += (
@@ -490,42 +535,16 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
                     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())
+        if self.enable_ensemble and ensemble_hyperparameters:
+            model_names = self.get_model_names(layer=0)
+            ensemble_names = self._fit_ensembles(
+                data_per_window=self._get_validation_windows(train_data, val_data),
+                predictions_per_window=self._get_base_model_predictions(model_names),
+                time_limit=None if time_limit is None else time_limit - (time.time() - time_start),
+                ensemble_hyperparameters=ensemble_hyperparameters,
+                num_windows_per_layer=self.num_val_windows,
+            )
+            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,73 +558,64 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
         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(
-        self, data_per_window: List[TimeSeriesDataFrame], model_names: List[str], time_limit: Optional[float] = None
-    ) -> str:
-        logger.info("Fitting simple weighted ensemble.")
-
-        model_preds: Dict[str, List[TimeSeriesDataFrame]] = {}
-        for model_name in model_names:
-            model_preds[model_name] = self._get_model_oof_predictions(model_name=model_name)
-
-        time_start = time.time()
-        ensemble = self.ensemble_model_type(
-            name=self._get_ensemble_model_name(),
+    def _fit_ensembles(
+        self,
+        *,
+        data_per_window: list[TimeSeriesDataFrame],
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        time_limit: float | None,
+        ensemble_hyperparameters: list[dict],
+        num_windows_per_layer: tuple[int, ...],
+    ) -> list[str]:
+        ensemble_composer = EnsembleComposer(
+            path=self.path,
+            prediction_length=self.prediction_length,
             eval_metric=self.eval_metric,
-            eval_metric_seasonal_period=self.eval_metric_seasonal_period,
             target=self.target,
-            prediction_length=self.prediction_length,
-            path=self.path,
-            freq=data_per_window[0].freq,
+            ensemble_hyperparameters=ensemble_hyperparameters,
+            num_windows_per_layer=num_windows_per_layer,
             quantile_levels=self.quantile_levels,
-            metadata=self.metadata,
+            model_graph=self.model_graph,
+        ).fit(
+            data_per_window=data_per_window,
+            predictions_per_window=predictions_per_window,
+            time_limit=time_limit,
         )
-        with warning_filter():
-            ensemble.fit_ensemble(model_preds, data_per_window=data_per_window, 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: model_preds[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,
+
+        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)
+
+        return ensembles_trained
+
+    def _get_validation_windows(self, train_data: TimeSeriesDataFrame, val_data: TimeSeriesDataFrame | None):
+        train_splitter = self._get_val_splitter(use_val_data=val_data is not None)
+        return [val_fold for _, val_fold in train_splitter.split(train_data)] + (
+            [] if val_data is None else [val_data]
+        )
+
+    def _get_val_splitter(self, use_val_data: bool = False) -> AbstractWindowSplitter:
+        num_windows_from_train = sum(self.num_val_windows[:-1]) if use_val_data else sum(self.num_val_windows)
+        return ExpandingWindowSplitter(
+            prediction_length=self.prediction_length,
+            num_val_windows=num_windows_from_train,
+            val_step_size=self.val_step_size,
         )
-        self._add_model(model=ensemble, base_models=ensemble.model_names)
-        self.save_model(model=ensemble)
-        return ensemble.name
+
+    def _get_base_model_predictions(self, model_names: list[str]) -> dict[str, list[TimeSeriesDataFrame]]:
+        """Get base model predictions for ensemble training / inference."""
+        predictions_per_window = {}
+        for model_name in model_names:
+            predictions_per_window[model_name] = self._get_model_oof_predictions(model_name)
+        return predictions_per_window
 
     def leaderboard(
         self,
-        data: Optional[TimeSeriesDataFrame] = None,
+        data: TimeSeriesDataFrame | None = None,
         extra_info: bool = False,
-        extra_metrics: Optional[List[Union[str, TimeSeriesScorer]]] = None,
+        extra_metrics: list[str | TimeSeriesScorer] | None = None,
         use_cache: bool = True,
     ) -> pd.DataFrame:
         logger.debug("Generating leaderboard for all models trained")
@@ -628,14 +638,15 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
                 if isinstance(model, MultiWindowBacktestingModel):
                     model = model.most_recent_model
                     assert model is not None
-                model_info[model_name]["hyperparameters"] = model.params
+                model_info[model_name]["hyperparameters"] = model.get_hyperparameters()
 
         if extra_metrics is None:
             extra_metrics = []
 
         if data is not None:
             past_data, known_covariates = data.get_model_inputs_for_scoring(
-                prediction_length=self.prediction_length, known_covariates_names=self.metadata.known_covariates
+                prediction_length=self.prediction_length,
+                known_covariates_names=self.covariate_metadata.known_covariates,
             )
             logger.info(
                 "Additional data provided, testing on additional data. Resulting leaderboard "
@@ -694,8 +705,8 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         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: 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":
@@ -719,7 +730,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
         return model_names
 
-    def unpersist(self, model_names: Union[Literal["all"], List[str]] = "all") -> List[str]:
+    def unpersist(self, model_names: Literal["all"] | list[str] = "all") -> list[str]:
         if model_names == "all":
             model_names = list(self.models.keys())
         if not isinstance(model_names, list):
@@ -731,9 +742,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
                 unpersisted_models.append(model)
         return unpersisted_models
 
-    def _get_model_for_prediction(
-        self, model: Optional[Union[str, AbstractTimeSeriesModel]] = None, verbose: bool = True
-    ) -> str:
+    def _get_model_for_prediction(self, model: str | TimeSeriesModelBase | None = None, verbose: bool = True) -> str:
         """Given an optional identifier or model object, return the name of the model with which to predict.
 
         If the model is not provided, this method will default to the best model according to the validation score.
@@ -749,18 +758,20 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
                 )
             return self.model_best
         else:
-            if isinstance(model, AbstractTimeSeriesModel):
+            if isinstance(model, TimeSeriesModelBase):
                 return model.name
             else:
+                if model not in self.get_model_names():
+                    raise KeyError(f"Model '{model}' not found. Available models: {self.get_model_names()}")
                 return model
 
     def predict(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
-        model: Optional[Union[str, AbstractTimeSeriesModel]] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
+        model: str | TimeSeriesModelBase | None = None,
         use_cache: bool = True,
-        random_seed: Optional[int] = None,
+        random_seed: int | None = None,
     ) -> TimeSeriesDataFrame:
         model_name = self._get_model_for_prediction(model)
         model_pred_dict, _ = self.get_model_pred_dict(
@@ -775,49 +786,57 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
             raise ValueError(f"Model {model_name} failed to predict. Please check the model's logs.")
         return predictions
 
+    def _get_eval_metric(self, metric: str | TimeSeriesScorer | None) -> TimeSeriesScorer:
+        if metric is None:
+            return self.eval_metric
+        else:
+            return check_get_evaluation_metric(
+                metric,
+                prediction_length=self.prediction_length,
+                seasonal_period=self.eval_metric.seasonal_period,
+                horizon_weight=self.eval_metric.horizon_weight,
+            )
+
     def _score_with_predictions(
         self,
         data: TimeSeriesDataFrame,
         predictions: TimeSeriesDataFrame,
-        metric: Union[str, TimeSeriesScorer, None] = None,
+        metric: str | TimeSeriesScorer | None = None,
     ) -> float:
         """Compute the score measuring how well the predictions align with the data."""
-        eval_metric = self.eval_metric if metric is None else check_get_evaluation_metric(metric)
-        return eval_metric.score(
+        return self._get_eval_metric(metric).score(
             data=data,
             predictions=predictions,
-            prediction_length=self.prediction_length,
             target=self.target,
-            seasonal_period=self.eval_metric_seasonal_period,
         )
 
     def score(
         self,
         data: TimeSeriesDataFrame,
-        model: Optional[Union[str, AbstractTimeSeriesModel]] = None,
-        metric: Union[str, TimeSeriesScorer, None] = None,
+        model: str | TimeSeriesModelBase | None = None,
+        metric: str | TimeSeriesScorer | None = None,
         use_cache: bool = True,
     ) -> float:
-        eval_metric = self.eval_metric if metric is None else check_get_evaluation_metric(metric)
+        eval_metric = self._get_eval_metric(metric)
         scores_dict = self.evaluate(data=data, model=model, metrics=[eval_metric], use_cache=use_cache)
         return scores_dict[eval_metric.name]
 
     def evaluate(
         self,
         data: TimeSeriesDataFrame,
-        model: Optional[Union[str, AbstractTimeSeriesModel]] = None,
-        metrics: Optional[Union[str, TimeSeriesScorer, List[Union[str, TimeSeriesScorer]]]] = None,
+        model: str | TimeSeriesModelBase | None = None,
+        metrics: str | TimeSeriesScorer | list[str | TimeSeriesScorer] | None = 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.metadata.known_covariates
+            prediction_length=self.prediction_length, known_covariates_names=self.covariate_metadata.known_covariates
         )
         predictions = self.predict(data=past_data, known_covariates=known_covariates, model=model, use_cache=use_cache)
 
         metrics_ = [metrics] if not isinstance(metrics, list) else metrics
         scores_dict = {}
         for metric in metrics_:
-            eval_metric = self.eval_metric if metric is None else check_get_evaluation_metric(metric)
+            eval_metric = self._get_eval_metric(metric)
             scores_dict[eval_metric.name] = self._score_with_predictions(
                 data=data, predictions=predictions, metric=eval_metric
             )
@@ -826,20 +845,20 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
     def get_feature_importance(
         self,
         data: TimeSeriesDataFrame,
-        features: List[str],
-        model: Optional[Union[str, AbstractTimeSeriesModel]] = None,
-        metric: Optional[Union[str, TimeSeriesScorer]] = None,
-        time_limit: Optional[float] = None,
+        features: list[str],
+        model: str | TimeSeriesModelBase | None = None,
+        metric: str | TimeSeriesScorer | None = None,
+        time_limit: float | None = None,
         method: Literal["naive", "permutation"] = "permutation",
         subsample_size: int = 50,
-        num_iterations: Optional[int] = None,
-        random_seed: Optional[int] = None,
+        num_iterations: int | None = None,
+        random_seed: int | None = None,
         relative_scores: bool = False,
         include_confidence_band: bool = True,
         confidence_level: float = 0.99,
     ) -> pd.DataFrame:
         assert method in ["naive", "permutation"], f"Invalid feature importance method {method}."
-        metric = check_get_evaluation_metric(metric) if metric is not None else self.eval_metric
+        eval_metric = self._get_eval_metric(metric)
 
         logger.info("Computing feature importance")
 
@@ -871,7 +890,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         )
 
         importance_transform = importance_transform_type(
-            covariate_metadata=self.metadata,
+            covariate_metadata=self.covariate_metadata,
             prediction_length=self.prediction_length,
             random_seed=random_seed,
         )
@@ -886,11 +905,13 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         for n in range(num_iterations):
             if subsample_size < data.num_items:
                 item_ids_sampled = data.item_ids.to_series().sample(subsample_size)  # noqa
-                data_sample: TimeSeriesDataFrame = data.query("item_id in @item_ids_sampled")  # type: ignore
+                data_sample: TimeSeriesDataFrame = data.query("item_id in @item_ids_sampled")
             else:
                 data_sample = data
 
-            base_score = self.evaluate(data=data_sample, model=model, metrics=metric, use_cache=False)[metric.name]
+            base_score = self.evaluate(data=data_sample, model=model, metrics=eval_metric, use_cache=False)[
+                eval_metric.name
+            ]
 
             for feature in features:
                 # override importance for unused features
@@ -898,9 +919,9 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
                     continue
                 else:
                     data_sample_replaced = importance_transform.transform(data_sample, feature_name=feature)
-                    score = self.evaluate(data=data_sample_replaced, model=model, metrics=metric, use_cache=False)[
-                        metric.name
-                    ]
+                    score = self.evaluate(
+                        data=data_sample_replaced, model=model, metrics=eval_metric, use_cache=False
+                    )[eval_metric.name]
 
                     importance = base_score - score
                     if relative_scores:
@@ -930,19 +951,85 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
         return importance_df
 
-    def _model_uses_feature(self, model: Union[str, AbstractTimeSeriesModel], feature: str) -> bool:
+    def _model_uses_feature(self, model: str | TimeSeriesModelBase, feature: str) -> bool:
         """Check if the given model uses the given feature."""
         models_with_ancestors = set(self.get_minimum_model_set(model))
 
-        if feature in self.metadata.static_features:
+        if feature in self.covariate_metadata.static_features:
             return any(self.load_model(m).supports_static_features for m in models_with_ancestors)
-        elif feature in self.metadata.known_covariates:
+        elif feature in self.covariate_metadata.known_covariates:
             return any(self.load_model(m).supports_known_covariates for m in models_with_ancestors)
-        elif feature in self.metadata.past_covariates:
+        elif feature in self.covariate_metadata.past_covariates:
             return any(self.load_model(m).supports_past_covariates for m in models_with_ancestors)
 
         return False
 
+    def backtest_predictions(
+        self,
+        data: TimeSeriesDataFrame | None,
+        model_names: list[str],
+        num_val_windows: int | None = None,
+        val_step_size: int | None = None,
+        use_cache: bool = True,
+    ) -> dict[str, list[TimeSeriesDataFrame]]:
+        if data is None:
+            assert num_val_windows is None, "num_val_windows must be None when data is None"
+            assert val_step_size is None, "val_step_size must be None when data is None"
+            return {model_name: self._get_model_oof_predictions(model_name) for model_name in model_names}
+
+        if val_step_size is None:
+            val_step_size = self.prediction_length
+        if num_val_windows is None:
+            num_val_windows = 1
+
+        splitter = ExpandingWindowSplitter(
+            prediction_length=self.prediction_length,
+            num_val_windows=num_val_windows,
+            val_step_size=val_step_size,
+        )
+
+        result: dict[str, list[TimeSeriesDataFrame]] = {model_name: [] for model_name in model_names}
+        for past_data, full_data in splitter.split(data):
+            known_covariates = full_data.slice_by_timestep(-self.prediction_length, None)[
+                self.covariate_metadata.known_covariates
+            ]
+            pred_dict, _ = self.get_model_pred_dict(
+                model_names=model_names,
+                data=past_data,
+                known_covariates=known_covariates,
+                use_cache=use_cache,
+            )
+            for model_name in model_names:
+                result[model_name].append(pred_dict[model_name])  # type: ignore
+
+        return result
+
+    def backtest_targets(
+        self,
+        data: TimeSeriesDataFrame | None,
+        num_val_windows: int | None = None,
+        val_step_size: int | None = None,
+    ) -> list[TimeSeriesDataFrame]:
+        if data is None:
+            assert num_val_windows is None, "num_val_windows must be None when data is None"
+            assert val_step_size is None, "val_step_size must be None when data is None"
+            train_data = self.load_train_data()
+            val_data = self.load_val_data()
+            return self._get_validation_windows(train_data=train_data, val_data=val_data)
+
+        if val_step_size is None:
+            val_step_size = self.prediction_length
+        if num_val_windows is None:
+            num_val_windows = 1
+
+        splitter = ExpandingWindowSplitter(
+            prediction_length=self.prediction_length,
+            num_val_windows=num_val_windows,
+            val_step_size=val_step_size,
+        )
+
+        return [val_fold for _, val_fold in splitter.split(data)]
+
     def _add_ci_to_feature_importance(
         self, importance_df: pd.DataFrame, confidence_level: float = 0.99
     ) -> pd.DataFrame:
@@ -972,10 +1059,10 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
     def _predict_model(
         self,
-        model: Union[str, AbstractTimeSeriesModel],
+        model: str | TimeSeriesModelBase,
         data: TimeSeriesDataFrame,
-        model_pred_dict: Dict[str, Optional[TimeSeriesDataFrame]],
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        model_pred_dict: dict[str, TimeSeriesDataFrame | None],
+        known_covariates: TimeSeriesDataFrame | None = None,
     ) -> TimeSeriesDataFrame:
         """Generate predictions using the given model.
 
@@ -988,10 +1075,10 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
     def _get_inputs_to_model(
         self,
-        model: Union[str, AbstractTimeSeriesModel],
+        model: str | TimeSeriesModelBase,
         data: TimeSeriesDataFrame,
-        model_pred_dict: Dict[str, Optional[TimeSeriesDataFrame]],
-    ) -> Union[TimeSeriesDataFrame, Dict[str, Optional[TimeSeriesDataFrame]]]:
+        model_pred_dict: dict[str, TimeSeriesDataFrame | None],
+    ) -> TimeSeriesDataFrame | dict[str, TimeSeriesDataFrame | None]:
         """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.
@@ -1007,13 +1094,13 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
     def get_model_pred_dict(
         self,
-        model_names: List[str],
+        model_names: list[str],
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         raise_exception_if_failed: bool = True,
         use_cache: bool = True,
-        random_seed: Optional[int] = None,
-    ) -> Tuple[Dict[str, Optional[TimeSeriesDataFrame]], Dict[str, float]]:
+        random_seed: int | None = None,
+    ) -> tuple[dict[str, TimeSeriesDataFrame | None], dict[str, float]]:
         """Return a dictionary with predictions of all models for the given dataset.
 
         Parameters
@@ -1033,20 +1120,20 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         use_cache
             If False, will ignore the cache even if it's available.
         """
-        # TODO: Unify design of the method with Tabular
-        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:
             model_set.update(self.get_minimum_model_set(model_name))
         if len(model_set) > 1:
-            model_to_level = self._get_model_levels()
-            model_set = sorted(model_set, key=model_to_level.get)  # type: ignore
+            model_to_layer = self._get_model_layers()
+            model_set = sorted(model_set, key=model_to_layer.get)  # type: ignore
         logger.debug(f"Prediction order: {model_set}")
 
         failed_models = []
@@ -1072,9 +1159,11 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
         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,
             )
@@ -1085,7 +1174,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
         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):
@@ -1093,57 +1182,8 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
                     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 _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. Otherwise returns empty dicts."""
-        if self._cached_predictions_path.exists():
-            cached_predictions = load_pkl.load(str(self._cached_predictions_path))
-            if dataset_hash in cached_predictions:
-                model_pred_dict = cached_predictions[dataset_hash]["model_pred_dict"]
-                pred_time_dict = cached_predictions[dataset_hash]["pred_time_dict"]
-                if model_pred_dict.keys() == pred_time_dict.keys():
-                    logger.debug(f"Loaded cached predictions for models {list(model_pred_dict.keys())}")
-                    return model_pred_dict, pred_time_dict
-                else:
-                    logger.warning(f"Found corrupted cached predictions in {self._cached_predictions_path}")
-        logger.debug("Found no cached predictions")
-        return {}, {}
-
-    def _save_cached_pred_dicts(
-        self,
-        dataset_hash: str,
-        model_pred_dict: Dict[str, Optional[TimeSeriesDataFrame]],
-        pred_time_dict: Dict[str, float],
-    ) -> None:
-        # TODO: Save separate file for each dataset if _cached_predictions file grows large?
-        if self._cached_predictions_path.exists():
-            logger.debug("Extending existing cached predictions")
-            cached_predictions = load_pkl.load(str(self._cached_predictions_path))
-        else:
-            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]
+        self, train_data: TimeSeriesDataFrame, val_data: TimeSeriesDataFrame | None
     ) -> TimeSeriesDataFrame:
         if val_data is None:
             return train_data
@@ -1153,10 +1193,10 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
 
     def refit_single_full(
         self,
-        train_data: Optional[TimeSeriesDataFrame] = None,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        models: Optional[List[str]] = None,
-    ) -> List[str]:
+        train_data: TimeSeriesDataFrame | None = None,
+        val_data: TimeSeriesDataFrame | None = None,
+        models: list[str] | None = 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)
@@ -1164,16 +1204,17 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         if models is None:
             models = self.get_model_names()
 
-        model_to_level = self._get_model_levels()
-        models_sorted_by_level = sorted(models, key=model_to_level.get)  # type: ignore
+        model_to_layer = self._get_model_layers()
+        models_sorted_by_layer = sorted(models, key=model_to_layer.get)  # type: ignore
 
         model_refit_map = {}
         models_trained_full = []
-        for model in models_sorted_by_level:
+        for model in models_sorted_by_layer:
             model = self.load_model(model)
             model_name = model.name
             if model._get_tags()["can_refit_full"]:
                 model_full = model.convert_to_refit_full_template()
+                assert isinstance(model_full, AbstractTimeSeriesModel)
                 logger.info(f"Fitting model: {model_full.name}")
                 models_trained = self._train_and_save(
                     train_data=refit_full_data,
@@ -1199,7 +1240,7 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         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":
@@ -1231,71 +1272,27 @@ class TimeSeriesTrainer(AbstractTrainer[AbstractTimeSeriesModel]):
         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: str | dict[str, Any],
         *,
         multi_window: bool = False,
-        freq: Optional[str] = None,
-        excluded_model_types: Optional[List[str]] = None,
+        freq: str | None = None,
+        excluded_model_types: list[str] | None = None,
         hyperparameter_tune: bool = False,
-    ) -> List[AbstractTimeSeriesModel]:
-        return get_preset_models(
+    ) -> list[AbstractTimeSeriesModel]:
+        return TrainableModelSetBuilder(
+            freq=freq,
+            prediction_length=self.prediction_length,
             path=self.path,
             eval_metric=self.eval_metric,
-            eval_metric_seasonal_period=self.eval_metric_seasonal_period,
-            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,
-            metadata=self.metadata,
-            excluded_model_types=excluded_model_types,
-            # if skip_model_selection = True, we skip backtesting
+            covariate_metadata=self.covariate_metadata,
             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(),
         )