Show progress bar while fitting to training data (#1606)

* Show progress bar while fitting to training data

* Minor fixes for progress bar

* Revert accidental changes to requirements.txt

* Document changes

* Skip type checks for tqdm

* Make progress bar more flexible with kwargs

* Fix link checker make command in CONTRIBUTE.md

* Update doc link to be sphinx compatible

* Switch to pytets-forked from pytest-xdist

Co-authored-by: Eddie Bergman <[email protected]>

Author: aron-bram

CONTRIBUTING.md

@@ -252,8 +252,8 @@ Lastly, if the feature really is a game changer or you're very proud of it, cons
     make doc
     ```
     *   If you're unfamiliar with sphinx, it's a documentation generator which can read comments and docstrings from within the code and generate html documentation.
-    *   If you've added documentation, we also have a command `links` for making sure 
-        all the links correctly go to some destination.
+    *   If you've added documentation, we also have a command `links` for making 
+        sure all the links correctly go to some destination.
         This helps tests for dead links or accidental typos.
     ```bash
     make links

autosklearn/automl.py

@@ -120,6 +120,7 @@
     warnings_to,
 )
 from autosklearn.util.parallel import preload_modules
+from autosklearn.util.progress_bar import ProgressBar
 from autosklearn.util.smac_wrap import SMACCallback, SmacRunCallback
 from autosklearn.util.stopwatch import StopWatch
 
@@ -239,6 +240,7 @@ def __init__(
         get_trials_callback: SMACCallback | None = None,
         dataset_compression: bool | Mapping[str, Any] = True,
         allow_string_features: bool = True,
+        disable_progress_bar: bool = False,
     ):
         super().__init__()
 
@@ -295,6 +297,7 @@ def __init__(
         self.logging_config = logging_config
         self.precision = precision
         self.allow_string_features = allow_string_features
+        self.disable_progress_bar = disable_progress_bar
         self._initial_configurations_via_metalearning = (
             initial_configurations_via_metalearning
         )
@@ -626,6 +629,12 @@ def fit(
         # By default try to use the TCP logging port or get a new port
         self._logger_port = logging.handlers.DEFAULT_TCP_LOGGING_PORT
 
+        progress_bar = ProgressBar(
+            total=self._time_for_task,
+            disable=self.disable_progress_bar,
+            desc="Fitting to the training data",
+            colour="green",
+        )
         # Once we start the logging server, it starts in a new process
         # If an error occurs then we want to make sure that we exit cleanly
         # and shut it down, else it might hang
@@ -961,6 +970,7 @@ def fit(
             self._logger.exception(e)
             raise e
         finally:
+            progress_bar.stop()
             self._fit_cleanup()
 
         self.fitted = True

autosklearn/estimators.py

@@ -76,6 +76,7 @@ def __init__(
         get_trials_callback: SMACCallback | None = None,
         dataset_compression: Union[bool, Mapping[str, Any]] = True,
         allow_string_features: bool = True,
+        disable_progress_bar: bool = False,
     ):
         """
         Parameters
@@ -381,6 +382,10 @@ def __init__(
             Whether autosklearn should process string features. By default the
             textpreprocessing is enabled.
 
+        disable_progress_bar: bool = False
+            Whether to disable the progress bar that is displayed in the console
+            while fitting to the training data.
+
         Attributes
         ----------
         cv_results_ : dict of numpy (masked) ndarrays
@@ -475,6 +480,7 @@ def __init__(
         self.get_trials_callback = get_trials_callback
         self.dataset_compression = dataset_compression
         self.allow_string_features = allow_string_features
+        self.disable_progress_bar = disable_progress_bar
 
         self.automl_ = None  # type: Optional[AutoML]
 
@@ -525,6 +531,7 @@ def build_automl(self):
             get_trials_callback=self.get_trials_callback,
             dataset_compression=self.dataset_compression,
             allow_string_features=self.allow_string_features,
+            disable_progress_bar=self.disable_progress_bar,
         )
 
         return automl

autosklearn/experimental/askl2.py

@@ -166,6 +166,7 @@ def __init__(
         load_models: bool = True,
         dataset_compression: Union[bool, Mapping[str, Any]] = True,
         allow_string_features: bool = True,
+        disable_progress_bar: bool = False,
     ):
 
         """
@@ -284,6 +285,10 @@ def __init__(
         load_models : bool, optional (True)
             Whether to load the models after fitting Auto-sklearn.
 
+        disable_progress_bar: bool = False
+            Whether to disable the progress bar that is displayed in the console
+            while fitting to the training data.
+
         Attributes
         ----------
 
@@ -337,6 +342,7 @@ def __init__(
             scoring_functions=scoring_functions,
             load_models=load_models,
             allow_string_features=allow_string_features,
+            disable_progress_bar=disable_progress_bar,
         )
 
     def train_selectors(self, selected_metric=None):

autosklearn/util/progress_bar.py

@@ -0,0 +1,68 @@
+from typing import Any
+
+import datetime
+import time
+from threading import Thread
+
+from tqdm import trange
+
+
+class ProgressBar(Thread):
+    """A Thread that displays a tqdm progress bar in the console.
+
+    It is specialized to display information relevant to fitting to the training data
+    with auto-sklearn.
+
+    Parameters
+    ----------
+    total : int
+        The total amount that should be reached by the progress bar once it finishes
+    update_interval : float
+        Specifies how frequently the progress bar is updated (in seconds)
+    disable : bool
+        Turns on or off the progress bar. If True, this thread won't be started or
+        initialized.
+    kwargs : Any
+        Keyword arguments that are passed into tqdm's constructor. Refer to:
+        `tqdm <https://tqdm.github.io/docs/tqdm/>`_. Note that postfix can not be
+        specified in the kwargs since it is already passed into tqdm by this class.
+    """
+
+    def __init__(
+        self,
+        total: int,
+        update_interval: float = 1.0,
+        disable: bool = False,
+        **kwargs: Any,
+    ):
+        self.disable = disable
+        if not disable:
+            super().__init__(name="_progressbar_")
+            self.total = total
+            self.update_interval = update_interval
+            self.terminated: bool = False
+            self.kwargs = kwargs
+            # start this thread
+            self.start()
+
+    def run(self) -> None:
+        """Display a tqdm progress bar in the console.
+
+        Additionally, it shows useful information related to the task. This method
+        overrides the run method of Thread.
+        """
+        if not self.disable:
+            for _ in trange(
+                self.total,
+                postfix=f"The total time budget for this task is "
+                f"{datetime.timedelta(seconds=self.total)}",
+                **self.kwargs,
+            ):
+                if not self.terminated:
+                    time.sleep(self.update_interval)
+
+    def stop(self) -> None:
+        """Terminates the thread."""
+        if not self.disable:
+            self.terminated = True
+            super().join()

pyproject.toml

@@ -155,7 +155,8 @@ module = [
     "setuptools.*",
     "pkg_resources.*",
     "yaml.*",
-    "psutil.*"
+    "psutil.*",
+    "tqdm.*",
 ]
 ignore_missing_imports = true
 

requirements.txt

@@ -14,8 +14,9 @@ pyyaml
 pandas>=1.0
 liac-arff
 threadpoolctl
+tqdm
 
 ConfigSpace>=0.4.21,<0.5
 pynisher>=0.6.3,<0.7
 pyrfr>=0.8.1,<0.9
-smac>=1.2,<1.3
+smac>=1.2,<1.3

setup.py

@@ -32,7 +32,7 @@
     "test": [
         "pytest>=4.6",
         "pytest-cov",
-        "pytest-xdist",
+        "pytest-forked",
         "pytest-timeout",
         "pytest-cases>=3.6.11",
         "mypy",