Experimental hf logger (#1456)

* first draft for an hflogger

* move to huggingface_hub

* copyright

* Add default tensorboard logger

* leverage  to push to hub in the background

* Push ot hub on context manager closing

* remove clunky HFLoggerMixin and focus on tensorboard

* mark HFSummaryWriter as experimental

* automatic push in HFSummaryWriter

* simplify logic if only 1 LFS file

* more verbose

* Append logdir to path_in_repo

* add apckage reference

* fix typing

* fix kwargs in doc

* fix type

docs/source/_toctree.yml

@@ -68,5 +68,7 @@
       title: Repo Cards and Repo Card Data
     - local: package_reference/space_runtime
       title: Space runtime
+    - local: package_reference/tensorboard
+      title: TensorBoard logger
     - local: package_reference/webhooks_server
       title: Webhooks server

docs/source/package_reference/tensorboard.mdx

@@ -0,0 +1,15 @@
+# TensorBoard logger
+
+TensorBoard is a visualization toolkit for machine learning experimentation. TensorBoard allows tracking and visualizing
+metrics such as loss and accuracy, visualizing the model graph, viewing histograms, displaying images and much more.
+TensorBoard is well integrated with the Hugging Face Hub. The Hub automatically detects TensorBoard traces (such as
+`tfevents`) when pushed to the Hub which starts an instance to visualize them. To get more information about TensorBoard
+integration on the Hub, check out [this guide](https://huggingface.co/docs/hub/tensorboard).
+
+To benefit from this integration, `huggingface_hub` provides a custom logger to push logs to the Hub. It works as a
+drop-in replacement for [SummaryWriter](https://tensorboardx.readthedocs.io/en/latest/tensorboard.html) with no extra
+code needed. Traces are still saved locally and a background job push them to the Hub at regular interval.
+
+## HFSummaryWriter
+
+[[autodoc]] HFSummaryWriter

src/huggingface_hub/__init__.py

@@ -77,6 +77,9 @@
         "SpaceRuntime",
         "SpaceStage",
     ],
+    "_tensorboard_logger": [
+        "HFSummaryWriter",
+    ],
     "_webhooks_payload": [
         "WebhookPayload",
         "WebhookPayloadComment",
@@ -371,6 +374,7 @@ def __dir__():
         SpaceRuntime,  # noqa: F401
         SpaceStage,  # noqa: F401
     )
+    from ._tensorboard_logger import HFSummaryWriter  # noqa: F401
     from ._webhooks_payload import (
         WebhookPayload,  # noqa: F401
         WebhookPayloadComment,  # noqa: F401

src/huggingface_hub/_commit_api.py

@@ -386,8 +386,11 @@ def _wrapped_lfs_upload(batch_action) -> None:
 
     if HF_HUB_ENABLE_HF_TRANSFER:
         logger.debug(f"Uploading {len(filtered_actions)} LFS files to the Hub using `hf_transfer`.")
-        for action in filtered_actions:
+        for action in hf_tqdm(filtered_actions):
             _wrapped_lfs_upload(action)
+    elif len(filtered_actions) == 1:
+        logger.debug("Uploading 1 LFS file to the Hub")
+        _wrapped_lfs_upload(filtered_actions[0])
     else:
         logger.debug(
             f"Uploading {len(filtered_actions)} LFS files to the Hub using up to {num_threads} threads concurrently"

src/huggingface_hub/_commit_scheduler.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import time
+from concurrent.futures import Future
 from dataclasses import dataclass
 from io import SEEK_END, SEEK_SET, BytesIO
 from pathlib import Path
@@ -143,11 +144,19 @@ def stop(self) -> None:
     def _run_scheduler(self) -> None:
         """Dumb thread waiting between each scheduled push to Hub."""
         while True:
-            self.last_future = self.api.run_as_future(self._push_to_hub)
+            self.last_future = self.trigger()
             time.sleep(self.every * 60)
             if self.__stopped:
                 break
 
+    def trigger(self) -> Future:
+        """Trigger a `push_to_hub` and return a future.
+
+        This method is automatically called every `every` minutes. You can also call it manually to trigger a commit
+        immediately, without waiting for the next scheduled commit.
+        """
+        return self.api.run_as_future(self._push_to_hub)
+
     def _push_to_hub(self) -> Optional[CommitInfo]:
         if self.__stopped:  # If stopped, already scheduled commits are ignored
             return None

src/huggingface_hub/_tensorboard_logger.py

@@ -0,0 +1,157 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# 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.
+"""Contains a logger to push training logs to the Hub, using Tensorboard."""
+from pathlib import Path
+from typing import TYPE_CHECKING, List, Optional, Union
+
+from huggingface_hub._commit_scheduler import CommitScheduler
+
+from .utils import experimental, is_tensorboard_available
+
+
+if is_tensorboard_available():
+    from tensorboardX import SummaryWriter
+
+    # TODO: clarify: should we import from torch.utils.tensorboard ?
+else:
+    SummaryWriter = object  # Dummy class to avoid failing at import. Will raise on instance creation.
+
+if TYPE_CHECKING:
+    from tensorboardX import SummaryWriter
+
+
+class HFSummaryWriter(SummaryWriter):
+    """
+    Wrapper around the tensorboard's `SummaryWriter` to push training logs to the Hub.
+
+    Data is logged locally and then pushed to the Hub asynchronously. Pushing data to the Hub is done in a separate
+    thread to avoid blocking the training script. In particular, if the upload fails for any reason (e.g. a connection
+    issue), the main script will not be interrupted. Data is automatically pushed to the Hub every `commit_every`
+    minutes (default to every 5 minutes).
+
+    <Tip warning={true}>
+
+    `HFSummaryWriter` is experimental. Its API is subject to change in the future without prior notice.
+
+    </Tip>
+
+    Args:
+        repo_id (`str`):
+            The id of the repo to which the logs will be pushed.
+        logdir (`str`, *optional*):
+            The directory where the logs will be written. If not specified, a local directory will be created by the
+            underlying `SummaryWriter` object.
+        commit_every (`int` or `float`, *optional*):
+            The frequency (in minutes) at which the logs will be pushed to the Hub. Defaults to 5 minutes.
+        repo_type (`str`, *optional*):
+            The type of the repo to which the logs will be pushed. Defaults to "model".
+        repo_revision (`str`, *optional*):
+            The revision of the repo to which the logs will be pushed. Defaults to "main".
+        repo_private (`bool`, *optional*):
+            Whether to create a private repo or not. Defaults to False. This argument is ignored if the repo already
+            exists.
+        path_in_repo (`str`, *optional*):
+            The path to the folder in the repo where the logs will be pushed. Defaults to "tensorboard/".
+        repo_allow_patterns (`List[str]` or `str`, *optional*):
+            A list of patterns to include in the upload. Defaults to `"*.tfevents.*"`. Check out the
+            [upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#upload-a-folder) for more details.
+        repo_ignore_patterns (`List[str]` or `str`, *optional*):
+            A list of patterns to exclude in the upload. Check out the
+            [upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#upload-a-folder) for more details.
+        token (`str`, *optional*):
+            Authentication token. Will default to the stored token. See https://huggingface.co/settings/token for more
+            details
+        kwargs:
+            Additional keyword arguments passed to `SummaryWriter`.
+
+    Examples:
+    ```py
+    >>> from huggingface_hub import HFSummaryWriter
+
+    # Logs are automatically pushed every 15 minutes
+    >>> logger = HFSummaryWriter(repo_id="test_hf_logger", commit_every=15)
+    >>> logger.add_scalar("a", 1)
+    >>> logger.add_scalar("b", 2)
+    ...
+
+    # You can also trigger a push manually
+    >>> logger.scheduler.trigger()
+    ```
+
+    ```py
+    >>> from huggingface_hub import HFSummaryWriter
+
+    # Logs are automatically pushed every 5 minutes (default) + when exiting the context manager
+    >>> with HFSummaryWriter(repo_id="test_hf_logger") as logger:
+    ...     logger.add_scalar("a", 1)
+    ...     logger.add_scalar("b", 2)
+    ```
+    """
+
+    @experimental
+    def __new__(cls, *args, **kwargs) -> "HFSummaryWriter":
+        if not is_tensorboard_available():
+            raise ImportError(
+                "You must have `tensorboard` installed to use `HFSummaryWriter`. Please run `pip install --upgrade"
+                " tensorboardX` first."
+            )
+        return super().__new__(cls)
+
+    def __init__(
+        self,
+        repo_id: str,
+        *,
+        logdir: Optional[str] = None,
+        commit_every: Union[int, float] = 5,
+        repo_type: Optional[str] = None,
+        repo_revision: Optional[str] = None,
+        repo_private: bool = False,
+        path_in_repo: Optional[str] = "tensorboard",
+        repo_allow_patterns: Optional[Union[List[str], str]] = "*.tfevents.*",
+        repo_ignore_patterns: Optional[Union[List[str], str]] = None,
+        token: Optional[str] = None,
+        **kwargs,
+    ):
+        # Initialize SummaryWriter
+        super().__init__(logdir=logdir, **kwargs)
+
+        # Check logdir has been correctly initialized and fail early otherwise. In practice, SummaryWriter takes care of it.
+        if not isinstance(self.logdir, str):
+            raise ValueError(f"`self.logdir` must be a string. Got '{self.logdir}' of type {type(self.logdir)}.")
+
+        # Append logdir name to `path_in_repo`
+        if path_in_repo is None or path_in_repo == "":
+            path_in_repo = Path(self.logdir).name
+        else:
+            path_in_repo = path_in_repo.strip("/") + "/" + Path(self.logdir).name
+
+        # Initialize scheduler
+        self.scheduler = CommitScheduler(
+            folder_path=self.logdir,
+            path_in_repo=path_in_repo,
+            repo_id=repo_id,
+            repo_type=repo_type,
+            revision=repo_revision,
+            private=repo_private,
+            token=token,
+            allow_patterns=repo_allow_patterns,
+            ignore_patterns=repo_ignore_patterns,
+            every=commit_every,
+        )
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Push to hub in a non-blocking way when exiting the logger's context manager."""
+        super().__exit__(exc_type, exc_val, exc_tb)
+        future = self.scheduler.trigger()
+        future.result()

src/huggingface_hub/utils/__init__.py

@@ -52,7 +52,6 @@
     get_fastai_version,
     get_fastcore_version,
     get_gradio_version,
-    is_gradio_available,
     get_graphviz_version,
     get_hf_hub_version,
     get_hf_transfer_version,
@@ -61,18 +60,21 @@
     get_pillow_version,
     get_pydot_version,
     get_python_version,
+    get_tensorboard_version,
     get_tf_version,
     get_torch_version,
     is_fastai_available,
     is_fastcore_available,
     is_numpy_available,
     is_google_colab,
+    is_gradio_available,
     is_graphviz_available,
     is_hf_transfer_available,
     is_jinja_available,
     is_notebook,
     is_pillow_available,
     is_pydot_available,
+    is_tensorboard_available,
     is_tf_available,
     is_torch_available,
 )

src/huggingface_hub/utils/_experimental.py

@@ -48,14 +48,16 @@ def experimental(fn: Callable) -> Callable:
     Hello world!
     ```
     """
+    # For classes, put the "experimental" around the "__new__" method => __new__ will be removed in warning message
+    name = fn.__qualname__[: -len(".__new__")] if fn.__qualname__.endswith(".__new__") else fn.__qualname__
 
     @wraps(fn)
     def _inner_fn(*args, **kwargs):
         if not constants.HF_HUB_DISABLE_EXPERIMENTAL_WARNING:
             warnings.warn(
                 (
-                    f"'{fn.__name__}' is experimental and might be subject to breaking changes in the future. You can"
-                    " disable this warning by setting `HF_HUB_DISABLE_EXPERIMENTAL_WARNING=1` as environment"
+                    f"'{name}' is experimental and might be subject to breaking changes in the future."
+                    " You can disable this warning by setting `HF_HUB_DISABLE_EXPERIMENTAL_WARNING=1` as environment"
                     " variable."
                 ),
                 UserWarning,

src/huggingface_hub/utils/_runtime.py

@@ -42,6 +42,7 @@
     "numpy": {"numpy"},
     "pillow": {"Pillow"},
     "pydot": {"pydot"},
+    "tensorboard": {"tensorboardX"},
     "tensorflow": (
         "tensorflow",
         "tensorflow-cpu",
@@ -167,6 +168,15 @@ def get_pydot_version() -> str:
     return _get_version("pydot")
 
 
+# Tensorboard
+def is_tensorboard_available() -> bool:
+    return _is_available("tensorboard")
+
+
+def get_tensorboard_version() -> str:
+    return _get_version("tensorboard")
+
+
 # Tensorflow
 def is_tf_available() -> bool:
     return _is_available("tensorflow")
@@ -275,6 +285,7 @@ def dump_environment_info() -> Dict[str, Any]:
     info["Pillow"] = get_pillow_version()
     info["hf_transfer"] = get_hf_transfer_version()
     info["gradio"] = get_gradio_version()
+    info["tensorboard"] = get_tensorboard_version()
     info["numpy"] = get_numpy_version()
 
     # Environment variables