Fix log_training_metric causing IndexError for time series models (#1469)

Co-authored-by: Li Jiang <lijiang1@microsoft.com>

.coveragerc

@@ -1,5 +1,7 @@
 [run]
 branch = True
-source = flaml
+source =
+  flaml
 omit =
-  *test*
+  */test/*
+  */flaml/autogen/*

.github/workflows/CD.yml

@@ -13,7 +13,7 @@ jobs:
     strategy:
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.10"]
+        python-version: ["3.12"]
     runs-on: ${{ matrix.os }}
     environment: package
     steps:

.github/workflows/deploy-website.yml

@@ -37,11 +37,11 @@ jobs:
       - name: setup python
         uses: actions/setup-python@v4
         with:
-          python-version: "3.10"
+          python-version: "3.12"
       - name: pydoc-markdown install
         run: |
           python -m pip install --upgrade pip
-          pip install pydoc-markdown==4.7.0
+          pip install pydoc-markdown==4.7.0 setuptools
       - name: pydoc-markdown run
         run: |
           pydoc-markdown
@@ -73,11 +73,11 @@ jobs:
       - name: setup python
         uses: actions/setup-python@v4
         with:
-          python-version: "3.10"
+          python-version: "3.12"
       - name: pydoc-markdown install
         run: |
           python -m pip install --upgrade pip
-          pip install pydoc-markdown==4.7.0
+          pip install pydoc-markdown==4.7.0 setuptools
       - name: pydoc-markdown run
         run: |
           pydoc-markdown

.github/workflows/openai.yml

@@ -4,14 +4,15 @@
 name: OpenAI
 
 on:
-  pull_request:
-    branches: ['main']
-    paths:
-      - 'flaml/autogen/**'
-      - 'test/autogen/**'
-      - 'notebook/autogen_openai_completion.ipynb'
-      - 'notebook/autogen_chatgpt_gpt4.ipynb'
-      - '.github/workflows/openai.yml'
+  workflow_dispatch:
+#   pull_request:
+#     branches: ['main']
+#     paths:
+#       - 'flaml/autogen/**'
+#       - 'test/autogen/**'
+#       - 'notebook/autogen_openai_completion.ipynb'
+#       - 'notebook/autogen_chatgpt_gpt4.ipynb'
+#       - '.github/workflows/openai.yml'
 
 permissions: {}
 

.github/workflows/python-package.yml

@@ -22,8 +22,12 @@ on:
       - 'setup.py'
   merge_group:
     types: [checks_requested]
+  schedule:
+    # Every other day at 02:00 UTC
+    - cron: '0 2 */2 * *'
 
-permissions: {}
+permissions:
+  contents: write
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref }}
   cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
@@ -36,7 +40,12 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.10", "3.11", "3.12"]
+        exclude:
+          - os: macos-latest
+            python-version: "3.10"  # macOS runners will hang on python 3.10 for unknown reasons
+          - os: macos-latest
+            python-version: "3.12"  # macOS runners will hang on python 3.12 for unknown reasons
     steps:
       - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
@@ -44,7 +53,7 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - name: On mac, install libomp to facilitate lgbm and xgboost install
-        if: matrix.os == 'macOS-latest'
+        if: matrix.os == 'macos-latest'
         run: |
           brew update
           brew install libomp
@@ -60,72 +69,70 @@ jobs:
           pip install -e .
           python -c "import flaml"
           pip install -e .[test]
-      - name: On Ubuntu python 3.10, install pyspark 3.4.1
-        if: matrix.python-version == '3.10' && matrix.os == 'ubuntu-latest'
-        run: |
-          pip install pyspark==3.4.1
-          pip list | grep "pyspark"
       - name: On Ubuntu python 3.11, install pyspark 3.5.1
         if: matrix.python-version == '3.11' && matrix.os == 'ubuntu-latest'
         run: |
           pip install pyspark==3.5.1
           pip list | grep "pyspark"
-      - name: If linux and python<3.11, install ray 2
-        if: matrix.os == 'ubuntu-latest' && matrix.python-version != '3.11'
+      - name: On Ubuntu python 3.12, install pyspark 4.0.1
+        if: matrix.python-version == '3.12' && matrix.os == 'ubuntu-latest'
         run: |
-          pip install "ray[tune]<2.5.0"
-      - name: If mac and python 3.10, install ray and xgboost 1
-        if: matrix.os == 'macOS-latest' && matrix.python-version == '3.10'
-        run: |
-          pip install -e .[ray]
-          # use macOS to test xgboost 1, but macOS also supports xgboost 2
-          pip install "xgboost<2"
-      - name: If linux, install prophet on python < 3.9
-        if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.8'
+          pip install pyspark==4.0.1
+          pip list | grep "pyspark"
+      # # TODO: support ray
+      # - name: If linux and python<3.11, install ray 2
+      #   if: matrix.os == 'ubuntu-latest' && matrix.python-version < '3.11'
+      #   run: |
+      #     pip install "ray[tune]<2.5.0"
+      - name: Install prophet when on linux
+        if: matrix.os == 'ubuntu-latest'
         run: |
           pip install -e .[forecast]
-      - name: Install vw on python < 3.10
-        if: matrix.python-version == '3.8' || matrix.python-version == '3.9'
+      # TODO: support vw for python 3.10+
+      - name:  If linux and python<3.10, install vw
+        if: matrix.os == 'ubuntu-latest' && matrix.python-version < '3.10'
         run: |
           pip install -e .[vw]
-      - name: Test with pytest
-        if: matrix.python-version != '3.10'
+      - name: Pip freeze
         run: |
-          pytest test/ --ignore=test/autogen
+          pip freeze
+      - name: Check dependencies
+        run: |
+          python test/check_dependency.py
+      - name: Clear pip cache
+        run: |
+          pip cache purge
+      - name: Test with pytest
+        if: matrix.python-version != '3.11'
+        run: |
+          pytest test/ --ignore=test/autogen --reruns 2 --reruns-delay 10
       - name: Coverage
-        if: matrix.python-version == '3.10'
+        if: matrix.python-version == '3.11'
         run: |
           pip install coverage
-          coverage run -a -m pytest test --ignore=test/autogen
+          coverage run -a -m pytest test --ignore=test/autogen --reruns 2 --reruns-delay 10
           coverage xml
       - name: Upload coverage to Codecov
-        if: matrix.python-version == '3.10'
+        if: matrix.python-version == '3.11'
         uses: codecov/codecov-action@v3
         with:
           file: ./coverage.xml
           flags: unittests
+      - name: Save dependencies
+        shell: bash
+        run: |
+          git config --global user.name 'github-actions[bot]'
+          git config --global user.email 'github-actions[bot]@users.noreply.github.com'
+          git config advice.addIgnoredFile false
 
-  # docs:
+          BRANCH=unit-tests-installed-dependencies
+          git fetch origin
+          git checkout -B "$BRANCH" "origin/$BRANCH"
 
-  #   runs-on: ubuntu-latest
-
-  #   steps:
-  #     - uses: actions/checkout@v3
-  #     - name: Setup Python
-  #       uses: actions/setup-python@v4
-  #       with:
-  #         python-version: '3.8'
-  #     - name: Compile documentation
-  #       run: |
-  #           pip install -e .
-  #           python -m pip install sphinx sphinx_rtd_theme
-  #           cd docs
-  #           make html
-  #     - name: Deploy to GitHub pages
-  #       if: ${{ github.ref == 'refs/heads/main' }}
-  #       uses: JamesIves/github-pages-deploy-action@3.6.2
-  #       with:
-  #         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-  #         BRANCH: gh-pages
-  #         FOLDER: docs/_build/html
-  #         CLEAN: true
+          pip freeze > installed_all_dependencies_${{ matrix.python-version }}_${{ matrix.os }}.txt
+          python test/check_dependency.py > installed_first_tier_dependencies_${{ matrix.python-version }}_${{ matrix.os }}.txt
+          git add installed_*dependencies*.txt
+          mv coverage.xml ./coverage_${{ matrix.python-version }}_${{ matrix.os }}.xml || true
+          git add -f ./coverage_${{ matrix.python-version }}_${{ matrix.os }}.xml || true
+          git commit -m "Update installed dependencies for Python ${{ matrix.python-version }} on ${{ matrix.os }}" || exit 0
+          git push origin "$BRANCH" --force

.gitignore

@@ -60,6 +60,7 @@ coverage.xml
 .hypothesis/
 .pytest_cache/
 cover/
+junit
 
 # Translations
 *.mo
@@ -172,7 +173,7 @@ test/default
 test/housing.json
 test/nlp/default/transformer_ms/seq-classification.json
 
-flaml/fabric/fanova/_fanova.c
+flaml/fabric/fanova/*fanova.c
 # local config files
 *.config.local
 
@@ -184,3 +185,7 @@ notebook/lightning_logs/
 lightning_logs/
 flaml/autogen/extensions/tmp/
 test/autogen/my_tmp/
+catboost_*
+
+# Internal configs
+.pypirc

README.md

@@ -14,15 +14,9 @@
     <br>
 </p>
 
-:fire: FLAML supports AutoML and Hyperparameter Tuning in [Microsoft Fabric Data Science](https://learn.microsoft.com/en-us/fabric/data-science/automated-machine-learning-fabric). In addition, we've introduced Python 3.11 support, along with a range of new estimators, and comprehensive integration with MLflow—thanks to contributions from the Microsoft Fabric product team.
+:fire: FLAML supports AutoML and Hyperparameter Tuning in [Microsoft Fabric Data Science](https://learn.microsoft.com/en-us/fabric/data-science/automated-machine-learning-fabric). In addition, we've introduced Python 3.11 and 3.12 support, along with a range of new estimators, and comprehensive integration with MLflow—thanks to contributions from the Microsoft Fabric product team.
 
-:fire: Heads-up: We have migrated [AutoGen](https://microsoft.github.io/autogen/) into a dedicated [github repository](https://github.com/microsoft/autogen). Alongside this move, we have also launched a dedicated [Discord](https://discord.gg/pAbnFJrkgZ) server and a [website](https://microsoft.github.io/autogen/) for comprehensive documentation.
-
-:fire: The automated multi-agent chat framework in [AutoGen](https://microsoft.github.io/autogen/) is in preview from v2.0.0.
-
-:fire: FLAML is highlighted in OpenAI's [cookbook](https://github.com/openai/openai-cookbook#related-resources-from-around-the-web).
-
-:fire: [autogen](https://microsoft.github.io/autogen/) is released with support for ChatGPT and GPT-4, based on [Cost-Effective Hyperparameter Optimization for Large Language Model Generation Inference](https://arxiv.org/abs/2303.04673).
+:fire: Heads-up: [AutoGen](https://microsoft.github.io/autogen/) has moved to a dedicated [GitHub repository](https://github.com/microsoft/autogen). FLAML no longer includes the `autogen` module—please use AutoGen directly.
 
 ## What is FLAML
 
@@ -30,7 +24,7 @@ FLAML is a lightweight Python library for efficient automation of machine
 learning and AI operations. It automates workflow based on large language models, machine learning models, etc.
 and optimizes their performance.
 
-- FLAML enables building next-gen GPT-X applications based on multi-agent conversations with minimal effort. It simplifies the orchestration, automation and optimization of a complex GPT-X workflow. It maximizes the performance of GPT-X models and augments their weakness.
+- FLAML enables economical automation and tuning for ML/AI workflows, including model selection and hyperparameter optimization under resource constraints.
 - For common machine learning tasks like classification and regression, it quickly finds quality models for user-provided data with low computational resources. It is easy to customize or extend. Users can find their desired customizability from a smooth range.
 - It supports fast and economical automatic tuning (e.g., inference hyperparameters for foundation models, configurations in MLOps/LMOps workflows, pipelines, mathematical/statistical models, algorithms, computing experiments, software configurations), capable of handling large search space with heterogeneous evaluation cost and complex constraints/guidance/early stopping.
 
@@ -46,10 +40,10 @@ FLAML requires **Python version >= 3.9**. It can be installed from pip:
 pip install flaml
 ```
 
-Minimal dependencies are installed without extra options. You can install extra options based on the feature you need. For example, use the following to install the dependencies needed by the [`autogen`](https://microsoft.github.io/autogen/) package.
+Minimal dependencies are installed without extra options. You can install extra options based on the feature you need. For example, use the following to install the dependencies needed by the [`automl`](https://microsoft.github.io/FLAML/docs/Use-Cases/Task-Oriented-AutoML) module.
 
 ```bash
-pip install "flaml[autogen]"
+pip install "flaml[automl]"
 ```
 
 Find more options in [Installation](https://microsoft.github.io/FLAML/docs/Installation).
@@ -57,39 +51,6 @@ Each of the [`notebook examples`](https://github.com/microsoft/FLAML/tree/main/n
 
 ## Quickstart
 
-- (New) The [autogen](https://microsoft.github.io/autogen/) package enables the next-gen GPT-X applications with a generic multi-agent conversation framework.
-  It offers customizable and conversable agents which integrate LLMs, tools and human.
-  By automating chat among multiple capable agents, one can easily make them collectively perform tasks autonomously or with human feedback, including tasks that require using tools via code. For example,
-
-```python
-from flaml import autogen
-
-assistant = autogen.AssistantAgent("assistant")
-user_proxy = autogen.UserProxyAgent("user_proxy")
-user_proxy.initiate_chat(
-    assistant,
-    message="Show me the YTD gain of 10 largest technology companies as of today.",
-)
-# This initiates an automated chat between the two agents to solve the task
-```
-
-Autogen also helps maximize the utility out of the expensive LLMs such as ChatGPT and GPT-4. It offers a drop-in replacement of `openai.Completion` or `openai.ChatCompletion` with powerful functionalites like tuning, caching, templating, filtering. For example, you can optimize generations by LLM with your own tuning data, success metrics and budgets.
-
-```python
-# perform tuning
-config, analysis = autogen.Completion.tune(
-    data=tune_data,
-    metric="success",
-    mode="max",
-    eval_func=eval_func,
-    inference_budget=0.05,
-    optimization_budget=3,
-    num_samples=-1,
-)
-# perform inference for a test instance
-response = autogen.Completion.create(context=test_instance, **config)
-```
-
 - With three lines of code, you can start using this economical and fast
   AutoML engine as a [scikit-learn style estimator](https://microsoft.github.io/FLAML/docs/Use-Cases/Task-Oriented-AutoML).
 
@@ -111,7 +72,10 @@ automl.fit(X_train, y_train, task="classification", estimator_list=["lgbm"])
 
 ```python
 from flaml import tune
-tune.run(evaluation_function, config={…}, low_cost_partial_config={…}, time_budget_s=3600)
+
+tune.run(
+    evaluation_function, config={…}, low_cost_partial_config={…}, time_budget_s=3600
+)
 ```
 
 - [Zero-shot AutoML](https://microsoft.github.io/FLAML/docs/Use-Cases/Zero-Shot-AutoML) allows using the existing training API from lightgbm, xgboost etc. while getting the benefit of AutoML in choosing high-performance hyperparameter configurations per task.

flaml/autogen/__init__.py

@@ -1,3 +1,12 @@
+import warnings
+
 from .agentchat import *
 from .code_utils import DEFAULT_MODEL, FAST_MODEL
 from .oai import *
+
+warnings.warn(
+    "The `flaml.autogen` module is deprecated and will be removed in a future release. "
+    "Please refer to `https://github.com/microsoft/autogen` for latest usage.",
+    DeprecationWarning,
+    stacklevel=2,
+)

flaml/automl/automl.py

@@ -4,6 +4,7 @@
 #  * project root for license information.
 from __future__ import annotations
 
+import inspect
 import json
 import logging
 import os
@@ -177,10 +178,11 @@ class AutoML(BaseEstimator):
                 ['auto', 'cv', 'holdout'].
             split_ratio: A float of the valiation data percentage for holdout.
             n_splits: An integer of the number of folds for cross - validation.
-            log_type: A string of the log type, one of
-                ['better', 'all'].
-                'better' only logs configs with better loss than previos iters
-                'all' logs all the tried configs.
+            log_type: Specifies which logs to save. One of ['better', 'all']. Default is 'better'.
+                - 'better': Logs configs and models (if `model_history` is True) only when the loss improves,
+                  to `log_file_name` and MLflow, respectively.
+                - 'all': Logs all configs and models (if `model_history` is True), regardless of performance.
+                Note: Configs are always logged to MLflow if MLflow logging is enabled.
             model_history: A boolean of whether to keep the best
                 model per estimator. Make sure memory is large enough if setting to True. Default False.
             log_training_metric: A boolean of whether to log the training
@@ -401,6 +403,24 @@ class AutoML(BaseEstimator):
         self._estimator_type = "classifier" if settings["task"] in CLASSIFICATION else "regressor"
         self.best_run_id = None
 
+    def __getstate__(self):
+        """Customize pickling to avoid serializing runtime-only objects.
+
+        MLflow's sklearn flavor serializes estimators via (cloud)pickle. During
+        AutoML fitting we may attach an internal mlflow integration instance
+        which holds `concurrent.futures.Future` objects and executors containing
+        thread locks, which are not picklable.
+        """
+
+        state = self.__dict__.copy()
+        state.pop("mlflow_integration", None)
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        # Ensure attribute exists post-unpickle.
+        self.mlflow_integration = None
+
     def get_params(self, deep: bool = False) -> dict:
         return self._settings.copy()
 
@@ -2156,7 +2176,7 @@ class AutoML(BaseEstimator):
                 use_spark=True,
                 force_cancel=self._force_cancel,
                 mlflow_exp_name=self._mlflow_exp_name,
-                automl_info=(mlflow_log_latency,),  # pass automl info to tune.run
+                automl_info=(mlflow_log_latency, self._log_type),  # pass automl info to tune.run
                 extra_tag=self.autolog_extra_tag,
                 # raise_on_failed_trial=False,
                 # keep_checkpoints_num=1,
@@ -2219,7 +2239,9 @@ class AutoML(BaseEstimator):
                 if better or self._log_type == "all":
                     self._log_trial(search_state, estimator)
                 if self.mlflow_integration:
-                    self.mlflow_integration.record_state(self, search_state, estimator)
+                    self.mlflow_integration.record_state(
+                        self, search_state, estimator, better or self._log_type == "all"
+                    )
 
     def _log_trial(self, search_state, estimator):
         if self._training_log:
@@ -2461,7 +2483,9 @@ class AutoML(BaseEstimator):
                 if better or self._log_type == "all":
                     self._log_trial(search_state, estimator)
                 if self.mlflow_integration:
-                    self.mlflow_integration.record_state(self, search_state, estimator)
+                    self.mlflow_integration.record_state(
+                        self, search_state, estimator, better or self._log_type == "all"
+                    )
 
                 logger.info(
                     " at {:.1f}s,\testimator {}'s best error={:.4f},\tbest estimator {}'s best error={:.4f}".format(

flaml/automl/data.py

@@ -5,6 +5,7 @@
 import json
 import os
 import random
+import re
 import uuid
 from datetime import datetime, timedelta
 from decimal import ROUND_HALF_UP, Decimal
@@ -50,7 +51,10 @@ def load_openml_dataset(dataset_id, data_dir=None, random_state=0, dataset_forma
     """
     import pickle
 
-    import openml
+    try:
+        import openml
+    except ImportError:
+        openml = None
     from sklearn.model_selection import train_test_split
 
     filename = "openml_ds" + str(dataset_id) + ".pkl"
@@ -61,15 +65,15 @@ def load_openml_dataset(dataset_id, data_dir=None, random_state=0, dataset_forma
             dataset = pickle.load(f)
     else:
         print("download dataset from openml")
-        dataset = openml.datasets.get_dataset(dataset_id)
+        dataset = openml.datasets.get_dataset(dataset_id) if openml else None
         if not os.path.exists(data_dir):
             os.makedirs(data_dir)
         with open(filepath, "wb") as f:
             pickle.dump(dataset, f, pickle.HIGHEST_PROTOCOL)
-    print("Dataset name:", dataset.name)
+    print("Dataset name:", dataset.name) if dataset else None
     try:
         X, y, *__ = dataset.get_data(target=dataset.default_target_attribute, dataset_format=dataset_format)
-    except ValueError:
+    except (ValueError, AttributeError, TypeError):
         from sklearn.datasets import fetch_openml
 
         X, y = fetch_openml(data_id=dataset_id, return_X_y=True)
@@ -705,6 +709,14 @@ def auto_convert_dtypes_pandas(
     """
     if na_values is None:
         na_values = {"NA", "na", "NULL", "null", ""}
+    # Remove the empty string separately (handled by the regex `^\s*$`)
+    vals = [re.escape(v) for v in na_values if v != ""]
+    # Build inner alternation group
+    inner = "|".join(vals) if vals else ""
+    if inner:
+        pattern = re.compile(rf"^\s*(?:{inner})?\s*$")
+    else:
+        pattern = re.compile(r"^\s*$")
 
     df_converted = df.convert_dtypes()
     schema = {}
@@ -718,7 +730,11 @@ def auto_convert_dtypes_pandas(
     for col in df.columns:
         series = df[col]
         # Replace NA-like values if string
-        series_cleaned = series.map(lambda x: np.nan if isinstance(x, str) and x.strip() in na_values else x)
+        if series.dtype == object:
+            mask = series.astype(str).str.match(pattern)
+            series_cleaned = series.where(~mask, np.nan)
+        else:
+            series_cleaned = series
 
         # Skip conversion if already non-object data type, except bool which can potentially be categorical
         if (

flaml/automl/ml.py

@@ -127,9 +127,21 @@ def metric_loss_score(
             import datasets
 
             datasets_metric_name = huggingface_submetric_to_metric.get(metric_name, metric_name.split(":")[0])
-            metric = datasets.load_metric(datasets_metric_name, trust_remote_code=True)
             metric_mode = huggingface_metric_to_mode[datasets_metric_name]
 
+            # datasets>=3 removed load_metric; prefer evaluate if available
+            try:
+                import evaluate
+
+                metric = evaluate.load(datasets_metric_name, trust_remote_code=True)
+            except Exception:
+                if hasattr(datasets, "load_metric"):
+                    metric = datasets.load_metric(datasets_metric_name, trust_remote_code=True)
+                else:
+                    from datasets import load_metric as _load_metric  # older datasets
+
+                    metric = _load_metric(datasets_metric_name, trust_remote_code=True)
+
             if metric_name.startswith("seqeval"):
                 y_processed_true = [[labels[tr] for tr in each_list] for each_list in y_processed_true]
             elif metric in ("pearsonr", "spearmanr"):
@@ -604,7 +616,12 @@ def _eval_estimator(
             logger.warning(f"ValueError {e} happened in `metric_loss_score`, set `val_loss` to `np.inf`")
         metric_for_logging = {"pred_time": pred_time}
         if log_training_metric:
-            train_pred_y = get_y_pred(estimator, X_train, eval_metric, task)
+            # For time series forecasting, X_train may be a sampled dataset whose
+            # test partition can be empty. Use the training partition from X_val
+            # (which is the dataset used to define y_train above) to keep shapes
+            # aligned and avoid empty prediction inputs.
+            X_train_for_metric = X_val.X_train if isinstance(X_val, TimeSeriesDataset) else X_train
+            train_pred_y = get_y_pred(estimator, X_train_for_metric, eval_metric, task)
             metric_for_logging["train_loss"] = metric_loss_score(
                 eval_metric,
                 train_pred_y,

flaml/automl/model.py

@@ -111,7 +111,7 @@ def limit_resource(memory_limit, time_limit):
                 pass
 
 
-class BaseEstimator:
+class BaseEstimator(sklearn.base.ClassifierMixin, sklearn.base.BaseEstimator):
     """The abstract class for all learners.
 
     Typical examples:
@@ -2347,8 +2347,11 @@ class SGDEstimator(SKLearnEstimator):
         params = super().config2params(config)
         params["tol"] = params.get("tol", 0.0001)
         params["loss"] = params.get("loss", None)
-        if params["loss"] is None and self._task.is_classification():
-            params["loss"] = "log_loss" if SKLEARN_VERSION >= "1.1" else "log"
+        if params["loss"] is None:
+            if self._task.is_classification():
+                params["loss"] = "log_loss" if SKLEARN_VERSION >= "1.1" else "log"
+            else:
+                params["loss"] = "squared_error"
         if not self._task.is_classification() and "n_jobs" in params:
             params.pop("n_jobs")
 
@@ -2820,7 +2823,7 @@ class suppress_stdout_stderr:
         # Open a pair of null files
         self.null_fds = [os.open(os.devnull, os.O_RDWR) for x in range(2)]
         # Save the actual stdout (1) and stderr (2) file descriptors.
-        self.save_fds = (os.dup(1), os.dup(2))
+        self.save_fds = [os.dup(1), os.dup(2)]
 
     def __enter__(self):
         # Assign the null pointers to stdout and stderr.
@@ -2832,5 +2835,5 @@ class suppress_stdout_stderr:
         os.dup2(self.save_fds[0], 1)
         os.dup2(self.save_fds[1], 2)
         # Close the null files
-        os.close(self.null_fds[0])
-        os.close(self.null_fds[1])
+        for fd in self.null_fds + self.save_fds:
+            os.close(fd)

flaml/automl/nlp/huggingface/training_args.py

@@ -77,6 +77,14 @@ class TrainingArgumentsForAuto(TrainingArguments):
 
     logging_steps: int = field(default=500, metadata={"help": "Log every X updates steps."})
 
+    # Newer versions of HuggingFace Transformers may access `TrainingArguments.generation_config`
+    # (e.g., in generation-aware trainers/callbacks). Keep this attribute to remain compatible
+    # while defaulting to None for non-generation tasks.
+    generation_config: Optional[object] = field(
+        default=None,
+        metadata={"help": "Optional generation config (or path) used by generation-aware trainers."},
+    )
+
     @staticmethod
     def load_args_from_console():
         from dataclasses import fields

flaml/automl/spark/__init__.py

@@ -1,3 +1,5 @@
+import atexit
+import logging
 import os
 
 os.environ["PYARROW_IGNORE_TIMEZONE"] = "1"
@@ -10,13 +12,14 @@ try:
     from pyspark.pandas import Series as psSeries
     from pyspark.pandas import set_option
     from pyspark.sql import DataFrame as sparkDataFrame
+    from pyspark.sql import SparkSession
     from pyspark.util import VersionUtils
 except ImportError:
 
     class psDataFrame:
         pass
 
-    F = T = ps = sparkDataFrame = psSeries = psDataFrame
+    F = T = ps = sparkDataFrame = SparkSession = psSeries = psDataFrame
     _spark_major_minor_version = set_option = None
     ERROR = ImportError(
         """Please run pip install flaml[spark]
@@ -32,3 +35,60 @@ try:
     from pandas import DataFrame, Series
 except ImportError:
     DataFrame = Series = pd = None
+
+
+logger = logging.getLogger(__name__)
+
+
+def disable_spark_ansi_mode():
+    """Disable Spark ANSI mode if it is enabled."""
+    spark = SparkSession.getActiveSession() if hasattr(SparkSession, "getActiveSession") else None
+    adjusted = False
+    try:
+        ps_conf = ps.get_option("compute.fail_on_ansi_mode")
+    except Exception:
+        ps_conf = None
+    ansi_conf = [None, ps_conf]  # ansi_conf and ps_conf original values
+    # Spark may store the config as string 'true'/'false' (or boolean in some contexts)
+    if spark is not None:
+        ansi_conf[0] = spark.conf.get("spark.sql.ansi.enabled")
+        ansi_enabled = (
+            (isinstance(ansi_conf[0], str) and ansi_conf[0].lower() == "true")
+            or (isinstance(ansi_conf[0], bool) and ansi_conf[0] is True)
+            or ansi_conf[0] is None
+        )
+        try:
+            if ansi_enabled:
+                logger.debug("Adjusting spark.sql.ansi.enabled to false")
+                spark.conf.set("spark.sql.ansi.enabled", "false")
+                adjusted = True
+        except Exception:
+            # If reading/setting options fail for some reason, keep going and let
+            # pandas-on-Spark raise a meaningful error later.
+            logger.exception("Failed to set spark.sql.ansi.enabled")
+
+    if ansi_conf[1]:
+        logger.debug("Adjusting pandas-on-Spark compute.fail_on_ansi_mode to False")
+        ps.set_option("compute.fail_on_ansi_mode", False)
+        adjusted = True
+
+    return spark, ansi_conf, adjusted
+
+
+def restore_spark_ansi_mode(spark, ansi_conf, adjusted):
+    """Restore Spark ANSI mode to its original setting."""
+    # Restore the original spark.sql.ansi.enabled to avoid persistent side-effects.
+    if adjusted and spark and ansi_conf[0] is not None:
+        try:
+            logger.debug(f"Restoring spark.sql.ansi.enabled to {ansi_conf[0]}")
+            spark.conf.set("spark.sql.ansi.enabled", ansi_conf[0])
+        except Exception:
+            logger.exception("Failed to restore spark.sql.ansi.enabled")
+
+    if adjusted and ansi_conf[1]:
+        logger.debug(f"Restoring pandas-on-Spark compute.fail_on_ansi_mode to {ansi_conf[1]}")
+        ps.set_option("compute.fail_on_ansi_mode", ansi_conf[1])
+
+
+spark, ansi_conf, adjusted = disable_spark_ansi_mode()
+atexit.register(restore_spark_ansi_mode, spark, ansi_conf, adjusted)

flaml/automl/spark/utils.py

@@ -59,17 +59,29 @@ def to_pandas_on_spark(
     ```
     """
     set_option("compute.default_index_type", default_index_type)
-    if isinstance(df, (DataFrame, Series)):
-        return ps.from_pandas(df)
-    elif isinstance(df, sparkDataFrame):
-        if _spark_major_minor_version[0] == 3 and _spark_major_minor_version[1] < 3:
-            return df.to_pandas_on_spark(index_col=index_col)
+    try:
+        orig_ps_conf = ps.get_option("compute.fail_on_ansi_mode")
+    except Exception:
+        orig_ps_conf = None
+    if orig_ps_conf:
+        ps.set_option("compute.fail_on_ansi_mode", False)
+
+    try:
+        if isinstance(df, (DataFrame, Series)):
+            return ps.from_pandas(df)
+        elif isinstance(df, sparkDataFrame):
+            if _spark_major_minor_version[0] == 3 and _spark_major_minor_version[1] < 3:
+                return df.to_pandas_on_spark(index_col=index_col)
+            else:
+                return df.pandas_api(index_col=index_col)
+        elif isinstance(df, (psDataFrame, psSeries)):
+            return df
         else:
-            return df.pandas_api(index_col=index_col)
-    elif isinstance(df, (psDataFrame, psSeries)):
-        return df
-    else:
-        raise TypeError(f"{type(df)} is not one of pandas.DataFrame, pandas.Series and pyspark.sql.DataFrame")
+            raise TypeError(f"{type(df)} is not one of pandas.DataFrame, pandas.Series and pyspark.sql.DataFrame")
+    finally:
+        # Restore original config
+        if orig_ps_conf:
+            ps.set_option("compute.fail_on_ansi_mode", orig_ps_conf)
 
 
 def train_test_split_pyspark(

flaml/automl/task/time_series_task.py

@@ -529,7 +529,7 @@ def remove_ts_duplicates(
     duplicates = X.duplicated()
 
     if any(duplicates):
-        logger.warning("Duplicate timestamp values found in timestamp column. " f"\n{X.loc[duplicates, X][time_col]}")
+        logger.warning("Duplicate timestamp values found in timestamp column. " f"\n{X.loc[duplicates, time_col]}")
         X = X.drop_duplicates()
         logger.warning("Removed duplicate rows based on all columns")
         assert (

flaml/automl/time_series/tcn.py

@@ -264,7 +264,8 @@ class TCNEstimator(TimeSeriesEstimator):
     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

@@ -1,3 +1,4 @@
+import inspect
 import time
 
 try:
@@ -106,12 +107,17 @@ class TemporalFusionTransformerEstimator(TimeSeriesEstimator):
     def fit(self, X_train, y_train, budget=None, **kwargs):
         import warnings
 
-        import pytorch_lightning as pl
+        try:
+            import lightning.pytorch as pl
+            from lightning.pytorch.callbacks import EarlyStopping, LearningRateMonitor
+            from lightning.pytorch.loggers import TensorBoardLogger
+        except ImportError:
+            import pytorch_lightning as pl
+            from pytorch_lightning.callbacks import EarlyStopping, LearningRateMonitor
+            from pytorch_lightning.loggers import TensorBoardLogger
         import torch
         from pytorch_forecasting import TemporalFusionTransformer
         from pytorch_forecasting.metrics import QuantileLoss
-        from pytorch_lightning.callbacks import EarlyStopping, LearningRateMonitor
-        from pytorch_lightning.loggers import TensorBoardLogger
 
         # a bit of monkey patching to fix the MacOS test
         # all the log_prediction method appears to do is plot stuff, which ?breaks github tests
@@ -132,12 +138,26 @@ class TemporalFusionTransformerEstimator(TimeSeriesEstimator):
         lr_logger = LearningRateMonitor()  # log the learning rate
         logger = TensorBoardLogger(kwargs.get("log_dir", "lightning_logs"))  # logging results to a tensorboard
         default_trainer_kwargs = dict(
-            gpus=self._kwargs.get("gpu_per_trial", [0]) if torch.cuda.is_available() else None,
             max_epochs=max_epochs,
             gradient_clip_val=gradient_clip_val,
             callbacks=[lr_logger, early_stop_callback],
             logger=logger,
         )
+
+        # PyTorch Lightning >=2.0 replaced `gpus` with `accelerator`/`devices`.
+        # Also, passing `gpus=None` is not accepted on newer versions.
+        trainer_sig_params = inspect.signature(pl.Trainer.__init__).parameters
+        if torch.cuda.is_available() and "gpus" in trainer_sig_params:
+            gpus = self._kwargs.get("gpu_per_trial", None)
+            if gpus is not None:
+                default_trainer_kwargs["gpus"] = gpus
+        elif torch.cuda.is_available() and "devices" in trainer_sig_params:
+            devices = self._kwargs.get("gpu_per_trial", None)
+            if devices == -1:
+                devices = "auto"
+            if devices is not None:
+                default_trainer_kwargs["accelerator"] = "gpu"
+                default_trainer_kwargs["devices"] = devices
         trainer = pl.Trainer(
             **default_trainer_kwargs,
         )
@@ -157,7 +177,14 @@ class TemporalFusionTransformerEstimator(TimeSeriesEstimator):
             val_dataloaders=val_dataloader,
         )
         best_model_path = trainer.checkpoint_callback.best_model_path
-        best_tft = TemporalFusionTransformer.load_from_checkpoint(best_model_path)
+        # PyTorch 2.6 changed `torch.load` default `weights_only` from False -> True.
+        # Some Lightning checkpoints (including those produced here) can require full unpickling.
+        # This path is generated locally during training, so it's trusted.
+        load_sig_params = inspect.signature(TemporalFusionTransformer.load_from_checkpoint).parameters
+        if "weights_only" in load_sig_params:
+            best_tft = TemporalFusionTransformer.load_from_checkpoint(best_model_path, weights_only=False)
+        else:
+            best_tft = TemporalFusionTransformer.load_from_checkpoint(best_model_path)
         train_time = time.time() - current_time
         self._model = best_tft
         return train_time
@@ -170,7 +197,11 @@ class TemporalFusionTransformerEstimator(TimeSeriesEstimator):
         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_data.py

@@ -9,6 +9,7 @@ import numpy as np
 try:
     import pandas as pd
     from pandas import DataFrame, Series, to_datetime
+    from pandas.api.types import is_datetime64_any_dtype
     from scipy.sparse import issparse
     from sklearn.compose import ColumnTransformer
     from sklearn.impute import SimpleImputer
@@ -392,6 +393,15 @@ class DataTransformerTS:
         assert len(self.num_columns) == 0, "Trying to call fit() twice, something is wrong"
 
         for column in X.columns:
+            # Never treat the time column as a feature for sklearn preprocessing
+            if column == self.time_col:
+                continue
+
+            # Robust datetime detection (covers datetime64[ms/us/ns], tz-aware, etc.)
+            if is_datetime64_any_dtype(X[column]):
+                self.datetime_columns.append(column)
+                continue
+
             # sklearn/utils/validation.py needs int/float values
             if X[column].dtype.name in ("object", "category", "string"):
                 if (

flaml/automl/time_series/ts_model.py

@@ -194,7 +194,13 @@ class Orbit(TimeSeriesEstimator):
 
         elif isinstance(X, TimeSeriesDataset):
             data = X
-            X = data.test_data[[self.time_col] + X.regressors]
+            # By default we predict on the dataset's test partition.
+            # Some internal call paths (e.g., training-metric logging) may pass a
+            # dataset whose test partition is empty; fall back to train partition.
+            if data.test_data is not None and len(data.test_data):
+                X = data.test_data[data.regressors + [data.time_col]]
+            else:
+                X = data.train_data[data.regressors + [data.time_col]]
 
         if self._model is not None:
             forecast = self._model.predict(X, **kwargs)
@@ -301,7 +307,13 @@ class Prophet(TimeSeriesEstimator):
 
         if isinstance(X, TimeSeriesDataset):
             data = X
-            X = data.test_data[data.regressors + [data.time_col]]
+            # By default we predict on the dataset's test partition.
+            # Some internal call paths (e.g., training-metric logging) may pass a
+            # dataset whose test partition is empty; fall back to train partition.
+            if data.test_data is not None and len(data.test_data):
+                X = data.test_data[data.regressors + [data.time_col]]
+            else:
+                X = data.train_data[data.regressors + [data.time_col]]
 
         X = X.rename(columns={self.time_col: "ds"})
         if self._model is not None:
@@ -327,11 +339,19 @@ class StatsModelsEstimator(TimeSeriesEstimator):
 
         if isinstance(X, TimeSeriesDataset):
             data = X
-            X = data.test_data[data.regressors + [data.time_col]]
+            # By default we predict on the dataset's test partition.
+            # Some internal call paths (e.g., training-metric logging) may pass a
+            # dataset whose test partition is empty; fall back to train partition.
+            if data.test_data is not None and len(data.test_data):
+                X = data.test_data[data.regressors + [data.time_col]]
+            else:
+                X = data.train_data[data.regressors + [data.time_col]]
         else:
             X = X[self.regressors + [self.time_col]]
 
         if isinstance(X, DataFrame):
+            if X.shape[0] == 0:
+                return pd.Series([], name=self.target_names[0], dtype=float)
             start = X[self.time_col].iloc[0]
             end = X[self.time_col].iloc[-1]
             if len(self.regressors):
@@ -829,6 +849,13 @@ class TS_SKLearn(TimeSeriesEstimator):
         if isinstance(X, TimeSeriesDataset):
             data = X
             X = data.test_data
+            # By default we predict on the dataset's test partition.
+            # Some internal call paths (e.g., training-metric logging) may pass a
+            # dataset whose test partition is empty; fall back to train partition.
+            if data.test_data is not None and len(data.test_data):
+                X = data.test_data
+            else:
+                X = data.train_data
 
         if self._model is not None:
             X = X[self.regressors]

flaml/default/greedy.py

@@ -32,6 +32,7 @@ def construct_portfolio(regret_matrix, meta_features, regret_bound):
     if meta_features is not None:
         scaler = RobustScaler()
         meta_features = meta_features.loc[tasks]
+        meta_features = meta_features.astype(float)
         meta_features.loc[:, :] = scaler.fit_transform(meta_features)
         nearest_task = {}
         for t in tasks:

flaml/default/portfolio.py

@@ -26,6 +26,7 @@ def config_predictor_tuple(tasks, configs, meta_features, regret_matrix):
     # pre-processing
     scaler = RobustScaler()
     meta_features_norm = meta_features.loc[tasks]  # this makes a copy
+    meta_features_norm = meta_features_norm.astype(float)
     meta_features_norm.loc[:, :] = scaler.fit_transform(meta_features_norm)
 
     proc = {

flaml/fabric/mlflow.py

@@ -567,7 +567,7 @@ class MLflowIntegration:
             try:
                 with open(pickle_fpath, "wb") as f:
                     pickle.dump(obj, f)
-                mlflow.log_artifact(pickle_fpath, artifact_name, run_id)
+                self.mlflow_client.log_artifact(run_id, pickle_fpath, artifact_name)
                 return True
             except Exception as e:
                 logger.debug(f"Failed to pickle and log {artifact_name}, error: {e}")
@@ -652,7 +652,7 @@ class MLflowIntegration:
         return f"Successfully pickle_and_log_automl_artifacts {estimator} to run_id {run_id}"
 
     @time_it
-    def record_state(self, automl, search_state, estimator):
+    def record_state(self, automl, search_state, estimator, is_log_model=True):
         _st = time.time()
         automl_metric_name = (
             automl._state.metric if isinstance(automl._state.metric, str) else automl._state.error_metric
@@ -727,7 +727,7 @@ class MLflowIntegration:
                 self.futures[future] = f"iter_{automl._track_iter}_log_info_to_run"
                 future = executor.submit(lambda: self._log_automl_configurations(child_run.info.run_id))
                 self.futures[future] = f"iter_{automl._track_iter}_log_automl_configurations"
-                if automl._state.model_history:
+                if automl._state.model_history and is_log_model:
                     if estimator.endswith("_spark"):
                         future = executor.submit(
                             lambda: self.log_model(
@@ -797,8 +797,10 @@ class MLflowIntegration:
                 conf = automl._config_history[automl._best_iteration][1].copy()
                 if "ml" in conf.keys():
                     conf = conf["ml"]
-
-                mlflow.log_params({**conf, "best_learner": automl._best_estimator}, run_id=self.parent_run_id)
+                params_arr = [
+                    Param(key, str(value)) for key, value in {**conf, "best_learner": automl._best_estimator}.items()
+                ]
+                self.mlflow_client.log_batch(run_id=self.parent_run_id, metrics=[], params=params_arr, tags=[])
                 if not self.has_summary:
                     logger.info(f"logging best model {automl.best_estimator}")
                     future = executor.submit(lambda: self.copy_mlflow_run(best_mlflow_run_id, self.parent_run_id))
@@ -894,6 +896,7 @@ class MLflowIntegration:
                 ),
             )
             self.child_counter = 0
+            num_infos = len(self.infos)
 
             # From latest to earliest, remove duplicate cross-validation runs
             _exist_child_run_params = []  # for deduplication of cross-validation child runs
@@ -958,22 +961,37 @@ class MLflowIntegration:
                             )
                         self.mlflow_client.set_tag(child_run_id, "flaml.child_counter", self.child_counter)
 
-                    # merge autolog child run and corresponding manual run
-                    flaml_info = self.infos[self.child_counter]
-                    child_run = self.mlflow_client.get_run(child_run_id)
-                    self._log_info_to_run(flaml_info, child_run_id, log_params=False)
+                    # Merge autolog child run and corresponding FLAML trial info (if available).
+                    # In nested scenarios (e.g., Tune -> AutoML -> MLflow autolog), MLflow can create
+                    # more child runs than the number of FLAML trials recorded in self.infos.
+                    # TODO: need more tests in nested scenarios.
+                    flaml_info = None
+                    child_run = None
+                    if self.child_counter < num_infos:
+                        flaml_info = self.infos[self.child_counter]
+                        child_run = self.mlflow_client.get_run(child_run_id)
+                        self._log_info_to_run(flaml_info, child_run_id, log_params=False)
 
-                    if self.experiment_type == "automl":
-                        if "learner" not in child_run.data.params:
-                            self.mlflow_client.log_param(child_run_id, "learner", flaml_info["params"]["learner"])
-                        if "sample_size" not in child_run.data.params:
-                            self.mlflow_client.log_param(
-                                child_run_id, "sample_size", flaml_info["params"]["sample_size"]
-                            )
+                        if self.experiment_type == "automl":
+                            if "learner" not in child_run.data.params:
+                                self.mlflow_client.log_param(child_run_id, "learner", flaml_info["params"]["learner"])
+                            if "sample_size" not in child_run.data.params:
+                                self.mlflow_client.log_param(
+                                    child_run_id, "sample_size", flaml_info["params"]["sample_size"]
+                                )
+                    else:
+                        logger.debug(
+                            "No corresponding FLAML info for MLflow child run %s (child_counter=%s, infos=%s); skipping merge.",
+                            child_run_id,
+                            self.child_counter,
+                            num_infos,
+                        )
 
-                    if self.child_counter == best_iteration:
+                    if flaml_info is not None and self.child_counter == best_iteration:
                         self.mlflow_client.set_tag(child_run_id, "flaml.best_run", True)
                         if result is not None:
+                            if child_run is None:
+                                child_run = self.mlflow_client.get_run(child_run_id)
                             result.best_run_id = child_run_id
                             result.best_run_name = child_run.info.run_name
                             self.best_run_id = child_run_id
@@ -997,7 +1015,7 @@ class MLflowIntegration:
         self.resume_mlflow()
 
 
-def register_automl_pipeline(automl, model_name=None, signature=None):
+def register_automl_pipeline(automl, model_name=None, signature=None, artifact_path="model"):
     pipeline = automl.automl_pipeline
     if pipeline is None:
         logger.warning("pipeline not found, cannot register it")
@@ -1007,7 +1025,7 @@ def register_automl_pipeline(automl, model_name=None, signature=None):
     if automl.best_run_id is None:
         mlflow.sklearn.log_model(
             pipeline,
-            "automl_pipeline",
+            artifact_path,
             registered_model_name=model_name,
             signature=automl.pipeline_signature if signature is None else signature,
         )
@@ -1017,5 +1035,5 @@ def register_automl_pipeline(automl, model_name=None, signature=None):
         return mvs[0]
     else:
         best_run = mlflow.get_run(automl.best_run_id)
-        model_uri = f"runs:/{best_run.info.run_id}/automl_pipeline"
+        model_uri = f"runs:/{best_run.info.run_id}/{artifact_path}"
         return mlflow.register_model(model_uri, model_name)

flaml/tune/searcher/blendsearch.py

@@ -244,13 +244,32 @@ class BlendSearch(Searcher):
                     evaluated_rewards=evaluated_rewards,
                 )
             except (AssertionError, ValueError):
-                self._gs = GlobalSearch(
-                    space=gs_space,
-                    metric=metric,
-                    mode=mode,
-                    seed=gs_seed,
-                    sampler=sampler,
-                )
+                try:
+                    self._gs = GlobalSearch(
+                        space=gs_space,
+                        metric=metric,
+                        mode=mode,
+                        seed=gs_seed,
+                        sampler=sampler,
+                    )
+                except ValueError:
+                    # Ray Tune's OptunaSearch converts Tune domains into Optuna
+                    # distributions. Optuna disallows integer log distributions
+                    # with step != 1 (e.g., qlograndint with q>1), which can
+                    # raise here. Fall back to FLAML's OptunaSearch wrapper,
+                    # which handles these spaces more permissively.
+                    if getattr(GlobalSearch, "__module__", "").startswith("ray.tune"):
+                        from .suggestion import OptunaSearch as _FallbackOptunaSearch
+
+                        self._gs = _FallbackOptunaSearch(
+                            space=gs_space,
+                            metric=metric,
+                            mode=mode,
+                            seed=gs_seed,
+                            sampler=sampler,
+                        )
+                    else:
+                        raise
             self._gs.space = space
         else:
             self._gs = None

flaml/tune/searcher/suggestion.py

@@ -35,6 +35,73 @@ from ..sample import (
     Quantized,
     Uniform,
 )
+
+# If Ray is installed, flaml.tune may re-export Ray Tune sampling functions.
+# In that case, the search space contains Ray Tune Domain/Sampler objects,
+# which should be accepted by our Optuna search-space conversion.
+try:
+    from ray import __version__ as _ray_version  # type: ignore
+
+    if str(_ray_version).startswith("1."):
+        from ray.tune.sample import (  # type: ignore
+            Categorical as _RayCategorical,
+        )
+        from ray.tune.sample import (
+            Domain as _RayDomain,
+        )
+        from ray.tune.sample import (
+            Float as _RayFloat,
+        )
+        from ray.tune.sample import (
+            Integer as _RayInteger,
+        )
+        from ray.tune.sample import (
+            LogUniform as _RayLogUniform,
+        )
+        from ray.tune.sample import (
+            Quantized as _RayQuantized,
+        )
+        from ray.tune.sample import (
+            Uniform as _RayUniform,
+        )
+    else:
+        from ray.tune.search.sample import (  # type: ignore
+            Categorical as _RayCategorical,
+        )
+        from ray.tune.search.sample import (
+            Domain as _RayDomain,
+        )
+        from ray.tune.search.sample import (
+            Float as _RayFloat,
+        )
+        from ray.tune.search.sample import (
+            Integer as _RayInteger,
+        )
+        from ray.tune.search.sample import (
+            LogUniform as _RayLogUniform,
+        )
+        from ray.tune.search.sample import (
+            Quantized as _RayQuantized,
+        )
+        from ray.tune.search.sample import (
+            Uniform as _RayUniform,
+        )
+
+    _FLOAT_TYPES = (Float, _RayFloat)
+    _INTEGER_TYPES = (Integer, _RayInteger)
+    _CATEGORICAL_TYPES = (Categorical, _RayCategorical)
+    _DOMAIN_TYPES = (Domain, _RayDomain)
+    _QUANTIZED_TYPES = (Quantized, _RayQuantized)
+    _UNIFORM_TYPES = (Uniform, _RayUniform)
+    _LOGUNIFORM_TYPES = (LogUniform, _RayLogUniform)
+except Exception:  # pragma: no cover
+    _FLOAT_TYPES = (Float,)
+    _INTEGER_TYPES = (Integer,)
+    _CATEGORICAL_TYPES = (Categorical,)
+    _DOMAIN_TYPES = (Domain,)
+    _QUANTIZED_TYPES = (Quantized,)
+    _UNIFORM_TYPES = (Uniform,)
+    _LOGUNIFORM_TYPES = (LogUniform,)
 from ..trial import flatten_dict, unflatten_dict
 from .variant_generator import parse_spec_vars
 
@@ -850,19 +917,22 @@ class OptunaSearch(Searcher):
         def resolve_value(domain: Domain) -> ot.distributions.BaseDistribution:
             quantize = None
 
-            sampler = domain.get_sampler()
-            if isinstance(sampler, Quantized):
+            # Ray Tune Domains and FLAML Domains both provide get_sampler(), but
+            # fall back to the .sampler attribute for robustness.
+            sampler = domain.get_sampler() if hasattr(domain, "get_sampler") else getattr(domain, "sampler", None)
+
+            if isinstance(sampler, _QUANTIZED_TYPES) or type(sampler).__name__ == "Quantized":
                 quantize = sampler.q
-                sampler = sampler.sampler
-                if isinstance(sampler, LogUniform):
+                sampler = getattr(sampler, "sampler", None) or sampler.get_sampler()
+                if isinstance(sampler, _LOGUNIFORM_TYPES) or type(sampler).__name__ == "LogUniform":
                     logger.warning(
                         "Optuna does not handle quantization in loguniform "
                         "sampling. The parameter will be passed but it will "
                         "probably be ignored."
                     )
 
-            if isinstance(domain, Float):
-                if isinstance(sampler, LogUniform):
+            if isinstance(domain, _FLOAT_TYPES) or type(domain).__name__ == "Float":
+                if isinstance(sampler, _LOGUNIFORM_TYPES) or type(sampler).__name__ == "LogUniform":
                     if quantize:
                         logger.warning(
                             "Optuna does not support both quantization and "
@@ -870,17 +940,17 @@ class OptunaSearch(Searcher):
                         )
                     return ot.distributions.LogUniformDistribution(domain.lower, domain.upper)
 
-                elif isinstance(sampler, Uniform):
+                elif isinstance(sampler, _UNIFORM_TYPES) or type(sampler).__name__ == "Uniform":
                     if quantize:
                         return ot.distributions.DiscreteUniformDistribution(domain.lower, domain.upper, quantize)
                     return ot.distributions.UniformDistribution(domain.lower, domain.upper)
 
-            elif isinstance(domain, Integer):
-                if isinstance(sampler, LogUniform):
+            elif isinstance(domain, _INTEGER_TYPES) or type(domain).__name__ == "Integer":
+                if isinstance(sampler, _LOGUNIFORM_TYPES) or type(sampler).__name__ == "LogUniform":
                     # ``step`` argument Deprecated in v2.0.0. ``step`` argument should be 1 in Log Distribution
                     # The removal of this feature is currently scheduled for v4.0.0,
                     return ot.distributions.IntLogUniformDistribution(domain.lower, domain.upper - 1, step=1)
-                elif isinstance(sampler, Uniform):
+                elif isinstance(sampler, _UNIFORM_TYPES) or type(sampler).__name__ == "Uniform":
                     # Upper bound should be inclusive for quantization and
                     # exclusive otherwise
                     return ot.distributions.IntUniformDistribution(
@@ -888,16 +958,16 @@ class OptunaSearch(Searcher):
                         domain.upper - int(bool(not quantize)),
                         step=quantize or 1,
                     )
-            elif isinstance(domain, Categorical):
-                if isinstance(sampler, Uniform):
+            elif isinstance(domain, _CATEGORICAL_TYPES) or type(domain).__name__ == "Categorical":
+                if isinstance(sampler, _UNIFORM_TYPES) or type(sampler).__name__ == "Uniform":
                     return ot.distributions.CategoricalDistribution(domain.categories)
 
             raise ValueError(
                 "Optuna search does not support parameters of type "
-                "`{}` with samplers of type `{}`".format(type(domain).__name__, type(domain.sampler).__name__)
+                "`{}` with samplers of type `{}`".format(type(domain).__name__, type(sampler).__name__)
             )
 
         # Parameter name is e.g. "a/b/c" for nested dicts
         values = {"/".join(path): resolve_value(domain) for path, domain in domain_vars}
 
-        return values
+        return values

flaml/tune/space.py

@@ -261,7 +261,7 @@ def add_cost_to_space(space: Dict, low_cost_point: Dict, choice_cost: Dict):
                         low_cost[i] = point
                 if len(low_cost) > len(domain.categories):
                     if domain.ordered:
-                        low_cost[-1] = int(np.where(ind == low_cost[-1])[0])
+                        low_cost[-1] = int(np.where(ind == low_cost[-1])[0].item())
                     domain.low_cost_point = low_cost[-1]
                 return
         if low_cost:

flaml/tune/tune.py

@@ -776,7 +776,7 @@ def run(
                         and (num_samples < 0 or num_trials < num_samples)
                         and num_failures < upperbound_num_failures
                     ):
-                        if automl_info and automl_info[0] > 0 and time_budget_s < np.inf:
+                        if automl_info and automl_info[1] == "all" and automl_info[0] > 0 and time_budget_s < np.inf:
                             time_budget_s -= automl_info[0] * n_concurrent_trials
                             logger.debug(f"Remaining time budget with mlflow log latency: {time_budget_s} seconds.")
                         while len(_runner.running_trials) < n_concurrent_trials:
@@ -802,9 +802,17 @@ def run(
                         )
                         results = None
                         with PySparkOvertimeMonitor(time_start, time_budget_s, force_cancel, parallel=parallel):
-                            results = parallel(
-                                delayed(evaluation_function)(trial_to_run.config) for trial_to_run in trials_to_run
-                            )
+                            try:
+                                results = parallel(
+                                    delayed(evaluation_function)(trial_to_run.config) for trial_to_run in trials_to_run
+                                )
+                            except RuntimeError as e:
+                                logger.warning(f"RuntimeError: {e}")
+                                results = None
+                                logger.info(
+                                    "Encountered RuntimeError. Waiting 10 seconds for Spark cluster to recover before retrying."
+                                )
+                                time.sleep(10)
                         # results = [evaluation_function(trial_to_run.config) for trial_to_run in trials_to_run]
                         while results:
                             result = results.pop(0)

flaml/version.py

@@ -1 +1 @@
-__version__ = "2.3.6"
+__version__ = "2.4.0"

pyproject.toml

@@ -2,7 +2,6 @@
 license_file = "LICENSE"
 description-file = "README.md"
 
-
 [tool.pytest.ini_options]
 addopts = '-m "not conda"'
 markers = [

setup.py

@@ -51,60 +51,59 @@ setuptools.setup(
             "joblib<=1.3.2",
         ],
         "test": [
-            "jupyter",
+            "numpy>=1.17,<2.0.0; python_version<'3.13'",
+            "numpy>2.0.0; python_version>='3.13'",
+            "jupyter; python_version<'3.13'",
             "lightgbm>=2.3.1",
-            "xgboost>=0.90,<2.0.0",
+            "xgboost>=0.90,<2.0.0; python_version<'3.11'",
+            "xgboost>=2.0.0; python_version>='3.11'",
             "scipy>=1.4.1",
             "pandas>=1.1.4,<2.0.0; python_version<'3.10'",
             "pandas>=1.1.4; python_version>='3.10'",
-            "scikit-learn>=1.0.0",
+            "scikit-learn>=1.2.0",
             "thop",
             "pytest>=6.1.1",
+            "pytest-rerunfailures>=13.0",
             "coverage>=5.3",
             "pre-commit",
             "torch",
             "torchvision",
-            "catboost>=0.26,<1.2; python_version<'3.11'",
-            "catboost>=0.26; python_version>='3.11'",
+            "catboost>=0.26; python_version<'3.13'",
             "rgf-python",
             "optuna>=2.8.0,<=3.6.1",
-            "openml",
+            "openml; python_version<'3.13'",
             "statsmodels>=0.12.2",
-            "psutil==5.8.0",
+            "psutil",
             "dataclasses",
-            "transformers[torch]==4.26",
-            "datasets<=3.5.0",
-            "nltk<=3.8.1",  # 3.8.2 doesn't work with mlflow
+            "transformers[torch]",
+            "datasets",
+            "evaluate",
+            "nltk!=3.8.2",  # 3.8.2 doesn't work with mlflow
             "rouge_score",
-            "hcrystalball==0.1.10",
+            "hcrystalball",
             "seqeval",
-            "pytorch-forecasting>=0.9.0,<=0.10.1; python_version<'3.11'",
-            # "pytorch-forecasting==0.10.1; python_version=='3.11'",
-            "mlflow==2.15.1",
+            "pytorch-forecasting; python_version<'3.13'",
+            "mlflow-skinny<=2.22.1",  # Refer to https://mvnrepository.com/artifact/org.mlflow/mlflow-spark
             "joblibspark>=0.5.0",
             "joblib<=1.3.2",
             "nbconvert",
             "nbformat",
             "ipykernel",
-            "pytorch-lightning<1.9.1",  # test_forecast_panel
-            "tensorboardX==2.6",  # test_forecast_panel
-            "requests<2.29.0",  # https://github.com/docker/docker-py/issues/3113
+            "pytorch-lightning",  # test_forecast_panel
+            "tensorboardX",  # test_forecast_panel
+            "requests",  # https://github.com/docker/docker-py/issues/3113
             "packaging",
-            "pydantic==1.10.9",
-            "sympy",
-            "wolframalpha",
             "dill",  # a drop in replacement of pickle
         ],
         "catboost": [
-            "catboost>=0.26,<1.2; python_version<'3.11'",
-            "catboost>=0.26,<=1.2.5; python_version>='3.11'",
+            "catboost>=0.26",
         ],
         "blendsearch": [
             "optuna>=2.8.0,<=3.6.1",
             "packaging",
         ],
         "ray": [
-            "ray[tune]~=1.13",
+            "ray[tune]>=1.13,<2.5.0",
         ],
         "azureml": [
             "azureml-mlflow",
@@ -131,33 +130,21 @@ setuptools.setup(
             "seqeval",
         ],
         "ts_forecast": [
-            "holidays<0.14",  # to prevent installation error for prophet
-            "prophet>=1.0.1",
+            "holidays",
+            "prophet>=1.1.5",
             "statsmodels>=0.12.2",
-            "hcrystalball==0.1.10",
+            "hcrystalball>=0.1.10",
         ],
         "forecast": [
-            "holidays<0.14",  # to prevent installation error for prophet
-            "prophet>=1.0.1",
+            "holidays",
+            "prophet>=1.1.5",
             "statsmodels>=0.12.2",
-            "hcrystalball==0.1.10",
-            "pytorch-forecasting>=0.9.0; python_version<'3.11'",
-            # "pytorch-forecasting==0.10.1; python_version=='3.11'",
-            "pytorch-lightning==1.9.0",
-            "tensorboardX==2.6",
+            "hcrystalball>=0.1.10",
+            "pytorch-forecasting>=0.10.4; python_version<'3.13'",
+            "pytorch-lightning>=1.9.0",
+            "tensorboardX>=2.6",
         ],
         "benchmark": ["catboost>=0.26", "psutil==5.8.0", "xgboost==1.3.3", "pandas==1.1.4"],
-        "openai": ["openai==0.27.8", "diskcache"],
-        "autogen": ["openai==0.27.8", "diskcache", "termcolor"],
-        "mathchat": ["openai==0.27.8", "diskcache", "termcolor", "sympy", "pydantic==1.10.9", "wolframalpha"],
-        "retrievechat": [
-            "openai==0.27.8",
-            "diskcache",
-            "termcolor",
-            "chromadb",
-            "tiktoken",
-            "sentence_transformers",
-        ],
         "synapse": [
             "joblibspark>=0.5.0",
             "optuna>=2.8.0,<=3.6.1",
@@ -170,9 +157,9 @@ setuptools.setup(
         "Operating System :: OS Independent",
         # Specify the Python versions you support here.
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
     ],
-    python_requires=">=3.9",
+    python_requires=">=3.10",
 )

test/automl/test_custom_hp.py

@@ -4,8 +4,17 @@ import pytest
 
 from flaml import AutoML, tune
 
+try:
+    import transformers
 
-@pytest.mark.skipif(sys.platform == "darwin", reason="do not run on mac os")
+    _transformers_installed = True
+except ImportError:
+    _transformers_installed = False
+
+
+@pytest.mark.skipif(
+    sys.platform == "darwin" or not _transformers_installed, reason="do not run on mac os or transformers not installed"
+)
 def test_custom_hp_nlp():
     from test.nlp.utils import get_automl_settings, get_toy_data_seqclassification
 

test/automl/test_extra_models.py

@@ -1,3 +1,4 @@
+import atexit
 import os
 import sys
 import unittest
@@ -15,8 +16,16 @@ from sklearn.model_selection import train_test_split
 
 from flaml import AutoML
 from flaml.automl.ml import sklearn_metric_loss_score
+from flaml.automl.spark import disable_spark_ansi_mode, restore_spark_ansi_mode
 from flaml.tune.spark.utils import check_spark
 
+try:
+    import pytorch_lightning
+
+    _pl_installed = True
+except ImportError:
+    _pl_installed = False
+
 pytestmark = pytest.mark.spark
 
 leaderboard = defaultdict(dict)
@@ -39,7 +48,7 @@ else:
             .config(
                 "spark.jars.packages",
                 (
-                    "com.microsoft.azure:synapseml_2.12:1.0.2,"
+                    "com.microsoft.azure:synapseml_2.12:1.1.0,"
                     "org.apache.hadoop:hadoop-azure:3.3.5,"
                     "com.microsoft.azure:azure-storage:8.6.6,"
                     f"org.mlflow:mlflow-spark_2.12:{mlflow.__version__}"
@@ -63,6 +72,9 @@ else:
     except ImportError:
         skip_spark = True
 
+spark, ansi_conf, adjusted = disable_spark_ansi_mode()
+atexit.register(restore_spark_ansi_mode, spark, ansi_conf, adjusted)
+
 
 def _test_regular_models(estimator_list, task):
     if isinstance(estimator_list, str):
@@ -271,7 +283,11 @@ class TestExtraModel(unittest.TestCase):
 
     @unittest.skipIf(skip_spark, reason="Spark is not installed. Skip all spark tests.")
     def test_default_spark(self):
-        _test_spark_models(None, "classification")
+        # TODO: remove the estimator assignment once SynapseML supports spark 4+.
+        from flaml.automl.spark.utils import _spark_major_minor_version
+
+        estimator_list = ["rf_spark"] if _spark_major_minor_version[0] >= 4 else None
+        _test_spark_models(estimator_list, "classification")
 
     def test_svc(self):
         _test_regular_models("svc", "classification")
@@ -302,7 +318,7 @@ class TestExtraModel(unittest.TestCase):
     def test_avg(self):
         _test_forecast("avg")
 
-    @unittest.skipIf(skip_spark, reason="Skip on Mac or Windows")
+    @unittest.skipIf(skip_spark or not _pl_installed, reason="Skip on Mac or Windows or no pytorch_lightning.")
     def test_tcn(self):
         _test_forecast("tcn")
 

test/automl/test_forecast.py

@@ -10,7 +10,7 @@ from flaml import AutoML
 from flaml.automl.task.time_series_task import TimeSeriesTask
 
 
-def test_forecast_automl(budget=10, estimators_when_no_prophet=["arima", "sarimax", "holt-winters"]):
+def test_forecast_automl(budget=20, estimators_when_no_prophet=["arima", "sarimax", "holt-winters"]):
     # using dataframe
     import statsmodels.api as sm
 
@@ -510,8 +510,12 @@ def get_stalliion_data():
     "3.11" in sys.version,
     reason="do not run on py 3.11",
 )
-def test_forecast_panel(budget=5):
-    data, special_days = get_stalliion_data()
+def test_forecast_panel(budget=30):
+    try:
+        data, special_days = get_stalliion_data()
+    except ImportError:
+        print("pytorch_forecasting not installed")
+        return
     time_horizon = 6  # predict six months
     training_cutoff = data["time_idx"].max() - time_horizon
     data["time_idx"] = data["time_idx"].astype("int")
@@ -677,11 +681,55 @@ def test_cv_step():
     print("yahoo!")
 
 
+def test_log_training_metric_ts_models():
+    """Test that log_training_metric=True works with time series models (arima, sarimax, holt-winters)."""
+    import statsmodels.api as sm
+
+    from flaml.automl.task.time_series_task import TimeSeriesTask
+
+    estimators_all = TimeSeriesTask("forecast").estimators.keys()
+    estimators_to_test = ["xgboost", "arima", "lassolars", "tcn", "snaive", "prophet", "orbit"]
+    estimators = [
+        est for est in estimators_to_test if est in estimators_all
+    ]  # not all estimators available in current python env
+    print(f"Testing estimators: {estimators}")
+
+    # Prepare data
+    data = sm.datasets.co2.load_pandas().data["co2"]
+    data = data.resample("MS").mean()
+    data = data.bfill().ffill()
+    data = data.to_frame().reset_index()
+    data = data.rename(columns={"index": "ds", "co2": "y"})
+    num_samples = data.shape[0]
+    time_horizon = 12
+    split_idx = num_samples - time_horizon
+    df = data[:split_idx]
+
+    # Test each time series model with log_training_metric=True
+    for estimator in estimators:
+        print(f"\nTesting {estimator} with log_training_metric=True")
+        automl = AutoML()
+        settings = {
+            "time_budget": 3,
+            "metric": "mape",
+            "task": "forecast",
+            "eval_method": "holdout",
+            "label": "y",
+            "log_training_metric": True,  # This should not cause errors
+            "estimator_list": [estimator],
+        }
+        automl.fit(dataframe=df, **settings, period=time_horizon, force_cancel=True)
+        print(f"  ✅ {estimator} SUCCESS with log_training_metric=True")
+        if automl.best_estimator:
+            assert automl.best_estimator == estimator
+
+
 if __name__ == "__main__":
     # test_forecast_automl(60)
     # test_multivariate_forecast_num(5)
     # test_multivariate_forecast_cat(5)
-    test_numpy()
+    # test_numpy()
     # test_forecast_classification(5)
     # test_forecast_panel(5)
     # test_cv_step()
+    test_log_training_metric_ts_models()

test/automl/test_notebook_example.py

@@ -1,8 +1,23 @@
 import sys
 
 import pytest
-from minio.error import ServerError
-from openml.exceptions import OpenMLServerException
+
+try:
+    from minio.error import ServerError
+except ImportError:
+
+    class ServerError(Exception):
+        pass
+
+
+try:
+    from openml.exceptions import OpenMLServerException
+except ImportError:
+
+    class OpenMLServerException(Exception):
+        pass
+
+
 from requests.exceptions import ChunkedEncodingError, SSLError
 
 
@@ -64,6 +79,9 @@ def test_automl(budget=5, dataset_format="dataframe", hpo_method=None):
     automl.fit(X_train=X_train, y_train=y_train, **settings)
     """ retrieve best config and best learner """
     print("Best ML leaner:", automl.best_estimator)
+    if not automl.best_estimator:
+        print("Training budget is not sufficient")
+        return
     print("Best hyperparmeter config:", automl.best_config)
     print(f"Best accuracy on validation data: {1 - automl.best_loss:.4g}")
     print(f"Training duration of best run: {automl.best_config_train_time:.4g} s")

test/automl/test_python_log.py

@@ -38,7 +38,7 @@ class TestLogging(unittest.TestCase):
                 "keep_search_state": True,
                 "learner_selector": "roundrobin",
             }
-            X_train, y_train = fetch_california_housing(return_X_y=True)
+            X_train, y_train = fetch_california_housing(return_X_y=True, data_home="test")
             n = len(y_train) >> 1
             print(automl.model, automl.classes_, automl.predict(X_train))
             automl.fit(

test/automl/test_regression.py

@@ -47,7 +47,7 @@ class TestRegression(unittest.TestCase):
             "n_jobs": 1,
             "model_history": True,
         }
-        X_train, y_train = fetch_california_housing(return_X_y=True)
+        X_train, y_train = fetch_california_housing(return_X_y=True, data_home="test")
         n = int(len(y_train) * 9 // 10)
         automl.fit(X_train=X_train[:n], y_train=y_train[:n], X_val=X_train[n:], y_val=y_train[n:], **automl_settings)
         assert automl._state.eval_method == "holdout"
@@ -141,7 +141,7 @@ class TestRegression(unittest.TestCase):
             "n_concurrent_trials": 10,
             "hpo_method": hpo_method,
         }
-        X_train, y_train = fetch_california_housing(return_X_y=True)
+        X_train, y_train = fetch_california_housing(return_X_y=True, data_home="test")
         try:
             automl_experiment.fit(X_train=X_train, y_train=y_train, **automl_settings)
             print(automl_experiment.predict(X_train))
@@ -268,7 +268,7 @@ def test_reproducibility_of_regression_models(estimator: str):
         "skip_transform": True,
         "retrain_full": True,
     }
-    X, y = fetch_california_housing(return_X_y=True, as_frame=True)
+    X, y = fetch_california_housing(return_X_y=True, as_frame=True, data_home="test")
     automl.fit(X_train=X, y_train=y, **automl_settings)
     best_model = automl.model
     assert best_model is not None
@@ -314,7 +314,7 @@ def test_reproducibility_of_catboost_regression_model():
         "skip_transform": True,
         "retrain_full": True,
     }
-    X, y = fetch_california_housing(return_X_y=True, as_frame=True)
+    X, y = fetch_california_housing(return_X_y=True, as_frame=True, data_home="test")
     automl.fit(X_train=X, y_train=y, **automl_settings)
     best_model = automl.model
     assert best_model is not None
@@ -360,7 +360,7 @@ def test_reproducibility_of_lgbm_regression_model():
         "skip_transform": True,
         "retrain_full": True,
     }
-    X, y = fetch_california_housing(return_X_y=True, as_frame=True)
+    X, y = fetch_california_housing(return_X_y=True, as_frame=True, data_home="test")
     automl.fit(X_train=X, y_train=y, **automl_settings)
     best_model = automl.model
     assert best_model is not None
@@ -424,7 +424,7 @@ def test_reproducibility_of_underlying_regression_models(estimator: str):
         "skip_transform": True,
         "retrain_full": False,
     }
-    X, y = fetch_california_housing(return_X_y=True, as_frame=True)
+    X, y = fetch_california_housing(return_X_y=True, as_frame=True, data_home="test")
     automl.fit(X_train=X, y_train=y, **automl_settings)
     best_model = automl.model
     assert best_model is not None

test/automl/test_score.py

@@ -142,7 +142,7 @@ class TestScore:
     def test_regression(self):
         automl_experiment = AutoML()
 
-        X_train, y_train = fetch_california_housing(return_X_y=True)
+        X_train, y_train = fetch_california_housing(return_X_y=True, data_home="test")
         n = int(len(y_train) * 9 // 10)
 
         for each_estimator in [

test/automl/test_training_log.py

@@ -30,7 +30,7 @@ class TestTrainingLog(unittest.TestCase):
                 "keep_search_state": True,
                 "estimator_list": estimator_list,
             }
-            X_train, y_train = fetch_california_housing(return_X_y=True)
+            X_train, y_train = fetch_california_housing(return_X_y=True, data_home="test")
             automl.fit(X_train=X_train, y_train=y_train, **automl_settings)
             # Check if the training log file is populated.
             self.assertTrue(os.path.exists(filename))

test/automl/test_warmstart.py

@@ -108,7 +108,14 @@ class TestWarmStart(unittest.TestCase):
 
     def test_FLAML_sample_size_in_starting_points(self):
         from minio.error import ServerError
-        from openml.exceptions import OpenMLServerException
+
+        try:
+            from openml.exceptions import OpenMLServerException
+        except ImportError:
+
+            class OpenMLServerException(Exception):
+                pass
+
         from requests.exceptions import ChunkedEncodingError, SSLError
 
         from flaml import AutoML

test/check_dependency.py

@@ -0,0 +1,60 @@
+import subprocess
+from importlib.metadata import distributions
+
+installed_libs = sorted(f"{dist.metadata['Name']}=={dist.version}" for dist in distributions())
+
+first_tier_dependencies = [
+    "numpy",
+    "jupyter",
+    "lightgbm",
+    "xgboost",
+    "scipy",
+    "pandas",
+    "scikit-learn",
+    "thop",
+    "pytest",
+    "pytest-rerunfailures",
+    "coverage",
+    "pre-commit",
+    "torch",
+    "torchvision",
+    "catboost",
+    "rgf-python",
+    "optuna",
+    "openml",
+    "statsmodels",
+    "psutil",
+    "dataclasses",
+    "transformers[torch]",
+    "transformers",
+    "datasets",
+    "evaluate",
+    "nltk",
+    "rouge_score",
+    "hcrystalball",
+    "seqeval",
+    "pytorch-forecasting",
+    "mlflow-skinny",
+    "joblibspark",
+    "joblib",
+    "nbconvert",
+    "nbformat",
+    "ipykernel",
+    "pytorch-lightning",
+    "tensorboardX",
+    "requests",
+    "packaging",
+    "dill",
+    "ray",
+    "prophet",
+]
+
+
+for lib in installed_libs:
+    lib_name = lib.split("==")[0]
+    if lib_name in first_tier_dependencies:
+        print(lib)
+
+# print current commit hash
+commit_hash = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("utf-8").strip()
+print(f"Current commit hash: {commit_hash}")

test/conftest.py

@@ -2,11 +2,24 @@ from typing import Any, Dict, List, Union
 
 import numpy as np
 import pandas as pd
-from catboost import CatBoostClassifier, CatBoostRegressor, Pool
+import pytest
 from sklearn.metrics import f1_score, r2_score
 
+try:
+    from catboost import CatBoostClassifier, CatBoostRegressor, Pool
+except ImportError:  # pragma: no cover
+    CatBoostClassifier = None
+    CatBoostRegressor = None
+    Pool = None
 
-def evaluate_cv_folds_with_underlying_model(X_train_all, y_train_all, kf, model: Any, task: str) -> pd.DataFrame:
+
+def _is_catboost_model_type(model_type: type) -> bool:
+    if CatBoostClassifier is not None and CatBoostRegressor is not None:
+        return model_type is CatBoostClassifier or model_type is CatBoostRegressor
+    return getattr(model_type, "__module__", "").startswith("catboost")
+
+
+def evaluate_cv_folds_with_underlying_model(X_train_all, y_train_all, kf, model: Any, task: str) -> List[float]:
     """Mimic the FLAML CV process to calculate the metrics across each fold.
 
     :param X_train_all: X training data
@@ -17,7 +30,7 @@ def evaluate_cv_folds_with_underlying_model(X_train_all, y_train_all, kf, model:
     :return: An array containing the metrics
     """
     rng = np.random.RandomState(2020)
-    all_fold_metrics: List[Dict[str, Union[int, float]]] = []
+    all_fold_metrics: List[float] = []
     for train_index, val_index in kf.split(X_train_all, y_train_all):
         X_train_split, y_train_split = X_train_all, y_train_all
         train_index = rng.permutation(train_index)
@@ -25,9 +38,11 @@ def evaluate_cv_folds_with_underlying_model(X_train_all, y_train_all, kf, model:
         X_val = X_train_split.iloc[val_index]
         y_train, y_val = y_train_split[train_index], y_train_split[val_index]
         model_type = type(model)
-        if model_type is not CatBoostClassifier and model_type is not CatBoostRegressor:
+        if not _is_catboost_model_type(model_type):
             model.fit(X_train, y_train)
         else:
+            if Pool is None:
+                pytest.skip("catboost is not installed")
             use_best_model = True
             n = max(int(len(y_train) * 0.9), len(y_train) - 1000) if use_best_model else len(y_train)
             X_tr, y_tr = (X_train)[:n], y_train[:n]
@@ -38,5 +53,5 @@ def evaluate_cv_folds_with_underlying_model(X_train_all, y_train_all, kf, model:
             reproduced_metric = 1 - f1_score(y_val, y_pred_classes)
         else:
             reproduced_metric = 1 - r2_score(y_val, y_pred_classes)
-        all_fold_metrics.append(reproduced_metric)
+        all_fold_metrics.append(float(reproduced_metric))
     return all_fold_metrics

test/default/test_defaults.py

@@ -60,7 +60,7 @@ 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)
+    X_train, y_train = fetch_california_housing(return_X_y=True, as_frame=as_frame, data_home="test")
     automl.fit(X_train, y_train, **automl_settings)
 
 
@@ -115,7 +115,7 @@ 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(return_X_y=True, as_frame=True, data_home="test")
     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 +137,7 @@ 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(return_X_y=True, as_frame=True, data_home="test")
     rf = RandomForestRegressor(default_location=location)
     rf.fit(X_train[:100], y_train[:100])
     rf.predict(X_train)
@@ -155,7 +155,7 @@ 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(return_X_y=True, as_frame=True, data_home="test")
     regressor = ExtraTreesRegressor(default_location=location)
     regressor.fit(X_train[:100], y_train[:100])
     regressor.predict(X_train)
@@ -175,7 +175,7 @@ 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(return_X_y=True, as_frame=True, data_home="test")
     regressor = LGBMRegressor(default_location=location)
     regressor.fit(X_train, y_train)
     regressor.predict(X_train)
@@ -194,7 +194,7 @@ 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(return_X_y=True, as_frame=True, data_home="test")
     regressor = XGBRegressor(default_location=location)
     regressor.fit(X_train[:100], y_train[:100])
     regressor.predict(X_train)

test/nlp/test_autohf_classificationhead.py

@@ -3,6 +3,12 @@ import shutil
 import sys
 
 import pytest
+
+try:
+    import transformers
+except ImportError:
+    pytest.skip("transformers not installed", allow_module_level=True)
+
 from utils import (
     get_automl_settings,
     get_toy_data_binclassification,

test/nlp/test_autohf_cv.py

@@ -5,10 +5,20 @@ import sys
 import pytest
 from utils import get_automl_settings, get_toy_data_seqclassification
 
+try:
+    import transformers
+
+    _transformers_installed = True
+except ImportError:
+    _transformers_installed = False
+
 pytestmark = pytest.mark.spark  # set to spark as parallel testing raised MlflowException of changing parameter
 
 
-@pytest.mark.skipif(sys.platform in ["darwin", "win32"], reason="do not run on mac os or windows")
+@pytest.mark.skipif(
+    sys.platform in ["darwin", "win32"] or not _transformers_installed,
+    reason="do not run on mac os or windows or transformers not installed",
+)
 def test_cv():
     import requests
 

test/nlp/test_autohf_multichoice_classification.py

@@ -5,8 +5,18 @@ import sys
 import pytest
 from utils import get_automl_settings, get_toy_data_multiplechoiceclassification
 
+try:
+    import transformers
 
-@pytest.mark.skipif(sys.platform in ["darwin", "win32"], reason="do not run on mac os or windows")
+    _transformers_installed = True
+except ImportError:
+    _transformers_installed = False
+
+
+@pytest.mark.skipif(
+    sys.platform in ["darwin", "win32"] or not _transformers_installed,
+    reason="do not run on mac os or windows or transformers not installed",
+)
 def test_mcc():
     import requests
 

test/nlp/test_default.py

@@ -7,8 +7,20 @@ from utils import get_automl_settings, get_toy_data_seqclassification
 
 from flaml.default import portfolio
 
-if sys.platform.startswith("darwin") and sys.version_info[0] == 3 and sys.version_info[1] == 11:
-    pytest.skip("skipping Python 3.11 on MacOS", allow_module_level=True)
+try:
+    import transformers
+
+    _transformers_installed = True
+except ImportError:
+    _transformers_installed = False
+
+if (
+    sys.platform.startswith("darwin")
+    and sys.version_info >= (3, 11)
+    or not _transformers_installed
+    or sys.platform == "win32"
+):
+    pytest.skip("skipping Python 3.11 on MacOS or without transformers or on Windows", allow_module_level=True)
 
 pytestmark = (
     pytest.mark.spark
@@ -28,23 +40,34 @@ def test_build_portfolio(path="./test/nlp/default", strategy="greedy"):
     portfolio.main()
 
 
-@pytest.mark.skipif(sys.platform == "win32", reason="do not run on windows")
 def test_starting_point_not_in_search_space():
-    from flaml import AutoML
+    """Regression test for invalid starting points and custom_hp.
+
+    This test must not require network access to Hugging Face.
+    """
 
     """
         test starting_points located outside of the search space, and custom_hp is not set
     """
+    from flaml.automl.state import SearchState
+    from flaml.automl.task.factory import task_factory
+
     this_estimator_name = "transformer"
-    X_train, y_train, X_val, y_val, _ = get_toy_data_seqclassification()
+    X_train, y_train, _, _, _ = get_toy_data_seqclassification()
+    task = task_factory("seq-classification", X_train, y_train)
+    estimator_class = task.estimator_class_from_str(this_estimator_name)
+    estimator_class.init()
 
-    automl = AutoML()
-    automl_settings = get_automl_settings(estimator_name=this_estimator_name)
-
-    automl_settings["starting_points"] = {this_estimator_name: [{"learning_rate": 2e-3}]}
-
-    automl.fit(X_train, y_train, **automl_settings)
-    assert automl._search_states[this_estimator_name].init_config[0]["learning_rate"] != 2e-3
+    # SearchState is where invalid starting points are filtered out when max_iter > 1.
+    search_state = SearchState(
+        learner_class=estimator_class,
+        data=X_train,
+        task=task,
+        starting_point={"learning_rate": 2e-3},
+        max_iter=3,
+        budget=10,
+    )
+    assert search_state.init_config and search_state.init_config[0].get("learning_rate") != 2e-3
 
     """
         test starting_points located outside of the search space, and custom_hp is set
@@ -52,39 +75,60 @@ def test_starting_point_not_in_search_space():
 
     from flaml import tune
 
-    X_train, y_train, X_val, y_val, _ = get_toy_data_seqclassification()
+    X_train, y_train, _, _, _ = get_toy_data_seqclassification()
 
     this_estimator_name = "transformer_ms"
-    automl = AutoML()
-    automl_settings = get_automl_settings(estimator_name=this_estimator_name)
+    task = task_factory("seq-classification", X_train, y_train)
+    estimator_class = task.estimator_class_from_str(this_estimator_name)
+    estimator_class.init()
 
-    automl_settings["custom_hp"] = {
-        this_estimator_name: {
-            "model_path": {
-                "domain": "albert-base-v2",
-            },
-            "learning_rate": {
-                "domain": tune.choice([1e-4, 1e-5]),
-            },
-            "per_device_train_batch_size": {
-                "domain": 2,
-            },
-        }
+    custom_hp = {
+        "model_path": {
+            "domain": "albert-base-v2",
+        },
+        "learning_rate": {
+            "domain": tune.choice([1e-4, 1e-5]),
+        },
+        "per_device_train_batch_size": {
+            "domain": 2,
+        },
     }
-    automl_settings["starting_points"] = "data:test/nlp/default/"
 
-    automl.fit(X_train, y_train, **automl_settings)
-    assert len(automl._search_states[this_estimator_name].init_config[0]) == len(
-        automl._search_states[this_estimator_name]._search_space_domain
-    ) - len(automl_settings["custom_hp"][this_estimator_name]), (
+    # Simulate a suggested starting point (e.g. from portfolio) which becomes invalid
+    # after custom_hp constrains the space.
+    invalid_starting_points = [
+        {
+            "learning_rate": 1e-5,
+            "num_train_epochs": 1.0,
+            "per_device_train_batch_size": 8,
+            "seed": 43,
+            "global_max_steps": 100,
+            "model_path": "google/electra-base-discriminator",
+        }
+    ]
+
+    search_state = SearchState(
+        learner_class=estimator_class,
+        data=X_train,
+        task=task,
+        starting_point=invalid_starting_points,
+        custom_hp=custom_hp,
+        max_iter=3,
+        budget=10,
+    )
+
+    assert search_state.init_config, "Expected a non-empty init_config list"
+    init_config0 = search_state.init_config[0]
+    assert init_config0 is not None
+    assert len(init_config0) == len(search_state._search_space_domain) - len(custom_hp), (
         "The search space is updated with the custom_hp on {} hyperparameters of "
         "the specified estimator without an initial value. Thus a valid init config "
         "should only contain the cardinality of the search space minus {}".format(
-            len(automl_settings["custom_hp"][this_estimator_name]),
-            len(automl_settings["custom_hp"][this_estimator_name]),
+            len(custom_hp),
+            len(custom_hp),
         )
     )
-    assert automl._search_states[this_estimator_name].search_space["model_path"] == "albert-base-v2"
+    assert search_state.search_space["model_path"] == "albert-base-v2"
 
     if os.path.exists("test/data/output/"):
         try:
@@ -93,7 +137,6 @@ def test_starting_point_not_in_search_space():
             print("PermissionError when deleting test/data/output/")
 
 
-@pytest.mark.skipif(sys.platform == "win32", reason="do not run on windows")
 def test_points_to_evaluate():
     from flaml import AutoML
 
@@ -106,7 +149,13 @@ def test_points_to_evaluate():
 
     automl_settings["custom_hp"] = {"transformer_ms": {"model_path": {"domain": "google/electra-small-discriminator"}}}
 
-    automl.fit(X_train, y_train, **automl_settings)
+    try:
+        automl.fit(X_train, y_train, **automl_settings)
+    except OSError as e:
+        message = str(e)
+        if "Too Many Requests" in message or "rate limit" in message.lower():
+            pytest.skip(f"Skipping HF model load/training: {message}")
+        raise
 
     if os.path.exists("test/data/output/"):
         try:
@@ -116,7 +165,6 @@ def test_points_to_evaluate():
 
 
 # TODO: implement _test_zero_shot_model
-@pytest.mark.skipif(sys.platform == "win32", reason="do not run on windows")
 def test_zero_shot_nomodel():
     from flaml.default import preprocess_and_suggest_hyperparams
 
@@ -141,7 +189,14 @@ def test_zero_shot_nomodel():
     fit_kwargs = automl_settings.pop("fit_kwargs_by_estimator", {}).get(estimator_name)
     fit_kwargs.update(automl_settings)
     pop_args(fit_kwargs)
-    model.fit(X_train, y_train, **fit_kwargs)
+
+    try:
+        model.fit(X_train, y_train, **fit_kwargs)
+    except OSError as e:
+        message = str(e)
+        if "Too Many Requests" in message or "rate limit" in message.lower():
+            pytest.skip(f"Skipping HF model load/training: {message}")
+        raise
 
     if os.path.exists("test/data/output/"):
         try:

test/object_store.py

@@ -7,7 +7,7 @@ from sklearn.model_selection import train_test_split
 from flaml import tune
 from flaml.automl.model import LGBMEstimator
 
-data = fetch_california_housing(return_X_y=False, as_frame=True)
+data = fetch_california_housing(return_X_y=False, as_frame=True, data_home="test")
 X, y = data.data, data.target
 X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
 X_train_ref = ray.put(X_train)

test/reg.py

@@ -11,7 +11,7 @@ automl_settings = {
     "task": "regression",
     "log_file_name": "test/california.log",
 }
-X_train, y_train = fetch_california_housing(return_X_y=True)
+X_train, y_train = fetch_california_housing(return_X_y=True, data_home="test")
 # Train with labeled input data
 automl.fit(X_train=X_train, y_train=y_train, **automl_settings)
 print(automl.model)

test/spark/test_0sparkml.py

@@ -1,3 +1,4 @@
+import atexit
 import os
 import sys
 import warnings
@@ -10,6 +11,7 @@ from packaging.version import Version
 
 from flaml import AutoML
 from flaml.automl.data import auto_convert_dtypes_pandas, auto_convert_dtypes_spark, get_random_dataframe
+from flaml.automl.spark import disable_spark_ansi_mode, restore_spark_ansi_mode
 from flaml.tune.spark.utils import check_spark
 
 warnings.simplefilter(action="ignore")
@@ -29,7 +31,7 @@ else:
             .config(
                 "spark.jars.packages",
                 (
-                    "com.microsoft.azure:synapseml_2.12:1.0.4,"
+                    "com.microsoft.azure:synapseml_2.12:1.1.0,"
                     "org.apache.hadoop:hadoop-azure:3.3.5,"
                     "com.microsoft.azure:azure-storage:8.6.6,"
                     f"org.mlflow:mlflow-spark_2.12:{mlflow.__version__}"
@@ -55,6 +57,9 @@ else:
     except ImportError:
         skip_spark = True
 
+spark, ansi_conf, adjusted = disable_spark_ansi_mode()
+atexit.register(restore_spark_ansi_mode, spark, ansi_conf, adjusted)
+
 if sys.version_info >= (3, 11):
     skip_py311 = True
 else:
@@ -64,6 +69,13 @@ pytestmark = [pytest.mark.skipif(skip_spark, reason="Spark is not installed. Ski
 
 
 def _test_spark_synapseml_lightgbm(spark=None, task="classification"):
+    # TODO: remove the estimator assignment once SynapseML supports spark 4+.
+    from flaml.automl.spark.utils import _spark_major_minor_version
+
+    if _spark_major_minor_version[0] >= 4:
+        # skip synapseml lightgbm test for spark 4+
+        return
+
     if task == "classification":
         metric = "accuracy"
         X_train, y_train = skds.load_iris(return_X_y=True, as_frame=True)
@@ -154,26 +166,31 @@ def test_spark_synapseml_rank():
 
 
 def test_spark_input_df():
-    df = (
-        spark.read.format("csv")
-        .option("header", True)
-        .option("inferSchema", True)
-        .load("wasbs://publicwasb@mmlspark.blob.core.windows.net/company_bankruptcy_prediction_data.csv")
-    )
+    import pandas as pd
+
+    file_url = "https://mmlspark.blob.core.windows.net/publicwasb/company_bankruptcy_prediction_data.csv"
+    df = pd.read_csv(file_url)
+    df = spark.createDataFrame(df)
     train, test = df.randomSplit([0.8, 0.2], seed=1)
     feature_cols = df.columns[1:]
     featurizer = VectorAssembler(inputCols=feature_cols, outputCol="features")
     train_data = featurizer.transform(train)["Bankrupt?", "features"]
     test_data = featurizer.transform(test)["Bankrupt?", "features"]
     automl = AutoML()
+
+    # TODO: remove the estimator assignment once SynapseML supports spark 4+.
+    from flaml.automl.spark.utils import _spark_major_minor_version
+
+    estimator_list = ["rf_spark"] if _spark_major_minor_version[0] >= 4 else None
+
     settings = {
         "time_budget": 30,  # total running time in seconds
         "metric": "roc_auc",
-        # "estimator_list": ["lgbm_spark"],  # list of ML learners; we tune lightgbm in this example
         "task": "classification",  # task type
         "log_file_name": "flaml_experiment.log",  # flaml log file
         "seed": 7654321,  # random seed
         "eval_method": "holdout",
+        "estimator_list": estimator_list,  # TODO: remove once SynapseML supports spark 4+
     }
     df = to_pandas_on_spark(to_pandas_on_spark(train_data).to_spark(index_col="index"))
 
@@ -184,6 +201,9 @@ def test_spark_input_df():
         **settings,
     )
 
+    if estimator_list == ["rf_spark"]:
+        return
+
     try:
         model = automl.model.estimator
         predictions = model.transform(test_data)

test/spark/test_exceptions.py

@@ -22,7 +22,7 @@ def base_automl(n_concurrent_trials=1, use_ray=False, use_spark=False, verbose=0
     except (ServerError, Exception):
         from sklearn.datasets import fetch_california_housing
 
-        X_train, y_train = fetch_california_housing(return_X_y=True)
+        X_train, y_train = fetch_california_housing(return_X_y=True, data_home="test")
     automl = AutoML()
     settings = {
         "time_budget": 3,  # total running time in seconds

test/spark/test_mlflow.py

@@ -1,3 +1,4 @@
+import atexit
 import importlib
 import os
 import sys
@@ -13,6 +14,7 @@ from sklearn.metrics import r2_score
 from sklearn.model_selection import train_test_split
 
 import flaml
+from flaml.automl.spark import disable_spark_ansi_mode, restore_spark_ansi_mode
 from flaml.automl.spark.utils import to_pandas_on_spark
 
 try:
@@ -120,6 +122,29 @@ def _check_mlflow_logging(possible_num_runs, metric, is_parent_run, experiment_i
     # mlflow.delete_experiment(experiment_id)
 
 
+@pytest.mark.skipif(skip_spark, reason="Spark is not installed. Skip all spark tests.")
+def test_automl_nonsparkdata_noautolog_noparentrun():
+    experiment_id = _test_automl_nonsparkdata(is_autolog=False, is_parent_run=False)
+    _check_mlflow_logging(0, "r2", False, experiment_id, is_automl=True)  # no logging
+
+
+@pytest.mark.skipif(skip_spark, reason="Spark is not installed. Skip all spark tests.")
+def test_automl_sparkdata_noautolog_noparentrun():
+    experiment_id = _test_automl_sparkdata(is_autolog=False, is_parent_run=False)
+    _check_mlflow_logging(0, "mse", False, experiment_id, is_automl=True)  # no logging
+
+
+@pytest.mark.skipif(skip_spark, reason="Spark is not installed. Skip all spark tests.")
+def test_tune_noautolog_noparentrun_parallel():
+    experiment_id = _test_tune(is_autolog=False, is_parent_run=False, is_parallel=True)
+    _check_mlflow_logging(0, "r2", False, experiment_id)
+
+
+def test_tune_noautolog_noparentrun_nonparallel():
+    experiment_id = _test_tune(is_autolog=False, is_parent_run=False, is_parallel=False)
+    _check_mlflow_logging(3, "r2", False, experiment_id, skip_tags=True)
+
+
 @pytest.mark.skipif(skip_spark, reason="Spark is not installed. Skip all spark tests.")
 def test_tune_autolog_parentrun_parallel():
     experiment_id = _test_tune(is_autolog=True, is_parent_run=True, is_parallel=True)
@@ -131,6 +156,16 @@ def test_tune_autolog_parentrun_nonparallel():
     _check_mlflow_logging(3, "r2", True, experiment_id)
 
 
+def test_tune_autolog_noparentrun_nonparallel():
+    experiment_id = _test_tune(is_autolog=True, is_parent_run=False, is_parallel=False)
+    _check_mlflow_logging(3, "r2", False, experiment_id)
+
+
+def test_tune_noautolog_parentrun_nonparallel():
+    experiment_id = _test_tune(is_autolog=False, is_parent_run=True, is_parallel=False)
+    _check_mlflow_logging(3, "r2", True, experiment_id)
+
+
 @pytest.mark.skipif(skip_spark, reason="Spark is not installed. Skip all spark tests.")
 def test_tune_autolog_noparentrun_parallel():
     experiment_id = _test_tune(is_autolog=True, is_parent_run=False, is_parallel=True)
@@ -143,28 +178,12 @@ def test_tune_noautolog_parentrun_parallel():
     _check_mlflow_logging([4, 3], "r2", True, experiment_id)
 
 
-def test_tune_autolog_noparentrun_nonparallel():
-    experiment_id = _test_tune(is_autolog=True, is_parent_run=False, is_parallel=False)
-    _check_mlflow_logging(3, "r2", False, experiment_id)
-
-
-def test_tune_noautolog_parentrun_nonparallel():
-    experiment_id = _test_tune(is_autolog=False, is_parent_run=True, is_parallel=False)
-    _check_mlflow_logging(3, "r2", True, experiment_id)
-
-
-@pytest.mark.skipif(skip_spark, reason="Spark is not installed. Skip all spark tests.")
-def test_tune_noautolog_noparentrun_parallel():
-    experiment_id = _test_tune(is_autolog=False, is_parent_run=False, is_parallel=True)
-    _check_mlflow_logging(0, "r2", False, experiment_id)
-
-
-def test_tune_noautolog_noparentrun_nonparallel():
-    experiment_id = _test_tune(is_autolog=False, is_parent_run=False, is_parallel=False)
-    _check_mlflow_logging(3, "r2", False, experiment_id, skip_tags=True)
-
-
 def _test_automl_sparkdata(is_autolog, is_parent_run):
+    # TODO: remove the estimator assignment once SynapseML supports spark 4+.
+    from flaml.automl.spark.utils import _spark_major_minor_version
+
+    estimator_list = ["rf_spark"] if _spark_major_minor_version[0] >= 4 else None
+
     mlflow.end_run()
     mlflow_exp_name = f"test_mlflow_integration_{int(time.time())}"
     mlflow_experiment = mlflow.set_experiment(mlflow_exp_name)
@@ -175,6 +194,9 @@ def _test_automl_sparkdata(is_autolog, is_parent_run):
     if is_parent_run:
         mlflow.start_run(run_name=f"automl_sparkdata_autolog_{is_autolog}")
     spark = pyspark.sql.SparkSession.builder.getOrCreate()
+    spark, ansi_conf, adjusted = disable_spark_ansi_mode()
+    atexit.register(restore_spark_ansi_mode, spark, ansi_conf, adjusted)
+
     pd_df = load_diabetes(as_frame=True).frame
     df = spark.createDataFrame(pd_df)
     df = df.repartition(4).cache()
@@ -193,6 +215,7 @@ def _test_automl_sparkdata(is_autolog, is_parent_run):
         "log_type": "all",
         "n_splits": 2,
         "model_history": True,
+        "estimator_list": estimator_list,
     }
     df = to_pandas_on_spark(to_pandas_on_spark(train_data).to_spark(index_col="index"))
     automl.fit(
@@ -252,12 +275,6 @@ def test_automl_sparkdata_noautolog_parentrun():
     _check_mlflow_logging(3, "mse", True, experiment_id, is_automl=True)
 
 
-@pytest.mark.skipif(skip_spark, reason="Spark is not installed. Skip all spark tests.")
-def test_automl_sparkdata_noautolog_noparentrun():
-    experiment_id = _test_automl_sparkdata(is_autolog=False, is_parent_run=False)
-    _check_mlflow_logging(0, "mse", False, experiment_id, is_automl=True)  # no logging
-
-
 @pytest.mark.skipif(skip_spark, reason="Spark is not installed. Skip all spark tests.")
 def test_automl_nonsparkdata_autolog_parentrun():
     experiment_id = _test_automl_nonsparkdata(is_autolog=True, is_parent_run=True)
@@ -276,12 +293,6 @@ def test_automl_nonsparkdata_noautolog_parentrun():
     _check_mlflow_logging([4, 3], "r2", True, experiment_id, is_automl=True)
 
 
-@pytest.mark.skipif(skip_spark, reason="Spark is not installed. Skip all spark tests.")
-def test_automl_nonsparkdata_noautolog_noparentrun():
-    experiment_id = _test_automl_nonsparkdata(is_autolog=False, is_parent_run=False)
-    _check_mlflow_logging(0, "r2", False, experiment_id, is_automl=True)  # no logging
-
-
 @pytest.mark.skipif(skip_spark, reason="Spark is not installed. Skip all spark tests.")
 def test_exit_pyspark_autolog():
     import pyspark
@@ -319,6 +330,9 @@ def _init_spark_for_main():
         "https://mmlspark.blob.core.windows.net/publicwasb/log_model_allowlist.txt",
     )
 
+    spark, ansi_conf, adjusted = disable_spark_ansi_mode()
+    atexit.register(restore_spark_ansi_mode, spark, ansi_conf, adjusted)
+
 
 if __name__ == "__main__":
     _init_spark_for_main()

test/spark/test_performance.py

@@ -2,8 +2,23 @@ import os
 import sys
 
 import pytest
-from minio.error import ServerError
-from openml.exceptions import OpenMLServerException
+
+try:
+    from minio.error import ServerError
+except ImportError:
+
+    class ServerError(Exception):
+        pass
+
+
+try:
+    from openml.exceptions import OpenMLServerException
+except ImportError:
+
+    class OpenMLServerException(Exception):
+        pass
+
+
 from requests.exceptions import ChunkedEncodingError, SSLError
 
 from flaml.tune.spark.utils import check_spark
@@ -16,14 +31,14 @@ pytestmark = [pytest.mark.skipif(skip_spark, reason="Spark is not installed. Ski
 os.environ["FLAML_MAX_CONCURRENT"] = "2"
 
 
-def run_automl(budget=3, dataset_format="dataframe", hpo_method=None):
+def run_automl(budget=30, dataset_format="dataframe", hpo_method=None):
     import urllib3
 
     from flaml.automl.data import load_openml_dataset
 
     performance_check_budget = 3600
     if sys.platform == "darwin" or "nt" in os.name or "3.10" not in sys.version:
-        budget = 3  # revise the buget if the platform is not linux + python 3.10
+        budget = 30  # revise the buget if the platform is not linux + python 3.10
     if budget >= performance_check_budget:
         max_iter = 60
         performance_check_budget = None
@@ -76,6 +91,11 @@ def run_automl(budget=3, dataset_format="dataframe", hpo_method=None):
     print("Best ML leaner:", automl.best_estimator)
     print("Best hyperparmeter config:", automl.best_config)
     print(f"Best accuracy on validation data: {1 - automl.best_loss:.4g}")
+    if performance_check_budget is not None and automl.best_estimator is None:
+        # skip the performance check if no model is trained
+        # this happens sometimes in github actions ubuntu python 3.12 environment
+        print("Warning: no model is trained, skip performance check")
+        return
     print(f"Training duration of best run: {automl.best_config_train_time:.4g} s")
     print(automl.model.estimator)
     print(automl.best_config_per_estimator)

test/spark/test_utils.py

@@ -1,3 +1,4 @@
+import atexit
 import os
 from functools import partial
 from timeit import timeit
@@ -14,6 +15,7 @@ try:
     from pyspark.sql import SparkSession
 
     from flaml.automl.ml import sklearn_metric_loss_score
+    from flaml.automl.spark import disable_spark_ansi_mode, restore_spark_ansi_mode
     from flaml.automl.spark.metrics import spark_metric_loss_score
     from flaml.automl.spark.utils import (
         iloc_pandas_on_spark,
@@ -24,6 +26,7 @@ try:
         unique_value_first_index,
     )
     from flaml.tune.spark.utils import (
+        _spark_major_minor_version,
         check_spark,
         get_broadcast_data,
         get_n_cpus,
@@ -35,10 +38,41 @@ try:
 except ImportError:
     print("Spark is not installed. Skip all spark tests.")
     skip_spark = True
+    _spark_major_minor_version = (0, 0)
+
 
 pytestmark = [pytest.mark.skipif(skip_spark, reason="Spark is not installed. Skip all spark tests."), pytest.mark.spark]
 
 
+@pytest.mark.skipif(_spark_major_minor_version[0] < 4, reason="Requires Spark 4.0+")
+def test_to_pandas_on_spark_temp_override():
+    import pyspark.pandas as ps
+    from pyspark.sql import Row
+
+    from flaml.automl.spark.utils import to_pandas_on_spark
+
+    spark_session = SparkSession.builder.getOrCreate()
+    spark, ansi_conf, adjusted = disable_spark_ansi_mode()
+    atexit.register(restore_spark_ansi_mode, spark, ansi_conf, adjusted)
+
+    # Ensure we can toggle options
+    orig = ps.get_option("compute.fail_on_ansi_mode")
+
+    try:
+        spark_session.conf.set("spark.sql.ansi.enabled", "true")
+        ps.set_option("compute.fail_on_ansi_mode", True)
+
+        # create tiny spark df
+        sdf = spark_session.createDataFrame([Row(a=1, b=2)])
+        # Should not raise as our function temporarily disables fail_on_ansi_mode
+        pds = to_pandas_on_spark(sdf)
+        assert "a" in pds.columns
+    finally:
+        # restore test environment
+        ps.set_option("compute.fail_on_ansi_mode", orig)
+        spark_session.conf.set("spark.sql.ansi.enabled", "false")
+
+
 def test_with_parameters_spark():
     def train(config, data=None):
         if isinstance(data, pyspark.broadcast.Broadcast):

test/test_autovw.py

@@ -5,17 +5,38 @@ import sys
 import unittest
 
 import numpy as np
-import openml
+
+try:
+    import openml
+except ImportError:
+    openml = None
 import pandas as pd
 import pytest
 import scipy.sparse
-from minio.error import ServerError
+
+try:
+    from minio.error import ServerError
+except ImportError:
+
+    class ServerError(Exception):
+        pass
+
+
 from requests.exceptions import SSLError
 from sklearn.metrics import mean_absolute_error, mean_squared_error
 
 from flaml import AutoVW
 from flaml.tune import loguniform, polynomial_expansion_set
 
+try:
+    from vowpalwabbit import pyvw
+except ImportError:
+    skip_vw_test = True
+else:
+    skip_vw_test = False
+
+pytest.skip("skipping if no openml", allow_module_level=True) if openml is None else None
+
 VW_DS_DIR = "test/data/"
 NS_LIST = list(string.ascii_lowercase) + list(string.ascii_uppercase)
 logger = logging.getLogger(__name__)
@@ -351,14 +372,9 @@ def get_vw_tuning_problem(tuning_hp="NamesapceInteraction"):
     return vw_oml_problem_args, vw_online_aml_problem
 
 
-@pytest.mark.skipif(
-    "3.10" in sys.version or "3.11" in sys.version,
-    reason="do not run on py >= 3.10",
-)
+@pytest.mark.skipif(skip_vw_test, reason="vowpalwabbit not installed")
 class TestAutoVW(unittest.TestCase):
     def test_vw_oml_problem_and_vanilla_vw(self):
-        from vowpalwabbit import pyvw
-
         try:
             vw_oml_problem_args, vw_online_aml_problem = get_vw_tuning_problem()
         except (SSLError, ServerError, Exception) as e:

test/tune/test_lexiflow.py

@@ -4,10 +4,17 @@ from collections import defaultdict
 
 import numpy as np
 import pytest
-import thop
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
+
+try:
+    import thop
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+except ImportError:
+    thop = None
+    torch = None
+    nn = None
+    F = None
 
 try:
     import torchvision
@@ -16,6 +23,11 @@ except ImportError:
 
 from flaml import tune
 
+if thop is None or torch is None or nn is None or F is None or torchvision is None:
+    pytest.skip(
+        "skipping test_lexiflow.py because torch, torchvision or thop is not installed.", allow_module_level=True
+    )
+
 DEVICE = torch.device("cpu")
 BATCHSIZE = 128
 N_TRAIN_EXAMPLES = BATCHSIZE * 30

test/tune/test_tune.py

@@ -53,6 +53,11 @@ def _easy_objective(config):
 
 
 def test_nested_run():
+    """
+    nested tuning example: Tune -> AutoML -> MLflow autolog
+    mlflow logging is complicated in nested tuning. It's better to turn off mlflow autologging to avoid
+    potential issues in FLAML's mlflow_integration.adopt_children() function.
+    """
     from flaml import AutoML, tune
 
     data, labels = sklearn.datasets.load_breast_cancer(return_X_y=True)

test/tune_example.py

@@ -6,12 +6,12 @@ from sklearn.model_selection import train_test_split
 from flaml import tune
 from flaml.automl.model import LGBMEstimator
 
-data = fetch_california_housing(return_X_y=False, as_frame=True)
+data = fetch_california_housing(return_X_y=False, as_frame=True, data_home="test")
 df, X, y = data.frame, data.data, data.target
 df_train, _, X_train, X_test, _, y_test = train_test_split(df, X, y, test_size=0.33, random_state=42)
 csv_file_name = "test/housing.csv"
 df_train.to_csv(csv_file_name, index=False)
-# X, y = fetch_california_housing(return_X_y=True, as_frame=True)
+# X, y = fetch_california_housing(return_X_y=True, as_frame=True, data_home="test")
 # X_train, X_test, y_train, y_test = train_test_split(
 #     X, y, test_size=0.33, random_state=42
 # )

website/docusaurus.config.js

@@ -99,6 +99,12 @@ module.exports = {
             'https://github.com/microsoft/FLAML/edit/main/website/',
           remarkPlugins: [math],
           rehypePlugins: [katex],
+          // Allow __init__.md and other underscore-prefixed markdown docs
+          exclude: [
+            '**/_*.{js,jsx,ts,tsx}',
+            '**/*.test.{js,jsx,ts,tsx}',
+            '**/__tests__/**',
+          ],
         },
         theme: {
           customCss: require.resolve('./src/css/custom.css'),

website/yarn.lock

@@ -2634,9 +2634,9 @@ ajv@^8.0.0, ajv@^8.8.0:
     uri-js "^4.2.2"
 
 algoliasearch-helper@^3.5.5:
-  version "3.11.1"
-  resolved "https://registry.npmmirror.com/algoliasearch-helper/-/algoliasearch-helper-3.11.1.tgz#d83ab7f1a2a374440686ef7a144b3c288b01188a"
-  integrity sha512-mvsPN3eK4E0bZG0/WlWJjeqe/bUD2KOEVOl0GyL/TGXn6wcpZU8NOuztGHCUKXkyg5gq6YzUakVTmnmSSO5Yiw==
+  version "3.26.0"
+  resolved "https://registry.yarnpkg.com/algoliasearch-helper/-/algoliasearch-helper-3.26.0.tgz#d6e283396a9fc5bf944f365dc3b712570314363f"
+  integrity sha512-Rv2x3GXleQ3ygwhkhJubhhYGsICmShLAiqtUuJTUkr9uOCOXyF2E71LVT4XDnVffbknv8XgScP4U0Oxtgm+hIw==
   dependencies:
     "@algolia/events" "^4.0.1"