Repository: openml/automlbenchmark (language: python)

Issue:
Not enough precision in predict_duration in results.csv
We are [rounding](https://github.com/openml/automlbenchmark/blob/aee089164e8236e581672732c8bdc5ee6393d0ac/amlb/results.py#L172) to the first decimal for `predict_duration` so a lot of the results are getting rounded to 0.

Here, `"{:.1f}".format`, we need .3g instead of .1f.
fix for mlr3automl
makes PR #424 obsolete


Discussion / hints:

Fix patch:
diff --git a/.dockerignore b/.dockerignore
index db7158276..7602493a2 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -8,13 +8,10 @@ results
 **/models
 **/predictions
 **/scores
-**/.setup
-**/.installed
-**/.setup_env
 **/__pycache__
-**/lib
-**/setup
 **/venv
+frameworks/*/.setup
+frameworks/*/lib
 
 # editors/IDEs
 .idea
diff --git a/.github/scripts/tag-set.sh b/.github/scripts/tag-set.sh
new file mode 100755
index 000000000..d2803a5a2
--- /dev/null
+++ b/.github/scripts/tag-set.sh
@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+TAG="$1"
+PUSH=${2:-"no-push"}
+
+function extract_version()
+{
+  local version=$(echo "$1" | sed -E 's/^(refs\/(heads|tags)\/)?(.*)/\3/')
+  echo "$version"
+}
+
+NEW_TAG=$(extract_version $TAG)
+
+git config user.name github-actions
+git config user.email github-actions@github.com
+git tag -f $NEW_TAG
+
+HIGHEST_vTAG=$(extract_version $(git ls-remote --tags --sort=-v:refname origin | grep -P "/v\d.*" | head -n 1 | awk '{print $2}'))
+#git fetch --all --tags
+#LATEST_vTAG=$(git describe --contains `git rev-list --tags="v*" --max-count=1`)
+#HIGHEST_vTAG=$(git tag --list 'v*' --sort=-v:refname | grep -P "v\d.*" | head -n 1)
+
+echo "New tag: '$NEW_TAG', Highest version tag: '$HIGHEST_vTAG'"
+if [ "$NEW_TAG" == "$HIGHEST_vTAG" ]; then
+  echo "Additionally setting 'stable' tag"
+  git tag -f stable
+fi
+if [[ "$PUSH" == "push" ]]; then
+  git push --tags -f
+fi
diff --git a/.github/scripts/version-set.sh b/.github/scripts/version-set.sh
new file mode 100755
index 000000000..734f64144
--- /dev/null
+++ b/.github/scripts/version-set.sh
@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+VERSION="$1"
+PUSH=${2:-"no-push"}
+ADD_COMMIT=${3:-"add-commit"}
+
+VERSION_FILE='app/__version__.py'
+VERSION_PLACEHOLDER='__version__ =\s+"(.*)"$'
+DEV_PLACEHOLDER='_dev_version =\s+"(.*)"$'
+
+function get_version()
+{
+  local version=$(grep -P "$2" "$1" | sed -E "s/$2/\1/")
+  echo "$version"
+}
+
+function set_version()
+{
+  sed -i -E "s/$2/__version__ = \"$3\"/g" "$1"
+  echo $(get_version "$1" "$2")
+}
+
+OLD_VERSION=$(get_version "$VERSION_FILE" "$VERSION_PLACEHOLDER")
+DEV_VERSION=$(get_version "$VERSION_FILE" "$DEV_PLACEHOLDER")
+NEW_VERSION=$(echo ${VERSION:-"$DEV_VERSION"} | sed -E 's/^v([0-9].*)/\1/')
+echo "setting version from \"$OLD_VERSION\" to \"$NEW_VERSION\""
+FINAL_VERSION=$(set_version "$VERSION_FILE" "$VERSION_PLACEHOLDER" "$NEW_VERSION")
+echo "version changed to \"$FINAL_VERSION\""
+echo "VERSION=$FINAL_VERSION" >> $GITHUB_ENV
+
+if [[ "$PUSH" == "push" ]]; then
+  if [[ -x "$(command -v git)" ]]; then
+    git config user.name github-actions
+    git config user.email github-actions@github.com
+    if [[ "$ADD_COMMIT" == "add-commit" ]]; then
+      git commit -am "Update version to $FINAL_VERSION"
+      git push
+    else
+      git commit -a --amend --no-edit
+      git push -f
+    fi
+  else
+    echo "Can not push the version changes as git is not available."
+    exit 1
+  fi
+fi
diff --git a/.github/workflows/versioning-reset.yml b/.github/workflows/versioning-reset.yml
new file mode 100644
index 000000000..63fbf6424
--- /dev/null
+++ b/.github/workflows/versioning-reset.yml
@@ -0,0 +1,27 @@
+name: Version reset to dev
+# reset the version if merging a release branch into main
+
+on:
+  push:
+    branches:
+      - main
+      - master
+    paths:
+      - 'app/__version__.py'
+
+  workflow_dispatch:
+
+env:
+  BRANCH: ${{ github.event.release.target_commitish }}
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          ref: ${{ env.BRANCH }}
+          fetch-depth: 2  # we need the previous commit as we amend the last one
+      - name: Reset version to dev
+        run: ./.github/scripts/version-set.sh "" push no-add-commit
diff --git a/.github/workflows/versioning.yml b/.github/workflows/versioning.yml
index b377caadc..fc05eb540 100644
--- a/.github/workflows/versioning.yml
+++ b/.github/workflows/versioning.yml
@@ -3,7 +3,7 @@ name: Version Bump
 on:
   release:
     types:
-      - created
+      - published
 
   workflow_dispatch:
 
@@ -12,30 +12,36 @@ env:
   TAG: ${{ github.event.release.tag_name }}
 
 jobs:
-  deploy:
+  validation:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Release only from stable-v* branch
+        if: ${{ !startsWith(env.BRANCH, 'stable-v') }}
+        uses: actions/github-script@v6
+        with:
+          script: core.setFailed("Releases can only be made from 'stable-v*' branches")
+  rollback:
+    needs: validation
+    if: ${{ failure() }}
     runs-on: ubuntu-latest
-
     steps:
-      - uses: actions/checkout@v2
+      - name: Rollback if invalid release
+        uses: author/action-rollback@stable
+        with:
+          tag: ${{ env.TAG }}
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  versioning:
+    needs: validation
+    if: ${{ success() }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
         with:
           ref: ${{ env.BRANCH }}
       - name: Bump version from tag
-        run: |
-          VERSION=$(echo $TAG | sed 's/^v//')
-          PLACEHOLDER='__version__ = _dev_version'
-          VERSION_FILE='amlb/__version__.py'
-          grep "$PLACEHOLDER" "$VERSION_FILE"
-          sed -i "s/$PLACEHOLDER/__version__ = \"${VERSION}\"/g" "$VERSION_FILE"
-          echo "VERSION=$VERSION" >> $GITHUB_ENV
-        shell: bash
-      - name: Commit version changes
-        run: |
-          git config user.name github-actions
-          git config user.email github-actions@github.com
-          git commit -am "Update version to $VERSION"
+        run: ./.github/scripts/version-set.sh $TAG push
       - name: Update tag
-        run: |
-          git tag $TAG
-          git push
-          git push --tags -f
-
+        run: ./.github/scripts/tag-set.sh $TAG push
+      - name: Restore dev version
+        run: ./.github/scripts/version-set.sh "" push
diff --git a/amlb/__version__.py b/amlb/__version__.py
index d3c8ad7e6..46945d24b 100644
--- a/amlb/__version__.py
+++ b/amlb/__version__.py
@@ -1,3 +1,3 @@
 # major.minor.patch[-label]
 _dev_version = "dev"
-__version__ = _dev_version
+__version__ = "2.0.6"
diff --git a/amlb/benchmark.py b/amlb/benchmark.py
index d38f8ad56..4fbc93a3f 100644
--- a/amlb/benchmark.py
+++ b/amlb/benchmark.py
@@ -72,6 +72,9 @@ def __init__(self, framework_name: str, benchmark_name: str, constraint_name: st
             return
 
         self._forward_params = locals()
+        if Benchmark.data_loader is None:
+            Benchmark.data_loader = DataLoader(rconfig())
+
         fsplits = framework_name.split(':', 1)
         framework_name = fsplits[0]
         tag = fsplits[1] if len(fsplits) > 1 else None
@@ -108,8 +111,6 @@ def setup(self, mode: SetupMode):
         and possibly download them if necessary.
         Delegates specific setup to the framework module
         """
-        Benchmark.data_loader = DataLoader(rconfig())
-
         if mode == SetupMode.skip or mode == SetupMode.auto and self._is_setup_done():
             return
 
diff --git a/amlb/datasets/openml.py b/amlb/datasets/openml.py
index e7d18e3a3..d037252d8 100644
--- a/amlb/datasets/openml.py
+++ b/amlb/datasets/openml.py
@@ -4,6 +4,7 @@
 """
 from abc import abstractmethod
 import copy
+import functools
 import logging
 import os
 import re
@@ -12,6 +13,7 @@
 import arff
 import pandas.api.types as pat
 import openml as oml
+import xmltodict
 
 from ..data import AM, DF, Dataset, DatasetType, Datasplit, Feature
 from ..resources import config as rconfig
@@ -23,6 +25,9 @@
 # hack (only adding a ? to the regexp pattern) to ensure that '?' values remain quoted when we save dataplits in arff format.
 arff._RE_QUOTE_CHARS = re.compile(r'[?"\'\\\s%,\000-\031]', re.UNICODE)
 
+# Fix a bug in openml-python<=0.12.2, see https://github.com/openml/automlbenchmark/issues/350
+xmltodict.parse = functools.partial(xmltodict.parse, strip_whitespace=False)
+
 
 class OpenmlLoader:
 
diff --git a/amlb/job.py b/amlb/job.py
index f89503ca7..01dd24f37 100644
--- a/amlb/job.py
+++ b/amlb/job.py
@@ -252,7 +252,8 @@ def reschedule(self, job, priority=None):
         if self.state not in [State.running]:
             return
         job.reschedule()
-        self.put(job, priority)
+        if job.state is State.rescheduling:
+            self.put(job, priority)
 
     def set_state(self, state: State):
         assert self.is_state_transition_ok(self.state, state), f"Illegal job runner transition from state {self.state} to {state}"
@@ -309,9 +310,11 @@ def _run(self):
             if job is None or self._interrupt.is_set():
                 break
             result = job.start()
-            if job.state is not State.rescheduling:
+            if job.state is State.rescheduling:
+                self.reschedule(job)
+            else:
                 self.results.append(result)
-            job.done()
+                job.done()
             self.stop_if_complete()
 
     def _on_state(self, state: State):
@@ -349,10 +352,12 @@ def worker():
                     if job is None or self._interrupt.is_set():
                         break
                     result = job.start()
-                    if job.state is not State.rescheduling:
+                    if job.state is State.rescheduling:
+                        self.reschedule(job)
+                    else:
                         self.results.append(result)
-                    if self._done_async:
-                        job.done()
+                        if self._done_async:
+                            job.done()
                     self.stop_if_complete()
                 finally:
                     q.task_done()
diff --git a/amlb/results.py b/amlb/results.py
index 2916dbcdf..2f547b4ec 100644
--- a/amlb/results.py
+++ b/amlb/results.py
@@ -169,7 +169,8 @@ def as_printable_data_frame(self, verbosity=3):
         for col in nanable_int_cols:
             df[col] = df[col].astype(np.object).map(int_print).astype(np.str)
         for col in low_precision_float_cols:
-            df[col] = df[col].astype(np.float).map(partial(num_print, "{:.1f}".format)).astype(np.float)
+            float_format = lambda f: ("{:.1g}" if f < 1 else "{:.1f}").format(f)
+            df[col] = df[col].astype(np.float).map(partial(num_print, float_format)).astype(np.float)
         for col in high_precision_float_cols:
             df[col] = df[col].map(partial(num_print, "{:.6g}".format)).astype(np.float)
 
diff --git a/amlb/runners/aws.py b/amlb/runners/aws.py
index 3bddb8a27..4368efde3 100644
--- a/amlb/runners/aws.py
+++ b/amlb/runners/aws.py
@@ -265,7 +265,7 @@ def _job_reschedule(self, job, reason=None, fallback=None):
                 raise JobError(reason)
         else:
             job.ext.wait_min_secs = wait
-        self.job_runner.reschedule(job)
+        job.reschedule()
 
     def _spot_fallback(self, job, reason):
         if 'Spot' in reason and rconfig().aws.ec2.spot.fallback_to_on_demand:
@@ -456,6 +456,10 @@ def log_console():
 
                 if state_code == 16:
                     if inst_desc['meta_info'] is None:
+                        volume_info = [
+                            dict(type=v.volume_type, size_gb=v.size, id=v.id)
+                            for v in instance.volumes.all()
+                        ]
                         meta_info = dict(
                             instance_type=instance.instance_type,
                             launch_time=str(instance.launch_time),
@@ -465,6 +469,7 @@ def log_console():
                             private_ip=instance.private_ip_address,
                             availability_zone=instance.placement['AvailabilityZone'],
                             subnet_id=instance.subnet_id,
+                            volumes=volume_info,
                         )
                         self._update_instance(job.ext.instance_id, meta_info=meta_info)
                         log.info("Running EC2 instance %s: %s", instance.id, meta_info)
@@ -521,7 +526,8 @@ def cpu_monitor():
                 return
             while not interrupt.is_set():
                 try:
-                    hanging_instances = list(filter(self._is_hanging, self.instances.keys()))
+                    active_instances = [iid for iid, info in self.instances.items() if info.instance is not None]
+                    hanging_instances = list(filter(self._is_hanging, active_instances))
                     for inst in hanging_instances:
                         if inst in self.instances:
                             inst_desc = self.instances[inst]
diff --git a/amlb/runners/container.py b/amlb/runners/container.py
index 2cf2f4307..f33f220ab 100644
--- a/amlb/runners/container.py
+++ b/amlb/runners/container.py
@@ -182,7 +182,15 @@ def _build_image(self, cache=True):
                 image = self._container_image_name(dev)
 
         if not image:
-            image = self._container_image_name()
+            tags = rget().git_info.tags
+            version_tags = [t for t in tags if re.match(r"v\d+(\d+.)*", t)]
+            if len(version_tags) > 1:
+                raise InvalidStateError(
+                    "The image can't be built as more than one version tag was found."
+                    f"Found tags: {version_tags}"
+                )
+            version_tag = next(iter(version_tags), None)
+            image = self._container_image_name(version_tag)
         self._run_container_build_command(image, cache)
         return image
 
diff --git a/amlb/utils/process.py b/amlb/utils/process.py
index c32fad3b1..593ed1ff8 100644
--- a/amlb/utils/process.py
+++ b/amlb/utils/process.py
@@ -286,11 +286,17 @@ def on_proc_terminated(proc):
         children.append(parent)
     for proc in children:
         log.warning("Terminating process %s.", proc)
-        proc.terminate()
+        try:
+            proc.terminate()
+        except psutil.NoSuchProcess:
+            pass
     terminated, alive = psutil.wait_procs(children, timeout=timeout, callback=on_proc_terminated)
     for proc in alive:
         log.warning("Killing process %s.", proc)
-        proc.kill()
+        try:
+            proc.kill()
+        except psutil.NoSuchProcess:
+            pass
 
 
 def call_in_subprocess(target, *args, **kwargs):
diff --git a/frameworks/RandomForest/__init__.py b/frameworks/RandomForest/__init__.py
index b024b051c..a25cbeee7 100644
--- a/frameworks/RandomForest/__init__.py
+++ b/frameworks/RandomForest/__init__.py
@@ -1,6 +1,6 @@
 from amlb.benchmark import TaskConfig
 from amlb.data import Dataset
-from amlb.utils import call_script_in_same_dir
+from amlb.utils import call_script_in_same_dir, unsparsify
 
 
 def setup(*args, **kwargs):
@@ -14,6 +14,7 @@ def run(dataset: Dataset, config: TaskConfig):
     encode = config.framework_params.get('_encode', True)
     X_train, X_test = impute_array(dataset.train.X_enc, dataset.test.X_enc) if encode else (dataset.train.X, dataset.test.X)
     y_train, y_test = (dataset.train.y_enc, dataset.test.y_enc) if encode else (dataset.train.y, dataset.test.y)
+    y_train, y_test = unsparsify(y_train, y_test)
     data = dict(
         train=dict(
             X=X_train,
diff --git a/frameworks/RandomForest/exec.py b/frameworks/RandomForest/exec.py
index 57fd8698e..77bdc99ef 100644
--- a/frameworks/RandomForest/exec.py
+++ b/frameworks/RandomForest/exec.py
@@ -1,12 +1,14 @@
 import logging
 import os
 import tempfile as tmp
+from typing import List
 
 os.environ['JOBLIB_TEMP_FOLDER'] = tmp.gettempdir()
 os.environ['OMP_NUM_THREADS'] = '1'
 os.environ['OPENBLAS_NUM_THREADS'] = '1'
 os.environ['MKL_NUM_THREADS'] = '1'
 
+import psutil
 import sklearn
 from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
 
@@ -16,29 +18,69 @@
 log = logging.getLogger(os.path.basename(__file__))
 
 
+def extrapolate_with_worst_case(values: List[float], n: int = 5) -> float:
+    """ Extrapolate the next value for `values`, based on the last `n` samples. """
+    n = min(len(values), n)
+    return values[-1] + max(v_next - v_prev for v_prev, v_next in zip(values[-n:], values[-n + 1:]))
+
+
 def run(dataset, config):
     log.info(f"\n**** Random Forest [sklearn v{sklearn.__version__}] ****\n")
 
     is_classification = config.type == 'classification'
+    this_process = psutil.Process(os.getpid())
 
     encode = config.framework_params.get('_encode', True)
     X_train, X_test = dataset.train.X, dataset.test.X
     y_train, y_test = dataset.train.y, dataset.test.y
 
-    training_params = {k: v for k, v in config.framework_params.items() if not k.startswith('_')}
+    training_params = {
+        k: v for k, v in config.framework_params.items()
+        if not (k.startswith('_') or k == "n_estimators")
+    }
     n_jobs = config.framework_params.get('_n_jobs', config.cores)  # useful to disable multicore, regardless of the dataset config
+    step_size = config.framework_params.get('_step_size', 10)
+    final_forest_size = config.framework_params.get('n_estimators', 2000)
+
+    # Default margins are conservative, because robustness is paramount for a baseline.
+    time_margin = config.framework_params.get('_time_margin', 0.9)
+    memory_margin = config.framework_params.get('_memory_margin', 0.9)
 
     log.info("Running RandomForest with a maximum time of {}s on {} cores.".format(config.max_runtime_seconds, n_jobs))
-    log.warning("We completely ignore the requirement to stay within the time limit.")
     log.warning("We completely ignore the advice to optimize towards metric: {}.".format(config.metric))
 
     estimator = RandomForestClassifier if is_classification else RandomForestRegressor
     rf = estimator(n_jobs=n_jobs,
                    random_state=config.seed,
+                   n_estimators=step_size,
+                   warm_start=True,
                    **training_params)
 
     with Timer() as training:
-        rf.fit(X_train, y_train)
+        training_times = [training.duration]
+        memory_usage = [this_process.memory_info()[0] / (2**20)]
+
+        while True:
+            rf.fit(X_train, y_train)
+
+            training_times.append(training.duration)
+            memory_usage.append(this_process.memory_info()[0] / (2**20))
+            log.info(f"Model trained {len(rf.estimators_):6d} trees in {int(training_times[-1]):6d} seconds using {int(memory_usage[-1]):6d}mb memory.")
+
+            will_run_out_of_memory = extrapolate_with_worst_case(memory_usage) >= config.max_mem_size_mb * memory_margin
+            will_run_out_of_time = extrapolate_with_worst_case(training_times) >= config.max_runtime_seconds * time_margin
+            if rf.n_estimators >= final_forest_size:
+                log.info("Stop training because desired forest size has been reached.")
+                break
+            elif will_run_out_of_time:
+                log.info("Stop training because it expects to exceed its time budget.")
+                break
+            elif will_run_out_of_memory:
+                log.info("Stop training because it expects to exceed its memory budget.")
+                break
+            else:
+                # https://stackoverflow.com/questions/42757892/how-to-use-warm-start/42763502
+                rf.n_estimators += step_size
 
     with Timer() as predict:
         predictions = rf.predict(X_test)
diff --git a/frameworks/RandomForest/setup.sh b/frameworks/RandomForest/setup.sh
index a6821c890..c678536f2 100755
--- a/frameworks/RandomForest/setup.sh
+++ b/frameworks/RandomForest/setup.sh
@@ -4,7 +4,7 @@ VERSION=${1:-"stable"}
 REPO=${2:-"https://github.com/scikit-learn/scikit-learn.git"}
 PKG=${3:-"scikit-learn"}
 if [[ "$VERSION" == "latest" ]]; then
-    VERSION="master"
+    VERSION="main"
 fi
 
 . "${HERE}/../shared/setup.sh" "${HERE}" true
diff --git a/frameworks/TPOT/__init__.py b/frameworks/TPOT/__init__.py
index 959c7072e..9828c6473 100644
--- a/frameworks/TPOT/__init__.py
+++ b/frameworks/TPOT/__init__.py
@@ -25,14 +25,13 @@ def run(dataset: Dataset, config: TaskConfig):
     )
 
     def process_results(results):
-        if results.probabilities is not None and not results.probabilities.shape:  # numpy load always return an array
-            prob_format = results.probabilities.item()
-            if prob_format == "predictions":
-                target_values_enc = dataset.target.label_encoder.transform(dataset.target.values)
-                results.probabilities = Encoder('one-hot', target=False, encoded_type=float).fit(target_values_enc).transform(results.predictions)
-            else:
-                raise ValueError(f"Unknown probabilities format: {prob_format}")
-        return results
+        if isinstance(results.probabilities, str) and results.probabilities == "predictions":
+            target_values_enc = dataset.target.label_encoder.transform(dataset.target.values)
+            results.probabilities = Encoder('one-hot', target=False, encoded_type=float).fit(target_values_enc).transform(results.predictions)
+        is_numpy_like = hasattr(results.probabilities, "shape") and results.probabilities.shape
+        if results.probabilities is None or is_numpy_like:
+            return results
+        raise ValueError(f"Unknown probabilities format: {results.probabilities}")
 
     return run_in_venv(__file__, "exec.py",
                        input_data=data, dataset=dataset, config=config,
diff --git a/frameworks/TunedRandomForest/__init__.py b/frameworks/TunedRandomForest/__init__.py
index a98dfdfac..561678497 100644
--- a/frameworks/TunedRandomForest/__init__.py
+++ b/frameworks/TunedRandomForest/__init__.py
@@ -1,6 +1,6 @@
 from amlb.benchmark import TaskConfig
 from amlb.data import Dataset
-from amlb.utils import call_script_in_same_dir
+from amlb.utils import call_script_in_same_dir, unsparsify
 
 
 def setup(*args, **kwargs):
@@ -12,7 +12,7 @@ def run(dataset: Dataset, config: TaskConfig):
     from frameworks.shared.caller import run_in_venv
 
     X_train, X_test = impute_array(dataset.train.X_enc, dataset.test.X_enc)
-    y_train, y_test = (dataset.train.y_enc, dataset.test.y_enc)
+    y_train, y_test = unsparsify(dataset.train.y_enc, dataset.test.y_enc)
     data = dict(
         train=dict(
             X=X_train,
diff --git a/frameworks/TunedRandomForest/custom_validate.py b/frameworks/TunedRandomForest/custom_validate.py
new file mode 100644
index 000000000..442a8bede
--- /dev/null
+++ b/frameworks/TunedRandomForest/custom_validate.py
@@ -0,0 +1,1936 @@
+"""
+This is a modification of scikit-learn model selection code to allow
+`cross_validate`  functionality together with `warm_start`.
+For more information, see: https://github.com/scikit-learn/scikit-learn/issues/22044
+
+The code uses scikit-learn internal modules and functions as of 1.0.2.
+
+The code is modified under the following license:
+
+BSD 3-Clause License
+
+Copyright (c) 2007-2021 The scikit-learn developers.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* 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.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+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 WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON 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.
+"""
+"""
+The :mod:`sklearn.model_selection._validation` module includes classes and
+functions to validate the model.
+"""
+
+# Author: Alexandre Gramfort <alexandre.gramfort@inria.fr>
+#         Gael Varoquaux <gael.varoquaux@normalesup.org>
+#         Olivier Grisel <olivier.grisel@ensta.org>
+#         Raghav RV <rvraghav93@gmail.com>
+#         Michal Karbownik <michakarbownik@gmail.com>
+# License: BSD 3 clause
+
+
+import warnings
+import numbers
+import time
+from traceback import format_exc
+from contextlib import suppress
+from collections import Counter
+
+import numpy as np
+import scipy.sparse as sp
+from joblib import Parallel, logger
+
+from sklearn.base import is_classifier, clone
+from sklearn.utils import indexable, check_random_state, _safe_indexing
+from sklearn.utils.validation import _check_fit_params
+from sklearn.utils.validation import _num_samples
+from sklearn.utils.fixes import delayed
+from sklearn.utils.metaestimators import _safe_split
+from sklearn.metrics import check_scoring
+from sklearn.metrics._scorer import _check_multimetric_scoring, _MultimetricScorer
+from sklearn.exceptions import FitFailedWarning, NotFittedError
+from sklearn.model_selection._split import check_cv
+from sklearn.preprocessing import LabelEncoder
+
+
+__all__ = [
+    "cross_validate",
+    "cross_val_score",
+    "cross_val_predict",
+    "permutation_test_score",
+    "learning_curve",
+    "validation_curve",
+]
+
+
+def cross_validate(
+    estimators,
+    X,
+    y=None,
+    *,
+    groups=None,
+    scoring=None,
+    cv=None,
+    n_jobs=None,
+    verbose=0,
+    fit_params=None,
+    pre_dispatch="2*n_jobs",
+    return_train_score=False,
+    return_estimator=False,
+    error_score=np.nan,
+):
+    """Evaluate metric(s) by cross-validation and also record fit/score times.
+
+    Read more in the :ref:`User Guide <multimetric_cross_validation>`.
+
+    Parameters
+    ----------
+    estimator : estimator object implementing 'fit'
+        The object to use to fit the data.
+
+    X : array-like of shape (n_samples, n_features)
+        The data to fit. Can be for example a list, or an array.
+
+    y : array-like of shape (n_samples,) or (n_samples, n_outputs), \
+            default=None
+        The target variable to try to predict in the case of
+        supervised learning.
+
+    groups : array-like of shape (n_samples,), default=None
+        Group labels for the samples used while splitting the dataset into
+        train/test set. Only used in conjunction with a "Group" :term:`cv`
+        instance (e.g., :class:`GroupKFold`).
+
+    scoring : str, callable, list, tuple, or dict, default=None
+        Strategy to evaluate the performance of the cross-validated model on
+        the test set.
+
+        If `scoring` represents a single score, one can use:
+
+        - a single string (see :ref:`scoring_parameter`);
+        - a callable (see :ref:`scoring`) that returns a single value.
+
+        If `scoring` represents multiple scores, one can use:
+
+        - a list or tuple of unique strings;
+        - a callable returning a dictionary where the keys are the metric
+          names and the values are the metric scores;
+        - a dictionary with metric names as keys and callables a values.
+
+        See :ref:`multimetric_grid_search` for an example.
+
+    cv : int, cross-validation generator or an iterable, default=None
+        Determines the cross-validation splitting strategy.
+        Possible inputs for cv are:
+
+        - None, to use the default 5-fold cross validation,
+        - int, to specify the number of folds in a `(Stratified)KFold`,
+        - :term:`CV splitter`,
+        - An iterable yielding (train, test) splits as arrays of indices.
+
+        For int/None inputs, if the estimator is a classifier and ``y`` is
+        either binary or multiclass, :class:`StratifiedKFold` is used. In all
+        other cases, :class:`.Fold` is used. These splitters are instantiated
+        with `shuffle=False` so the splits will be the same across calls.
+
+        Refer :ref:`User Guide <cross_validation>` for the various
+        cross-validation strategies that can be used here.
+
+        .. versionchanged:: 0.22
+            ``cv`` default value if None changed from 3-fold to 5-fold.
+
+    n_jobs : int, default=None
+        Number of jobs to run in parallel. Training the estimator and computing
+        the score are parallelized over the cross-validation splits.
+        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
+        ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
+        for more details.
+
+    verbose : int, default=0
+        The verbosity level.
+
+    fit_params : dict, default=None
+        Parameters to pass to the fit method of the estimator.
+
+    pre_dispatch : int or str, default='2*n_jobs'
+        Controls the number of jobs that get dispatched during parallel
+        execution. Reducing this number can be useful to avoid an
+        explosion of memory consumption when more jobs get dispatched
+        than CPUs can process. This parameter can be:
+
+            - None, in which case all the jobs are immediately
+              created and spawned. Use this for lightweight and
+              fast-running jobs, to avoid delays due to on-demand
+              spawning of the jobs
+
+            - An int, giving the exact number of total jobs that are
+              spawned
+
+            - A str, giving an expression as a function of n_jobs,
+              as in '2*n_jobs'
+
+    return_train_score : bool, default=False
+        Whether to include train scores.
+        Computing training scores is used to get insights on how different
+        parameter settings impact the overfitting/underfitting trade-off.
+        However computing the scores on the training set can be computationally
+        expensive and is not strictly required to select the parameters that
+        yield the best generalization performance.
+
+        .. versionadded:: 0.19
+
+        .. versionchanged:: 0.21
+            Default value was changed from ``True`` to ``False``
+
+    return_estimator : bool, default=False
+        Whether to return the estimators fitted on each split.
+
+        .. versionadded:: 0.20
+
+    error_score : 'raise' or numeric, default=np.nan
+        Value to assign to the score if an error occurs in estimator fitting.
+        If set to 'raise', the error is raised.
+        If a numeric value is given, FitFailedWarning is raised.
+
+        .. versionadded:: 0.20
+
+    Returns
+    -------
+    scores : dict of float arrays of shape (n_splits,)
+        Array of scores of the estimator for each run of the cross validation.
+
+        A dict of arrays containing the score/time arrays for each scorer is
+        returned. The possible keys for this ``dict`` are:
+
+            ``test_score``
+                The score array for test scores on each cv split.
+                Suffix ``_score`` in ``test_score`` changes to a specific
+                metric like ``test_r2`` or ``test_auc`` if there are
+                multiple scoring metrics in the scoring parameter.
+            ``train_score``
+                The score array for train scores on each cv split.
+                Suffix ``_score`` in ``train_score`` changes to a specific
+                metric like ``train_r2`` or ``train_auc`` if there are
+                multiple scoring metrics in the scoring parameter.
+                This is available only if ``return_train_score`` parameter
+                is ``True``.
+            ``fit_time``
+                The time for fitting the estimator on the train
+                set for each cv split.
+            ``score_time``
+                The time for scoring the estimator on the test set for each
+                cv split. (Note time for scoring on the train set is not
+                included even if ``return_train_score`` is set to ``True``
+            ``estimator``
+                The estimator objects for each cv split.
+                This is available only if ``return_estimator`` parameter
+                is set to ``True``.
+
+    Examples
+    --------
+    >>> from sklearn import datasets, linear_model
+    >>> from sklearn.model_selection import cross_validate
+    >>> from sklearn.metrics import make_scorer
+    >>> from sklearn.metrics import confusion_matrix
+    >>> from sklearn.svm import LinearSVC
+    >>> diabetes = datasets.load_diabetes()
+    >>> X = diabetes.data[:150]
+    >>> y = diabetes.target[:150]
+    >>> lasso = linear_model.Lasso()
+
+    Single metric evaluation using ``cross_validate``
+
+    >>> cv_results = cross_validate(lasso, X, y, cv=3)
+    >>> sorted(cv_results.keys())
+    ['fit_time', 'score_time', 'test_score']
+    >>> cv_results['test_score']
+    array([0.33150734, 0.08022311, 0.03531764])
+
+    Multiple metric evaluation using ``cross_validate``
+    (please refer the ``scoring`` parameter doc for more information)
+
+    >>> scores = cross_validate(lasso, X, y, cv=3,
+    ...                         scoring=('r2', 'neg_mean_squared_error'),
+    ...                         return_train_score=True)
+    >>> print(scores['test_neg_mean_squared_error'])
+    [-3635.5... -3573.3... -6114.7...]
+    >>> print(scores['train_r2'])
+    [0.28010158 0.39088426 0.22784852]
+
+    See Also
+    ---------
+    cross_val_score : Run cross-validation for single metric evaluation.
+
+    cross_val_predict : Get predictions from each split of cross-validation for
+        diagnostic purposes.
+
+    sklearn.metrics.make_scorer : Make a scorer from a performance metric or
+        loss function.
+
+    """
+    # Because we don't have to support a general use-case in this adaptation,
+    # we simply use the first estimator of the list for general checks.
+    if isinstance(estimators, list):
+        estimator = estimators[0]
+    else:
+        estimator = estimators
+
+    X, y, groups = indexable(X, y, groups)
+
+    cv = check_cv(cv, y, classifier=is_classifier(estimator))
+    if not isinstance(estimators, list):
+        estimators = [clone(estimator) for _ in range(cv.n_splits)]
+
+    if callable(scoring):
+        scorers = scoring
+    elif scoring is None or isinstance(scoring, str):
+        scorers = check_scoring(estimator, scoring)
+    else:
+        scorers = _check_multimetric_scoring(estimator, scoring)
+
+    # We clone the estimator to make sure that all the folds are
+    # independent, and that it is pickle-able.
+
+    from sklearn.utils import parallel_backend
+    with parallel_backend("threading"):
+        parallel = Parallel(n_jobs=n_jobs, verbose=verbose, pre_dispatch=pre_dispatch)
+        results = parallel(
+            delayed(_fit_and_score)(
+                estimator,
+                X,
+                y,
+                scorers,
+                train,
+                test,
+                verbose,
+                None,
+                fit_params,
+                return_train_score=return_train_score,
+                return_times=True,
+                return_estimator=return_estimator,
+                error_score=error_score,
+            )
+            for estimator, (train, test) in zip(estimators, cv.split(X, y, groups))
+        )
+
+    _warn_about_fit_failures(results, error_score)
+
+    # For callabe scoring, the return type is only know after calling. If the
+    # return type is a dictionary, the error scores can now be inserted with
+    # the correct key.
+    if callable(scoring):
+        _insert_error_scores(results, error_score)
+
+    results = _aggregate_score_dicts(results)
+
+    ret = {}
+    ret["fit_time"] = results["fit_time"]
+    ret["score_time"] = results["score_time"]
+
+    if return_estimator:
+        ret["estimator"] = results["estimator"]
+
+    test_scores_dict = _normalize_score_results(results["test_scores"])
+    if return_train_score:
+        train_scores_dict = _normalize_score_results(results["train_scores"])
+
+    for name in test_scores_dict:
+        ret["test_%s" % name] = test_scores_dict[name]
+        if return_train_score:
+            key = "train_%s" % name
+            ret[key] = train_scores_dict[name]
+
+    return ret
+
+
+def _insert_error_scores(results, error_score):
+    """Insert error in `results` by replacing them inplace with `error_score`.
+
+    This only applies to multimetric scores because `_fit_and_score` will
+    handle the single metric case.
+    """
+    successful_score = None
+    failed_indices = []
+    for i, result in enumerate(results):
+        if result["fit_error"] is not None:
+            failed_indices.append(i)
+        elif successful_score is None:
+            successful_score = result["test_scores"]
+
+    if successful_score is None:
+        raise NotFittedError("All estimators failed to fit")
+
+    if isinstance(successful_score, dict):
+        formatted_error = {name: error_score for name in successful_score}
+        for i in failed_indices:
+            results[i]["test_scores"] = formatted_error.copy()
+            if "train_scores" in results[i]:
+                results[i]["train_scores"] = formatted_error.copy()
+
+
+def _normalize_score_results(scores, scaler_score_key="score"):
+    """Creates a scoring dictionary based on the type of `scores`"""
+    if isinstance(scores[0], dict):
+        # multimetric scoring
+        return _aggregate_score_dicts(scores)
+    # scaler
+    return {scaler_score_key: scores}
+
+
+def _warn_about_fit_failures(results, error_score):
+    fit_errors = [
+        result["fit_error"] for result in results if result["fit_error"] is not None
+    ]
+    if fit_errors:
+        num_failed_fits = len(fit_errors)
+        num_fits = len(results)
+        fit_errors_counter = Counter(fit_errors)
+        delimiter = "-" * 80 + "\n"
+        fit_errors_summary = "\n".join(
+            f"{delimiter}{n} fits failed with the following error:\n{error}"
+            for error, n in fit_errors_counter.items()
+        )
+
+        some_fits_failed_message = (
+            f"\n{num_failed_fits} fits failed out of a total of {num_fits}.\n"
+            "The score on these train-test partitions for these parameters"
+            f" will be set to {error_score}.\n"
+            "If these failures are not expected, you can try to debug them "
+            "by setting error_score='raise'.\n\n"
+            f"Below are more details about the failures:\n{fit_errors_summary}"
+        )
+        warnings.warn(some_fits_failed_message, FitFailedWarning)
+
+
+def cross_val_score(
+    estimator,
+    X,
+    y=None,
+    *,
+    groups=None,
+    scoring=None,
+    cv=None,
+    n_jobs=None,
+    verbose=0,
+    fit_params=None,
+    pre_dispatch="2*n_jobs",
+    error_score=np.nan,
+):
+    """Evaluate a score by cross-validation.
+
+    Read more in the :ref:`User Guide <cross_validation>`.
+
+    Parameters
+    ----------
+    estimator : estimator object implementing 'fit'
+        The object to use to fit the data.
+
+    X : array-like of shape (n_samples, n_features)
+        The data to fit. Can be for example a list, or an array.
+
+    y : array-like of shape (n_samples,) or (n_samples, n_outputs), \
+            default=None
+        The target variable to try to predict in the case of
+        supervised learning.
+
+    groups : array-like of shape (n_samples,), default=None
+        Group labels for the samples used while splitting the dataset into
+        train/test set. Only used in conjunction with a "Group" :term:`cv`
+        instance (e.g., :class:`GroupKFold`).
+
+    scoring : str or callable, default=None
+        A str (see model evaluation documentation) or
+        a scorer callable object / function with signature
+        ``scorer(estimator, X, y)`` which should return only
+        a single value.
+
+        Similar to :func:`cross_validate`
+        but only a single metric is permitted.
+
+        If `None`, the estimator's default scorer (if available) is used.
+
+    cv : int, cross-validation generator or an iterable, default=None
+        Determines the cross-validation splitting strategy.
+        Possible inputs for cv are:
+
+        - `None`, to use the default 5-fold cross validation,
+        - int, to specify the number of folds in a `(Stratified)KFold`,
+        - :term:`CV splitter`,
+        - An iterable that generates (train, test) splits as arrays of indices.
+
+        For `int`/`None` inputs, if the estimator is a classifier and `y` is
+        either binary or multiclass, :class:`StratifiedKFold` is used. In all
+        other cases, :class:`KFold` is used. These splitters are instantiated
+        with `shuffle=False` so the splits will be the same across calls.
+
+        Refer :ref:`User Guide <cross_validation>` for the various
+        cross-validation strategies that can be used here.
+
+        .. versionchanged:: 0.22
+            `cv` default value if `None` changed from 3-fold to 5-fold.
+
+    n_jobs : int, default=None
+        Number of jobs to run in parallel. Training the estimator and computing
+        the score are parallelized over the cross-validation splits.
+        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
+        ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
+        for more details.
+
+    verbose : int, default=0
+        The verbosity level.
+
+    fit_params : dict, default=None
+        Parameters to pass to the fit method of the estimator.
+
+    pre_dispatch : int or str, default='2*n_jobs'
+        Controls the number of jobs that get dispatched during parallel
+        execution. Reducing this number can be useful to avoid an
+        explosion of memory consumption when more jobs get dispatched
+        than CPUs can process. This parameter can be:
+
+            - ``None``, in which case all the jobs are immediately
+              created and spawned. Use this for lightweight and
+              fast-running jobs, to avoid delays due to on-demand
+              spawning of the jobs
+
+            - An int, giving the exact number of total jobs that are
+              spawned
+
+            - A str, giving an expression as a function of n_jobs,
+              as in '2*n_jobs'
+
+    error_score : 'raise' or numeric, default=np.nan
+        Value to assign to the score if an error occurs in estimator fitting.
+        If set to 'raise', the error is raised.
+        If a numeric value is given, FitFailedWarning is raised.
+
+        .. versionadded:: 0.20
+
+    Returns
+    -------
+    scores : ndarray of float of shape=(len(list(cv)),)
+        Array of scores of the estimator for each run of the cross validation.
+
+    Examples
+    --------
+    >>> from sklearn import datasets, linear_model
+    >>> from sklearn.model_selection import cross_val_score
+    >>> diabetes = datasets.load_diabetes()
+    >>> X = diabetes.data[:150]
+    >>> y = diabetes.target[:150]
+    >>> lasso = linear_model.Lasso()
+    >>> print(cross_val_score(lasso, X, y, cv=3))
+    [0.33150734 0.08022311 0.03531764]
+
+    See Also
+    ---------
+    cross_validate : To run cross-validation on multiple metrics and also to
+        return train scores, fit times and score times.
+
+    cross_val_predict : Get predictions from each split of cross-validation for
+        diagnostic purposes.
+
+    sklearn.metrics.make_scorer : Make a scorer from a performance metric or
+        loss function.
+    """
+    # To ensure multimetric format is not supported
+    scorer = check_scoring(estimator, scoring=scoring)
+
+    cv_results = cross_validate(
+        estimator=estimator,
+        X=X,
+        y=y,
+        groups=groups,
+        scoring={"score": scorer},
+        cv=cv,
+        n_jobs=n_jobs,
+        verbose=verbose,
+        fit_params=fit_params,
+        pre_dispatch=pre_dispatch,
+        error_score=error_score,
+    )
+    return cv_results["test_score"]
+
+
+def _fit_and_score(
+    estimator,
+    X,
+    y,
+    scorer,
+    train,
+    test,
+    verbose,
+    parameters,
+    fit_params,
+    return_train_score=False,
+    return_parameters=False,
+    return_n_test_samples=False,
+    return_times=False,
+    return_estimator=False,
+    split_progress=None,
+    candidate_progress=None,
+    error_score=np.nan,
+):
+
+    """Fit estimator and compute scores for a given dataset split.
+
+    Parameters
+    ----------
+    estimator : estimator object implementing 'fit'
+        The object to use to fit the data.
+
+    X : array-like of shape (n_samples, n_features)
+        The data to fit.
+
+    y : array-like of shape (n_samples,) or (n_samples, n_outputs) or None
+        The target variable to try to predict in the case of
+        supervised learning.
+
+    scorer : A single callable or dict mapping scorer name to the callable
+        If it is a single callable, the return value for ``train_scores`` and
+        ``test_scores`` is a single float.
+
+        For a dict, it should be one mapping the scorer name to the scorer
+        callable object / function.
+
+        The callable object / fn should have signature
+        ``scorer(estimator, X, y)``.
+
+    train : array-like of shape (n_train_samples,)
+        Indices of training samples.
+
+    test : array-like of shape (n_test_samples,)
+        Indices of test samples.
+
+    verbose : int
+        The verbosity level.
+
+    error_score : 'raise' or numeric, default=np.nan
+        Value to assign to the score if an error occurs in estimator fitting.
+        If set to 'raise', the error is raised.
+        If a numeric value is given, FitFailedWarning is raised.
+
+    parameters : dict or None
+        Parameters to be set on the estimator.
+
+    fit_params : dict or None
+        Parameters that will be passed to ``estimator.fit``.
+
+    return_train_score : bool, default=False
+        Compute and return score on training set.
+
+    return_parameters : bool, default=False
+        Return parameters that has been used for the estimator.
+
+    split_progress : {list, tuple} of int, default=None
+        A list or tuple of format (<current_split_id>, <total_num_of_splits>).
+
+    candidate_progress : {list, tuple} of int, default=None
+        A list or tuple of format
+        (<current_candidate_id>, <total_number_of_candidates>).
+
+    return_n_test_samples : bool, default=False
+        Whether to return the ``n_test_samples``.
+
+    return_times : bool, default=False
+        Whether to return the fit/score times.
+
+    return_estimator : bool, default=False
+        Whether to return the fitted estimator.
+
+    Returns
+    -------
+    result : dict with the following attributes
+        train_scores : dict of scorer name -> float
+            Score on training set (for all the scorers),
+            returned only if `return_train_score` is `True`.
+        test_scores : dict of scorer name -> float
+            Score on testing set (for all the scorers).
+        n_test_samples : int
+            Number of test samples.
+        fit_time : float
+            Time spent for fitting in seconds.
+        score_time : float
+            Time spent for scoring in seconds.
+        parameters : dict or None
+            The parameters that have been evaluated.
+        estimator : estimator object
+            The fitted estimator.
+        fit_error : str or None
+            Traceback str if the fit failed, None if the fit succeeded.
+    """
+    if not isinstance(error_score, numbers.Number) and error_score != "raise":
+        raise ValueError(
+            "error_score must be the string 'raise' or a numeric value. "
+            "(Hint: if using 'raise', please make sure that it has been "
+            "spelled correctly.)"
+        )
+
+    progress_msg = ""
+    if verbose > 2:
+        if split_progress is not None:
+            progress_msg = f" {split_progress[0]+1}/{split_progress[1]}"
+        if candidate_progress and verbose > 9:
+            progress_msg += f"; {candidate_progress[0]+1}/{candidate_progress[1]}"
+
+    if verbose > 1:
+        if parameters is None:
+            params_msg = ""
+        else:
+            sorted_keys = sorted(parameters)  # Ensure deterministic o/p
+            params_msg = ", ".join(f"{k}={parameters[k]}" for k in sorted_keys)
+    if verbose > 9:
+        start_msg = f"[CV{progress_msg}] START {params_msg}"
+        print(f"{start_msg}{(80 - len(start_msg)) * '.'}")
+
+    # Adjust length of sample weights
+    fit_params = fit_params if fit_params is not None else {}
+    fit_params = _check_fit_params(X, fit_params, train)
+
+    if parameters is not None:
+        # clone after setting parameters in case any parameters
+        # are estimators (like pipeline steps)
+        # because pipeline doesn't clone steps in fit
+        cloned_parameters = {}
+        for k, v in parameters.items():
+            cloned_parameters[k] = clone(v, safe=False)
+
+        estimator = estimator.set_params(**cloned_parameters)
+
+    start_time = time.time()
+
+    X_train, y_train = _safe_split(estimator, X, y, train)
+    X_test, y_test = _safe_split(estimator, X, y, test, train)
+
+    result = {}
+    try:
+        if y_train is None:
+            estimator.fit(X_train, **fit_params)
+        else:
+            estimator.fit(X_train, y_train, **fit_params)
+
+    except Exception:
+        # Note fit time as time until error
+        fit_time = time.time() - start_time
+        score_time = 0.0
+        if error_score == "raise":
+            raise
+        elif isinstance(error_score, numbers.Number):
+            if isinstance(scorer, dict):
+                test_scores = {name: error_score for name in scorer}
+                if return_train_score:
+                    train_scores = test_scores.copy()
+            else:
+                test_scores = error_score
+                if return_train_score:
+                    train_scores = error_score
+        result["fit_error"] = format_exc()
+    else:
+        result["fit_error"] = None
+
+        fit_time = time.time() - start_time
+        test_scores = _score(estimator, X_test, y_test, scorer, error_score)
+        score_time = time.time() - start_time - fit_time
+        if return_train_score:
+            train_scores = _score(estimator, X_train, y_train, scorer, error_score)
+
+    if verbose > 1:
+        total_time = score_time + fit_time
+        end_msg = f"[CV{progress_msg}] END "
+        result_msg = params_msg + (";" if params_msg else "")
+        if verbose > 2:
+            if isinstance(test_scores, dict):
+                for scorer_name in sorted(test_scores):
+                    result_msg += f" {scorer_name}: ("
+                    if return_train_score:
+                        scorer_scores = train_scores[scorer_name]
+                        result_msg += f"train={scorer_scores:.3f}, "
+                    result_msg += f"test={test_scores[scorer_name]:.3f})"
+            else:
+                result_msg += ", score="
+                if return_train_score:
+                    result_msg += f"(train={train_scores:.3f}, test={test_scores:.3f})"
+                else:
+                    result_msg += f"{test_scores:.3f}"
+        result_msg += f" total time={logger.short_format_time(total_time)}"
+
+        # Right align the result_msg
+        end_msg += "." * (80 - len(end_msg) - len(result_msg))
+        end_msg += result_msg
+        print(end_msg)
+
+    result["test_scores"] = test_scores
+    if return_train_score:
+        result["train_scores"] = train_scores
+    if return_n_test_samples:
+        result["n_test_samples"] = _num_samples(X_test)
+    if return_times:
+        result["fit_time"] = fit_time
+        result["score_time"] = score_time
+    if return_parameters:
+        result["parameters"] = parameters
+    if return_estimator:
+        result["estimator"] = estimator
+    return result
+
+
+def _score(estimator, X_test, y_test, scorer, error_score="raise"):
+    """Compute the score(s) of an estimator on a given test set.
+
+    Will return a dict of floats if `scorer` is a dict, otherwise a single
+    float is returned.
+    """
+    if isinstance(scorer, dict):
+        # will cache method calls if needed. scorer() returns a dict
+        scorer = _MultimetricScorer(**scorer)
+
+    try:
+        if y_test is None:
+            scores = scorer(estimator, X_test)
+        else:
+            scores = scorer(estimator, X_test, y_test)
+    except Exception:
+        if error_score == "raise":
+            raise
+        else:
+            if isinstance(scorer, _MultimetricScorer):
+                scores = {name: error_score for name in scorer._scorers}
+            else:
+                scores = error_score
+            warnings.warn(
+                "Scoring failed. The score on this train-test partition for "
+                f"these parameters will be set to {error_score}. Details: \n"
+                f"{format_exc()}",
+                UserWarning,
+            )
+
+    error_msg = "scoring must return a number, got %s (%s) instead. (scorer=%s)"
+    if isinstance(scores, dict):
+        for name, score in scores.items():
+            if hasattr(score, "item"):
+                with suppress(ValueError):
+                    # e.g. unwrap memmapped scalars
+                    score = score.item()
+            if not isinstance(score, numbers.Number):
+                raise ValueError(error_msg % (score, type(score), name))
+            scores[name] = score
+    else:  # scalar
+        if hasattr(scores, "item"):
+            with suppress(ValueError):
+                # e.g. unwrap memmapped scalars
+                scores = scores.item()
+        if not isinstance(scores, numbers.Number):
+            raise ValueError(error_msg % (scores, type(scores), scorer))
+    return scores
+
+
+def cross_val_predict(
+    estimator,
+    X,
+    y=None,
+    *,
+    groups=None,
+    cv=None,
+    n_jobs=None,
+    verbose=0,
+    fit_params=None,
+    pre_dispatch="2*n_jobs",
+    method="predict",
+):
+    """Generate cross-validated estimates for each input data point.
+
+    The data is split according to the cv parameter. Each sample belongs
+    to exactly one test set, and its prediction is computed with an
+    estimator fitted on the corresponding training set.
+
+    Passing these predictions into an evaluation metric may not be a valid
+    way to measure generalization performance. Results can differ from
+    :func:`cross_validate` and :func:`cross_val_score` unless all tests sets
+    have equal size and the metric decomposes over samples.
+
+    Read more in the :ref:`User Guide <cross_validation>`.
+
+    Parameters
+    ----------
+    estimator : estimator object implementing 'fit' and 'predict'
+        The object to use to fit the data.
+
+    X : array-like of shape (n_samples, n_features)
+        The data to fit. Can be, for example a list, or an array at least 2d.
+
+    y : array-like of shape (n_samples,) or (n_samples, n_outputs), \
+            default=None
+        The target variable to try to predict in the case of
+        supervised learning.
+
+    groups : array-like of shape (n_samples,), default=None
+        Group labels for the samples used while splitting the dataset into
+        train/test set. Only used in conjunction with a "Group" :term:`cv`
+        instance (e.g., :class:`GroupKFold`).
+
+    cv : int, cross-validation generator or an iterable, default=None
+        Determines the cross-validation splitting strategy.
+        Possible inputs for cv are:
+
+        - None, to use the default 5-fold cross validation,
+        - int, to specify the number of folds in a `(Stratified)KFold`,
+        - :term:`CV splitter`,
+        - An iterable that generates (train, test) splits as arrays of indices.
+
+        For int/None inputs, if the estimator is a classifier and ``y`` is
+        either binary or multiclass, :class:`StratifiedKFold` is used. In all
+        other cases, :class:`KFold` is used. These splitters are instantiated
+        with `shuffle=False` so the splits will be the same across calls.
+
+        Refer :ref:`User Guide <cross_validation>` for the various
+        cross-validation strategies that can be used here.
+
+        .. versionchanged:: 0.22
+            ``cv`` default value if None changed from 3-fold to 5-fold.
+
+    n_jobs : int, default=None
+        Number of jobs to run in parallel. Training the estimator and
+        predicting are parallelized over the cross-validation splits.
+        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
+        ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
+        for more details.
+
+    verbose : int, default=0
+        The verbosity level.
+
+    fit_params : dict, default=None
+        Parameters to pass to the fit method of the estimator.
+
+    pre_dispatch : int or str, default='2*n_jobs'
+        Controls the number of jobs that get dispatched during parallel
+        execution. Reducing this number can be useful to avoid an
+        explosion of memory consumption when more jobs get dispatched
+        than CPUs can process. This parameter can be:
+
+            - None, in which case all the jobs are immediately
+              created and spawned. Use this for lightweight and
+              fast-running jobs, to avoid delays due to on-demand
+              spawning of the jobs
+
+            - An int, giving the exact number of total jobs that are
+              spawned
+
+            - A str, giving an expression as a function of n_jobs,
+              as in '2*n_jobs'
+
+    method : {'predict', 'predict_proba', 'predict_log_proba', \
+              'decision_function'}, default='predict'
+        The method to be invoked by `estimator`.
+
+    Returns
+    -------
+    predictions : ndarray
+        This is the result of calling `method`. Shape:
+
+            - When `method` is 'predict' and in special case where `method` is
+              'decision_function' and the target is binary: (n_samples,)
+            - When `method` is one of {'predict_proba', 'predict_log_proba',
+              'decision_function'} (unless special case above):
+              (n_samples, n_classes)
+            - If `estimator` is :term:`multioutput`, an extra dimension
+              'n_outputs' is added to the end of each shape above.
+
+    See Also
+    --------
+    cross_val_score : Calculate score for each CV split.
+    cross_validate : Calculate one or more scores and timings for each CV
+        split.
+
+    Notes
+    -----
+    In the case that one or more classes are absent in a training portion, a
+    default score needs to be assigned to all instances for that class if
+    ``method`` produces columns per class, as in {'decision_function',
+    'predict_proba', 'predict_log_proba'}.  For ``predict_proba`` this value is
+    0.  In order to ensure finite output, we approximate negative infinity by
+    the minimum finite float value for the dtype in other cases.
+
+    Examples
+    --------
+    >>> from sklearn import datasets, linear_model
+    >>> from sklearn.model_selection import cross_val_predict
+    >>> diabetes = datasets.load_diabetes()
+    >>> X = diabetes.data[:150]
+    >>> y = diabetes.target[:150]
+    >>> lasso = linear_model.Lasso()
+    >>> y_pred = cross_val_predict(lasso, X, y, cv=3)
+    """
+    X, y, groups = indexable(X, y, groups)
+
+    cv = check_cv(cv, y, classifier=is_classifier(estimator))
+    splits = list(cv.split(X, y, groups))
+
+    test_indices = np.concatenate([test for _, test in splits])
+    if not _check_is_permutation(test_indices, _num_samples(X)):
+        raise ValueError("cross_val_predict only works for partitions")
+
+    # If classification methods produce multiple columns of output,
+    # we need to manually encode classes to ensure consistent column ordering.
+    encode = (
+        method in ["decision_function", "predict_proba", "predict_log_proba"]
+        and y is not None
+    )
+    if encode:
+        y = np.asarray(y)
+        if y.ndim == 1:
+            le = LabelEncoder()
+            y = le.fit_transform(y)
+        elif y.ndim == 2:
+            y_enc = np.zeros_like(y, dtype=int)
+            for i_label in range(y.shape[1]):
+                y_enc[:, i_label] = LabelEncoder().fit_transform(y[:, i_label])
+            y = y_enc
+
+    # We clone the estimator to make sure that all the folds are
+    # independent, and that it is pickle-able.
+    parallel = Parallel(n_jobs=n_jobs, verbose=verbose, pre_dispatch=pre_dispatch)
+    predictions = parallel(
+        delayed(_fit_and_predict)(
+            clone(estimator), X, y, train, test, verbose, fit_params, method
+        )
+        for train, test in splits
+    )
+
+    inv_test_indices = np.empty(len(test_indices), dtype=int)
+    inv_test_indices[test_indices] = np.arange(len(test_indices))
+
+    if sp.issparse(predictions[0]):
+        predictions = sp.vstack(predictions, format=predictions[0].format)
+    elif encode and isinstance(predictions[0], list):
+        # `predictions` is a list of method outputs from each fold.
+        # If each of those is also a list, then treat this as a
+        # multioutput-multiclass task. We need to separately concatenate
+        # the method outputs for each label into an `n_labels` long list.
+        n_labels = y.shape[1]
+        concat_pred = []
+        for i_label in range(n_labels):
+            label_preds = np.concatenate([p[i_label] for p in predictions])
+            concat_pred.append(label_preds)
+        predictions = concat_pred
+    else:
+        predictions = np.concatenate(predictions)
+
+    if isinstance(predictions, list):
+        return [p[inv_test_indices] for p in predictions]
+    else:
+        return predictions[inv_test_indices]
+
+
+def _fit_and_predict(estimator, X, y, train, test, verbose, fit_params, method):
+    """Fit estimator and predict values for a given dataset split.
+
+    Read more in the :ref:`User Guide <cross_validation>`.
+
+    Parameters
+    ----------
+    estimator : estimator object implementing 'fit' and 'predict'
+        The object to use to fit the data.
+
+    X : array-like of shape (n_samples, n_features)
+        The data to fit.
+
+        .. versionchanged:: 0.20
+            X is only required to be an object with finite length or shape now
+
+    y : array-like of shape (n_samples,) or (n_samples, n_outputs) or None
+        The target variable to try to predict in the case of
+        supervised learning.
+
+    train : array-like of shape (n_train_samples,)
+        Indices of training samples.
+
+    test : array-like of shape (n_test_samples,)
+        Indices of test samples.
+
+    verbose : int
+        The verbosity level.
+
+    fit_params : dict or None
+        Parameters that will be passed to ``estimator.fit``.
+
+    method : str
+        Invokes the passed method name of the passed estimator.
+
+    Returns
+    -------
+    predictions : sequence
+        Result of calling 'estimator.method'
+    """
+    # Adjust length of sample weights
+    fit_params = fit_params if fit_params is not None else {}
+    fit_params = _check_fit_params(X, fit_params, train)
+
+    X_train, y_train = _safe_split(estimator, X, y, train)
+    X_test, _ = _safe_split(estimator, X, y, test, train)
+
+    if y_train is None:
+        estimator.fit(X_train, **fit_params)
+    else:
+        estimator.fit(X_train, y_train, **fit_params)
+    func = getattr(estimator, method)
+    predictions = func(X_test)
+
+    encode = (
+        method in ["decision_function", "predict_proba", "predict_log_proba"]
+        and y is not None
+    )
+
+    if encode:
+        if isinstance(predictions, list):
+            predictions = [
+                _enforce_prediction_order(
+                    estimator.classes_[i_label],
+                    predictions[i_label],
+                    n_classes=len(set(y[:, i_label])),
+                    method=method,
+                )
+                for i_label in range(len(predictions))
+            ]
+        else:
+            # A 2D y array should be a binary label indicator matrix
+            n_classes = len(set(y)) if y.ndim == 1 else y.shape[1]
+            predictions = _enforce_prediction_order(
+                estimator.classes_, predictions, n_classes, method
+            )
+    return predictions
+
+
+def _enforce_prediction_order(classes, predictions, n_classes, method):
+    """Ensure that prediction arrays have correct column order
+
+    When doing cross-validation, if one or more classes are
+    not present in the subset of data used for training,
+    then the output prediction array might not have the same
+    columns as other folds. Use the list of class names
+    (assumed to be ints) to enforce the correct column order.
+
+    Note that `classes` is the list of classes in this fold
+    (a subset of the classes in the full training set)
+    and `n_classes` is the number of classes in the full training set.
+    """
+    if n_classes != len(classes):
+        recommendation = (
+            "To fix this, use a cross-validation "
+            "technique resulting in properly "
+            "stratified folds"
+        )
+        warnings.warn(
+            "Number of classes in training fold ({}) does "
+            "not match total number of classes ({}). "
+            "Results may not be appropriate for your use case. "
+            "{}".format(len(classes), n_classes, recommendation),
+            RuntimeWarning,
+        )
+        if method == "decision_function":
+            if predictions.ndim == 2 and predictions.shape[1] != len(classes):
+                # This handles the case when the shape of predictions
+                # does not match the number of classes used to train
+                # it with. This case is found when sklearn.svm.SVC is
+                # set to `decision_function_shape='ovo'`.
+                raise ValueError(
+                    "Output shape {} of {} does not match "
+                    "number of classes ({}) in fold. "
+                    "Irregular decision_function outputs "
+                    "are not currently supported by "
+                    "cross_val_predict".format(predictions.shape, method, len(classes))
+                )
+            if len(classes) <= 2:
+                # In this special case, `predictions` contains a 1D array.
+                raise ValueError(
+                    "Only {} class/es in training fold, but {} "
+                    "in overall dataset. This "
+                    "is not supported for decision_function "
+                    "with imbalanced folds. {}".format(
+                        len(classes), n_classes, recommendation
+                    )
+                )
+
+        float_min = np.finfo(predictions.dtype).min
+        default_values = {
+            "decision_function": float_min,
+            "predict_log_proba": float_min,
+            "predict_proba": 0,
+        }
+        predictions_for_all_classes = np.full(
+            (_num_samples(predictions), n_classes),
+            default_values[method],
+            dtype=predictions.dtype,
+        )
+        predictions_for_all_classes[:, classes] = predictions
+        predictions = predictions_for_all_classes
+    return predictions
+
+
+def _check_is_permutation(indices, n_samples):
+    """Check whether indices is a reordering of the array np.arange(n_samples)
+
+    Parameters
+    ----------
+    indices : ndarray
+        int array to test
+    n_samples : int
+        number of expected elements
+
+    Returns
+    -------
+    is_partition : bool
+        True iff sorted(indices) is np.arange(n)
+    """
+    if len(indices) != n_samples:
+        return False
+    hit = np.zeros(n_samples, dtype=bool)
+    hit[indices] = True
+    if not np.all(hit):
+        return False
+    return True
+
+
+def permutation_test_score(
+    estimator,
+    X,
+    y,
+    *,
+    groups=None,
+    cv=None,
+    n_permutations=100,
+    n_jobs=None,
+    random_state=0,
+    verbose=0,
+    scoring=None,
+    fit_params=None,
+):
+    """Evaluate the significance of a cross-validated score with permutations
+
+    Permutes targets to generate 'randomized data' and compute the empirical
+    p-value against the null hypothesis that features and targets are
+    independent.
+
+    The p-value represents the fraction of randomized data sets where the
+    estimator performed as well or better than in the original data. A small
+    p-value suggests that there is a real dependency between features and
+    targets which has been used by the estimator to give good predictions.
+    A large p-value may be due to lack of real dependency between features
+    and targets or the estimator was not able to use the dependency to
+    give good predictions.
+
+    Read more in the :ref:`User Guide <permutation_test_score>`.
+
+    Parameters
+    ----------
+    estimator : estimator object implementing 'fit'
+        The object to use to fit the data.
+
+    X : array-like of shape at least 2D
+        The data to fit.
+
+    y : array-like of shape (n_samples,) or (n_samples, n_outputs) or None
+        The target variable to try to predict in the case of
+        supervised learning.
+
+    groups : array-like of shape (n_samples,), default=None
+        Labels to constrain permutation within groups, i.e. ``y`` values
+        are permuted among samples with the same group identifier.
+        When not specified, ``y`` values are permuted among all samples.
+
+        When a grouped cross-validator is used, the group labels are
+        also passed on to the ``split`` method of the cross-validator. The
+        cross-validator uses them for grouping the samples  while splitting
+        the dataset into train/test set.
+
+    scoring : str or callable, default=None
+        A single str (see :ref:`scoring_parameter`) or a callable
+        (see :ref:`scoring`) to evaluate the predictions on the test set.
+
+        If `None` the estimator's score method is used.
+
+    cv : int, cross-validation generator or an iterable, default=None
+        Determines the cross-validation splitting strategy.
+        Possible inputs for cv are:
+
+        - `None`, to use the default 5-fold cross validation,
+        - int, to specify the number of folds in a `(Stratified)KFold`,
+        - :term:`CV splitter`,
+        - An iterable yielding (train, test) splits as arrays of indices.
+
+        For `int`/`None` inputs, if the estimator is a classifier and `y` is
+        either binary or multiclass, :class:`StratifiedKFold` is used. In all
+        other cases, :class:`KFold` is used. These splitters are instantiated
+        with `shuffle=False` so the splits will be the same across calls.
+
+        Refer :ref:`User Guide <cross_validation>` for the various
+        cross-validation strategies that can be used here.
+
+        .. versionchanged:: 0.22
+            `cv` default value if `None` changed from 3-fold to 5-fold.
+
+    n_permutations : int, default=100
+        Number of times to permute ``y``.
+
+    n_jobs : int, default=None
+        Number of jobs to run in parallel. Training the estimator and computing
+        the cross-validated score are parallelized over the permutations.
+        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
+        ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
+        for more details.
+
+    random_state : int, RandomState instance or None, default=0
+        Pass an int for reproducible output for permutation of
+        ``y`` values among samples. See :term:`Glossary <random_state>`.
+
+    verbose : int, default=0
+        The verbosity level.
+
+    fit_params : dict, default=None
+        Parameters to pass to the fit method of the estimator.
+
+        .. versionadded:: 0.24
+
+    Returns
+    -------
+    score : float
+        The true score without permuting targets.
+
+    permutation_scores : array of shape (n_permutations,)
+        The scores obtained for each permutations.
+
+    pvalue : float
+        The p-value, which approximates the probability that the score would
+        be obtained by chance. This is calculated as:
+
+        `(C + 1) / (n_permutations + 1)`
+
+        Where C is the number of permutations whose score >= the true score.
+
+        The best possible p-value is 1/(n_permutations + 1), the worst is 1.0.
+
+    Notes
+    -----
+    This function implements Test 1 in:
+
+        Ojala and Garriga. `Permutation Tests for Studying Classifier
+        Performance
+        <http://www.jmlr.org/papers/volume11/ojala10a/ojala10a.pdf>`_. The
+        Journal of Machine Learning Research (2010) vol. 11
+
+    """
+    X, y, groups = indexable(X, y, groups)
+
+    cv = check_cv(cv, y, classifier=is_classifier(estimator))
+    scorer = check_scoring(estimator, scoring=scoring)
+    random_state = check_random_state(random_state)
+
+    # We clone the estimator to make sure that all the folds are
+    # independent, and that it is pickle-able.
+    score = _permutation_test_score(
+        clone(estimator), X, y, groups, cv, scorer, fit_params=fit_params
+    )
+    permutation_scores = Parallel(n_jobs=n_jobs, verbose=verbose)(
+        delayed(_permutation_test_score)(
+            clone(estimator),
+            X,
+            _shuffle(y, groups, random_state),
+            groups,
+            cv,
+            scorer,
+            fit_params=fit_params,
+        )
+        for _ in range(n_permutations)
+    )
+    permutation_scores = np.array(permutation_scores)
+    pvalue = (np.sum(permutation_scores >= score) + 1.0) / (n_permutations + 1)
+    return score, permutation_scores, pvalue
+
+
+def _permutation_test_score(estimator, X, y, groups, cv, scorer, fit_params):
+    """Auxiliary function for permutation_test_score"""
+    # Adjust length of sample weights
+    fit_params = fit_params if fit_params is not None else {}
+    avg_score = []
+    for train, test in cv.split(X, y, groups):
+        X_train, y_train = _safe_split(estimator, X, y, train)
+        X_test, y_test = _safe_split(estimator, X, y, test, train)
+        fit_params = _check_fit_params(X, fit_params, train)
+        estimator.fit(X_train, y_train, **fit_params)
+        avg_score.append(scorer(estimator, X_test, y_test))
+    return np.mean(avg_score)
+
+
+def _shuffle(y, groups, random_state):
+    """Return a shuffled copy of y eventually shuffle among same groups."""
+    if groups is None:
+        indices = random_state.permutation(len(y))
+    else:
+        indices = np.arange(len(groups))
+        for group in np.unique(groups):
+            this_mask = groups == group
+            indices[this_mask] = random_state.permutation(indices[this_mask])
+    return _safe_indexing(y, indices)
+
+
+def learning_curve(
+    estimator,
+    X,
+    y,
+    *,
+    groups=None,
+    train_sizes=np.linspace(0.1, 1.0, 5),
+    cv=None,
+    scoring=None,
+    exploit_incremental_learning=False,
+    n_jobs=None,
+    pre_dispatch="all",
+    verbose=0,
+    shuffle=False,
+    random_state=None,
+    error_score=np.nan,
+    return_times=False,
+    fit_params=None,
+):
+    """Learning curve.
+
+    Determines cross-validated training and test scores for different training
+    set sizes.
+
+    A cross-validation generator splits the whole dataset k times in training
+    and test data. Subsets of the training set with varying sizes will be used
+    to train the estimator and a score for each training subset size and the
+    test set will be computed. Afterwards, the scores will be averaged over
+    all k runs for each training subset size.
+
+    Read more in the :ref:`User Guide <learning_curve>`.
+
+    Parameters
+    ----------
+    estimator : object type that implements the "fit" and "predict" methods
+        An object of that type which is cloned for each validation.
+
+    X : array-like of shape (n_samples, n_features)
+        Training vector, where `n_samples` is the number of samples and
+        `n_features` is the number of features.
+
+    y : array-like of shape (n_samples,) or (n_samples, n_outputs)
+        Target relative to X for classification or regression;
+        None for unsupervised learning.
+
+    groups : array-like of  shape (n_samples,), default=None
+        Group labels for the samples used while splitting the dataset into
+        train/test set. Only used in conjunction with a "Group" :term:`cv`
+        instance (e.g., :class:`GroupKFold`).
+
+    train_sizes : array-like of shape (n_ticks,), \
+            default=np.linspace(0.1, 1.0, 5)
+        Relative or absolute numbers of training examples that will be used to
+        generate the learning curve. If the dtype is float, it is regarded as a
+        fraction of the maximum size of the training set (that is determined
+        by the selected validation method), i.e. it has to be within (0, 1].
+        Otherwise it is interpreted as absolute sizes of the training sets.
+        Note that for classification the number of samples usually have to
+        be big enough to contain at least one sample from each class.
+
+    cv : int, cross-validation generator or an iterable, default=None
+        Determines the cross-validation splitting strategy.
+        Possible inputs for cv are:
+
+        - None, to use the default 5-fold cross validation,
+        - int, to specify the number of folds in a `(Stratified)KFold`,
+        - :term:`CV splitter`,
+        - An iterable yielding (train, test) splits as arrays of indices.
+
+        For int/None inputs, if the estimator is a classifier and ``y`` is
+        either binary or multiclass, :class:`StratifiedKFold` is used. In all
+        other cases, :class:`KFold` is used. These splitters are instantiated
+        with `shuffle=False` so the splits will be the same across calls.
+
+        Refer :ref:`User Guide <cross_validation>` for the various
+        cross-validation strategies that can be used here.
+
+        .. versionchanged:: 0.22
+            ``cv`` default value if None changed from 3-fold to 5-fold.
+
+    scoring : str or callable, default=None
+        A str (see model evaluation documentation) or
+        a scorer callable object / function with signature
+        ``scorer(estimator, X, y)``.
+
+    exploit_incremental_learning : bool, default=False
+        If the estimator supports incremental learning, this will be
+        used to speed up fitting for different training set sizes.
+
+    n_jobs : int, default=None
+        Number of jobs to run in parallel. Training the estimator and computing
+        the score are parallelized over the different training and test sets.
+        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
+        ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
+        for more details.
+
+    pre_dispatch : int or str, default='all'
+        Number of predispatched jobs for parallel execution (default is
+        all). The option can reduce the allocated memory. The str can
+        be an expression like '2*n_jobs'.
+
+    verbose : int, default=0
+        Controls the verbosity: the higher, the more messages.
+
+    shuffle : bool, default=False
+        Whether to shuffle training data before taking prefixes of it
+        based on``train_sizes``.
+
+    random_state : int, RandomState instance or None, default=None
+        Used when ``shuffle`` is True. Pass an int for reproducible
+        output across multiple function calls.
+        See :term:`Glossary <random_state>`.
+
+    error_score : 'raise' or numeric, default=np.nan
+        Value to assign to the score if an error occurs in estimator fitting.
+        If set to 'raise', the error is raised.
+        If a numeric value is given, FitFailedWarning is raised.
+
+        .. versionadded:: 0.20
+
+    return_times : bool, default=False
+        Whether to return the fit and score times.
+
+    fit_params : dict, default=None
+        Parameters to pass to the fit method of the estimator.
+
+        .. versionadded:: 0.24
+
+    Returns
+    -------
+    train_sizes_abs : array of shape (n_unique_ticks,)
+        Numbers of training examples that has been used to generate the
+        learning curve. Note that the number of ticks might be less
+        than n_ticks because duplicate entries will be removed.
+
+    train_scores : array of shape (n_ticks, n_cv_folds)
+        Scores on training sets.
+
+    test_scores : array of shape (n_ticks, n_cv_folds)
+        Scores on test set.
+
+    fit_times : array of shape (n_ticks, n_cv_folds)
+        Times spent for fitting in seconds. Only present if ``return_times``
+        is True.
+
+    score_times : array of shape (n_ticks, n_cv_folds)
+        Times spent for scoring in seconds. Only present if ``return_times``
+        is True.
+
+    Notes
+    -----
+    See :ref:`examples/model_selection/plot_learning_curve.py
+    <sphx_glr_auto_examples_model_selection_plot_learning_curve.py>`
+    """
+    if exploit_incremental_learning and not hasattr(estimator, "partial_fit"):
+        raise ValueError(
+            "An estimator must support the partial_fit interface "
+            "to exploit incremental learning"
+        )
+    X, y, groups = indexable(X, y, groups)
+
+    cv = check_cv(cv, y, classifier=is_classifier(estimator))
+    # Store it as list as we will be iterating over the list multiple times
+    cv_iter = list(cv.split(X, y, groups))
+
+    scorer = check_scoring(estimator, scoring=scoring)
+
+    n_max_training_samples = len(cv_iter[0][0])
+    # Because the lengths of folds can be significantly different, it is
+    # not guaranteed that we use all of the available training data when we
+    # use the first 'n_max_training_samples' samples.
+    train_sizes_abs = _translate_train_sizes(train_sizes, n_max_training_samples)
+    n_unique_ticks = train_sizes_abs.shape[0]
+    if verbose > 0:
+        print("[learning_curve] Training set sizes: " + str(train_sizes_abs))
+
+    parallel = Parallel(n_jobs=n_jobs, pre_dispatch=pre_dispatch, verbose=verbose)
+
+    if shuffle:
+        rng = check_random_state(random_state)
+        cv_iter = ((rng.permutation(train), test) for train, test in cv_iter)
+
+    if exploit_incremental_learning:
+        classes = np.unique(y) if is_classifier(estimator) else None
+        out = parallel(
+            delayed(_incremental_fit_estimator)(
+                clone(estimator),
+                X,
+                y,
+                classes,
+                train,
+                test,
+                train_sizes_abs,
+                scorer,
+                verbose,
+                return_times,
+                error_score=error_score,
+                fit_params=fit_params,
+            )
+            for train, test in cv_iter
+        )
+        out = np.asarray(out).transpose((2, 1, 0))
+    else:
+        train_test_proportions = []
+        for train, test in cv_iter:
+            for n_train_samples in train_sizes_abs:
+                train_test_proportions.append((train[:n_train_samples], test))
+
+        results = parallel(
+            delayed(_fit_and_score)(
+                clone(estimator),
+                X,
+                y,
+                scorer,
+                train,
+                test,
+                verbose,
+                parameters=None,
+                fit_params=fit_params,
+                return_train_score=True,
+                error_score=error_score,
+                return_times=return_times,
+            )
+            for train, test in train_test_proportions
+        )
+        results = _aggregate_score_dicts(results)
+        train_scores = results["train_scores"].reshape(-1, n_unique_ticks).T
+        test_scores = results["test_scores"].reshape(-1, n_unique_ticks).T
+        out = [train_scores, test_scores]
+
+        if return_times:
+            fit_times = results["fit_time"].reshape(-1, n_unique_ticks).T
+            score_times = results["score_time"].reshape(-1, n_unique_ticks).T
+            out.extend([fit_times, score_times])
+
+    ret = train_sizes_abs, out[0], out[1]
+
+    if return_times:
+        ret = ret + (out[2], out[3])
+
+    return ret
+
+
+def _translate_train_sizes(train_sizes, n_max_training_samples):
+    """Determine absolute sizes of training subsets and validate 'train_sizes'.
+
+    Examples:
+        _translate_train_sizes([0.5, 1.0], 10) -> [5, 10]
+        _translate_train_sizes([5, 10], 10) -> [5, 10]
+
+    Parameters
+    ----------
+    train_sizes : array-like of shape (n_ticks,)
+        Numbers of training examples that will be used to generate the
+        learning curve. If the dtype is float, it is regarded as a
+        fraction of 'n_max_training_samples', i.e. it has to be within (0, 1].
+
+    n_max_training_samples : int
+        Maximum number of training samples (upper bound of 'train_sizes').
+
+    Returns
+    -------
+    train_sizes_abs : array of shape (n_unique_ticks,)
+        Numbers of training examples that will be used to generate the
+        learning curve. Note that the number of ticks might be less
+        than n_ticks because duplicate entries will be removed.
+    """
+    train_sizes_abs = np.asarray(train_sizes)
+    n_ticks = train_sizes_abs.shape[0]
+    n_min_required_samples = np.min(train_sizes_abs)
+    n_max_required_samples = np.max(train_sizes_abs)
+    if np.issubdtype(train_sizes_abs.dtype, np.floating):
+        if n_min_required_samples <= 0.0 or n_max_required_samples > 1.0:
+            raise ValueError(
+                "train_sizes has been interpreted as fractions "
+                "of the maximum number of training samples and "
+                "must be within (0, 1], but is within [%f, %f]."
+                % (n_min_required_samples, n_max_required_samples)
+            )
+        train_sizes_abs = (train_sizes_abs * n_max_training_samples).astype(
+            dtype=int, copy=False
+        )
+        train_sizes_abs = np.clip(train_sizes_abs, 1, n_max_training_samples)
+    else:
+        if (
+            n_min_required_samples <= 0
+            or n_max_required_samples > n_max_training_samples
+        ):
+            raise ValueError(
+                "train_sizes has been interpreted as absolute "
+                "numbers of training samples and must be within "
+                "(0, %d], but is within [%d, %d]."
+                % (
+                    n_max_training_samples,
+                    n_min_required_samples,
+                    n_max_required_samples,
+                )
+            )
+
+    train_sizes_abs = np.unique(train_sizes_abs)
+    if n_ticks > train_sizes_abs.shape[0]:
+        warnings.warn(
+            "Removed duplicate entries from 'train_sizes'. Number "
+            "of ticks will be less than the size of "
+            "'train_sizes': %d instead of %d." % (train_sizes_abs.shape[0], n_ticks),
+            RuntimeWarning,
+        )
+
+    return train_sizes_abs
+
+
+def _incremental_fit_estimator(
+    estimator,
+    X,
+    y,
+    classes,
+    train,
+    test,
+    train_sizes,
+    scorer,
+    verbose,
+    return_times,
+    error_score,
+    fit_params,
+):
+    """Train estimator on training subsets incrementally and compute scores."""
+    train_scores, test_scores, fit_times, score_times = [], [], [], []
+    partitions = zip(train_sizes, np.split(train, train_sizes)[:-1])
+    if fit_params is None:
+        fit_params = {}
+    for n_train_samples, partial_train in partitions:
+        train_subset = train[:n_train_samples]
+        X_train, y_train = _safe_split(estimator, X, y, train_subset)
+        X_partial_train, y_partial_train = _safe_split(estimator, X, y, partial_train)
+        X_test, y_test = _safe_split(estimator, X, y, test, train_subset)
+        start_fit = time.time()
+        if y_partial_train is None:
+            estimator.partial_fit(X_partial_train, classes=classes, **fit_params)
+        else:
+            estimator.partial_fit(
+                X_partial_train, y_partial_train, classes=classes, **fit_params
+            )
+        fit_time = time.time() - start_fit
+        fit_times.append(fit_time)
+
+        start_score = time.time()
+
+        test_scores.append(_score(estimator, X_test, y_test, scorer, error_score))
+        train_scores.append(_score(estimator, X_train, y_train, scorer, error_score))
+
+        score_time = time.time() - start_score
+        score_times.append(score_time)
+
+    ret = (
+        (train_scores, test_scores, fit_times, score_times)
+        if return_times
+        else (train_scores, test_scores)
+    )
+
+    return np.array(ret).T
+
+
+def validation_curve(
+    estimator,
+    X,
+    y,
+    *,
+    param_name,
+    param_range,
+    groups=None,
+    cv=None,
+    scoring=None,
+    n_jobs=None,
+    pre_dispatch="all",
+    verbose=0,
+    error_score=np.nan,
+    fit_params=None,
+):
+    """Validation curve.
+
+    Determine training and test scores for varying parameter values.
+
+    Compute scores for an estimator with different values of a specified
+    parameter. This is similar to grid search with one parameter. However, this
+    will also compute training scores and is merely a utility for plotting the
+    results.
+
+    Read more in the :ref:`User Guide <validation_curve>`.
+
+    Parameters
+    ----------
+    estimator : object type that implements the "fit" and "predict" methods
+        An object of that type which is cloned for each validation.
+
+    X : array-like of shape (n_samples, n_features)
+        Training vector, where `n_samples` is the number of samples and
+        `n_features` is the number of features.
+
+    y : array-like of shape (n_samples,) or (n_samples, n_outputs) or None
+        Target relative to X for classification or regression;
+        None for unsupervised learning.
+
+    param_name : str
+        Name of the parameter that will be varied.
+
+    param_range : array-like of shape (n_values,)
+        The values of the parameter that will be evaluated.
+
+    groups : array-like of shape (n_samples,), default=None
+        Group labels for the samples used while splitting the dataset into
+        train/test set. Only used in conjunction with a "Group" :term:`cv`
+        instance (e.g., :class:`GroupKFold`).
+
+    cv : int, cross-validation generator or an iterable, default=None
+        Determines the cross-validation splitting strategy.
+        Possible inputs for cv are:
+
+        - None, to use the default 5-fold cross validation,
+        - int, to specify the number of folds in a `(Stratified)KFold`,
+        - :term:`CV splitter`,
+        - An iterable yielding (train, test) splits as arrays of indices.
+
+        For int/None inputs, if the estimator is a classifier and ``y`` is
+        either binary or multiclass, :class:`StratifiedKFold` is used. In all
+        other cases, :class:`KFold` is used. These splitters are instantiated
+        with `shuffle=False` so the splits will be the same across calls.
+
+        Refer :ref:`User Guide <cross_validation>` for the various
+        cross-validation strategies that can be used here.
+
+        .. versionchanged:: 0.22
+            ``cv`` default value if None changed from 3-fold to 5-fold.
+
+    scoring : str or callable, default=None
+        A str (see model evaluation documentation) or
+        a scorer callable object / function with signature
+        ``scorer(estimator, X, y)``.
+
+    n_jobs : int, default=None
+        Number of jobs to run in parallel. Training the estimator and computing
+        the score are parallelized over the combinations of each parameter
+        value and each cross-validation split.
+        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
+        ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
+        for more details.
+
+    pre_dispatch : int or str, default='all'
+        Number of predispatched jobs for parallel execution (default is
+        all). The option can reduce the allocated memory. The str can
+        be an expression like '2*n_jobs'.
+
+    verbose : int, default=0
+        Controls the verbosity: the higher, the more messages.
+
+    fit_params : dict, default=None
+        Parameters to pass to the fit method of the estimator.
+
+        .. versionadded:: 0.24
+
+    error_score : 'raise' or numeric, default=np.nan
+        Value to assign to the score if an error occurs in estimator fitting.
+        If set to 'raise', the error is raised.
+        If a numeric value is given, FitFailedWarning is raised.
+
+        .. versionadded:: 0.20
+
+    Returns
+    -------
+    train_scores : array of shape (n_ticks, n_cv_folds)
+        Scores on training sets.
+
+    test_scores : array of shape (n_ticks, n_cv_folds)
+        Scores on test set.
+
+    Notes
+    -----
+    See :ref:`sphx_glr_auto_examples_model_selection_plot_validation_curve.py`
+
+    """
+    X, y, groups = indexable(X, y, groups)
+
+    cv = check_cv(cv, y, classifier=is_classifier(estimator))
+    scorer = check_scoring(estimator, scoring=scoring)
+
+    parallel = Parallel(n_jobs=n_jobs, pre_dispatch=pre_dispatch, verbose=verbose)
+    results = parallel(
+        delayed(_fit_and_score)(
+            clone(estimator),
+            X,
+            y,
+            scorer,
+            train,
+            test,
+            verbose,
+            parameters={param_name: v},
+            fit_params=fit_params,
+            return_train_score=True,
+            error_score=error_score,
+        )
+        # NOTE do not change order of iteration to allow one time cv splitters
+        for train, test in cv.split(X, y, groups)
+        for v in param_range
+    )
+    n_params = len(param_range)
+
+    results = _aggregate_score_dicts(results)
+    train_scores = results["train_scores"].reshape(-1, n_params).T
+    test_scores = results["test_scores"].reshape(-1, n_params).T
+
+    return train_scores, test_scores
+
+
+def _aggregate_score_dicts(scores):
+    """Aggregate the list of dict to dict of np ndarray
+
+    The aggregated output of _aggregate_score_dicts will be a list of dict
+    of form [{'prec': 0.1, 'acc':1.0}, {'prec': 0.1, 'acc':1.0}, ...]
+    Convert it to a dict of array {'prec': np.array([0.1 ...]), ...}
+
+    Parameters
+    ----------
+
+    scores : list of dict
+        List of dicts of the scores for all scorers. This is a flat list,
+        assumed originally to be of row major order.
+
+    Example
+    -------
+
+    >>> scores = [{'a': 1, 'b':10}, {'a': 2, 'b':2}, {'a': 3, 'b':3},
+    ...           {'a': 10, 'b': 10}]                         # doctest: +SKIP
+    >>> _aggregate_score_dicts(scores)                        # doctest: +SKIP
+    {'a': array([1, 2, 3, 10]),
+     'b': array([10, 2, 3, 10])}
+    """
+    return {
+        key: np.asarray([score[key] for score in scores])
+        if isinstance(scores[0][key], numbers.Number)
+        else [score[key] for score in scores]
+        for key in scores[0]
+    }
diff --git a/frameworks/TunedRandomForest/exec.py b/frameworks/TunedRandomForest/exec.py
index badc79a7a..7c7a7dc15 100644
--- a/frameworks/TunedRandomForest/exec.py
+++ b/frameworks/TunedRandomForest/exec.py
@@ -8,21 +8,22 @@
 import os
 import statistics
 import tempfile as tmp
+from collections import defaultdict
+from typing import List
 
 os.environ['JOBLIB_TEMP_FOLDER'] = tmp.gettempdir()
 os.environ['OMP_NUM_THREADS'] = '1'
 os.environ['OPENBLAS_NUM_THREADS'] = '1'
 os.environ['MKL_NUM_THREADS'] = '1'
 
+import psutil
+import pandas as pd
 import sklearn
 from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
-from sklearn.pipeline import Pipeline
-from sklearn.impute import SimpleImputer
-from sklearn.model_selection import cross_val_score
-import stopit
 
 from frameworks.shared.callee import call_run, result
 from frameworks.shared.utils import Timer
+from custom_validate import cross_validate
 
 log = logging.getLogger(__name__)
 
@@ -33,14 +34,31 @@ def pick_values_uniform(start: int, end: int, length: int):
     return sorted(set([int(f) for f in uniform_floats]))
 
 
+def extrapolate_with_worst_case(values: List[float], n: int = 5) -> float:
+    """ Extrapolate the next value for `values`, based on the last `n` samples. """
+    n = min(len(values), n)
+    return values[-1] + max(v_next - v_prev for v_prev, v_next in zip(values[-n:], values[-n+1:]))
+
+
 def run(dataset, config):
     log.info(f"\n**** Tuned Random Forest [sklearn v{sklearn.__version__}] ****\n")
 
     is_classification = config.type == 'classification'
 
-    training_params = {k: v for k, v in config.framework_params.items() if not k.startswith('_')}
-    tuning_params = config.framework_params.get('_tuning', training_params)
+    training_params = {
+        k: v for k, v in config.framework_params.items()
+        if not (k.startswith('_') or k == "n_estimators")
+    }
+    if "max_features" in training_params:
+        raise ValueError("`max_features` may not be specified for Tuned Random Forest.")
+
     n_jobs = config.framework_params.get('_n_jobs', config.cores)  # useful to disable multicore, regardless of the dataset config
+    k_folds = config.framework_params.get('_k_folds', 5)
+    step_size = config.framework_params.get('_step_size', 10)
+    memory_margin = config.framework_params.get('_memory_margin', 0.9)
+    final_forest_size = config.framework_params.get('n_estimators', 2000)
+    tune_fraction = config.framework_params.get('_tune_fraction', 0.9)
+    tune_time = config.max_runtime_seconds * tune_fraction
 
     X_train, X_test = dataset.train.X, dataset.test.X
     y_train, y_test = dataset.train.y, dataset.test.y
@@ -64,60 +82,130 @@ def run(dataset, config):
     default_value = max(1, int(math.sqrt(n_features)))
     below_default = pick_values_uniform(start=1, end=default_value, length=5+1)[:-1]   # 5 below
     above_default = pick_values_uniform(start=default_value, end=n_features, length=10+1 - len(below_default))[1:]  # 5 above
-    # Mix up the order of `max_features` to try, so that a fair range is tried even if we have too little time
-    # to try all possible values. Order: [sqrt(p), 1, p, random order for remaining values]
-    # max_features_to_try = below_default[1:] + above_default[:-1]
-    # max_features_values = ([default_value, 1, n_features]
-    #                        + random.sample(max_features_to_try, k=len(max_features_to_try)))
+
     max_features_values = [default_value] + below_default + above_default
     # Define up to how much of total time we spend 'optimizing' `max_features`.
     # (the remainder if used for fitting the final model).
-    safety_factor = 0.85
-    with stopit.ThreadingTimeout(seconds=int(config.max_runtime_seconds * safety_factor)):
-        log.info("Evaluating multiple values for `max_features`: %s.", max_features_values)
-        max_feature_scores = []
-        tuning_durations = []
-        for i, max_features_value in enumerate(max_features_values):
-            log.info("[{:2d}/{:2d}] Evaluating max_features={}"
-                     .format(i + 1, len(max_features_values), max_features_value))
-            imputation = SimpleImputer()
-            random_forest = estimator(n_jobs=n_jobs,
-                                      random_state=config.seed,
-                                      max_features=max_features_value,
-                                      **tuning_params)
-            pipeline = Pipeline(steps=[
-                ('preprocessing', imputation),
-                ('learning', random_forest)
-            ])
-            with Timer() as cv_scoring:
+
+    log.info("Evaluating multiple values for `max_features`: %s.", max_features_values)
+    max_feature_scores = defaultdict(list)
+    tuning_durations = defaultdict(list)
+    memory_usage_by = defaultdict(list)
+    this_process = psutil.Process(os.getpid())
+    last_initial_fit_time = 0
+
+    with Timer() as training:
+        while max_features_values:
+            time_left = tune_time - training.duration
+            time_per_value = time_left / len(max_features_values)
+            if time_per_value < last_initial_fit_time:
+                log.info("Expect to exceed time constraints on next first fit, "
+                         f"budget is {time_per_value}s and last first fit took "
+                         f"{last_initial_fit_time}s.")
+                log.info(f"Did not try max_features={max_features_values}.")
+                break
+
+            value = max_features_values.pop(0)
+            log.info(f"Evaluating max_features={value} in {time_per_value} seconds.")
+
+            random_forests = [
+                estimator(
+                    n_jobs=n_jobs,
+                    random_state=config.seed,
+                    max_features=value,
+                    n_estimators=step_size,
+                    warm_start=True,
+                    **training_params
+                ) for _ in range(k_folds)
+            ]
+
+            training_times = [training.duration]
+            memory_usage = [this_process.memory_info()[0] / (2 ** 20)]
+
+            while True:
                 try:
-                    scores = cross_val_score(estimator=pipeline,
-                                             X=X_train,
-                                             y=y_train,
-                                             scoring=metric,
-                                             error_score='raise',
-                                             cv=5)
-                    max_feature_scores.append((statistics.mean(scores), max_features_value))
-                except stopit.utils.TimeoutException as toe:
-                    log.error("Failed CV scoring for max_features=%s : Timeout", max_features_value)
-                    tuning_durations.append((max_features_value, cv_scoring.duration))
-                    raise toe
+                    cv_result = cross_validate(
+                        estimators=random_forests,
+                        X=X_train,
+                        y=y_train,
+                        scoring=metric,
+                        error_score='raise',
+                        cv=k_folds,
+                        return_estimator=True,
+                    )
+                    max_feature_scores[value].append(statistics.mean(cv_result["test_score"]))
                 except Exception as e:
-                    log.error("Failed CV scoring for max_features=%s :\n%s", max_features_value, e)
+                    log.error("Failed CV scoring for max_features=%s :\n%s", value, e)
                     log.debug("Exception:", exc_info=True)
-                    max_feature_scores.append((math.nan, max_features_value))
-            tuning_durations.append((max_features_value, cv_scoring.duration))
-
-    log.info("Tuning scores:\n%s", sorted(max_feature_scores))
-    log.info("Tuning durations:\n%s", sorted(tuning_durations))
-    _, best_max_features_value = max(max_feature_scores) if len(max_feature_scores) > 0 else (math.nan, 'auto')
-    log.info("Training final model with `max_features={}`.".format(best_max_features_value))
-    rf = estimator(n_jobs=n_jobs,
-                   random_state=config.seed,
-                   max_features=best_max_features_value,
-                   **training_params)
-    with Timer() as training:
-        rf.fit(X_train, y_train)
+                    max_feature_scores[value].append(math.nan)
+                    break
+                finally:
+                    training_times.append(training.duration)
+                    memory_usage.append(this_process.memory_info()[0] / (2 ** 20))
+
+                if random_forests[0].n_estimators >= final_forest_size:
+                    log.info("Stop training because desired forest size has been reached.")
+                    break
+                if extrapolate_with_worst_case(training_times) - training_times[0] >= time_per_value:
+                    log.info(f"Stop training after fitting {random_forests[0].n_estimators} trees because it expects to exceed its time budget.")
+                    break
+                elif extrapolate_with_worst_case(memory_usage) >= config.max_mem_size_mb * memory_margin:
+                    log.info(f"Stop training after fitting {random_forests[0].n_estimators} trees because it expects to exceed its memory budget.")
+                    break
+
+                for rf in random_forests:
+                    rf.n_estimators += step_size
+
+            tuning_durations[value] = training_times
+            memory_usage_by[value] = memory_usage
+            last_initial_fit_time = training_times[1] - training_times[0]
+
+        summary = []
+        for max_feature, scores in max_feature_scores.items():
+            # remember: duration and memory have a record before the first fit too
+            stats_for_max_feature = zip(scores, tuning_durations[max_feature][1:], memory_usage_by[max_feature][1:])
+            for i, (score, duration, memory) in enumerate(stats_for_max_feature, start=1):
+                summary.append({
+                    "max_features": max_feature,
+                    "trees": (step_size) * i,
+                    metric: score,
+                    "duration total (s)": duration,
+                    "memory total (mb)": memory,
+                    "duration delta (s)": duration - tuning_durations[max_feature][i - 1],
+                    "memory delta (mb)": memory - memory_usage_by[max_feature][i - 1],
+                })
+        summary = pd.DataFrame.from_records(summary)
+        log.debug("Report of tuning iterations:\n %s \n" % summary.to_string())
+        log.info("Report final tunings:\n %s \n" % summary.groupby("max_features").last().sort_values(by="duration total (s)").to_string())
+
+        # Iteratively fit a random forest with our "optimal" `max_features`
+        _, best_value = max((scores[-1], value) for value, scores in max_feature_scores.items())
+        log.info("Training final model with `max_features=%s`.", best_value)
+        rf = estimator(n_jobs=n_jobs,
+                       random_state=config.seed,
+                       max_features=best_value,
+                       warm_start=True,
+                       n_estimators=step_size,
+                       **training_params)
+
+        training_times = [training.duration]
+        memory_usage = [this_process.memory_info()[0] / (2 ** 20)]
+        while True:
+            rf.fit(X_train, y_train)
+            training_times.append(training.duration)
+            memory_usage.append(this_process.memory_info()[0] / (2 ** 20))
+            if rf.n_estimators == final_forest_size:
+                log.info("Stop training because desired forest size has been reached.")
+                break
+            if extrapolate_with_worst_case(training_times) >= config.max_runtime_seconds:
+                log.info("Stop training because it expects to exceed its time budget.")
+                break
+            elif extrapolate_with_worst_case(
+                    memory_usage) >= config.max_mem_size_mb * memory_margin:
+                log.info(
+                    "Stop training because it expects to exceed its memory budget.")
+                break
+            rf.n_estimators += step_size
 
     with Timer() as predict:
         predictions = rf.predict(X_test)
@@ -130,7 +218,7 @@ def run(dataset, config):
         probabilities=probabilities,
         target_is_encoded=is_classification,
         models_count=len(rf),
-        training_duration=training.duration+sum(map(lambda t: t[1], tuning_durations)),
+        training_duration=training.duration,
         predict_duration=predict.duration
     )
 
diff --git a/frameworks/TunedRandomForest/requirements.txt b/frameworks/TunedRandomForest/requirements.txt
index 33bb952d1..fb6c7ed7e 100644
--- a/frameworks/TunedRandomForest/requirements.txt
+++ b/frameworks/TunedRandomForest/requirements.txt
@@ -1,1 +1,1 @@
-stopit==1.1.2
+pandas
diff --git a/frameworks/TunedRandomForest/setup.sh b/frameworks/TunedRandomForest/setup.sh
index d547c3864..89c2d2d0d 100755
--- a/frameworks/TunedRandomForest/setup.sh
+++ b/frameworks/TunedRandomForest/setup.sh
@@ -4,7 +4,7 @@ VERSION=${1:-"stable"}
 REPO=${2:-"https://github.com/scikit-learn/scikit-learn.git"}
 PKG=${3:-"scikit-learn"}
 if [[ "$VERSION" == "latest" ]]; then
-    VERSION="master"
+    VERSION="main"
 fi
 
 # create venv
diff --git a/frameworks/hyperoptsklearn/__init__.py b/frameworks/hyperoptsklearn/__init__.py
index 97ecd49b2..7edcb225b 100644
--- a/frameworks/hyperoptsklearn/__init__.py
+++ b/frameworks/hyperoptsklearn/__init__.py
@@ -25,14 +25,13 @@ def run(dataset: Dataset, config: TaskConfig):
     )
 
     def process_results(results):
-        if results.probabilities is not None and not results.probabilities.shape:  # numpy load always return an array
-            prob_format = results.probabilities.item()
-            if prob_format == "predictions":
-                target_values_enc = dataset.target.label_encoder.transform(dataset.target.values)
-                results.probabilities = Encoder('one-hot', target=False, encoded_type=float).fit(target_values_enc).transform(results.predictions)
-            else:
-                raise ValueError(f"Unknown probabilities format: {prob_format}")
-        return results
+        if isinstance(results.probabilities, str) and results.probabilities == "predictions":
+            target_values_enc = dataset.target.label_encoder.transform(dataset.target.values)
+            results.probabilities = Encoder('one-hot', target=False, encoded_type=float).fit(target_values_enc).transform(results.predictions)
+        is_numpy_like = hasattr(results.probabilities, "shape") and results.probabilities.shape
+        if results.probabilities is None or is_numpy_like:
+            return results
+        raise ValueError(f"Unknown probabilities format: {results.probabilities}")
 
     return run_in_venv(__file__, "exec.py",
                        input_data=data, dataset=dataset, config=config,
diff --git a/frameworks/mlr3automl/setup.sh b/frameworks/mlr3automl/setup.sh
index 05fa7593e..cd5225c04 100755
--- a/frameworks/mlr3automl/setup.sh
+++ b/frameworks/mlr3automl/setup.sh
@@ -41,10 +41,11 @@ fi
 LIB="${HERE}/lib/"
 mkdir "${LIB}"
 
-Rscript -e 'options(install.packages.check.source="no"); install.packages(c("mlr3", "mlr3pipelines", "mlr3misc", "mlr3oml", "mlr3hyperband", "mlr3tuning", "paradox"), repos="https://cloud.r-project.org/", lib="'"${LIB}"'")'
-Rscript -e 'options(install.packages.check.source="no"); install.packages(c("remotes", "checkmate", "R6", "xgboost", "ranger", "LiblineaR", "emoa", "e1071", "glmnet"), repos="https://cloud.r-project.org/", lib="'"${LIB}"'")'
-Rscript -e '.libPaths("'"${LIB}"'"); remotes::install_github("'"${MLR_REPO}"'/mlr3extralearners", lib="'"${LIB}"'")'
-Rscript -e '.libPaths("'"${LIB}"'"); remotes::install_github("'"${REPO}"'", ref="'"${VERSION}"'", lib="'"${LIB}"'")'
+Rscript -e 'options(install.packages.check.source="no"); install.packages("remotes", repos="https://cloud.r-project.org/", lib="'"${LIB}"'")'
+Rscript -e 'options(install.packages.check.source="no"); install.packages(c("xgboost", "ranger", "LiblineaR", "emoa", "e1071", "glmnet"), repos="https://cloud.r-project.org/", lib="'"${LIB}"'")'
+Rscript -e '.libPaths("'"${LIB}"'"); remotes::install_github("mlr-org/mlr3learners", dependencies=TRUE)'
+Rscript -e '.libPaths("'"${LIB}"'"); remotes::install_github("mlr-org/mlr3extralearners")'
+Rscript -e '.libPaths("'"${LIB}"'"); remotes::install_github("'"${REPO}"'", lib="'"${LIB}"'", dependencies=TRUE)'
 
 OFFICIAL_VERSION=$(Rscript -e '.libPaths("'"${LIB}"'"); packageVersion("mlr3automl")' | awk '{print $2}' | sed "s/[‘’]//g")
 echo "${OFFICIAL_VERSION}#${VERSION:0:7}" >> "${HERE}/.setup/installed"
diff --git a/frameworks/oboe/__init__.py b/frameworks/oboe/__init__.py
index d21694d78..3e6f9bfa8 100644
--- a/frameworks/oboe/__init__.py
+++ b/frameworks/oboe/__init__.py
@@ -25,14 +25,13 @@ def run(dataset: Dataset, config: TaskConfig):
     )
 
     def process_results(results):
-        if results.probabilities is not None and not results.probabilities.shape:  # numpy load always return an array
-            prob_format = results.probabilities.item()
-            if prob_format == "predictions":
-                target_values_enc = dataset.target.label_encoder.transform(dataset.target.values)
-                results.probabilities = Encoder('one-hot', target=False, encoded_type=float).fit(target_values_enc).transform(results.predictions)
-            else:
-                raise ValueError(f"Unknown probabilities format: {prob_format}")
-        return results
+        if isinstance(results.probabilities, str) and results.probabilities == "predictions":
+            target_values_enc = dataset.target.label_encoder.transform(dataset.target.values)
+            results.probabilities = Encoder('one-hot', target=False, encoded_type=float).fit(target_values_enc).transform(results.predictions)
+        is_numpy_like = hasattr(results.probabilities, "shape") and results.probabilities.shape
+        if results.probabilities is None or is_numpy_like:
+            return results
+        raise ValueError(f"Unknown probabilities format: {results.probabilities}")
 
     return run_in_venv(__file__, "exec.py",
                        input_data=data, dataset=dataset, config=config,
diff --git a/frameworks/shared/caller.py b/frameworks/shared/caller.py
index ee60b42d4..b4c9b0210 100644
--- a/frameworks/shared/caller.py
+++ b/frameworks/shared/caller.py
@@ -93,7 +93,7 @@ def run_in_venv(caller_file, script_file: str, *args,
     ser_config = options['serialization']
     env = options['env'] or ns()
 
-    with TemporaryDirectory() as tmpdir:
+    with TemporaryDirectory(prefix='amlb_', suffix='_xproc') as tmpdir:
 
         ds = _make_input_dataset(input_data, dataset, tmpdir, serialization=ser_config)
 
diff --git a/resources/constraints.yaml b/resources/constraints.yaml
index ab2b3c831..369328efb 100644
--- a/resources/constraints.yaml
+++ b/resources/constraints.yaml
@@ -27,3 +27,17 @@ test:
   max_runtime_seconds: 28800
   cores: 8
   min_vol_size_mb: 131072
+
+1h8c_gp3:
+  folds: 10
+  max_runtime_seconds: 3600
+  cores: 8
+  min_vol_size_mb: 100000
+  ec2_volume_type: gp3
+
+4h8c_gp3:
+  folds: 10
+  max_runtime_seconds: 14400
+  cores: 8
+  min_vol_size_mb: 100000
+  ec2_volume_type: gp3
diff --git a/resources/frameworks_2021Q3.yaml b/resources/frameworks_2021Q3.yaml
index e36e361ba..cea36f5ba 100644
--- a/resources/frameworks_2021Q3.yaml
+++ b/resources/frameworks_2021Q3.yaml
@@ -15,7 +15,7 @@ AutoGluon_benchmark:
     presets: best_quality
 
 autosklearn:
-  version: '0.12.7'
+  version: '0.14.0'
 
 autosklearn2:
   extends: autosklearn
@@ -32,7 +32,13 @@ flaml:
   version: '0.6.2'
 
 GAMA:
-  version: '21.0.0'
+  abstract: true
+  version: '21.0.1'
+
+GAMA_benchmark:
+  extends: GAMA
+  params:
+    goal: performance
 
 H2OAutoML:
   version: '3.34.0.1'
@@ -50,7 +56,7 @@ mljarsupervised_benchmark:
   params:
     mode: Compete
 
-MLNET:
+MLNet:
   version: '16.5.26'
 
 MLPlan:
@@ -63,7 +69,7 @@ MLPlanWEKA:
     _backend: weka
 
 mlr3automl:
-  version: '#dc4031fac57e1d1caea0b48c8526528852572e7c'
+  version: '#f667900292e3ded64bb419285c71cd5d1d2c4301'
   project: https://github.com/a-hanf/mlr3automl
 
 TPOT:
diff --git a/runbenchmark.py b/runbenchmark.py
index 536dc1a1a..7ff8c6065 100644
--- a/runbenchmark.py
+++ b/runbenchmark.py
@@ -8,7 +8,7 @@
 # prevent asap other modules from defining the root logger using basicConfig
 import amlb.logger
 
-from openml.config import get_cache_directory
+from openml.config import cache_directory
 
 import amlb
 from amlb.utils import Namespace as ns, config_load, datetime_iso, str2bool, str_sanitize, zip_path
@@ -16,7 +16,7 @@
 
 
 default_dirs = ns(
-    input_dir=get_cache_directory(),
+    input_dir=cache_directory,
     output_dir="./results",
     user_dir="~/.config/automlbenchmark",
     root_dir=os.path.dirname(__file__)


Test patch:
diff --git a/tests/unit/amlb/job/dummy.py b/tests/unit/amlb/job/dummy.py
index 74f8a67b5..e34cd0ebe 100644
--- a/tests/unit/amlb/job/dummy.py
+++ b/tests/unit/amlb/job/dummy.py
@@ -11,6 +11,7 @@ def __init__(self, name="", timeout_secs=None, priority=None, raise_exceptions=F
         self._verbose = verbose
         self._duration_secs = duration_secs
         self._result = result
+        self.ext = {}
         super().__init__(name, timeout_secs, priority, raise_exceptions)
 
     def _add_step(self, step):
diff --git a/tests/unit/amlb/job/test_MultiThreadingJobRunner.py b/tests/unit/amlb/job/test_MultiThreadingJobRunner.py
index 65533afa9..543c2285c 100644
--- a/tests/unit/amlb/job/test_MultiThreadingJobRunner.py
+++ b/tests/unit/amlb/job/test_MultiThreadingJobRunner.py
@@ -1,10 +1,11 @@
 import functools as ft
 import random
+import time
 from unittest.mock import patch, DEFAULT
 
 import pytest
 
-from amlb.job import MultiThreadingJobRunner
+from amlb.job import MultiThreadingJobRunner, State
 from amlb.utils import Timeout
 
 from dummy import DummyJob
@@ -84,7 +85,7 @@ def test_reschedule_job_default():
 
     def _run(self, mock, ori):
         if mock.call_count < 3:
-            runner.reschedule(self)
+            self.reschedule()
             return
         return ori()
 
@@ -125,7 +126,7 @@ def test_reschedule_job_enforce_job_priority():
                                      queueing_strategy=MultiThreadingJobRunner.QueueingStrategy.enforce_job_priority)
     def _run(self, mock, ori):
         if mock.call_count < 3:
-            runner.reschedule(self)
+            self.reschedule()
             return
         return ori()
         # return DEFAULT  # ensures that the wrapped function is called after the side effect
@@ -171,7 +172,7 @@ def test_reschedule_job_high_parallelism():
 
     def _run(self, mock, ori):
         if mock.call_count < 3:
-            runner.reschedule(self)
+            self.reschedule()
             return
         return ori()
 
@@ -199,3 +200,57 @@ def _run(self, mock, ori):
     assert len(list(filter(lambda s: s == 'stopping', rescheduled_job_steps))) == 1 * len(rescheduled_job_names)
     assert len(list(filter(lambda s: s == 'stopped', rescheduled_job_steps))) == 1 * len(rescheduled_job_names)
 
+
+@pytest.mark.slow
+@pytest.mark.stress
+def test_reschedule_does_not_lead_to_runner_cancellation_on_high_parallelism():
+    seq_steps = []
+    attempts = 5
+    n_jobs = 200
+    jobs = [DummyJob(name=f"job_{i}", duration_secs=random.randint(10, 20)/10, result=i,
+                     steps=seq_steps, verbose=True)
+            for i in range(n_jobs)]
+
+    def _run(self, ori=None):
+        attempt = self.ext['attempt'] = self.ext.get('attempt', 0) + 1
+        if attempt < attempts:
+            self.reschedule()
+            time.sleep(self._duration_secs/10)
+            return
+        return ori()
+
+    def _on_state(self, state, ori=None):
+        ori(state)
+        if state is State.rescheduling:
+            time.sleep(random.randint(1, 10)/10)
+        elif state is State.stopped:
+            print(f"#results = {len(runner.results)}")
+
+    rescheduled_jobs = [j for i, j in enumerate(jobs) if i % 5 == 0]
+    for job in rescheduled_jobs:
+        job._run = ft.partialmethod(_run, ori=job._run).__get__(job)
+    for job in jobs:
+        job._on_state = ft.partialmethod(_on_state, ori=job._on_state).__get__(job)
+
+    runner = MultiThreadingJobRunner(jobs, parallel_jobs=50, delay_secs=0.01,
+                                     queueing_strategy=MultiThreadingJobRunner.QueueingStrategy.enforce_job_priority)
+    results = runner.start()
+    # print(results)
+
+    rescheduled_job_names = [j.name for j in rescheduled_jobs]
+    normal_job_steps = [s for n, s in seq_steps if n not in rescheduled_job_names]
+    rescheduled_job_steps = [s for n, s in seq_steps if n in rescheduled_job_names]
+
+    assert len(results) == n_jobs
+    assert len(list(filter(lambda s: s == 'cancelling', rescheduled_job_steps))) == 0
+    for state in ['created', 'starting', 'running', 'completing', 'stopping', 'stopped']:
+        assert len(list(filter(lambda s: s == state, normal_job_steps))) == n_jobs - len(rescheduled_job_names)
+
+    assert len(list(filter(lambda s: s == 'created', rescheduled_job_steps))) == 1 * len(rescheduled_job_names)
+    assert len(list(filter(lambda s: s == 'starting', rescheduled_job_steps))) == attempts * len(rescheduled_job_names)
+    assert len(list(filter(lambda s: s == 'running', rescheduled_job_steps))) == attempts * len(rescheduled_job_names)
+    assert len(list(filter(lambda s: s == 'rescheduling', rescheduled_job_steps))) == (attempts-1) * len(rescheduled_job_names)
+    assert len(list(filter(lambda s: s == 'completing', rescheduled_job_steps))) == 1 * len(rescheduled_job_names)
+    assert len(list(filter(lambda s: s == 'stopping', rescheduled_job_steps))) == 1 * len(rescheduled_job_names)
+    assert len(list(filter(lambda s: s == 'stopped', rescheduled_job_steps))) == 1 * len(rescheduled_job_names)
+
diff --git a/tests/unit/amlb/job/test_SimpleJobRunner.py b/tests/unit/amlb/job/test_SimpleJobRunner.py
index 29bc849cf..331335e12 100644
--- a/tests/unit/amlb/job/test_SimpleJobRunner.py
+++ b/tests/unit/amlb/job/test_SimpleJobRunner.py
@@ -77,7 +77,7 @@ def test_reschedule_job():
 
     def _run(self, mock, ori):
         if mock.call_count < 3:
-            runner.reschedule(self)
+            self.reschedule()
             return
         return ori()