Fix IndexError in time series models when log_training_metric=True and improve test reliability

flaml/automl/time_series/tcn.py

@@ -264,7 +264,8 @@ def fit(self, X_train: TimeSeriesDataset, y_train=None, budget=None, **kwargs):
     def predict(self, X):
         X = self.enrich(X)
         if isinstance(X, TimeSeriesDataset):
-            df = X.X_val
+            # Use X_train if X_val is empty (e.g., when computing training metrics)
+            df = X.X_val if len(X.test_data) > 0 else X.X_train
         else:
             df = X
         dataset = DataframeDataset(

flaml/automl/time_series/tft.py

@@ -170,7 +170,11 @@ def predict(self, X):
         last_data_cols = self.group_ids.copy()
         last_data_cols.append(self.target_names[0])
         last_data = self.data[lambda x: x.time_idx == x.time_idx.max()][last_data_cols]
-        decoder_data = X.X_val if isinstance(X, TimeSeriesDataset) else X
+        # Use X_train if test_data is empty (e.g., when computing training metrics)
+        if isinstance(X, TimeSeriesDataset):
+            decoder_data = X.X_val if len(X.test_data) > 0 else X.X_train
+        else:
+            decoder_data = X
         if "time_idx" not in decoder_data:
             decoder_data = add_time_idx_col(decoder_data)
         decoder_data["time_idx"] += encoder_data["time_idx"].max() + 1 - decoder_data["time_idx"].min()

flaml/automl/time_series/ts_model.py

@@ -194,7 +194,11 @@ def predict(self, X: Union[TimeSeriesDataset, DataFrame], **kwargs):
 
         elif isinstance(X, TimeSeriesDataset):
             data = X
-            X = data.test_data[[self.time_col] + X.regressors]
+            # Use train_data if test_data is empty (e.g., when computing training metrics)
+            if len(data.test_data) == 0:
+                X = data.train_data[[self.time_col] + X.regressors]
+            else:
+                X = data.test_data[[self.time_col] + X.regressors]
 
         if self._model is not None:
             forecast = self._model.predict(X, **kwargs)
@@ -301,7 +305,11 @@ def predict(self, X, **kwargs):
 
         if isinstance(X, TimeSeriesDataset):
             data = X
-            X = data.test_data[data.regressors + [data.time_col]]
+            # Use train_data if test_data is empty (e.g., when computing training metrics)
+            if len(data.test_data) == 0:
+                X = data.train_data[data.regressors + [data.time_col]]
+            else:
+                X = data.test_data[data.regressors + [data.time_col]]
 
         X = X.rename(columns={self.time_col: "ds"})
         if self._model is not None:
@@ -327,7 +335,11 @@ def predict(self, X, **kwargs) -> pd.Series:
 
         if isinstance(X, TimeSeriesDataset):
             data = X
-            X = data.test_data[data.regressors + [data.time_col]]
+            # Use train_data if test_data is empty (e.g., when computing training metrics)
+            if len(data.test_data) == 0:
+                X = data.train_data[data.regressors + [data.time_col]]
+            else:
+                X = data.test_data[data.regressors + [data.time_col]]
         else:
             X = X[self.regressors + [self.time_col]]
 
@@ -828,7 +840,11 @@ def predict(self, X, **kwargs):
         X = self.enrich(X)
         if isinstance(X, TimeSeriesDataset):
             data = X
-            X = data.test_data
+            # Use train_data if test_data is empty (e.g., when computing training metrics)
+            if len(data.test_data) == 0:
+                X = data.train_data
+            else:
+                X = data.test_data
 
         if self._model is not None:
             X = X[self.regressors]

setup.py

@@ -46,9 +46,10 @@
             "jupyter",
         ],
         "spark": [
-            "pyspark>=3.2.0",
+            "pyspark>=3.2.0,<3.5.0; python_version<'3.11'",
+            "pyspark>=3.5.0; python_version>='3.11'",
             "joblibspark>=0.5.0",
-            "joblib<=1.3.2",
+            "joblib>=1.2.0,<=1.4.2",
         ],
         "test": [
             "jupyter",
@@ -82,7 +83,7 @@
             # "pytorch-forecasting==0.10.1; python_version=='3.11'",
             "mlflow==2.15.1",
             "joblibspark>=0.5.0",
-            "joblib<=1.3.2",
+            "joblib>=1.2.0,<=1.4.2",
             "nbconvert",
             "nbformat",
             "ipykernel",
@@ -161,7 +162,8 @@
         "synapse": [
             "joblibspark>=0.5.0",
             "optuna>=2.8.0,<=3.6.1",
-            "pyspark>=3.2.0",
+            "pyspark>=3.2.0,<3.5.0; python_version<'3.11'",
+            "pyspark>=3.5.0; python_version>='3.11'",
         ],
         "autozero": ["scikit-learn", "pandas", "packaging"],
     },

test/default/test_defaults.py

@@ -1,7 +1,10 @@
 import pickle
 import sys
+import time
+import urllib.error
 
 import pandas as pd
+import pytest
 from sklearn.datasets import fetch_california_housing, load_breast_cancer, load_iris
 from sklearn.model_selection import train_test_split
 
@@ -15,6 +18,32 @@
 )
 
 
+def fetch_california_housing_with_retry(max_retries=3, retry_delay=2, **kwargs):
+    """Fetch California Housing dataset with retry logic for HTTP errors.
+
+    Args:
+        max_retries: Maximum number of retry attempts (default: 3)
+        retry_delay: Initial delay between retries in seconds (default: 2)
+        **kwargs: Arguments to pass to fetch_california_housing
+
+    Returns:
+        Dataset from fetch_california_housing
+
+    Raises:
+        Exception: If all retry attempts fail
+    """
+    for attempt in range(max_retries):
+        try:
+            return fetch_california_housing(**kwargs)
+        except (urllib.error.HTTPError, urllib.error.URLError, Exception) as e:
+            if attempt == max_retries - 1:
+                # Last attempt failed, skip the test
+                pytest.skip(f"Failed to fetch California Housing dataset after {max_retries} attempts: {e}")
+            # Exponential backoff
+            wait_time = retry_delay * (2 ** attempt)
+            time.sleep(wait_time)
+
+
 def test_greedy_feedback(path="test/default", strategy="greedy-feedback"):
     # sys.argv = f"portfolio.py --output {path} --input {path} --metafeatures {path}/all/metafeatures.csv --task binary --estimator lgbm xgboost xgb_limitdepth rf extra_tree --strategy {strategy}".split()
     # portfolio.main()
@@ -50,6 +79,8 @@ def test_iris(as_frame=True):
 
 
 def test_housing(as_frame=True):
+    X_train, y_train = fetch_california_housing_with_retry(return_X_y=True, as_frame=as_frame)
+
     automl = AutoML()
     automl_settings = {
         "time_budget": 2,
@@ -60,7 +91,6 @@ def test_housing(as_frame=True):
         "starting_points": "data",
         "max_iter": 0,
     }
-    X_train, y_train = fetch_california_housing(return_X_y=True, as_frame=as_frame)
     automl.fit(X_train, y_train, **automl_settings)
 
 
@@ -115,7 +145,8 @@ def test_suggest_classification():
 
 def test_suggest_regression():
     location = "test/default"
-    X_train, y_train = fetch_california_housing(return_X_y=True, as_frame=True)
+    X_train, y_train = fetch_california_housing_with_retry(return_X_y=True, as_frame=True)
+
     suggested = suggest_hyperparams("regression", X_train, y_train, "lgbm", location=location)
     print(suggested)
     suggested = preprocess_and_suggest_hyperparams("regression", X_train, y_train, "xgboost", location=location)
@@ -137,7 +168,8 @@ def test_rf():
     print(rf)
 
     location = "test/default"
-    X_train, y_train = fetch_california_housing(return_X_y=True, as_frame=True)
+    X_train, y_train = fetch_california_housing_with_retry(return_X_y=True, as_frame=True)
+
     rf = RandomForestRegressor(default_location=location)
     rf.fit(X_train[:100], y_train[:100])
     rf.predict(X_train)
@@ -155,7 +187,8 @@ def test_extratrees():
     print(classifier)
 
     location = "test/default"
-    X_train, y_train = fetch_california_housing(return_X_y=True, as_frame=True)
+    X_train, y_train = fetch_california_housing_with_retry(return_X_y=True, as_frame=True)
+
     regressor = ExtraTreesRegressor(default_location=location)
     regressor.fit(X_train[:100], y_train[:100])
     regressor.predict(X_train)
@@ -175,7 +208,8 @@ def test_lgbm():
     print(classifier.classes_)
 
     location = "test/default"
-    X_train, y_train = fetch_california_housing(return_X_y=True, as_frame=True)
+    X_train, y_train = fetch_california_housing_with_retry(return_X_y=True, as_frame=True)
+
     regressor = LGBMRegressor(default_location=location)
     regressor.fit(X_train, y_train)
     regressor.predict(X_train)
@@ -194,7 +228,8 @@ def test_xgboost():
     print(classifier.classes_)
 
     location = "test/default"
-    X_train, y_train = fetch_california_housing(return_X_y=True, as_frame=True)
+    X_train, y_train = fetch_california_housing_with_retry(return_X_y=True, as_frame=True)
+
     regressor = XGBRegressor(default_location=location)
     regressor.fit(X_train[:100], y_train[:100])
     regressor.predict(X_train)