Repository: PennyLaneAI/pennylane
Pull request #5560 (2024-04-23 15:49:30)

add progress bar when downloading datasets
**Context:**
The beloved progress bar was neglected when I made the data module because it wasn't an MVP feature, but adding it is quite simple.

**Description of the Change:**
- use `rich.progress` while downloading datasets to show a progress bar. If `rich` isn't installed, just print progress to stdout
- refactor `ThreadPoolExecutor` usage to call `map` instead of `submit`. note that it [does the error collection](https://github.com/python/cpython/blob/30e25d4ff9e4adcbeb775b2125196017d0d19b23/Lib/concurrent/futures/_base.py#L619) on its own

**Benefits:**
progress bar 😄

**Possible Drawbacks:**
users must install `rich` to see the progress bar, otherwise they get a basic progress update on stdout (which isn't bad). the real problem is that if you don't have `rich` installed, you won't know that you need it

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5560?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.68%. Comparing base [(`a8bba50`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/a8bba509e16a8582d57a0baa978ddb89a6d4d75a?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`95e2a16`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5560?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 1 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5560      +/-   ##
==========================================
- Coverage   99.69%   99.68%   -0.01%     
==========================================
  Files         410      410              
  Lines       38230    37940     -290     
==========================================
- Hits        38113    37822     -291     
- Misses        117      118       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5560?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

trbromley: @DSGuala what do you recommend here? Are we happy to merge this in as a net improvement? It looks like the drawbacks are just that users might not have `rich` installed, but nothing breaks.
DSGuala: > @DSGuala what do you recommend here? Are we happy to merge this in as a net improvement? It looks like the drawbacks are just that users might not have `rich` installed, but nothing breaks.

I'm not confident in merging yet. I would like to keep it as a prototype and merge when we have something that also tracks the download. @brownj85 will be looking into this in the next iteration if there is time. Otherwise we can plan it for Q3.
brownj85: @timmysilv, @DSGuala there are some changes I've introduced since taking over this PR. The default implementation provides more detailed information about the files being downloaded:
```
>>> from pennylane.data import load
>>> ds = load("qchem", molname="H2", bondlength="full", basis="STO-3G")
qchem/H2/STO-3G/0.5/H2_STO-3G_0.5.h5     1.11/1.11 MB
qchem/H2/STO-3G/0.54/H2_STO-3G_0.54.h5   1.11/1.11 MB
qchem/H2/STO-3G/0.58/H2_STO-3G_0.58.h5   1.11/1.11 MB
qchem/H2/STO-3G/0.62/H2_STO-3G_0.62.h5   1.13/1.13 MB
qchem/H2/STO-3G/0.66/H2_STO-3G_0.66.h5   1.11/1.11 MB
qchem/H2/STO-3G/0.7/H2_STO-3G_0.7.h5     1.11/1.11 MB
...
```

Getting information about in-progress downloads required a change to how `_download_full` works, it now streams data in chunks to disk rather than loading everything into memory first. This should provide better performance in any case.

timmysilv: I think it looks great!
brownj85: > Looks great so far! Left a couple of comments, but the main thing is that it would be nice for it to display in Jupyter notebooks. Is it implemented and just not displaying on my side?

The `rich` implementation should now work when running in a Jupyter notebook


timmysilv: what environment did you test it in? I'm seeing things working fine in terminal with/without rich, and with/without files already being present. I tried it all in jupyter lab and saw no loading outputs.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 7088f0723..81b0a21d1 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -262,6 +262,9 @@
 * Added `is_leaf` parameter to function `flatten` in the `qml.pytrees` module. This is to allow node flattening to be stopped for any node where the `is_leaf` optional argument evaluates to being `True`.
   [(#6107)](https://github.com/PennyLaneAI/pennylane/issues/6107)
 
+* Added a progress bar when downloading datasets with `qml.data.load()`
+  [(#5560)](https://github.com/PennyLaneAI/pennylane/pull/5560)
+
 <h3>Breaking changes 💔</h3>
 
 * `GlobalPhase` is considered non-differentiable with tape transforms.
@@ -418,6 +421,7 @@ Ali Asadi,
 Utkarsh Azad,
 Tonmoy T. Bhattacharya,
 Gabriel Bottrill,
+Jack Brown,
 Ahmed Darwish,
 Astral Cai,
 Yushao Chen,
diff --git a/pennylane/data/data_manager/__init__.py b/pennylane/data/data_manager/__init__.py
index b07e3017f..e8f274c30 100644
--- a/pennylane/data/data_manager/__init__.py
+++ b/pennylane/data/data_manager/__init__.py
@@ -17,17 +17,17 @@ them.
 """
 
 import urllib.parse
-from collections.abc import Mapping, Iterable
-from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait
+from concurrent import futures
 from functools import lru_cache
 from pathlib import Path
 from time import sleep
-from typing import Optional, Union
-
-from requests import get
+from typing import Iterable, Mapping, Optional, Union
+import sys
+from requests import get, head
 
 from pennylane.data.base import Dataset
 from pennylane.data.base.hdf5 import open_hdf5_s3
+from pennylane.data.data_manager import progress
 
 from .foldermap import DataPath, FolderMapView, ParamArg
 from .params import DEFAULT, FULL, format_params
@@ -55,12 +55,13 @@ def _get_data_struct():
     return response.json()
 
 
-def _download_partial(
+def _download_partial(  # pylint: disable=too-many-arguments
     s3_url: str,
     dest: Path,
     attributes: Optional[Iterable[str]],
     overwrite: bool,
     block_size: int,
+    pbar_task: Optional[progress.Task],
 ) -> None:
     """Download the requested attributes of the Dataset at ``s3_path`` into ``dest``.
 
@@ -98,23 +99,33 @@ def _download_partial(
     del remote_dataset
     del dest_dataset
 
+    if pbar_task:
+        file_size = dest.stat().st_size
+        pbar_task.update(completed=file_size, total=file_size)
+
 
-def _download_full(s3_url: str, dest: Path):
+def _download_full(s3_url: str, dest: Path, block_size: int, pbar_task: Optional[progress.Task]):
     """Download the full dataset file at ``s3_url`` to ``path``."""
+    resp = get(s3_url, timeout=5.0, stream=True)
+    resp.raise_for_status()
 
     with open(dest, "wb") as f:
-        resp = get(s3_url, timeout=5.0)
-        resp.raise_for_status()
-
-        f.write(resp.content)
+        if pbar_task is not None:
+            for block in resp.iter_content(chunk_size=block_size):
+                f.write(block)
+                pbar_task.update(advance=len(block))
+        else:
+            for block in resp.iter_content(chunk_size=block_size):
+                f.write(block)
 
 
-def _download_dataset(
-    data_path: DataPath,
+def _download_dataset(  # pylint:disable=too-many-arguments
+    s3_url: str,
     dest: Path,
     attributes: Optional[Iterable[str]],
     block_size: int,
-    force: bool = False,
+    force: bool,
+    pbar_task: Optional[progress.Task],
 ) -> None:
     """Downloads the dataset at ``data_path`` to ``dest``, optionally downloading
     only requested attributes. If ``attributes`` is not provided, every attribute
@@ -122,18 +133,80 @@ def _download_dataset(
 
     If any of the attributes of the remote dataset are already downloaded locally,
     they will not be overwritten unless ``force`` is True.
-    """
 
-    # URL-escape special characters like '+', '$', and '%' in the data path
-    url_safe_datapath = urllib.parse.quote(str(data_path))
-    s3_url = f"{S3_URL}/{url_safe_datapath}"
+    If ``pbar_task`` is provided, it will be updated with the download progress.
+    """
 
     if attributes is not None or dest.exists():
         _download_partial(
-            s3_url, dest=dest, attributes=attributes, overwrite=force, block_size=block_size
+            s3_url,
+            dest=dest,
+            attributes=attributes,
+            overwrite=force,
+            block_size=block_size,
+            pbar_task=pbar_task,
         )
     else:
-        _download_full(s3_url, dest=dest)
+        _download_full(s3_url, dest=dest, block_size=block_size, pbar_task=pbar_task)
+
+
+def _download_datasets(  # pylint: disable=too-many-arguments
+    s3_base_url: str,
+    folder_path: Path,
+    data_paths: list[DataPath],
+    attributes: Optional[Iterable[str]],
+    force: bool,
+    block_size: int,
+    num_threads: int,
+    pbar: Optional[progress.Progress],
+) -> list[Path]:
+    """Downloads the datasets with given ``data_paths`` to ``folder_path``, copying the
+    directory structure of the bucket at ``s3_base_url``.
+
+    If ``pbar`` is provided, a progress task will be added for each requested dataset.
+
+    Returns:
+        list[Path]: List of downloaded dataset paths
+    """
+    # URL-escape special characters like '+', '$', and '%' in the data path
+    s3_urls = [f"{s3_base_url}/{urllib.parse.quote(str(data_path))}" for data_path in data_paths]
+
+    dest_paths = [folder_path / data_path for data_path in data_paths]
+    for path_parents in set(path.parent for path in dest_paths):
+        path_parents.mkdir(parents=True, exist_ok=True)
+
+    if pbar is not None:
+        if attributes is None:
+            file_sizes = [int(head(s3_url).headers["Content-Length"]) for s3_url in s3_urls]
+        else:
+            # Can't get file sizes for partial downloads
+            file_sizes = (None for _ in s3_urls)
+
+        pbar_tasks = [
+            pbar.add_task(str(dest_path.relative_to(folder_path)), total=file_size)
+            for dest_path, file_size in zip(dest_paths, file_sizes)
+        ]
+    else:
+        pbar_tasks = (None for _ in dest_paths)
+
+    with futures.ThreadPoolExecutor(min(num_threads, len(dest_paths))) as pool:
+        for s3_url, dest_path, pbar_task in zip(s3_urls, dest_paths, pbar_tasks):
+            futs = [
+                pool.submit(
+                    _download_dataset,
+                    s3_url,
+                    dest_path,
+                    attributes=attributes,
+                    force=force,
+                    block_size=block_size,
+                    pbar_task=pbar_task,
+                )
+            ]
+            for result in futures.wait(futs, return_when=futures.FIRST_EXCEPTION).done:
+                if result.exception() is not None:
+                    raise result.exception()
+
+    return dest_paths
 
 
 def _validate_attributes(data_struct: dict, data_name: str, attributes: Iterable[str]):
@@ -160,6 +233,7 @@ def load(  # pylint: disable=too-many-arguments
     force: bool = False,
     num_threads: int = 50,
     block_size: int = 8388608,
+    progress_bar: Optional[bool] = None,
     **params: Union[ParamArg, str, list[str]],
 ):
     r"""Downloads the data if it is not already present in the directory and returns it as a list of
@@ -175,6 +249,8 @@ def load(  # pylint: disable=too-many-arguments
         block_size (int)  : The number of bytes to fetch per read operation when fetching datasets from S3.
             Larger values may improve performance for large datasets, but will slow down small reads. Defaults
             to 8MB
+        progress_bar (bool) : Whether to show a progress bars for downloads. Defaults to True if running
+            in an interactive terminal, False otherwise.
         params (kwargs)   : Keyword arguments exactly matching the parameters required for the data type.
             Note that these are not optional
 
@@ -252,33 +328,37 @@ def load(  # pylint: disable=too-many-arguments
     if attributes:
         _validate_attributes(data_struct, data_name, attributes)
 
-    folder_path = Path(folder_path)
-
+    folder_path = Path(folder_path).resolve()
     data_paths = [data_path for _, data_path in foldermap.find(data_name, **params)]
 
-    dest_paths = [folder_path / data_path for data_path in data_paths]
-
-    for path_parents in set(path.parent for path in dest_paths):
-        path_parents.mkdir(parents=True, exist_ok=True)
+    progress_bar = progress_bar if progress_bar is not None else sys.stdout.isatty()
 
-    with ThreadPoolExecutor(min(num_threads, len(dest_paths))) as pool:
-        futures = [
-            pool.submit(
-                _download_dataset,
-                data_path,
-                dest_path,
+    if progress_bar:
+        with progress.Progress() as pbar:
+            download_paths = _download_datasets(
+                S3_URL,
+                folder_path,
+                data_paths,
                 attributes,
                 force=force,
                 block_size=block_size,
+                num_threads=num_threads,
+                pbar=pbar,
             )
-            for data_path, dest_path in zip(data_paths, dest_paths)
-        ]
-        results = wait(futures, return_when=FIRST_EXCEPTION)
-        for result in results.done:
-            if result.exception() is not None:
-                raise result.exception()
 
-    return [Dataset.open(Path(dest_path), "a") for dest_path in dest_paths]
+    else:
+        download_paths = _download_datasets(
+            S3_URL,
+            folder_path,
+            data_paths,
+            attributes,
+            force=force,
+            block_size=block_size,
+            num_threads=num_threads,
+            pbar=None,
+        )
+
+    return [Dataset.open(path, "a") for path in download_paths]
 
 
 def list_datasets() -> dict:
@@ -461,7 +541,11 @@ def load_interactive():
         return None
 
     return load(
-        data_name, attributes=attributes, folder_path=dest_folder, force=force, **description
+        data_name,
+        attributes=attributes,
+        folder_path=dest_folder,
+        force=force,
+        **description,
     )[0]
 
 
diff --git a/pennylane/data/data_manager/foldermap.py b/pennylane/data/data_manager/foldermap.py
index 48eba6e97..f7324fa34 100644
--- a/pennylane/data/data_manager/foldermap.py
+++ b/pennylane/data/data_manager/foldermap.py
@@ -16,7 +16,6 @@ Contains ``FolderMapView`` for reading the ``foldermap.json`` file in the
 datasets bucket.
 """
 
-
 from collections.abc import Iterable, Iterator, Mapping
 from pathlib import PurePosixPath
 from typing import Any, Literal, Optional, Union
diff --git a/pennylane/data/data_manager/params.py b/pennylane/data/data_manager/params.py
index 0e467777c..16c3fe8e5 100644
--- a/pennylane/data/data_manager/params.py
+++ b/pennylane/data/data_manager/params.py
@@ -15,7 +15,6 @@
 Contains types and functions for dataset parameters.
 """
 
-
 import enum
 from collections.abc import Iterable, Iterator, Mapping
 from functools import lru_cache
diff --git a/pennylane/data/data_manager/progress/__init__.py b/pennylane/data/data_manager/progress/__init__.py
new file mode 100644
index 000000000..f11d8c259
--- /dev/null
+++ b/pennylane/data/data_manager/progress/__init__.py
@@ -0,0 +1,94 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""A library for showing loading progress, using ``rich`` or basic stdout."""
+
+from typing import Any, Optional
+
+from pennylane.data.data_manager.progress._default import (
+    make_progress as make_progress_default,
+)
+
+try:
+    from pennylane.data.data_manager.progress._rich import (
+        make_progress as make_progress_rich,
+    )
+except ImportError:  # pragma: no cover
+    make_progress_rich = None
+
+
+class Task:  # pylint: disable=too-few-public-methods
+    """Represents progress display for a single dataset download."""
+
+    def __init__(self, _task_id: Any, _progress: Any):
+        """Private constructor."""
+        self._progress = _progress
+        self._task_id = _task_id
+
+    def update(
+        self,
+        *,
+        completed: Optional[float] = None,
+        advance: Optional[float] = None,
+        total: Optional[float] = None,
+    ):
+        """Update download state.
+
+        Args:
+            advance: Adds to number of bytes downloaded so far
+            completed: Sets the number of bytes downloaded so far
+            total: Sets the total number of bytes for the download
+        """
+        self._progress.update(
+            self._task_id, completed=completed, total=total, advance=advance, refresh=True
+        )
+
+
+class Progress:
+    """Displays dataset download progress on the terminal. Will use
+    ``rich.progress.Progress`` if available, otherwise it will fall back to the
+    default implementation.
+
+    Must be used as a context manager to ensure correct output.
+    """
+
+    def __init__(self) -> None:
+        """Initialize progress."""
+        if make_progress_rich is not None:
+            self.progress = make_progress_rich()
+        else:
+            self.progress = make_progress_default()
+
+    def __enter__(self) -> "Progress":
+        """Enter progress context."""
+        self.progress.__enter__()
+
+        return self
+
+    def __exit__(self, *args):
+        """Exit progress context."""
+        return self.progress.__exit__(*args)
+
+    def add_task(self, description: str, total: Optional[float] = None) -> Task:
+        """Add a task.
+
+        Args:
+            description: Description for the task
+            total: Total size of the dataset download in bytes, if available.
+        """
+        task_id = self.progress.add_task(description=description, total=total)
+
+        return Task(task_id, self.progress)
+
+
+__all__ = ["Progress", "Task"]
diff --git a/pennylane/data/data_manager/progress/_default/__init__.py b/pennylane/data/data_manager/progress/_default/__init__.py
new file mode 100644
index 000000000..fc121a931
--- /dev/null
+++ b/pennylane/data/data_manager/progress/_default/__init__.py
@@ -0,0 +1,193 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Fallback implementation of a progress bar, used when ``rich`` is not available."""
+
+import shutil
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from pennylane.data.data_manager.progress._default import term
+
+
+@dataclass
+class Task:
+    """Data class for progress bar state."""
+
+    description: str
+    completed: float = 0
+    total: Optional[float] = None
+
+    def update(
+        self,
+        *,
+        advance: Optional[float] = None,
+        completed: Optional[float] = None,
+        total: Optional[float] = None,
+    ) -> None:
+        """Update the state of the progress bar and set the display
+        string."""
+        if completed:
+            self.completed = completed
+        if advance:
+            self.completed += advance
+        if total:
+            self.total = total
+
+
+class TerminalInfo:
+    """Contains information on  the dimensions of
+    the terminal."""
+
+    # pylint: disable=too-few-public-methods
+
+    def __init__(self):
+        self.columns, self.lines = shutil.get_terminal_size()
+        self.max_display_lines = self.lines - 2
+        self.description_len_max = int(self.columns * 0.6) - 1
+
+
+class DefaultProgress:
+    """Implements a progress bar."""
+
+    tasks: list[Task]
+
+    def __init__(self):
+        self.tasks = []
+
+        self._active = False
+        self._term_info = TerminalInfo()
+        self._task_display_lines = []
+        self._curr_longest_description = 0
+
+    def __enter__(self) -> "DefaultProgress":
+        self._active = True
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Any,
+    ) -> None:
+        self._print_final()
+        self._active = False
+
+    def add_task(self, description: str, total: Optional[float] = None) -> int:
+        """Add a task."""
+        description = _truncate(description, self._term_info.description_len_max)
+        self.tasks.append(Task(description=description, total=total))
+
+        self._curr_longest_description = max(self._curr_longest_description, len(description))
+        self.refresh()
+
+        return len(self.tasks) - 1
+
+    def refresh(self, task_id: Optional[int] = None):
+        """Refresh display liens for one or all tasks."""
+        if task_id is None:
+            self._task_display_lines.clear()
+            self._task_display_lines.extend(
+                self._get_task_display_line(task)
+                for task in self.tasks[: self._term_info.max_display_lines]
+            )
+
+            if len(self.tasks) > self._term_info.max_display_lines:
+                self._task_display_lines.append(f"{term.erase_line()}...")
+
+        elif task_id < self._term_info.max_display_lines:
+            self._task_display_lines[task_id] = self._get_task_display_line(self.tasks[task_id])
+
+    def update(
+        self,
+        task_id: int,
+        *,
+        completed: Optional[float] = None,
+        total: Optional[float] = None,
+        advance: Optional[float] = None,
+        refresh: bool = False,
+    ):
+        """Update task with given ``task_id`` and refresh its progress.
+
+        Args:
+            task_id: ID of task
+            completed: Set the completed state of the task
+            total: Set the total for the task
+            advance: Advance the completion state of the task
+            refresh: Included for compatability with ``rich.Progress``, has no effect
+        """
+        del refresh
+        task = self.tasks[task_id]
+        task.update(advance=advance, completed=completed, total=total)
+
+        self.refresh(task_id)
+
+        if self._active:
+            self._print()
+
+    def _get_task_display_line(self, task: Task) -> str:
+        """Get display line for the task."""
+        if task.total is None:
+            progress_column = f"{task.completed / 1e6:.2f} MB"
+        else:
+            progress_column = f"{task.completed / 1e6:.2f}/{task.total / 1e6:.2f} MB"
+
+        display = _truncate(
+            f"{task.description.ljust(self._curr_longest_description + 1)}" f"{progress_column}",
+            self._term_info.columns,
+        )
+
+        return f"{term.erase_line()}{display}"
+
+    def _print_final(self):
+        """Print all tasks without any terminal control. Should be called when
+        progress is complete."""
+        print(*(self._get_task_display_line(task) for task in self.tasks), sep="\n", flush=True)
+
+    def _print(self):
+        """Prints up to ``_task_display_lines_max`` lines and returns the terminal cursor to
+        the starting point."""
+        if len(self._task_display_lines) > 1:
+            end = f"\r{term.cursor_up(len(self._task_display_lines) - 1)}"
+        else:
+            end = "\r"
+
+        print(
+            *self._task_display_lines,
+            sep="\n",
+            end=end,
+            flush=True,
+        )
+
+
+def make_progress() -> DefaultProgress:
+    """Factory function for a progress instance."""
+    return DefaultProgress()
+
+
+def _truncate(s: str, maxlen: int) -> str:
+    """If ``s`` is longer than ``maxlen``, truncate
+    it and replace the last 3 characters with '...'.
+
+    >>> _truncate("abcdef", 6)
+    "abcdef"
+    >>> _truncate("abcdef", 5)
+    "ab..."
+    """
+    if len(s) > maxlen:
+        return f"{s[:maxlen - 3]}..."
+
+    return s
+
+
+__all__ = ["make_progress", "DefaultProgress"]
diff --git a/pennylane/data/data_manager/progress/_default/term.py b/pennylane/data/data_manager/progress/_default/term.py
new file mode 100644
index 000000000..0590ecca5
--- /dev/null
+++ b/pennylane/data/data_manager/progress/_default/term.py
@@ -0,0 +1,36 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Functions wrapping ANSI terminal control sequences. See:
+https://en.wikipedia.org/wiki/ANSI_escape_code#CSI_(Control_Sequence_Introducer)_sequences
+"""
+
+
+def cursor_up(n: int) -> str:
+    """Return 'A' control sequence, which to moves the cursor up ``n`` columns.
+
+    >>> cursor_up(2)
+    '\x1B[2;A'
+
+    """
+    return f"\x1B[{n};A"
+
+
+def erase_line() -> str:
+    """Return 'K' control sequence, which erases the current line starting from
+    the cursor.
+
+    >>> erase_line()
+    '\x1B[0;K'
+    """
+    return "\x1B[0;K"
diff --git a/pennylane/data/data_manager/progress/_rich.py b/pennylane/data/data_manager/progress/_rich.py
new file mode 100644
index 000000000..f4f40144b
--- /dev/null
+++ b/pennylane/data/data_manager/progress/_rich.py
@@ -0,0 +1,33 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Progress bar using ``rich``."""
+
+import rich.progress
+from rich.progress import Progress as RichProgress
+
+
+def make_progress() -> RichProgress:
+    """Factory function for a progress instance."""
+    return rich.progress.Progress(
+        rich.progress.TextColumn("[progress.description]{task.description}"),
+        rich.progress.BarColumn(),
+        rich.progress.TaskProgressColumn(),
+        rich.progress.TransferSpeedColumn(),
+        rich.progress.DownloadColumn(),
+        refresh_per_second=10,
+        transient=False,
+    )
+
+
+__all__ = ["make_progress", "RichProgress"]
diff --git a/requirements-ci.txt b/requirements-ci.txt
index ef83bbab2..083beaae2 100644
--- a/requirements-ci.txt
+++ b/requirements-ci.txt
@@ -11,4 +11,5 @@ packaging
 autoray>=0.6.1,<0.6.10
 matplotlib
 requests
+rich
 tomli # Drop once minimum Python version is 3.11
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 496d76a6b..3ad3ba1e6 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -10,3 +10,4 @@ black>=21
 tomli~=2.0.0 # Drop once minimum Python version is 3.11
 isort==5.13.2
 pylint==2.7.4
+rich>=13.7.1


Test patch:
diff --git a/tests/data/data_manager/test_dataset_access.py b/tests/data/data_manager/test_dataset_access.py
index e7fdcea43..7157c0cd6 100644
--- a/tests/data/data_manager/test_dataset_access.py
+++ b/tests/data/data_manager/test_dataset_access.py
@@ -16,7 +16,8 @@ Unit tests for the :class:`pennylane.data.data_manager` functions.
 """
 import os
 from pathlib import Path, PosixPath
-from unittest.mock import MagicMock, patch
+from typing import NamedTuple
+from unittest.mock import MagicMock, call, patch
 
 import pytest
 import requests
@@ -24,7 +25,15 @@ import requests
 import pennylane as qml
 import pennylane.data.data_manager
 from pennylane.data import Dataset
-from pennylane.data.data_manager import S3_URL, DataPath, _validate_attributes
+from pennylane.data.data_manager import S3_URL, _validate_attributes
+
+has_rich = False
+try:
+    import rich  # pylint:disable=unused-import
+
+    has_rich = True
+except ImportError:
+    pass
 
 # pylint:disable=protected-access,redefined-outer-name
 
@@ -80,6 +89,11 @@ def get_mock(url, timeout=1.0):
     return resp
 
 
+def head_mock(url):
+    """Return a fake header stating content-length is 1."""
+    return NamedTuple("Head", headers=dict)(headers={"Content-Length": 10000})
+
+
 @pytest.fixture
 def mock_get_args():
     """A Mock object that tracks the arguments passed to ``mock_requests_get``."""
@@ -105,7 +119,15 @@ def mock_requests_get(request, monkeypatch, mock_get_args):
 
         mock_resp.json.return_value = json_data
         if hasattr(request, "param"):
-            mock_resp.content = request.param
+            content = request.param
+            mock_resp.content = content
+
+            def mock_iter_content(chunk_size: int):
+                """Mock for Response.iter_content()."""
+                for i in range(0, len(content), chunk_size):
+                    yield content[i : i + chunk_size]
+
+            mock_resp.iter_content = mock_iter_content
 
         return mock_resp
 
@@ -138,6 +160,7 @@ def mock_load(monkeypatch):
 
 
 @patch.object(requests, "get", get_mock)
+@patch.object(pennylane.data.data_manager, "head", head_mock)
 @patch("pennylane.data.data_manager.sleep")
 @patch("builtins.input")
 class TestLoadInteractive:
@@ -234,7 +257,8 @@ class TestMiscHelpers:
 
 @pytest.fixture
 def mock_download_dataset(monkeypatch):
-    def mock(data_path, dest, attributes, force, block_size):
+    # pylint:disable=too-many-arguments
+    def mock(data_path, dest, attributes, force, block_size, pbar_task):
         dset = Dataset.open(Path(dest), "w")
         dset.close()
 
@@ -243,6 +267,8 @@ def mock_download_dataset(monkeypatch):
     return mock
 
 
+# pylint: disable=too-many-arguments
+@patch.object(pennylane.data.data_manager, "head", head_mock)
 @pytest.mark.usefixtures("mock_download_dataset")
 @pytest.mark.parametrize(
     "data_name, params, expect_paths",
@@ -254,7 +280,9 @@ def mock_download_dataset(monkeypatch):
         )
     ],
 )
-def test_load(tmp_path, data_name, params, expect_paths):
+@pytest.mark.parametrize("progress_bar", [True, False])
+@pytest.mark.parametrize("attributes", [None, ["molecule"]])
+def test_load(tmp_path, data_name, params, expect_paths, progress_bar, attributes):
     """Test that load fetches the correct datasets at the
     expected paths."""
 
@@ -263,6 +291,8 @@ def test_load(tmp_path, data_name, params, expect_paths):
         data_name=data_name,
         folder_path=folder_path,
         block_size=1,
+        progress_bar=progress_bar,
+        attributes=attributes,
         **params,
     )
 
@@ -271,6 +301,7 @@ def test_load(tmp_path, data_name, params, expect_paths):
     }
 
 
+@patch.object(pennylane.data.data_manager, "head", head_mock)
 def test_load_except(monkeypatch, tmp_path):
     """Test that an exception raised by _download_dataset is propagated."""
     monkeypatch.setattr(
@@ -302,7 +333,7 @@ def test_download_dataset_full_or_partial(
     dest.exists.return_value = dest_exists
 
     pennylane.data.data_manager._download_dataset(
-        "dataset/path", attributes=attributes, dest=dest, force=force, block_size=1
+        "dataset/path", attributes=attributes, dest=dest, force=force, block_size=1, pbar_task=None
     )
 
     assert download_partial.called is called_partial
@@ -317,12 +348,20 @@ def test_download_dataset_full_call(download_full, force):
     """
     dest = MagicMock()
     dest.exists.return_value = False
+    pbar_task = MagicMock()
 
     pennylane.data.data_manager._download_dataset(
-        "dataset/path", attributes=None, dest=dest, force=force, block_size=1
+        f"{S3_URL}/dataset/path",
+        attributes=None,
+        dest=dest,
+        force=force,
+        block_size=1,
+        pbar_task=pbar_task,
     )
 
-    download_full.assert_called_once_with(f"{S3_URL}/dataset/path", dest=dest)
+    download_full.assert_called_once_with(
+        f"{S3_URL}/dataset/path", block_size=1, dest=dest, pbar_task=pbar_task
+    )
 
 
 @pytest.mark.parametrize("attributes", [None, ["x"]])
@@ -334,13 +373,24 @@ def test_download_dataset_partial_call(download_partial, attributes, force):
     """
     dest = MagicMock()
     dest.exists.return_value = True
+    pbar_task = MagicMock()
 
     pennylane.data.data_manager._download_dataset(
-        "dataset/path", attributes=attributes, dest=dest, force=force, block_size=1
+        f"{S3_URL}/dataset/path",
+        attributes=attributes,
+        dest=dest,
+        force=force,
+        block_size=1,
+        pbar_task=pbar_task,
     )
 
     download_partial.assert_called_once_with(
-        f"{S3_URL}/dataset/path", dest=dest, attributes=attributes, overwrite=force, block_size=1
+        f"{S3_URL}/dataset/path",
+        dest=dest,
+        attributes=attributes,
+        overwrite=force,
+        block_size=1,
+        pbar_task=pbar_task,
     )
 
 
@@ -349,16 +399,31 @@ def test_download_dataset_partial_call(download_partial, attributes, force):
 def test_download_full(tmp_path):
     """Tests that _download_dataset will fetch the dataset file
     at ``s3_url`` into ``dest``."""
-
     pennylane.data.data_manager._download_full(
-        "dataset/path",
-        tmp_path / "dataset",
+        "dataset/path", tmp_path / "dataset", block_size=1, pbar_task=None
     )
 
     with open(tmp_path / "dataset", "rb") as f:
         assert f.read() == b"This is binary data"
 
 
+@pytest.mark.usefixtures("mock_requests_get")
+@pytest.mark.parametrize("mock_requests_get", [b"0123456789"], indirect=True)
+def test_download_full_with_progress(tmp_path):
+    """Tests that _download_dataset will fetch the dataset file
+    at ``s3_url`` into ``dest`` and call the ``update`` method
+    of the progress bar task."""
+    pbar_task = MagicMock()
+    pennylane.data.data_manager._download_full(
+        "dataset/path", tmp_path / "dataset", block_size=4, pbar_task=pbar_task
+    )
+
+    with open(tmp_path / "dataset", "rb") as f:
+        assert f.read() == b"0123456789"
+
+    pbar_task.update.assert_has_calls([call(advance=4), call(advance=4), call(advance=2)])
+
+
 @pytest.mark.parametrize("overwrite", [True, False])
 @pytest.mark.parametrize(
     "attributes, expect_attrs",
@@ -383,6 +448,7 @@ def test_download_partial_dest_not_exists(
         attributes=attributes,
         overwrite=overwrite,
         block_size=1,
+        pbar_task=MagicMock(),
     )
 
     local = Dataset.open(tmp_path / "dataset")
@@ -425,6 +491,7 @@ def test_download_partial_dest_exists(tmp_path, monkeypatch, attributes, overwri
         attributes=attributes,
         overwrite=overwrite,
         block_size=1,
+        pbar_task=MagicMock(),
     )
 
     local = Dataset.open(tmp_path / "dataset")
@@ -443,25 +510,38 @@ def test_download_partial_no_check_remote(open_hdf5_s3, tmp_path):
     local_dataset.close()
 
     pennylane.data.data_manager._download_partial(
-        "dataset_url", tmp_path / "dataset", ["x", "y"], overwrite=False, block_size=1
+        "dataset_url",
+        tmp_path / "dataset",
+        ["x", "y"],
+        overwrite=False,
+        block_size=1,
+        pbar_task=MagicMock(),
     )
 
     open_hdf5_s3.assert_not_called()
 
 
 @patch("builtins.open")
+@patch.object(pennylane.data.data_manager, "head", head_mock)
 @pytest.mark.parametrize(
     "datapath, escaped",
     [("data/NH3+/data.h5", "data/NH3%2B/data.h5"), ("data/CA$H/money.h5", "data/CA%24H/money.h5")],
 )
-def test_download_dataset_escapes_url(_, mock_get_args, datapath, escaped):
-    """Tests that _download_dataset escapes special characters in a URL when doing a full download."""
+def test_download_datasets_escapes_url(_, tmp_path, mock_get_args, datapath, escaped):
+    """Tests that _download_datasets escapes special characters in a URL when doing a full download."""
 
     dest = MagicMock()
     dest.exists.return_value = False
 
-    pennylane.data.data_manager._download_dataset(
-        DataPath(datapath), dest=dest, attributes=None, block_size=1
+    pennylane.data.data_manager._download_datasets(
+        S3_URL,
+        folder_path=tmp_path,
+        data_paths=[datapath],
+        attributes=None,
+        force=True,
+        block_size=1,
+        num_threads=1,
+        pbar=MagicMock(),
     )
 
     mock_get_args.assert_called_once()
@@ -473,19 +553,33 @@ def test_download_dataset_escapes_url(_, mock_get_args, datapath, escaped):
     "datapath, escaped",
     [("data/NH3+/data.h5", "data/NH3%2B/data.h5"), ("data/CA$H/money.h5", "data/CA%24H/money.h5")],
 )
-def test_download_dataset_escapes_url_partial(download_partial, datapath, escaped):
-    """Tests that _download_dataset escapes special characters in a URL when doing a partial
+def test_download_datasets_escapes_url_partial(download_partial, tmp_path, datapath, escaped):
+    """Tests that _download_datasets escapes special characters in a URL when doing a partial
     download."""
-    dest = Path("dest")
     attributes = ["attr"]
     force = False
-
-    pennylane.data.data_manager._download_dataset(
-        DataPath(datapath), dest=dest, attributes=attributes, force=force, block_size=1
+    pbar = MagicMock()
+    pbar_task = MagicMock()
+    pbar.add_task.return_value = pbar_task
+
+    pennylane.data.data_manager._download_datasets(
+        S3_URL,
+        folder_path=tmp_path,
+        data_paths=[datapath],
+        attributes=attributes,
+        force=force,
+        block_size=1,
+        num_threads=1,
+        pbar=pbar,
     )
 
     download_partial.assert_called_once_with(
-        f"{S3_URL}/{escaped}", dest=dest, attributes=attributes, overwrite=force, block_size=1
+        f"{S3_URL}/{escaped}",
+        dest=tmp_path / datapath,
+        attributes=attributes,
+        overwrite=force,
+        block_size=1,
+        pbar_task=pbar_task,
     )
 
 
diff --git a/tests/data/data_manager/test_progress.py b/tests/data/data_manager/test_progress.py
new file mode 100644
index 000000000..6939daf5d
--- /dev/null
+++ b/tests/data/data_manager/test_progress.py
@@ -0,0 +1,210 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Tests for :class:`pennylane.data.progress`
+"""
+import shutil
+
+import pytest
+
+import pennylane.data.data_manager.progress
+import pennylane.data.data_manager.progress._default
+from pennylane.data.data_manager.progress import Progress, Task
+from pennylane.data.data_manager.progress._default import term
+
+# pylint: disable=redefined-outer-name
+
+
+@pytest.fixture
+def disable_rich(request, monkeypatch):
+    """Sets ``progress.make_progress_rich`` to ``None``, mocking the result of
+    a failed import.
+
+    Also skips the calling test if rich is not disabled, but is not
+    installed.
+    """
+    if getattr(request, "param", True):
+        monkeypatch.setattr(pennylane.data.data_manager.progress, "make_progress_rich", None)
+        return True
+
+    try:
+        from pennylane.data.data_manager.progress._rich import RichProgress
+
+        del RichProgress
+    except ImportError:
+        pytest.skip("'rich' is not installed")
+
+    return False
+
+
+class TestProgress:
+    """Tests for :class:`pennylane.data.progress.Progress`."""
+
+    @pytest.mark.parametrize(
+        "disable_rich, expect_cls",
+        [
+            (True, "pennylane.data.data_manager.progress._default.DefaultProgress"),
+            (False, "rich.progress.Progress"),
+        ],
+        indirect=["disable_rich"],
+    )
+    def test_init(self, disable_rich, expect_cls):
+        """Test that ``__init__()`` uses the correct implementation based
+        on the value of ``disable_rich``."""
+        del disable_rich
+
+        prog_cls = type(Progress().progress)
+        assert f"{prog_cls.__module__}.{prog_cls.__name__}" == expect_cls
+
+    @pytest.mark.usefixtures("disable_rich")
+    @pytest.mark.parametrize("disable_rich", [True, False], indirect=True)
+    @pytest.mark.parametrize("total", [100, None])
+    def test_add_task(self, total):
+        """Test that ``add_task()`` returns a new Task instance."""
+        progress = Progress()
+        task = progress.add_task(description="abc", total=total)
+
+        assert isinstance(task, Task)
+
+    @pytest.mark.usefixtures("disable_rich")
+    @pytest.mark.parametrize("disable_rich", [True, False], indirect=True)
+    def test_context(self):
+        """Test that ``__enter__()`` returns the instance."""
+        progress = Progress()
+        with progress as progress_ctx:
+            assert progress_ctx is progress
+
+
+@pytest.mark.usefixtures("disable_rich")
+class TestDefaultProgress:
+    """Tests for :class:`pennylane.data.data_manger.progress._default.DefaultProgress`."""
+
+    @pytest.fixture(autouse=True)
+    def terminal_size(self, monkeypatch):
+        """Patches terminal size for testing."""
+
+        def get_terminal_size(fallback=None):
+            # pylint: disable=unused-argument
+            return (40, 4)
+
+        monkeypatch.setattr(shutil, "get_terminal_size", get_terminal_size)
+
+        yield (40, 4)
+
+    @pytest.fixture()
+    def progress(self):
+        """Progress bar fixture."""
+        yield Progress()
+
+    def test_task_update_with_total(self, progress: Progress, capsys: pytest.CaptureFixture):
+        """Tests for updating a task with a total."""
+        task_1 = progress.add_task(description="Task-1", total=100 * 1e6)
+        progress.add_task(description="Task-2", total=None)
+
+        with progress:
+            task_1.update(advance=50 * 1e6)
+            out, _ = capsys.readouterr()
+
+            assert out == (
+                f"{term.erase_line()}Task-1 50.00/100.00 MB\n"
+                f"{term.erase_line()}Task-2 0.00 MB\r{term.cursor_up(1)}"
+            )
+
+    def test_task_update_one_task(self, progress: Progress, capsys: pytest.CaptureFixture):
+        """Test for updating with only one task."""
+        task = progress.add_task(description="Task-1", total=100 * 1e6)
+
+        with progress:
+            task.update(advance=50 * 1e6)
+            out, _ = capsys.readouterr()
+
+            assert out == f"{term.erase_line()}Task-1 50.00/100.00 MB\r"
+
+    @pytest.mark.parametrize(
+        "kwds, numbers",
+        [
+            ({"advance": 50 * 1e6}, "50.00"),
+            ({"completed": 100 * 1e6, "total": 100 * 1e6}, "100.00/100.00"),
+        ],
+    )
+    def test_task_update_without_total(
+        self, progress: Progress, capsys: pytest.CaptureFixture, kwds, numbers
+    ):
+        """Tests for updating a task without a total."""
+        progress.add_task(description="Task-1", total=100 * 1e6)
+        task_2 = progress.add_task(description="Task-2", total=None)
+
+        with progress:
+            task_2.update(**kwds)
+            out, _ = capsys.readouterr()
+
+            assert out == (
+                f"{term.erase_line()}Task-1 0.00/100.00 MB\n"
+                f"{term.erase_line()}Task-2 {numbers} MB\r{term.cursor_up(1)}"
+            )
+
+    def test_task_lines_truncated(self, progress: Progress, capsys: pytest.CaptureFixture):
+        """Tests that task lines will not be printed if there is not enough
+        terminal lines available."""
+        task_1 = progress.add_task(description="Task-1", total=100e6)
+        progress.add_task(description="Task-2", total=300e6)
+
+        invisible = progress.add_task(description="Task-3", total=100e6)
+        progress.add_task(description="Task-4", total=100e6)
+
+        with progress:
+            task_1.update(advance=50e6)
+            out, _ = capsys.readouterr()
+
+            invisible.update(advance=50e6)
+
+            assert out.encode("utf-8") == (
+                f"{term.erase_line()}Task-1 50.00/100.00 MB\n"
+                f"{term.erase_line()}Task-2 0.00/300.00 MB\n"
+                f"{term.erase_line()}...\r{term.cursor_up(2)}"
+            ).encode("utf-8")
+
+    def test_description_truncated(self, progress: Progress, capsys: pytest.CaptureFixture):
+        """Test that long task descriptions will be truncated."""
+        task_1 = progress.add_task(description="Task-1-with-a-too-long-name", total=100e6)
+        progress.add_task(description="Task-2", total=300e6)
+
+        with progress:
+            task_1.update(advance=50e6)
+            out, _ = capsys.readouterr()
+
+            assert out == (
+                f"{term.erase_line()}Task-1-with-a-too-lo... 50.00/100.00 MB\n"
+                f"{term.erase_line()}Task-2                  0.00/300.00 MB\r{term.cursor_up(1)}"
+            )
+
+    def test_final_print(self, progress: Progress, capsys: pytest.CaptureFixture):
+        """Test that all task lines are printed, without cursor control codes,
+        when the progress bar exits."""
+        progress.add_task(description="Task-1", total=100e6)
+        progress.add_task(description="Task-2", total=300e6)
+
+        task_3 = progress.add_task(description="Task-3", total=100e6)
+
+        with progress:
+            task_3.update(advance=50e6)
+            capsys.readouterr()
+
+        out, _ = capsys.readouterr()
+
+        assert out == (
+            f"{term.erase_line()}Task-1 0.00/100.00 MB\n"
+            f"{term.erase_line()}Task-2 0.00/300.00 MB\n"
+            f"{term.erase_line()}Task-3 50.00/100.00 MB\n"
+        )