autogluon.timeseries 1.4.1b20250820__py3-none-any.whl → 1.4.1b20250901__py3-none-any.whl

autogluon/timeseries/trainer/prediction_cache.py

@@ -0,0 +1,149 @@
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Optional
+
+from autogluon.common.utils.utils import hash_pandas_df
+from autogluon.core.utils.loaders import load_pkl
+from autogluon.core.utils.savers import save_pkl
+from autogluon.timeseries import TimeSeriesDataFrame
+
+logger = logging.getLogger(__name__)
+
+
+class PredictionCache(ABC):
+    """A prediction cache is an abstract key-value store for time series predictions. The storage is keyed by
+    (data, known_covariates) pairs and stores (model_pred_dict, pred_time_dict) pair values. In this stored pair,
+    (model_pred_dict, pred_time_dict), both dictionaries are keyed by model names.
+    """
+
+    def __init__(self, root_path: str):
+        self.root_path = Path(root_path)
+
+    @abstractmethod
+    def get(
+        self, data: TimeSeriesDataFrame, known_covariates: Optional[TimeSeriesDataFrame]
+    ) -> tuple[dict[str, Optional[TimeSeriesDataFrame]], dict[str, float]]:
+        pass
+
+    @abstractmethod
+    def put(
+        self,
+        data: TimeSeriesDataFrame,
+        known_covariates: Optional[TimeSeriesDataFrame],
+        model_pred_dict: dict[str, Optional[TimeSeriesDataFrame]],
+        pred_time_dict: dict[str, float],
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def clear(self) -> None:
+        pass
+
+
+def get_prediction_cache(use_cache: bool, root_path: str) -> PredictionCache:
+    if use_cache:
+        return FileBasedPredictionCache(root_path=root_path)
+    else:
+        return NoOpPredictionCache(root_path=root_path)
+
+
+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
+
+
+class NoOpPredictionCache(PredictionCache):
+    """A dummy (no-op) prediction cache."""
+
+    def get(
+        self, data: TimeSeriesDataFrame, known_covariates: Optional[TimeSeriesDataFrame]
+    ) -> tuple[dict[str, Optional[TimeSeriesDataFrame]], dict[str, float]]:
+        return {}, {}
+
+    def put(
+        self,
+        data: TimeSeriesDataFrame,
+        known_covariates: Optional[TimeSeriesDataFrame],
+        model_pred_dict: dict[str, Optional[TimeSeriesDataFrame]],
+        pred_time_dict: dict[str, float],
+    ) -> None:
+        pass
+
+    def clear(self) -> None:
+        pass
+
+
+class FileBasedPredictionCache(PredictionCache):
+    """A file-backed cache of model predictions."""
+
+    _cached_predictions_filename = "cached_predictions.pkl"
+
+    @property
+    def path(self) -> Path:
+        return Path(self.root_path) / self._cached_predictions_filename
+
+    def get(
+        self, data: TimeSeriesDataFrame, known_covariates: Optional[TimeSeriesDataFrame]
+    ) -> tuple[dict[str, Optional[TimeSeriesDataFrame]], dict[str, float]]:
+        dataset_hash = compute_dataset_hash(data, known_covariates)
+        return self._get_cached_pred_dicts(dataset_hash)
+
+    def put(
+        self,
+        data: TimeSeriesDataFrame,
+        known_covariates: Optional[TimeSeriesDataFrame],
+        model_pred_dict: dict[str, Optional[TimeSeriesDataFrame]],
+        pred_time_dict: dict[str, float],
+    ) -> None:
+        dataset_hash = compute_dataset_hash(data, known_covariates)
+        self._save_cached_pred_dicts(dataset_hash, model_pred_dict, pred_time_dict)
+
+    def clear(self) -> None:
+        if self.path.exists():
+            logger.debug(f"Removing existing cached predictions file {self.path}")
+            self.path.unlink()
+
+    def _load_cached_predictions(self) -> dict[str, dict[str, dict[str, Any]]]:
+        if self.path.exists():
+            try:
+                cached_predictions = load_pkl.load(str(self.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.path), object=cached_predictions)
+        logger.debug(f"Cached predictions saved to {self.path}")

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, Type, 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
@@ -22,7 +22,6 @@ from autogluon.timeseries.metrics import TimeSeriesScorer, check_get_evaluation_
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel, TimeSeriesModelBase
 from autogluon.timeseries.models.ensemble import AbstractTimeSeriesEnsembleModel, GreedyEnsemble
 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.utils.features import (
     ConstantReplacementFeatureImportanceTransform,
@@ -31,12 +30,13 @@ from autogluon.timeseries.utils.features import (
 )
 from autogluon.timeseries.utils.warning_filters import disable_tqdm, warning_filter
 
+from .model_set_builder import TrainableModelSetBuilder, contains_searchspace
+from .prediction_cache import PredictionCache, get_prediction_cache
+
 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
@@ -81,7 +81,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
         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 = {}
 
@@ -91,12 +91,10 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         assert isinstance(val_splitter, AbstractWindowSplitter), "val_splitter must be of type AbstractWindowSplitter"
         self.val_splitter = val_splitter
         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 +119,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 +134,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 +142,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 +174,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 +195,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 +228,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
@@ -339,12 +337,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 = []
@@ -397,13 +396,13 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
     def _train_multi(
         self,
         train_data: TimeSeriesDataFrame,
-        hyperparameters: Union[str, Dict],
+        hyperparameters: Union[str, dict],
         val_data: Optional[TimeSeriesDataFrame] = 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]:
+    ) -> list[str]:
         logger.info(f"\nStarting training. Start time is {time.strftime('%Y-%m-%d %H:%M:%S')}")
 
         time_start = time.time()
@@ -415,7 +414,7 @@ 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,
@@ -439,8 +438,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
@@ -541,7 +538,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
     def _get_ensemble_oof_data(
         self, train_data: TimeSeriesDataFrame, val_data: Optional[TimeSeriesDataFrame]
-    ) -> List[TimeSeriesDataFrame]:
+    ) -> list[TimeSeriesDataFrame]:
         if val_data is None:
             return [val_fold for _, val_fold in self.val_splitter.split(train_data)]
         else:
@@ -558,13 +555,13 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
     def fit_ensemble(
         self,
-        data_per_window: List[TimeSeriesDataFrame],
-        model_names: List[str],
+        data_per_window: list[TimeSeriesDataFrame],
+        model_names: list[str],
         time_limit: Optional[float] = None,
     ) -> str:
         logger.info("Fitting simple weighted ensemble.")
 
-        predictions_per_window: Dict[str, List[TimeSeriesDataFrame]] = {}
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]] = {}
         base_model_scores = self.get_models_attribute_dict(attribute="val_score", models=self.get_model_names(0))
 
         for model_name in model_names:
@@ -614,7 +611,7 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
         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 +701,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 +726,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 +823,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 +843,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 +993,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 +1009,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 +1026,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 +1052,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 +1091,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 +1106,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 +1114,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 +1127,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 +1172,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,38 +1204,38 @@ 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,
+        ).get_model_set(
+            hyperparameters=hyperparameters,
+            hyperparameter_tune=hyperparameter_tune,
+            excluded_model_types=excluded_model_types,
+            banned_model_names=self._get_banned_model_names(),
         )
 
     def fit(
         self,
         train_data: TimeSeriesDataFrame,
-        hyperparameters: Union[str, Dict[Any, Dict]],
+        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,
+        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,
     ):
@@ -1301,13 +1245,13 @@ class TimeSeriesTrainer(AbstractTrainer[TimeSeriesModelBase]):
 
         Parameters
         ----------
-        train_data: TimeSeriesDataFrame
+        train_data
             Training data for fitting time series timeseries models.
-        hyperparameters: str or Dict
+        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: TimeSeriesDataFrame
+        val_data
             Optional validation data set to report validation scores on.
         hyperparameter_tune_kwargs
             Args for hyperparameter tuning

autogluon/timeseries/transforms/covariate_scaler.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Dict, List, Literal, Optional, Protocol, overload, runtime_checkable
+from typing import Literal, Optional, Protocol, overload, runtime_checkable
 
 import numpy as np
 import pandas as pd
@@ -53,7 +53,7 @@ class GlobalCovariateScaler(CovariateScaler):
         self.use_past_covariates = use_past_covariates
         self.use_static_features = use_static_features
         self.skew_threshold = skew_threshold
-        self._column_transformers: Optional[Dict[Literal["known", "past", "static"], ColumnTransformer]] = None
+        self._column_transformers: Optional[dict[Literal["known", "past", "static"], ColumnTransformer]] = None
 
     def is_fit(self) -> bool:
         return self._column_transformers is not None
@@ -117,7 +117,7 @@ class GlobalCovariateScaler(CovariateScaler):
             known_covariates[columns] = self._column_transformers["known"].transform(known_covariates[columns])
         return known_covariates
 
-    def _get_transformer_for_columns(self, df: pd.DataFrame, columns: List[str]) -> ColumnTransformer:
+    def _get_transformer_for_columns(self, df: pd.DataFrame, columns: list[str]) -> ColumnTransformer:
         """Passthrough bool features, use QuantileTransform for skewed features, and use StandardScaler for the rest.
 
         The preprocessing logic is similar to the TORCH_NN model from Tabular.