Install wheel and setuptools (#1443)

* Fix CD actions

* Skip Build if no relevant changes

.github/ISSUE_TEMPLATE.md

@@ -0,0 +1,73 @@
+### Description
+
+<!-- A clear and concise description of the issue or feature request. -->
+
+### Environment
+
+- FLAML version: <!-- Specify the FLAML version (e.g., v0.2.0) -->
+- Python version: <!-- Specify the Python version (e.g., 3.8) -->
+- Operating System: <!-- Specify the OS (e.g., Windows 10, Ubuntu 20.04) -->
+
+### Steps to Reproduce (for bugs)
+
+<!-- Provide detailed steps to reproduce the issue. Include code snippets, configuration files, or any other relevant information. -->
+
+1. Step 1
+1. Step 2
+1. ...
+
+### Expected Behavior
+
+<!-- Describe what you expected to happen. -->
+
+### Actual Behavior
+
+<!-- Describe what actually happened. Include any error messages, stack traces, or unexpected behavior. -->
+
+### Screenshots / Logs (if applicable)
+
+<!-- If relevant, include screenshots or logs that help illustrate the issue. -->
+
+### Additional Information
+
+<!-- Include any additional information that might be helpful, such as specific configurations, data samples, or context about the environment. -->
+
+### Possible Solution (if you have one)
+
+<!-- If you have suggestions on how to address the issue, provide them here. -->
+
+### Is this a Bug or Feature Request?
+
+<!-- Choose one: Bug | Feature Request -->
+
+### Priority
+
+<!-- Choose one: High | Medium | Low -->
+
+### Difficulty
+
+<!-- Choose one: Easy | Moderate | Hard -->
+
+### Any related issues?
+
+<!-- If this is related to another issue, reference it here. -->
+
+### Any relevant discussions?
+
+<!-- If there are any discussions or forum threads related to this issue, provide links. -->
+
+### Checklist
+
+<!-- Please check the items that you have completed -->
+
+- [ ] I have searched for similar issues and didn't find any duplicates.
+- [ ] I have provided a clear and concise description of the issue.
+- [ ] I have included the necessary environment details.
+- [ ] I have outlined the steps to reproduce the issue.
+- [ ] I have included any relevant logs or screenshots.
+- [ ] I have indicated whether this is a bug or a feature request.
+- [ ] I have set the priority and difficulty levels.
+
+### Additional Comments
+
+<!-- Any additional comments or context that you think would be helpful. -->

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -0,0 +1,53 @@
+name: Bug Report
+description: File a bug report
+title: "[Bug]: "
+labels: ["bug"]
+
+body:
+  - type: textarea
+    id: description
+    attributes:
+      label: Describe the bug
+      description: A clear and concise description of what the bug is.
+      placeholder: What went wrong?
+  - type: textarea
+    id: reproduce
+    attributes:
+      label: Steps to reproduce
+      description: |
+        Steps to reproduce the behavior:
+
+        1. Step 1
+        2. Step 2
+        3. ...
+        4. See error
+      placeholder: How can we replicate the issue?
+  - type: textarea
+    id: modelused
+    attributes:
+      label: Model Used
+      description: A description of the model that was used when the error was encountered
+      placeholder: gpt-4, mistral-7B etc
+  - type: textarea
+    id: expected_behavior
+    attributes:
+      label: Expected Behavior
+      description: A clear and concise description of what you expected to happen.
+      placeholder: What should have happened?
+  - type: textarea
+    id: screenshots
+    attributes:
+      label: Screenshots and logs
+      description: If applicable, add screenshots and logs to help explain your problem.
+      placeholder: Add screenshots here
+  - type: textarea
+    id: additional_information
+    attributes:
+      label: Additional Information
+      description: |
+        - FLAML Version: <!-- Specify the FLAML version (e.g., v0.2.0) -->
+        - Operating System: <!-- Specify the OS (e.g., Windows 10, Ubuntu 20.04) -->
+        - Python Version: <!-- Specify the Python version (e.g., 3.8) -->
+        - Related Issues: <!-- Link to any related issues here (e.g., #1) -->
+        - Any other relevant information.
+      placeholder: Any additional details

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1 @@
+blank_issues_enabled: true

.github/ISSUE_TEMPLATE/feature_request.yml

@@ -0,0 +1,26 @@
+name: Feature Request
+description: File a feature request
+labels: ["enhancement"]
+title: "[Feature Request]: "
+
+body:
+  - type: textarea
+    id: problem_description
+    attributes:
+      label: Is your feature request related to a problem? Please describe.
+      description: A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+      placeholder: What problem are you trying to solve?
+
+  - type: textarea
+    id: solution_description
+    attributes:
+      label: Describe the solution you'd like
+      description: A clear and concise description of what you want to happen.
+      placeholder: How do you envision the solution?
+
+  - type: textarea
+    id: additional_context
+    attributes:
+      label: Additional context
+      description: Add any other context or screenshots about the feature request here.
+      placeholder: Any additional information

.github/ISSUE_TEMPLATE/general_issue.yml

@@ -0,0 +1,41 @@
+name: General Issue
+description: File a general issue
+title: "[Issue]: "
+labels: []
+
+body:
+  - type: textarea
+    id: description
+    attributes:
+      label: Describe the issue
+      description: A clear and concise description of what the issue is.
+      placeholder: What went wrong?
+  - type: textarea
+    id: reproduce
+    attributes:
+      label: Steps to reproduce
+      description: |
+        Steps to reproduce the behavior:
+
+        1. Step 1
+        2. Step 2
+        3. ...
+        4. See error
+      placeholder: How can we replicate the issue?
+  - type: textarea
+    id: screenshots
+    attributes:
+      label: Screenshots and logs
+      description: If applicable, add screenshots and logs to help explain your problem.
+      placeholder: Add screenshots here
+  - type: textarea
+    id: additional_information
+    attributes:
+      label: Additional Information
+      description: |
+        - FLAML Version: <!-- Specify the FLAML version (e.g., v0.2.0) -->
+        - Operating System: <!-- Specify the OS (e.g., Windows 10, Ubuntu 20.04) -->
+        - Python Version: <!-- Specify the Python version (e.g., 3.8) -->
+        - Related Issues: <!-- Link to any related issues here (e.g., #1) -->
+        - Any other relevant information.
+      placeholder: Any additional details

.github/PULL_REQUEST_TEMPLATE.md

@@ -12,7 +12,7 @@
 
 ## Checks
 
-<!-- - I've used [pre-commit](https://microsoft.github.io/FLAML/docs/Contribute#pre-commit) to lint the changes in this PR (note the same in integrated in our CI checks). -->
+- [ ] I've used [pre-commit](https://microsoft.github.io/FLAML/docs/Contribute#pre-commit) to lint the changes in this PR (note the same in integrated in our CI checks).
 - [ ] I've included any doc changes needed for https://microsoft.github.io/FLAML/. See https://microsoft.github.io/FLAML/docs/Contribute#documentation to build and test documentation locally.
 - [ ] I've added tests (if relevant) corresponding to the changes introduced in this PR.
 - [ ] I've made sure all auto checks have passed.

.github/workflows/CD.yml

@@ -12,26 +12,17 @@ jobs:
   deploy:
     strategy:
       matrix:
-        os: ['ubuntu-latest']
-        python-version: [3.8]
+        os: ["ubuntu-latest"]
+        python-version: ["3.10"]
     runs-on: ${{ matrix.os }}
     environment: package
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
-      - name: Cache conda
-        uses: actions/cache@v3
+        uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
         with:
-          path: ~/conda_pkgs_dir
-          key: conda-${{ matrix.os }}-python-${{ matrix.python-version }}-${{ hashFiles('environment.yml') }}
-      - name: Setup Miniconda
-        uses: conda-incubator/setup-miniconda@v2
-        with:
-          auto-update-conda: true
-          auto-activate-base: false
-          activate-environment: hcrystalball
           python-version: ${{ matrix.python-version }}
-          use-only-tar-bz2: true
       - name: Install from source
         # This is required for the pre-commit tests
         shell: pwsh
@@ -42,7 +33,7 @@ jobs:
       - name: Build
         shell: pwsh
         run: |
-          pip install twine
+          pip install twine wheel setuptools
           python setup.py sdist bdist_wheel
       - name: Publish to PyPI
         env:

.github/workflows/deploy-website.yml

@@ -17,6 +17,9 @@ on:
   merge_group:
     types: [checks_requested]
 
+permissions:
+  contents: write
+
 jobs:
   checks:
     if: github.event_name != 'push'
@@ -34,11 +37,11 @@ jobs:
       - name: setup python
         uses: actions/setup-python@v4
         with:
-          python-version: "3.8"
+          python-version: "3.10"
       - name: pydoc-markdown install
         run: |
           python -m pip install --upgrade pip
-          pip install pydoc-markdown==4.5.0
+          pip install pydoc-markdown==4.7.0
       - name: pydoc-markdown run
         run: |
           pydoc-markdown
@@ -70,11 +73,11 @@ jobs:
       - name: setup python
         uses: actions/setup-python@v4
         with:
-          python-version: "3.8"
+          python-version: "3.10"
       - name: pydoc-markdown install
         run: |
           python -m pip install --upgrade pip
-          pip install pydoc-markdown==4.5.0
+          pip install pydoc-markdown==4.7.0
       - name: pydoc-markdown run
         run: |
           pydoc-markdown

.github/workflows/openai.yml

@@ -13,6 +13,8 @@ on:
       - 'notebook/autogen_chatgpt_gpt4.ipynb'
       - '.github/workflows/openai.yml'
 
+permissions: {}
+
 jobs:
   test:
     strategy:

.github/workflows/pre-commit.yml

@@ -10,6 +10,7 @@ defaults:
   run:
     shell: bash
 
+permissions: {}
 jobs:
 
   pre-commit-check:

.github/workflows/python-package.yml

@@ -14,9 +14,16 @@ on:
       - 'setup.py'
   pull_request:
     branches: ['main']
+    paths:
+      - 'flaml/**'
+      - 'test/**'
+      - 'notebook/**'
+      - '.github/workflows/python-package.yml'
+      - 'setup.py'
   merge_group:
     types: [checks_requested]
 
+permissions: {}
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref }}
   cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
@@ -29,19 +36,17 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-latest, windows-2019]
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.9", "3.10", "3.11"]
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
-      - name: On mac + python 3.10, install libomp to facilitate lgbm and xgboost install
-        if: matrix.os == 'macOS-latest' && matrix.python-version == '3.10'
+      - name: On mac, install libomp to facilitate lgbm and xgboost install
+        if: matrix.os == 'macOS-latest'
         run: |
-          # remove libomp version constraint after xgboost works with libomp>11.1.0 on python 3.10
-          wget https://raw.githubusercontent.com/Homebrew/homebrew-core/679923b4eb48a8dc7ecc1f05d06063cd79b3fc00/Formula/libomp.rb -O $(find $(brew --repository) -name libomp.rb)
-          brew unlink libomp
+          brew update
           brew install libomp
           export CC=/usr/bin/clang
           export CXX=/usr/bin/clang++
@@ -51,40 +56,42 @@ jobs:
           export LDFLAGS="$LDFLAGS -Wl,-rpath,/usr/local/opt/libomp/lib -L/usr/local/opt/libomp/lib -lomp"
       - name: Install packages and dependencies
         run: |
-          python -m pip install --upgrade pip wheel
+          python -m pip install --upgrade pip wheel setuptools
           pip install -e .
           python -c "import flaml"
           pip install -e .[test]
-      - name: On Ubuntu python 3.8, install pyspark 3.2.3
-        if: matrix.python-version == '3.8' && matrix.os == 'ubuntu-latest'
+      - 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.2.3
+          pip install pyspark==3.4.1
           pip list | grep "pyspark"
-      - name: If linux, install ray 2
-        if: matrix.os == 'ubuntu-latest'
+      - 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'
         run: |
           pip install "ray[tune]<2.5.0"
-      - name: If mac, install ray
-        if: matrix.os == 'macOS-latest'
+      - 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]
-      - name: If linux or mac, install prophet on python < 3.9
-        if: (matrix.os == 'macOS-latest' || matrix.os == 'ubuntu-latest') && matrix.python-version != '3.9' && matrix.python-version != '3.10'
+          # 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'
         run: |
           pip install -e .[forecast]
       - name: Install vw on python < 3.10
-        if: matrix.python-version != '3.10'
+        if: matrix.python-version == '3.8' || matrix.python-version == '3.9'
         run: |
           pip install -e .[vw]
-      - name: Uninstall pyspark on (python 3.9) or (python 3.8 + windows)
-        if: matrix.python-version == '3.9' || (matrix.python-version == '3.8' && matrix.os == 'windows-2019')
-        run: |
-          # Uninstall pyspark to test env without pyspark
-          pip uninstall -y pyspark
       - name: Test with pytest
         if: matrix.python-version != '3.10'
         run: |
-          pytest test
+          pytest test/
       - name: Coverage
         if: matrix.python-version == '3.10'
         run: |

.gitignore

@@ -163,5 +163,24 @@ output/
 flaml/tune/spark/mylearner.py
 *.pkl
 
+data/
+benchmark/pmlb/csv_datasets
+benchmark/*.csv
+
+checkpoints/
+test/default
+test/housing.json
+test/nlp/default/transformer_ms/seq-classification.json
+
+flaml/fabric/fanova/_fanova.c
 # local config files
 *.config.local
+
+local_debug/
+patch.diff
+
+# Test things
+notebook/lightning_logs/
+lightning_logs/
+flaml/autogen/extensions/tmp/
+test/autogen/my_tmp/

.pre-commit-config.yaml

@@ -22,10 +22,28 @@ repos:
     - id: trailing-whitespace
     - id: end-of-file-fixer
     - id: no-commit-to-branch
+
+  - repo: https://github.com/asottile/pyupgrade
+    rev: v2.31.1
+    hooks:
+      - id: pyupgrade
+        args: [--py38-plus]
+        name: Upgrade code
+
   - repo: https://github.com/psf/black
     rev: 23.3.0
     hooks:
     - id: black
+
+  - repo: https://github.com/executablebooks/mdformat
+    rev: 0.7.17
+    hooks:
+      - id: mdformat
+        additional_dependencies:
+          - mdformat-gfm
+          - mdformat-black
+          - mdformat_frontmatter
+
   - repo: https://github.com/charliermarsh/ruff-pre-commit
     rev: v0.0.261
     hooks:

Dockerfile

@@ -1,5 +1,5 @@
 # basic setup
-FROM python:3.7
+FROM mcr.microsoft.com/devcontainers/python:3.10
 RUN apt-get update && apt-get -y update
 RUN apt-get install -y sudo git npm
 

NOTICE.md

@@ -1,221 +1,222 @@
-NOTICES
+# NOTICES
 
 This repository incorporates material as listed below or described in the code.
 
-#
 ## Component. Ray.
 
-Code in tune/[analysis.py, sample.py, trial.py, result.py],
-searcher/[suggestion.py, variant_generator.py], and scheduler/trial_scheduler.py is adapted from
+Code in tune/\[analysis.py, sample.py, trial.py, result.py\],
+searcher/\[suggestion.py, variant_generator.py\], and scheduler/trial_scheduler.py is adapted from
 https://github.com/ray-project/ray/blob/master/python/ray/tune/
 
-
-
 ## Open Source License/Copyright Notice.
 
- Apache License
-                           Version 2.0, January 2004
-                        http://www.apache.org/licenses/
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/
 
-   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
 
-   1. Definitions.
+1. Definitions.
 
-      "License" shall mean the terms and conditions for use, reproduction,
-      and distribution as defined by Sections 1 through 9 of this document.
+   "License" shall mean the terms and conditions for use, reproduction,
+   and distribution as defined by Sections 1 through 9 of this document.
 
-      "Licensor" shall mean the copyright owner or entity authorized by
-      the copyright owner that is granting the License.
+   "Licensor" shall mean the copyright owner or entity authorized by
+   the copyright owner that is granting the License.
 
-      "Legal Entity" shall mean the union of the acting entity and all
-      other entities that control, are controlled by, or are under common
-      control with that entity. For the purposes of this definition,
-      "control" means (i) the power, direct or indirect, to cause the
-      direction or management of such entity, whether by contract or
-      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-      outstanding shares, or (iii) beneficial ownership of such entity.
+   "Legal Entity" shall mean the union of the acting entity and all
+   other entities that control, are controlled by, or are under common
+   control with that entity. For the purposes of this definition,
+   "control" means (i) the power, direct or indirect, to cause the
+   direction or management of such entity, whether by contract or
+   otherwise, or (ii) ownership of fifty percent (50%) or more of the
+   outstanding shares, or (iii) beneficial ownership of such entity.
 
-      "You" (or "Your") shall mean an individual or Legal Entity
-      exercising permissions granted by this License.
+   "You" (or "Your") shall mean an individual or Legal Entity
+   exercising permissions granted by this License.
 
-      "Source" form shall mean the preferred form for making modifications,
-      including but not limited to software source code, documentation
-      source, and configuration files.
+   "Source" form shall mean the preferred form for making modifications,
+   including but not limited to software source code, documentation
+   source, and configuration files.
 
-      "Object" form shall mean any form resulting from mechanical
-      transformation or translation of a Source form, including but
-      not limited to compiled object code, generated documentation,
-      and conversions to other media types.
+   "Object" form shall mean any form resulting from mechanical
+   transformation or translation of a Source form, including but
+   not limited to compiled object code, generated documentation,
+   and conversions to other media types.
 
-      "Work" shall mean the work of authorship, whether in Source or
-      Object form, made available under the License, as indicated by a
-      copyright notice that is included in or attached to the work
-      (an example is provided in the Appendix below).
+   "Work" shall mean the work of authorship, whether in Source or
+   Object form, made available under the License, as indicated by a
+   copyright notice that is included in or attached to the work
+   (an example is provided in the Appendix below).
 
-      "Derivative Works" shall mean any work, whether in Source or Object
-      form, that is based on (or derived from) the Work and for which the
-      editorial revisions, annotations, elaborations, or other modifications
-      represent, as a whole, an original work of authorship. For the purposes
-      of this License, Derivative Works shall not include works that remain
-      separable from, or merely link (or bind by name) to the interfaces of,
-      the Work and Derivative Works thereof.
+   "Derivative Works" shall mean any work, whether in Source or Object
+   form, that is based on (or derived from) the Work and for which the
+   editorial revisions, annotations, elaborations, or other modifications
+   represent, as a whole, an original work of authorship. For the purposes
+   of this License, Derivative Works shall not include works that remain
+   separable from, or merely link (or bind by name) to the interfaces of,
+   the Work and Derivative Works thereof.
 
-      "Contribution" shall mean any work of authorship, including
-      the original version of the Work and any modifications or additions
-      to that Work or Derivative Works thereof, that is intentionally
-      submitted to Licensor for inclusion in the Work by the copyright owner
-      or by an individual or Legal Entity authorized to submit on behalf of
-      the copyright owner. For the purposes of this definition, "submitted"
-      means any form of electronic, verbal, or written communication sent
-      to the Licensor or its representatives, including but not limited to
-      communication on electronic mailing lists, source code control systems,
-      and issue tracking systems that are managed by, or on behalf of, the
-      Licensor for the purpose of discussing and improving the Work, but
-      excluding communication that is conspicuously marked or otherwise
-      designated in writing by the copyright owner as "Not a Contribution."
+   "Contribution" shall mean any work of authorship, including
+   the original version of the Work and any modifications or additions
+   to that Work or Derivative Works thereof, that is intentionally
+   submitted to Licensor for inclusion in the Work by the copyright owner
+   or by an individual or Legal Entity authorized to submit on behalf of
+   the copyright owner. For the purposes of this definition, "submitted"
+   means any form of electronic, verbal, or written communication sent
+   to the Licensor or its representatives, including but not limited to
+   communication on electronic mailing lists, source code control systems,
+   and issue tracking systems that are managed by, or on behalf of, the
+   Licensor for the purpose of discussing and improving the Work, but
+   excluding communication that is conspicuously marked or otherwise
+   designated in writing by the copyright owner as "Not a Contribution."
 
-      "Contributor" shall mean Licensor and any individual or Legal Entity
-      on behalf of whom a Contribution has been received by Licensor and
-      subsequently incorporated within the Work.
+   "Contributor" shall mean Licensor and any individual or Legal Entity
+   on behalf of whom a Contribution has been received by Licensor and
+   subsequently incorporated within the Work.
 
-   2. Grant of Copyright License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      copyright license to reproduce, prepare Derivative Works of,
-      publicly display, publicly perform, sublicense, and distribute the
-      Work and such Derivative Works in Source or Object form.
+1. Grant of Copyright License. Subject to the terms and conditions of
+   this License, each Contributor hereby grants to You a perpetual,
+   worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+   copyright license to reproduce, prepare Derivative Works of,
+   publicly display, publicly perform, sublicense, and distribute the
+   Work and such Derivative Works in Source or Object form.
 
-   3. Grant of Patent License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      (except as stated in this section) patent license to make, have made,
-      use, offer to sell, sell, import, and otherwise transfer the Work,
-      where such license applies only to those patent claims licensable
-      by such Contributor that are necessarily infringed by their
-      Contribution(s) alone or by combination of their Contribution(s)
-      with the Work to which such Contribution(s) was submitted. If You
-      institute patent litigation against any entity (including a
-      cross-claim or counterclaim in a lawsuit) alleging that the Work
-      or a Contribution incorporated within the Work constitutes direct
-      or contributory patent infringement, then any patent licenses
-      granted to You under this License for that Work shall terminate
-      as of the date such litigation is filed.
+1. Grant of Patent License. Subject to the terms and conditions of
+   this License, each Contributor hereby grants to You a perpetual,
+   worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+   (except as stated in this section) patent license to make, have made,
+   use, offer to sell, sell, import, and otherwise transfer the Work,
+   where such license applies only to those patent claims licensable
+   by such Contributor that are necessarily infringed by their
+   Contribution(s) alone or by combination of their Contribution(s)
+   with the Work to which such Contribution(s) was submitted. If You
+   institute patent litigation against any entity (including a
+   cross-claim or counterclaim in a lawsuit) alleging that the Work
+   or a Contribution incorporated within the Work constitutes direct
+   or contributory patent infringement, then any patent licenses
+   granted to You under this License for that Work shall terminate
+   as of the date such litigation is filed.
 
-   4. Redistribution. You may reproduce and distribute copies of the
-      Work or Derivative Works thereof in any medium, with or without
-      modifications, and in Source or Object form, provided that You
-      meet the following conditions:
+1. Redistribution. You may reproduce and distribute copies of the
+   Work or Derivative Works thereof in any medium, with or without
+   modifications, and in Source or Object form, provided that You
+   meet the following conditions:
 
-      (a) You must give any other recipients of the Work or
-          Derivative Works a copy of this License; and
+   (a) You must give any other recipients of the Work or
+   Derivative Works a copy of this License; and
 
-      (b) You must cause any modified files to carry prominent notices
-          stating that You changed the files; and
+   (b) You must cause any modified files to carry prominent notices
+   stating that You changed the files; and
 
-      (c) You must retain, in the Source form of any Derivative Works
-          that You distribute, all copyright, patent, trademark, and
-          attribution notices from the Source form of the Work,
-          excluding those notices that do not pertain to any part of
-          the Derivative Works; and
+   (c) You must retain, in the Source form of any Derivative Works
+   that You distribute, all copyright, patent, trademark, and
+   attribution notices from the Source form of the Work,
+   excluding those notices that do not pertain to any part of
+   the Derivative Works; and
 
-      (d) If the Work includes a "NOTICE" text file as part of its
-          distribution, then any Derivative Works that You distribute must
-          include a readable copy of the attribution notices contained
-          within such NOTICE file, excluding those notices that do not
-          pertain to any part of the Derivative Works, in at least one
-          of the following places: within a NOTICE text file distributed
-          as part of the Derivative Works; within the Source form or
-          documentation, if provided along with the Derivative Works; or,
-          within a display generated by the Derivative Works, if and
-          wherever such third-party notices normally appear. The contents
-          of the NOTICE file are for informational purposes only and
-          do not modify the License. You may add Your own attribution
-          notices within Derivative Works that You distribute, alongside
-          or as an addendum to the NOTICE text from the Work, provided
-          that such additional attribution notices cannot be construed
-          as modifying the License.
+   (d) If the Work includes a "NOTICE" text file as part of its
+   distribution, then any Derivative Works that You distribute must
+   include a readable copy of the attribution notices contained
+   within such NOTICE file, excluding those notices that do not
+   pertain to any part of the Derivative Works, in at least one
+   of the following places: within a NOTICE text file distributed
+   as part of the Derivative Works; within the Source form or
+   documentation, if provided along with the Derivative Works; or,
+   within a display generated by the Derivative Works, if and
+   wherever such third-party notices normally appear. The contents
+   of the NOTICE file are for informational purposes only and
+   do not modify the License. You may add Your own attribution
+   notices within Derivative Works that You distribute, alongside
+   or as an addendum to the NOTICE text from the Work, provided
+   that such additional attribution notices cannot be construed
+   as modifying the License.
 
-      You may add Your own copyright statement to Your modifications and
-      may provide additional or different license terms and conditions
-      for use, reproduction, or distribution of Your modifications, or
-      for any such Derivative Works as a whole, provided Your use,
-      reproduction, and distribution of the Work otherwise complies with
-      the conditions stated in this License.
+   You may add Your own copyright statement to Your modifications and
+   may provide additional or different license terms and conditions
+   for use, reproduction, or distribution of Your modifications, or
+   for any such Derivative Works as a whole, provided Your use,
+   reproduction, and distribution of the Work otherwise complies with
+   the conditions stated in this License.
 
-   5. Submission of Contributions. Unless You explicitly state otherwise,
-      any Contribution intentionally submitted for inclusion in the Work
-      by You to the Licensor shall be under the terms and conditions of
-      this License, without any additional terms or conditions.
-      Notwithstanding the above, nothing herein shall supersede or modify
-      the terms of any separate license agreement you may have executed
-      with Licensor regarding such Contributions.
+1. Submission of Contributions. Unless You explicitly state otherwise,
+   any Contribution intentionally submitted for inclusion in the Work
+   by You to the Licensor shall be under the terms and conditions of
+   this License, without any additional terms or conditions.
+   Notwithstanding the above, nothing herein shall supersede or modify
+   the terms of any separate license agreement you may have executed
+   with Licensor regarding such Contributions.
 
-   6. Trademarks. This License does not grant permission to use the trade
-      names, trademarks, service marks, or product names of the Licensor,
-      except as required for reasonable and customary use in describing the
-      origin of the Work and reproducing the content of the NOTICE file.
+1. Trademarks. This License does not grant permission to use the trade
+   names, trademarks, service marks, or product names of the Licensor,
+   except as required for reasonable and customary use in describing the
+   origin of the Work and reproducing the content of the NOTICE file.
 
-   7. Disclaimer of Warranty. Unless required by applicable law or
-      agreed to in writing, Licensor provides the Work (and each
-      Contributor provides its Contributions) on an "AS IS" BASIS,
-      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-      implied, including, without limitation, any warranties or conditions
-      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-      PARTICULAR PURPOSE. You are solely responsible for determining the
-      appropriateness of using or redistributing the Work and assume any
-      risks associated with Your exercise of permissions under this License.
+1. Disclaimer of Warranty. Unless required by applicable law or
+   agreed to in writing, Licensor provides the Work (and each
+   Contributor provides its Contributions) on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+   implied, including, without limitation, any warranties or conditions
+   of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+   PARTICULAR PURPOSE. You are solely responsible for determining the
+   appropriateness of using or redistributing the Work and assume any
+   risks associated with Your exercise of permissions under this License.
 
-   8. Limitation of Liability. In no event and under no legal theory,
-      whether in tort (including negligence), contract, or otherwise,
-      unless required by applicable law (such as deliberate and grossly
-      negligent acts) or agreed to in writing, shall any Contributor be
-      liable to You for damages, including any direct, indirect, special,
-      incidental, or consequential damages of any character arising as a
-      result of this License or out of the use or inability to use the
-      Work (including but not limited to damages for loss of goodwill,
-      work stoppage, computer failure or malfunction, or any and all
-      other commercial damages or losses), even if such Contributor
-      has been advised of the possibility of such damages.
+1. Limitation of Liability. In no event and under no legal theory,
+   whether in tort (including negligence), contract, or otherwise,
+   unless required by applicable law (such as deliberate and grossly
+   negligent acts) or agreed to in writing, shall any Contributor be
+   liable to You for damages, including any direct, indirect, special,
+   incidental, or consequential damages of any character arising as a
+   result of this License or out of the use or inability to use the
+   Work (including but not limited to damages for loss of goodwill,
+   work stoppage, computer failure or malfunction, or any and all
+   other commercial damages or losses), even if such Contributor
+   has been advised of the possibility of such damages.
 
-   9. Accepting Warranty or Additional Liability. While redistributing
-      the Work or Derivative Works thereof, You may choose to offer,
-      and charge a fee for, acceptance of support, warranty, indemnity,
-      or other liability obligations and/or rights consistent with this
-      License. However, in accepting such obligations, You may act only
-      on Your own behalf and on Your sole responsibility, not on behalf
-      of any other Contributor, and only if You agree to indemnify,
-      defend, and hold each Contributor harmless for any liability
-      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
+1. Accepting Warranty or Additional Liability. While redistributing
+   the Work or Derivative Works thereof, You may choose to offer,
+   and charge a fee for, acceptance of support, warranty, indemnity,
+   or other liability obligations and/or rights consistent with this
+   License. However, in accepting such obligations, You may act only
+   on Your own behalf and on Your sole responsibility, not on behalf
+   of any other Contributor, and only if You agree to indemnify,
+   defend, and hold each Contributor harmless for any liability
+   incurred by, or claims asserted against, such Contributor by reason
+   of your accepting any such warranty or additional liability.
 
-   END OF TERMS AND CONDITIONS
+END OF TERMS AND CONDITIONS
 
-   APPENDIX: How to apply the Apache License to your work.
+APPENDIX: How to apply the Apache License to your work.
 
-      To apply the Apache License to your work, attach the following
-      boilerplate notice, with the fields enclosed by brackets "{}"
-      replaced with your own identifying information. (Don't include
-      the brackets!)  The text should be enclosed in the appropriate
-      comment syntax for the file format. We also recommend that a
-      file or class name and description of purpose be included on the
-      same "printed page" as the copyright notice for easier
-      identification within third-party archives.
+```
+  To apply the Apache License to your work, attach the following
+  boilerplate notice, with the fields enclosed by brackets "{}"
+  replaced with your own identifying information. (Don't include
+  the brackets!)  The text should be enclosed in the appropriate
+  comment syntax for the file format. We also recommend that a
+  file or class name and description of purpose be included on the
+  same "printed page" as the copyright notice for easier
+  identification within third-party archives.
+```
 
-   Copyright {yyyy} {name of copyright owner}
+Copyright {yyyy} {name of copyright owner}
 
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
 
-       http://www.apache.org/licenses/LICENSE-2.0
+```
+   http://www.apache.org/licenses/LICENSE-2.0
+```
 
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License.
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
 
---------------------------------------------------------------------------------
+______________________________________________________________________
 
 Code in python/ray/rllib/{evolution_strategies, dqn} adapted from
 https://github.com/openai (MIT License)
@@ -240,7 +241,7 @@ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 
---------------------------------------------------------------------------------
+______________________________________________________________________
 
 Code in python/ray/rllib/impala/vtrace.py from
 https://github.com/deepmind/scalable_agent
@@ -251,7 +252,9 @@ Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
 
-    https://www.apache.org/licenses/LICENSE-2.0
+```
+https://www.apache.org/licenses/LICENSE-2.0
+```
 
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
@@ -259,7 +262,8 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 
---------------------------------------------------------------------------------
+______________________________________________________________________
+
 Code in python/ray/rllib/ars is adapted from https://github.com/modestyachts/ARS
 
 Copyright (c) 2018, ARS contributors (Horia Mania, Aurelia Guy, Benjamin Recht)
@@ -269,11 +273,11 @@ Redistribution and use of ARS in source and binary forms, with or without
 modification, are permitted provided that the following conditions are met:
 
 1. Redistributions of source code must retain the above copyright notice, this
-list of conditions and the following disclaimer.
+   list of conditions and the following disclaimer.
 
-2. Redistributions in binary form must reproduce the above copyright notice,
-this list of conditions and the following disclaimer in the documentation and/or
-other materials provided with the distribution.
+1. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation and/or
+   other materials provided with the distribution.
 
 THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
@@ -286,5 +290,6 @@ ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-------------------
-Code in python/ray/_private/prometheus_exporter.py is adapted from https://github.com/census-instrumentation/opencensus-python/blob/master/contrib/opencensus-ext-prometheus/opencensus/ext/prometheus/stats_exporter/__init__.py
+______________________________________________________________________
+
+Code in python/ray/\_private/prometheus_exporter.py is adapted from https://github.com/census-instrumentation/opencensus-python/blob/master/contrib/opencensus-ext-prometheus/opencensus/ext/prometheus/stats_exporter/__init__.py

README.md

@@ -1,11 +1,11 @@
 [![PyPI version](https://badge.fury.io/py/FLAML.svg)](https://badge.fury.io/py/FLAML)
 ![Conda version](https://img.shields.io/conda/vn/conda-forge/flaml)
 [![Build](https://github.com/microsoft/FLAML/actions/workflows/python-package.yml/badge.svg)](https://github.com/microsoft/FLAML/actions/workflows/python-package.yml)
-![Python Version](https://img.shields.io/badge/3.8%20%7C%203.9%20%7C%203.10-blue)
+[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/FLAML)](https://pypi.org/project/FLAML/)
 [![Downloads](https://pepy.tech/badge/flaml)](https://pepy.tech/project/flaml)
 [![](https://img.shields.io/discord/1025786666260111483?logo=discord&style=flat)](https://discord.gg/Cppx2vSPVP)
-<!-- [![Join the chat at https://gitter.im/FLAMLer/community](https://badges.gitter.im/FLAMLer/community.svg)](https://gitter.im/FLAMLer/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge) -->
 
+<!-- [![Join the chat at https://gitter.im/FLAMLer/community](https://badges.gitter.im/FLAMLer/community.svg)](https://gitter.im/FLAMLer/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge) -->
 
 # A Fast Library for Automated Machine Learning & Tuning
 
@@ -14,37 +14,40 @@
     <br>
 </p>
 
-:fire: The automated multi-agent chat framework in [autogen](https://microsoft.github.io/FLAML/docs/Use-Cases/Autogen) is in preview from v2.0.0.
+: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: 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/FLAML/docs/Use-Cases/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: FLAML supports AutoML and Hyperparameter Tuning features in [Microsoft Fabric](https://learn.microsoft.com/en-us/fabric/get-started/microsoft-fabric-overview) private preview. Sign up for these features at: https://aka.ms/fabric/data-science/sign-up.
-
+: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).
 
 ## What is FLAML
+
 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.
-* 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.
+- 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.
+- 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.
 
-FLAML is powered by a series of [research studies](/docs/Research) from Microsoft Research and collaborators such as Penn State University, Stevens Institute of Technology, University of Washington, and University of Waterloo.
+FLAML is powered by a series of [research studies](https://microsoft.github.io/FLAML/docs/Research/) from Microsoft Research and collaborators such as Penn State University, Stevens Institute of Technology, University of Washington, and University of Waterloo.
 
 FLAML has a .NET implementation in [ML.NET](http://dot.net/ml), an open-source, cross-platform machine learning framework for .NET.
 
 ## Installation
 
-FLAML requires **Python version >= 3.8**. It can be installed from pip:
+FLAML requires **Python version >= 3.9**. It can be installed from pip:
 
 ```bash
 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/FLAML/docs/Use-Cases/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 [`autogen`](https://microsoft.github.io/autogen/) package.
+
 ```bash
 pip install "flaml[autogen]"
 ```
@@ -54,18 +57,24 @@ Each of the [`notebook examples`](https://github.com/microsoft/FLAML/tree/main/n
 
 ## Quickstart
 
-* (New) The [autogen](https://microsoft.github.io/FLAML/docs/Use-Cases/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,
+- (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.")
+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(
@@ -80,30 +89,32 @@ config, analysis = autogen.Completion.tune(
 # 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).
+
+- 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).
 
 ```python
 from flaml import AutoML
+
 automl = AutoML()
 automl.fit(X_train, y_train, task="classification")
 ```
 
-* You can restrict the learners and use FLAML as a fast hyperparameter tuning
-tool for XGBoost, LightGBM, Random Forest etc. or a [customized learner](https://microsoft.github.io/FLAML/docs/Use-Cases/Task-Oriented-AutoML#estimator-and-search-space).
+- You can restrict the learners and use FLAML as a fast hyperparameter tuning
+  tool for XGBoost, LightGBM, Random Forest etc. or a [customized learner](https://microsoft.github.io/FLAML/docs/Use-Cases/Task-Oriented-AutoML#estimator-and-search-space).
 
 ```python
 automl.fit(X_train, y_train, task="classification", estimator_list=["lgbm"])
 ```
 
-* You can also run generic hyperparameter tuning for a [custom function](https://microsoft.github.io/FLAML/docs/Use-Cases/Tune-User-Defined-Function).
+- You can also run generic hyperparameter tuning for a [custom function](https://microsoft.github.io/FLAML/docs/Use-Cases/Tune-User-Defined-Function).
 
 ```python
 from flaml import tune
 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.
+- [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.
 
 ```python
 from flaml.default import LGBMRegressor
@@ -143,3 +154,9 @@ provided by the bot. You will only need to do this once across all repos using o
 This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
 For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or
 contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments.
+
+## Contributors Wall
+
+<a href="https://github.com/microsoft/flaml/graphs/contributors">
+  <img src="https://contrib.rocks/image?repo=microsoft/flaml&max=204" />
+</a>

SECURITY.md

@@ -4,7 +4,7 @@
 
 Microsoft takes the security of our software products and services seriously, which includes all source code repositories managed through our GitHub organizations, which include [Microsoft](https://github.com/Microsoft), [Azure](https://github.com/Azure), [DotNet](https://github.com/dotnet), [AspNet](https://github.com/aspnet), [Xamarin](https://github.com/xamarin), and [our GitHub organizations](https://opensource.microsoft.com/).
 
-If you believe you have found a security vulnerability in any Microsoft-owned repository that meets [Microsoft's definition of a security vulnerability](https://docs.microsoft.com/en-us/previous-versions/tn-archive/cc751383(v=technet.10)), please report it to us as described below.
+If you believe you have found a security vulnerability in any Microsoft-owned repository that meets [Microsoft's definition of a security vulnerability](<https://docs.microsoft.com/en-us/previous-versions/tn-archive/cc751383(v=technet.10)>), please report it to us as described below.
 
 ## Reporting Security Issues
 
@@ -18,13 +18,13 @@ You should receive a response within 24 hours. If for some reason you do not, pl
 
 Please include the requested information listed below (as much as you can provide) to help us better understand the nature and scope of the possible issue:
 
-  * Type of issue (e.g. buffer overflow, SQL injection, cross-site scripting, etc.)
-  * Full paths of source file(s) related to the manifestation of the issue
-  * The location of the affected source code (tag/branch/commit or direct URL)
-  * Any special configuration required to reproduce the issue
-  * Step-by-step instructions to reproduce the issue
-  * Proof-of-concept or exploit code (if possible)
-  * Impact of the issue, including how an attacker might exploit the issue
+- Type of issue (e.g. buffer overflow, SQL injection, cross-site scripting, etc.)
+- Full paths of source file(s) related to the manifestation of the issue
+- The location of the affected source code (tag/branch/commit or direct URL)
+- Any special configuration required to reproduce the issue
+- Step-by-step instructions to reproduce the issue
+- Proof-of-concept or exploit code (if possible)
+- Impact of the issue, including how an attacker might exploit the issue
 
 This information will help us triage your report more quickly.
 

flaml/__init__.py

@@ -1,10 +1,20 @@
 import logging
-from flaml.automl import AutoML, logger_formatter
-from flaml.tune.searcher import CFO, BlendSearch, FLOW2, BlendSearchTuner, RandomSearch
-from flaml.onlineml.autovw import AutoVW
-from flaml.version import __version__
+import warnings
 
+try:
+    from flaml.automl import AutoML, logger_formatter
+
+    has_automl = True
+except ImportError:
+    has_automl = False
+from flaml.onlineml.autovw import AutoVW
+from flaml.tune.searcher import CFO, FLOW2, BlendSearch, BlendSearchTuner, RandomSearch
+from flaml.version import __version__
 
 # Set the root logger.
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
+if logger.level == logging.NOTSET:
+    logger.setLevel(logging.INFO)
+
+if not has_automl:
+    warnings.warn("flaml.automl is not available. Please install flaml[automl] to enable AutoML functionalities.")

flaml/autogen/__init__.py

@@ -1,3 +1,3 @@
-from .oai import *
 from .agentchat import *
 from .code_utils import DEFAULT_MODEL, FAST_MODEL
+from .oai import *

flaml/autogen/agentchat/__init__.py

@@ -1,12 +1,12 @@
 from .agent import Agent
-from .responsive_agent import ResponsiveAgent
 from .assistant_agent import AssistantAgent
-from .user_proxy_agent import UserProxyAgent
+from .conversable_agent import ConversableAgent
 from .groupchat import GroupChat, GroupChatManager
+from .user_proxy_agent import UserProxyAgent
 
 __all__ = [
     "Agent",
-    "ResponsiveAgent",
+    "ConversableAgent",
     "AssistantAgent",
     "UserProxyAgent",
     "GroupChat",

flaml/autogen/agentchat/agent.py

@@ -25,10 +25,10 @@ class Agent:
         return self._name
 
     def send(self, message: Union[Dict, str], recipient: "Agent", request_reply: Optional[bool] = None):
-        """(Aabstract method) Send a message to another agent."""
+        """(Abstract method) Send a message to another agent."""
 
     async def a_send(self, message: Union[Dict, str], recipient: "Agent", request_reply: Optional[bool] = None):
-        """(Aabstract async method) Send a message to another agent."""
+        """(Abstract async method) Send a message to another agent."""
 
     def receive(self, message: Union[Dict, str], sender: "Agent", request_reply: Optional[bool] = None):
         """(Abstract method) Receive a message from another agent."""

flaml/autogen/agentchat/assistant_agent.py

@@ -1,26 +1,30 @@
-from .responsive_agent import ResponsiveAgent
 from typing import Callable, Dict, Optional, Union
 
+from .conversable_agent import ConversableAgent
 
-class AssistantAgent(ResponsiveAgent):
-    """(In preview) Assistant agent, designed to solve a task with LLM.
 
-    AssistantAgent is a subclass of ResponsiveAgent configured with a default system message.
-    The default system message is designed to solve a task with LLM,
-    including suggesting python code blocks and debugging.
-    `human_input_mode` is default to "NEVER"
-    and `code_execution_config` is default to False.
-    This agent doesn't execute code by default, and expects the user to execute the code.
+class AssistantAgent(ConversableAgent):
+    """(In preview) Assistant agent, designed to solve tasks with LLM.
+
+    AssistantAgent is a subclass of ConversableAgent configured with a default system message.
+    The default system message is designed to solve tasks with LLM,
+    including suggesting Python code blocks and debugging.
+    `human_input_mode` defaults to "NEVER"
+    and `code_execution_config` defaults to False.
+    This agent doesn't execute code by default and expects the user to execute the code.
     """
 
     DEFAULT_SYSTEM_MESSAGE = """You are a helpful AI assistant.
-    In the following cases, suggest python code (in a python coding block) or shell script (in a sh coding block) for the user to execute. You must indicate the script type in the code block. The user cannot provide any other feedback or perform any other action beyond executing the code you suggest. The user can't modify your code. So do not suggest incomplete code which requires users to modify. Don't use a code block if it's not intended to be executed by the user.
-    1. When you need to collect info, use the code to output the info you need, for example, browse or search the web, download/read a file, print the content of a webpage or a file, get the current date/time.
-    2. When you need to perform some task with code, use the code to perform the task and output the result. Finish the task smartly. Solve the task step by step if you need to.
-    If you want the user to save the code in a file before executing it, put # filename: <filename> inside the code block as the first line. Don't include multiple code blocks in one response. Do not ask users to copy and paste the result. Instead, use 'print' function for the output when relevant. Check the execution result returned by the user.
-    If the result indicates there is an error, fix the error and output the code again. Suggest the full code instead of partial code or code changes. If the error can't be fixed or if the task is not solved even after the code is executed successfully, analyze the problem, revisit your assumption, collect additional info you need, and think of a different approach to try.
-    When you find an answer, verify the answer carefully. If a function for planning is provided, call the function to make plans and verify the execution.
-    Reply "TERMINATE" in the end when everything is done.
+Solve tasks using your coding and language skills.
+In the following cases, suggest Python code (in a Python coding block) or shell script (in an sh coding block) for the user to execute.
+    1. When you need to collect info, use the code to output the info you need, for example, browse or search the web, download/read a file, print the content of a webpage or a file, get the current date/time. After sufficient info is printed and the task is ready to be solved based on your language skill, you can solve the task by yourself.
+    2. When you need to perform some task with code, use the code to perform the task and output the result. Finish the task smartly.
+Solve the task step by step if you need to. If a plan is not provided, explain your plan first. Be clear which step uses code, and which step uses your language skill.
+When using code, you must indicate the script type in the code block. The user cannot provide any other feedback or perform any other action beyond executing the code you suggest. The user can't modify your code. So do not suggest incomplete code which requires users to modify. Don't use a code block if it's not intended to be executed by the user.
+If you want the user to save the code in a file before executing it, put # filename: <filename> inside the code block as the first line. Don't include multiple code blocks in one response. Do not ask users to copy and paste the result. Instead, use the 'print' function for the output when relevant. Check the execution result returned by the user.
+If the result indicates there is an error, fix the error and output the code again. Suggest the full code instead of partial code or code changes. If the error can't be fixed or if the task is not solved even after the code is executed successfully, analyze the problem, revisit your assumption, collect additional info you need, and think of a different approach to try.
+When you find an answer, verify the answer carefully. Include verifiable evidence in your response if possible.
+Reply "TERMINATE" in the end when everything is done.
     """
 
     def __init__(
@@ -32,24 +36,24 @@ class AssistantAgent(ResponsiveAgent):
         max_consecutive_auto_reply: Optional[int] = None,
         human_input_mode: Optional[str] = "NEVER",
         code_execution_config: Optional[Union[Dict, bool]] = False,
-        **kwargs,
+        **kwargs: Dict,
     ):
         """
         Args:
-            name (str): agent name.
-            system_message (str): system message for the ChatCompletion inference.
-                Please override this attribute if you want to reprogram the agent.
-            llm_config (dict): llm inference configuration.
-                Please refer to [autogen.Completion.create](/docs/reference/autogen/oai/completion#create)
+            name (str): Agent name.
+            system_message (Optional[str]): System message for the ChatCompletion inference.
+                Override this attribute if you want to reprogram the agent.
+            llm_config (Optional[Union[Dict, bool]]): LLM inference configuration.
+                Refer to [autogen.Completion.create](/docs/reference/autogen/oai/completion#create)
                 for available options.
-            is_termination_msg (function): a function that takes a message in the form of a dictionary
+            is_termination_msg (Optional[Callable[[Dict], bool]]): A function that takes a message in the form of a dictionary
                 and returns a boolean value indicating if this received message is a termination message.
                 The dict can contain the following keys: "content", "role", "name", "function_call".
-            max_consecutive_auto_reply (int): the maximum number of consecutive auto replies.
-                default to None (no limit provided, class attribute MAX_CONSECUTIVE_AUTO_REPLY will be used as the limit in this case).
+            max_consecutive_auto_reply (Optional[int]): The maximum number of consecutive auto replies.
+                Defaults to None (no limit provided, class attribute MAX_CONSECUTIVE_AUTO_REPLY will be used as the limit in this case).
                 The limit only plays a role when human_input_mode is not "ALWAYS".
-            **kwargs (dict): Please refer to other kwargs in
-                [ResponsiveAgent](responsive_agent#__init__).
+            **kwargs (Dict): Additional keyword arguments. Refer to other kwargs in
+                [ConversableAgent](conversable_agent#__init__).
         """
         super().__init__(
             name,

flaml/autogen/agentchat/contrib/math_user_proxy_agent.py

@@ -1,14 +1,14 @@
-import re
 import os
-from pydantic import BaseModel, Extra, root_validator
-from typing import Any, Callable, Dict, List, Optional, Union
+import re
 from time import sleep
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Extra, root_validator
 
 from flaml.autogen.agentchat import Agent, UserProxyAgent
-from flaml.autogen.code_utils import UNKNOWN, extract_code, execute_code, infer_lang
+from flaml.autogen.code_utils import UNKNOWN, execute_code, extract_code, infer_lang
 from flaml.autogen.math_utils import get_answer
 
-
 PROMPTS = {
     # default
     "default": """Let's use Python to solve a math problem.
@@ -156,7 +156,7 @@ class MathUserProxyAgent(UserProxyAgent):
                     when the number of auto reply reaches the max_consecutive_auto_reply or when is_termination_msg is True.
             default_auto_reply (str or dict or None): the default auto reply message when no code execution or llm based reply is generated.
             max_invalid_q_per_step (int): (ADDED) the maximum number of invalid queries per step.
-            **kwargs (dict): other kwargs in [UserProxyAgent](user_proxy_agent#__init__).
+            **kwargs (dict): other kwargs in [UserProxyAgent](../user_proxy_agent#__init__).
         """
         super().__init__(
             name=name,
@@ -165,7 +165,7 @@ class MathUserProxyAgent(UserProxyAgent):
             default_auto_reply=default_auto_reply,
             **kwargs,
         )
-        self.register_auto_reply(Agent, MathUserProxyAgent._generate_math_reply, 1)
+        self.register_reply([Agent, None], MathUserProxyAgent._generate_math_reply, 1)
         # fixed var
         self._max_invalid_q_per_step = max_invalid_q_per_step
 

flaml/autogen/agentchat/contrib/retrieve_assistant_agent.py

@@ -1,6 +1,7 @@
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
 from flaml.autogen.agentchat.agent import Agent
 from flaml.autogen.agentchat.assistant_agent import AssistantAgent
-from typing import Callable, Dict, Optional, Union, List, Tuple, Any
 
 
 class RetrieveAssistantAgent(AssistantAgent):
@@ -16,7 +17,7 @@ class RetrieveAssistantAgent(AssistantAgent):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.register_auto_reply(Agent, RetrieveAssistantAgent._generate_retrieve_assistant_reply)
+        self.register_reply(Agent, RetrieveAssistantAgent._generate_retrieve_assistant_reply)
 
     def _generate_retrieve_assistant_reply(
         self,

flaml/autogen/agentchat/contrib/retrieve_user_proxy_agent.py

@@ -1,12 +1,13 @@
-import chromadb
-from flaml.autogen.agentchat.agent import Agent
-from flaml.autogen.agentchat import UserProxyAgent
-from flaml.autogen.retrieve_utils import create_vector_db_from_dir, query_vector_db, num_tokens_from_text
-from flaml.autogen.code_utils import extract_code
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
-from typing import Callable, Dict, Optional, Union, List, Tuple, Any
+import chromadb
 from IPython import get_ipython
 
+from flaml.autogen.agentchat import UserProxyAgent
+from flaml.autogen.agentchat.agent import Agent
+from flaml.autogen.code_utils import extract_code
+from flaml.autogen.retrieve_utils import create_vector_db_from_dir, num_tokens_from_text, query_vector_db
+
 try:
     from termcolor import colored
 except ImportError:
@@ -122,7 +123,7 @@ class RetrieveUserProxyAgent(UserProxyAgent):
                     can be found at `https://www.sbert.net/docs/pretrained_models.html`. The default model is a
                     fast model. If you want to use a high performance model, `all-mpnet-base-v2` is recommended.
                 - customized_prompt (Optional, str): the customized prompt for the retrieve chat. Default is None.
-            **kwargs (dict): other kwargs in [UserProxyAgent](user_proxy_agent#__init__).
+            **kwargs (dict): other kwargs in [UserProxyAgent](../user_proxy_agent#__init__).
         """
         super().__init__(
             name=name,
@@ -148,7 +149,7 @@ class RetrieveUserProxyAgent(UserProxyAgent):
         self._ipython = get_ipython()
         self._doc_idx = -1  # the index of the current used doc
         self._results = {}  # the results of the current query
-        self.register_auto_reply(Agent, RetrieveUserProxyAgent._generate_retrieve_user_reply)
+        self.register_reply(Agent, RetrieveUserProxyAgent._generate_retrieve_user_reply)
 
     @staticmethod
     def get_max_tokens(model="gpt-3.5-turbo"):

flaml/autogen/agentchat/conversable_agent.py

@@ -1,10 +1,10 @@
 import asyncio
-from collections import defaultdict
 import copy
 import json
+from collections import defaultdict
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
+
 from flaml.autogen import oai
-from .agent import Agent
 from flaml.autogen.code_utils import (
     DEFAULT_MODEL,
     UNKNOWN,
@@ -13,6 +13,8 @@ from flaml.autogen.code_utils import (
     infer_lang,
 )
 
+from .agent import Agent
+
 try:
     from termcolor import colored
 except ImportError:
@@ -21,11 +23,11 @@ except ImportError:
         return x
 
 
-class ResponsiveAgent(Agent):
-    """(Experimental) A class for generic responsive agents which can be configured as assistant or user proxy.
+class ConversableAgent(Agent):
+    """(In preview) A class for generic conversable agents which can be configured as assistant or user proxy.
 
     After receiving each message, the agent will send a reply to the sender unless the msg is a termination msg.
-    For example, AssistantAgent and UserProxyAgent are subclasses of ResponsiveAgent,
+    For example, AssistantAgent and UserProxyAgent are subclasses of this class,
     configured with different default settings.
 
     To modify auto reply, override `generate_reply` method.
@@ -119,12 +121,12 @@ class ResponsiveAgent(Agent):
         self._default_auto_reply = default_auto_reply
         self._reply_func_list = []
         self.reply_at_receive = defaultdict(bool)
-        self.register_auto_reply(Agent, ResponsiveAgent.generate_oai_reply)
-        self.register_auto_reply(Agent, ResponsiveAgent.generate_code_execution_reply)
-        self.register_auto_reply(Agent, ResponsiveAgent.generate_function_call_reply)
-        self.register_auto_reply(Agent, ResponsiveAgent.check_termination_and_human_reply)
+        self.register_reply([Agent, None], ConversableAgent.generate_oai_reply)
+        self.register_reply([Agent, None], ConversableAgent.generate_code_execution_reply)
+        self.register_reply([Agent, None], ConversableAgent.generate_function_call_reply)
+        self.register_reply([Agent, None], ConversableAgent.check_termination_and_human_reply)
 
-    def register_auto_reply(
+    def register_reply(
         self,
         trigger: Union[Type[Agent], str, Agent, Callable[[Agent], bool], List],
         reply_func: Callable,
@@ -145,11 +147,13 @@ class ResponsiveAgent(Agent):
                 - If an agent instance is provided, the reply function will be called when the sender is the agent instance.
                 - If a callable is provided, the reply function will be called when the callable returns True.
                 - If a list is provided, the reply function will be called when any of the triggers in the list is activated.
+                - If None is provided, the reply function will be called only when the sender is None.
+                Note: Be sure to register `None` as a trigger if you would like to trigger an auto-reply function with non-empty messages and `sender=None`.
             reply_func (Callable): the reply function.
                 The function takes a recipient agent, a list of messages, a sender agent and a config as input and returns a reply message.
         ```python
         def reply_func(
-            recipient: ResponsiveAgent,
+            recipient: ConversableAgent,
             messages: Optional[List[Dict]] = None,
             sender: Optional[Agent] = None,
             config: Optional[Any] = None,
@@ -497,7 +501,7 @@ class ResponsiveAgent(Agent):
 
     def initiate_chat(
         self,
-        recipient: "ResponsiveAgent",
+        recipient: "ConversableAgent",
         clear_history: Optional[bool] = True,
         silent: Optional[bool] = False,
         **context,
@@ -520,7 +524,7 @@ class ResponsiveAgent(Agent):
 
     async def a_initiate_chat(
         self,
-        recipient: "ResponsiveAgent",
+        recipient: "ConversableAgent",
         clear_history: Optional[bool] = True,
         silent: Optional[bool] = False,
         **context,
@@ -609,7 +613,7 @@ class ResponsiveAgent(Agent):
         if messages is None:
             messages = self._oai_messages[sender]
         last_n_messages = code_execution_config.pop("last_n_messages", 1)
-        for i in range(last_n_messages):
+        for i in range(min(len(messages), last_n_messages)):
             message = messages[-(i + 1)]
             code_blocks = extract_code(message["content"])
             if len(code_blocks) == 1 and code_blocks[0][0] == UNKNOWN:
@@ -726,6 +730,7 @@ class ResponsiveAgent(Agent):
         """Reply based on the conversation history and the sender.
 
         Either messages or sender must be provided.
+        Register a reply_func with `None` as one trigger for it to be activated when `messages` is non-empty and `sender` is `None`.
         Use registered auto reply functions to generate replies.
         By default, the following functions are checked in order:
         1. check_termination_and_human_reply
@@ -748,17 +753,19 @@ class ResponsiveAgent(Agent):
             str or dict or None: reply. None if no reply is generated.
         """
         assert messages is not None or sender is not None, "Either messages or sender must be provided."
-        if sender is not None:
-            for reply_func_tuple in self._reply_func_list:
-                reply_func = reply_func_tuple["reply_func"]
-                if exclude and reply_func in exclude:
-                    continue
-                if asyncio.coroutines.iscoroutinefunction(reply_func):
-                    continue
-                if self._match_trigger(reply_func_tuple["trigger"], sender):
-                    final, reply = reply_func(self, messages=messages, sender=sender, config=reply_func_tuple["config"])
-                    if final:
-                        return reply
+        if messages is None:
+            messages = self._oai_messages[sender]
+
+        for reply_func_tuple in self._reply_func_list:
+            reply_func = reply_func_tuple["reply_func"]
+            if exclude and reply_func in exclude:
+                continue
+            if asyncio.coroutines.iscoroutinefunction(reply_func):
+                continue
+            if self._match_trigger(reply_func_tuple["trigger"], sender):
+                final, reply = reply_func(self, messages=messages, sender=sender, config=reply_func_tuple["config"])
+                if final:
+                    return reply
         return self._default_auto_reply
 
     async def a_generate_reply(
@@ -770,6 +777,7 @@ class ResponsiveAgent(Agent):
         """(async) Reply based on the conversation history and the sender.
 
         Either messages or sender must be provided.
+        Register a reply_func with `None` as one trigger for it to be activated when `messages` is non-empty and `sender` is `None`.
         Use registered auto reply functions to generate replies.
         By default, the following functions are checked in order:
         1. check_termination_and_human_reply
@@ -792,27 +800,29 @@ class ResponsiveAgent(Agent):
             str or dict or None: reply. None if no reply is generated.
         """
         assert messages is not None or sender is not None, "Either messages or sender must be provided."
-        if sender is not None:
-            for reply_func_tuple in self._reply_func_list:
-                reply_func = reply_func_tuple["reply_func"]
-                if exclude and reply_func in exclude:
-                    continue
-                if self._match_trigger(reply_func_tuple["trigger"], sender):
-                    if asyncio.coroutines.iscoroutinefunction(reply_func):
-                        final, reply = await reply_func(
-                            self, messages=messages, sender=sender, config=reply_func_tuple["config"]
-                        )
-                    else:
-                        final, reply = reply_func(
-                            self, messages=messages, sender=sender, config=reply_func_tuple["config"]
-                        )
-                    if final:
-                        return reply
+        if messages is None:
+            messages = self._oai_messages[sender]
+
+        for reply_func_tuple in self._reply_func_list:
+            reply_func = reply_func_tuple["reply_func"]
+            if exclude and reply_func in exclude:
+                continue
+            if self._match_trigger(reply_func_tuple["trigger"], sender):
+                if asyncio.coroutines.iscoroutinefunction(reply_func):
+                    final, reply = await reply_func(
+                        self, messages=messages, sender=sender, config=reply_func_tuple["config"]
+                    )
+                else:
+                    final, reply = reply_func(self, messages=messages, sender=sender, config=reply_func_tuple["config"])
+                if final:
+                    return reply
         return self._default_auto_reply
 
     def _match_trigger(self, trigger, sender):
         """Check if the sender matches the trigger."""
-        if isinstance(trigger, str):
+        if trigger is None:
+            return sender is None
+        elif isinstance(trigger, str):
             return trigger == sender.name
         elif isinstance(trigger, type):
             return isinstance(sender, trigger)
@@ -887,10 +897,11 @@ class ResponsiveAgent(Agent):
                 exitcode, logs, image = (
                     1,
                     f"unknown language {lang}",
-                    self._code_execution_config["use_docker"],
+                    None,
                 )
                 # raise NotImplementedError
-            self._code_execution_config["use_docker"] = image
+            if image is not None:
+                self._code_execution_config["use_docker"] = image
             logs_all += "\n" + logs
             if exitcode != 0:
                 return exitcode, logs_all
@@ -953,7 +964,7 @@ class ResponsiveAgent(Agent):
                 content = f"Error: {e}\n You argument should follow json format."
 
             # Try to execute the function
-            if arguments:
+            if arguments is not None:
                 print(
                     colored(f"\n>>>>>>>> EXECUTING FUNCTION {func_name}...", "magenta"),
                     flush=True,

flaml/autogen/agentchat/groupchat.py

@@ -1,8 +1,9 @@
-from dataclasses import dataclass
 import sys
+from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
+
 from .agent import Agent
-from .responsive_agent import ResponsiveAgent
+from .conversable_agent import ConversableAgent
 
 
 @dataclass
@@ -12,6 +13,7 @@ class GroupChat:
     agents: List[Agent]
     messages: List[Dict]
     max_round: int = 10
+    admin_name: str = "Admin"  # the name of the admin agent
 
     @property
     def agent_names(self) -> List[str]:
@@ -38,10 +40,18 @@ class GroupChat:
 Read the following conversation.
 Then select the next role from {self.agent_names} to play. Only return the role."""
 
-    def select_speaker(self, last_speaker: Agent, selctor: ResponsiveAgent):
+    def select_speaker(self, last_speaker: Agent, selector: ConversableAgent):
         """Select the next speaker."""
-        selctor.update_system_message(self.select_speaker_msg())
-        final, name = selctor.generate_oai_reply(self.messages)
+        selector.update_system_message(self.select_speaker_msg())
+        final, name = selector.generate_oai_reply(
+            self.messages
+            + [
+                {
+                    "role": "system",
+                    "content": f"Read the above conversation. Then select the next role from {self.agent_names} to play. Only return the role.",
+                }
+            ]
+        )
         if not final:
             # i = self._random.randint(0, len(self._agent_names) - 1)  # randomly pick an id
             return self.next_agent(last_speaker)
@@ -54,8 +64,8 @@ Then select the next role from {self.agent_names} to play. Only return the role.
         return "\n".join([f"{agent.name}: {agent.system_message}" for agent in self.agents])
 
 
-class GroupChatManager(ResponsiveAgent):
-    """(WIP) A chat manager agent that can manage a group chat of multiple agents."""
+class GroupChatManager(ConversableAgent):
+    """(In preview) A chat manager agent that can manage a group chat of multiple agents."""
 
     def __init__(
         self,
@@ -75,7 +85,7 @@ class GroupChatManager(ResponsiveAgent):
             system_message=system_message,
             **kwargs,
         )
-        self.register_auto_reply(Agent, GroupChatManager.run_chat, config=groupchat, reset_config=GroupChat.reset)
+        self.register_reply(Agent, GroupChatManager.run_chat, config=groupchat, reset_config=GroupChat.reset)
         # self._random = random.Random(seed)
 
     def run_chat(
@@ -89,21 +99,36 @@ class GroupChatManager(ResponsiveAgent):
             messages = self._oai_messages[sender]
         message = messages[-1]
         speaker = sender
-        for i in range(config.max_round):
+        groupchat = config
+        for i in range(groupchat.max_round):
             # set the name to speaker's name if the role is not function
             if message["role"] != "function":
                 message["name"] = speaker.name
-            config.messages.append(message)
+            groupchat.messages.append(message)
             # broadcast the message to all agents except the speaker
-            for agent in config.agents:
+            for agent in groupchat.agents:
                 if agent != speaker:
                     self.send(message, agent, request_reply=False, silent=True)
-            if i != config.max_round - 1:
-                # speaker selection msg from an agent
-                speaker = config.select_speaker(speaker, self)
+            if i == groupchat.max_round - 1:
+                # the last round
+                break
+            try:
+                # select the next speaker
+                speaker = groupchat.select_speaker(speaker, self)
+                # let the speaker speak
                 reply = speaker.generate_reply(sender=self)
-                if reply is None:
-                    break
-                speaker.send(reply, self, request_reply=False)
-                message = self.last_message(speaker)
+            except KeyboardInterrupt:
+                # let the admin agent speak if interrupted
+                if groupchat.admin_name in groupchat.agent_names:
+                    # admin agent is one of the participants
+                    speaker = groupchat.agent_by_name(groupchat.admin_name)
+                    reply = speaker.generate_reply(sender=self)
+                else:
+                    # admin agent is not found in the participants
+                    raise
+            if reply is None:
+                break
+            # The speaker sends the message without requesting a reply
+            speaker.send(reply, self, request_reply=False)
+            message = self.last_message(speaker)
         return True, None

flaml/autogen/agentchat/user_proxy_agent.py

@@ -1,14 +1,15 @@
-from .responsive_agent import ResponsiveAgent
 from typing import Callable, Dict, Optional, Union
 
+from .conversable_agent import ConversableAgent
 
-class UserProxyAgent(ResponsiveAgent):
+
+class UserProxyAgent(ConversableAgent):
     """(In preview) A proxy agent for the user, that can execute code and provide feedback to the other agents.
 
-    UserProxyAgent is a subclass of ResponsiveAgent configured with `human_input_mode` to ALWAYS
+    UserProxyAgent is a subclass of ConversableAgent configured with `human_input_mode` to ALWAYS
     and `llm_config` to False. By default, the agent will prompt for human input every time a message is received.
     Code execution is enabled by default. LLM-based auto reply is disabled by default.
-    To modify auto reply, register a method with (`register_auto_reply`)[responsive_agent#register_auto_reply].
+    To modify auto reply, register a method with (`register_reply`)[conversable_agent#register_reply].
     To modify the way to get human input, override `get_human_input` method.
     To modify the way to execute code blocks, single code block, or function call, override `execute_code_blocks`,
     `run_code`, and `execute_function` methods respectively.

flaml/autogen/code_utils.py

@@ -1,13 +1,14 @@
+import logging
+import os
+import pathlib
+import re
 import signal
 import subprocess
 import sys
-import os
-import pathlib
-from typing import List, Dict, Tuple, Optional, Union, Callable
-import re
 import time
 from hashlib import md5
-import logging
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
 from flaml.autogen import oai
 
 try:
@@ -124,7 +125,7 @@ def improve_function(file_name, func_name, objective, **config):
     """(work in progress) Improve the function to achieve the objective."""
     params = {**_IMPROVE_FUNCTION_CONFIG, **config}
     # read the entire file into a str
-    with open(file_name, "r") as f:
+    with open(file_name) as f:
         file_string = f.read()
     response = oai.Completion.create(
         {"func_name": func_name, "objective": objective, "file_string": file_string}, **params
@@ -157,7 +158,7 @@ def improve_code(files, objective, suggest_only=True, **config):
     code = ""
     for file_name in files:
         # read the entire file into a string
-        with open(file_name, "r") as f:
+        with open(file_name) as f:
             file_string = f.read()
         code += f"""{file_name}:
 {file_string}
@@ -347,9 +348,9 @@ def execute_code(
         # extract the exit code from the logs
         pattern = re.compile(f"{exit_code_str}(\\d+){exit_code_str}")
         match = pattern.search(logs)
-        exit_code = int(match.group(1))
+        exit_code = 1 if match is None else int(match.group(1))
         # remove the exit code from the logs
-        logs = pattern.sub("", logs)
+        logs = logs if match is None else pattern.sub("", logs)
 
     if original_filename is None:
         os.remove(filepath)

flaml/autogen/math_utils.py

@@ -1,5 +1,6 @@
 from typing import Optional
-from flaml.autogen import oai, DEFAULT_MODEL
+
+from flaml.autogen import DEFAULT_MODEL, oai
 
 _MATH_PROMPT = "{problem} Solve the problem carefully. Simplify your answer as much as possible. Put the final answer in \\boxed{{}}."
 _MATH_CONFIG = {
@@ -129,7 +130,7 @@ def _fix_a_slash_b(string: str) -> str:
     try:
         a = int(a_str)
         b = int(b_str)
-        assert string == "{}/{}".format(a, b)
+        assert string == f"{a}/{b}"
         new_string = "\\frac{" + str(a) + "}{" + str(b) + "}"
         return new_string
     except Exception:

flaml/autogen/oai/__init__.py

@@ -1,10 +1,10 @@
-from flaml.autogen.oai.completion import Completion, ChatCompletion
+from flaml.autogen.oai.completion import ChatCompletion, Completion
 from flaml.autogen.oai.openai_utils import (
-    get_config_list,
+    config_list_from_json,
+    config_list_from_models,
     config_list_gpt4_gpt35,
     config_list_openai_aoai,
-    config_list_from_models,
-    config_list_from_json,
+    get_config_list,
 )
 
 __all__ = [

flaml/autogen/oai/completion.py

@@ -1,28 +1,31 @@
-from time import sleep
 import logging
-import time
-from typing import List, Optional, Dict, Callable, Union
-import sys
 import shutil
+import sys
+import time
+from time import sleep
+from typing import Callable, Dict, List, Optional, Union
+
 import numpy as np
-from flaml import tune, BlendSearch
-from flaml.tune.space import is_constant
+
+from flaml import BlendSearch, tune
 from flaml.automl.logger import logger_formatter
+from flaml.tune.space import is_constant
+
 from .openai_utils import get_key
 
 try:
-    import openai
-    from openai.error import (
-        ServiceUnavailableError,
-        RateLimitError,
-        APIError,
-        InvalidRequestError,
-        APIConnectionError,
-        Timeout,
-        AuthenticationError,
-    )
-    from openai import Completion as openai_Completion
     import diskcache
+    import openai
+    from openai import Completion as openai_Completion
+    from openai.error import (
+        APIConnectionError,
+        APIError,
+        AuthenticationError,
+        InvalidRequestError,
+        RateLimitError,
+        ServiceUnavailableError,
+        Timeout,
+    )
 
     ERROR = None
 except ImportError:
@@ -48,6 +51,7 @@ class Completion(openai_Completion):
         "gpt-3.5-turbo-0301",  # deprecate in Sep
         "gpt-3.5-turbo-0613",
         "gpt-3.5-turbo-16k",
+        "gpt-3.5-turbo-16k-0613",
         "gpt-35-turbo",
         "gpt-4",
         "gpt-4-32k",
@@ -70,6 +74,7 @@ class Completion(openai_Completion):
         "gpt-3.5-turbo-0301": (0.0015, 0.002),  # deprecate in Sep
         "gpt-3.5-turbo-0613": (0.0015, 0.002),
         "gpt-3.5-turbo-16k": (0.003, 0.004),
+        "gpt-3.5-turbo-16k-0613": (0.003, 0.004),
         "gpt-35-turbo": 0.002,
         "gpt-4": (0.03, 0.06),
         "gpt-4-32k": (0.06, 0.12),
@@ -695,7 +700,7 @@ class Completion(openai_Completion):
                 E.g., `prompt="Complete the following sentence: {prefix}, context={"prefix": "Today I feel"}`.
                 The actual prompt will be:
                 "Complete the following sentence: Today I feel".
-                More examples can be found at [templating](/docs/Use-Cases/Autogen#templating).
+                More examples can be found at [templating](https://microsoft.github.io/autogen/docs/Use-Cases/enhanced_inference#templating).
             use_cache (bool, Optional): Whether to use cached responses.
             config_list (List, Optional): List of configurations for the completion to try.
                 The first one that does not raise an error will be used.
@@ -746,7 +751,11 @@ class Completion(openai_Completion):
                 Also, the "prompt" or "messages" parameter can contain a template (str or Callable) which will be instantiated with the context.
 
         Returns:
-            Responses from OpenAI API.
+            Responses from OpenAI API, with additional fields.
+                - `cost`: the total cost.
+            When `config_list` is provided, the response will contain a few more fields:
+                - `config_id`: the index of the config in the config_list that is used to generate the response.
+                - `pass_filter`: whether the response passes the filter function. None if no filter is provided.
         """
         if ERROR:
             raise ERROR

flaml/autogen/oai/openai_utils.py

@@ -1,7 +1,7 @@
-import os
 import json
-from typing import List, Optional, Dict, Set, Union
 import logging
+import os
+from typing import Dict, List, Optional, Set, Union
 
 NON_CACHE_KEY = ["api_key", "api_base", "api_type", "api_version"]
 

flaml/autogen/retrieve_utils.py

@@ -1,13 +1,14 @@
-from typing import List, Union, Dict, Tuple
-import os
-import requests
-from urllib.parse import urlparse
 import glob
-import tiktoken
-import chromadb
-from chromadb.api import API
-import chromadb.utils.embedding_functions as ef
 import logging
+import os
+from typing import Dict, List, Tuple, Union
+from urllib.parse import urlparse
+
+import chromadb
+import chromadb.utils.embedding_functions as ef
+import requests
+import tiktoken
+from chromadb.api import API
 
 logger = logging.getLogger(__name__)
 TEXT_FORMATS = ["txt", "json", "csv", "tsv", "md", "html", "htm", "rtf", "rst", "jsonl", "log", "xml", "yaml", "yml"]
@@ -125,7 +126,7 @@ def split_files_to_chunks(
     """Split a list of files into chunks of max_tokens."""
     chunks = []
     for file in files:
-        with open(file, "r") as f:
+        with open(file) as f:
             text = f.read()
         chunks += split_text_to_chunks(text, max_tokens, chunk_mode, must_break_at_empty_line)
     return chunks

flaml/automl/__init__.py

@@ -1,5 +1,9 @@
-from flaml.automl.automl import AutoML, size
 from flaml.automl.logger import logger_formatter
-from flaml.automl.state import SearchState, AutoMLState
 
-__all__ = ["AutoML", "AutoMLState", "SearchState", "logger_formatter", "size"]
+try:
+    from flaml.automl.automl import AutoML, size
+    from flaml.automl.state import AutoMLState, SearchState
+
+    __all__ = ["AutoML", "AutoMLState", "SearchState", "logger_formatter", "size"]
+except ImportError:
+    __all__ = ["logger_formatter"]

flaml/automl/automl.py

@@ -3,40 +3,43 @@
 #  * Licensed under the MIT License. See LICENSE file in the
 #  * project root for license information.
 from __future__ import annotations
-import time
-import os
-import sys
-from typing import Callable, List, Union, Optional
-from functools import partial
-import numpy as np
-import logging
+
 import json
+import logging
+import os
+import random
+import sys
+import time
+from concurrent.futures import as_completed
+from functools import partial
+from typing import Callable, List, Optional, Union
 
-from flaml.automl.state import SearchState, AutoMLState
-from flaml.automl.ml import train_estimator
+import numpy as np
 
-from flaml.automl.time_series import TimeSeriesDataset
-from flaml.config import (
-    MIN_SAMPLE_TRAIN,
-    MEM_THRES,
-    RANDOM_SEED,
-    SMALL_LARGE_THRES,
-    CV_HOLDOUT_THRESHOLD,
-    SPLIT_RATIO,
-    N_SPLITS,
-    SAMPLE_MULTIPLY_FACTOR,
-)
+from flaml import tune
+from flaml.automl.logger import logger, logger_formatter
+from flaml.automl.ml import huggingface_metric_to_mode, sklearn_metric_name_set, spark_metric_name_dict, train_estimator
+from flaml.automl.spark import DataFrame, Series, psDataFrame, psSeries
+from flaml.automl.state import AutoMLState, SearchState
+from flaml.automl.task.factory import task_factory
 
 # TODO check to see when we can remove these
 from flaml.automl.task.task import CLASSIFICATION, Task
-from flaml.automl.task.factory import task_factory
-from flaml import tune
-from flaml.automl.logger import logger, logger_formatter
+from flaml.automl.time_series import TimeSeriesDataset
 from flaml.automl.training_log import training_log_reader, training_log_writer
+from flaml.config import (
+    CV_HOLDOUT_THRESHOLD,
+    MEM_THRES,
+    MIN_SAMPLE_TRAIN,
+    N_SPLITS,
+    RANDOM_SEED,
+    SAMPLE_MULTIPLY_FACTOR,
+    SMALL_LARGE_THRES,
+    SPLIT_RATIO,
+)
 from flaml.default import suggest_learner
-from flaml.version import __version__ as flaml_version
-from flaml.automl.spark import psDataFrame, psSeries, DataFrame, Series
 from flaml.tune.spark.utils import check_spark, get_broadcast_data
+from flaml.version import __version__ as flaml_version
 
 ERROR = (
     DataFrame is None and ImportError("please install flaml[automl] option to use the flaml.automl package.") or None
@@ -44,6 +47,7 @@ ERROR = (
 
 try:
     from sklearn.base import BaseEstimator
+    from sklearn.pipeline import Pipeline
 except ImportError:
     BaseEstimator = object
     ERROR = ERROR or ImportError("please install flaml[automl] option to use the flaml.automl package.")
@@ -53,6 +57,14 @@ try:
 except ImportError:
     mlflow = None
 
+try:
+    from flaml.fabric.mlflow import MLflowIntegration, get_mlflow_log_latency, infer_signature, is_autolog_enabled
+
+    internal_mlflow = True
+except ImportError:
+    internal_mlflow = False
+
+
 try:
     from ray import __version__ as ray_version
 
@@ -170,15 +182,22 @@ class AutoML(BaseEstimator):
                 'better' only logs configs with better loss than previos iters
                 'all' logs all the tried configs.
             model_history: A boolean of whether to keep the best
-                model per estimator. Make sure memory is large enough if setting to True.
+                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
                 metric for each model.
             mem_thres: A float of the memory size constraint in bytes.
             pred_time_limit: A float of the prediction latency constraint in seconds.
                 It refers to the average prediction time per row in validation data.
-            train_time_limit: A float of the training time constraint in seconds.
+            train_time_limit: None or a float of the training time constraint in seconds for each trial.
+                Only valid for sequential search.
             verbose: int, default=3 | Controls the verbosity, higher means more
                 messages.
+                verbose=0: logger level = CRITICAL
+                verbose=1: logger level = ERROR
+                verbose=2: logger level = WARNING
+                verbose=3: logger level = INFO
+                verbose=4: logger level = DEBUG
+                verbose>5: logger level = NOTSET
             retrain_full: bool or str, default=True | whether to retrain the
                 selected model on the full training data when using holdout.
                 True - retrain only after search finishes; False - no retraining;
@@ -192,7 +211,7 @@ class AutoML(BaseEstimator):
                 * Valid str options depend on different tasks.
                 For classification tasks, valid choices are
                     ["auto", 'stratified', 'uniform', 'time', 'group']. "auto" -> stratified.
-                For regression tasks, valid choices are ["auto", 'uniform', 'time'].
+                For regression tasks, valid choices are ["auto", 'uniform', 'time', 'group'].
                     "auto" -> uniform.
                 For time series forecast tasks, must be "auto" or 'time'.
                 For ranking task, must be "auto" or 'group'.
@@ -211,9 +230,9 @@ class AutoML(BaseEstimator):
                     - if "data:path" use data-dependent defaults which are stored at path;
                     - if "static", use data-independent defaults.
                 If dict, keys are the name of the estimators, and values are the starting
-                hyperparamter configurations for the corresponding estimators.
-                The value can be a single hyperparamter configuration dict or a list
-                of hyperparamter configuration dicts.
+                hyperparameter configurations for the corresponding estimators.
+                The value can be a single hyperparameter configuration dict or a list
+                of hyperparameter configuration dicts.
                 In the following code example, we get starting_points from the
                 `automl` object and use them in the `new_automl` object.
                 e.g.,
@@ -246,6 +265,9 @@ class AutoML(BaseEstimator):
                 search is considered to converge.
             force_cancel: boolean, default=False | Whether to forcely cancel Spark jobs if the
                 search time exceeded the time budget.
+            mlflow_exp_name: str, default=None | The name of the mlflow experiment. This should be specified if
+                enable mlflow autologging on Spark. Otherwise it will log all the results into the experiment of the
+                same name as the basename of main entry file.
             append_log: boolean, default=False | Whetehr to directly append the log
                 records to the input log file if it exists.
             auto_augment: boolean, default=True | Whether to automatically
@@ -319,9 +341,7 @@ class AutoML(BaseEstimator):
             }
         }
         ```
-            mlflow_logging: boolean, default=True | Whether to log the training results to mlflow.
-                This requires mlflow to be installed and to have an active mlflow run.
-                FLAML will create nested runs.
+            mlflow_logging: boolean, default=True | Whether to log the training results to mlflow. Not valid if mlflow is not installed.
 
         """
         if ERROR:
@@ -330,6 +350,8 @@ class AutoML(BaseEstimator):
         self._state = AutoMLState()
         self._state.learner_classes = {}
         self._settings = settings
+        self._automl_user_configurations = settings.copy()
+        self._settings.pop("automl_user_configurations", None)
         # no budget by default
         settings["time_budget"] = settings.get("time_budget", -1)
         settings["task"] = settings.get("task", "classification")
@@ -361,6 +383,7 @@ class AutoML(BaseEstimator):
         settings["preserve_checkpoint"] = settings.get("preserve_checkpoint", True)
         settings["early_stop"] = settings.get("early_stop", False)
         settings["force_cancel"] = settings.get("force_cancel", False)
+        settings["mlflow_exp_name"] = settings.get("mlflow_exp_name", None)
         settings["append_log"] = settings.get("append_log", False)
         settings["min_sample_size"] = settings.get("min_sample_size", MIN_SAMPLE_TRAIN)
         settings["use_ray"] = settings.get("use_ray", False)
@@ -376,6 +399,7 @@ class AutoML(BaseEstimator):
         settings["mlflow_logging"] = settings.get("mlflow_logging", True)
 
         self._estimator_type = "classifier" if settings["task"] in CLASSIFICATION else "regressor"
+        self.best_run_id = None
 
     def get_params(self, deep: bool = False) -> dict:
         return self._settings.copy()
@@ -408,6 +432,8 @@ class AutoML(BaseEstimator):
             If `model_history` was set to True, then the returned model is trained.
         """
         state = self._search_states.get(estimator_name)
+        if state and estimator_name == self._best_estimator:
+            return self.model
         return state and getattr(state, "trained_estimator", None)
 
     @property
@@ -474,10 +500,25 @@ class AutoML(BaseEstimator):
         with open(filename, "w") as f:
             json.dump(best, f)
 
+    @property
+    def supported_metrics(self):
+        """
+        Returns a tuple of supported metrics for the task.
+
+            Returns:
+                    metrics (Tuple): sklearn metrics from sklearn package;
+                                    huggingface metrics from datasets package;
+                                    spark metrics from pyspark package
+
+        """
+
+        return sklearn_metric_name_set, huggingface_metric_to_mode.keys(), spark_metric_name_dict
+
     @property
     def feature_transformer(self):
         """Returns AutoML Transformer"""
-        return getattr(self, "_transformer", None)
+        data_precessor = getattr(self, "_transformer", None)
+        return data_precessor
 
     @property
     def label_transformer(self):
@@ -520,8 +561,8 @@ class AutoML(BaseEstimator):
 
     def score(
         self,
-        X: Union[DataFrame, psDataFrame],
-        y: Union[Series, psSeries],
+        X: DataFrame | psDataFrame,
+        y: Series | psSeries,
         **kwargs,
     ):
         estimator = getattr(self, "_trained_estimator", None)
@@ -535,7 +576,7 @@ class AutoML(BaseEstimator):
 
     def predict(
         self,
-        X: Union[np.array, DataFrame, List[str], List[List[str]], psDataFrame],
+        X: np.array | DataFrame | list[str] | list[list[str]] | psDataFrame,
         **pred_kwargs,
     ):
         """Predict label from features.
@@ -606,11 +647,11 @@ class AutoML(BaseEstimator):
 
         Args:
             learner_name: A string of the learner's name.
-            learner_class: A subclass of flaml.model.BaseEstimator.
+            learner_class: A subclass of flaml.automl.model.BaseEstimator.
         """
         self._state.learner_classes[learner_name] = learner_class
 
-    def get_estimator_from_log(self, log_file_name: str, record_id: int, task: Union[str, Task]):
+    def get_estimator_from_log(self, log_file_name: str, record_id: int, task: str | Task):
         """Get the estimator from log file.
 
         Args:
@@ -652,7 +693,7 @@ class AutoML(BaseEstimator):
         dataframe=None,
         label=None,
         time_budget=np.inf,
-        task: Optional[Union[str, Task]] = None,
+        task: str | Task | None = None,
         eval_method=None,
         split_ratio=None,
         n_splits=None,
@@ -708,7 +749,7 @@ class AutoML(BaseEstimator):
                 * Valid str options depend on different tasks.
                 For classification tasks, valid choices are
                     ["auto", 'stratified', 'uniform', 'time', 'group']. "auto" -> stratified.
-                For regression tasks, valid choices are ["auto", 'uniform', 'time'].
+                For regression tasks, valid choices are ["auto", 'uniform', 'time', 'group'].
                     "auto" -> uniform.
                 For time series forecast tasks, must be "auto" or 'time'.
                 For ranking task, must be "auto" or 'group'.
@@ -778,7 +819,7 @@ class AutoML(BaseEstimator):
                     max_epochs: int, default = 20 | Maximum number of epochs to run training,
                         only used by TemporalFusionTransformerEstimator.
                     batch_size: int, default = 64 | Batch size for training model, only
-                        used by TemporalFusionTransformerEstimator.
+                        used by TemporalFusionTransformerEstimator and TCNEstimator.
         """
         task = task or self._settings.get("task")
         if isinstance(task, str):
@@ -801,7 +842,7 @@ class AutoML(BaseEstimator):
         )
         task.validate_data(self, self._state, X_train, y_train, dataframe, label, groups=groups)
 
-        logger.info("log file name {}".format(log_file_name))
+        logger.info(f"log file name {log_file_name}")
 
         best_config = None
         best_val_loss = float("+inf")
@@ -854,9 +895,7 @@ class AutoML(BaseEstimator):
         else:
             self._state.fit_kwargs_by_estimator[best_estimator] = self._state.fit_kwargs
 
-        logger.info(
-            "estimator = {}, config = {}, #training instances = {}".format(best_estimator, best_config, sample_size)
-        )
+        logger.info(f"estimator = {best_estimator}, config = {best_config}, #training instances = {sample_size}")
         # Partially copied from fit() function
         # Initilize some attributes required for retrain_from_log
         self._split_type = task.decide_split_type(
@@ -1027,7 +1066,7 @@ class AutoML(BaseEstimator):
         return points
 
     @property
-    def resource_attr(self) -> Optional[str]:
+    def resource_attr(self) -> str | None:
         """Attribute of the resource dimension.
 
         Returns:
@@ -1037,7 +1076,7 @@ class AutoML(BaseEstimator):
         return "FLAML_sample_size" if self._sample else None
 
     @property
-    def min_resource(self) -> Optional[float]:
+    def min_resource(self) -> float | None:
         """Attribute for pruning.
 
         Returns:
@@ -1046,7 +1085,7 @@ class AutoML(BaseEstimator):
         return self._min_sample_size if self._sample else None
 
     @property
-    def max_resource(self) -> Optional[float]:
+    def max_resource(self) -> float | None:
         """Attribute for pruning.
 
         Returns:
@@ -1068,7 +1107,7 @@ class AutoML(BaseEstimator):
             pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)
 
     @property
-    def trainable(self) -> Callable[[dict], Optional[float]]:
+    def trainable(self) -> Callable[[dict], float | None]:
         """Training function.
         Returns:
             A function that evaluates each config and returns the loss.
@@ -1154,7 +1193,7 @@ class AutoML(BaseEstimator):
         dataframe=None,
         label=None,
         metric=None,
-        task: Optional[Union[str, Task]] = None,
+        task: str | Task | None = None,
         n_jobs=None,
         # gpu_per_trial=0,
         log_file_name=None,
@@ -1202,6 +1241,7 @@ class AutoML(BaseEstimator):
         skip_transform=None,
         mlflow_logging=None,
         fit_kwargs_by_estimator=None,
+        mlflow_exp_name=None,
         **fit_kwargs,
     ):
         """Find a model for a given task.
@@ -1295,14 +1335,15 @@ class AutoML(BaseEstimator):
                 'all' logs all the tried configs.
             model_history: A boolean of whether to keep the trained best
                 model per estimator. Make sure memory is large enough if setting to True.
-                Default value is False: best_model_for_estimator would return a
+                Default value is False. If False, best_model_for_estimator would return a
                 untrained model for non-best learner.
             log_training_metric: A boolean of whether to log the training
                 metric for each model.
             mem_thres: A float of the memory size constraint in bytes.
             pred_time_limit: A float of the prediction latency constraint in seconds.
                 It refers to the average prediction time per row in validation data.
-            train_time_limit: None or a float of the training time constraint in seconds.
+            train_time_limit: None or a float of the training time constraint in seconds for each trial.
+                Only valid for sequential search.
             X_val: None or a numpy array or a pandas dataframe of validation data.
             y_val: None or a numpy array or a pandas series of validation labels.
             sample_weight_val: None or a numpy array of the sample weight of
@@ -1315,6 +1356,12 @@ class AutoML(BaseEstimator):
                 for training data.
             verbose: int, default=3 | Controls the verbosity, higher means more
                 messages.
+                verbose=0: logger level = CRITICAL
+                verbose=1: logger level = ERROR
+                verbose=2: logger level = WARNING
+                verbose=3: logger level = INFO
+                verbose=4: logger level = DEBUG
+                verbose>5: logger level = NOTSET
             retrain_full: bool or str, default=True | whether to retrain the
                 selected model on the full training data when using holdout.
                 True - retrain only after search finishes; False - no retraining;
@@ -1328,7 +1375,7 @@ class AutoML(BaseEstimator):
                 * Valid str options depend on different tasks.
                 For classification tasks, valid choices are
                     ["auto", 'stratified', 'uniform', 'time', 'group']. "auto" -> stratified.
-                For regression tasks, valid choices are ["auto", 'uniform', 'time'].
+                For regression tasks, valid choices are ["auto", 'uniform', 'time', 'group'].
                     "auto" -> uniform.
                 For time series forecast tasks, must be "auto" or 'time'.
                 For ranking task, must be "auto" or 'group'.
@@ -1347,9 +1394,9 @@ class AutoML(BaseEstimator):
                     - if "data:path" use data-dependent defaults which are stored at path;
                     - if "static", use data-independent defaults.
                 If dict, keys are the name of the estimators, and values are the starting
-                hyperparamter configurations for the corresponding estimators.
-                The value can be a single hyperparamter configuration dict or a list
-                of hyperparamter configuration dicts.
+                hyperparameter configurations for the corresponding estimators.
+                The value can be a single hyperparameter configuration dict or a list
+                of hyperparameter configuration dicts.
                 In the following code example, we get starting_points from the
                 `automl` object and use them in the `new_automl` object.
                 e.g.,
@@ -1381,6 +1428,9 @@ class AutoML(BaseEstimator):
             early_stop: boolean, default=False | Whether to stop early if the
                 search is considered to converge.
             force_cancel: boolean, default=False | Whether to forcely cancel the PySpark job if overtime.
+            mlflow_exp_name: str, default=None | The name of the mlflow experiment. This should be specified if
+                enable mlflow autologging on Spark. Otherwise it will log all the results into the experiment of the
+                same name as the basename of main entry file.
             append_log: boolean, default=False | Whetehr to directly append the log
                 records to the input log file if it exists.
             auto_augment: boolean, default=True | Whether to automatically
@@ -1466,9 +1516,7 @@ class AutoML(BaseEstimator):
             skip_transform: boolean, default=False | Whether to pre-process data prior to modeling.
             mlflow_logging: boolean, default=None | Whether to log the training results to mlflow.
                 Default value is None, which means the logging decision is made based on
-                AutoML.__init__'s mlflow_logging argument.
-                This requires mlflow to be installed and to have an active mlflow run.
-                FLAML will create nested runs.
+                AutoML.__init__'s mlflow_logging argument. Not valid if mlflow is not installed.
             fit_kwargs_by_estimator: dict, default=None | The user specified keywords arguments, grouped by estimator name.
                 For TransformersEstimator, available fit_kwargs can be found from
                 [TrainingArgumentsForAuto](nlp/huggingface/training_args).
@@ -1518,7 +1566,7 @@ class AutoML(BaseEstimator):
                     max_epochs: int, default = 20 | Maximum number of epochs to run training,
                         only used by TemporalFusionTransformerEstimator.
                     batch_size: int, default = 64 | Batch size for training model, only
-                        used by TemporalFusionTransformerEstimator.
+                        used by TemporalFusionTransformerEstimator and TCNEstimator.
         """
 
         self._state._start_time_flag = self._start_time_flag = time.time()
@@ -1569,6 +1617,7 @@ class AutoML(BaseEstimator):
         )
         early_stop = self._settings.get("early_stop") if early_stop is None else early_stop
         force_cancel = self._settings.get("force_cancel") if force_cancel is None else force_cancel
+        mlflow_exp_name = self._settings.get("mlflow_exp_name") if mlflow_exp_name is None else mlflow_exp_name
         # no search budget is provided?
         no_budget = time_budget < 0 and max_iter is None and not early_stop
         append_log = self._settings.get("append_log") if append_log is None else append_log
@@ -1591,6 +1640,13 @@ class AutoML(BaseEstimator):
             _ch.setFormatter(logger_formatter)
             logger.addHandler(_ch)
 
+        if model_history:
+            logger.warning(
+                "With `model_history` set to `True` by default, all intermediate models are retained in memory, "
+                "which may significantly increase memory usage and slow down training. "
+                "Consider setting `model_history=False` to optimize memory and accelerate the training process."
+            )
+
         if not use_ray and not use_spark and n_concurrent_trials > 1:
             if ray_available:
                 logger.warning(
@@ -1621,7 +1677,6 @@ class AutoML(BaseEstimator):
         self._use_ray = use_ray
         # use the following condition if we have an estimation of average_trial_time and average_trial_overhead
         # self._use_ray = use_ray or n_concurrent_trials > ( average_trial_time + average_trial_overhead) / (average_trial_time)
-
         if self._use_ray is not False:
             import ray
 
@@ -1655,11 +1710,29 @@ class AutoML(BaseEstimator):
         self._state.fit_kwargs = fit_kwargs
         custom_hp = custom_hp or self._settings.get("custom_hp")
         self._skip_transform = self._settings.get("skip_transform") if skip_transform is None else skip_transform
-        self._mlflow_logging = self._settings.get("mlflow_logging") if mlflow_logging is None else mlflow_logging
+        self._mlflow_logging = (
+            False
+            if mlflow is None
+            else self._settings.get("mlflow_logging")
+            if mlflow_logging is None
+            else mlflow_logging
+        )
         fit_kwargs_by_estimator = fit_kwargs_by_estimator or self._settings.get("fit_kwargs_by_estimator")
         self._state.fit_kwargs_by_estimator = fit_kwargs_by_estimator.copy()  # shallow copy of fit_kwargs_by_estimator
         self._state.weight_val = sample_weight_val
-
+        self._mlflow_exp_name = mlflow_exp_name
+        self.mlflow_integration = None
+        self.autolog_extra_tag = {
+            "extra_tag.sid": f"flaml_{flaml_version}_{int(time.time())}_{random.randint(1001, 9999)}"
+        }
+        if internal_mlflow and self._mlflow_logging and (mlflow.active_run() or is_autolog_enabled()):
+            try:
+                self.mlflow_integration = MLflowIntegration("automl", mlflow_exp_name, extra_tag=self.autolog_extra_tag)
+                self._mlflow_exp_name = self.mlflow_integration.experiment_name
+                if not (mlflow.active_run() is not None or is_autolog_enabled()):
+                    self.mlflow_integration.only_history = True
+            except KeyError:
+                logger.info("Not in Fabric, Skipped")
         task.validate_data(
             self,
             self._state,
@@ -1687,7 +1760,7 @@ class AutoML(BaseEstimator):
             logger.info(f"Data split method: {self._split_type}")
         eval_method = self._decide_eval_method(eval_method, time_budget)
         self._state.eval_method = eval_method
-        logger.info("Evaluation method: {}".format(eval_method))
+        logger.info(f"Evaluation method: {eval_method}")
         self._state.cv_score_agg_func = cv_score_agg_func or self._settings.get("cv_score_agg_func")
 
         self._retrain_in_budget = retrain_full == "budget" and (eval_method == "holdout" and self._state.X_val is None)
@@ -1704,13 +1777,9 @@ class AutoML(BaseEstimator):
                 if sample_size:
                     _sample_size_from_starting_points[_estimator] = sample_size
                 elif _point_per_estimator and isinstance(_point_per_estimator, list):
-                    _sample_size_set = set(
-                        [
-                            config["FLAML_sample_size"]
-                            for config in _point_per_estimator
-                            if "FLAML_sample_size" in config
-                        ]
-                    )
+                    _sample_size_set = {
+                        config["FLAML_sample_size"] for config in _point_per_estimator if "FLAML_sample_size" in config
+                    }
                     if _sample_size_set:
                         _sample_size_from_starting_points[_estimator] = min(_sample_size_set)
                     if len(_sample_size_set) > 1:
@@ -1728,6 +1797,11 @@ class AutoML(BaseEstimator):
         self._min_sample_size_input = min_sample_size
         self._prepare_data(eval_method, split_ratio, n_splits)
 
+        # infer the signature of the input/output data
+        if self.mlflow_integration is not None:
+            self.estimator_signature = infer_signature(self._state.X_train, self._state.y_train)
+            self.pipeline_signature = infer_signature(X_train, y_train, dataframe, label)
+
         # TODO pull this to task as decide_sample_size
         if isinstance(self._min_sample_size, dict):
             self._sample = {
@@ -1826,6 +1900,11 @@ class AutoML(BaseEstimator):
             and (max_iter > 0 or retrain_full is True)
             or max_iter == 1
         )
+        if self.mlflow_integration is not None and all(
+            [self.mlflow_integration.parent_run_id is None, not self.mlflow_integration.only_history]
+        ):
+            # force not retrain if no active run
+            self._state.retrain_final = False
         # add custom learner
         for estimator_name in estimator_list:
             if estimator_name not in self._state.learner_classes:
@@ -1897,7 +1976,7 @@ class AutoML(BaseEstimator):
                 max_iter=max_iter / len(estimator_list) if self._learner_selector == "roundrobin" else max_iter,
                 budget=self._state.time_budget,
             )
-        logger.info("List of ML learners in AutoML Run: {}".format(estimator_list))
+        logger.info(f"List of ML learners in AutoML Run: {estimator_list}")
         self.estimator_list = estimator_list
         self._active_estimators = estimator_list.copy()
         self._ensemble = ensemble
@@ -1939,7 +2018,7 @@ class AutoML(BaseEstimator):
                 )
             ):
                 logger.warning(
-                    "Time taken to find the best model is {0:.0f}% of the "
+                    "Time taken to find the best model is {:.0f}% of the "
                     "provided time budget and not all estimators' hyperparameter "
                     "search converged. Consider increasing the time budget.".format(
                         self._time_taken_best_iter / self._state.time_budget * 100
@@ -1958,6 +2037,8 @@ class AutoML(BaseEstimator):
             )  # NOTE: this is after kwargs is updated to fit_kwargs_by_estimator
             del self._state.groups, self._state.groups_all, self._state.groups_val
         logger.setLevel(old_level)
+        if self.mlflow_integration is not None:
+            self.mlflow_integration.resume_mlflow()
 
     def _search_parallel(self):
         if self._use_ray is not False:
@@ -2054,6 +2135,14 @@ class AutoML(BaseEstimator):
 
         if self._use_spark:
             # use spark as parallel backend
+            mlflow_log_latency = (
+                get_mlflow_log_latency(model_history=self._state.model_history) if self.mlflow_integration else 0
+            )
+            (
+                logger.info(f"Estimated mlflow_log_latency: {mlflow_log_latency} seconds.")
+                if mlflow_log_latency > 0
+                else None
+            )
             analysis = tune.run(
                 self.trainable,
                 search_alg=search_alg,
@@ -2066,6 +2155,9 @@ class AutoML(BaseEstimator):
                 use_ray=False,
                 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
+                extra_tag=self.autolog_extra_tag,
                 # raise_on_failed_trial=False,
                 # keep_checkpoints_num=1,
                 # checkpoint_score_attr="min-val_loss",
@@ -2126,6 +2218,8 @@ class AutoML(BaseEstimator):
                     self._search_states[estimator].best_config = config
                 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)
 
     def _log_trial(self, search_state, estimator):
         if self._training_log:
@@ -2139,36 +2233,6 @@ class AutoML(BaseEstimator):
                 estimator,
                 search_state.sample_size,
             )
-        if self._mlflow_logging and mlflow is not None and mlflow.active_run():
-            with mlflow.start_run(nested=True):
-                mlflow.log_metric("iter_counter", self._track_iter)
-                if (search_state.metric_for_logging is not None) and (
-                    "intermediate_results" in search_state.metric_for_logging
-                ):
-                    for each_entry in search_state.metric_for_logging["intermediate_results"]:
-                        with mlflow.start_run(nested=True):
-                            mlflow.log_metrics(each_entry)
-                            mlflow.log_metric("iter_counter", self._iter_per_learner[estimator])
-                    del search_state.metric_for_logging["intermediate_results"]
-                if search_state.metric_for_logging:
-                    mlflow.log_metrics(search_state.metric_for_logging)
-                mlflow.log_metric("trial_time", search_state.trial_time)
-                mlflow.log_metric("wall_clock_time", self._state.time_from_start)
-                mlflow.log_metric("validation_loss", search_state.val_loss)
-                mlflow.log_params(search_state.config)
-                mlflow.log_param("learner", estimator)
-                mlflow.log_param("sample_size", search_state.sample_size)
-                mlflow.log_metric("best_validation_loss", search_state.best_loss)
-                mlflow.log_param("best_config", search_state.best_config)
-                mlflow.log_param("best_learner", self._best_estimator)
-                mlflow.log_metric(
-                    self._state.metric if isinstance(self._state.metric, str) else self._state.error_metric,
-                    1 - search_state.val_loss
-                    if self._state.error_metric.startswith("1-")
-                    else -search_state.val_loss
-                    if self._state.error_metric.startswith("-")
-                    else search_state.val_loss,
-                )
 
     def _search_sequential(self):
         try:
@@ -2322,10 +2386,19 @@ class AutoML(BaseEstimator):
                 verbose=max(self.verbose - 3, 0),
                 use_ray=False,
                 use_spark=False,
+                force_cancel=self._force_cancel,
+                mlflow_exp_name=self._mlflow_exp_name,
+                automl_info=(0,),  # pass automl info to tune.run
+                extra_tag=self.autolog_extra_tag,
             )
             time_used = time.time() - start_run_time
             better = False
-            if analysis.trials:
+            (
+                logger.debug(f"result in automl: {analysis.trials}, {analysis.trials[-1].last_result}")
+                if analysis.trials
+                else logger.debug("result in automl: [], None")
+            )
+            if analysis.trials and analysis.trials[-1].last_result:
                 result = analysis.trials[-1].last_result
                 search_state.update(result, time_used=time_used)
                 if self._estimator_index is None:
@@ -2387,6 +2460,8 @@ class AutoML(BaseEstimator):
                     search_state.trained_estimator.cleanup()
                 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)
 
                 logger.info(
                     " at {:.1f}s,\testimator {}'s best error={:.4f},\tbest estimator {}'s best error={:.4f}".format(
@@ -2439,7 +2514,7 @@ class AutoML(BaseEstimator):
                     state.best_config,
                     self.data_size_full,
                 )
-                logger.info("retrain {} for {:.1f}s".format(self._best_estimator, retrain_time))
+                logger.info(f"retrain {self._best_estimator} for {retrain_time:.1f}s")
                 self._retrained_config[best_config_sig] = state.best_config_train_time = retrain_time
                 est_retrain_time = 0
             self._state.time_from_start = time.time() - self._start_time_flag
@@ -2461,8 +2536,8 @@ class AutoML(BaseEstimator):
         self._time_taken_best_iter = 0
         self._config_history = {}
         self._max_iter_per_learner = 10000
-        self._iter_per_learner = dict([(e, 0) for e in self.estimator_list])
-        self._iter_per_learner_fullsize = dict([(e, 0) for e in self.estimator_list])
+        self._iter_per_learner = {e: 0 for e in self.estimator_list}
+        self._iter_per_learner_fullsize = {e: 0 for e in self.estimator_list}
         self._fullsize_reached = False
         self._trained_estimator = None
         self._best_estimator = None
@@ -2478,6 +2553,21 @@ class AutoML(BaseEstimator):
             self._selected = state = self._search_states[estimator]
             state.best_config_sample_size = self._state.data_size[0]
             state.best_config = state.init_config[0] if state.init_config else {}
+            self._track_iter = 0
+            self._config_history[self._track_iter] = (estimator, state.best_config, self._state.time_from_start)
+            self._best_iteration = self._track_iter
+            state.val_loss = getattr(state, "val_loss", float("inf"))
+            state.best_loss = getattr(state, "best_loss", float("inf"))
+            state.config = getattr(state, "config", state.best_config.copy())
+            state.metric_for_logging = getattr(state, "metric_for_logging", None)
+            state.sample_size = getattr(state, "sample_size", self._state.data_size[0])
+            state.learner_class = getattr(state, "learner_class", self._state.learner_classes.get(estimator))
+            if hasattr(self, "mlflow_integration") and self.mlflow_integration:
+                self.mlflow_integration.record_state(
+                    automl=self,
+                    search_state=state,
+                    estimator=estimator,
+                )
         elif self._use_ray is False and self._use_spark is False:
             self._search_sequential()
         else:
@@ -2487,6 +2577,12 @@ class AutoML(BaseEstimator):
             self._training_log.checkpoint()
         self._state.time_from_start = time.time() - self._start_time_flag
         if self._best_estimator:
+            if self.mlflow_integration:
+                self.mlflow_integration.log_automl(self)
+                if mlflow.active_run() is None:
+                    if self.mlflow_integration.parent_run_id is not None and self.mlflow_integration.autolog:
+                        # ensure result of retrain autolog to parent run
+                        mlflow.start_run(run_id=self.mlflow_integration.parent_run_id)
             self._selected = self._search_states[self._best_estimator]
             self.modelcount = sum(search_state.total_iter for search_state in self._search_states.values())
             if self._trained_estimator:
@@ -2623,13 +2719,67 @@ class AutoML(BaseEstimator):
                         self._best_estimator,
                         state.best_config,
                         self.data_size_full,
+                        is_retrain=True,
                     )
-                    logger.info("retrain {} for {:.1f}s".format(self._best_estimator, retrain_time))
+                    logger.info(f"retrain {self._best_estimator} for {retrain_time:.1f}s")
                     state.best_config_train_time = retrain_time
                     if self._trained_estimator:
                         logger.info(f"retrained model: {self._trained_estimator.model}")
+                        if self.best_run_id is not None:
+                            logger.info(f"Best MLflow run name: {self.best_run_name}")
+                            logger.info(f"Best MLflow run id: {self.best_run_id}")
+                        if self.mlflow_integration is not None:
+                            # try log retrained model
+                            if all(
+                                [
+                                    self.mlflow_integration.manual_log,
+                                    not self.mlflow_integration.has_model,
+                                    self.mlflow_integration.parent_run_id is not None,
+                                ]
+                            ):
+                                if mlflow.active_run() is None:
+                                    mlflow.start_run(run_id=self.mlflow_integration.parent_run_id)
+                                if self.best_estimator.endswith("_spark"):
+                                    self.mlflow_integration.log_model(
+                                        self._trained_estimator.model,
+                                        self.best_estimator,
+                                        signature=self.estimator_signature,
+                                        run_id=self.mlflow_integration.parent_run_id,
+                                    )
+                                else:
+                                    self.mlflow_integration.pickle_and_log_automl_artifacts(
+                                        self,
+                                        self.model,
+                                        self.best_estimator,
+                                        signature=self.pipeline_signature,
+                                        run_id=self.mlflow_integration.parent_run_id,
+                                    )
                 else:
-                    logger.info("not retraining because the time budget is too small.")
+                    logger.warning("not retraining because the time budget is too small.")
+        self.wait_futures()
+
+    def wait_futures(self):
+        if self.mlflow_integration is not None:
+            logger.debug("Collecting results from submitted record_state tasks")
+            t1 = time.perf_counter()
+            for future in as_completed(self.mlflow_integration.futures):
+                _task = self.mlflow_integration.futures[future]
+                try:
+                    result = future.result()
+                    logger.debug(f"Result for record_state task {_task}: {result}")
+                except Exception as e:
+                    logger.warning(f"Exception for record_state task {_task}: {e}")
+            for future in as_completed(self.mlflow_integration.futures_log_model):
+                _task = self.mlflow_integration.futures_log_model[future]
+                try:
+                    result = future.result()
+                    logger.debug(f"Result for log_model task {_task}: {result}")
+                except Exception as e:
+                    logger.warning(f"Exception for log_model task {_task}: {e}")
+            t2 = time.perf_counter()
+            logger.debug(f"Collecting results from tasks submitted to executors costs {t2-t1} seconds.")
+        else:
+            logger.debug("No futures to wait for.")
 
     def __del__(self):
         if (
@@ -2647,7 +2797,7 @@ class AutoML(BaseEstimator):
             if self._estimator_index == len(estimator_list):
                 self._estimator_index = 0
             return estimator_list[self._estimator_index]
-        min_estimated_cost, selected = np.Inf, None
+        min_estimated_cost, selected = np.inf, None
         inv = []
         untried_exists = False
         for i, estimator in enumerate(estimator_list):
@@ -2701,3 +2851,7 @@ class AutoML(BaseEstimator):
                 q += inv[i] / s
                 if p < q:
                     return estimator_list[i]
+
+    @property
+    def automl_pipeline(self):
+        return None

flaml/automl/contrib/__init__.py

@@ -0,0 +1 @@
+from .histgb import HistGradientBoostingEstimator

flaml/automl/contrib/histgb.py

@@ -0,0 +1,75 @@
+try:
+    from sklearn.ensemble import HistGradientBoostingClassifier, HistGradientBoostingRegressor
+except ImportError as e:
+    print(f"scikit-learn is required for HistGradientBoostingEstimator. Please install it; error: {e}")
+
+from flaml import tune
+from flaml.automl.model import SKLearnEstimator
+from flaml.automl.task import Task
+
+
+class HistGradientBoostingEstimator(SKLearnEstimator):
+    """The class for tuning Histogram Gradient Boosting."""
+
+    ITER_HP = "max_iter"
+    HAS_CALLBACK = False
+    DEFAULT_ITER = 100
+
+    @classmethod
+    def search_space(cls, data_size: int, task, **params) -> dict:
+        upper = max(5, min(32768, int(data_size[0])))  # upper must be larger than lower
+        return {
+            "n_estimators": {
+                "domain": tune.lograndint(lower=4, upper=upper),
+                "init_value": 4,
+                "low_cost_init_value": 4,
+            },
+            "max_leaves": {
+                "domain": tune.lograndint(lower=4, upper=upper),
+                "init_value": 4,
+                "low_cost_init_value": 4,
+            },
+            "min_samples_leaf": {
+                "domain": tune.lograndint(lower=2, upper=2**7 + 1),
+                "init_value": 20,
+            },
+            "learning_rate": {
+                "domain": tune.loguniform(lower=1 / 1024, upper=1.0),
+                "init_value": 0.1,
+            },
+            "log_max_bin": {  # log transformed with base 2, <= 256
+                "domain": tune.lograndint(lower=3, upper=9),
+                "init_value": 8,
+            },
+            "l2_regularization": {
+                "domain": tune.loguniform(lower=1 / 1024, upper=1024),
+                "init_value": 1.0,
+            },
+        }
+
+    def config2params(self, config: dict) -> dict:
+        params = super().config2params(config)
+        if "log_max_bin" in params:
+            params["max_bins"] = (1 << params.pop("log_max_bin")) - 1
+        if "max_leaves" in params:
+            params["max_leaf_nodes"] = params.get("max_leaf_nodes", params.pop("max_leaves"))
+        if "n_estimators" in params:
+            params["max_iter"] = params.get("max_iter", params.pop("n_estimators"))
+        if "random_state" not in params:
+            params["random_state"] = 24092023
+        if "n_jobs" in params:
+            params.pop("n_jobs")
+        return params
+
+    def __init__(
+        self,
+        task: Task,
+        **config,
+    ):
+        super().__init__(task, **config)
+        self.params["verbose"] = 0
+
+        if self._task.is_classification():
+            self.estimator_class = HistGradientBoostingClassifier
+        else:
+            self.estimator_class = HistGradientBoostingRegressor

flaml/automl/data.py

@@ -2,21 +2,28 @@
 #  * Copyright (c) Microsoft Corporation. All rights reserved.
 #  * Licensed under the MIT License. See LICENSE file in the
 #  * project root for license information.
-import numpy as np
-from datetime import datetime
-from typing import TYPE_CHECKING, Union
+import json
 import os
+import random
+import uuid
+from datetime import datetime, timedelta
+from decimal import ROUND_HALF_UP, Decimal
+from typing import TYPE_CHECKING, Union
+
+import numpy as np
+
+from flaml.automl.spark import DataFrame, F, Series, T, pd, ps, psDataFrame, psSeries
 from flaml.automl.training_log import training_log_reader
-from flaml.automl.spark import ps, psDataFrame, psSeries, DataFrame, Series, pd
 
 try:
-    from scipy.sparse import vstack, issparse
+    from scipy.sparse import issparse, vstack
 except ImportError:
     pass
 
 if TYPE_CHECKING:
     from flaml.automl.task import Task
 
+
 TS_TIMESTAMP_COL = "ds"
 TS_VALUE_COL = "y"
 
@@ -41,8 +48,9 @@ def load_openml_dataset(dataset_id, data_dir=None, random_state=0, dataset_forma
         y_train: A series or array of labels for training data.
         y_test:  A series or array of labels for test data.
     """
-    import openml
     import pickle
+
+    import openml
     from sklearn.model_selection import train_test_split
 
     filename = "openml_ds" + str(dataset_id) + ".pkl"
@@ -93,9 +101,10 @@ def load_openml_task(task_id, data_dir):
         y_train: A series of labels for training data.
         y_test:  A series of labels for test data.
     """
-    import openml
     import pickle
 
+    import openml
+
     task = openml.tasks.get_task(task_id)
     filename = "openml_task" + str(task_id) + ".pkl"
     filepath = os.path.join(data_dir, filename)
@@ -289,7 +298,7 @@ class DataTransformer:
                     y = y.rename(TS_VALUE_COL)
             for column in X.columns:
                 # sklearn\utils\validation.py needs int/float values
-                if X[column].dtype.name in ("object", "category"):
+                if X[column].dtype.name in ("object", "category", "string"):
                     if X[column].nunique() == 1 or X[column].nunique(dropna=True) == n - X[column].isnull().sum():
                         X.drop(columns=column, inplace=True)
                         drop = True
@@ -341,8 +350,8 @@ class DataTransformer:
                     drop = True
                 else:
                     drop = False
-                from sklearn.impute import SimpleImputer
                 from sklearn.compose import ColumnTransformer
+                from sklearn.impute import SimpleImputer
 
                 self.transformer = ColumnTransformer(
                     [
@@ -441,3 +450,331 @@ class DataTransformer:
 def group_counts(groups):
     _, i, c = np.unique(groups, return_counts=True, return_index=True)
     return c[np.argsort(i)]
+
+
+def get_random_dataframe(n_rows: int = 200, ratio_none: float = 0.1, seed: int = 42) -> DataFrame:
+    """Generate a random pandas DataFrame with various data types for testing.
+    This function creates a DataFrame with multiple column types including:
+    - Timestamps
+    - Integers
+    - Floats
+    - Categorical values
+    - Booleans
+    - Lists (tags)
+    - Decimal strings
+    - UUIDs
+    - Binary data (as hex strings)
+    - JSON blobs
+    - Nullable text fields
+    Parameters
+    ----------
+    n_rows : int, default=200
+        Number of rows in the generated DataFrame
+    ratio_none : float, default=0.1
+        Probability of generating None values in applicable columns
+    seed : int, default=42
+        Random seed for reproducibility
+    Returns
+    -------
+    pd.DataFrame
+        A DataFrame with 14 columns of various data types
+    Examples
+    --------
+    >>> df = get_random_dataframe(100, 0.05, 123)
+    >>> df.shape
+    (100, 14)
+    >>> df.dtypes
+    timestamp       datetime64[ns]
+    id                       int64
+    score                  float64
+    status                  object
+    flag                    object
+    count                   object
+    value                   object
+    tags                    object
+    rating                  object
+    uuid                    object
+    binary                  object
+    json_blob               object
+    category              category
+    nullable_text           object
+    dtype: object
+    """
+
+    np.random.seed(seed)
+    random.seed(seed)
+
+    def random_tags():
+        tags = ["AI", "ML", "data", "robotics", "vision"]
+        return random.sample(tags, k=random.randint(1, 3)) if random.random() > ratio_none else None
+
+    def random_decimal():
+        return (
+            str(Decimal(random.uniform(1, 5)).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP))
+            if random.random() > ratio_none
+            else None
+        )
+
+    def random_json_blob():
+        blob = {"a": random.randint(1, 10), "b": random.random()}
+        return json.dumps(blob) if random.random() > ratio_none else None
+
+    def random_binary():
+        return bytes(random.randint(0, 255) for _ in range(4)).hex() if random.random() > ratio_none else None
+
+    data = {
+        "timestamp": [
+            datetime(2020, 1, 1) + timedelta(days=np.random.randint(0, 1000)) if np.random.rand() > ratio_none else None
+            for _ in range(n_rows)
+        ],
+        "id": range(1, n_rows + 1),
+        "score": np.random.uniform(0, 100, n_rows),
+        "status": np.random.choice(
+            ["active", "inactive", "pending", None],
+            size=n_rows,
+            p=[(1 - ratio_none) / 3, (1 - ratio_none) / 3, (1 - ratio_none) / 3, ratio_none],
+        ),
+        "flag": np.random.choice(
+            [True, False, None], size=n_rows, p=[(1 - ratio_none) / 2, (1 - ratio_none) / 2, ratio_none]
+        ),
+        "count": [np.random.randint(0, 100) if np.random.rand() > ratio_none else None for _ in range(n_rows)],
+        "value": [round(np.random.normal(50, 15), 2) if np.random.rand() > ratio_none else None for _ in range(n_rows)],
+        "tags": [random_tags() for _ in range(n_rows)],
+        "rating": [random_decimal() for _ in range(n_rows)],
+        "uuid": [str(uuid.uuid4()) if np.random.rand() > ratio_none else None for _ in range(n_rows)],
+        "binary": [random_binary() for _ in range(n_rows)],
+        "json_blob": [random_json_blob() for _ in range(n_rows)],
+        "category": pd.Categorical(
+            np.random.choice(
+                ["A", "B", "C", None],
+                size=n_rows,
+                p=[(1 - ratio_none) / 3, (1 - ratio_none) / 3, (1 - ratio_none) / 3, ratio_none],
+            )
+        ),
+        "nullable_text": [random.choice(["Good", "Bad", "Average", None]) for _ in range(n_rows)],
+    }
+
+    return pd.DataFrame(data)
+
+
+def auto_convert_dtypes_spark(
+    df: psDataFrame,
+    na_values: list = None,
+    category_threshold: float = 0.3,
+    convert_threshold: float = 0.6,
+    sample_ratio: float = 0.1,
+) -> tuple[psDataFrame, dict]:
+    """Automatically convert data types in a PySpark DataFrame using heuristics.
+
+    This function analyzes a sample of the DataFrame to infer appropriate data types
+    and applies the conversions. It handles timestamps, numeric values, booleans,
+    and categorical fields.
+
+    Args:
+        df: A PySpark DataFrame to convert.
+        na_values: List of strings to be considered as NA/NaN. Defaults to
+            ['NA', 'na', 'NULL', 'null', ''].
+        category_threshold: Maximum ratio of unique values to total values
+            to consider a column categorical. Defaults to 0.3.
+        convert_threshold: Minimum ratio of successfully converted values required
+            to apply a type conversion. Defaults to 0.6.
+        sample_ratio: Fraction of data to sample for type inference. Defaults to 0.1.
+
+    Returns:
+        tuple: (The DataFrame with converted types, A dictionary mapping column names to
+                their inferred types as strings)
+
+    Note:
+        - 'category' in the schema dict is conceptual as PySpark doesn't have a true
+            category type like pandas
+        - The function uses sampling for efficiency with large datasets
+    """
+    n_rows = df.count()
+    if na_values is None:
+        na_values = ["NA", "na", "NULL", "null", ""]
+
+    # Normalize NA-like values
+    for colname, coltype in df.dtypes:
+        if coltype == "string":
+            df = df.withColumn(
+                colname,
+                F.when(F.trim(F.lower(F.col(colname))).isin([v.lower() for v in na_values]), None).otherwise(
+                    F.col(colname)
+                ),
+            )
+
+    schema = {}
+    for colname in df.columns:
+        # Sample once at an appropriate ratio
+        sample_ratio_to_use = min(1.0, sample_ratio if n_rows * sample_ratio > 100 else 100 / n_rows)
+        col_sample = df.select(colname).sample(withReplacement=False, fraction=sample_ratio_to_use).dropna()
+        sample_count = col_sample.count()
+
+        inferred_type = "string"  # Default
+
+        if col_sample.dtypes[0][1] != "string":
+            schema[colname] = col_sample.dtypes[0][1]
+            continue
+
+        if sample_count == 0:
+            schema[colname] = "string"
+            continue
+
+        # Check if timestamp
+        ts_col = col_sample.withColumn("parsed", F.to_timestamp(F.col(colname)))
+
+        # Check numeric
+        if (
+            col_sample.withColumn("n", F.col(colname).cast("double")).filter("n is not null").count()
+            >= sample_count * convert_threshold
+        ):
+            # All whole numbers?
+            all_whole = (
+                col_sample.withColumn("n", F.col(colname).cast("double"))
+                .filter("n is not null")
+                .withColumn("frac", F.abs(F.col("n") % 1))
+                .filter("frac > 0.000001")
+                .count()
+                == 0
+            )
+            inferred_type = "int" if all_whole else "double"
+
+        # Check low-cardinality (category-like)
+        elif (
+            sample_count > 0
+            and col_sample.select(F.countDistinct(F.col(colname))).collect()[0][0] / sample_count <= category_threshold
+        ):
+            inferred_type = "category"  # Will just be string, but marked as such
+
+        # Check if timestamp
+        elif ts_col.filter(F.col("parsed").isNotNull()).count() >= sample_count * convert_threshold:
+            inferred_type = "timestamp"
+
+        schema[colname] = inferred_type
+
+    # Apply inferred schema
+    for colname, inferred_type in schema.items():
+        if inferred_type == "int":
+            df = df.withColumn(colname, F.col(colname).cast(T.IntegerType()))
+        elif inferred_type == "double":
+            df = df.withColumn(colname, F.col(colname).cast(T.DoubleType()))
+        elif inferred_type == "boolean":
+            df = df.withColumn(
+                colname,
+                F.when(F.lower(F.col(colname)).isin("true", "yes", "1"), True)
+                .when(F.lower(F.col(colname)).isin("false", "no", "0"), False)
+                .otherwise(None),
+            )
+        elif inferred_type == "timestamp":
+            df = df.withColumn(colname, F.to_timestamp(F.col(colname)))
+        elif inferred_type == "category":
+            df = df.withColumn(colname, F.col(colname).cast(T.StringType()))  # Marked conceptually
+
+        # otherwise keep as string (or original type)
+
+    return df, schema
+
+
+def auto_convert_dtypes_pandas(
+    df: DataFrame,
+    na_values: list = None,
+    category_threshold: float = 0.3,
+    convert_threshold: float = 0.6,
+    sample_ratio: float = 1.0,
+) -> tuple[DataFrame, dict]:
+    """Automatically convert data types in a pandas DataFrame using heuristics.
+
+    This function analyzes the DataFrame to infer appropriate data types
+    and applies the conversions. It handles timestamps, timedeltas, numeric values,
+    and categorical fields.
+
+    Args:
+        df: A pandas DataFrame to convert.
+        na_values: List of strings to be considered as NA/NaN. Defaults to
+            ['NA', 'na', 'NULL', 'null', ''].
+        category_threshold: Maximum ratio of unique values to total values
+            to consider a column categorical. Defaults to 0.3.
+        convert_threshold: Minimum ratio of successfully converted values required
+            to apply a type conversion. Defaults to 0.6.
+        sample_ratio: Fraction of data to sample for type inference. Not used in pandas version
+            but included for API compatibility. Defaults to 1.0.
+
+    Returns:
+        tuple: (The DataFrame with converted types, A dictionary mapping column names to
+                their inferred types as strings)
+    """
+    if na_values is None:
+        na_values = {"NA", "na", "NULL", "null", ""}
+
+    df_converted = df.convert_dtypes()
+    schema = {}
+
+    # Sample if needed (for API compatibility)
+    if sample_ratio < 1.0:
+        df = df.sample(frac=sample_ratio)
+
+    n_rows = len(df)
+
+    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)
+
+        # Skip conversion if already non-object data type, except bool which can potentially be categorical
+        if (
+            not isinstance(series_cleaned.dtype, pd.BooleanDtype)
+            and not isinstance(series_cleaned.dtype, pd.StringDtype)
+            and series_cleaned.dtype != "object"
+        ):
+            # Keep the original data type for non-object dtypes
+            df_converted[col] = series
+            schema[col] = str(series_cleaned.dtype)
+            continue
+
+        # print(f"type: {series_cleaned.dtype}, column: {series_cleaned.name}")
+
+        if not isinstance(series_cleaned.dtype, pd.BooleanDtype):
+            # Try numeric (int or float)
+            numeric = pd.to_numeric(series_cleaned, errors="coerce")
+            if numeric.notna().sum() >= n_rows * convert_threshold:
+                if (numeric.dropna() % 1 == 0).all():
+                    try:
+                        df_converted[col] = numeric.astype("int")  # Nullable integer
+                        schema[col] = "int"
+                        continue
+                    except Exception:
+                        pass
+                df_converted[col] = numeric.astype("double")
+                schema[col] = "double"
+                continue
+
+            # Try datetime
+            datetime_converted = pd.to_datetime(series_cleaned, errors="coerce")
+            if datetime_converted.notna().sum() >= n_rows * convert_threshold:
+                df_converted[col] = datetime_converted
+                schema[col] = "timestamp"
+                continue
+
+            # Try timedelta
+            try:
+                timedelta_converted = pd.to_timedelta(series_cleaned, errors="coerce")
+                if timedelta_converted.notna().sum() >= n_rows * convert_threshold:
+                    df_converted[col] = timedelta_converted
+                    schema[col] = "timedelta"
+                    continue
+            except TypeError:
+                pass
+
+        # Try category
+        try:
+            unique_ratio = series_cleaned.nunique(dropna=True) / n_rows if n_rows > 0 else 1.0
+            if unique_ratio <= category_threshold:
+                df_converted[col] = series_cleaned.astype("category")
+                schema[col] = "category"
+                continue
+        except Exception:
+            pass
+        df_converted[col] = series_cleaned.astype("string")
+        schema[col] = "string"
+
+    return df_converted, schema

flaml/automl/logger.py

@@ -1,7 +1,37 @@
 import logging
+import os
+
+
+class ColoredFormatter(logging.Formatter):
+    # ANSI escape codes for colors
+    COLORS = {
+        # logging.DEBUG: "\033[36m",  # Cyan
+        # logging.INFO: "\033[32m",   # Green
+        logging.WARNING: "\033[33m",  # Yellow
+        logging.ERROR: "\033[31m",  # Red
+        logging.CRITICAL: "\033[1;31m",  # Bright Red
+    }
+    RESET = "\033[0m"  # Reset to default
+
+    def __init__(self, fmt, datefmt, use_color=True):
+        super().__init__(fmt, datefmt)
+        self.use_color = use_color
+
+    def format(self, record):
+        formatted = super().format(record)
+        if self.use_color:
+            color = self.COLORS.get(record.levelno, "")
+            if color:
+                return f"{color}{formatted}{self.RESET}"
+        return formatted
+
 
 logger = logging.getLogger(__name__)
-logger_formatter = logging.Formatter(
-    "[%(name)s: %(asctime)s] {%(lineno)d} %(levelname)s - %(message)s", "%m-%d %H:%M:%S"
+use_color = True
+if os.getenv("FLAML_LOG_NO_COLOR"):
+    use_color = False
+
+logger_formatter = ColoredFormatter(
+    "[%(name)s: %(asctime)s] {%(lineno)d} %(levelname)s - %(message)s", "%m-%d %H:%M:%S", use_color
 )
 logger.propagate = False

flaml/automl/ml.py

@@ -2,30 +2,31 @@
 #  * Copyright (c) FLAML authors. All rights reserved.
 #  * Licensed under the MIT License. See LICENSE file in the
 #  * project root for license information.
-import time
-from typing import Union, Callable, TypeVar, Optional, Tuple
 import logging
+import time
+from typing import Callable, Optional, Tuple, TypeVar, Union
 
 import numpy as np
 
-
 from flaml.automl.data import group_counts
-from flaml.automl.task.task import Task
 from flaml.automl.model import BaseEstimator, TransformersEstimator
-from flaml.automl.spark import psDataFrame, psSeries, ERROR as SPARK_ERROR, Series, DataFrame
+from flaml.automl.spark import ERROR as SPARK_ERROR
+from flaml.automl.spark import DataFrame, Series, psDataFrame, psSeries
+from flaml.automl.task.task import Task
+from flaml.automl.time_series import TimeSeriesDataset
 
 try:
     from sklearn.metrics import (
-        mean_squared_error,
-        r2_score,
-        roc_auc_score,
         accuracy_score,
-        mean_absolute_error,
-        log_loss,
         average_precision_score,
         f1_score,
+        log_loss,
+        mean_absolute_error,
         mean_absolute_percentage_error,
+        mean_squared_error,
         ndcg_score,
+        r2_score,
+        roc_auc_score,
     )
 except ImportError:
     pass
@@ -33,7 +34,6 @@ except ImportError:
 if SPARK_ERROR is None:
     from flaml.automl.spark.metrics import spark_metric_loss_score
 
-from flaml.automl.time_series import TimeSeriesDataset
 
 logger = logging.getLogger(__name__)
 
@@ -89,6 +89,11 @@ huggingface_metric_to_mode = {
     "wer": "min",
 }
 huggingface_submetric_to_metric = {"rouge1": "rouge", "rouge2": "rouge"}
+spark_metric_name_dict = {
+    "Regression": ["r2", "rmse", "mse", "mae", "var"],
+    "Binary Classification": ["pr_auc", "roc_auc"],
+    "Multi-class Classification": ["accuracy", "log_loss", "f1", "micro_f1", "macro_f1"],
+}
 
 
 def metric_loss_score(
@@ -122,7 +127,7 @@ 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)
+            metric = datasets.load_metric(datasets_metric_name, trust_remote_code=True)
             metric_mode = huggingface_metric_to_mode[datasets_metric_name]
 
             if metric_name.startswith("seqeval"):
@@ -323,7 +328,7 @@ def compute_estimator(
     estimator_name: str,
     eval_method: str,
     eval_metric: Union[str, Callable],
-    best_val_loss=np.Inf,
+    best_val_loss=np.inf,
     n_jobs: Optional[int] = 1,  # some estimators of EstimatorSubclass don't accept n_jobs. Should be None in that case.
     estimator_class: Optional[EstimatorSubclass] = None,
     cv_score_agg_func: Optional[callable] = None,
@@ -334,6 +339,14 @@ def compute_estimator(
     if fit_kwargs is None:
         fit_kwargs = {}
 
+    fe_params = {}
+    for param, value in config_dic.items():
+        if param.startswith("fe."):
+            fe_params[param] = value
+
+    for param, value in fe_params.items():
+        config_dic.pop(param)
+
     estimator_class = estimator_class or task.estimator_class_from_str(estimator_name)
     estimator = estimator_class(
         **config_dic,
@@ -401,12 +414,21 @@ def train_estimator(
     free_mem_ratio=0,
 ) -> Tuple[EstimatorSubclass, float]:
     start_time = time.time()
+    fe_params = {}
+    for param, value in config_dic.items():
+        if param.startswith("fe."):
+            fe_params[param] = value
+
+    for param, value in fe_params.items():
+        config_dic.pop(param)
+
     estimator_class = estimator_class or task.estimator_class_from_str(estimator_name)
     estimator = estimator_class(
         **config_dic,
         task=task,
         n_jobs=n_jobs,
     )
+
     if fit_kwargs is None:
         fit_kwargs = {}
 
@@ -567,14 +589,19 @@ def _eval_estimator(
 
         pred_time = (time.time() - pred_start) / num_val_rows
 
-        val_loss = metric_loss_score(
-            eval_metric,
-            y_processed_predict=val_pred_y,
-            y_processed_true=y_val,
-            labels=labels,
-            sample_weight=weight_val,
-            groups=groups_val,
-        )
+        try:
+            val_loss = metric_loss_score(
+                eval_metric,
+                y_processed_predict=val_pred_y,
+                y_processed_true=y_val,
+                labels=labels,
+                sample_weight=weight_val,
+                groups=groups_val,
+            )
+        except ValueError as e:
+            # `r2_score` and other metrics may raise a `ValueError` when a model returns `inf` or `nan` values. In this case, we set the val_loss to infinity.
+            val_loss = np.inf
+            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)

flaml/automl/nlp/README.md

@@ -4,16 +4,15 @@ This directory contains utility functions used by AutoNLP. Currently we support
 
 Please refer to this [link](https://microsoft.github.io/FLAML/docs/Examples/AutoML-NLP) for examples.
 
-
 # Troubleshooting fine-tuning HPO for pre-trained language models
 
 The frequent updates of transformers may lead to fluctuations in the results of tuning. To help users quickly troubleshoot the result of AutoNLP when a tuning failure occurs (e.g., failing to reproduce previous results), we have provided the following jupyter notebook:
 
-* [Troubleshooting HPO for fine-tuning pre-trained language models](https://github.com/microsoft/FLAML/blob/main/notebook/research/acl2021.ipynb)
+- [Troubleshooting HPO for fine-tuning pre-trained language models](https://github.com/microsoft/FLAML/blob/main/notebook/research/acl2021.ipynb)
 
 Our findings on troubleshooting fine-tuning the Electra and RoBERTa model for the GLUE dataset can be seen in the following paper published in ACL 2021:
 
-* [An Empirical Study on Hyperparameter Optimization for Fine-Tuning Pre-trained Language Models](https://arxiv.org/abs/2106.09204). Xueqing Liu, Chi Wang. ACL-IJCNLP 2021.
+- [An Empirical Study on Hyperparameter Optimization for Fine-Tuning Pre-trained Language Models](https://arxiv.org/abs/2106.09204). Xueqing Liu, Chi Wang. ACL-IJCNLP 2021.
 
 ```bibtex
 @inproceedings{liu2021hpo,

flaml/automl/nlp/huggingface/data_collator.py

@@ -1,17 +1,18 @@
-from dataclasses import dataclass
-from transformers.data.data_collator import (
-    DataCollatorWithPadding,
-    DataCollatorForTokenClassification,
-    DataCollatorForSeq2Seq,
-)
 from collections import OrderedDict
+from dataclasses import dataclass
+
+from transformers.data.data_collator import (
+    DataCollatorForSeq2Seq,
+    DataCollatorForTokenClassification,
+    DataCollatorWithPadding,
+)
 
 from flaml.automl.task.task import (
-    TOKENCLASSIFICATION,
     MULTICHOICECLASSIFICATION,
-    SUMMARIZATION,
     SEQCLASSIFICATION,
     SEQREGRESSION,
+    SUMMARIZATION,
+    TOKENCLASSIFICATION,
 )
 
 
@@ -19,6 +20,7 @@ from flaml.automl.task.task import (
 class DataCollatorForMultipleChoiceClassification(DataCollatorWithPadding):
     def __call__(self, features):
         from itertools import chain
+
         import torch
 
         label_name = "label" if "label" in features[0].keys() else "labels"
@@ -30,7 +32,7 @@ class DataCollatorForMultipleChoiceClassification(DataCollatorWithPadding):
             [{k: v[i] for k, v in feature.items()} for i in range(num_choices)] for feature in features
         ]
         flattened_features = list(chain(*flattened_features))
-        batch = super(DataCollatorForMultipleChoiceClassification, self).__call__(flattened_features)
+        batch = super().__call__(flattened_features)
         # Un-flatten
         batch = {k: v.view(batch_size, num_choices, -1) for k, v in batch.items()}
         # Add back labels

flaml/automl/nlp/huggingface/training_args.py

@@ -1,6 +1,7 @@
 import argparse
 from dataclasses import dataclass, field
-from typing import Optional, List
+from typing import List, Optional
+
 from flaml.automl.task.task import NLG_TASKS
 
 try:

flaml/automl/nlp/huggingface/utils.py

@@ -1,14 +1,16 @@
 from itertools import chain
+
 import numpy as np
-from flaml.automl.task.task import (
-    SUMMARIZATION,
-    SEQREGRESSION,
-    SEQCLASSIFICATION,
-    MULTICHOICECLASSIFICATION,
-    TOKENCLASSIFICATION,
-    NLG_TASKS,
-)
+
 from flaml.automl.data import pd
+from flaml.automl.task.task import (
+    MULTICHOICECLASSIFICATION,
+    NLG_TASKS,
+    SEQCLASSIFICATION,
+    SEQREGRESSION,
+    SUMMARIZATION,
+    TOKENCLASSIFICATION,
+)
 
 
 def todf(X, Y, column_name):
@@ -243,7 +245,7 @@ def tokenize_row(
     return_column_name=False,
 ):
     if prefix:
-        this_row = tuple(["".join(x) for x in zip(prefix, this_row)])
+        this_row = tuple("".join(x) for x in zip(prefix, this_row))
 
     # tokenizer.pad_token = tokenizer.eos_token
     tokenized_example = tokenizer(
@@ -377,6 +379,7 @@ def load_model(checkpoint_path, task, num_labels=None):
     transformers.logging.set_verbosity_error()
 
     from transformers import AutoConfig
+
     from flaml.automl.task.task import (
         SEQCLASSIFICATION,
         SEQREGRESSION,
@@ -384,10 +387,12 @@ def load_model(checkpoint_path, task, num_labels=None):
     )
 
     def get_this_model(checkpoint_path, task, model_config):
-        from transformers import AutoModelForSequenceClassification
-        from transformers import AutoModelForSeq2SeqLM
-        from transformers import AutoModelForMultipleChoice
-        from transformers import AutoModelForTokenClassification
+        from transformers import (
+            AutoModelForMultipleChoice,
+            AutoModelForSeq2SeqLM,
+            AutoModelForSequenceClassification,
+            AutoModelForTokenClassification,
+        )
 
         if task in (SEQCLASSIFICATION, SEQREGRESSION):
             return AutoModelForSequenceClassification.from_pretrained(

flaml/automl/nlp/utils.py

@@ -1,11 +1,12 @@
-from typing import Dict, Any
+from typing import Any, Dict
+
 import numpy as np
 
 from flaml.automl.task.task import (
-    SUMMARIZATION,
-    SEQREGRESSION,
-    SEQCLASSIFICATION,
     MULTICHOICECLASSIFICATION,
+    SEQCLASSIFICATION,
+    SEQREGRESSION,
+    SUMMARIZATION,
     TOKENCLASSIFICATION,
 )
 
@@ -31,7 +32,7 @@ def is_a_list_of_str(this_obj):
 
 def _clean_value(value: Any) -> str:
     if isinstance(value, float):
-        return "{:.5}".format(value)
+        return f"{value:.5}"
     else:
         return str(value).replace("/", "_")
 
@@ -85,7 +86,7 @@ class Counter:
     @staticmethod
     def get_trial_fold_name(local_dir, trial_config, trial_id):
         Counter.counter += 1
-        experiment_tag = "{0}_{1}".format(str(Counter.counter), format_vars(trial_config))
+        experiment_tag = f"{str(Counter.counter)}_{format_vars(trial_config)}"
         logdir = get_logdir_name(_generate_dirname(experiment_tag, trial_id=trial_id), local_dir)
         return logdir
 

flaml/automl/spark/__init__.py

@@ -6,8 +6,10 @@ try:
     import pyspark.pandas as ps
     import pyspark.sql.functions as F
     import pyspark.sql.types as T
+    from pyspark.pandas import DataFrame as psDataFrame
+    from pyspark.pandas import Series as psSeries
+    from pyspark.pandas import set_option
     from pyspark.sql import DataFrame as sparkDataFrame
-    from pyspark.pandas import DataFrame as psDataFrame, Series as psSeries, set_option
     from pyspark.util import VersionUtils
 except ImportError:
 

flaml/automl/spark/configs.py

@@ -1,97 +0,0 @@
-ParamList_LightGBM_Base = [
-    "baggingFraction",
-    "baggingFreq",
-    "baggingSeed",
-    "binSampleCount",
-    "boostFromAverage",
-    "boostingType",
-    "catSmooth",
-    "categoricalSlotIndexes",
-    "categoricalSlotNames",
-    "catl2",
-    "chunkSize",
-    "dataRandomSeed",
-    "defaultListenPort",
-    "deterministic",
-    "driverListenPort",
-    "dropRate",
-    "dropSeed",
-    "earlyStoppingRound",
-    "executionMode",
-    "extraSeed" "featureFraction",
-    "featureFractionByNode",
-    "featureFractionSeed",
-    "featuresCol",
-    "featuresShapCol",
-    "fobj" "improvementTolerance",
-    "initScoreCol",
-    "isEnableSparse",
-    "isProvideTrainingMetric",
-    "labelCol",
-    "lambdaL1",
-    "lambdaL2",
-    "leafPredictionCol",
-    "learningRate",
-    "matrixType",
-    "maxBin",
-    "maxBinByFeature",
-    "maxCatThreshold",
-    "maxCatToOnehot",
-    "maxDeltaStep",
-    "maxDepth",
-    "maxDrop",
-    "metric",
-    "microBatchSize",
-    "minDataInLeaf",
-    "minDataPerBin",
-    "minDataPerGroup",
-    "minGainToSplit",
-    "minSumHessianInLeaf",
-    "modelString",
-    "monotoneConstraints",
-    "monotoneConstraintsMethod",
-    "monotonePenalty",
-    "negBaggingFraction",
-    "numBatches",
-    "numIterations",
-    "numLeaves",
-    "numTasks",
-    "numThreads",
-    "objectiveSeed",
-    "otherRate",
-    "parallelism",
-    "passThroughArgs",
-    "posBaggingFraction",
-    "predictDisableShapeCheck",
-    "predictionCol",
-    "repartitionByGroupingColumn",
-    "seed",
-    "skipDrop",
-    "slotNames",
-    "timeout",
-    "topK",
-    "topRate",
-    "uniformDrop",
-    "useBarrierExecutionMode",
-    "useMissing",
-    "useSingleDatasetMode",
-    "validationIndicatorCol",
-    "verbosity",
-    "weightCol",
-    "xGBoostDartMode",
-    "zeroAsMissing",
-    "objective",
-]
-ParamList_LightGBM_Classifier = ParamList_LightGBM_Base + [
-    "isUnbalance",
-    "probabilityCol",
-    "rawPredictionCol",
-    "thresholds",
-]
-ParamList_LightGBM_Regressor = ParamList_LightGBM_Base + ["tweedieVariancePower"]
-ParamList_LightGBM_Ranker = ParamList_LightGBM_Base + [
-    "groupCol",
-    "evalAt",
-    "labelGain",
-    "maxPosition",
-]

flaml/automl/spark/metrics.py

@@ -1,14 +1,17 @@
-import numpy as np
+import json
 from typing import Union
-from flaml.automl.spark import psSeries, F
+
+import numpy as np
 from pyspark.ml.evaluation import (
     BinaryClassificationEvaluator,
-    RegressionEvaluator,
     MulticlassClassificationEvaluator,
     MultilabelClassificationEvaluator,
     RankingEvaluator,
+    RegressionEvaluator,
 )
 
+from flaml.automl.spark import F, T, psDataFrame, psSeries, sparkDataFrame
+
 
 def ps_group_counts(groups: Union[psSeries, np.ndarray]) -> np.ndarray:
     if isinstance(groups, np.ndarray):
@@ -34,6 +37,16 @@ def _compute_label_from_probability(df, probability_col, prediction_col):
     return df
 
 
+def string_to_array(s):
+    try:
+        return json.loads(s)
+    except json.JSONDecodeError:
+        return []
+
+
+string_to_array_udf = F.udf(string_to_array, T.ArrayType(T.DoubleType()))
+
+
 def spark_metric_loss_score(
     metric_name: str,
     y_predict: psSeries,
@@ -133,6 +146,11 @@ def spark_metric_loss_score(
         )
     elif metric_name == "log_loss":
         # For log_loss, prediction_col should be probability, and we need to convert it to label
+        # handle data like "{'type': '1', 'values': '[1, 2, 3]'}"
+        # Fix cannot resolve "array_max(prediction)" due to data type mismatch: Parameter 1 requires the "ARRAY" type,
+        # however "prediction" has the type "STRUCT<type: TINYINT, size: INT, indices: ARRAY<INT>, values: ARRAY<DOUBLE>>"
+        df = df.withColumn(prediction_col, df[prediction_col].cast(T.StringType()))
+        df = df.withColumn(prediction_col, string_to_array_udf(df[prediction_col]))
         df = _compute_label_from_probability(df, prediction_col, prediction_col + "_label")
         evaluator = MulticlassClassificationEvaluator(
             metricName="logLoss",

flaml/automl/spark/utils.py

@@ -1,17 +1,19 @@
 import logging
-from typing import Union, List, Optional, Tuple
+from typing import List, Optional, Tuple, Union
+
 import numpy as np
+
 from flaml.automl.spark import (
-    sparkDataFrame,
-    ps,
+    DataFrame,
     F,
+    Series,
     T,
+    _spark_major_minor_version,
+    ps,
     psDataFrame,
     psSeries,
-    _spark_major_minor_version,
-    DataFrame,
-    Series,
     set_option,
+    sparkDataFrame,
 )
 
 logger = logging.getLogger(__name__)

flaml/automl/state.py

@@ -1,13 +1,15 @@
-import inspect
 import copy
+import inspect
 import time
 from typing import Any, Optional
+
 import numpy as np
+
 from flaml import tune
 from flaml.automl.logger import logger
 from flaml.automl.ml import compute_estimator, train_estimator
+from flaml.automl.spark import DataFrame, Series, psDataFrame, psSeries
 from flaml.automl.time_series.ts_data import TimeSeriesDataset
-from flaml.automl.spark import psDataFrame, psSeries, DataFrame, Series
 
 
 class SearchState:
@@ -63,6 +65,7 @@ class SearchState:
         custom_hp=None,
         max_iter=None,
         budget=None,
+        featurization="auto",
     ):
         self.init_eci = learner_class.cost_relative2lgbm() if budget >= 0 else 1
         self._search_space_domain = {}
@@ -80,6 +83,7 @@ class SearchState:
         else:
             data_size = data.shape
             search_space = learner_class.search_space(data_size=data_size, task=task)
+
         self.data_size = data_size
 
         if custom_hp is not None:
@@ -89,9 +93,7 @@ class SearchState:
             starting_point = AutoMLState.sanitize(starting_point)
             if max_iter > 1 and not self.valid_starting_point(starting_point, search_space):
                 # If the number of iterations is larger than 1, remove invalid point
-                logger.warning(
-                    "Starting point {} removed because it is outside of the search space".format(starting_point)
-                )
+                logger.warning(f"Starting point {starting_point} removed because it is outside of the search space")
                 starting_point = None
         elif isinstance(starting_point, list):
             starting_point = [AutoMLState.sanitize(x) for x in starting_point]
@@ -206,7 +208,7 @@ class SearchState:
         self.val_loss, self.config = obj, config
 
     def get_hist_config_sig(self, sample_size, config):
-        config_values = tuple([config[k] for k in self._hp_names if k in config])
+        config_values = tuple(config[k] for k in self._hp_names if k in config)
         config_sig = str(sample_size) + "_" + str(config_values)
         return config_sig
 
@@ -288,9 +290,11 @@ class AutoMLState:
         budget = (
             None
             if state.time_budget < 0
-            else state.time_budget - state.time_from_start
-            if sample_size == state.data_size[0]
-            else (state.time_budget - state.time_from_start) / 2 * sample_size / state.data_size[0]
+            else (
+                state.time_budget - state.time_from_start
+                if sample_size == state.data_size[0]
+                else (state.time_budget - state.time_from_start) / 2 * sample_size / state.data_size[0]
+            )
         )
 
         (
@@ -351,6 +355,7 @@ class AutoMLState:
         estimator: str,
         config_w_resource: dict,
         sample_size: Optional[int] = None,
+        is_retrain: bool = False,
     ):
         if not sample_size:
             sample_size = config_w_resource.get("FLAML_sample_size", len(self.y_train_all))
@@ -376,9 +381,8 @@ class AutoMLState:
             this_estimator_kwargs[
                 "groups"
             ] = groups  # NOTE: _train_with_config is after kwargs is updated to fit_kwargs_by_estimator
-
+        this_estimator_kwargs.update({"is_retrain": is_retrain})
         budget = None if self.time_budget < 0 else self.time_budget - self.time_from_start
-
         estimator, train_time = train_estimator(
             X_train=sampled_X_train,
             y_train=sampled_y_train,

flaml/automl/task/factory.py

@@ -1,8 +1,9 @@
 from typing import Optional, Union
+
 import numpy as np
 
 from flaml.automl.data import DataFrame, Series
-from flaml.automl.task.task import Task, TS_FORECAST
+from flaml.automl.task.task import TS_FORECAST, Task
 
 
 def task_factory(

flaml/automl/task/generic_task.py

@@ -1,43 +1,39 @@
 import logging
 import time
 from typing import List, Optional
-import numpy as np
-from flaml.automl.data import TS_TIMESTAMP_COL, concat
-from flaml.automl.ml import EstimatorSubclass, get_val_loss, default_cv_score_agg_func
 
-from flaml.automl.task.task import (
-    Task,
-    get_classification_objective,
-    TS_FORECAST,
-    TS_FORECASTPANEL,
-)
-from flaml.config import RANDOM_SEED
-from flaml.automl.spark import ps, psDataFrame, psSeries, pd
+import numpy as np
+
+from flaml.automl.data import TS_TIMESTAMP_COL, concat
+from flaml.automl.ml import EstimatorSubclass, default_cv_score_agg_func, get_val_loss
+from flaml.automl.spark import pd, ps, psDataFrame, psSeries
 from flaml.automl.spark.utils import (
     iloc_pandas_on_spark,
+    len_labels,
+    set_option,
     spark_kFold,
     train_test_split_pyspark,
     unique_pandas_on_spark,
     unique_value_first_index,
-    len_labels,
-    set_option,
 )
+from flaml.automl.task.task import TS_FORECAST, TS_FORECASTPANEL, Task, get_classification_objective
+from flaml.config import RANDOM_SEED
 
 try:
     from scipy.sparse import issparse
 except ImportError:
     pass
 try:
-    from sklearn.utils import shuffle
     from sklearn.model_selection import (
-        train_test_split,
-        RepeatedStratifiedKFold,
-        RepeatedKFold,
         GroupKFold,
-        TimeSeriesSplit,
         GroupShuffleSplit,
+        RepeatedKFold,
+        RepeatedStratifiedKFold,
         StratifiedGroupKFold,
+        TimeSeriesSplit,
+        train_test_split,
     )
+    from sklearn.utils import shuffle
 except ImportError:
     pass
 
@@ -49,19 +45,31 @@ class GenericTask(Task):
     def estimators(self):
         if self._estimators is None:
             # put this into a function to avoid circular dependency
+            from flaml.automl.contrib.histgb import HistGradientBoostingEstimator
             from flaml.automl.model import (
-                XGBoostSklearnEstimator,
-                XGBoostLimitDepthEstimator,
-                RandomForestEstimator,
+                CatBoostEstimator,
+                ElasticNetEstimator,
+                ExtraTreesEstimator,
+                KNeighborsEstimator,
+                LassoLarsEstimator,
                 LGBMEstimator,
                 LRL1Classifier,
                 LRL2Classifier,
-                CatBoostEstimator,
-                ExtraTreesEstimator,
-                KNeighborsEstimator,
+                RandomForestEstimator,
+                SGDEstimator,
+                SparkAFTSurvivalRegressionEstimator,
+                SparkGBTEstimator,
+                SparkGLREstimator,
+                SparkLGBMEstimator,
+                SparkLinearRegressionEstimator,
+                SparkLinearSVCEstimator,
+                SparkNaiveBayesEstimator,
+                SparkRandomForestEstimator,
+                SVCEstimator,
                 TransformersEstimator,
                 TransformersEstimatorModelSelection,
-                SparkLGBMEstimator,
+                XGBoostLimitDepthEstimator,
+                XGBoostSklearnEstimator,
             )
 
             self._estimators = {
@@ -70,6 +78,7 @@ class GenericTask(Task):
                 "rf": RandomForestEstimator,
                 "lgbm": LGBMEstimator,
                 "lgbm_spark": SparkLGBMEstimator,
+                "rf_spark": SparkRandomForestEstimator,
                 "lrl1": LRL1Classifier,
                 "lrl2": LRL2Classifier,
                 "catboost": CatBoostEstimator,
@@ -77,6 +86,17 @@ class GenericTask(Task):
                 "kneighbor": KNeighborsEstimator,
                 "transformer": TransformersEstimator,
                 "transformer_ms": TransformersEstimatorModelSelection,
+                "histgb": HistGradientBoostingEstimator,
+                "svc": SVCEstimator,
+                "sgd": SGDEstimator,
+                "nb_spark": SparkNaiveBayesEstimator,
+                "enet": ElasticNetEstimator,
+                "lassolars": LassoLarsEstimator,
+                "glr_spark": SparkGLREstimator,
+                "lr_spark": SparkLinearRegressionEstimator,
+                "svc_spark": SparkLinearSVCEstimator,
+                "gbt_spark": SparkGBTEstimator,
+                "aft_spark": SparkAFTSurvivalRegressionEstimator,
             }
         return self._estimators
 
@@ -268,8 +288,8 @@ class GenericTask(Task):
             seed=RANDOM_SEED,
         )
         columns_to_drop = [c for c in df_all_train.columns if c in [stratify_column, "sample_weight"]]
-        X_train = df_all_train.drop(columns_to_drop)
-        X_val = df_all_val.drop(columns_to_drop)
+        X_train = df_all_train.drop(columns=columns_to_drop)
+        X_val = df_all_val.drop(columns=columns_to_drop)
         y_train = df_all_train[stratify_column]
         y_val = df_all_val[stratify_column]
 
@@ -422,8 +442,8 @@ class GenericTask(Task):
                 X_train_all, y_train_all = shuffle(X_train_all, y_train_all, random_state=RANDOM_SEED)
             if data_is_df:
                 X_train_all.reset_index(drop=True, inplace=True)
-            if isinstance(y_train_all, pd.Series):
-                y_train_all.reset_index(drop=True, inplace=True)
+        if isinstance(y_train_all, pd.Series):
+            y_train_all.reset_index(drop=True, inplace=True)
 
         X_train, y_train = X_train_all, y_train_all
         state.groups_all = state.groups
@@ -494,14 +514,37 @@ class GenericTask(Task):
                     last = first[i] + 1
                 rest.extend(range(last, len(y_train_all)))
                 X_first = X_train_all.iloc[first] if data_is_df else X_train_all[first]
-                X_rest = X_train_all.iloc[rest] if data_is_df else X_train_all[rest]
-                y_rest = (
-                    y_train_all[rest]
-                    if isinstance(y_train_all, np.ndarray)
-                    else iloc_pandas_on_spark(y_train_all, rest)
-                    if is_spark_dataframe
-                    else y_train_all.iloc[rest]
-                )
+                if len(first) < len(y_train_all) / 2:
+                    # Get X_rest and y_rest with drop, sparse matrix can't apply np.delete
+                    X_rest = (
+                        np.delete(X_train_all, first, axis=0)
+                        if isinstance(X_train_all, np.ndarray)
+                        else X_train_all.drop(first.tolist())
+                        if data_is_df
+                        else X_train_all[rest]
+                    )
+                    y_rest = (
+                        np.delete(y_train_all, first, axis=0)
+                        if isinstance(y_train_all, np.ndarray)
+                        else y_train_all.drop(first.tolist())
+                        if data_is_df
+                        else y_train_all[rest]
+                    )
+                else:
+                    X_rest = (
+                        iloc_pandas_on_spark(X_train_all, rest)
+                        if is_spark_dataframe
+                        else X_train_all.iloc[rest]
+                        if data_is_df
+                        else X_train_all[rest]
+                    )
+                    y_rest = (
+                        iloc_pandas_on_spark(y_train_all, rest)
+                        if is_spark_dataframe
+                        else y_train_all.iloc[rest]
+                        if data_is_df
+                        else y_train_all[rest]
+                    )
                 stratify = y_rest if split_type == "stratified" else None
                 X_train, X_val, y_train, y_val = self._train_test_split(
                     state, X_rest, y_rest, first, rest, split_ratio, stratify
@@ -510,6 +553,12 @@ class GenericTask(Task):
                 y_train = concat(label_set, y_train) if data_is_df else np.concatenate([label_set, y_train])
                 X_val = concat(X_first, X_val)
                 y_val = concat(label_set, y_val) if data_is_df else np.concatenate([label_set, y_val])
+
+                if isinstance(y_train, (psDataFrame, pd.DataFrame)) and y_train.shape[1] == 1:
+                    y_train = y_train[y_train.columns[0]]
+                    y_val = y_val[y_val.columns[0]]
+                    y_train.name = y_val.name = y_rest.name
+
             elif self.is_regression():
                 X_train, X_val, y_train, y_val = self._train_test_split(
                     state, X_train_all, y_train_all, split_ratio=split_ratio
@@ -656,7 +705,6 @@ class GenericTask(Task):
             fit_kwargs = {}
         if cv_score_agg_func is None:
             cv_score_agg_func = default_cv_score_agg_func
-        start_time = time.time()
         val_loss_folds = []
         log_metric_folds = []
         metric = None
@@ -721,10 +769,10 @@ class GenericTask(Task):
             if not is_spark_dataframe:
                 y_train, y_val = y_train_split[train_index], y_train_split[val_index]
                 if weight is not None:
-                    fit_kwargs["sample_weight"], weight_val = (
-                        weight[train_index],
-                        weight[val_index],
+                    fit_kwargs["sample_weight"] = (
+                        weight[train_index] if isinstance(weight, np.ndarray) else weight.iloc[train_index]
                     )
+                    weight_val = weight[val_index] if isinstance(weight, np.ndarray) else weight.iloc[val_index]
                 if groups is not None:
                     fit_kwargs["groups"] = (
                         groups[train_index] if isinstance(groups, np.ndarray) else groups.iloc[train_index]
@@ -763,8 +811,6 @@ class GenericTask(Task):
             if is_spark_dataframe:
                 X_train.spark.unpersist()  # uncache data to free memory
                 X_val.spark.unpersist()  # uncache data to free memory
-            if budget and time.time() - start_time >= budget:
-                break
         val_loss, metric = cv_score_agg_func(val_loss_folds, log_metric_folds)
         n = total_fold_num
         pred_time /= n
@@ -807,27 +853,23 @@ class GenericTask(Task):
         elif self.is_ts_forecastpanel():
             estimator_list = ["tft"]
         else:
+            estimator_list = [
+                "lgbm",
+                "rf",
+                "xgboost",
+                "extra_tree",
+                "xgb_limitdepth",
+                "lgbm_spark",
+                "rf_spark",
+                "sgd",
+            ]
             try:
                 import catboost
 
-                estimator_list = [
-                    "lgbm",
-                    "rf",
-                    "catboost",
-                    "xgboost",
-                    "extra_tree",
-                    "xgb_limitdepth",
-                    "lgbm_spark",
-                ]
+                estimator_list += ["catboost"]
             except ImportError:
-                estimator_list = [
-                    "lgbm",
-                    "rf",
-                    "xgboost",
-                    "extra_tree",
-                    "xgb_limitdepth",
-                    "lgbm_spark",
-                ]
+                pass
+
             # if self.is_ts_forecast():
             #     # catboost is removed because it has a `name` parameter, making it incompatible with hcrystalball
             #     if "catboost" in estimator_list:
@@ -859,9 +901,7 @@ class GenericTask(Task):
             return metric
 
         if self.is_nlp():
-            from flaml.automl.nlp.utils import (
-                load_default_huggingface_metric_for_task,
-            )
+            from flaml.automl.nlp.utils import load_default_huggingface_metric_for_task
 
             return load_default_huggingface_metric_for_task(self.name)
         elif self.is_binary():

flaml/automl/task/task.py

@@ -1,6 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, List, Optional, Tuple, Union
+
 import numpy as np
+
 from flaml.automl.data import DataFrame, Series, psDataFrame, psSeries
 
 if TYPE_CHECKING:
@@ -190,7 +192,7 @@ class Task(ABC):
                 * Valid str options depend on different tasks.
                 For classification tasks, valid choices are
                     ["auto", 'stratified', 'uniform', 'time', 'group']. "auto" -> stratified.
-                For regression tasks, valid choices are ["auto", 'uniform', 'time'].
+                For regression tasks, valid choices are ["auto", 'uniform', 'time', 'group'].
                     "auto" -> uniform.
                 For time series forecast tasks, must be "auto" or 'time'.
                 For ranking task, must be "auto" or 'group'.

flaml/automl/task/time_series_task.py

@@ -2,26 +2,25 @@ import logging
 import time
 from typing import List
 
-import pandas as pd
 import numpy as np
+import pandas as pd
 from scipy.sparse import issparse
 from sklearn.model_selection import (
     GroupKFold,
     TimeSeriesSplit,
 )
 
-from flaml.automl.ml import get_val_loss, default_cv_score_agg_func
-from flaml.automl.time_series.ts_data import (
-    TimeSeriesDataset,
-    DataTransformerTS,
-    normalize_ts_data,
-)
-
+from flaml.automl.ml import default_cv_score_agg_func, get_val_loss
 from flaml.automl.task.task import (
-    Task,
-    get_classification_objective,
     TS_FORECAST,
     TS_FORECASTPANEL,
+    Task,
+    get_classification_objective,
+)
+from flaml.automl.time_series.ts_data import (
+    DataTransformerTS,
+    TimeSeriesDataset,
+    normalize_ts_data,
 )
 
 logger = logging.getLogger(__name__)
@@ -33,18 +32,24 @@ class TimeSeriesTask(Task):
         if self._estimators is None:
             # put this into a function to avoid circular dependency
             from flaml.automl.time_series import (
+                ARIMA,
+                LGBM_TS,
+                RF_TS,
+                SARIMAX,
+                Average,
+                CatBoost_TS,
+                ExtraTrees_TS,
+                HoltWinters,
+                LassoLars_TS,
+                Naive,
+                Orbit,
+                Prophet,
+                SeasonalAverage,
+                SeasonalNaive,
+                TCNEstimator,
+                TemporalFusionTransformerEstimator,
                 XGBoost_TS,
                 XGBoostLimitDepth_TS,
-                RF_TS,
-                LGBM_TS,
-                ExtraTrees_TS,
-                CatBoost_TS,
-                Prophet,
-                Orbit,
-                ARIMA,
-                SARIMAX,
-                TemporalFusionTransformerEstimator,
-                HoltWinters,
             )
 
             self._estimators = {
@@ -58,8 +63,19 @@ class TimeSeriesTask(Task):
                 "holt-winters": HoltWinters,
                 "catboost": CatBoost_TS,
                 "tft": TemporalFusionTransformerEstimator,
+                "lassolars": LassoLars_TS,
+                "tcn": TCNEstimator,
+                "snaive": SeasonalNaive,
+                "naive": Naive,
+                "savg": SeasonalAverage,
+                "avg": Average,
             }
 
+            if self._estimators["tcn"] is None:
+                # remove TCN if import failed
+                del self._estimators["tcn"]
+                logger.info("Couldn't import pytorch_lightning, skipping TCN estimator")
+
             try:
                 from prophet import Prophet as foo
 
@@ -72,7 +88,7 @@ class TimeSeriesTask(Task):
 
                 self._estimators["orbit"] = Orbit
             except ImportError:
-                logger.info("Couldn't import Prophet, skipping")
+                logger.info("Couldn't import orbit, skipping")
 
         return self._estimators
 

flaml/automl/time_series/__init__.py

@@ -1,17 +1,27 @@
-from .ts_model import (
-    Prophet,
-    Orbit,
-    ARIMA,
-    SARIMAX,
-    HoltWinters,
-    LGBM_TS,
-    XGBoost_TS,
-    RF_TS,
-    ExtraTrees_TS,
-    XGBoostLimitDepth_TS,
-    CatBoost_TS,
-    TimeSeriesEstimator,
-)
 from .tft import TemporalFusionTransformerEstimator
+from .ts_model import (
+    ARIMA,
+    LGBM_TS,
+    RF_TS,
+    SARIMAX,
+    Average,
+    CatBoost_TS,
+    ExtraTrees_TS,
+    HoltWinters,
+    LassoLars_TS,
+    Naive,
+    Orbit,
+    Prophet,
+    SeasonalAverage,
+    SeasonalNaive,
+    TimeSeriesEstimator,
+    XGBoost_TS,
+    XGBoostLimitDepth_TS,
+)
+
+try:
+    from .tcn import TCNEstimator
+except ImportError:
+    TCNEstimator = None
 
 from .ts_data import TimeSeriesDataset

flaml/automl/time_series/feature.py

@@ -1,5 +1,5 @@
-import math
 import datetime
+import math
 from functools import lru_cache
 
 import pandas as pd

flaml/automl/time_series/sklearn.py

@@ -12,8 +12,8 @@ except ImportError:
     DataFrame = Series = None
 
 import numpy as np
-from sklearn.preprocessing import StandardScaler
 from sklearn.decomposition import PCA
+from sklearn.preprocessing import StandardScaler
 
 
 def make_lag_features(X: pd.DataFrame, y: pd.Series, lags: int):

flaml/automl/time_series/tcn.py

@@ -0,0 +1,285 @@
+# This file is adapted from
+# https://github.com/locuslab/TCN/blob/master/TCN/tcn.py
+# https://github.com/locuslab/TCN/blob/master/TCN/adding_problem/add_test.py
+
+import datetime
+import logging
+import time
+
+import pandas as pd
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from pytorch_lightning.callbacks import EarlyStopping, LearningRateMonitor
+from pytorch_lightning.loggers import TensorBoardLogger
+from torch.nn.utils import weight_norm
+from torch.utils.data import DataLoader, TensorDataset
+
+from flaml import tune
+from flaml.automl.data import add_time_idx_col
+from flaml.automl.logger import logger, logger_formatter
+from flaml.automl.time_series.ts_data import TimeSeriesDataset
+from flaml.automl.time_series.ts_model import TimeSeriesEstimator
+
+
+class Chomp1d(nn.Module):
+    def __init__(self, chomp_size):
+        super().__init__()
+        self.chomp_size = chomp_size
+
+    def forward(self, x):
+        return x[:, :, : -self.chomp_size].contiguous()
+
+
+class TemporalBlock(nn.Module):
+    def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, padding, dropout=0.2):
+        super().__init__()
+        self.conv1 = weight_norm(
+            nn.Conv1d(n_inputs, n_outputs, kernel_size, stride=stride, padding=padding, dilation=dilation)
+        )
+        self.chomp1 = Chomp1d(padding)
+        self.relu1 = nn.ReLU()
+        self.dropout1 = nn.Dropout(dropout)
+
+        self.conv2 = weight_norm(
+            nn.Conv1d(n_outputs, n_outputs, kernel_size, stride=stride, padding=padding, dilation=dilation)
+        )
+        self.chomp2 = Chomp1d(padding)
+        self.relu2 = nn.ReLU()
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.net = nn.Sequential(
+            self.conv1, self.chomp1, self.relu1, self.dropout1, self.conv2, self.chomp2, self.relu2, self.dropout2
+        )
+        self.downsample = nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None
+        self.relu = nn.ReLU()
+        self.init_weights()
+
+    def init_weights(self):
+        self.conv1.weight.data.normal_(0, 0.01)
+        self.conv2.weight.data.normal_(0, 0.01)
+        if self.downsample is not None:
+            self.downsample.weight.data.normal_(0, 0.01)
+
+    def forward(self, x):
+        out = self.net(x)
+        res = x if self.downsample is None else self.downsample(x)
+        return self.relu(out + res)
+
+
+class TCNForecaster(nn.Module):
+    def __init__(
+        self,
+        input_feature_num,
+        num_outputs,
+        num_channels,
+        kernel_size=2,
+        dropout=0.2,
+    ):
+        super().__init__()
+        layers = []
+        num_levels = len(num_channels)
+        for i in range(num_levels):
+            dilation_size = 2**i
+            in_channels = input_feature_num if i == 0 else num_channels[i - 1]
+            out_channels = num_channels[i]
+            layers += [
+                TemporalBlock(
+                    in_channels,
+                    out_channels,
+                    kernel_size,
+                    stride=1,
+                    dilation=dilation_size,
+                    padding=(kernel_size - 1) * dilation_size,
+                    dropout=dropout,
+                )
+            ]
+
+        self.network = nn.Sequential(*layers)
+        self.linear = nn.Linear(num_channels[-1], num_outputs)
+
+    def forward(self, x):
+        y1 = self.network(x)
+        return self.linear(y1[:, :, -1])
+
+
+class TCNForecasterLightningModule(pl.LightningModule):
+    def __init__(self, model: TCNForecaster, learning_rate: float = 1e-3):
+        super().__init__()
+        self.model = model
+        self.learning_rate = learning_rate
+        self.loss_fn = nn.MSELoss()
+
+    def forward(self, x):
+        return self.model(x)
+
+    def step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self.model(x)
+        loss = self.loss_fn(y_hat, y)
+        return loss
+
+    def training_step(self, batch, batch_idx):
+        loss = self.step(batch, batch_idx)
+        self.log("train_loss", loss)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        loss = self.step(batch, batch_idx)
+        self.log("val_loss", loss)
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=self.learning_rate)
+
+
+class DataframeDataset(torch.utils.data.Dataset):
+    def __init__(self, dataframe, target_column, features_columns, sequence_length, train=True):
+        self.data = torch.tensor(dataframe[features_columns].to_numpy(), dtype=torch.float)
+        self.sequence_length = sequence_length
+        if train:
+            self.labels = torch.tensor(dataframe[target_column].to_numpy(), dtype=torch.float)
+        self.is_train = train
+
+    def __len__(self):
+        return len(self.data) - self.sequence_length + 1
+
+    def __getitem__(self, idx):
+        data = self.data[idx : idx + self.sequence_length]
+        data = data.permute(1, 0)
+        if self.is_train:
+            label = self.labels[idx : idx + self.sequence_length]
+            return data, label
+        else:
+            return data
+
+
+class TCNEstimator(TimeSeriesEstimator):
+    """The class for tuning TCN Forecaster"""
+
+    @classmethod
+    def search_space(cls, data, task, pred_horizon, **params):
+        space = {
+            "num_levels": {
+                "domain": tune.randint(lower=4, upper=20),  # hidden = 2^num_hidden
+                "init_value": 4,
+            },
+            "num_hidden": {
+                "domain": tune.randint(lower=4, upper=8),  # hidden = 2^num_hidden
+                "init_value": 5,
+            },
+            "kernel_size": {
+                "domain": tune.choice([2, 3, 5, 7]),  # common choices for kernel size
+                "init_value": 3,
+            },
+            "dropout": {
+                "domain": tune.uniform(lower=0.0, upper=0.5),  # standard range for dropout
+                "init_value": 0.1,
+            },
+            "learning_rate": {
+                "domain": tune.loguniform(lower=1e-4, upper=1e-1),  # typical range for learning rate
+                "init_value": 1e-3,
+            },
+        }
+        return space
+
+    def __init__(self, task="ts_forecast", n_jobs=1, **params):
+        super().__init__(task, **params)
+        logging.getLogger("pytorch_lightning").setLevel(logging.WARNING)
+
+    def fit(self, X_train: TimeSeriesDataset, y_train=None, budget=None, **kwargs):
+        start_time = time.time()
+        if budget is not None:
+            deltabudget = datetime.timedelta(seconds=budget)
+        else:
+            deltabudget = None
+        X_train = self.enrich(X_train)
+        super().fit(X_train, y_train, budget, **kwargs)
+
+        self.batch_size = kwargs.get("batch_size", 64)
+        self.horizon = kwargs.get("period", 1)
+        self.feature_cols = X_train.time_varying_known_reals
+        self.target_col = X_train.target_names[0]
+
+        train_dataset = DataframeDataset(
+            X_train.train_data,
+            self.target_col,
+            self.feature_cols,
+            self.horizon,
+        )
+        train_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=False)
+        if not X_train.test_data.empty:
+            val_dataset = DataframeDataset(
+                X_train.test_data,
+                self.target_col,
+                self.feature_cols,
+                self.horizon,
+            )
+        else:
+            val_dataset = DataframeDataset(
+                X_train.train_data.sample(frac=0.2, random_state=kwargs.get("random_state", 0)),
+                self.target_col,
+                self.feature_cols,
+                self.horizon,
+            )
+
+        val_loader = DataLoader(val_dataset, batch_size=self.batch_size, shuffle=False)
+
+        model = TCNForecaster(
+            len(self.feature_cols),
+            self.horizon,
+            [2 ** self.params["num_hidden"]] * self.params["num_levels"],
+            self.params["kernel_size"],
+            self.params["dropout"],
+        )
+
+        pl_module = TCNForecasterLightningModule(model, self.params["learning_rate"])
+
+        # Training loop
+        # gpus is deprecated in v1.7 and removed in v2.0
+        # accelerator="auto" can cast all condition.
+        trainer = pl.Trainer(
+            max_epochs=kwargs.get("max_epochs", 10),
+            accelerator="auto",
+            callbacks=[
+                EarlyStopping(monitor="val_loss", min_delta=1e-4, patience=10, verbose=False, mode="min"),
+                LearningRateMonitor(),
+            ],
+            logger=TensorBoardLogger(kwargs.get("log_dir", "logs/lightning_logs")),  # logging results to a tensorboard
+            max_time=deltabudget,
+            enable_model_summary=False,
+            enable_progress_bar=False,
+        )
+        trainer.fit(
+            pl_module,
+            train_dataloaders=train_loader,
+            val_dataloaders=val_loader,
+        )
+        best_model = trainer.model
+        self._model = best_model
+        train_time = time.time() - start_time
+        return train_time
+
+    def predict(self, X):
+        X = self.enrich(X)
+        if isinstance(X, TimeSeriesDataset):
+            df = X.X_val
+        else:
+            df = X
+        dataset = DataframeDataset(
+            df,
+            self.target_col,
+            self.feature_cols,
+            self.horizon,
+            train=False,
+        )
+        data_loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False)
+        self._model.eval()
+        raw_preds = []
+        for batch_x in data_loader:
+            raw_pred = self._model(batch_x)
+            raw_preds.append(raw_pred)
+        raw_preds = torch.cat(raw_preds, dim=0)
+        preds = pd.Series(raw_preds.detach().numpy().ravel())
+        return preds

flaml/automl/time_series/tft.py

@@ -105,6 +105,7 @@ class TemporalFusionTransformerEstimator(TimeSeriesEstimator):
 
     def fit(self, X_train, y_train, budget=None, **kwargs):
         import warnings
+
         import pytorch_lightning as pl
         import torch
         from pytorch_forecasting import TemporalFusionTransformer

flaml/automl/time_series/ts_data.py

@@ -2,7 +2,7 @@ import copy
 import datetime
 import math
 from dataclasses import dataclass, field
-from typing import List, Optional, Callable, Dict, Generator, Union
+from typing import Callable, Dict, Generator, List, Optional, Union
 
 import numpy as np
 
@@ -10,9 +10,9 @@ try:
     import pandas as pd
     from pandas import DataFrame, Series, to_datetime
     from scipy.sparse import issparse
-    from sklearn.preprocessing import LabelEncoder
-    from sklearn.impute import SimpleImputer
     from sklearn.compose import ColumnTransformer
+    from sklearn.impute import SimpleImputer
+    from sklearn.preprocessing import LabelEncoder
 
     from .feature import monthly_fourier_features
 except ImportError:
@@ -26,6 +26,8 @@ except ImportError:
     DataFrame = Series = None
 
 
+# dataclass will remove empty default value even with field(default_factory=lambda: [])
+# Change into default=None to place the attr
 @dataclass
 class TimeSeriesDataset:
     train_data: pd.DataFrame
@@ -34,10 +36,10 @@ class TimeSeriesDataset:
     target_names: List[str]
     frequency: str
     test_data: pd.DataFrame
-    time_varying_known_categoricals: List[str] = field(default_factory=lambda: [])
-    time_varying_known_reals: List[str] = field(default_factory=lambda: [])
-    time_varying_unknown_categoricals: List[str] = field(default_factory=lambda: [])
-    time_varying_unknown_reals: List[str] = field(default_factory=lambda: [])
+    time_varying_known_categoricals: List[str] = field(default=None)
+    time_varying_known_reals: List[str] = field(default=None)
+    time_varying_unknown_categoricals: List[str] = field(default=None)
+    time_varying_unknown_reals: List[str] = field(default=None)
 
     def __init__(
         self,
@@ -391,7 +393,7 @@ class DataTransformerTS:
 
         for column in X.columns:
             # sklearn/utils/validation.py needs int/float values
-            if X[column].dtype.name in ("object", "category"):
+            if X[column].dtype.name in ("object", "category", "string"):
                 if (
                     # drop columns where all values are the same
                     X[column].nunique() == 1
@@ -403,7 +405,7 @@ class DataTransformerTS:
                     self.cat_columns.append(column)
             elif X[column].nunique(dropna=True) < 2:
                 self.drop_columns.append(column)
-            elif X[column].dtype.name == "datetime64[ns]":
+            elif X[column].dtype.name in ["datetime64[ns]", "datetime64[s]"]:
                 pass  # these will be processed at model level,
                 # so they can also be done in the predict method
             else:

flaml/automl/time_series/ts_model.py

@@ -1,8 +1,8 @@
-import time
 import logging
-import os
-from datetime import datetime
 import math
+import os
+import time
+from datetime import datetime
 from typing import List, Optional, Union
 
 try:
@@ -22,26 +22,27 @@ except ImportError:
 import numpy as np
 
 from flaml import tune
-from flaml.model import (
-    suppress_stdout_stderr,
-    SKLearnEstimator,
-    logger,
-    LGBMEstimator,
-    XGBoostSklearnEstimator,
-    RandomForestEstimator,
-    ExtraTreesEstimator,
-    XGBoostLimitDepthEstimator,
+from flaml.automl.data import TS_TIMESTAMP_COL, TS_VALUE_COL
+from flaml.automl.model import (
     CatBoostEstimator,
-)
-from flaml.data import TS_TIMESTAMP_COL, TS_VALUE_COL
-from flaml.automl.time_series.ts_data import (
-    TimeSeriesDataset,
-    enrich_dataset,
-    enrich_dataframe,
-    normalize_ts_data,
-    create_forward_frame,
+    ExtraTreesEstimator,
+    LassoLarsEstimator,
+    LGBMEstimator,
+    RandomForestEstimator,
+    SKLearnEstimator,
+    XGBoostLimitDepthEstimator,
+    XGBoostSklearnEstimator,
+    logger,
+    suppress_stdout_stderr,
 )
 from flaml.automl.task import Task
+from flaml.automl.time_series.ts_data import (
+    TimeSeriesDataset,
+    create_forward_frame,
+    enrich_dataframe,
+    enrich_dataset,
+    normalize_ts_data,
+)
 
 
 class TimeSeriesEstimator(SKLearnEstimator):
@@ -143,6 +144,7 @@ class TimeSeriesEstimator(SKLearnEstimator):
 
     def score(self, X_val: DataFrame, y_val: Series, **kwargs):
         from sklearn.metrics import r2_score
+
         from ..ml import metric_loss_score
 
         y_pred = self.predict(X_val, **kwargs)
@@ -610,15 +612,13 @@ class HoltWinters(StatsModelsEstimator):
         ):  # this would prevent heuristic initialization to work properly
             self.params["seasonal"] = None
         if (
-            self.params["seasonal"] == "mul" and (train_df.y == 0).sum() > 0
+            self.params["seasonal"] == "mul" and (train_df[target_col] == 0).sum() > 0
         ):  # cannot have multiplicative seasonality in this case
             self.params["seasonal"] = "add"
-        if self.params["trend"] == "mul" and (train_df.y == 0).sum() > 0:
+        if self.params["trend"] == "mul" and (train_df[target_col] == 0).sum() > 0:
             self.params["trend"] = "add"
-
         if not self.params["seasonal"] or self.params["trend"] not in ["mul", "add"]:
             self.params["damped_trend"] = False
-
         model = HWExponentialSmoothing(
             train_df[[target_col]],
             damped_trend=self.params["damped_trend"],
@@ -632,6 +632,125 @@ class HoltWinters(StatsModelsEstimator):
         return train_time
 
 
+class SimpleForecaster(StatsModelsEstimator):
+    """Base class for Naive Forecaster like Seasonal Naive, Naive, Seasonal Average, Average"""
+
+    @classmethod
+    def _search_space(cls, data: TimeSeriesDataset, task: Task, pred_horizon: int, **params):
+        return {
+            "season": {
+                "domain": tune.randint(1, pred_horizon),
+                "init_value": pred_horizon,
+            }
+        }
+
+    def joint_preprocess(self, X_train, y_train=None):
+        X_train = self.enrich(X_train)
+
+        self.regressors = []
+
+        if isinstance(X_train, TimeSeriesDataset):
+            data = X_train
+            target_col = data.target_names[0]
+            # this class only supports univariate regression
+            train_df = data.train_data[self.regressors + [target_col]]
+            train_df.index = to_datetime(data.train_data[data.time_col])
+        else:
+            target_col = TS_VALUE_COL
+            train_df = self._join(X_train, y_train)
+
+        self.time_col = data.time_col
+        self.target_names = data.target_names
+
+        train_df = self._preprocess(train_df)
+        return train_df, target_col
+
+    def fit(self, X_train, y_train=None, budget=None, **kwargs):
+        import warnings
+
+        warnings.filterwarnings("ignore")
+        from statsmodels.tsa.holtwinters import SimpleExpSmoothing
+
+        self.season = self.params.get("season", 1)
+        current_time = time.time()
+        super().fit(X_train, y_train, budget=budget, **kwargs)
+
+        train_df, target_col = self.joint_preprocess(X_train, y_train)
+
+        model = SimpleExpSmoothing(
+            train_df[[target_col]],
+        )
+        with suppress_stdout_stderr():
+            model = model.fit(smoothing_level=self.smoothing_level)
+        train_time = time.time() - current_time
+        self._model = model
+        return train_time
+
+
+class SeasonalNaive(SimpleForecaster):
+    smoothing_level = 1.0
+
+    def predict(self, X, **kwargs):
+        if isinstance(X, int):
+            forecasts = []
+            for i in range(X):
+                forecast = self._model.forecast(steps=self.season)[0]
+                forecasts.append(forecast)
+            return pd.Series(forecasts)
+        else:
+            return super().predict(X, **kwargs)
+
+
+class Naive(SimpleForecaster):
+    smoothing_level = 0.0
+
+    @classmethod
+    def _search_space(cls, data: TimeSeriesDataset, task: Task, pred_horizon: int, **params):
+        return {}
+
+    def predict(self, X, **kwargs):
+        if isinstance(X, int):
+            last_observation = self._model.params["initial_level"]
+            return pd.Series([last_observation] * X)
+        else:
+            return super().predict(X, **kwargs)
+
+
+class SeasonalAverage(SimpleForecaster):
+    def fit(self, X_train, y_train=None, budget=None, **kwargs):
+        from statsmodels.tsa.ar_model import AutoReg, ar_select_order
+
+        start_time = time.time()
+
+        self.season = kwargs.get("season", 1)  # seasonality period
+        train_df, target_col = self.joint_preprocess(X_train, y_train)
+        selection_res = ar_select_order(train_df[target_col], maxlag=self.season)
+
+        # Fit autoregressive model with optimal order
+        model = AutoReg(train_df[target_col], lags=selection_res.ar_lags)
+        self._model = model.fit()
+        end_time = time.time()
+
+        return end_time - start_time
+
+
+class Average(SimpleForecaster):
+    @classmethod
+    def _search_space(cls, data: TimeSeriesDataset, task: Task, pred_horizon: int, **params):
+        return {}
+
+    def fit(self, X_train, y_train=None, budget=None, **kwargs):
+        from statsmodels.tsa.ar_model import AutoReg
+
+        start_time = time.time()
+        train_df, target_col = self.joint_preprocess(X_train, y_train)
+        model = AutoReg(train_df[target_col], lags=0)
+        self._model = model.fit()
+        end_time = time.time()
+
+        return end_time - start_time
+
+
 class TS_SKLearn(TimeSeriesEstimator):
     """The class for tuning SKLearn Regressors for time-series forecasting"""
 
@@ -758,3 +877,7 @@ class XGBoostLimitDepth_TS(TS_SKLearn):
 # catboost regressor is invalid because it has a `name` parameter, making it incompatible with hcrystalball
 class CatBoost_TS(TS_SKLearn):
     base_class = CatBoostEstimator
+
+
+class LassoLars_TS(TS_SKLearn):
+    base_class = LassoLarsEstimator

flaml/automl/training_log.py

@@ -4,14 +4,14 @@
 """
 
 import json
-from typing import IO
-from contextlib import contextmanager
 import logging
+from contextlib import contextmanager
+from typing import IO
 
 logger = logging.getLogger("flaml.automl")
 
 
-class TrainingLogRecord(object):
+class TrainingLogRecord:
     def __init__(
         self,
         record_id: int,
@@ -52,7 +52,7 @@ class TrainingLogCheckPoint(TrainingLogRecord):
         self.curr_best_record_id = curr_best_record_id
 
 
-class TrainingLogWriter(object):
+class TrainingLogWriter:
     def __init__(self, output_filename: str):
         self.output_filename = output_filename
         self.file = None
@@ -79,7 +79,7 @@ class TrainingLogWriter(object):
         sample_size,
     ):
         if self.file is None:
-            raise IOError("Call open() to open the output file first.")
+            raise OSError("Call open() to open the output file first.")
         if validation_loss is None:
             raise ValueError("TEST LOSS NONE ERROR!!!")
         record = TrainingLogRecord(
@@ -109,7 +109,7 @@ class TrainingLogWriter(object):
 
     def checkpoint(self):
         if self.file is None:
-            raise IOError("Call open() to open the output file first.")
+            raise OSError("Call open() to open the output file first.")
         if self.current_best_loss_record_id is None:
             logger.warning("flaml.training_log: checkpoint() called before any record is written, skipped.")
             return
@@ -124,7 +124,7 @@ class TrainingLogWriter(object):
         self.file = None  # for pickle
 
 
-class TrainingLogReader(object):
+class TrainingLogReader:
     def __init__(self, filename: str):
         self.filename = filename
         self.file = None
@@ -134,7 +134,7 @@ class TrainingLogReader(object):
 
     def records(self):
         if self.file is None:
-            raise IOError("Call open() before reading log file.")
+            raise OSError("Call open() before reading log file.")
         for line in self.file:
             data = json.loads(line)
             if len(data) == 1:
@@ -149,7 +149,7 @@ class TrainingLogReader(object):
 
     def get_record(self, record_id) -> TrainingLogRecord:
         if self.file is None:
-            raise IOError("Call open() before reading log file.")
+            raise OSError("Call open() before reading log file.")
         for rec in self.records():
             if rec.record_id == record_id:
                 return rec

flaml/data.py

@@ -1,9 +0,0 @@
-import warnings
-
-from flaml.automl.data import *
-
-
-warnings.warn(
-    "Importing from `flaml.data` is deprecated. Please use `flaml.automl.data`.",
-    DeprecationWarning,
-)

flaml/default/README.md

@@ -14,7 +14,6 @@ estimator.fit(X_train, y_train)
 estimator.predict(X_test, y_test)
 ```
 
-
 1. Use AutoML.fit(). set `starting_points="data"` and `max_iter=0`.
 
 ```python
@@ -36,10 +35,17 @@ automl.fit(X_train, y_train, **automl_settings)
 from flaml.default import preprocess_and_suggest_hyperparams
 
 X, y = load_iris(return_X_y=True, as_frame=True)
-X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
-hyperparams, estimator_class, X_transformed, y_transformed, feature_transformer, label_transformer = preprocess_and_suggest_hyperparams(
-    "classification", X_train, y_train, "lgbm"
+X_train, X_test, y_train, y_test = train_test_split(
+    X, y, test_size=0.33, random_state=42
 )
+(
+    hyperparams,
+    estimator_class,
+    X_transformed,
+    y_transformed,
+    feature_transformer,
+    label_transformer,
+) = preprocess_and_suggest_hyperparams("classification", X_train, y_train, "lgbm")
 model = estimator_class(**hyperparams)  # estimator_class is LGBMClassifier
 model.fit(X_transformed, y_train)  # LGBMClassifier can handle raw labels
 X_test = feature_transformer.transform(X_test)  # preprocess test data
@@ -172,7 +178,7 @@ Change "binary" into "multiclass" or "regression" for the other tasks.
 
 For more technical details, please check our research paper.
 
-* [Mining Robust Default Configurations for Resource-constrained AutoML](https://arxiv.org/abs/2202.09927). Moe Kayali, Chi Wang. arXiv preprint arXiv:2202.09927 (2022).
+- [Mining Robust Default Configurations for Resource-constrained AutoML](https://arxiv.org/abs/2202.09927). Moe Kayali, Chi Wang. arXiv preprint arXiv:2202.09927 (2022).
 
 ```bibtex
 @article{Kayali2022default,

flaml/default/__init__.py

@@ -1,18 +1,18 @@
-from .suggest import (
-    suggest_config,
-    suggest_learner,
-    suggest_hyperparams,
-    preprocess_and_suggest_hyperparams,
-    meta_feature,
-)
 from .estimator import (
-    flamlize_estimator,
-    LGBMClassifier,
-    LGBMRegressor,
-    XGBClassifier,
-    XGBRegressor,
-    RandomForestClassifier,
-    RandomForestRegressor,
     ExtraTreesClassifier,
     ExtraTreesRegressor,
+    LGBMClassifier,
+    LGBMRegressor,
+    RandomForestClassifier,
+    RandomForestRegressor,
+    XGBClassifier,
+    XGBRegressor,
+    flamlize_estimator,
+)
+from .suggest import (
+    meta_feature,
+    preprocess_and_suggest_hyperparams,
+    suggest_config,
+    suggest_hyperparams,
+    suggest_learner,
 )

flaml/default/estimator.py

@@ -1,5 +1,7 @@
 from functools import wraps
+
 from flaml.automl.task.task import CLASSIFICATION
+
 from .suggest import preprocess_and_suggest_hyperparams
 
 DEFAULT_LOCATION = "default_location"
@@ -105,7 +107,12 @@ def flamlize_estimator(super_class, name: str, task: str, alternatives=None):
                 # if hasattr(self, "_classes"):
                 #     self._classes = self._label_transformer.classes_
                 # else:
-                self.classes_ = self._label_transformer.classes_
+                try:
+                    self.classes_ = self._label_transformer.classes_
+                except AttributeError:
+                    # xgboost 2: AttributeError: can't set attribute
+                    if "xgb" not in estimator_name:
+                        raise
                 if "xgb" not in estimator_name:
                     # rf and et would do inverse transform automatically; xgb doesn't
                     self._label_transformer = None

flaml/default/greedy.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pandas as pd
-from sklearn.preprocessing import RobustScaler
 from sklearn.metrics import pairwise_distances
+from sklearn.preprocessing import RobustScaler
 
 
 def _augment(row):
@@ -12,7 +12,7 @@ def _augment(row):
 def construct_portfolio(regret_matrix, meta_features, regret_bound):
     """The portfolio construction algorithm.
 
-    (Reference)[https://arxiv.org/abs/2202.09927].
+    Reference: [Mining Robust Default Configurations for Resource-constrained AutoML](https://arxiv.org/abs/2202.09927).
 
     Args:
         regret_matrix: A dataframe of regret matrix.

flaml/default/portfolio.py

@@ -1,11 +1,13 @@
-import pandas as pd
-import numpy as np
 import argparse
-from pathlib import Path
 import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
 from sklearn.preprocessing import RobustScaler
+
 from flaml.default import greedy
-from flaml.default.regret import load_result, build_regret
+from flaml.default.regret import build_regret, load_result
 from flaml.version import __version__
 
 regret_bound = 0.01
@@ -67,7 +69,7 @@ def build_portfolio(meta_features, regret, strategy):
 
 def load_json(filename):
     """Returns the contents of json file filename."""
-    with open(filename, "r") as f:
+    with open(filename) as f:
         return json.load(f)
 
 

flaml/default/regret.py

@@ -1,5 +1,6 @@
 import argparse
 from os import path
+
 import pandas as pd
 
 

flaml/default/suggest.py

@@ -1,11 +1,13 @@
-import numpy as np
+import json
 import logging
 import pathlib
-import json
+
+import numpy as np
+
 from flaml.automl.data import DataTransformer
-from flaml.automl.task.task import CLASSIFICATION, get_classification_objective
-from flaml.automl.task.generic_task import len_labels
 from flaml.automl.task.factory import task_factory
+from flaml.automl.task.generic_task import len_labels
+from flaml.automl.task.task import CLASSIFICATION, get_classification_objective
 from flaml.version import __version__
 
 try:
@@ -41,7 +43,7 @@ def meta_feature(task, X_train, y_train, meta_feature_names):
                 # 'numpy.ndarray' object has no attribute 'select_dtypes'
                 this_feature.append(1)  # all features are numeric
         else:
-            raise ValueError("Feature {} not implemented. ".format(each_feature_name))
+            raise ValueError(f"Feature {each_feature_name} not implemented. ")
 
     return this_feature
 
@@ -55,7 +57,7 @@ def load_config_predictor(estimator_name, task, location=None):
     task = "multiclass" if task == "multi" else task  # TODO: multi -> multiclass?
     try:
         location = location or LOCATION
-        with open(f"{location}/{estimator_name}/{task}.json", "r") as f:
+        with open(f"{location}/{estimator_name}/{task}.json") as f:
             CONFIG_PREDICTORS[key] = predictor = json.load(f)
     except FileNotFoundError:
         raise FileNotFoundError(f"Portfolio has not been built for {estimator_name} on {task} task.")

flaml/ml.py

@@ -2,7 +2,6 @@ import warnings
 
 from flaml.automl.ml import *
 
-
 warnings.warn(
     "Importing from `flaml.ml` is deprecated. Please use `flaml.automl.ml`.",
     DeprecationWarning,

flaml/model.py

@@ -1,9 +0,0 @@
-import warnings
-
-from flaml.automl.model import *
-
-
-warnings.warn(
-    "Importing from `flaml.model` is deprecated. Please use `flaml.automl.model`.",
-    DeprecationWarning,
-)

flaml/onlineml/README.md

@@ -4,7 +4,8 @@ FLAML includes *ChaCha* which is an automatic hyperparameter tuning solution for
 
 For more technical details about *ChaCha*, please check our paper.
 
-* [ChaCha for Online AutoML](https://www.microsoft.com/en-us/research/publication/chacha-for-online-automl/). Qingyun Wu, Chi Wang, John Langford, Paul Mineiro and Marco Rossi. ICML 2021.
+- [ChaCha for Online AutoML](https://www.microsoft.com/en-us/research/publication/chacha-for-online-automl/). Qingyun Wu, Chi Wang, John Langford, Paul Mineiro and Marco Rossi. ICML 2021.
+
 ```
 @inproceedings{wu2021chacha,
     title={ChaCha for online AutoML},
@@ -23,8 +24,9 @@ An example of online namespace interactions tuning in VW:
 ```python
 # require: pip install flaml[vw]
 from flaml import AutoVW
-'''create an AutoVW instance for tuning namespace interactions'''
-autovw = AutoVW(max_live_model_num=5, search_space={'interactions': AutoVW.AUTOMATIC})
+
+"""create an AutoVW instance for tuning namespace interactions"""
+autovw = AutoVW(max_live_model_num=5, search_space={"interactions": AutoVW.AUTOMATIC})
 ```
 
 An example of online tuning of both namespace interactions and learning rate in VW:
@@ -33,12 +35,18 @@ An example of online tuning of both namespace interactions and learning rate in
 # require: pip install flaml[vw]
 from flaml import AutoVW
 from flaml.tune import loguniform
-''' create an AutoVW instance for tuning namespace interactions and learning rate'''
+
+""" create an AutoVW instance for tuning namespace interactions and learning rate"""
 # set up the search space and init config
-search_space_nilr = {'interactions': AutoVW.AUTOMATIC, 'learning_rate': loguniform(lower=2e-10, upper=1.0)}
-init_config_nilr = {'interactions': set(), 'learning_rate': 0.5}
+search_space_nilr = {
+    "interactions": AutoVW.AUTOMATIC,
+    "learning_rate": loguniform(lower=2e-10, upper=1.0),
+}
+init_config_nilr = {"interactions": set(), "learning_rate": 0.5}
 # create an AutoVW instance
-autovw = AutoVW(max_live_model_num=5, search_space=search_space_nilr, init_config=init_config_nilr)
+autovw = AutoVW(
+    max_live_model_num=5, search_space=search_space_nilr, init_config=init_config_nilr
+)
 ```
 
 A user can use the resulting AutoVW instances `autovw` in a similar way to a vanilla Vowpal Wabbit instance, i.e., `pyvw.vw`, to perform online learning by iteratively calling its `predict(data_example)` and `learn(data_example)` functions at each data example.

flaml/onlineml/autovw.py

@@ -1,16 +1,17 @@
-from typing import Optional, Union
 import logging
+from typing import Optional, Union
+
+from flaml.onlineml import OnlineTrialRunner
+from flaml.onlineml.trial import get_ns_feature_dim_from_vw_example
 from flaml.tune import (
-    Trial,
     Categorical,
     Float,
     PolynomialExpansionSet,
+    Trial,
     polynomial_expansion_set,
 )
-from flaml.onlineml import OnlineTrialRunner
 from flaml.tune.scheduler import ChaChaScheduler
 from flaml.tune.searcher import ChampionFrontierSearcher
-from flaml.onlineml.trial import get_ns_feature_dim_from_vw_example
 
 logger = logging.getLogger(__name__)
 
@@ -140,7 +141,7 @@ class AutoVW:
             max_live_model_num=self._max_live_model_num,
             searcher=searcher,
             scheduler=scheduler,
-            **self._automl_runner_args
+            **self._automl_runner_args,
         )
 
     def predict(self, data_sample):

flaml/onlineml/trial.py

@@ -1,14 +1,16 @@
-import numpy as np
-import logging
-import time
-import math
-import copy
 import collections
+import copy
+import logging
+import math
+import time
 from typing import Optional, Union
+
+import numpy as np
+
 from flaml.tune import Trial
 
 try:
-    from sklearn.metrics import mean_squared_error, mean_absolute_error
+    from sklearn.metrics import mean_absolute_error, mean_squared_error
 except ImportError:
     pass
 

flaml/onlineml/trial_runner.py

@@ -1,10 +1,11 @@
-import numpy as np
+import logging
 import math
+
+import numpy as np
+
 from flaml.tune import Trial
 from flaml.tune.scheduler import TrialScheduler
 
-import logging
-
 logger = logging.getLogger(__name__)
 
 

flaml/tune/README.md

@@ -5,45 +5,47 @@ It can be used standalone, or together with ray tune or nni. Please find detaile
 
 Below are some quick examples.
 
-* Example for sequential tuning (recommended when compute resource is limited and each trial can consume all the resources):
+- Example for sequential tuning (recommended when compute resource is limited and each trial can consume all the resources):
 
 ```python
 # require: pip install flaml[blendsearch]
 from flaml import tune
 import time
 
+
 def evaluate_config(config):
-    '''evaluate a hyperparameter configuration'''
+    """evaluate a hyperparameter configuration"""
     # we uss a toy example with 2 hyperparameters
-    metric = (round(config['x'])-85000)**2 - config['x']/config['y']
+    metric = (round(config["x"]) - 85000) ** 2 - config["x"] / config["y"]
     # usually the evaluation takes an non-neglible cost
     # and the cost could be related to certain hyperparameters
     # in this example, we assume it's proportional to x
-    time.sleep(config['x']/100000)
+    time.sleep(config["x"] / 100000)
     # use tune.report to report the metric to optimize
     tune.report(metric=metric)
 
+
 analysis = tune.run(
-    evaluate_config,    # the function to evaluate a config
+    evaluate_config,  # the function to evaluate a config
     config={
-        'x': tune.lograndint(lower=1, upper=100000),
-        'y': tune.randint(lower=1, upper=100000)
-    }, # the search space
-    low_cost_partial_config={'x':1},    # a initial (partial) config with low cost
-    metric='metric',    # the name of the metric used for optimization
-    mode='min',         # the optimization mode, 'min' or 'max'
-    num_samples=-1,    # the maximal number of configs to try, -1 means infinite
-    time_budget_s=60,   # the time budget in seconds
-    local_dir='logs/',  # the local directory to store logs
+        "x": tune.lograndint(lower=1, upper=100000),
+        "y": tune.randint(lower=1, upper=100000),
+    },  # the search space
+    low_cost_partial_config={"x": 1},  # a initial (partial) config with low cost
+    metric="metric",  # the name of the metric used for optimization
+    mode="min",  # the optimization mode, 'min' or 'max'
+    num_samples=-1,  # the maximal number of configs to try, -1 means infinite
+    time_budget_s=60,  # the time budget in seconds
+    local_dir="logs/",  # the local directory to store logs
     # verbose=0,          # verbosity
     # use_ray=True, # uncomment when performing parallel tuning using ray
-    )
+)
 
 print(analysis.best_trial.last_result)  # the best trial's result
-print(analysis.best_config) # the best config
+print(analysis.best_config)  # the best config
 ```
 
-* Example for using ray tune's API:
+- Example for using ray tune's API:
 
 ```python
 # require: pip install flaml[blendsearch,ray]
@@ -51,36 +53,39 @@ from ray import tune as raytune
 from flaml import CFO, BlendSearch
 import time
 
+
 def evaluate_config(config):
-    '''evaluate a hyperparameter configuration'''
+    """evaluate a hyperparameter configuration"""
     # we use a toy example with 2 hyperparameters
-    metric = (round(config['x'])-85000)**2 - config['x']/config['y']
+    metric = (round(config["x"]) - 85000) ** 2 - config["x"] / config["y"]
     # usually the evaluation takes a non-neglible cost
     # and the cost could be related to certain hyperparameters
     # in this example, we assume it's proportional to x
-    time.sleep(config['x']/100000)
+    time.sleep(config["x"] / 100000)
     # use tune.report to report the metric to optimize
     tune.report(metric=metric)
 
+
 # provide a time budget (in seconds) for the tuning process
 time_budget_s = 60
 # provide the search space
 config_search_space = {
-        'x': tune.lograndint(lower=1, upper=100000),
-        'y': tune.randint(lower=1, upper=100000)
-    }
+    "x": tune.lograndint(lower=1, upper=100000),
+    "y": tune.randint(lower=1, upper=100000),
+}
 # provide the low cost partial config
-low_cost_partial_config={'x':1}
+low_cost_partial_config = {"x": 1}
 
 # set up CFO
 cfo = CFO(low_cost_partial_config=low_cost_partial_config)
 
 # set up BlendSearch
 blendsearch = BlendSearch(
-    metric="metric", mode="min",
+    metric="metric",
+    mode="min",
     space=config_search_space,
     low_cost_partial_config=low_cost_partial_config,
-    time_budget_s=time_budget_s
+    time_budget_s=time_budget_s,
 )
 # NOTE: when using BlendSearch as a search_alg in ray tune, you need to
 # configure the 'time_budget_s' for BlendSearch accordingly such that
@@ -89,28 +94,28 @@ blendsearch = BlendSearch(
 # automatically in flaml.
 
 analysis = raytune.run(
-    evaluate_config,    # the function to evaluate a config
+    evaluate_config,  # the function to evaluate a config
     config=config_search_space,
-    metric='metric',    # the name of the metric used for optimization
-    mode='min',         # the optimization mode, 'min' or 'max'
-    num_samples=-1,     # the maximal number of configs to try, -1 means infinite
-    time_budget_s=time_budget_s,   # the time budget in seconds
-    local_dir='logs/',  # the local directory to store logs
-    search_alg=blendsearch  # or cfo
+    metric="metric",  # the name of the metric used for optimization
+    mode="min",  # the optimization mode, 'min' or 'max'
+    num_samples=-1,  # the maximal number of configs to try, -1 means infinite
+    time_budget_s=time_budget_s,  # the time budget in seconds
+    local_dir="logs/",  # the local directory to store logs
+    search_alg=blendsearch,  # or cfo
 )
 
 print(analysis.best_trial.last_result)  # the best trial's result
 print(analysis.best_config)  # the best config
 ```
 
-* Example for using NNI: An example of using BlendSearch with NNI can be seen in [test](https://github.com/microsoft/FLAML/tree/main/test/nni). CFO can be used as well in a similar manner. To run the example, first make sure you have [NNI](https://nni.readthedocs.io/en/stable/) installed, then run:
+- Example for using NNI: An example of using BlendSearch with NNI can be seen in [test](https://github.com/microsoft/FLAML/tree/main/test/nni). CFO can be used as well in a similar manner. To run the example, first make sure you have [NNI](https://nni.readthedocs.io/en/stable/) installed, then run:
 
 ```shell
 $nnictl create --config ./config.yml
 ```
 
-* For more examples, please check out
-[notebooks](https://github.com/microsoft/FLAML/tree/main/notebook/).
+- For more examples, please check out
+  [notebooks](https://github.com/microsoft/FLAML/tree/main/notebook/).
 
 `flaml` offers two HPO methods: CFO and BlendSearch.
 `flaml.tune` uses BlendSearch by default.
@@ -185,16 +190,16 @@ tune.run(...
 )
 ```
 
-* Recommended scenario: cost-related hyperparameters exist, a low-cost
-initial point is known, and the search space is complex such that local search
-is prone to be stuck at local optima.
+- Recommended scenario: cost-related hyperparameters exist, a low-cost
+  initial point is known, and the search space is complex such that local search
+  is prone to be stuck at local optima.
 
-* Suggestion about using larger search space in BlendSearch:
-In hyperparameter optimization, a larger search space is desirable because it is more likely to include the optimal configuration (or one of the optimal configurations) in hindsight. However the performance (especially anytime performance) of most existing HPO methods is undesirable if the cost of the configurations in the search space has a large variation. Thus hand-crafted small search spaces (with relatively homogeneous cost) are often used in practice for these methods, which is subject to idiosyncrasy. BlendSearch combines the benefits of local search and global search, which enables a smart (economical) way of deciding where to explore in the search space even though it is larger than necessary. This allows users to specify a larger search space in BlendSearch, which is often easier and a better practice than narrowing down the search space by hand.
+- Suggestion about using larger search space in BlendSearch:
+  In hyperparameter optimization, a larger search space is desirable because it is more likely to include the optimal configuration (or one of the optimal configurations) in hindsight. However the performance (especially anytime performance) of most existing HPO methods is undesirable if the cost of the configurations in the search space has a large variation. Thus hand-crafted small search spaces (with relatively homogeneous cost) are often used in practice for these methods, which is subject to idiosyncrasy. BlendSearch combines the benefits of local search and global search, which enables a smart (economical) way of deciding where to explore in the search space even though it is larger than necessary. This allows users to specify a larger search space in BlendSearch, which is often easier and a better practice than narrowing down the search space by hand.
 
 For more technical details, please check our papers.
 
-* [Frugal Optimization for Cost-related Hyperparameters](https://arxiv.org/abs/2005.01571). Qingyun Wu, Chi Wang, Silu Huang. AAAI 2021.
+- [Frugal Optimization for Cost-related Hyperparameters](https://arxiv.org/abs/2005.01571). Qingyun Wu, Chi Wang, Silu Huang. AAAI 2021.
 
 ```bibtex
 @inproceedings{wu2021cfo,
@@ -205,7 +210,7 @@ For more technical details, please check our papers.
 }
 ```
 
-* [Economical Hyperparameter Optimization With Blended Search Strategy](https://www.microsoft.com/en-us/research/publication/economical-hyperparameter-optimization-with-blended-search-strategy/). Chi Wang, Qingyun Wu, Silu Huang, Amin Saied. ICLR 2021.
+- [Economical Hyperparameter Optimization With Blended Search Strategy](https://www.microsoft.com/en-us/research/publication/economical-hyperparameter-optimization-with-blended-search-strategy/). Chi Wang, Qingyun Wu, Silu Huang, Amin Saied. ICLR 2021.
 
 ```bibtex
 @inproceedings{wang2021blendsearch,

flaml/tune/__init__.py

@@ -3,16 +3,16 @@ try:
 
     assert ray_version >= "1.10.0"
     from ray.tune import (
-        uniform,
+        lograndint,
+        loguniform,
+        qlograndint,
+        qloguniform,
+        qrandint,
+        qrandn,
         quniform,
         randint,
-        qrandint,
         randn,
-        qrandn,
-        loguniform,
-        qloguniform,
-        lograndint,
-        qlograndint,
+        uniform,
     )
 
     if ray_version.startswith("1."):
@@ -20,21 +20,20 @@ try:
     else:
         from ray.tune.search import sample
 except (ImportError, AssertionError):
+    from . import sample
     from .sample import (
-        uniform,
+        lograndint,
+        loguniform,
+        qlograndint,
+        qloguniform,
+        qrandint,
+        qrandn,
         quniform,
         randint,
-        qrandint,
         randn,
-        qrandn,
-        loguniform,
-        qloguniform,
-        lograndint,
-        qlograndint,
+        uniform,
     )
-    from . import sample
-from .tune import run, report, INCUMBENT_RESULT
-from .sample import polynomial_expansion_set
-from .sample import PolynomialExpansionSet, Categorical, Float
+from .sample import Categorical, Float, PolynomialExpansionSet, polynomial_expansion_set
 from .trial import Trial
+from .tune import INCUMBENT_RESULT, report, run
 from .utils import choice

flaml/tune/analysis.py

@@ -15,10 +15,12 @@
 # This source file is adapted here because ray does not fully support Windows.
 
 # Copyright (c) Microsoft Corporation.
-from typing import Dict, Optional
-import numpy as np
-from .trial import Trial
 import logging
+from typing import Dict, Optional
+
+import numpy as np
+
+from .trial import Trial
 
 logger = logging.getLogger(__name__)
 

flaml/tune/logger.py

@@ -0,0 +1,37 @@
+import logging
+import os
+
+
+class ColoredFormatter(logging.Formatter):
+    # ANSI escape codes for colors
+    COLORS = {
+        # logging.DEBUG: "\033[36m",  # Cyan
+        # logging.INFO: "\033[32m",   # Green
+        logging.WARNING: "\033[33m",  # Yellow
+        logging.ERROR: "\033[31m",  # Red
+        logging.CRITICAL: "\033[1;31m",  # Bright Red
+    }
+    RESET = "\033[0m"  # Reset to default
+
+    def __init__(self, fmt, datefmt, use_color=True):
+        super().__init__(fmt, datefmt)
+        self.use_color = use_color
+
+    def format(self, record):
+        formatted = super().format(record)
+        if self.use_color:
+            color = self.COLORS.get(record.levelno, "")
+            if color:
+                return f"{color}{formatted}{self.RESET}"
+        return formatted
+
+
+logger = logging.getLogger(__name__)
+use_color = True
+if os.getenv("FLAML_LOG_NO_COLOR"):
+    use_color = False
+
+logger_formatter = ColoredFormatter(
+    "[%(name)s: %(asctime)s] {%(lineno)d} %(levelname)s - %(message)s", "%m-%d %H:%M:%S", use_color
+)
+logger.propagate = False

flaml/tune/sample.py

@@ -19,6 +19,7 @@ import logging
 from copy import copy
 from math import isclose
 from typing import Any, Dict, List, Optional, Sequence, Union
+
 import numpy as np
 
 # Backwards compatibility

flaml/tune/scheduler/__init__.py

@@ -1,6 +1,6 @@
-from .trial_scheduler import TrialScheduler
 from .online_scheduler import (
+    ChaChaScheduler,
     OnlineScheduler,
     OnlineSuccessiveDoublingScheduler,
-    ChaChaScheduler,
 )
+from .trial_scheduler import TrialScheduler

flaml/tune/scheduler/online_scheduler.py

@@ -1,9 +1,12 @@
-import numpy as np
 import logging
 from typing import Dict
-from flaml.tune.scheduler import TrialScheduler
+
+import numpy as np
+
 from flaml.tune import Trial
 
+from .trial_scheduler import TrialScheduler
+
 logger = logging.getLogger(__name__)
 
 

flaml/tune/searcher/blendsearch.py

@@ -2,10 +2,11 @@
 #  * Copyright (c) Microsoft Corporation. All rights reserved.
 #  * Licensed under the MIT License. See LICENSE file in the
 #  * project root for license information.
-from typing import Dict, Optional, List, Tuple, Callable, Union
-import numpy as np
-import time
 import pickle
+import time
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
 
 try:
     from ray import __version__ as ray_version
@@ -18,17 +19,17 @@ try:
         from ray.tune.search import Searcher
         from ray.tune.search.optuna import OptunaSearch as GlobalSearch
 except (ImportError, AssertionError):
-    from .suggestion import Searcher
     from .suggestion import OptunaSearch as GlobalSearch
-from ..trial import unflatten_dict, flatten_dict
-from .. import INCUMBENT_RESULT
-from .search_thread import SearchThread
-from .flow2 import FLOW2
-from ..space import add_cost_to_space, indexof, normalize, define_by_run_func
-from ..result import TIME_TOTAL_S
-
+    from .suggestion import Searcher
 import logging
 
+from .. import INCUMBENT_RESULT
+from ..result import TIME_TOTAL_S
+from ..space import add_cost_to_space, define_by_run_func, indexof, normalize
+from ..trial import flatten_dict, unflatten_dict
+from .flow2 import FLOW2
+from .search_thread import SearchThread
+
 SEARCH_THREAD_EPS = 1.0
 PENALTY = 1e10  # penalty term for constraints
 logger = logging.getLogger(__name__)
@@ -931,27 +932,27 @@ try:
 
     assert ray_version >= "1.10.0"
     from ray.tune import (
-        uniform,
-        quniform,
         choice,
-        randint,
-        qrandint,
-        randn,
-        qrandn,
         loguniform,
         qloguniform,
+        qrandint,
+        qrandn,
+        quniform,
+        randint,
+        randn,
+        uniform,
     )
 except (ImportError, AssertionError):
     from ..sample import (
-        uniform,
-        quniform,
         choice,
-        randint,
-        qrandint,
-        randn,
-        qrandn,
         loguniform,
         qloguniform,
+        qrandint,
+        qrandn,
+        quniform,
+        randint,
+        randn,
+        uniform,
     )
 
 try:
@@ -978,7 +979,7 @@ class BlendSearchTuner(BlendSearch, NNITuner):
         result = {
             "config": parameters,
             self._metric: extract_scalar_reward(value),
-            self.cost_attr: 1 if isinstance(value, float) else value.get(self.cost_attr, value.get("sequence", 1))
+            self.cost_attr: 1 if isinstance(value, float) else value.get(self.cost_attr, value.get("sequence", 1)),
             # if nni does not report training cost,
             # using sequence as an approximation.
             # if no sequence, using a constant 1

flaml/tune/searcher/cfo_cat.py

@@ -2,8 +2,8 @@
 #  * Copyright (c) Microsoft Corporation. All rights reserved.
 #  * Licensed under the MIT License. See LICENSE file in the
 #  * project root for license information.
-from .flow2 import FLOW2
 from .blendsearch import CFO
+from .flow2 import FLOW2
 
 
 class FLOW2Cat(FLOW2):

flaml/tune/searcher/flow2.py

@@ -2,31 +2,34 @@
 #  * Copyright (c) Microsoft Corporation. All rights reserved.
 #  * Licensed under the MIT License. See LICENSE file in the
 #  * project root for license information.
-from typing import Dict, Optional, Tuple
-import numpy as np
 import logging
 from collections import defaultdict
+from typing import Dict, Optional, Tuple
+
+import numpy as np
 
 try:
     from ray import __version__ as ray_version
 
     assert ray_version >= "1.0.0"
     if ray_version.startswith("1."):
-        from ray.tune.suggest import Searcher
         from ray.tune import sample
+        from ray.tune.suggest import Searcher
     else:
         from ray.tune.search import Searcher, sample
     from ray.tune.utils.util import flatten_dict, unflatten_dict
 except (ImportError, AssertionError):
-    from .suggestion import Searcher
     from flaml.tune import sample
+
     from ..trial import flatten_dict, unflatten_dict
+    from .suggestion import Searcher
 from flaml.config import SAMPLE_MULTIPLY_FACTOR
+
 from ..space import (
     complete_config,
     denormalize,
-    normalize,
     generate_variants_compatible,
+    normalize,
 )
 
 logger = logging.getLogger(__name__)
@@ -106,7 +109,7 @@ class FLOW2(Searcher):
         else:
             mode = "min"
 
-        super(FLOW2, self).__init__(metric=metric, mode=mode)
+        super().__init__(metric=metric, mode=mode)
         # internally minimizes, so "max" => -1
         if mode == "max":
             self.metric_op = -1.0
@@ -135,7 +138,7 @@ class FLOW2(Searcher):
         self.max_resource = max_resource
         self._resource = None
         self._f_best = None  # only use for lexico_comapre. It represent the best value achieved by lexico_flow.
-        self._step_lb = np.Inf
+        self._step_lb = np.inf
         self._histories = None  # only use for lexico_comapre. It records the result of historical configurations.
         if space is not None:
             self._init_search()
@@ -347,7 +350,7 @@ class FLOW2(Searcher):
             else:
                 assert (
                     self.lexico_objectives["tolerances"][k_metric][-1] == "%"
-                ), "String tolerance of {} should use %% as the suffix".format(k_metric)
+                ), f"String tolerance of {k_metric} should use %% as the suffix"
                 tolerance_bound = self._f_best[k_metric] * (
                     1 + 0.01 * float(self.lexico_objectives["tolerances"][k_metric].replace("%", ""))
                 )
@@ -382,7 +385,7 @@ class FLOW2(Searcher):
                 else:
                     assert (
                         self.lexico_objectives["tolerances"][k_metric][-1] == "%"
-                    ), "String tolerance of {} should use %% as the suffix".format(k_metric)
+                    ), f"String tolerance of {k_metric} should use %% as the suffix"
                     tolerance_bound = self._f_best[k_metric] * (
                         1 + 0.01 * float(self.lexico_objectives["tolerances"][k_metric].replace("%", ""))
                     )

flaml/tune/searcher/online_searcher.py

@@ -1,9 +1,11 @@
-import numpy as np
-import logging
 import itertools
-from typing import Dict, Optional, List
-from flaml.tune import Categorical, Float, PolynomialExpansionSet, Trial
+import logging
+from typing import Dict, List, Optional
+
+import numpy as np
+
 from flaml.onlineml import VowpalWabbitTrial
+from flaml.tune import Categorical, Float, PolynomialExpansionSet, Trial
 from flaml.tune.searcher import CFO
 
 logger = logging.getLogger(__name__)
@@ -64,7 +66,7 @@ class ChampionFrontierSearcher(BaseSearcher):
     POLY_EXPANSION_ADDITION_NUM = 1
     # the order of polynomial expansions to add based on the given seed interactions
     EXPANSION_ORDER = 2
-    # the number of new challengers with new numerical hyperparamter configs
+    # the number of new challengers with new numerical hyperparameter configs
     NUMERICAL_NUM = 2
 
     # In order to use CFO, a loss name and loss values of configs are need
@@ -78,7 +80,7 @@ class ChampionFrontierSearcher(BaseSearcher):
     CFO_SEARCHER_METRIC_NAME = "pseudo_loss"
     CFO_SEARCHER_LARGE_LOSS = 1e6
 
-    # the random seed used in generating numerical hyperparamter configs (when CFO is not used)
+    # the random seed used in generating numerical hyperparameter configs (when CFO is not used)
     NUM_RANDOM_SEED = 111
 
     CHAMPION_TRIAL_NAME = "champion_trial"
@@ -317,7 +319,7 @@ class ChampionFrontierSearcher(BaseSearcher):
         candidate_configs = [set(seed_interactions) | set(item) for item in space]
         final_candidate_configs = []
         for c in candidate_configs:
-            new_c = set([e for e in c if len(e) > 1])
+            new_c = {e for e in c if len(e) > 1}
             final_candidate_configs.append(new_c)
         return final_candidate_configs
 

flaml/tune/searcher/search_thread.py

@@ -3,6 +3,7 @@
 #  * Licensed under the MIT License. See LICENSE file in the
 #  * project root for license information.
 from typing import Dict, Optional
+
 import numpy as np
 
 try:
@@ -15,11 +16,12 @@ try:
         from ray.tune.search import Searcher
 except (ImportError, AssertionError):
     from .suggestion import Searcher
-from .flow2 import FLOW2
-from ..space import add_cost_to_space, unflatten_hierarchical
-from ..result import TIME_TOTAL_S
 import logging
 
+from ..result import TIME_TOTAL_S
+from ..space import add_cost_to_space, unflatten_hierarchical
+from .flow2 import FLOW2
+
 logger = logging.getLogger(__name__)
 
 

flaml/tune/searcher/suggestion.py

@@ -15,15 +15,17 @@
 # This source file is adapted here because ray does not fully support Windows.
 
 # Copyright (c) Microsoft Corporation.
-import time
-import functools
-import warnings
 import copy
-import numpy as np
+import functools
 import logging
-from typing import Any, Dict, Optional, Union, List, Tuple, Callable
 import pickle
-from .variant_generator import parse_spec_vars
+import time
+import warnings
+from collections import defaultdict
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+
 from ..sample import (
     Categorical,
     Domain,
@@ -34,7 +36,7 @@ from ..sample import (
     Uniform,
 )
 from ..trial import flatten_dict, unflatten_dict
-from collections import defaultdict
+from .variant_generator import parse_spec_vars
 
 logger = logging.getLogger(__name__)
 
@@ -183,13 +185,13 @@ class ConcurrencyLimiter(Searcher):
     """
 
     def __init__(self, searcher: Searcher, max_concurrent: int, batch: bool = False):
-        assert type(max_concurrent) is int and max_concurrent > 0
+        assert isinstance(max_concurrent, int) and max_concurrent > 0
         self.searcher = searcher
         self.max_concurrent = max_concurrent
         self.batch = batch
         self.live_trials = set()
         self.cached_results = {}
-        super(ConcurrencyLimiter, self).__init__(metric=self.searcher.metric, mode=self.searcher.mode)
+        super().__init__(metric=self.searcher.metric, mode=self.searcher.mode)
 
     def suggest(self, trial_id: str) -> Optional[Dict]:
         assert trial_id not in self.live_trials, f"Trial ID {trial_id} must be unique: already found in set."
@@ -252,8 +254,8 @@ try:
     import optuna as ot
     from optuna.distributions import BaseDistribution as OptunaDistribution
     from optuna.samplers import BaseSampler
-    from optuna.trial import TrialState as OptunaTrialState
     from optuna.trial import Trial as OptunaTrial
+    from optuna.trial import TrialState as OptunaTrialState
 except ImportError:
     ot = None
     OptunaDistribution = None
@@ -283,25 +285,21 @@ def validate_warmstart(
     """
     if points_to_evaluate:
         if not isinstance(points_to_evaluate, list):
-            raise TypeError("points_to_evaluate expected to be a list, got {}.".format(type(points_to_evaluate)))
+            raise TypeError(f"points_to_evaluate expected to be a list, got {type(points_to_evaluate)}.")
         for point in points_to_evaluate:
             if not isinstance(point, (dict, list)):
                 raise TypeError(f"points_to_evaluate expected to include list or dict, " f"got {point}.")
 
             if validate_point_name_lengths and (not len(point) == len(parameter_names)):
-                raise ValueError(
-                    "Dim of point {}".format(point)
-                    + " and parameter_names {}".format(parameter_names)
-                    + " do not match."
-                )
+                raise ValueError(f"Dim of point {point}" + f" and parameter_names {parameter_names}" + " do not match.")
 
     if points_to_evaluate and evaluated_rewards:
         if not isinstance(evaluated_rewards, list):
-            raise TypeError("evaluated_rewards expected to be a list, got {}.".format(type(evaluated_rewards)))
+            raise TypeError(f"evaluated_rewards expected to be a list, got {type(evaluated_rewards)}.")
         if not len(evaluated_rewards) == len(points_to_evaluate):
             raise ValueError(
-                "Dim of evaluated_rewards {}".format(evaluated_rewards)
-                + " and points_to_evaluate {}".format(points_to_evaluate)
+                f"Dim of evaluated_rewards {evaluated_rewards}"
+                + f" and points_to_evaluate {points_to_evaluate}"
                 + " do not match."
             )
 
@@ -545,7 +543,7 @@ class OptunaSearch(Searcher):
         evaluated_rewards: Optional[List] = None,
     ):
         assert ot is not None, "Optuna must be installed! Run `pip install optuna`."
-        super(OptunaSearch, self).__init__(metric=metric, mode=mode)
+        super().__init__(metric=metric, mode=mode)
 
         if isinstance(space, dict) and space:
             resolved_vars, domain_vars, grid_vars = parse_spec_vars(space)
@@ -559,7 +557,15 @@ class OptunaSearch(Searcher):
         self._space = space
 
         self._points_to_evaluate = points_to_evaluate or []
-        self._evaluated_rewards = evaluated_rewards
+        # rewards should be a list of floats, not a dict
+        # After Optuna > 3.5.0, there is a check for NaN in the list "any(math.isnan(x) for x in self._values)"
+        # which will raise an error when encountering a dict
+        if evaluated_rewards is not None:
+            self._evaluated_rewards = [
+                list(item.values())[0] if isinstance(item, dict) else item for item in evaluated_rewards
+            ]
+        else:
+            self._evaluated_rewards = evaluated_rewards
 
         self._study_name = "optuna"  # Fixed study name for in-memory storage
 
@@ -871,9 +877,9 @@ class OptunaSearch(Searcher):
 
             elif isinstance(domain, Integer):
                 if isinstance(sampler, LogUniform):
-                    return ot.distributions.IntLogUniformDistribution(
-                        domain.lower, domain.upper - 1, step=quantize or 1
-                    )
+                    # ``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):
                     # Upper bound should be inclusive for quantization and
                     # exclusive otherwise

flaml/tune/searcher/variant_generator.py

@@ -17,9 +17,11 @@
 # Copyright (c) Microsoft Corporation.
 import copy
 import logging
-from typing import Any, Dict, Generator, List, Tuple
-import numpy
 import random
+from typing import Any, Dict, Generator, List, Tuple
+
+import numpy
+
 from ..sample import Categorical, Domain, RandomState
 
 try:
@@ -250,7 +252,7 @@ def _try_resolve(v) -> Tuple[bool, Any]:
         # Grid search values
         grid_values = v["grid_search"]
         if not isinstance(grid_values, list):
-            raise TuneError("Grid search expected list of values, got: {}".format(grid_values))
+            raise TuneError(f"Grid search expected list of values, got: {grid_values}")
         return False, Categorical(grid_values).grid()
     return True, v
 
@@ -300,13 +302,13 @@ def has_unresolved_values(spec: Dict) -> bool:
 
 class _UnresolvedAccessGuard(dict):
     def __init__(self, *args, **kwds):
-        super(_UnresolvedAccessGuard, self).__init__(*args, **kwds)
+        super().__init__(*args, **kwds)
         self.__dict__ = self
 
     def __getattribute__(self, item):
         value = dict.__getattribute__(self, item)
         if not _is_resolved(value):
-            raise RecursiveDependencyError("`{}` recursively depends on {}".format(item, value))
+            raise RecursiveDependencyError(f"`{item}` recursively depends on {value}")
         elif isinstance(value, dict):
             return _UnresolvedAccessGuard(value)
         else:

flaml/tune/space.py

@@ -11,9 +11,10 @@ try:
 except (ImportError, AssertionError):
     from . import sample
     from .searcher.variant_generator import generate_variants
-from typing import Dict, Optional, Any, Tuple, Generator, List, Union
-import numpy as np
 import logging
+from typing import Any, Dict, Generator, List, Optional, Tuple, Union
+
+import numpy as np
 
 logger = logging.getLogger(__name__)
 
@@ -489,7 +490,7 @@ def complete_config(
             elif domain.bounded:
                 up, low, gauss_std = 1, 0, 1.0
             else:
-                up, low, gauss_std = np.Inf, -np.Inf, 1.0
+                up, low, gauss_std = np.inf, -np.inf, 1.0
             if domain.bounded:
                 if isinstance(up, list):
                     up[-1] = min(up[-1], 1)

flaml/tune/spark/__init__.py

@@ -1,8 +1,8 @@
 from flaml.tune.spark.utils import (
+    broadcast_code,
     check_spark,
     get_n_cpus,
     with_parameters,
-    broadcast_code,
 )
 
 __all__ = ["check_spark", "get_n_cpus", "with_parameters", "broadcast_code"]

flaml/tune/spark/utils.py

@@ -5,7 +5,6 @@ import threading
 import time
 from functools import lru_cache, partial
 
-
 logger = logging.getLogger(__name__)
 logger_formatter = logging.Formatter(
     "[%(name)s: %(asctime)s] {%(lineno)d} %(levelname)s - %(message)s", "%m-%d %H:%M:%S"
@@ -13,10 +12,10 @@ logger_formatter = logging.Formatter(
 logger.propagate = False
 os.environ["PYARROW_IGNORE_TIMEZONE"] = "1"
 try:
+    import py4j
     import pyspark
     from pyspark.sql import SparkSession
     from pyspark.util import VersionUtils
-    import py4j
 except ImportError:
     _have_spark = False
     py4j = None
@@ -163,6 +162,10 @@ def broadcast_code(custom_code="", file_name="mylearner"):
     assert isinstance(MyLargeLGBM(), LGBMEstimator)
     ```
     """
+    # Check if Spark is available
+    spark_available, _ = check_spark()
+
+    # Write to local driver file system
     flaml_path = os.path.dirname(os.path.abspath(__file__))
     custom_code = textwrap.dedent(custom_code)
     custom_path = os.path.join(flaml_path, file_name + ".py")
@@ -170,6 +173,24 @@ def broadcast_code(custom_code="", file_name="mylearner"):
     with open(custom_path, "w") as f:
         f.write(custom_code)
 
+    # If using Spark, broadcast the code content to executors
+    if spark_available:
+        spark = SparkSession.builder.getOrCreate()
+        bc_code = spark.sparkContext.broadcast(custom_code)
+
+        # Execute a job to ensure the code is distributed to all executors
+        def _write_code(bc):
+            code = bc.value
+            import os
+
+            module_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), file_name + ".py")
+            os.makedirs(os.path.dirname(module_path), exist_ok=True)
+            with open(module_path, "w") as f:
+                f.write(code)
+            return True
+
+        spark.sparkContext.parallelize(range(1)).map(lambda _: _write_code(bc_code)).collect()
+
     return custom_path
 
 
@@ -286,6 +307,7 @@ class PySparkOvertimeMonitor:
     def __exit__(self, exc_type, exc_value, exc_traceback):
         """Exit the context manager.
         This will wait for the monitor thread to nicely exit."""
+        logger.debug(f"monitor exited: {exc_type}, {exc_value}, {exc_traceback}")
         if self._force_cancel and _have_spark:
             self._finished_flag = True
             self._monitor_daemon.join()
@@ -296,6 +318,11 @@ class PySparkOvertimeMonitor:
             if not exc_type:
                 return True
             elif exc_type == py4j.protocol.Py4JJavaError:
+                logger.debug("Py4JJavaError Exception: %s", exc_value)
+                return True
+            elif exc_type == TypeError:
+                # When force cancel, joblib>1.2.0 will raise joblib.externals.loky.process_executor._ExceptionWithTraceback
+                logger.debug("TypeError Exception: %s", exc_value)
                 return True
             else:
                 return False

flaml/tune/trial.py

@@ -15,10 +15,10 @@
 # This source file is adapted here because ray does not fully support Windows.
 
 # Copyright (c) Microsoft Corporation.
-import uuid
 import time
-from numbers import Number
+import uuid
 from collections import deque
+from numbers import Number
 
 
 def flatten_dict(dt, delimiter="/", prevent_delimiter=False):
@@ -110,7 +110,7 @@ class Trial:
                     }
                     self.metric_n_steps[metric] = {}
                     for n in self.n_steps:
-                        key = "last-{:d}-avg".format(n)
+                        key = f"last-{n:d}-avg"
                         self.metric_analysis[metric][key] = value
                         # Store n as string for correct restore.
                         self.metric_n_steps[metric][str(n)] = deque([value], maxlen=n)
@@ -124,7 +124,7 @@ class Trial:
                     self.metric_analysis[metric]["last"] = value
 
                     for n in self.n_steps:
-                        key = "last-{:d}-avg".format(n)
+                        key = f"last-{n:d}-avg"
                         self.metric_n_steps[metric][str(n)].append(value)
                         self.metric_analysis[metric][key] = sum(self.metric_n_steps[metric][str(n)]) / len(
                             self.metric_n_steps[metric][str(n)]

flaml/tune/trial_runner.py

@@ -2,6 +2,7 @@
 #  * Copyright (c) Microsoft Corporation. All rights reserved.
 #  * Licensed under the MIT License. See LICENSE file in the
 #  * project root for license information.
+import logging
 from typing import Optional
 
 # try:
@@ -10,7 +11,6 @@ from typing import Optional
 #     from ray.tune.trial import Trial
 # except (ImportError, AssertionError):
 from .trial import Trial
-import logging
 
 logger = logging.getLogger(__name__)
 

flaml/tune/tune.py

@@ -2,13 +2,14 @@
 #  * Copyright (c) FLAML authors. All rights reserved.
 #  * Licensed under the MIT License. See LICENSE file in the
 #  * project root for license information.
-from typing import Optional, Union, List, Callable, Tuple, Dict
-import numpy as np
 import datetime
-import time
 import os
 import sys
+import time
 from collections import defaultdict
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
 
 try:
     from ray import __version__ as ray_version
@@ -20,14 +21,26 @@ except (ImportError, AssertionError):
     from .analysis import ExperimentAnalysis as EA
 else:
     ray_available = True
-
-from .trial import Trial
-from .result import DEFAULT_METRIC
 import logging
+
 from flaml.tune.spark.utils import PySparkOvertimeMonitor, check_spark
 
-logger = logging.getLogger(__name__)
-logger.propagate = False
+from .logger import logger, logger_formatter
+from .result import DEFAULT_METRIC
+from .trial import Trial
+
+try:
+    import mlflow
+except ImportError:
+    mlflow = None
+try:
+    from flaml.fabric.mlflow import MLflowIntegration, is_autolog_enabled
+
+    internal_mlflow = True
+except ImportError:
+    internal_mlflow = False
+
+
 _use_ray = True
 _runner = None
 _verbose = 0
@@ -41,6 +54,7 @@ class ExperimentAnalysis(EA):
     """Class for storing the experiment results."""
 
     def __init__(self, trials, metric, mode, lexico_objectives=None):
+        self.best_run_id = None
         try:
             super().__init__(self, None, trials, metric, mode)
             self.lexico_objectives = lexico_objectives
@@ -92,10 +106,12 @@ class ExperimentAnalysis(EA):
             feasible_index_filter = np.where(
                 feasible_value
                 <= max(
-                    f_best[k_metric] + self.lexico_objectives["tolerances"][k_metric]
-                    if not isinstance(self.lexico_objectives["tolerances"][k_metric], str)
-                    else f_best[k_metric]
-                    * (1 + 0.01 * float(self.lexico_objectives["tolerances"][k_metric].replace("%", ""))),
+                    (
+                        f_best[k_metric] + self.lexico_objectives["tolerances"][k_metric]
+                        if not isinstance(self.lexico_objectives["tolerances"][k_metric], str)
+                        else f_best[k_metric]
+                        * (1 + 0.01 * float(self.lexico_objectives["tolerances"][k_metric].replace("%", "")))
+                    ),
                     k_target,
                 )
             )[0]
@@ -123,6 +139,16 @@ class ExperimentAnalysis(EA):
         else:
             return self.best_trial.last_result
 
+    @property
+    def best_iteration(self) -> List[str]:
+        """Help better navigate"""
+        best_trial = self.best_trial
+        best_trial_id = best_trial.trial_id
+        for i, trial in enumerate(self.trials):
+            if trial.trial_id == best_trial_id:
+                return i
+        return None
+
 
 def report(_metric=None, **kwargs):
     """A function called by the HPO application to report final or intermediate
@@ -169,9 +195,16 @@ def report(_metric=None, **kwargs):
     global _training_iteration
     if _use_ray:
         try:
-            from ray import tune
+            from ray import __version__ as ray_version
 
-            return tune.report(_metric, **kwargs)
+            if ray_version.startswith("1."):
+                from ray import tune
+
+                return tune.report(_metric, **kwargs)
+            else:  # ray>=2
+                from ray.air import session
+
+                return session.report(metrics={"metric": _metric, **kwargs})
         except ImportError:
             # calling tune.report() outside tune.run()
             return
@@ -229,6 +262,11 @@ def run(
     lexico_objectives: Optional[dict] = None,
     force_cancel: Optional[bool] = False,
     n_concurrent_trials: Optional[int] = 0,
+    mlflow_exp_name: Optional[str] = None,
+    automl_info: Optional[Tuple[float]] = None,
+    extra_tag: Optional[dict] = None,
+    cost_attr: Optional[str] = "auto",
+    cost_budget: Optional[float] = None,
     **ray_args,
 ):
     """The function-based way of performing HPO.
@@ -419,6 +457,10 @@ def run(
     }
     ```
         force_cancel: boolean, default=False | Whether to forcely cancel the PySpark job if overtime.
+        mlflow_exp_name: str, default=None | The name of the mlflow experiment. This should be specified if
+            enable mlflow autologging on Spark. Otherwise it will log all the results into the experiment of the
+            same name as the basename of main entry file.
+        automl_info: tuple, default=None | The information of the automl run. It should be a tuple of (mlflow_log_latency,).
         n_concurrent_trials: int, default=0 | The number of concurrent trials when perform hyperparameter
             tuning with Spark. Only valid when use_spark=True and spark is required:
             `pip install flaml[spark]`. Please check
@@ -426,6 +468,13 @@ def run(
             for more details about installing Spark. When tune.run() is called from AutoML, it will be
             overwritten by the value of `n_concurrent_trials` in AutoML. When <= 0, the concurrent trials
             will be set to the number of executors.
+        extra_tag: dict, default=None | Extra tags to be added to the mlflow runs created by autologging.
+        cost_attr: None or str to specify the attribute to evaluate the cost of different trials.
+            Default is "auto", which means that we will automatically choose the cost attribute to use (depending
+            on the nature of the resource budget). When cost_attr is set to None, cost differences between different trials will be omitted
+            in our search algorithm. When cost_attr is set to a str different from "auto" and "time_total_s",
+            this cost_attr must be available in the result dict of the trial.
+        cost_budget: A float of the cost budget. Only valid when cost_attr is a str different from "auto" and "time_total_s".
         **ray_args: keyword arguments to pass to ray.tune.run().
             Only valid when use_ray=True.
     """
@@ -433,10 +482,12 @@ def run(
     global _verbose
     global _running_trial
     global _training_iteration
+    global internal_mlflow
     old_use_ray = _use_ray
     old_verbose = _verbose
     old_running_trial = _running_trial
     old_training_iteration = _training_iteration
+
     if log_file_name:
         dir_name = os.path.dirname(log_file_name)
         if dir_name:
@@ -468,10 +519,6 @@ def run(
             elif not logger.hasHandlers():
                 # Add the console handler.
                 _ch = logging.StreamHandler(stream=sys.stdout)
-                logger_formatter = logging.Formatter(
-                    "[%(name)s: %(asctime)s] {%(lineno)d} %(levelname)s - %(message)s",
-                    "%m-%d %H:%M:%S",
-                )
                 _ch.setFormatter(logger_formatter)
                 logger.addHandler(_ch)
             if verbose <= 2:
@@ -481,7 +528,14 @@ def run(
         else:
             logger.setLevel(logging.CRITICAL)
 
-    from .searcher.blendsearch import BlendSearch, CFO, RandomSearch
+    if internal_mlflow and not automl_info and (mlflow.active_run() or is_autolog_enabled()):
+        mlflow_integration = MLflowIntegration("tune", mlflow_exp_name, extra_tag)
+        evaluation_function = mlflow_integration.wrap_evaluation_function(evaluation_function)
+        _internal_mlflow = not automl_info  # True if mlflow_integration will be used for logging
+    else:
+        _internal_mlflow = False
+
+    from .searcher.blendsearch import CFO, BlendSearch, RandomSearch
 
     if lexico_objectives is not None:
         if "modes" not in lexico_objectives.keys():
@@ -526,7 +580,7 @@ def run(
                     import optuna as _
 
                     SearchAlgorithm = BlendSearch
-                    logger.info("Using search algorithm {}.".format(SearchAlgorithm.__name__))
+                    logger.info(f"Using search algorithm {SearchAlgorithm.__name__}.")
                 except ImportError:
                     if search_alg == "BlendSearch":
                         raise ValueError("To use BlendSearch, run: pip install flaml[blendsearch]")
@@ -535,7 +589,7 @@ def run(
                         logger.warning("Using CFO for search. To use BlendSearch, run: pip install flaml[blendsearch]")
             else:
                 SearchAlgorithm = locals()[search_alg]
-                logger.info("Using search algorithm {}.".format(SearchAlgorithm.__name__))
+                logger.info(f"Using search algorithm {SearchAlgorithm.__name__}.")
             metric = metric or DEFAULT_METRIC
         search_alg = SearchAlgorithm(
             metric=metric,
@@ -555,6 +609,8 @@ def run(
             metric_constraints=metric_constraints,
             use_incumbent_result_in_evaluation=use_incumbent_result_in_evaluation,
             lexico_objectives=lexico_objectives,
+            cost_attr=cost_attr,
+            cost_budget=cost_budget,
         )
     else:
         if metric is None or mode is None:
@@ -650,12 +706,13 @@ def run(
         if not spark_available:
             raise spark_error_msg
         try:
-            from pyspark.sql import SparkSession
             from joblib import Parallel, delayed, parallel_backend
             from joblibspark import register_spark
+            from pyspark.sql import SparkSession
         except ImportError as e:
             raise ImportError(f"{e}. Try pip install flaml[spark] or set use_spark=False.")
         from flaml.tune.searcher.suggestion import ConcurrencyLimiter
+
         from .trial_runner import SparkTrialRunner
 
         register_spark()
@@ -689,10 +746,16 @@ def run(
             max_concurrent = max(1, search_alg.max_concurrent)
         else:
             max_concurrent = max(1, max_spark_parallelism)
+        passed_in_n_concurrent_trials = max(n_concurrent_trials, max_concurrent)
         n_concurrent_trials = min(
             n_concurrent_trials if n_concurrent_trials > 0 else num_executors,
             max_concurrent,
         )
+        if n_concurrent_trials < passed_in_n_concurrent_trials:
+            logger.warning(
+                f"The actual concurrent trials is {n_concurrent_trials}. You can set the environment "
+                f"variable `FLAML_MAX_CONCURRENT` to '{passed_in_n_concurrent_trials}' to override the detected num of executors."
+            )
         with parallel_backend("spark"):
             with Parallel(n_jobs=n_concurrent_trials, verbose=max(0, (verbose - 1) * 50)) as parallel:
                 try:
@@ -707,11 +770,15 @@ def run(
                         time_budget_s = np.inf
                     num_failures = 0
                     upperbound_num_failures = (len(evaluated_rewards) if evaluated_rewards else 0) + max_failure
+                    logger.debug(f"automl_info: {automl_info}")
                     while (
                         time.time() - time_start < time_budget_s
                         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:
+                            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:
                             # suggest trials for spark
                             trial_next = _runner.step()
@@ -744,6 +811,9 @@ def run(
                             trial_to_run = trials_to_run[0]
                             _runner.running_trial = trial_to_run
                             if result is not None:
+                                if _internal_mlflow:
+                                    mlflow_integration.record_trial(result, trial_to_run, metric)
+
                                 if isinstance(result, dict):
                                     if result:
                                         logger.info(f"Brief result: {result}")
@@ -752,7 +822,7 @@ def run(
                                         # When the result returned is an empty dict, set the trial status to error
                                         trial_to_run.set_status(Trial.ERROR)
                                 else:
-                                    logger.info("Brief result: {}".format({metric: result}))
+                                    logger.info("Brief result: {metric: result}")
                                     report(_metric=result)
                             _runner.stop_trial(trial_to_run)
                         num_failures = 0
@@ -762,6 +832,20 @@ def run(
                         mode=mode,
                         lexico_objectives=lexico_objectives,
                     )
+                    analysis.search_space = config
+
+                    if _internal_mlflow:
+                        mlflow_integration.log_tune(analysis, metric)
+                        # try:
+                        #     _best_config = analysis.best_config
+                        # except Exception:
+                        #     _best_config = None
+                        # if _best_config:
+                        #     parallel(
+                        #         delayed(mlflow_integration.retrain)(evaluation_function, analysis.best_config)
+                        #         for dummy in [0]
+                        #     )
+
                     return analysis
                 finally:
                     # recover the global variables in case of nested run
@@ -773,6 +857,8 @@ def run(
                         _runner = old_runner
                         logger.handlers = old_handlers
                         logger.setLevel(old_level)
+                    if _internal_mlflow:
+                        mlflow_integration.adopt_children()
 
     # simple sequential run without using tune.run() from ray
     time_start = time.time()
@@ -806,7 +892,11 @@ def run(
                 result = None
                 with PySparkOvertimeMonitor(time_start, time_budget_s, force_cancel):
                     result = evaluation_function(trial_to_run.config)
+                logger.debug(f"result in tune: {trial_to_run}, {result}")
                 if result is not None:
+                    if _internal_mlflow:
+                        mlflow_integration.record_trial(result, trial_to_run, metric)
+
                     if isinstance(result, dict):
                         if result:
                             report(**result)
@@ -832,6 +922,19 @@ def run(
             mode=mode,
             lexico_objectives=lexico_objectives,
         )
+        analysis.search_space = config
+        if _internal_mlflow:
+            mlflow_integration.log_tune(analysis, metric)
+            if analysis.best_run_id is not None:
+                logger.info(f"Best MLflow run name: {analysis.best_run_name}")
+                logger.info(f"Best MLflow run id: {analysis.best_run_id}")
+            # try:
+            #     _best_config = analysis.best_config
+            # except Exception:
+            #     _best_config = None
+            # if _best_config:
+            #     mlflow_integration.retrain(evaluation_function, analysis.best_config)
+
         return analysis
     finally:
         # recover the global variables in case of nested run
@@ -843,6 +946,8 @@ def run(
             _runner = old_runner
             logger.handlers = old_handlers
             logger.setLevel(old_level)
+        if _internal_mlflow:
+            mlflow_integration.adopt_children()
 
 
 class Tuner: