autogluon.timeseries 1.2.1b20250219__tar.gz → 1.2.1b20250221__tar.gz

{autogluon.timeseries-1.2.1b20250219 → autogluon.timeseries-1.2.1b20250221}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: autogluon.timeseries
-Version: 1.2.1b20250219
+Version: 1.2.1b20250221
 Summary: Fast and Accurate ML in 3 Lines of Code
 Home-page: https://github.com/autogluon/autogluon
 Author: AutoGluon Community

{autogluon.timeseries-1.2.1b20250219 → autogluon.timeseries-1.2.1b20250221}/src/autogluon/timeseries/models/abstract/abstract_timeseries_model.py

@@ -1,9 +1,12 @@
+from __future__ import annotations
+
+import copy
 import logging
 import os
 import re
 import time
 from contextlib import nullcontext
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import pandas as pd
 from typing_extensions import Self
@@ -11,10 +14,16 @@ from typing_extensions import Self
 from autogluon.common import space
 from autogluon.common.loaders import load_pkl
 from autogluon.common.savers import save_pkl
-from autogluon.core.constants import AG_ARGS_FIT, REFIT_FULL_SUFFIX
+from autogluon.common.utils.distribute_utils import DistributedContext
+from autogluon.common.utils.log_utils import DuplicateFilter
+from autogluon.common.utils.resource_utils import get_resource_manager
+from autogluon.common.utils.try_import import try_import_ray
+from autogluon.common.utils.utils import setup_outputdir
+from autogluon.core.constants import AG_ARG_PREFIX, AG_ARGS_FIT, REFIT_FULL_SUFFIX
+from autogluon.core.hpo.constants import CUSTOM_BACKEND, RAY_BACKEND
 from autogluon.core.hpo.exceptions import EmptySearchSpace
-from autogluon.core.hpo.executors import HpoExecutor, RayHpoExecutor
-from autogluon.core.models import AbstractModel
+from autogluon.core.hpo.executors import HpoExecutor, HpoExecutorFactory, RayHpoExecutor
+from autogluon.core.models import ModelBase
 from autogluon.core.utils.exceptions import TimeLimitExceeded
 from autogluon.timeseries.dataset import TimeSeriesDataFrame
 from autogluon.timeseries.metrics import TimeSeriesScorer, check_get_evaluation_metric
@@ -32,9 +41,88 @@ from autogluon.timeseries.utils.warning_filters import disable_stdout, warning_f
 from .model_trial import model_trial, skip_hpo
 
 logger = logging.getLogger(__name__)
+dup_filter = DuplicateFilter()
+logger.addFilter(dup_filter)
+
 
+# TODO: refactor and move to util. We do not need to use "params_aux" in time series
+def check_and_split_hyperparameters(
+    params: Optional[Dict[str, Any]] = None,
+    ag_args_fit: str = AG_ARGS_FIT,
+    ag_arg_prefix: str = AG_ARG_PREFIX,
+) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """
+    Given the user-specified hyperparameters, split into `params` and `params_aux`.
 
-class AbstractTimeSeriesModel(AbstractModel):
+    Parameters
+    ----------
+    params : Optional[Dict[str, Any]], default = None
+        The model hyperparameters dictionary
+    ag_args_fit : str, default = "ag_args_fit"
+        The params key to look for that contains params_aux.
+        If the key is present, the value is used for params_aux and popped from params.
+        If no such key is found, then initialize params_aux as an empty dictionary.
+    ag_arg_prefix : str, default = "ag."
+        The key prefix to look for that indicates a parameter is intended for params_aux.
+        If None, this logic is skipped.
+        If a key starts with this prefix, it is popped from params and added to params_aux with the prefix removed.
+        For example:
+            input:  params={'ag.foo': 2, 'abc': 7}, params_aux={'bar': 3}, and ag_arg_prefix='.ag',
+            output: params={'abc': 7}, params_aux={'bar': 3, 'foo': 2}
+        In cases where the key is specified multiple times, the value of the key with the prefix will always take priority.
+        A warning will be logged if a key is present multiple times.
+        For example, given the most complex scenario:
+            input:  params={'ag.foo': 1, 'foo': 2, 'ag_args_fit': {'ag.foo': 3, 'foo': 4}}
+            output: params={'foo': 2}, params_aux={'foo': 1}
+
+    Returns
+    -------
+    params, params_aux : (Dict[str, Any], Dict[str, Any])
+        params will contain the native model hyperparameters
+        params_aux will contain special auxiliary hyperparameters
+    """
+    params = copy.deepcopy(params) if params is not None else dict()
+    assert isinstance(params, dict), f"Invalid dtype of params! Expected dict, but got {type(params)}"
+    for k in params.keys():
+        if not isinstance(k, str):
+            logger.warning(
+                f"Warning: Specified hyperparameter key is not of type str: {k} (type={type(k)}). "
+                f"There might be a bug in your configuration."
+            )
+
+    params_aux = params.pop(ag_args_fit, dict())
+    if params_aux is None:
+        params_aux = dict()
+    assert isinstance(params_aux, dict), f"Invalid dtype of params_aux! Expected dict, but got {type(params_aux)}"
+    if ag_arg_prefix is not None:
+        param_aux_keys = list(params_aux.keys())
+        for k in param_aux_keys:
+            if isinstance(k, str) and k.startswith(ag_arg_prefix):
+                k_no_prefix = k[len(ag_arg_prefix) :]
+                if k_no_prefix in params_aux:
+                    logger.warning(
+                        f'Warning: hyperparameter "{k}" is present '
+                        f'in `ag_args_fit` as both "{k}" and "{k_no_prefix}". '
+                        f'Will use "{k}" and ignore "{k_no_prefix}".'
+                    )
+                params_aux[k_no_prefix] = params_aux.pop(k)
+        param_keys = list(params.keys())
+        for k in param_keys:
+            if isinstance(k, str) and k.startswith(ag_arg_prefix):
+                k_no_prefix = k[len(ag_arg_prefix) :]
+                if k_no_prefix in params_aux:
+                    logger.warning(
+                        f'Warning: hyperparameter "{k}" is present '
+                        f"in both `ag_args_fit` and `hyperparameters`. "
+                        f"Will use `hyperparameters` value."
+                    )
+                params_aux[k_no_prefix] = params.pop(k)
+    return params, params_aux
+
+
+# TODO: refactor. remove params_aux, etc. make class inherit from ABC, make overrides and abstract
+# methods clear, change name to TimeSeriesModel, et al.
+class AbstractTimeSeriesModel(ModelBase):
     """Abstract class for all `Model` objects in autogluon.timeseries.
 
     Parameters
@@ -67,49 +155,51 @@ class AbstractTimeSeriesModel(AbstractModel):
         If None, model defaults are used. This is identical to passing an empty dictionary.
     """
 
+    model_file_name = "model.pkl"
+    model_info_name = "info.pkl"
     _oof_filename = "oof.pkl"
+
     # TODO: For which models should we override this parameter?
     _covariate_regressor_fit_time_fraction: float = 0.5
     default_max_time_limit_ratio: float = 0.9
 
-    # TODO: This is a hack to override the AbstractModel method, which the HPO module
-    # also circumvents with an ugly None-check.
-    estimate_memory_usage: Callable = None  # type: ignore
-
     _supports_known_covariates: bool = False
     _supports_past_covariates: bool = False
     _supports_static_features: bool = False
 
     def __init__(
         self,
-        freq: Optional[str] = None,
-        prediction_length: int = 1,
         path: Optional[str] = None,
         name: Optional[str] = None,
+        hyperparameters: Optional[Dict[str, Any]] = None,
+        freq: Optional[str] = None,
+        prediction_length: int = 1,
         metadata: Optional[CovariateMetadata] = None,
+        target: str = "target",
+        quantile_levels: Sequence[float] = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9),
         eval_metric: Union[str, TimeSeriesScorer, None] = None,
         eval_metric_seasonal_period: Optional[int] = None,
-        hyperparameters: Optional[Dict[str, Union[int, float, str, space.Space]]] = None,
-        **kwargs,
     ):
-        name = name or re.sub(r"Model$", "", self.__class__.__name__)
-        super().__init__(
-            path=path,
-            name=name,
-            problem_type=None,
-            eval_metric=None,
-            hyperparameters=hyperparameters,
-        )
+        self.name = name or re.sub(r"Model$", "", self.__class__.__name__)
+
+        self.path_root = path
+        if self.path_root is None:
+            path_suffix = self.name
+            # TODO: Would be ideal to not create dir, but still track that it is unique. However, this isn't possible to do without a global list of used dirs or using UUID.
+            path_cur = setup_outputdir(path=None, create_dir=True, path_suffix=path_suffix)
+            self.path_root = path_cur.rsplit(self.name, 1)[0]
+            logger.log(20, f"Warning: No path was specified for model, defaulting to: {self.path_root}")
+
+        self.path = os.path.join(self.path_root, self.name)
+
         self.eval_metric: TimeSeriesScorer = check_get_evaluation_metric(eval_metric)
         self.eval_metric_seasonal_period = eval_metric_seasonal_period
-        self.problem_type = "timeseries"
-        self.conformalize = False
-        self.target: str = kwargs.get("target", "target")
+        self.target: str = target
         self.metadata = metadata or CovariateMetadata()
 
         self.freq: Optional[str] = freq
         self.prediction_length: int = 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.quantile_levels: list[float] = list(quantile_levels)
 
         if not all(0 < q < 1 for q in self.quantile_levels):
             raise ValueError("Invalid quantile_levels specified. Quantiles must be between 0 and 1 (exclusive).")
@@ -126,29 +216,69 @@ class AbstractTimeSeriesModel(AbstractModel):
         self.target_scaler: Optional[LocalTargetScaler] = None
         self.covariate_scaler: Optional[CovariateScaler] = None
         self.covariate_regressor: Optional[CovariateRegressor] = None
-        self.fit_time: Optional[float]
+
+        # TODO: remove the variables below
+        self.model = None
+
+        self._is_initialized = False
+        self._user_params, self._user_params_aux = check_and_split_hyperparameters(hyperparameters)
+
+        self.params = {}
+        self.params_aux = {}
+        self.nondefault_params: List[str] = []
+
+        self.fit_time: Optional[float] = None  # Time taken to fit in seconds (Training data)
+        self.predict_time: Optional[float] = None  # Time taken to predict in seconds (Validation data)
+        self.predict_1_time: Optional[float] = (
+            None  # Time taken to predict 1 row of data in seconds (with batch size `predict_1_batch_size` in params_aux)
+        )
+        self.val_score: Optional[float] = None  # Score with eval_metric (Validation data)
 
     def __repr__(self) -> str:
         return self.name
 
+    def rename(self, name: str) -> None:
+        if self.name is not None and len(self.name) > 0:
+            self.path = os.path.join(os.path.dirname(self.path), name)
+        else:
+            self.path = os.path.join(self.path, name)
+        self.name = name
+
+    def set_contexts(self, path_context):
+        self.path = path_context
+        self.path_root = self.path.rsplit(self.name, 1)[0]
+
     def save(self, path: Optional[str] = None, verbose=True) -> str:
+        if path is None:
+            path = self.path
+
         # Save self._oof_predictions as a separate file, not model attribute
         if self._oof_predictions is not None:
             save_pkl.save(
-                path=os.path.join(self.path, "utils", self._oof_filename),
+                path=os.path.join(path, "utils", self._oof_filename),
                 object=self._oof_predictions,
                 verbose=verbose,
             )
         oof_predictions = self._oof_predictions
         self._oof_predictions = None
-        save_path = super().save(path=path, verbose=verbose)
+
+        file_path = os.path.join(path, self.model_file_name)
+        _model = self.model
+        save_pkl.save(path=file_path, object=self, verbose=verbose)
+        self.model = _model
+
         self._oof_predictions = oof_predictions
-        return save_path
+        return path
 
     @classmethod
     def load(cls, path: str, reset_paths: bool = True, load_oof: bool = False, verbose: bool = True) -> Self:  # type: ignore
-        # TODO: align method signature in new AbstractModel class
-        model = super().load(path=path, reset_paths=reset_paths, verbose=verbose)
+        file_path = os.path.join(path, cls.model_file_name)
+        model = load_pkl.load(path=file_path, verbose=verbose)
+        if reset_paths:
+            model.set_contexts(path)
+        if hasattr(model, "_compiler"):
+            if model._compiler is not None and not model._compiler.save_in_pkl:
+                model.model = model._compiler.load(path=path)
         if load_oof and model._oof_predictions is None:
             model._oof_predictions = cls.load_oof_predictions(path=path, verbose=verbose)
         return model
@@ -181,7 +311,6 @@ class AbstractTimeSeriesModel(AbstractModel):
         return self._oof_predictions
 
     def _get_default_auxiliary_params(self) -> dict:
-        # TODO: refine to values that are absolutely necessary
         return dict(
             # ratio of given time_limit to use during fit(). If time_limit == 10 and max_time_limit_ratio=0.3,
             # time_limit would be changed to 3.
@@ -191,27 +320,45 @@ class AbstractTimeSeriesModel(AbstractModel):
             max_time_limit=None,
         )
 
-    def initialize(self, **kwargs) -> dict:
-        # TODO: remove **kwargs from method signature
-        # TODO: do we even need deferred initialization?
+    # TODO: remove
+    @classmethod
+    def _get_default_ag_args(cls) -> dict:
+        return {}
+
+    def _init_params(self):
+        """Initializes model hyperparameters"""
+        hyperparameters = self._user_params
+        self.nondefault_params = []
+        if hyperparameters is not None:
+            self.params.update(hyperparameters)
+            # These are hyperparameters that user has specified.
+            self.nondefault_params = list(hyperparameters.keys())[:]
+        self.params_trained = {}
+
+    def _init_params_aux(self):
+        """
+        Initializes auxiliary hyperparameters.
+        These parameters are generally not model specific and can have a wide variety of effects.
+        For documentation on some of the available options and their defaults, refer to `self._get_default_auxiliary_params`.
+        """
+        hyperparameters_aux = self._user_params_aux or {}
+        self.params_aux = {**self._get_default_auxiliary_params(), **hyperparameters_aux}
 
+    def initialize(self) -> None:
         if not self._is_initialized:
             self._init_params_aux()
             self._init_params()
             self._initialize_transforms()
             self._is_initialized = True
 
-        # TODO: remove
-        kwargs.pop("feature_metadata", None)
-        kwargs.pop("num_classes", None)
-
-        return kwargs
-
     def _initialize_transforms(self) -> None:
         self.target_scaler = self._create_target_scaler()
         self.covariate_scaler = self._create_covariate_scaler()
         self.covariate_regressor = self._create_covariate_regressor()
 
+    def _get_model_params(self) -> dict:
+        return self.params.copy()
+
     def get_params(self) -> dict:
         # TODO: do not extract to AbstractModel if this is only used for getting a
         # prototype of the object for HPO.
@@ -222,7 +369,6 @@ class AbstractTimeSeriesModel(AbstractModel):
         return dict(
             path=self.path_root,
             name=self.name,
-            problem_type=self.problem_type,
             eval_metric=self.eval_metric,
             hyperparameters=hyperparameters,
             freq=self.freq,
@@ -232,6 +378,19 @@ class AbstractTimeSeriesModel(AbstractModel):
             target=self.target,
         )
 
+    @classmethod
+    def load_info(cls, path: str, load_model_if_required: bool = True) -> dict:
+        # TODO: remove?
+        load_path = os.path.join(path, cls.model_info_name)
+        try:
+            return load_pkl.load(path=load_path)
+        except:
+            if load_model_if_required:
+                model = cls.load(path=path, reset_paths=True)
+                return model.get_info()
+            else:
+                raise
+
     def get_info(self) -> dict:
         """
         Returns a dictionary of numerous fields describing the model.
@@ -256,6 +415,7 @@ class AbstractTimeSeriesModel(AbstractModel):
         train_data: TimeSeriesDataFrame,
         val_data: Optional[TimeSeriesDataFrame] = None,
         time_limit: Optional[float] = None,
+        verbosity: int = 2,
         **kwargs,
     ) -> Self:
         """Fit timeseries model.
@@ -291,9 +451,8 @@ class AbstractTimeSeriesModel(AbstractModel):
         model: AbstractTimeSeriesModel
             The fitted model object
         """
-        # TODO: align method signature in new AbstractModel as fit(*args, **kwargs)
         start_time = time.monotonic()
-        self.initialize(**kwargs)
+        self.initialize()
 
         if self.target_scaler is not None:
             train_data = self.target_scaler.fit_transform(train_data)
@@ -308,7 +467,7 @@ class AbstractTimeSeriesModel(AbstractModel):
             self.covariate_regressor.fit(
                 train_data,
                 time_limit=covariate_regressor_time_limit,
-                verbosity=kwargs.get("verbosity", 2) - 1,
+                verbosity=verbosity,
             )
 
         if self._get_tags()["can_use_train_data"]:
@@ -334,15 +493,12 @@ class AbstractTimeSeriesModel(AbstractModel):
                 )
                 raise TimeLimitExceeded
 
-        # TODO: disentangle fit_resources computation and validation from tabular logic
-        kwargs = self._preprocess_fit_resources(**kwargs)
-        self.validate_fit_resources(**kwargs)
-
         self._fit(
             train_data=train_data,
             val_data=val_data,
             time_limit=time_limit,
-            **kwargs,
+            verbosity=verbosity,
+            **(self._get_system_resources() | kwargs),
         )
 
         return self
@@ -369,6 +525,32 @@ class AbstractTimeSeriesModel(AbstractModel):
 
         return time_limit
 
+    def _fit(  # type: ignore
+        self,
+        train_data: TimeSeriesDataFrame,
+        val_data: Optional[TimeSeriesDataFrame] = None,
+        time_limit: Optional[float] = None,
+        num_cpus: Optional[int] = None,
+        num_gpus: Optional[int] = None,
+        verbosity: int = 2,
+        **kwargs,
+    ) -> None:
+        """Private method for `fit`. See `fit` for documentation of arguments. Apart from
+        the model training logic, `fit` additionally implements other logic such as keeping
+        track of the time limit, etc.
+        """
+        # TODO: Make the models respect `num_cpus` and `num_gpus` parameters
+        raise NotImplementedError
+
+    # TODO: perform this check inside fit() ?
+    def _check_fit_params(self):
+        # gracefully handle hyperparameter specifications if they are provided to fit instead
+        if any(isinstance(v, space.Space) for v in self.params.values()):
+            raise ValueError(
+                "Hyperparameter spaces provided to `fit`. Please provide concrete values "
+                "as hyperparameters when initializing or use `hyperparameter_tune` instead."
+            )
+
     @property
     def allowed_hyperparameters(self) -> List[str]:
         """List of hyperparameters allowed by the model."""
@@ -425,33 +607,6 @@ class AbstractTimeSeriesModel(AbstractModel):
         else:
             return None
 
-    def _fit(  # type: ignore
-        self,
-        train_data: TimeSeriesDataFrame,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[float] = None,
-        num_cpus: Optional[int] = None,
-        num_gpus: Optional[int] = None,
-        verbosity: int = 2,
-        **kwargs,
-    ) -> None:
-        """Private method for `fit`. See `fit` for documentation of arguments. Apart from
-        the model training logic, `fit` additionally implements other logic such as keeping
-        track of the time limit, etc.
-        """
-        # TODO: will not be extracted to new AbstractModel
-
-        # TODO: Make the models respect `num_cpus` and `num_gpus` parameters
-        raise NotImplementedError
-
-    def _check_fit_params(self):
-        # gracefully handle hyperparameter specifications if they are provided to fit instead
-        if any(isinstance(v, space.Space) for v in self.params.values()):
-            raise ValueError(
-                "Hyperparameter spaces provided to `fit`. Please provide concrete values "
-                "as hyperparameters when initializing or use `hyperparameter_tune` instead."
-            )
-
     def predict(  # type: ignore
         self,
         data: Union[TimeSeriesDataFrame, Dict[str, Optional[TimeSeriesDataFrame]]],
@@ -480,8 +635,6 @@ class AbstractTimeSeriesModel(AbstractModel):
             data is given as a separate forecast item in the dictionary, keyed by the `item_id`s
             of input items.
         """
-        # TODO: align method signature in new AbstractModel as predict(*args, **kwargs)
-
         # TODO: the method signature is not aligned with the model interface in general as it allows dict
         assert isinstance(data, TimeSeriesDataFrame)
 
@@ -617,56 +770,40 @@ class AbstractTimeSeriesModel(AbstractModel):
     def _is_gpu_available(self) -> bool:
         return False
 
+    @staticmethod
+    def _get_system_resources() -> Dict[str, Any]:
+        resource_manager = get_resource_manager()
+        system_num_cpus = resource_manager.get_cpu_count()
+        system_num_gpus = resource_manager.get_gpu_count()
+        return {
+            "num_cpus": system_num_cpus,
+            "num_gpus": system_num_gpus,
+        }
+
     def hyperparameter_tune(
         self,
+        train_data: TimeSeriesDataFrame,
+        val_data: Optional[TimeSeriesDataFrame],
+        val_splitter: Any = None,
+        default_num_trials: Optional[int] = 1,
+        refit_every_n_windows: Optional[int] = 1,
         hyperparameter_tune_kwargs: Union[str, dict] = "auto",
-        hpo_executor: Optional[HpoExecutor] = None,
         time_limit: Optional[float] = None,
-        **kwargs,
-    ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
-        if hpo_executor is None:
-            hpo_executor = self._get_default_hpo_executor()
-            default_num_trials = kwargs.pop("default_num_trials", None)
-            hpo_executor.initialize(
-                hyperparameter_tune_kwargs, default_num_trials=default_num_trials, time_limit=time_limit
-            )
-
-        kwargs = self.initialize(time_limit=time_limit, **kwargs)
-
-        kwargs = self._preprocess_fit_resources(
-            parallel_hpo=hpo_executor.executor_type == "ray", silent=True, **kwargs
+    ) -> Tuple[Dict[str, Any], Any]:
+        hpo_executor = self._get_default_hpo_executor()
+        hpo_executor.initialize(
+            hyperparameter_tune_kwargs, default_num_trials=default_num_trials, time_limit=time_limit
         )
-        self.validate_fit_resources(**kwargs)
 
-        # autogluon.core runs a complicated logic to determine the final number of gpus
-        # used in trials, which results in unintended setting of num_gpus=0. We override this
-        # logic here, and set to minimum num_gpus to 1 if it is set to 0 when GPUs are available
-        kwargs["num_gpus"] = 0 if not self._is_gpu_available() else max(kwargs.get("num_gpus", 1), 1)
+        self.initialize()
 
         # we use k_fold=1 to circumvent autogluon.core logic to manage resources during parallelization
         # of different folds
-        hpo_executor.register_resources(self, k_fold=1, **kwargs)
+        # FIXME: we pass in self which currently does not inherit from AbstractModel
+        hpo_executor.register_resources(self, k_fold=1, **self._get_system_resources())  # type: ignore
 
-        # TODO: Clean up call to _hyperparameter_tune
-        return self._hyperparameter_tune(hpo_executor=hpo_executor, **kwargs)  # type: ignore
-
-    def persist(self) -> Self:
-        """Ask the model to persist its assets in memory, i.e., to predict with low latency. In practice
-        this is used for pretrained models that have to lazy-load model parameters to device memory at
-        prediction time.
-        """
-        return self
-
-    def _hyperparameter_tune(  # type: ignore
-        self,
-        train_data: TimeSeriesDataFrame,
-        val_data: TimeSeriesDataFrame,
-        hpo_executor: HpoExecutor,
-        **kwargs,
-    ):
-        # TODO: do not extract to new AbstractModel
         time_start = time.time()
-        logger.debug(f"\tStarting AbstractTimeSeriesModel hyperparameter tuning for {self.name}")
+        logger.debug(f"\tStarting hyperparameter tuning for {self.name}")
         search_space = self._get_search_space()
 
         try:
@@ -674,35 +811,22 @@ class AbstractTimeSeriesModel(AbstractModel):
         except EmptySearchSpace:
             return skip_hpo(self, train_data, val_data, time_limit=hpo_executor.time_limit)
 
-        self.set_contexts(os.path.abspath(self.path))
-        directory = self.path
-        dataset_train_filename = "dataset_train.pkl"
-        train_path = os.path.join(self.path, dataset_train_filename)
-        save_pkl.save(path=train_path, object=train_data)
-
-        dataset_val_filename = "dataset_val.pkl"
-        val_path = os.path.join(self.path, dataset_val_filename)
-        save_pkl.save(path=val_path, object=val_data)
+        train_path, val_path = self._save_with_data(train_data, val_data)
 
-        fit_kwargs = dict(
-            val_splitter=kwargs.get("val_splitter"),
-            refit_every_n_windows=kwargs.get("refit_every_n_windows", 1),
-        )
         train_fn_kwargs = self._get_hpo_train_fn_kwargs(
             model_cls=self.__class__,
             init_params=self.get_params(),
             time_start=time_start,
             time_limit=hpo_executor.time_limit,
-            fit_kwargs=fit_kwargs,
+            fit_kwargs=dict(
+                val_splitter=val_splitter,
+                refit_every_n_windows=refit_every_n_windows,
+            ),
             train_path=train_path,
             val_path=val_path,
             hpo_executor=hpo_executor,
         )
 
-        model_estimate_memory_usage = None
-        if self.estimate_memory_usage is not None:
-            model_estimate_memory_usage = self.estimate_memory_usage(**kwargs)
-
         minimum_resources = self.get_minimum_resources(is_gpu_available=self._is_gpu_available())
         hpo_context = disable_stdout if isinstance(hpo_executor, RayHpoExecutor) else nullcontext
 
@@ -718,10 +842,10 @@ class AbstractTimeSeriesModel(AbstractModel):
             hpo_executor.execute(
                 model_trial=model_trial,
                 train_fn_kwargs=train_fn_kwargs,
-                directory=directory,
+                directory=self.path,
                 minimum_cpu_per_trial=minimum_cpu_per_trial,
                 minimum_gpu_per_trial=minimum_resources.get("num_gpus", 0),
-                model_estimate_memory_usage=model_estimate_memory_usage,  # type: ignore
+                model_estimate_memory_usage=None,
                 adapter_type="timeseries",
             )
 
@@ -734,6 +858,53 @@ class AbstractTimeSeriesModel(AbstractModel):
 
         return hpo_models, analysis
 
+    @property
+    def is_ensemble(self) -> bool:
+        """Return True if the model is an ensemble model or a container of multiple models."""
+        return self._get_model_base() is self
+
+    def _get_default_hpo_executor(self) -> HpoExecutor:
+        backend = (
+            self._get_model_base()._get_hpo_backend()
+        )  # If ensemble, will use the base model to determine backend
+        if backend == RAY_BACKEND:
+            try:
+                try_import_ray()
+            except Exception as e:
+                warning_msg = f"Will use custom hpo logic because ray import failed. Reason: {str(e)}"
+                dup_filter.attach_filter_targets(warning_msg)
+                logger.warning(warning_msg)
+                backend = CUSTOM_BACKEND
+        hpo_executor = HpoExecutorFactory.get_hpo_executor(backend)()  # type: ignore
+        return hpo_executor
+
+    def _get_model_base(self) -> AbstractTimeSeriesModel:
+        return self
+
+    def _get_hpo_backend(self) -> str:
+        """Choose which backend("ray" or "custom") to use for hpo"""
+        if DistributedContext.is_distributed_mode():
+            return RAY_BACKEND
+        return CUSTOM_BACKEND
+
+    def _get_search_space(self):
+        """Sets up default search space for HPO. Each hyperparameter which user did not specify is converted from
+        default fixed value to default search space.
+        """
+        params = self.params.copy()
+        return params
+
+    def _save_with_data(self, train_data, val_data):
+        self.set_contexts(os.path.abspath(self.path))
+        dataset_train_filename = "dataset_train.pkl"
+        train_path = os.path.join(self.path, dataset_train_filename)
+        save_pkl.save(path=train_path, object=train_data)
+
+        dataset_val_filename = "dataset_val.pkl"
+        val_path = os.path.join(self.path, dataset_val_filename)
+        save_pkl.save(path=val_path, object=val_data)
+        return train_path, val_path
+
     def preprocess(  # type: ignore
         self,
         data: TimeSeriesDataFrame,
@@ -745,9 +916,12 @@ class AbstractTimeSeriesModel(AbstractModel):
         # TODO: move to new AbstractModel
         return data, known_covariates
 
-    def get_memory_size(self, allow_exception: bool = False) -> Optional[int]:
-        # TODO: move to new AbstractModel
-        return None
+    def persist(self) -> Self:
+        """Ask the model to persist its assets in memory, i.e., to predict with low latency. In practice
+        this is used for pretrained models that have to lazy-load model parameters to device memory at
+        prediction time.
+        """
+        return self
 
     def convert_to_refit_full_via_copy(self) -> Self:
         # save the model as a new model on disk
@@ -764,6 +938,33 @@ class AbstractTimeSeriesModel(AbstractModel):
 
         return refit_model
 
+    def convert_to_refit_full_template(self):
+        """
+        After calling this function, returned model should be able to be fit without X_val, y_val using the iterations trained by the original model.
+
+        Increase max_memory_usage_ratio by 25% to reduce the chance that the refit model will trigger NotEnoughMemoryError and skip training.
+        This can happen without the 25% increase since the refit model generally will use more training data and thus require more memory.
+        """
+        params = copy.deepcopy(self.get_params())
+
+        if "hyperparameters" not in params:
+            params["hyperparameters"] = dict()
+
+        if AG_ARGS_FIT not in params["hyperparameters"]:
+            params["hyperparameters"][AG_ARGS_FIT] = dict()
+
+        # TODO: remove
+        # Increase memory limit by 25% to avoid memory restrictions during fit
+        params["hyperparameters"][AG_ARGS_FIT]["max_memory_usage_ratio"] = (
+            params["hyperparameters"][AG_ARGS_FIT].get("max_memory_usage_ratio", 1.0) * 1.25
+        )
+
+        params["hyperparameters"].update(self.params_trained)
+        params["name"] = params["name"] + REFIT_FULL_SUFFIX
+        template = self.__class__(**params)
+
+        return template
+
     def get_user_params(self) -> dict:
         """Used to access user-specified parameters for the model before initialization."""
         if self._user_params is None:

{autogluon.timeseries-1.2.1b20250219 → autogluon.timeseries-1.2.1b20250221}/src/autogluon/timeseries/version.py

@@ -1,4 +1,4 @@
 """This is the autogluon version file."""
 
-__version__ = "1.2.1b20250219"
+__version__ = "1.2.1b20250221"
 __lite__ = False

{autogluon.timeseries-1.2.1b20250219 → autogluon.timeseries-1.2.1b20250221}/src/autogluon.timeseries.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: autogluon.timeseries
-Version: 1.2.1b20250219
+Version: 1.2.1b20250221
 Summary: Fast and Accurate ML in 3 Lines of Code
 Home-page: https://github.com/autogluon/autogluon
 Author: AutoGluon Community

{autogluon.timeseries-1.2.1b20250219 → autogluon.timeseries-1.2.1b20250221}/src/autogluon.timeseries.egg-info/requires.txt

@@ -17,9 +17,9 @@ fugue>=0.9.0
 tqdm<5,>=4.38
 orjson~=3.9
 tensorboard<3,>=2.9
-autogluon.core[raytune]==1.2.1b20250219
-autogluon.common==1.2.1b20250219
-autogluon.tabular[catboost,lightgbm,xgboost]==1.2.1b20250219
+autogluon.core[raytune]==1.2.1b20250221
+autogluon.common==1.2.1b20250221
+autogluon.tabular[catboost,lightgbm,xgboost]==1.2.1b20250221
 
 [all]