Factor out time series-related functionality into a time series Task object

.flake8

@@ -1,5 +1,5 @@
 [flake8]
-ignore = E203, E266, E501, W503, F403, F401, C901
+ignore = E203, E266, E402, E501, W503, F403, F401, C901
 max-line-length = 127
 max-complexity = 10
 select = B,C,E,F,W,T4,B9

flaml/automl/automl.py

@@ -15,10 +15,9 @@
 import json
 
 from flaml.automl.state import SearchState, AutoMLState
-from flaml.automl.ml import (
- train_estimator,
- get_estimator_class,
-)
+from flaml.automl.ml import train_estimator
+
+from flaml.automl.time_series import TimeSeriesDataset
 from flaml.config import (
  MIN_SAMPLE_TRAIN,
  MEM_THRES,
@@ -31,7 +30,7 @@
 )
 
 # TODO check to see when we can remove these
-from flaml.automl.task.task import CLASSIFICATION, TS_FORECAST, Task
+from flaml.automl.task.task import CLASSIFICATION, Task
 from flaml.automl.task.factory import task_factory
 from flaml import tune
 from flaml.automl.logger import logger, logger_formatter
@@ -914,7 +913,7 @@ def _decide_eval_method(self, eval_method, time_budget):
  ], "eval_method must be 'auto' or 'cv' for custom data splitter."
  assert self._state.X_val is None, "custom splitter and custom validation data can't be used together."
  return "cv"
- if self._state.X_val is not None:
+ if self._state.X_val is not None and not isinstance(self._state.X_val, TimeSeriesDataset):
  assert eval_method in [
  "auto",
  "holdout",
@@ -1159,7 +1158,7 @@ def _prepare_data(self, eval_method, split_ratio, n_splits):
  self._df,
  self._sample_weight_full,
  )
- self.data_size_full = len(self._state.y_train_all)
+ self.data_size_full = self._state.data_size_full
 
  def fit(
  self,
@@ -1211,6 +1210,7 @@ def fit(
  free_mem_ratio=0,
  metric_constraints=None,
  custom_hp=None,
+ time_col=None,
  cv_score_agg_func=None,
  skip_transform=None,
  fit_kwargs_by_estimator=None,
@@ -1531,6 +1531,7 @@ def cv_score_agg_func(val_loss_folds, log_metrics_folds):
  if isinstance(task, str):
  task = task_factory(task, X_train, y_train)
  self._state.task = task
+ self._state.task.time_col = time_col
  self._estimator_type = "classifier" if task.is_classification() else "regressor"
  time_budget = time_budget or self._settings.get("time_budget")
  n_jobs = n_jobs or self._settings.get("n_jobs")
@@ -1832,7 +1833,7 @@ def is_to_reverse_metric(metric, task):
  if estimator_name not in self._state.learner_classes:
  self.add_learner(
  estimator_name,
- get_estimator_class(self._state.task, estimator_name),
+ self._state.task.estimator_class_from_str(estimator_name),
  )
  # set up learner search space
  if isinstance(starting_points, str) and starting_points.startswith("data"):
@@ -1888,6 +1889,7 @@ def is_to_reverse_metric(metric, task):
  self._search_states[estimator_name] = SearchState(
  learner_class=estimator_class,
  data_size=self._state.data_size,
+ data=self._state.X_val,
  task=self._state.task,
  starting_point=starting_points.get(estimator_name),
  period=self._state.fit_kwargs.get(
@@ -2597,7 +2599,7 @@ def _search(self):
  if self._max_iter > 1:
  self._state.time_budget = -1
  if (
- self._state.task in TS_FORECAST
+ self._state.task.is_ts_forecast()
  or self._trained_estimator is None
  or self._trained_estimator.model is None
  or (